#!/usr/bin/env python3
"""Zoo mô hình mở rộng cho EXP3 — nhiều họ hơn và nhiều biến thể ANN/CNN hơn.

Ý tưởng người dùng: tầng 1 của EXP3 chỉ cần trả lời MỘT câu hỏi -- "mẫu này gần
như chắc chắn nEHDE?" -- và câu trả lời đó là p = M/N.  Với N nhỏ, p chỉ nhận vài
giá trị rời rạc nên cổng rất thô; tăng N bằng những họ mô hình THỰC SỰ khác nhau
làm p mịn hơn và, quan trọng hơn, làm sai số của các thành viên bớt tương quan.

Hai nhóm thành viên có bản chất đầu vào khác nhau:
  * "tab" -- học trên vector đặc trưng đã gộp tuần (trung bình/cực đại trong tuần);
  * "seq" -- CNN 1-D chạy dọc CHUỖI LEAD DAY trong tuần, tức nhìn thấy diễn biến
    ngày-qua-ngày mà bản gộp trung bình/cực đại đã xoá mất.  Đây là lý do đáng kể
    nhất để thêm CNN vào bài toán này chứ không phải vì nó "mạnh hơn".

Mọi họ đều hồi quy Sev_week (liên tục, có dấu) đúng đặc tả EXP3; mỗi họ sinh 9
thành viên theo mức phân vị của phần dư.
"""

from __future__ import annotations

import numpy as np


# --------------------------------------------------------------- torch ANN
class TorchANN:
    """Mạng truyền thẳng cho hồi quy Sev, cấu hình được theo lớp/node/activation."""

    def __init__(self, hidden=(64, 32), act="relu", dropout=0.0, epochs=150,
                 lr=1e-3, wd=1e-4, batch=256, seed=0):
        self.hidden, self.act, self.dropout = hidden, act, dropout
        self.epochs, self.lr, self.wd, self.batch, self.seed = \
            epochs, lr, wd, batch, seed

    def _act(self):
        import torch.nn as nn
        return {"relu": nn.ReLU, "tanh": nn.Tanh, "gelu": nn.GELU,
                "elu": nn.ELU, "silu": nn.SiLU}[self.act]()

    def fit(self, X, y):
        import torch
        import torch.nn as nn
        torch.set_num_threads(1)
        torch.manual_seed(self.seed)
        X = np.asarray(X, dtype="f4"); y = np.asarray(y, dtype="f4").ravel()
        self.mu_, self.sd_ = X.mean(0), X.std(0)
        self.sd_[self.sd_ < 1e-6] = 1.0
        self.ym_, self.ys_ = float(y.mean()), float(y.std()) or 1.0

        layers, prev = [], X.shape[1]
        for h in self.hidden:
            layers += [nn.Linear(prev, h), self._act()]
            if self.dropout > 0:
                layers.append(nn.Dropout(self.dropout))
            prev = h
        layers.append(nn.Linear(prev, 1))
        net = nn.Sequential(*layers)
        opt = torch.optim.Adam(net.parameters(), lr=self.lr, weight_decay=self.wd)
        lossf = nn.MSELoss()

        Xt = torch.from_numpy((X - self.mu_) / self.sd_)
        yt = torch.from_numpy((y - self.ym_) / self.ys_).unsqueeze(1)
        n = X.shape[0]
        rng = np.random.default_rng(self.seed)
        net.train()
        for _ in range(self.epochs):
            order = rng.permutation(n)
            for a in range(0, n, self.batch):
                idx = torch.from_numpy(order[a:a + self.batch])
                opt.zero_grad()
                lossf(net(Xt[idx]), yt[idx]).backward()
                opt.step()
        net.eval()
        self.net_ = net
        return self

    def predict(self, X):
        import torch
        X = np.asarray(X, dtype="f4")
        with torch.no_grad():
            z = self.net_(torch.from_numpy((X - self.mu_) / self.sd_))
        return z.numpy().ravel() * self.ys_ + self.ym_


