"""Deep multi-criteria collaborative filtering: Nassar, Jafar & Rahhal, Knowledge-Based Systems 187 (2020) 104811.

Everything the experiments need: data split, the two DNNs of the paper, the single-DNN baseline,
the Surprise baselines, and the five evaluation metrics (MAE, F-beta, FCP, MAP, MRR).

The paper used Keras 2 / TensorFlow; this is a PyTorch port of the same settings (Section 4.3):
  Criteria Ratings DNN : user & item embeddings (64 each) -> concatenate -> [128, 64, 32, 16, 8] ReLU
                         -> 7 ReLU outputs (one per criterion)
  Overall Rating DNN   : z-scored criteria (training mean/std) -> [64, 32, 16, 8] ReLU -> 1 ReLU output
  Both                 : weights ~ N(0, 0.05), Adam (lr 1e-3, default betas), batch 512, MAE loss,
                         trained separately ("without knowing each other").
  Single DNN baseline  : embeddings 256 -> [512, 256, 128, 64, 32, 16, 8] ReLU -> 1 ReLU output.
  Evaluation           : 5-fold CV; each fold: 80 % train (of which 1 % validation), 20 % test.
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field, replace
from typing import Callable

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.model_selection import KFold

CRITERIA = ["value", "rooms", "location", "cleanliness", "check_in", "service", "business_service"]
CRITERIA_LABELS = ["Value", "Rooms", "Location", "Cleanliness", "Check in / front desk", "Service", "Business service"]


# --------------------------------------------------------------------------------------------------
# Data
# --------------------------------------------------------------------------------------------------
@dataclass
class Fold:
    k: int
    train: np.ndarray  # row indices used to fit the DNNs (79 % of the data)
    val: np.ndarray    # 1 % validation set, for early stopping
    test: np.ndarray   # 20 % test set

    @property
    def train_full(self) -> np.ndarray:  # the full 80 %, used by the non-DNN baselines
        return np.concatenate([self.train, self.val])


def load_ratings(path) -> pd.DataFrame:
    df = pd.read_csv(path, dtype={"user": str, "hotel": str})
    df["u"] = pd.factorize(df["user"])[0]
    df["i"] = pd.factorize(df["hotel"])[0]
    return df


def make_folds(n: int, n_splits: int = 5, val_frac: float = 0.01, seed: int = 42) -> list[Fold]:
    """Random 5-fold split. val_frac is a share of the whole data set (paper: 80 % train incl. 1 % val)."""
    rng = np.random.default_rng(seed)
    folds = []
    for k, (tr, te) in enumerate(KFold(n_splits, shuffle=True, random_state=seed).split(np.arange(n))):
        tr = rng.permutation(tr)
        n_val = int(round(val_frac * n))
        folds.append(Fold(k, np.sort(tr[n_val:]), np.sort(tr[:n_val]), np.sort(te)))
    return folds


# --------------------------------------------------------------------------------------------------
# Models
# --------------------------------------------------------------------------------------------------
def _init_normal(module: nn.Module, std: float) -> None:
    """Paper: 'randomly initialized the DNN parameters using a normal distribution with mean 0 and
    standard deviation 0.05'. Applied to embeddings and dense kernels; biases start at 0 (Keras default)."""
    for m in module.modules():
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, 0.0, std)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.zeros_(m.bias)


def _mlp(sizes: list[int]) -> nn.Sequential:
    layers = []
    for a, b in zip(sizes[:-1], sizes[1:]):
        layers += [nn.Linear(a, b), nn.ReLU()]
    return nn.Sequential(*layers)


class CriteriaRatingsDNN(nn.Module):
    """Eqs. (1)-(5): x = concat(v_u, v_i); h_l = ReLU(W_l h_{l-1} + b_l); y_ui = ReLU(W_L h_{L-1} + b_L)."""

    def __init__(self, n_users, n_items, n_criteria=7, emb=64, hidden=(128, 64, 32, 16, 8), init_std=0.05):
        super().__init__()
        self.user_emb = nn.Embedding(n_users, emb)
        self.item_emb = nn.Embedding(n_items, emb)
        self.mlp = _mlp([2 * emb, *hidden])
        self.out = nn.Linear(hidden[-1], n_criteria)
        _init_normal(self, init_std)

    def forward(self, u, i):
        x = torch.cat([self.user_emb(u), self.item_emb(i)], dim=1)
        return torch.relu(self.out(self.mlp(x)))


