"""
Joins the raw feature matrix with source/genre/generator metadata.

The training pipeline (train_classifier.py) reads DataSet/features.csv, which
contains only the 47 acoustic features plus file_path/label_int. It has no
notion of which AI generator or genre a sample came from, so none of the
reviewer-requested analyses (Leave-One-Generator-Out, dataset/genre bias,
segment-leakage checks) can be run against it directly.

DataSet/metadata.csv carries that provenance (source=generator, genre,
ai_model) but was produced by a separate pass and only covers a subset of
files. This script joins the two by filename and writes
DataSet/features_with_meta.csv, reporting anything that fails to join so the
gap is visible rather than silently dropped.

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

from __future__ import annotations

import csv
import re
import sys
from pathlib import Path

DATASET_DIR = Path("D:/CrownCode/DataSet")
FEATURES_CSV = DATASET_DIR / "features.csv"
METADATA_CSV = DATASET_DIR / "metadata.csv"
OUTPUT_CSV = DATASET_DIR / "features_with_meta.csv"
UNMATCHED_REPORT_CSV = DATASET_DIR / "features_with_meta_unmatched_report.csv"
LABEL_MISMATCH_CSV = DATASET_DIR / "features_with_meta_label_mismatches.csv"

# Files under DataSet/*/ai/echoes_<generator>_*.{wav,mp3} have no row in
# metadata.csv, but the generator is recoverable from the filename prefix.
_ECHOES_PREFIX = re.compile(r"^echoes_([a-zA-Z0-9]+)_", re.IGNORECASE)


def _filename(path: str) -> str:
    return path.replace("\\", "/").rsplit("/", 1)[-1]


def _infer_generator_from_filename(filename: str, label_int: int) -> tuple[str, str] | None:
    """Best-effort (generator, ai_model) recovery for files absent from metadata.csv."""
    m = _ECHOES_PREFIX.match(filename)
    if m:
        return "echoes", f"echoes_{m.group(1).lower()}"
    if label_int == 0:
        return "archive_org", ""
    return None


def build() -> None:
    with open(FEATURES_CSV, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        feature_fieldnames = reader.fieldnames or []
        feature_rows = list(reader)

    with open(METADATA_CSV, "r", encoding="utf-8") as f:
        meta_rows = list(csv.DictReader(f))

    meta_by_filename: dict[str, dict] = {}
    for row in meta_rows:
        meta_by_filename[_filename(row["path"])] = row

    matched_direct = 0
    matched_inferred = 0
    truly_unmatched: list[dict] = []
    label_mismatches: list[dict] = []
    out_fieldnames = feature_fieldnames + [
        "genre", "generator", "ai_model", "meta_sample_rate", "meta_duration_sec",
        "match_source",
    ]

    with open(OUTPUT_CSV, "w", newline="", encoding="utf-8") as f_out:
        writer = csv.DictWriter(f_out, fieldnames=out_fieldnames)
        writer.writeheader()

        for row in feature_rows:
            fname = _filename(row["file_path"])
            meta = meta_by_filename.get(fname)

            if meta is not None:
                meta_label_int = 1 if meta["label"] == "ai" else 0
                if int(row["label_int"]) != meta_label_int:
                    mismatch_row = dict(row)
                    mismatch_row["features_label_int"] = row["label_int"]
                    mismatch_row["metadata_label"] = meta["label"]
                    mismatch_row["generator"] = meta["source"]
                    mismatch_row["genre"] = meta["genre"]
                    label_mismatches.append(mismatch_row)
                    # features.csv (the file the training pipeline actually
                    # reads) is treated as the ground-truth label; metadata's
                    # label disagreement is recorded but does not override it.

                out_row = dict(row)
                out_row["genre"] = meta["genre"]
                out_row["generator"] = meta["source"]
                out_row["ai_model"] = meta["ai_model"]
                out_row["meta_sample_rate"] = meta["sample_rate"]
                out_row["meta_duration_sec"] = meta["duration_sec"]
                out_row["match_source"] = "metadata_csv"
                writer.writerow(out_row)
                matched_direct += 1
                continue

            inferred = _infer_generator_from_filename(fname, int(row["label_int"]))
            if inferred is not None:
                generator, ai_model = inferred
                out_row = dict(row)
                out_row["genre"] = ""
                out_row["generator"] = generator
                out_row["ai_model"] = ai_model
                out_row["meta_sample_rate"] = ""
                out_row["meta_duration_sec"] = ""
                out_row["match_source"] = "filename_inferred"
                writer.writerow(out_row)
                matched_inferred += 1
                continue

            truly_unmatched.append(row)

    if truly_unmatched:
        with open(UNMATCHED_REPORT_CSV, "w", newline="", encoding="utf-8") as f_out:
            writer = csv.DictWriter(f_out, fieldnames=feature_fieldnames)
            writer.writeheader()
            writer.writerows(truly_unmatched)

    if label_mismatches:
        mismatch_fieldnames = feature_fieldnames + [
            "features_label_int", "metadata_label", "generator", "genre",
        ]
        with open(LABEL_MISMATCH_CSV, "w", newline="", encoding="utf-8") as f_out:
            writer = csv.DictWriter(f_out, fieldnames=mismatch_fieldnames)
            writer.writeheader()
            writer.writerows(label_mismatches)

    total_matched = matched_direct + matched_inferred
    print(f"features.csv rows:            {len(feature_rows)}")
    print(f"metadata.csv rows:            {len(meta_rows)}")
    print(f"matched via metadata.csv:     {matched_direct}")
    print(f"matched via filename inference: {matched_inferred}")
    print(f"total matched:                {total_matched} / {len(feature_rows)}")
    print(f"truly unmatched:              {len(truly_unmatched)}")
    print(f"label mismatches (features.csv vs metadata.csv label): {len(label_mismatches)}")
    print(f"\nOutput:            {OUTPUT_CSV}")
    if truly_unmatched:
        print(f"Unmatched report:  {UNMATCHED_REPORT_CSV}")
    if label_mismatches:
        print(f"Label mismatches:  {LABEL_MISMATCH_CSV}")
        gens = {}
        for m in label_mismatches:
            gens[m["generator"]] = gens.get(m["generator"], 0) + 1
        print(f"  by generator: {gens}")


if __name__ == "__main__":
    build()