# ---------------------------------------------------------------- torch CNN
class TorchCNN1D:
    """CNN 1-D dọc trục lead day trong tuần: đầu vào (n, T ngày, F biến/ngày).

    Trung bình và cực đại tuần bóp cả tuần thành hai con số; CNN giữ được thứ tự
    các ngày, nên phân biệt được một tuần nóng dần lên với một tuần nóng rồi nguội
    -- hai hình thái có cùng trung bình và cùng cực đại nhưng rủi ro rất khác nhau.
    """

    def __init__(self, chan=(32, 32), kernel=3, dense=32, act="relu", dropout=0.1,
                 epochs=120, lr=1e-3, wd=1e-4, batch=256, seed=0):
        self.chan, self.kernel, self.dense = chan, kernel, dense
        self.act, self.dropout = act, dropout
        self.epochs, self.lr, self.wd, self.batch, self.seed = \
            epochs, lr, wd, batch, seed

    def _act(self):
        import torch.nn as nn
        return {"relu": nn.ReLU, "tanh": nn.Tanh, "gelu": nn.GELU}[self.act]()

    def fit(self, X, y):
        import torch
        import torch.nn as nn
        torch.set_num_threads(1)
        torch.manual_seed(self.seed)
        X = np.asarray(X, dtype="f4")                      # (n, T, F)
        y = np.asarray(y, dtype="f4").ravel()
        self.mu_ = X.mean(axis=(0, 1)); self.sd_ = X.std(axis=(0, 1))
        self.sd_[self.sd_ < 1e-6] = 1.0
        self.ym_, self.ys_ = float(y.mean()), float(y.std()) or 1.0
        nT, nF = X.shape[1], X.shape[2]

        class Net(nn.Module):
            def __init__(s, chan, k, dense, actf, drop):
                super().__init__()
                blocks, prev = [], nF
                for c in chan:
                    pad = k // 2
                    blocks += [nn.Conv1d(prev, c, k, padding=pad), actf()]
                    prev = c
                s.conv = nn.Sequential(*blocks)
                s.head = nn.Sequential(nn.Linear(2 * prev, dense), actf(),
                                       nn.Dropout(drop), nn.Linear(dense, 1))

            def forward(s, x):                              # x: (n, F, T)
                h = s.conv(x)
                return s.head(torch.cat([h.mean(dim=2), h.amax(dim=2)], dim=1))

        actf = type(self._act())
        net = Net(self.chan, self.kernel, self.dense, actf, self.dropout)
        opt = torch.optim.Adam(net.parameters(), lr=self.lr, weight_decay=self.wd)
        lossf = nn.MSELoss()

        Xs = ((X - self.mu_) / self.sd_).transpose(0, 2, 1).copy()   # (n, F, T)
        Xt = torch.from_numpy(Xs)
        yt = torch.from_numpy((y - self.ym_) / self.ys_).unsqueeze(1)
        n = X.shape[0]
        rng = np.random.default_rng(self.seed)
        net.train()
        for _ in range(self.epochs):
            order = rng.permutation(n)
            for a in range(0, n, self.batch):
                idx = torch.from_numpy(order[a:a + self.batch])
                opt.zero_grad()
                lossf(net(Xt[idx]), yt[idx]).backward()
                opt.step()
        net.eval()
        self.net_ = net
        self.nT_, self.nF_ = nT, nF
        return self

    def predict(self, X):
        import torch
        X = np.asarray(X, dtype="f4")
        Xs = ((X - self.mu_) / self.sd_).transpose(0, 2, 1).copy()
        with torch.no_grad():
            z = self.net_(torch.from_numpy(Xs))
        return z.numpy().ravel() * self.ys_ + self.ym_