class OverallRatingDNN(nn.Module):
    """Eqs. (6)-(9): input = z-scored criteria ratings, dense ReLU layers, 1 ReLU output (the overall rating)."""

    def __init__(self, n_criteria=7, hidden=(64, 32, 16, 8), init_std=0.05):
        super().__init__()
        self.mlp = _mlp([n_criteria, *hidden])
        self.out = nn.Linear(hidden[-1], 1)
        _init_normal(self, init_std)

    def forward(self, z):
        return torch.relu(self.out(self.mlp(z))).squeeze(1)


class SingleDNN(nn.Module):
    """The paper's single-rating DNN baseline: predicts the overall rating straight from the IDs."""

    def __init__(self, n_users, n_items, emb=256, hidden=(512, 256, 128, 64, 32, 16, 8), init_std=0.05):
        super().__init__()
        self.user_emb = nn.Embedding(n_users, emb)
        self.item_emb = nn.Embedding(n_items, emb)
        self.mlp = _mlp([2 * emb, *hidden])
        self.out = nn.Linear(hidden[-1], 1)
        _init_normal(self, init_std)

    def forward(self, u, i):
        x = torch.cat([self.user_emb(u), self.item_emb(i)], dim=1)
        return torch.relu(self.out(self.mlp(x))).squeeze(1)


# --------------------------------------------------------------------------------------------------
# Training
# --------------------------------------------------------------------------------------------------
@dataclass
class TrainConfig:
    lr: float | None = 1e-3      # None = the optimizer's Keras 2 default
    batch_size: int = 512
    optimizer: str = "adam"      # Table 3: adadelta, adagrad, adam, adamax, rmsprop, sgd
    max_epochs: int = 100        # not given in the paper
    patience: int = 10           # early stopping on the 1 % validation MAE (not given in the paper)
    max_restarts: int = 5        # re-initialise if the ReLU output layer is dead (constant 0)
    seed: int = 0


# Keras 2 default learning rates, used by the optimizer comparison (Table 3) when TrainConfig.lr is None.
KERAS2_DEFAULT_LR = {"adam": 1e-3, "adamax": 2e-3, "adagrad": 1e-2, "adadelta": 1.0, "rmsprop": 1e-3, "sgd": 1e-2}
OPTIMIZERS: dict[str, Callable] = {  # Keras 2 defaults for the other optimizer arguments
    "adam": lambda p, lr: torch.optim.Adam(p, lr=lr),
    "adamax": lambda p, lr: torch.optim.Adamax(p, lr=lr),
    "adagrad": lambda p, lr: torch.optim.Adagrad(p, lr=lr, initial_accumulator_value=0.1),
    "adadelta": lambda p, lr: torch.optim.Adadelta(p, lr=lr, rho=0.95, eps=1e-7),
    "rmsprop": lambda p, lr: torch.optim.RMSprop(p, lr=lr, alpha=0.9, eps=1e-7),
    "sgd": lambda p, lr: torch.optim.SGD(p, lr=lr),
}


def _batches(n, bs, gen):
    idx = torch.randperm(n, generator=gen)
    for s in range(0, n, bs):
        yield idx[s:s + bs]


