"""
Head-to-head: a teammate's pipeline against ours, on two common yardsticks.

NOTE FOR THIS BUNDLE. This file is here because tools/full_score_combo.py imports
it to rebuild the zero-shot anchor of FINAL-2 (its XGBoost, DANN, region
normalisation and smoothing functions). The head-to-head in main() needs the
teammate's lockbox block list, which is not in this repo, so main() cannot be run
from the bundle.

  lockbox    their own Madrid test: train on their 80% dev blocks, score once on
             their 29 locked 2 km blocks (15,157 px). Rebuilt from our pixel keys
             with their block rule. Their own lockbox scores, which are not in
             this repo, were the fidelity check.
  amsterdam  train on ALL of Madrid, score ALL 25,992 Amsterdam pixels with no
             Amsterdam labels used except to score (zero-shot: the setting their
             pipeline was built for).

The functions here are adapted from the teammate's code and published with their
agreement. One difference: XGBoost and the DANN run
on CPU because this PC has no CUDA (different random streams; their DANN seeds
vary by 0.002-0.004). Their features are the official 60 on the same pixels:
pass `first` as the third argument to use the organisers' exact construction
(tools/build_official60_first.py); the default is our median-per-year rebuild.

Every model is scored under two decision rules:
  argmax        what their pipeline uses
  tau+plugin    ours: tau=0.5 logit adjustment then the macro-F1 plug-in, Madrid prior
and with/without their neighbour smoothing (r = 6 px = 180 m, alpha 0.5).

    python -I tools/compare_teammate.py <lockbox_blocks.json> reports/teammate_compare.csv
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))
os.chdir(ROOT)

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from lightgbm import LGBMClassifier
from scipy.spatial import cKDTree
from sklearn.ensemble import RandomForestClassifier
from sklearn.utils.class_weight import compute_sample_weight
from xgboost import XGBClassifier

import config as C
import evaluation as E
import gbdt as G

torch.set_num_threads(max(1, (os.cpu_count() or 4) - 2))
DEV = torch.device("cpu")
BLOCK = 67
SMOOTH_ALPHA = 0.5
R_PX = 6
XGB_PARAMS = {"n_estimators": 500, "learning_rate": 0.08, "max_depth": 8,
              "min_child_weight": 3.0, "subsample": 0.8, "colsample_bytree": 1.0,
              "reg_lambda": 10.0, "gamma": 1.0}            # their best_config.json
DANN_EPOCHS, DANN_SEEDS = 40, (0, 1, 2)


# ---- copied from their compare_models.py ---------------------------------- #
def block_ids(pid):
    return (pid[:, 0] // BLOCK) * 100000 + (pid[:, 1] // BLOCK)


def region_norm(X):
    mu, sd = X.mean(0), X.std(0)
    sd[sd == 0] = 1
    return (X - mu) / sd


def smooth(P, coords, r):
    tree = cKDTree(coords)
    nbrs = tree.query_ball_point(coords, r=r)
    out = np.empty_like(P)
    for i, nb in enumerate(nbrs):
        out[i] = SMOOTH_ALPHA * P[i] + (1 - SMOOTH_ALPHA) * P[nb].mean(0)
    return out


# ---- copied from their dann.py --------------------------------------------- #
class GradReverse(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, lam):
        ctx.lam = lam
        return x.view_as(x)

    @staticmethod
    def backward(ctx, g):
        return -ctx.lam * g, None


class DANN(nn.Module):
    def __init__(self, d_in=60, h=256, z=64):
        super().__init__()
        self.feat = nn.Sequential(nn.Linear(d_in, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(0.2),
                                  nn.Linear(h, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(0.2),
                                  nn.Linear(h, z), nn.ReLU())
        self.age = nn.Sequential(nn.Linear(z, 64), nn.ReLU(), nn.Linear(64, 4))
        self.city = nn.Sequential(nn.Linear(z, 64), nn.ReLU(), nn.Linear(64, 1))

    def forward(self, x, lam=0.0):
        f = self.feat(x)
        return self.age(f), self.city(GradReverse.apply(f, lam)), f


def train_dann(Xs, ys, Xt, epochs, adversarial=True, seed=0):
    torch.manual_seed(seed)
    np.random.seed(seed)
    net = DANN(Xs.shape[1]).to(DEV)
    opt = torch.optim.AdamW(net.parameters(), lr=1e-3, weight_decay=1e-4)
    counts = np.bincount(ys, minlength=5)[1:]
    w = torch.tensor(len(ys) / (4 * counts), dtype=torch.float32, device=DEV)
    ce, bce = nn.CrossEntropyLoss(weight=w), nn.BCEWithLogitsLoss()
    S, Y = torch.tensor(Xs, dtype=torch.float32, device=DEV), torch.tensor(ys - 1, device=DEV)
    T = torch.tensor(Xt, dtype=torch.float32, device=DEV)
    bs, steps = 512, int(np.ceil(len(S) / 512))
    for ep in range(epochs):
        net.train()
        perm_s = torch.randperm(len(S), device=DEV)
        for k in range(steps):
            p = (ep * steps + k) / (epochs * steps)
            lam = (2 / (1 + np.exp(-10 * p)) - 1) if adversarial else 0.0
            i = perm_s[k * bs:(k + 1) * bs]
            j = torch.randint(0, len(T), (len(i),), device=DEV)
            ya, cs, _ = net(S[i], lam)
            _, ct, _ = net(T[j], lam)
            loss = ce(ya, Y[i])
            if adversarial:
                loss = loss + bce(cs.squeeze(1), torch.zeros(len(i), device=DEV)) \
                            + bce(ct.squeeze(1), torch.ones(len(j), device=DEV))
            opt.zero_grad()
            loss.backward()
            opt.step()
    net.eval()
    return net


@torch.no_grad()
def dann_proba(net, X):
    a, _, _ = net(torch.tensor(X, dtype=torch.float32, device=DEV))
    return torch.softmax(a, 1).cpu().numpy()


# ---- their tree models (final_lockbox.py / tune_boosting.py) --------------- #
def their_trees(Xtr, ytr, Xte):
    sw = compute_sample_weight("balanced", ytr)
    xgb = XGBClassifier(objective="multi:softprob", tree_method="hist", device="cpu",
                        eval_metric="mlogloss", random_state=42, **XGB_PARAMS)
    xgb.fit(Xtr, ytr - 1, sample_weight=sw)
    lgb = LGBMClassifier(n_estimators=600, learning_rate=0.05, num_leaves=63, subsample=0.8,
                         subsample_freq=1, colsample_bytree=0.8, class_weight="balanced",
                         n_jobs=6, random_state=42, verbose=-1).fit(Xtr, ytr)
    rf = RandomForestClassifier(n_estimators=300, class_weight="balanced", n_jobs=6,
                                random_state=42).fit(Xtr, ytr)
    return {"T: XGBoost (tuned)": xgb.predict_proba(Xte),
            "T: LightGBM": lgb.predict_proba(Xte),
            "T: Random Forest": rf.predict_proba(Xte)}


# ---- ours ------------------------------------------------------------------ #
def our_lgbm(Xtr, ytr, Xte):
    """notebook 3's 'zero-shot, no normalisation': full LGBM_PARAMS, balanced."""
    cw = G.balanced_class_weights(ytr)
    m = G.fit_lgbm(Xtr, ytr, class_weight=dict(zip(range(4), cw)))
    return m.predict_proba(Xte)


