"""
Dataset and genre bias analysis for AURIS (reviewer priority #2), plus a
segment/track-level leakage check that feeds into the nested-CV / independent
test split design (reviewer priority #4).

Three outputs:
  1. dataset_bias_table.csv   — per-source composition (n, class, genre mix,
                                 duration stats, sample rate) so a reader can
                                 see whether sources differ in ways unrelated
                                 to AI-vs-human content.
  2. genre_bias_test.csv      — LightGBM performance on the full (genre-
                                 imbalanced) test set vs. a genre-balanced
                                 subsample, to check for a genre-as-shortcut
                                 effect.
  3. segment_leakage_report.txt — whether any track appears to be split
                                 across multiple audio files (same base name
                                 with numbered suffixes), which would leak
                                 information across CV folds.

Usage:
    python -m app.training.bias_analysis
"""

from __future__ import annotations

import csv
import re
import sys
import warnings
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

import lightgbm as lgb
from sklearn.exceptions import ConvergenceWarning
from sklearn.metrics import balanced_accuracy_score, f1_score, roc_auc_score, roc_curve
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

DATASET_DIR = Path("D:/CrownCode/DataSet")
FEATURES_WITH_META = DATASET_DIR / "features_with_meta.csv"
TABLES_DIR = Path(__file__).resolve().parents[3] / "docs/academic/paper/real_tables"

_EXCLUDED_COLUMNS = {
    "file_path", "label_int", "duration_sec", "sample_rate",
    "genre", "generator", "ai_model", "meta_sample_rate", "meta_duration_sec",
    "match_source",
}

_LGBM_PARAMS = dict(
    n_estimators=300, max_depth=-1, learning_rate=0.05, num_leaves=31,
    subsample=0.8, colsample_bytree=0.8, min_child_samples=20,
    reg_alpha=0.1, reg_lambda=1.0, class_weight="balanced",
    random_state=42, verbose=-1,
)