def fit(make_model: Callable[[], nn.Module], X_tr: tuple, y_tr: torch.Tensor, X_va: tuple, y_va: torch.Tensor,
        cfg: TrainConfig) -> tuple[nn.Module, dict]:
    """Train with MAE loss and early stopping on validation MAE; restore the best epoch.

    A ReLU output that is negative for every input passes no gradient and never recovers. With the
    paper's N(0, 0.05) initialisation this happens in some runs, so such a run is restarted with a new
    seed (counted in info['restarts']).
    """
    info = {"restarts": 0}
    for attempt in range(cfg.max_restarts + 1):
        torch.manual_seed(cfg.seed * 1000 + attempt)
        gen = torch.Generator().manual_seed(cfg.seed * 1000 + attempt)
        model = make_model()
        lr = KERAS2_DEFAULT_LR[cfg.optimizer] if cfg.lr is None else cfg.lr
        opt = OPTIMIZERS[cfg.optimizer](model.parameters(), lr)
        loss_fn = nn.L1Loss()
        best, best_state, best_epoch, wait = float("inf"), None, 0, 0
        n = len(y_tr)
        for epoch in range(1, cfg.max_epochs + 1):
            model.train()
            for b in _batches(n, cfg.batch_size, gen):
                opt.zero_grad()
                loss = loss_fn(model(*(x[b] for x in X_tr)), y_tr[b])
                loss.backward()
                opt.step()
            model.eval()
            with torch.no_grad():
                pv = model(*X_va)
                val = (pv - y_va).abs().mean().item()
            if val < best - 1e-5:
                best, best_epoch, wait = val, epoch, 0
                best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}
            else:
                wait += 1
                if wait >= cfg.patience:
                    break
        model.load_state_dict(best_state)
        model.eval()
        with torch.no_grad():
            dead = (model(*X_va) <= 0).all(dim=0).any().item()  # some output unit is 0 for every input
        if not dead or attempt == cfg.max_restarts:
            info.update(best_val_mae=best, best_epoch=best_epoch, epochs_run=epoch, dead_output=bool(dead))
            return model, info
        info["restarts"] += 1


@torch.no_grad()
def predict(model, *X, batch=65536):
    model.eval()
    n = len(X[0])
    return torch.cat([model(*(x[s:s + batch] for x in X)) for s in range(0, n, batch)]).numpy()


def _t(a, dtype=torch.float32):
    return torch.as_tensor(np.array(a), dtype=dtype)  # np.array copies, so the tensor is writable


@dataclass
class DMCCFConfig:
    emb: int = 64
    criteria_hidden: tuple = (128, 64, 32, 16, 8)
    overall_hidden: tuple = (64, 32, 16, 8)
    init_std: float = 0.05
    train: TrainConfig = field(default_factory=TrainConfig)


def train_overall_dnn(df, fold: Fold, cfg: DMCCFConfig):
    """Learn the aggregation function f on the TRUE training criteria (Section 3.2)."""
    C = df[CRITERIA].to_numpy(np.float32)
    mean, std = C[fold.train].mean(0), C[fold.train].std(0)
    z = lambda rows: _t((C[rows] - mean) / std)
    y = df["overall"].to_numpy(np.float32)
    model, info = fit(lambda: OverallRatingDNN(len(CRITERIA), cfg.overall_hidden, cfg.init_std),
                      (z(fold.train),), _t(y[fold.train]), (z(fold.val),), _t(y[fold.val]),
                      replace(cfg.train, seed=cfg.train.seed + 100))
    return model, (mean, std), info


def train_criteria_dnn(df, fold: Fold, cfg: DMCCFConfig):
    """Predict the 7 criteria ratings from user and item IDs (Section 3.1)."""
    nu, ni = df.u.max() + 1, df.i.max() + 1
    U, I = _t(df.u.values, torch.long), _t(df.i.values, torch.long)
    C = _t(df[CRITERIA].to_numpy(np.float32))
    return fit(lambda: CriteriaRatingsDNN(nu, ni, len(CRITERIA), cfg.emb, cfg.criteria_hidden, cfg.init_std),
               (U[fold.train], I[fold.train]), C[fold.train], (U[fold.val], I[fold.val]), C[fold.val], cfg.train)


def run_dmccf(df, fold: Fold, cfg: DMCCFConfig = DMCCFConfig(), overall=None) -> dict:
    """The full two-stage model on one fold. `overall` lets ablations reuse a trained Overall DNN."""
    t0 = time.time()
    if overall is None:
        overall = train_overall_dnn(df, fold, cfg)
    o_model, (mean, std), o_info = overall
    c_model, c_info = train_criteria_dnn(df, fold, cfg)
    te = fold.test
    crit_pred = predict(c_model, _t(df.u.values[te], torch.long), _t(df.i.values[te], torch.long))
    pred = predict(o_model, _t((crit_pred - mean) / std))            # Recommendation steps (a) and (b)
    crit_true = df[CRITERIA].to_numpy(np.float32)[te]
    return {
        "pred": pred,
        "criteria_mae": np.abs(crit_pred - crit_true).mean(0),
        "criteria_info": c_info,
        "overall_info": o_info,
        "overall": overall,
        "seconds": time.time() - t0,
    }


