"""DMCCF+ : an improved version of the paper's two-stage model.

It keeps the paper's structure (a Criteria Ratings DNN, then an Overall Rating DNN that aggregates the
criteria) and changes how it is trained, guided by what the reproduction found:

  1. Cold users.  84 % of test ratings come from users with no training rating, whose embedding in the
     paper's model is an untrained random vector. Here every user without training ratings maps to one
     shared "unknown user" embedding, and during training each user ID is replaced by that token with
     probability `id_dropout`, so the token learns what an average new reviewer gives a hotel.
  2. Small embeddings, regularised.  16 instead of 64, with dropout and weight decay: most users have
     one rating, so a big free embedding only memorises it. (Per-criterion user/item biases were also
     tried and gave no gain.)
  3. Trainable networks.  He initialisation and a linear output instead of N(0, 0.05) and a ReLU output:
     no dead output units, so no restarts are needed.
  4. Stage alignment.  The paper trains the Overall DNN on true criteria but feeds it predicted ones,
     which are much smoother. After both are trained, the two stages are fine-tuned together on the
     overall rating (with the criteria loss kept as an auxiliary term).
  5. Predictions are clipped to the 1..5 rating range, and several seeds are averaged.

All selection was done on an inner split of fold 0's training data (see `dev_select.py`), never on test.
"""
from __future__ import annotations

import time
from dataclasses import dataclass

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

import deepmccf as D


@dataclass
class PlusConfig:
    emb: int = 16
    hidden: tuple = (128, 64, 32)
    overall_hidden: tuple = (64, 32, 16, 8)
    dropout: float = 0.1
    id_dropout: float = 0.3      # chance a training user ID is replaced by the unknown-user token
    biases: bool = False         # tried; no gain on the inner split, so left out
    weight_decay: float = 1e-4
    lr: float = 1e-3
    batch_size: int = 512
    max_epochs: int = 100
    patience: int = 8
    finetune: bool = True
    finetune_lr: float = 3e-4
    aux_weight: float = 0.5      # criteria loss weight during fine-tuning
    clip: bool = True
    loss: str = "l1"           # l1 (MAE), mse, or huber; applies to every stage
    seeds: int = 5               # ensemble size used by run_plus
    val_frac: float = 0.05       # early-stopping share of the outer training data (paper: 1 %)


class CriteriaNetPlus(nn.Module):
    def __init__(self, n_users, n_items, k, cfg: PlusConfig):
        super().__init__()
        self.unk = n_users                      # index of the unknown-user token
        self.user_emb = nn.Embedding(n_users + 1, cfg.emb)
        self.item_emb = nn.Embedding(n_items, cfg.emb)
        layers, d = [], 2 * cfg.emb
        for h in cfg.hidden:
            layers += [nn.Linear(d, h), nn.ReLU(), nn.Dropout(cfg.dropout)]
            d = h
        self.mlp = nn.Sequential(*layers)
        self.out = nn.Linear(d, k)
        self.biases = cfg.biases
        if cfg.biases:
            self.user_bias = nn.Embedding(n_users + 1, k)
            self.item_bias = nn.Embedding(n_items, k)
            nn.init.zeros_(self.user_bias.weight); nn.init.zeros_(self.item_bias.weight)
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity="relu"); nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Embedding) and m.embedding_dim == cfg.emb:
                nn.init.normal_(m.weight, 0, 0.1)

    def forward(self, u, i):                    # standardised criteria
        y = self.out(self.mlp(torch.cat([self.user_emb(u), self.item_emb(i)], 1)))
        if self.biases:
            y = y + self.user_bias(u) + self.item_bias(i)
        return y


class OverallNetPlus(nn.Module):
    def __init__(self, k, hidden):
        super().__init__()
        layers, d = [], k
        for h in hidden:
            layers += [nn.Linear(d, h), nn.ReLU()]
            d = h
        self.mlp = nn.Sequential(*layers)
        self.out = nn.Linear(d, 1)
        self.skip = nn.Linear(k, 1)             # linear path: overall is close to a weighted mean of criteria
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity="relu"); nn.init.zeros_(m.bias)

    def forward(self, z):                       # z: standardised criteria -> standardised overall
        return (self.out(self.mlp(z)) + self.skip(z)).squeeze(1)


def _loss(diff, kind):
    if kind == "mse":
        return (diff ** 2).mean()
    if kind == "huber":
        return nn.functional.huber_loss(diff, torch.zeros_like(diff), delta=0.5)
    return diff.abs().mean()


