"""Choose the DMCCF+ settings on an inner split, never touching the outer test folds.

Inner split: fold 0's 80 % training data is split again into 80 % inner-train (of which 5 % early-stopping
validation) and 20 % inner-test. Each step adds one change; results go to results/dev_selection.csv.
"""
from __future__ import annotations

import sys
import warnings
from dataclasses import replace
from pathlib import Path

import numpy as np
import pandas as pd
import torch

import deepmccf as D
import improved as I

warnings.filterwarnings("ignore", category=UserWarning)
HERE = Path(__file__).resolve().parent


def inner_split(df, seed=7):
    outer = D.make_folds(len(df))[0]
    rng = np.random.default_rng(seed)
    pool = rng.permutation(outer.train_full)
    n_te = len(pool) // 5
    te, rest = pool[:n_te], pool[n_te:]
    n_va = int(0.05 * len(rest))
    return np.sort(rest[n_va:]), np.sort(rest[:n_va]), np.sort(te)


def main(seeds=(0, 1)):
    torch.set_num_threads(8)
    df = D.load_ratings(HERE / "data" / "tripadvisor_mc.csv")
    tr, va, te = inner_split(df)
    y = df.overall.values[te]
    rows = []

    def record(name, preds):
        maes = [float(np.abs(p - y).mean()) for p in preds]
        rows.append({"step": name, "inner_MAE": np.mean(maes), "spread": np.ptp(maes)})
        print(f"{name:55s} {np.mean(maes):.4f}  (seeds {', '.join(f'{m:.4f}' for m in maes)})", flush=True)

    # Reference points on the same inner split
    fold = D.Fold(0, tr, va, te)
    record("paper model (as published)", [D.run_dmccf(df, fold, D.DMCCFConfig(train=D.TrainConfig(seed=s)))["pred"]
                                          for s in seeds])
    record("paper model + clip to 1..5", [np.clip(p, 1, 5) for p in [D.run_dmccf(
        df, fold, D.DMCCFConfig(train=D.TrainConfig(seed=s)))["pred"] for s in seeds]])

    final = I.PlusConfig()
    base = replace(final, hidden=(64, 32), id_dropout=0.0, finetune=False, clip=False)
    steps = [
        ("1  He init, linear outputs, emb 16, [64,32], dropout", base),
        ("2  + unknown-user token (ID dropout 0.3)", replace(base, id_dropout=0.3)),
        ("3  + joint fine-tuning of both stages", replace(base, id_dropout=0.3, finetune=True)),
        ("4  + clip to 1..5", replace(final, hidden=(64, 32))),
        ("5  + deeper tower [128,64,32]  (= DMCCF+)", final),
        ("   tried: per-criterion user/item biases", replace(final, biases=True)),
        ("   tried: ID dropout 0.1", replace(final, id_dropout=0.1)),
        ("   tried: ID dropout 0.5", replace(final, id_dropout=0.5)),
        ("   tried: emb 8", replace(final, emb=8)),
        ("   tried: emb 32", replace(final, emb=32)),
        ("   tried: no criteria loss while fine-tuning", replace(final, aux_weight=0.0)),
        ("   tried: wider tower [256,128,64]", replace(final, hidden=(256, 128, 64))),
        ("   tried: dropout 0.2", replace(final, dropout=0.2)),
    ]
    for name, cfg in steps:
        record(name, [I.fit_predict(df, tr, va, te, cfg, seed=s)["pred"] for s in seeds])
    ens = I.fit_predict  # 5-seed ensemble of the chosen config
    record("6  5-seed ensemble of DMCCF+", [np.mean([ens(df, tr, va, te, final, seed=s)["pred"]
                                                    for s in range(5)], 0)])
    pd.DataFrame(rows).to_csv(HERE / "results" / "dev_selection.csv", index=False)


if __name__ == "__main__":
    main()