#!/usr/bin/env python
"""Сводка follow-up экспериментов и парные сравнения по повторам.

Скрипт читает каталоги results/followup/<label>/, созданные
scripts/run_followup_experiments.sh, и формирует:

- summary.csv / summary.md — среднее macro-F1 по повторам, а не только repeat 0;
- comparisons.csv / comparisons.md — парные сравнения с McNemar и bootstrap;
- comparisons/<A>__vs__<B>.json — детали по каждому повтору.

Пример:
    python scripts/analyze_followup.py \
        --followup-dir results/followup \
        --pairs final_rf:ref_logreg final_rf:ref_no_lsa
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

from stylometry.evaluation import classification_metrics  # noqa: E402
from stylometry.stats import bootstrap_difference, mcnemar_exact  # noqa: E402


def _json_load(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _read_run(run_dir: Path) -> tuple[dict[str, Any], pd.DataFrame]:
    metadata = _json_load(run_dir / "run_metadata.json")
    pooled_path = run_dir / "pooled_metrics.csv"
    if pooled_path.exists():
        pooled = pd.read_csv(pooled_path)
    else:
        pooled = pd.DataFrame([metadata.get("metrics", {})])
    return metadata, pooled


def _feature_summary(config: dict[str, Any]) -> str:
    keys = (
        "groups",
        "lsa_variants",
        "lsa_components",
        "min_df",
        "char_max_features",
        "pos_max_features",
        "char_ngram_range",
        "pos_ngram_range",
        "block_weights",
    )
    features = config.get("feature_config", {})
    return json.dumps({key: features.get(key) for key in keys if key in features}, ensure_ascii=False)


def build_summary(followup_dir: Path) -> pd.DataFrame:
    rows: list[dict[str, Any]] = []
    for run_dir in sorted(followup_dir.iterdir()):
        if not run_dir.is_dir() or not (run_dir / "run_metadata.json").exists():
            continue
        metadata, pooled = _read_run(run_dir)
        config = metadata.get("config", {})
        bootstrap = metadata.get("bootstrap", {}).get("macro_f1", {})
        macro = pd.to_numeric(pooled.get("macro_f1"), errors="coerce")
        accuracy = pd.to_numeric(pooled.get("accuracy"), errors="coerce")
        balanced = pd.to_numeric(pooled.get("balanced_accuracy"), errors="coerce")
        top3 = pd.to_numeric(pooled.get("top3_accuracy"), errors="coerce")
        command_path = run_dir / "command.txt"
        rows.append(
            {
                "label": run_dir.name,
                "preset": config.get("preset"),
                "model": config.get("model"),
                "folds": config.get("folds"),
                "n_repeats": int(len(pooled)),
                "n_samples": config.get("n_samples"),
                "mean_macro_f1": float(macro.mean()),
                "std_macro_f1": float(macro.std(ddof=1)) if len(macro) > 1 else 0.0,
                "min_macro_f1": float(macro.min()),
                "max_macro_f1": float(macro.max()),
                "repeat0_macro_f1": float(macro.iloc[0]),
                "mean_accuracy": float(accuracy.mean()),
                "mean_balanced_accuracy": float(balanced.mean()),
                "mean_top3_accuracy": float(top3.mean()),
                "ci_lower_repeat0": bootstrap.get("lower"),
                "ci_upper_repeat0": bootstrap.get("upper"),
                "model_params": json.dumps(config.get("model_params", {}), ensure_ascii=False),
                "feature_config": _feature_summary(config),
                "command": command_path.read_text(encoding="utf-8").strip()
                if command_path.exists()
                else "",
            }
        )
    summary = pd.DataFrame(rows)
    if summary.empty:
        return summary
    reference = summary.loc[summary["label"] == "ref_logreg", "mean_macro_f1"]
    if not reference.empty:
        summary["delta_vs_ref_logreg"] = summary["mean_macro_f1"] - float(reference.iloc[0])
    else:
        summary["delta_vs_ref_logreg"] = np.nan
    summary = summary.sort_values("mean_macro_f1", ascending=False).reset_index(drop=True)
    summary.insert(0, "rank", np.arange(1, len(summary) + 1))
    return summary


def _compare_pair(
    followup_dir: Path,
    label_a: str,
    label_b: str,
    n_boot: int,
    seed: int,
) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]:
    dir_a = followup_dir / label_a
    dir_b = followup_dir / label_b
    if not (dir_a / "oof_predictions.csv").exists() or not (dir_b / "oof_predictions.csv").exists():
        print(f"[compare] пропуск {label_a} vs {label_b}: нет oof_predictions.csv")
        return None, []

    oof_a = pd.read_csv(dir_a / "oof_predictions.csv")
    oof_b = pd.read_csv(dir_b / "oof_predictions.csv")
    common_repeats = sorted(set(oof_a["repeat"].unique()) & set(oof_b["repeat"].unique()))
    if not common_repeats:
        print(f"[compare] пропуск {label_a} vs {label_b}: нет общих повторов")
        return None, []

    details: list[dict[str, Any]] = []
    for repeat in common_repeats:
        first = oof_a.loc[oof_a["repeat"] == repeat, ["line_id", "y_true", "y_pred"]]
        second = oof_b.loc[oof_b["repeat"] == repeat, ["line_id", "y_true", "y_pred"]]
        merged = first.merge(
            second,
            on=["line_id", "y_true"],
            suffixes=("_a", "_b"),
            how="inner",
        )
        if merged.empty:
            continue
        labels = sorted(merged["y_true"].unique())
        metrics_a = classification_metrics(merged["y_true"], merged["y_pred_a"], labels=labels)
        metrics_b = classification_metrics(merged["y_true"], merged["y_pred_b"], labels=labels)
        mcnemar = mcnemar_exact(merged["y_true"], merged["y_pred_a"], merged["y_pred_b"])
        difference = bootstrap_difference(
            merged["y_true"],
            merged["y_pred_a"],
            merged["y_pred_b"],
            metric="macro_f1",
            labels=labels,
            n_boot=n_boot,
            seed=seed + int(repeat),
        )
        details.append(
            {
                "repeat": int(repeat),
                "n_samples": int(len(merged)),
                "macro_f1_a": metrics_a["macro_f1"],
                "macro_f1_b": metrics_b["macro_f1"],
                "macro_f1_diff": metrics_a["macro_f1"] - metrics_b["macro_f1"],
                "accuracy_a": metrics_a["accuracy"],
                "accuracy_b": metrics_b["accuracy"],
                "accuracy_diff": metrics_a["accuracy"] - metrics_b["accuracy"],
                "mcnemar_p": mcnemar["p_value"],
                "mcnemar_b": mcnemar["b"],
                "mcnemar_c": mcnemar["c"],
                "bootstrap_mean_diff": difference["mean_difference"],
                "bootstrap_lower": difference["lower"],
                "bootstrap_upper": difference["upper"],
            }
        )

    if not details:
        print(f"[compare] пропуск {label_a} vs {label_b}: не удалось сопоставить строки")
        return None, []

    diffs = np.asarray([row["macro_f1_diff"] for row in details], dtype=float)
    p_values = np.asarray([row["mcnemar_p"] for row in details], dtype=float)
    ci_lower = np.asarray([row["bootstrap_lower"] for row in details], dtype=float)
    ci_upper = np.asarray([row["bootstrap_upper"] for row in details], dtype=float)
    summary = {
        "pair": f"{label_a} vs {label_b}",
        "label_a": label_a,
        "label_b": label_b,
        "n_repeats": len(details),
        "mean_macro_f1_a": float(np.mean([row["macro_f1_a"] for row in details])),
        "mean_macro_f1_b": float(np.mean([row["macro_f1_b"] for row in details])),
        "mean_diff": float(diffs.mean()),
        "std_diff": float(diffs.std(ddof=1)) if len(diffs) > 1 else 0.0,
        "min_diff": float(diffs.min()),
        "max_diff": float(diffs.max()),
        "n_positive": int((diffs > 0).sum()),
        "n_negative": int((diffs < 0).sum()),
        "mean_mcnemar_p": float(p_values.mean()),
        "min_mcnemar_p": float(p_values.min()),
        "mean_ci_lower": float(ci_lower.mean()),
        "mean_ci_upper": float(ci_upper.mean()),
        "n_ci_excludes_zero": int(((ci_lower > 0) | (ci_upper < 0)).sum()),
    }
    return summary, details


def _format_markdown(summary: pd.DataFrame) -> str:
    columns = [
        "rank",
        "label",
        "model",
        "preset",
        "n_repeats",
        "mean_macro_f1",
        "std_macro_f1",
        "mean_accuracy",
        "mean_top3_accuracy",
        "delta_vs_ref_logreg",
    ]
    table = summary[columns].copy()
    return table.to_markdown(index=False, floatfmt=".4f")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--followup-dir",
        default=str(PROJECT_ROOT / "results" / "followup"),
        help="Каталог с подкаталогами запусков.",
    )
    parser.add_argument(
        "--pairs",
        nargs="*",
        default=[],
        help="Пары для сравнения в формате label_a:label_b.",
    )
    parser.add_argument("--n-boot", type=int, default=1000)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args(argv)

    followup_dir = Path(args.followup_dir)
    if not followup_dir.exists():
        raise FileNotFoundError(f"Нет каталога: {followup_dir}")

    summary = build_summary(followup_dir)
    if summary.empty:
        print(f"[summary] в {followup_dir} нет запусков с run_metadata.json")
        return 1

    summary_path = followup_dir / "summary.csv"
    summary.to_csv(summary_path, index=False, encoding="utf-8-sig")
    markdown = [
        "# Follow-up: сводка запусков",
        "",
        "Среднее и стандартное отклонение посчитаны по pooled-метрикам каждого",
        "повтора (`pooled_metrics.csv`), а не только по repeat 0.",
        "",
        _format_markdown(summary),
        "",
    ]
    (followup_dir / "summary.md").write_text("\n".join(markdown), encoding="utf-8")
    print(f"[summary] {summary_path}")
    print(
        summary[
            ["rank", "label", "model", "n_repeats", "mean_macro_f1", "std_macro_f1"]
        ].to_string(index=False)
    )

    comparison_rows: list[dict[str, Any]] = []
    comparisons_dir = followup_dir / "comparisons"
    comparisons_dir.mkdir(exist_ok=True)
    for pair in args.pairs:
        if ":" not in pair:
            print(f"[compare] некорректная пара: {pair}")
            continue
        label_a, label_b = pair.split(":", 1)
        summary_row, details = _compare_pair(
            followup_dir,
            label_a.strip(),
            label_b.strip(),
            n_boot=args.n_boot,
            seed=args.seed,
        )
        if summary_row is None:
            continue
        comparison_rows.append(summary_row)
        payload = {"summary": summary_row, "details": details}
        output = comparisons_dir / f"{label_a.strip()}__vs__{label_b.strip()}.json"
        output.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        print(
            f"[compare] {summary_row['pair']}: "
            f"diff={summary_row['mean_diff']:+.4f}, "
            f"{summary_row['n_positive']}/{summary_row['n_repeats']} повторов в плюс, "
            f"mean McNemar p={summary_row['mean_mcnemar_p']:.4f}"
        )

    if comparison_rows:
        comparisons = pd.DataFrame(comparison_rows).sort_values("mean_diff", ascending=False)
        comparisons.to_csv(followup_dir / "comparisons.csv", index=False, encoding="utf-8-sig")
        columns = [
            "pair",
            "n_repeats",
            "mean_macro_f1_a",
            "mean_macro_f1_b",
            "mean_diff",
            "std_diff",
            "n_positive",
            "mean_mcnemar_p",
            "mean_ci_lower",
            "mean_ci_upper",
            "n_ci_excludes_zero",
        ]
        comparison_markdown = [
            "# Follow-up: парные сравнения",
            "",
            "`mean_diff` — средняя разница macro-F1 (A − B) по общим повторам.",
            "`n_positive` — сколько повторов дали положительную разницу.",
            "`n_ci_excludes_zero` — сколько повторов имеют bootstrap-интервал,",
            "не содержащий ноль (описательно, без поправки на множественность).",
            "",
            comparisons[columns].to_markdown(index=False, floatfmt=".4f"),
            "",
        ]
        (followup_dir / "comparisons.md").write_text(
            "\n".join(comparison_markdown), encoding="utf-8"
        )
        print(f"[compare] {followup_dir / 'comparisons.csv'}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())