def load(official: str = "median"):
    mad = pd.read_parquet(C.PROC_DIR / "madrid_features.parquet")
    ams = pd.read_parquet(C.PROC_DIR / "amsterdam_features.parquet")
    drop = {"px_key", "py_key", "age_class", "weighted_mean_year", "city"}
    feats = [c for c in mad.columns if c not in drop and mad[c].dtype.kind in "fiu"]
    feats = [c for c in feats if not (mad[c].std() == 0 or ams[c].std() == 0)]
    assert len(feats) == 393, len(feats)
    f60 = list(C.FEATURE_COLS_60)                      # all 60, as their pkl has
    if official == "first":
        # The organisers' exact construction (first QA-valid observation per
        # year), from tools/build_official60_first.py, aligned row by row.
        def aligned(city, ref):
            o = pd.read_parquet(C.PROC_DIR / f"{city}_official60_first.parquet")
            has_label = "age_class" in o.columns       # the published tables carry no label
            o = ref[C.PID + ["age_class"]].merge(o, on=C.PID, how="left",
                                                 suffixes=("", "_o"))
            assert o[f60].notna().all().all() and len(o) == len(ref)
            assert not has_label or (o["age_class"] == o["age_class_o"]).all()
            return o[f60].to_numpy(float)
        Xm60, Xa60 = aligned("madrid", mad), aligned("amsterdam", ams)
    else:
        Xm60, Xa60 = mad[f60].to_numpy(float), ams[f60].to_numpy(float)
    print(f"official 60 built with agg={official}", flush=True)
    return dict(
        Xm60=Xm60, Xa60=Xa60,
        Xm=mad[feats].to_numpy(float), Xa=ams[feats].to_numpy(float),
        ym=mad["age_class"].to_numpy().astype(int), ya=ams["age_class"].to_numpy().astype(int),
        pm=mad[C.PID].to_numpy().astype(np.int64), pa=ams[C.PID].to_numpy().astype(np.int64))