def _load_rows() -> list[dict]:
    with open(FEATURES_WITH_META, "r", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def dataset_bias_table(rows: list[dict]) -> None:
    """Per-source composition: n, class, genre mix, duration, sample rate."""
    by_source: dict[str, list[dict]] = defaultdict(list)
    for r in rows:
        by_source[r["generator"]].append(r)

    out_rows = []
    for source, source_rows in sorted(by_source.items()):
        durations = [float(r["duration_sec"]) for r in source_rows]
        sample_rates = Counter(r["sample_rate"] for r in source_rows)
        genres = Counter(r["genre"] for r in source_rows if r["genre"])
        labels = Counter(r["label_int"] for r in source_rows)
        top_genre = genres.most_common(1)[0] if genres else ("n/a", 0)

        out_rows.append({
            "source": source,
            "n": len(source_rows),
            "label": "AI" if labels.get("1", 0) > labels.get("0", 0) else "Human",
            "n_ai": labels.get("1", 0),
            "n_human": labels.get("0", 0),
            "n_distinct_genres": len(genres),
            "top_genre": top_genre[0],
            "top_genre_frac": round(top_genre[1] / len(source_rows), 3) if source_rows else 0,
            "duration_mean_sec": round(float(np.mean(durations)), 2),
            "duration_std_sec": round(float(np.std(durations)), 2),
            "duration_min_sec": round(float(np.min(durations)), 2),
            "duration_max_sec": round(float(np.max(durations)), 2),
            "sample_rate_modes": ";".join(f"{sr}={n}" for sr, n in sample_rates.most_common(3)),
        })

    out_path = TABLES_DIR / "dataset_bias_table.csv"
    with open(out_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(out_rows[0].keys()))
        writer.writeheader()
        writer.writerows(out_rows)

    print(f"Dataset bias table written: {out_path}")
    for r in out_rows:
        print(
            f"  {r['source']:25s} n={r['n']:5d} label={r['label']:5s} "
            f"dur={r['duration_mean_sec']:.1f}±{r['duration_std_sec']:.1f}s "
            f"genres={r['n_distinct_genres']:2d} top={r['top_genre']}({r['top_genre_frac']})"
        )


def genre_bias_test(rows: list[dict]) -> None:
    """
    Compare LightGBM performance on the natural (genre-imbalanced) test split
    vs. a genre-balanced test split, to check whether the model is exploiting
    genre as a shortcut for the AI/human label rather than learning
    AI-vs-human acoustic differences.
    """
    genre_rows = [r for r in rows if r["genre"]]
    print(f"\nRows with genre metadata: {len(genre_rows)} / {len(rows)}")

    feature_cols = [c for c in rows[0].keys() if c not in _EXCLUDED_COLUMNS]

    def to_xy(row_subset: list[dict]) -> tuple[np.ndarray, np.ndarray]:
        X = np.array([[float(r[c]) for c in feature_cols] for r in row_subset], dtype=np.float32)
        X = np.nan_to_num(X, nan=0.0, posinf=1.0, neginf=-1.0)
        y = np.array([int(r["label_int"]) for r in row_subset], dtype=np.int32)
        return X, y

    # Train on everything EXCEPT genre_rows' test portion — use a standard
    # stratified split of genre_rows only, so genre is controlled for on
    # both sides.
    train_rows, test_rows = train_test_split(
        genre_rows, test_size=0.3, stratify=[r["label_int"] for r in genre_rows], random_state=42,
    )

    X_train, y_train = to_xy(train_rows)
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)

    model = lgb.LGBMClassifier(**_LGBM_PARAMS)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=ConvergenceWarning)
        model.fit(X_train_scaled, y_train)

    def eval_subset(subset_rows: list[dict], label: str) -> dict:
        if len(subset_rows) < 10 or len(set(r["label_int"] for r in subset_rows)) < 2:
            return {"subset": label, "n": len(subset_rows), "note": "too small / single class"}
        X, y = to_xy(subset_rows)
        X_scaled = scaler.transform(X)
        y_prob = model.predict_proba(X_scaled)[:, 1]
        fpr, tpr, thr = roc_curve(y, y_prob)
        threshold = float(thr[np.argmax(tpr - fpr)])
        y_pred = (y_prob >= threshold).astype(int)
        return {
            "subset": label,
            "n": len(subset_rows),
            "roc_auc": round(float(roc_auc_score(y, y_prob)), 4),
            "f1": round(float(f1_score(y, y_pred, zero_division=0)), 4),
            "balanced_accuracy": round(float(balanced_accuracy_score(y, y_pred)), 4),
        }

    results = [eval_subset(test_rows, "natural_test_split")]

    # Genre-balanced subsample: for each genre, cap AI and human counts to
    # min(n_ai, n_human) in that genre so the AI/human ratio is ~1:1 within
    # every genre — removes genre as a usable shortcut.
    by_genre_label: dict[tuple[str, str], list[dict]] = defaultdict(list)
    for r in test_rows:
        by_genre_label[(r["genre"], r["label_int"])].append(r)

    balanced_rows: list[dict] = []
    rng = np.random.RandomState(42)
    genres_present = set(g for g, _ in by_genre_label.keys())
    for genre in genres_present:
        ai_rows = by_genre_label.get((genre, "1"), [])
        human_rows = by_genre_label.get((genre, "0"), [])
        cap = min(len(ai_rows), len(human_rows))
        if cap == 0:
            continue
        ai_sample = list(rng.choice(len(ai_rows), size=cap, replace=False))
        human_sample = list(rng.choice(len(human_rows), size=cap, replace=False))
        balanced_rows.extend(ai_rows[i] for i in ai_sample)
        balanced_rows.extend(human_rows[i] for i in human_sample)

    results.append(eval_subset(balanced_rows, "genre_balanced_test_split"))

    out_path = TABLES_DIR / "genre_bias_test.csv"
    fieldnames = ["subset", "n", "roc_auc", "f1", "balanced_accuracy", "note"]
    with open(out_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for r in results:
            writer.writerow({k: r.get(k) for k in fieldnames})

    print(f"\nGenre bias test written: {out_path}")
    for r in results:
        print(f"  {r}")

    # Per-genre breakdown for the paper's genre table
    per_genre_path = TABLES_DIR / "genre_distribution.csv"
    genre_label_counts = Counter((r["genre"], r["label_int"]) for r in genre_rows)
    genres_sorted = sorted(set(g for g, _ in genre_label_counts.keys()))
    with open(per_genre_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["genre", "n_human", "n_ai", "ai_fraction"])
        writer.writeheader()
        for genre in genres_sorted:
            n_human = genre_label_counts.get((genre, "0"), 0)
            n_ai = genre_label_counts.get((genre, "1"), 0)
            total = n_human + n_ai
            writer.writerow({
                "genre": genre, "n_human": n_human, "n_ai": n_ai,
                "ai_fraction": round(n_ai / total, 3) if total else 0,
            })
    print(f"Genre distribution table written: {per_genre_path}")


_SEGMENT_PATTERN = re.compile(
    r"^(?P<base>.+?)[_\-](?:part|seg|segment|chunk|p)[_\-]?0*(?P<idx>\d+)$",
    re.IGNORECASE,
)


def segment_leakage_report(rows: list[dict]) -> None:
    """
    Detect filenames that look like numbered segments of the same source
    track (e.g. track_part1.wav, track_part2.wav) — if such groups exist and
    the current pipeline does group-unaware stratified CV, segments of the
    same track could land in both train and validation folds.
    """
    stems = []
    for r in rows:
        fname = Path(r["file_path"]).stem
        stems.append(fname)

    groups: dict[str, list[str]] = defaultdict(list)
    ungrouped = 0
    for stem in stems:
        m = _SEGMENT_PATTERN.match(stem)
        if m:
            groups[m.group("base")].append(stem)
        else:
            ungrouped += 1

    multi_segment_groups = {k: v for k, v in groups.items() if len(v) > 1}

    lines = []
    lines.append("Segment/track-level leakage check")
    lines.append("=" * 50)
    lines.append(f"Total files: {len(stems)}")
    lines.append(f"Files matching a segment-numbering pattern (base_part1, base_seg2, ...): "
                 f"{sum(len(v) for v in groups.values())}")
    lines.append(f"Distinct base tracks with >1 segment file: {len(multi_segment_groups)}")
    lines.append(f"Files with no recognizable segment suffix: {ungrouped}")
    lines.append("")
    if multi_segment_groups:
        lines.append("Sample multi-segment groups (up to 10):")
        for base, members in list(multi_segment_groups.items())[:10]:
            lines.append(f"  {base}: {members}")
        lines.append("")
        lines.append(
            "FINDING: multi-segment groups exist. The current pipeline uses "
            "StandardScaler + StratifiedKFold with no GroupKFold — segments "
            "of the same base track can be split across train/validation, "
            "which would leak track-level information (recording conditions, "
            "mastering, etc.) across the CV boundary. This affects the "
            "independent-test-set design in nested_cv.py, which groups by "
            "base track before splitting."
        )
    else:
        lines.append(
            "FINDING: no filename-based segment groups detected. Filenames "
            "appear to be one-file-per-track (e.g. gtzan_blues_0000.wav), "
            "so standard stratified splitting does not leak track identity "
            "via segment numbering. This does not rule out leakage from "
            "duplicate or near-duplicate source audio uploaded under "
            "different filenames, which is not filename-detectable."
        )

    out_path = TABLES_DIR / "segment_leakage_report.txt"
    out_path.write_text("\n".join(lines), encoding="utf-8")
    print(f"\nSegment leakage report written: {out_path}")
    print("\n".join(lines))


def run() -> None:
    rows = _load_rows()
    TABLES_DIR.mkdir(parents=True, exist_ok=True)

    print("=" * 70)
    print("STEP 1/3 — Dataset bias table (per-source composition)")
    print("=" * 70)
    dataset_bias_table(rows)

    print("\n" + "=" * 70)
    print("STEP 2/3 — Genre bias test (natural vs. genre-balanced)")
    print("=" * 70)
    genre_bias_test(rows)

    print("\n" + "=" * 70)
    print("STEP 3/3 — Segment/track leakage check")
    print("=" * 70)
    segment_leakage_report(rows)


if __name__ == "__main__":
    run()