# ------------------------------------------------------------------- danh mục
def zoo_base():
    """Bảy họ của bản EXP3 gốc — giữ nguyên để đối chiếu công bằng."""
    from sklearn.linear_model import Ridge, ElasticNet
    from sklearn.ensemble import ExtraTreesRegressor
    from sklearn.neighbors import KNeighborsRegressor
    from sklearn.neural_network import MLPRegressor
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler
    import lightgbm as lgb
    return [
        ("Ridge", "tab", "hồi quy tuyến tính, phạt L2",
         lambda: make_pipeline(StandardScaler(), Ridge(alpha=10.0))),
        ("ElasticNet", "tab", "hồi quy tuyến tính thưa, phạt L1+L2",
         lambda: make_pipeline(StandardScaler(),
                               ElasticNet(alpha=0.01, l1_ratio=0.5, max_iter=2000))),
        ("LGBM_gbdt", "tab", "cây tăng cường độ dốc",
         lambda: lgb.LGBMRegressor(n_estimators=300, num_leaves=15,
                                   learning_rate=0.05, min_child_samples=100,
                                   subsample=0.8, subsample_freq=1,
                                   colsample_bytree=0.8, reg_lambda=5.0,
                                   n_jobs=1, verbose=-1)),
        ("LGBM_rf", "tab", "rừng ngẫu nhiên dựng bằng LightGBM",
         lambda: lgb.LGBMRegressor(boosting_type="rf", n_estimators=200,
                                   num_leaves=63, subsample=0.7, subsample_freq=1,
                                   colsample_bytree=0.7, min_child_samples=50,
                                   n_jobs=1, verbose=-1)),
        ("ExtraTrees", "tab", "cây cực kỳ ngẫu nhiên",
         lambda: ExtraTreesRegressor(n_estimators=100, min_samples_leaf=20,
                                     max_features=0.5, n_jobs=1, random_state=0)),
        ("kNN25", "tab", "láng giềng gần nhất, k = 25",
         lambda: make_pipeline(StandardScaler(),
                               KNeighborsRegressor(n_neighbors=25, n_jobs=1))),
        ("MLP_sk_64_32", "tab", "MLP sklearn 64-32, relu",
         lambda: make_pipeline(StandardScaler(),
                               MLPRegressor(hidden_layer_sizes=(64, 32), max_iter=60,
                                            early_stopping=True, random_state=0))),
    ]