@dataclass
class SingleDNNConfig:
    emb: int = 256
    hidden: tuple = (512, 256, 128, 64, 32, 16, 8)
    init_std: float = 0.05
    train: TrainConfig = field(default_factory=TrainConfig)


def run_single_dnn(df, fold: Fold, cfg: SingleDNNConfig = SingleDNNConfig()) -> dict:
    t0 = time.time()
    nu, ni = df.u.max() + 1, df.i.max() + 1
    U, I = _t(df.u.values, torch.long), _t(df.i.values, torch.long)
    y = _t(df["overall"].to_numpy(np.float32))
    model, info = fit(lambda: SingleDNN(nu, ni, cfg.emb, cfg.hidden, cfg.init_std),
                      (U[fold.train], I[fold.train]), y[fold.train], (U[fold.val], I[fold.val]), y[fold.val],
                      replace(cfg.train, seed=cfg.train.seed + 200))
    return {"pred": predict(model, U[fold.test], I[fold.test]), "info": info, "seconds": time.time() - t0}


# --------------------------------------------------------------------------------------------------
# Single-rating baselines (scikit-surprise, default hyper-parameters)
# --------------------------------------------------------------------------------------------------
def surprise_algorithms() -> dict:
    from surprise import SVD, SVDpp, SlopeOne, BaselineOnly, KNNBaseline
    return {
        "SVD": lambda: SVD(random_state=0),
        "SVD++": lambda: SVDpp(random_state=0, cache_ratings=True),
        "SlopeOne": SlopeOne,
        "Baseline estimates": lambda: BaselineOnly(verbose=False),
        # 70k users make a user-user similarity matrix too large (~40 GB), so KNN is item-based.
        "KNN with baseline": lambda: KNNBaseline(sim_options={"name": "pearson_baseline", "user_based": False},
                                                 verbose=False),
    }


def run_surprise(df, fold: Fold, name: str) -> dict:
    from surprise import Dataset, Reader
    t0 = time.time()
    tr = df.iloc[fold.train_full]
    data = Dataset.load_from_df(tr[["user", "hotel", "overall"]], Reader(rating_scale=(1, 5)))
    algo = surprise_algorithms()[name]()
    algo.fit(data.build_full_trainset())
    te = df.iloc[fold.test]
    pred = np.array([algo.predict(u, i).est for u, i in zip(te.user, te.hotel)], dtype=np.float32)
    return {"pred": pred, "seconds": time.time() - t0}


# --------------------------------------------------------------------------------------------------
# Metrics (Section 4.2)
# --------------------------------------------------------------------------------------------------
@dataclass
class MetricConfig:
    k: int = 10             # top-K cutoff: not stated in the paper
    threshold: float = 4.0  # an item is relevant if its true overall rating >= threshold: not stated
    min_items: int = 1      # only users with at least this many test ratings count in the top-K metrics
    undefined: float = 1.0  # precision (recall) of a user with no recommended (relevant) item. 1.0 is the
                            # 2018 scikit-surprise FAQ recipe; it reproduces the paper's baseline F1/F2.