def _train(params, loss_of_batch, n, val_loss, cfg, lr, seed, modules):
    opt = torch.optim.AdamW(params, lr=lr, weight_decay=cfg.weight_decay)
    gen = torch.Generator().manual_seed(seed)
    best, best_state, wait = float("inf"), None, 0
    for _ in range(cfg.max_epochs):
        for m in modules: m.train()
        for b in torch.randperm(n, generator=gen).split(cfg.batch_size):
            opt.zero_grad(); loss_of_batch(b).backward(); opt.step()
        for m in modules: m.eval()
        with torch.no_grad():
            v = val_loss()
        if v < best - 1e-5:
            best, wait = v, 0
            best_state = [{k: t.clone() for k, t in m.state_dict().items()} for m in modules]
        else:
            wait += 1
            if wait >= cfg.patience:
                break
    for m, s in zip(modules, best_state): m.load_state_dict(s)
    return best


def fit_predict(df, train, val, test, cfg: PlusConfig = PlusConfig(), seed: int = 0) -> dict:
    """Train DMCCF+ on `train` (early stopping on `val`) and predict the overall rating for `test`."""
    torch.manual_seed(seed)
    C = df[D.CRITERIA].to_numpy(np.float32)
    y = df["overall"].to_numpy(np.float32)
    cm, cs = C[train].mean(0), C[train].std(0)
    ym, ys = y[train].mean(), y[train].std()
    n_users, n_items, k = df.u.max() + 1, df.i.max() + 1, len(D.CRITERIA)

    known = np.zeros(n_users + 1, bool); known[df.u.values[train]] = True
    uidx = df.u.values.copy(); uidx[~known[uidx]] = n_users          # unseen users -> unknown token
    U = torch.as_tensor(uidx, dtype=torch.long); I = torch.as_tensor(df.i.values, dtype=torch.long)
    Cz = torch.as_tensor((C - cm) / cs); Yz = torch.as_tensor((y - ym) / ys)
    tr, va, te = (torch.as_tensor(np.asarray(a)) for a in (train, val, test))

    crit = CriteriaNetPlus(n_users, n_items, k, cfg)
    over = OverallNetPlus(k, cfg.overall_hidden)

    def drop_ids(u):
        if cfg.id_dropout <= 0 or not crit.training:
            return u
        mask = torch.rand(len(u)) < cfg.id_dropout
        return torch.where(mask, torch.full_like(u, n_users), u)

    # Stage 1: criteria network (MAE on standardised criteria)
    _train(crit.parameters(),
           lambda b: _loss(crit(drop_ids(U[tr[b]]), I[tr[b]]) - Cz[tr[b]], cfg.loss), len(tr),
           lambda: _loss(crit(U[va], I[va]) - Cz[va], cfg.loss).item(), cfg, cfg.lr, seed, [crit])
    # Stage 2: aggregation network on TRUE criteria, as in the paper
    _train(over.parameters(),
           lambda b: _loss(over(Cz[tr[b]]) - Yz[tr[b]], cfg.loss), len(tr),
           lambda: _loss(over(Cz[va]) - Yz[va], cfg.loss).item(), cfg, cfg.lr, seed + 1, [over])
    # Stage 3: fine-tune both on the overall rating, fed with PREDICTED criteria
    if cfg.finetune:
        def joint(b):
            cp = crit(drop_ids(U[tr[b]]), I[tr[b]])
            return _loss(over(cp) - Yz[tr[b]], cfg.loss) + cfg.aux_weight * _loss(cp - Cz[tr[b]], cfg.loss)
        _train(list(crit.parameters()) + list(over.parameters()), joint, len(tr),
               lambda: _loss(over(crit(U[va], I[va])) - Yz[va], cfg.loss).item(), cfg, cfg.finetune_lr, seed + 2,
               [crit, over])

    crit.eval(); over.eval()
    with torch.no_grad():
        cp = crit(U[te], I[te])
        pred = over(cp).numpy() * ys + ym
        crit_pred = cp.numpy() * cs + cm
    if cfg.clip:
        pred, crit_pred = np.clip(pred, 1, 5), np.clip(crit_pred, 1, 5)
    return {"pred": pred, "criteria_pred": crit_pred}


def run_plus(df, fold: D.Fold, cfg: PlusConfig = PlusConfig()) -> dict:
    """Seed ensemble on one outer fold: train on the fold's 80 % (5 % of it held out for early stopping),
    predict the fold's 20 % test set."""
    t0 = time.time()
    pool = np.random.default_rng(fold.k).permutation(fold.train_full)   # larger validation set than 1 %
    n_va = int(cfg.val_frac * len(pool))
    tr, va = np.sort(pool[n_va:]), np.sort(pool[:n_va])
    outs = [fit_predict(df, tr, va, fold.test, cfg, seed=s) for s in range(cfg.seeds)]
    pred = np.mean([o["pred"] for o in outs], 0)
    cp = np.mean([o["criteria_pred"] for o in outs], 0)
    crit_true = df[D.CRITERIA].to_numpy(np.float32)[fold.test]
    return {"pred": pred, "criteria_mae": np.abs(cp - crit_true).mean(0), "seconds": time.time() - t0}