def zoo_extra():
    """Các họ THÊM VÀO: tuyến tính bền vững, cây khác cơ chế, ANN nhiều cấu hình,
    CNN 1-D dọc chuỗi lead day."""
    from sklearn.linear_model import BayesianRidge, HuberRegressor
    from sklearn.cross_decomposition import PLSRegression
    from sklearn.ensemble import (RandomForestRegressor,
                                  HistGradientBoostingRegressor)
    from sklearn.neighbors import KNeighborsRegressor
    from sklearn.neural_network import MLPRegressor
    from sklearn.svm import SVR
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler
    import lightgbm as lgb

    def pls():
        class _P(PLSRegression):
            def predict(self, X):
                return super().predict(X).ravel()
        return make_pipeline(StandardScaler(), _P(n_components=8))

    return [
        # -- tuyến tính khác hàm mất mát / khác cách chính quy hoá
        ("BayesRidge", "tab", "hồi quy tuyến tính Bayes, chính quy hoá tự động",
         lambda: make_pipeline(StandardScaler(), BayesianRidge())),
        ("Huber", "tab", "hồi quy tuyến tính bền vững với ngoại lai (mất mát Huber)",
         lambda: make_pipeline(StandardScaler(),
                               HuberRegressor(alpha=1e-3, max_iter=300))),
        ("PLS8", "tab", "bình phương tối thiểu riêng phần, 8 thành phần",
         pls),
        # -- cây, ba cơ chế tăng cường khác nhau
        ("LGBM_dart", "tab", "tăng cường độ dốc kiểu DART (bỏ ngẫu nhiên cây)",
         lambda: lgb.LGBMRegressor(boosting_type="dart", n_estimators=250,
                                   num_leaves=31, learning_rate=0.05,
                                   min_child_samples=40, colsample_bytree=0.8,
                                   n_jobs=1, verbose=-1)),
        ("HistGB", "tab", "tăng cường độ dốc theo histogram (sklearn)",
         lambda: HistGradientBoostingRegressor(max_iter=250, max_leaf_nodes=15,
                                               learning_rate=0.06,
                                               min_samples_leaf=40,
                                               l2_regularization=1.0,
                                               random_state=0)),
        ("RandomForest", "tab", "rừng ngẫu nhiên CART đầy đủ",
         lambda: RandomForestRegressor(n_estimators=200, min_samples_leaf=10,
                                       max_features=0.4, n_jobs=1, random_state=0)),
        # -- phi tham số / hạt nhân
        ("kNN10", "tab", "láng giềng gần nhất, k = 10",
         lambda: make_pipeline(StandardScaler(),
                               KNeighborsRegressor(n_neighbors=10, n_jobs=1))),
        ("kNN60", "tab", "láng giềng gần nhất, k = 60",
         lambda: make_pipeline(StandardScaler(),
                               KNeighborsRegressor(n_neighbors=60, n_jobs=1))),
        ("SVR_rbf", "tab", "máy vector hỗ trợ, hạt nhân RBF",
         lambda: make_pipeline(StandardScaler(),
                               SVR(kernel="rbf", C=3.0, gamma="scale",
                                   epsilon=0.2))),
        # -- ANN sklearn, hai cấu hình khác activation
        ("MLP_sk_128_64_tanh", "tab", "MLP sklearn 128-64, tanh",
         lambda: make_pipeline(StandardScaler(),
                               MLPRegressor(hidden_layer_sizes=(128, 64),
                                            activation="tanh", max_iter=80,
                                            early_stopping=True, random_state=1))),
        ("MLP_sk_256_lbfgs", "tab", "MLP sklearn 256 nút, tối ưu L-BFGS",
         lambda: make_pipeline(StandardScaler(),
                               MLPRegressor(hidden_layer_sizes=(256,),
                                            solver="lbfgs", alpha=1e-2,
                                            max_iter=400, random_state=2))),
        # -- ANN torch: quét số lớp, số nút, hàm kích hoạt, dropout
        ("ANN_32_relu", "tab", "ANN torch 1 lớp ẩn 32 nút, relu",
         lambda: TorchANN(hidden=(32,), act="relu", epochs=150, seed=0)),
        ("ANN_64_32_relu", "tab", "ANN torch 2 lớp 64-32, relu",
         lambda: TorchANN(hidden=(64, 32), act="relu", epochs=150, seed=1)),
        ("ANN_128_64_32_relu", "tab", "ANN torch 3 lớp 128-64-32, relu, dropout 0.1",
         lambda: TorchANN(hidden=(128, 64, 32), act="relu", dropout=0.1,
                          epochs=150, seed=2)),
        ("ANN_64_64_tanh", "tab", "ANN torch 2 lớp 64-64, tanh",
         lambda: TorchANN(hidden=(64, 64), act="tanh", epochs=150, seed=3)),
        ("ANN_256_128_gelu", "tab", "ANN torch 2 lớp 256-128, gelu, dropout 0.2",
         lambda: TorchANN(hidden=(256, 128), act="gelu", dropout=0.2,
                          epochs=150, seed=4)),
        ("ANN_96_48_elu", "tab", "ANN torch 2 lớp 96-48, elu, phạt trọng số mạnh",
         lambda: TorchANN(hidden=(96, 48), act="elu", wd=1e-3, epochs=150, seed=5)),
        ("ANN_512_silu", "tab", "ANN torch 1 lớp rộng 512 nút, silu, dropout 0.3",
         lambda: TorchANN(hidden=(512,), act="silu", dropout=0.3, epochs=150,
                          seed=6)),
        # -- CNN 1-D dọc chuỗi lead day trong tuần
        ("CNN_32_32_k3", "seq", "CNN 1-D 2 tầng 32-32 kênh, kernel 3, relu",
         lambda: TorchCNN1D(chan=(32, 32), kernel=3, dense=32, act="relu",
                            epochs=120, seed=7)),
        ("CNN_64_32_k3", "seq", "CNN 1-D 2 tầng 64-32 kênh, kernel 3, gelu",
         lambda: TorchCNN1D(chan=(64, 32), kernel=3, dense=64, act="gelu",
                            epochs=120, seed=8)),
        ("CNN_48_k5", "seq", "CNN 1-D 1 tầng 48 kênh, kernel 5 (cửa sổ rộng hơn)",
         lambda: TorchCNN1D(chan=(48,), kernel=5, dense=32, act="relu",
                            epochs=120, seed=9)),
    ]


def zoo_full():
    return zoo_base() + zoo_extra()