def evaluate(users: np.ndarray, y_true: np.ndarray, y_pred: np.ndarray, mc: MetricConfig = MetricConfig()) -> dict:
    """MAE (Eq. 10), F1 and F2 (Eqs. 11-13), FCP (Eqs. 14-18), MAP (Eq. 19) and MRR (Eq. 20).

    Precision/recall follow the scikit-surprise FAQ recipe: per user, the recommended items are the
    top-K by prediction whose predicted rating is >= threshold; precision and recall are set to
    `mc.undefined` when their denominator is 0; F-beta uses the user-averaged precision and recall (Eq. 11).
    MAP and MRR rank each user's test items by prediction and average over users with >= 1 relevant item.
    """
    out = {"MAE": float(np.mean(np.abs(y_true - y_pred))), "RMSE": float(np.sqrt(np.mean((y_true - y_pred) ** 2)))}
    order = np.lexsort((-y_pred, users))           # group by user, highest prediction first
    u, t, p = users[order], y_true[order], y_pred[order]
    starts = np.flatnonzero(np.r_[True, u[1:] != u[:-1]])
    ends = np.r_[starts[1:], len(u)]
    precs, recs, aps, rrs = [], [], [], []
    nc = nd = 0
    n_users = 0
    for s, e in zip(starts, ends):
        if e - s < mc.min_items:
            continue
        n_users += 1
        tu, pu = t[s:e], p[s:e]
        rel = tu >= mc.threshold
        top = slice(0, mc.k)
        rec_k = pu[top] >= mc.threshold
        n_rel, n_rec = rel.sum(), rec_k.sum()
        n_hit = (rel[top] & rec_k).sum()
        precs.append(n_hit / n_rec if n_rec else mc.undefined)
        recs.append(n_hit / n_rel if n_rel else mc.undefined)
        if n_rel:
            hits = rel[top]
            prec_at = np.cumsum(hits) / np.arange(1, len(hits) + 1)
            aps.append(float((prec_at * hits).sum() / min(mc.k, n_rel)))
            first = np.flatnonzero(hits)
            rrs.append(1.0 / (first[0] + 1) if len(first) else 0.0)
        if e - s > 1:  # Eqs. (17)-(18) over ordered pairs of this user's test items
            dp = pu[:, None] - pu[None, :]
            dt = tu[:, None] - tu[None, :]
            nc += int(((dp > 0) & (dt > 0)).sum())
            nd += int(((dp >= 0) & (dt < 0)).sum())
    P, R = float(np.mean(precs)), float(np.mean(recs))
    f = lambda b: (1 + b * b) * P * R / (b * b * P + R) if (P + R) > 0 else 0.0
    out.update(Precision=P, Recall=R, F1=f(1.0), F2=f(2.0),
               FCP=nc / (nc + nd) if (nc + nd) else float("nan"),
               MAP=float(np.mean(aps)) if aps else float("nan"),
               MRR=float(np.mean(rrs)) if rrs else float("nan"),
               users=n_users, pairs=nc + nd)
    return out


def summarise(per_fold: pd.DataFrame, metrics: list[str]) -> pd.DataFrame:
    """mean ± std over folds, one row per method, in the order methods first appear."""
    g = per_fold.groupby("method", sort=False)[metrics]
    m, s = g.mean(), g.std(ddof=0)
    return pd.DataFrame({c: [f"{a:.4f} ± {b:.4f}" for a, b in zip(m[c], s[c])] for c in metrics}, index=m.index)


# The paper's reported numbers (Tables 3-6), for side-by-side comparison.
PAPER_TABLE6 = pd.DataFrame(
    {
        "SVD": [0.8071, 0.8370, 0.7951, 0.5024, 0.3935, 0.7272],
        "SVD++": [0.8222, 0.8390, 0.7991, 0.5092, 0.3483, 0.7513],
        "SlopeOne": [0.9160, 0.8666, 0.9248, 0.5096, 0.3006, 0.7502],
        "Baseline estimates": [0.8094, 0.8391, 0.7986, 0.5025, 0.3881, 0.7374],
        "KNN with baseline": [0.8115, 0.8388, 0.7984, 0.5097, 0.3865, 0.7360],
        "Single DNN": [0.7855, 0.8572, 0.9176, 0.4957, 0.4481, 0.7244],
        "Proposed (DMCCF)": [0.7552, 0.8769, 0.9295, 0.5126, 0.4643, 0.7664],
    },
    index=["MAE", "F1", "F2", "FCP", "MAP", "MRR"],
).T
PAPER_TABLE5 = pd.Series([0.8195, 0.7612, 0.6581, 0.6324, 0.7582, 0.7401, 0.8774, 0.7495],
                         index=CRITERIA_LABELS + ["Mean"])
PAPER_TABLE3 = pd.Series({"adadelta": 0.9112, "adagrad": 0.7565, "adam": 0.7552, "adamax": 0.7919,
                          "rmsprop": 0.7612, "sgd": 3.2805})
PAPER_TABLE4 = pd.Series({"[8]": 3.4703, "[8,16]": 3.0678, "[8,16,32]": 0.9481, "[8,16,32,64]": 0.7776,
                          "[8,16,32,64,128]": 0.7552, "[8,16,32,64,128,256]": 0.7830})