# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Task A neural models (PyTorch). Never import lightgbm / xgboost in the same process (OpenMP conflict on macOS)."""
from __future__ import annotations

import os

import numpy as np
import torch
import torch.nn as nn

from wp23.models_a import TAUS

torch.set_num_threads(int(os.environ.get("WP23_TORCH_THREADS", "6")))


def _pinball(pred, y, taus):
    d = y[:, None] - pred; return torch.mean(torch.maximum(taus * d, (taus - 1) * d))


class MLP:
    """Multilayer perceptron with a quantile head (pinball loss on 5 quantiles; the median is the point forecast)."""
    name = "mlp"
    def __init__(self, hidden=(256, 128), epochs=8, lr=2e-3, batch=4096, seed=0):
        self.hidden, self.epochs, self.lr, self.batch, self.seed = hidden, epochs, lr, batch, seed
    def _net(self, d):
        layers = []; last = d
        for h in self.hidden:
            layers += [nn.Linear(last, h), nn.GELU(), nn.Dropout(0.1)]; last = h
        layers += [nn.Linear(last, len(TAUS))]; return nn.Sequential(*layers)
    def fit(self, X, y, meta=None):
        torch.manual_seed(self.seed); self.mu = X.mean(0); self.sd = X.std(0) + 1e-9
        Xt = torch.tensor((X - self.mu) / self.sd, dtype=torch.float32); yt = torch.tensor(y, dtype=torch.float32); taus = torch.tensor(TAUS, dtype=torch.float32)
        self.net = self._net(X.shape[1]); opt = torch.optim.AdamW(self.net.parameters(), lr=self.lr, weight_decay=1e-4)
        sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=self.lr, total_steps=self.epochs * ((len(Xt) + self.batch - 1) // self.batch))
        self.net.train()
        for ep in range(self.epochs):
            perm = torch.randperm(len(Xt))
            for i in range(0, len(Xt), self.batch):
                b = perm[i:i + self.batch]; opt.zero_grad(); loss = _pinball(self.net(Xt[b]), yt[b], taus); loss.backward(); opt.step(); sched.step()
        self.net.eval(); return self
    def predict_quantiles(self, X, meta=None, taus=TAUS):
        with torch.no_grad():
            Q = self.net(torch.tensor((X - self.mu) / self.sd, dtype=torch.float32)).numpy()
        return np.sort(Q, axis=1)
    def predict(self, X, meta=None):
        return self.predict_quantiles(X, meta)[:, 2]


class HSHN(MLP):
    """Hierarchical spatiotemporal hedonic network: the MLP of the hedonic + market features plus learned embeddings of the
    municipality, MRC and region (nested: the municipality embedding is a *deviation* penalised toward zero, so thin
    municipalities fall back on their MRC and region — partial pooling by weight decay), a horizon embedding, and an
    interaction block between the horizon/market state and the hedonic representation. pooling ∈ {partial, full, none}."""
    name = "hshn"
    def __init__(self, pooling="partial", emb=8, hidden=(256, 128), epochs=8, lr=2e-3, batch=4096, seed=0, mun_decay=5e-3):
        super().__init__(hidden, epochs, lr, batch, seed); self.pooling, self.emb, self.mun_decay = pooling, emb, mun_decay
        self.name = f"hshn_{pooling}"
    def _codes(self, meta, fit=False):
        out = []
        for k in ("mun_code", "mrc_code", "region_code"):
            v = np.asarray(meta[k]).astype(str)
            if fit:
                self.maps = getattr(self, "maps", {}); self.maps[k] = {c: i + 1 for i, c in enumerate(sorted(set(v)))}
            out.append(torch.tensor([self.maps[k].get(c, 0) for c in v]))
        h = torch.tensor(np.clip(np.asarray(meta["h"]).astype(int), 0, 12)); return out + [h]
    def fit(self, X, y, meta=None):
        torch.manual_seed(self.seed); self.mu = X.mean(0); self.sd = X.std(0) + 1e-9
        Xt = torch.tensor((X - self.mu) / self.sd, dtype=torch.float32); yt = torch.tensor(y, dtype=torch.float32); taus = torch.tensor(TAUS, dtype=torch.float32)
        codes = self._codes(meta, fit=True)
        n_mun, n_mrc, n_reg = (len(self.maps["mun_code"]) + 1, len(self.maps["mrc_code"]) + 1, len(self.maps["region_code"]) + 1)
        e = self.emb
        self.emb_mun = nn.Embedding(n_mun, e); self.emb_mrc = nn.Embedding(n_mrc, e); self.emb_reg = nn.Embedding(n_reg, e); self.emb_h = nn.Embedding(13, 4)
        for em in (self.emb_mun, self.emb_mrc, self.emb_reg, self.emb_h):
            nn.init.normal_(em.weight, std=0.02)
        d_in = X.shape[1] + e + 4
        self.enc = nn.Sequential(nn.Linear(d_in, self.hidden[0]), nn.GELU(), nn.Dropout(0.1), nn.Linear(self.hidden[0], self.hidden[1]), nn.GELU())
        self.inter = nn.Sequential(nn.Linear(self.hidden[1] + e + 4, 64), nn.GELU())
        self.head = nn.Linear(self.hidden[1] + 64, len(TAUS))
        params = [p for m in (self.emb_mrc, self.emb_reg, self.emb_h, self.enc, self.inter, self.head) for p in m.parameters()]
        groups = [{"params": params, "weight_decay": 1e-4}]
        if self.pooling != "full":
            groups.append({"params": self.emb_mun.parameters(), "weight_decay": self.mun_decay if self.pooling == "partial" else 0.0})
        opt = torch.optim.AdamW(groups, lr=self.lr)
        steps = self.epochs * ((len(Xt) + self.batch - 1) // self.batch); sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=self.lr, total_steps=steps)
        for ep in range(self.epochs):
            perm = torch.randperm(len(Xt))
            for i in range(0, len(Xt), self.batch):
                b = perm[i:i + self.batch]; opt.zero_grad(); loss = _pinball(self._forward(Xt[b], [c[b] for c in codes]), yt[b], taus); loss.backward(); opt.step(); sched.step()
        return self
    def _forward(self, x, codes):
        mun, mrc, reg, h = codes
        g = self.emb_reg(reg) + self.emb_mrc(mrc) + (self.emb_mun(mun) if self.pooling != "full" else 0.0)
        hv = self.emb_h(h); z = self.enc(torch.cat([x, g, hv], 1)); it = self.inter(torch.cat([z, g, hv], 1))
        return self.head(torch.cat([z, it], 1))
    def predict_quantiles(self, X, meta=None, taus=TAUS):
        with torch.no_grad():
            for m in (self.enc, self.inter, self.head):
                m.eval()
            Q = self._forward(torch.tensor((X - self.mu) / self.sd, dtype=torch.float32), self._codes(meta)).numpy()
        return np.sort(Q, axis=1)