"""Stack LightGCN scores with explicit graph/meta-path features.

The validation estimate uses out-of-fold predictions on the notebook-style
dynamic validation split, so the second-stage model is not evaluated on rows it
was trained on.
"""

from __future__ import annotations

import argparse
import importlib.util
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import precision_recall_curve, roc_auc_score
from sklearn.model_selection import StratifiedKFold


def load_lgcn_module(path: Path):
    spec = importlib.util.spec_from_file_location("train_val_lgcn_ensemble", path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def read_txt(path: Path) -> list[list[int]]:
    return [list(map(int, line.strip().split())) for line in path.open()]


def best_f1(y: np.ndarray, score: np.ndarray):
    p, r, t = precision_recall_curve(y, score)
    f = 2 * p * r / (p + r + 1e-12)
    i = int(np.argmax(f))
    th = float(t[i]) if i < len(t) else 0.5
    return float(f[i]), th, float(roc_auc_score(y, score))


def rank01(x: np.ndarray) -> np.ndarray:
    order = np.argsort(x, kind="mergesort")
    out = np.empty(len(x), dtype=np.float32)
    out[order] = np.linspace(0, 1, len(x), dtype=np.float32)
    return out


def zscore(x: np.ndarray) -> np.ndarray:
    return ((x - x.mean()) / (x.std() + 1e-8)).astype(np.float32)


class ExplicitGraphFeatures:
    def __init__(self, root: Path, train_refs: pd.DataFrame, num_authors: int = 6611, num_papers: int = 79937):
        data_dir = root / "data_and_docs"
        self.num_authors = num_authors
        self.num_papers = num_papers
        self.train = train_refs[["source", "target"]].to_numpy(np.int64)
        citation = np.array(read_txt(data_dir / "paper_file_ann.txt"), dtype=np.int64)
        coauthor = np.array(read_txt(data_dir / "author_file_ann.txt"), dtype=np.int64)

        self.author_papers: list[set[int]] = [set() for _ in range(num_authors)]
        self.paper_readers: list[set[int]] = [set() for _ in range(num_papers)]
        self.author_degree = np.zeros(num_authors, dtype=np.float32)
        self.paper_degree = np.zeros(num_papers, dtype=np.float32)
        for a, p in self.train:
            a = int(a)
            p = int(p)
            self.author_papers[a].add(p)
            self.paper_readers[p].add(a)
            self.author_degree[a] += 1
            self.paper_degree[p] += 1

        self.coauthors: list[set[int]] = [set() for _ in range(num_authors)]
        for a, b in coauthor:
            self.coauthors[int(a)].add(int(b))
            self.coauthors[int(b)].add(int(a))

        self.paper_refs: list[set[int]] = [set() for _ in range(num_papers)]
        self.paper_cited_by: list[set[int]] = [set() for _ in range(num_papers)]
        self.cite_out_degree = np.zeros(num_papers, dtype=np.float32)
        self.cite_in_degree = np.zeros(num_papers, dtype=np.float32)
        for s, t in citation:
            s = int(s)
            t = int(t)
            self.paper_refs[s].add(t)
            self.paper_cited_by[t].add(s)
            self.cite_out_degree[s] += 1
            self.cite_in_degree[t] += 1

        # A-P-A neighborhood: authors sharing at least one historical paper.
        self.shared_paper_authors: list[set[int]] = [set() for _ in range(num_authors)]
        for a in range(num_authors):
            neigh = set()
            for p in self.author_papers[a]:
                neigh.update(self.paper_readers[p])
            neigh.discard(a)
            self.shared_paper_authors[a] = neigh

        # Coauthor paper union is reused by A-A-P style counts.
        self.coauthor_paper_union: list[set[int]] = [set() for _ in range(num_authors)]
        for a in range(num_authors):
            papers = set()
            for c in self.coauthors[a]:
                papers.update(self.author_papers[c])
            self.coauthor_paper_union[a] = papers

    def transform(self, pairs: np.ndarray) -> np.ndarray:
        out = np.zeros((len(pairs), 18), dtype=np.float32)
        for i, (a_raw, p_raw) in enumerate(pairs):
            a = int(a_raw)
            p = int(p_raw)
            hist = self.author_papers[a]
            coauthors = self.coauthors[a]
            co_papers = self.coauthor_paper_union[a]
            refs = self.paper_refs[p]
            cited_by = self.paper_cited_by[p]
            readers = self.paper_readers[p]

            co_read_count = sum(1 for c in coauthors if p in self.author_papers[c])
            hist_ref_overlap = len(hist & refs)
            hist_cited_by_overlap = len(hist & cited_by)
            ref_union = len(hist | refs)
            cited_by_union = len(hist | cited_by)
            shared_author_read_count = len(self.shared_paper_authors[a] & readers)

            out[i, 0] = self.author_degree[a]
            out[i, 1] = self.paper_degree[p]
            out[i, 2] = len(coauthors)
            out[i, 3] = co_read_count
            out[i, 4] = co_read_count / max(1.0, float(len(coauthors)))
            out[i, 5] = self.cite_in_degree[p]
            out[i, 6] = self.cite_out_degree[p]
            out[i, 7] = hist_ref_overlap
            out[i, 8] = hist_cited_by_overlap
            out[i, 9] = hist_ref_overlap / max(1.0, float(ref_union))
            out[i, 10] = hist_cited_by_overlap / max(1.0, float(cited_by_union))
            out[i, 11] = float(p in co_papers)  # A-A-P binary.
            out[i, 12] = co_read_count  # A-A-P count.
            out[i, 13] = hist_ref_overlap + hist_cited_by_overlap  # A-P-P count.
            out[i, 14] = shared_author_read_count  # A-P-A-P count.
            out[i, 15] = shared_author_read_count / max(1.0, float(len(self.shared_paper_authors[a])))
            out[i, 16] = np.log1p(self.author_degree[a])
            out[i, 17] = np.log1p(self.paper_degree[p])
        return out


def add_rank_features(pairs: np.ndarray, score: np.ndarray) -> np.ndarray:
    global_rank = rank01(score)
    author_pct = np.zeros(len(score), dtype=np.float32)
    author_rank = np.zeros(len(score), dtype=np.float32)
    df = pd.DataFrame({"idx": np.arange(len(score)), "author": pairs[:, 0], "score": score})
    for _, g in df.groupby("author", sort=False):
        order = np.argsort(g["score"].to_numpy(), kind="mergesort")
        idx = g["idx"].to_numpy()
        n = len(idx)
        vals = np.linspace(0, 1, n, dtype=np.float32) if n > 1 else np.array([1.0], dtype=np.float32)
        author_pct[idx[order]] = vals
        author_rank[idx[order]] = np.arange(n, dtype=np.float32)
    return np.column_stack([score.astype(np.float32), global_rank, author_pct, author_rank])


def fit_oof(X: np.ndarray, y: np.ndarray, model_kind: str, seed: int, n_splits: int) -> np.ndarray:
    oof = np.zeros(len(y), dtype=np.float32)
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
    for fold, (tr, va) in enumerate(skf.split(X, y), start=1):
        if model_kind == "logreg":
            clf = LogisticRegression(C=0.5, max_iter=1000, solver="lbfgs")
        else:
            clf = lgb.LGBMClassifier(
                n_estimators=1200,
                learning_rate=0.025,
                num_leaves=31,
                subsample=0.9,
                colsample_bytree=0.9,
                reg_lambda=5.0,
                min_child_samples=80,
                objective="binary",
                verbose=-1,
                random_state=seed + fold,
            )
        clf.fit(X[tr], y[tr])
        oof[va] = clf.predict_proba(X[va])[:, 1]
    return oof


def boundary_rerank(y: np.ndarray, lgcn: np.ndarray, stack: np.ndarray, raw_th: float):
    best = None
    dist = np.abs(lgcn - raw_th)
    for frac in [0.05, 0.10, 0.15, 0.20, 0.30, 0.40]:
        cutoff = np.quantile(dist, frac)
        mask = dist <= cutoff
        for alpha in np.linspace(0.0, 1.0, 11):
            score = zscore(lgcn)
            mixed = alpha * zscore(lgcn) + (1.0 - alpha) * zscore(stack)
            score[mask] = mixed[mask]
            f1, th, auc = best_f1(y, score)
            row = {"frac": frac, "alpha_lgcn": float(alpha), "f1": f1, "threshold": th, "auc": auc}
            if best is None or f1 > best["f1"]:
                best = row
    return best


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--package-root", type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument("--split-seed", type=int, required=True)
    parser.add_argument("--lgcn-score-file", type=Path, required=True)
    parser.add_argument("--model-kind", choices=["lgb", "logreg"], default="lgb")
    parser.add_argument("--n-splits", type=int, default=5)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--test-score-file", type=Path, default=None)
    parser.add_argument("--test-feature-source", choices=["split", "full"], default="full")
    args = parser.parse_args()

    root = args.package_root
    lgcn_mod = load_lgcn_module(root / "code" / "train_val_lgcn_ensemble.py")
    train_refs, val_pairs = lgcn_mod.make_notebook_style_split(root, args.split_seed, 0.9)
    builder = ExplicitGraphFeatures(root, train_refs)

    val_arr = val_pairs[["source", "target"]].to_numpy(np.int64)
    y = val_pairs["label"].to_numpy(np.int8)
    lgcn_score = np.load(args.lgcn_score_file).astype(np.float32)
    if len(lgcn_score) != len(y):
        raise ValueError(f"score length {len(lgcn_score)} != labels {len(y)}")

    print("computing validation explicit graph features", val_arr.shape)
    X_hand = builder.transform(val_arr)
    X_rank = add_rank_features(val_arr, lgcn_score)
    X_stack = np.column_stack([X_rank, X_hand]).astype(np.float32)

    raw_f1, raw_th, raw_auc = best_f1(y, lgcn_score)
    hand_oof = fit_oof(X_hand, y, args.model_kind, args.seed, args.n_splits)
    hand_f1, hand_th, hand_auc = best_f1(y, hand_oof)
    stack_oof = fit_oof(X_stack, y, args.model_kind, args.seed, args.n_splits)
    stack_f1, stack_th, stack_auc = best_f1(y, stack_oof)
    rerank = boundary_rerank(y, lgcn_score, stack_oof, raw_th)

    out_dir = root / "validation_runs" / f"dynamic_seed{args.split_seed}" / "stack_rank_calibration"
    out_dir.mkdir(parents=True, exist_ok=True)
    np.save(out_dir / "val_handcrafted_oof.npy", hand_oof)
    np.save(out_dir / "val_stack_oof.npy", stack_oof)
    rows = [
        {"method": "lgcn_raw", "f1": raw_f1, "threshold": raw_th, "auc": raw_auc},
        {"method": f"handcrafted_{args.model_kind}_oof", "f1": hand_f1, "threshold": hand_th, "auc": hand_auc},
        {"method": f"stack_lgcn_hand_{args.model_kind}_oof", "f1": stack_f1, "threshold": stack_th, "auc": stack_auc},
        {"method": "boundary_rerank", **rerank},
    ]
    result = pd.DataFrame(rows).sort_values("f1", ascending=False)
    result.to_csv(out_dir / "result.csv", index=False)
    print(result.to_string(index=False))

    if args.test_score_file is not None:
        test_pairs = np.array(read_txt(root / "data_and_docs" / "bipartite_test_ann.txt"), dtype=np.int64)
        test_score = np.load(args.test_score_file).astype(np.float32)
        if len(test_score) != len(test_pairs):
            raise ValueError(f"test score length {len(test_score)} != test pairs {len(test_pairs)}")
        test_builder = builder
        if args.test_feature_source == "full":
            full_refs = pd.DataFrame(
                read_txt(root / "data_and_docs" / "bipartite_train_ann.txt"),
                columns=["source", "target"],
            )
            test_builder = ExplicitGraphFeatures(root, full_refs)
        print("computing test explicit graph features", test_pairs.shape)
        X_test = np.column_stack([add_rank_features(test_pairs, test_score), test_builder.transform(test_pairs)]).astype(np.float32)
        clf = lgb.LGBMClassifier(
            n_estimators=1200,
            learning_rate=0.025,
            num_leaves=31,
            subsample=0.9,
            colsample_bytree=0.9,
            reg_lambda=5.0,
            min_child_samples=80,
            objective="binary",
            verbose=-1,
            random_state=args.seed,
        )
        clf.fit(X_stack, y)
        test_pred = clf.predict_proba(X_test)[:, 1].astype(np.float32)
        np.save(out_dir / "test_stack_pred.npy", test_pred)
        for ratio in [0.505, 0.515, 0.521, 0.530, 0.540]:
            n_pos = int(round(len(test_pred) * ratio))
            pred = np.zeros(len(test_pred), dtype=np.int8)
            pred[np.argsort(test_pred)[-n_pos:]] = 1
            known = np.load(root / "cached_scores" / "test_known_mask.npy").astype(bool)
            pred[known] = 1
            sub = pd.DataFrame({"Id": np.arange(len(pred)), "Probability": pred})
            sub.to_csv(out_dir / f"submission_stack_r{ratio:.3f}.csv", index=False)
        print(f"saved test predictions and ratio submissions under {out_dir}")


if __name__ == "__main__":
    main()