def run_setting(name, d, tr, te_X60, te_X, te_y, te_xy, Xadv60):
    """Fit every model on Madrid rows `tr`, score on the given test set."""
    t0 = time.time()
    Xtr60, Xtr, ytr = d["Xm60"][tr], d["Xm"][tr], d["ym"][tr]
    prior = np.array([(ytr == c).mean() for c in E.CLASSES])
    P = their_trees(Xtr60, ytr, te_X60)
    print(f"[{name}] their trees done {time.time() - t0:.0f}s", flush=True)
    P["O: LightGBM, 393 features (our zero-shot)"] = our_lgbm(Xtr, ytr, te_X)
    P["O: LightGBM, official 60 (our params)"] = our_lgbm(Xtr60, ytr, te_X60)
    print(f"[{name}] our LightGBM done {time.time() - t0:.0f}s", flush=True)
    seeds = {}
    Xtr_n, Xte_n, Xadv_n = region_norm(Xtr60), region_norm(te_X60), region_norm(Xadv60)
    for s in DANN_SEEDS:
        net = train_dann(Xtr_n, ytr, Xadv_n, DANN_EPOCHS, adversarial=True, seed=s)
        Pd = dann_proba(net, Xte_n)
        seeds[s] = {"T: DANN": Pd,
                    "T: XGBoost + DANN (their FINAL)": (P["T: XGBoost (tuned)"] + Pd) / 2}
        print(f"[{name}] DANN seed {s} done {time.time() - t0:.0f}s", flush=True)

    coords = te_xy.astype(float)
    rows = []

    def score(model, seed, PP):
        for sm in (False, True):
            Q = smooth(PP, coords, R_PX) if sm else PP
            for rule in ("argmax", "tau+plugin"):
                lab = (np.array(E.CLASSES)[Q.argmax(1)] if rule == "argmax"
                       else E.adapt_pipeline(Q, prior, tau=0.5)[0])
                pc = E.per_class_f1(te_y, lab)
                rows.append(dict(setting=name, model=model, seed=seed, smooth=sm, rule=rule,
                                 macro_f1=E.macro_f1(te_y, lab),
                                 **{f"f1_c{c}": float(v) for c, v in zip(E.CLASSES, pc)},
                                 n_test=len(te_y)))

    for k, PP in P.items():
        score(k, 0, PP)
    for s, dd in seeds.items():
        for k, PP in dd.items():
            score(k, s, PP)
    print(f"[{name}] scored, {time.time() - t0:.0f}s total", flush=True)
    return rows


def main():
    lock_json, out_csv = Path(sys.argv[1]), Path(sys.argv[2])
    d = load(sys.argv[3] if len(sys.argv) > 3 else "median")
    g = block_ids(d["pm"])
    lock = np.isin(g, json.loads(lock_json.read_text())["blocks"])
    print(f"lockbox pixels {lock.sum():,} (theirs: 15,157); dev {(~lock).sum():,} "
          f"(theirs: 61,106)", flush=True)
    assert lock.sum() == 15157 and (~lock).sum() == 61106

    rows = []
    dev, lk = np.flatnonzero(~lock), np.flatnonzero(lock)
    rows += run_setting("lockbox", d, dev, d["Xm60"][lk], d["Xm"][lk], d["ym"][lk],
                        d["pm"][lk], d["Xa60"])
    pd.DataFrame(rows).to_csv(out_csv, index=False)
    allm = np.arange(len(d["ym"]))
    rows += run_setting("amsterdam", d, allm, d["Xa60"], d["Xa"], d["ya"], d["pa"], d["Xa60"])
    df = pd.DataFrame(rows)
    df.to_csv(out_csv, index=False)

    summ = (df.groupby(["setting", "model", "smooth", "rule"])["macro_f1"]
              .agg(["mean", "std", "size"]).round(4).reset_index())
    for st in ("lockbox", "amsterdam"):
        t = summ[summ.setting == st].pivot_table(index="model", columns=["rule", "smooth"],
                                                 values="mean")
        print(f"\n=== {st} (macro F1; DANN rows = mean of 3 seeds) ===")
        print(t.sort_values(("argmax", True), ascending=False).to_string())


if __name__ == "__main__":
    main()