#!/usr/bin/env python3
"""Audit EgoSpeak EasyCom feature rates, shapes, labels, and raw alignment.

Example:
    python audit_egospeak_features.py \
      --feature-root /path/to/EgoSpeak/dataset/EasyCom \
      --raw-manifest manifests/easycom.json \
      --output reports/egospeak_easycom_audit.json
"""

from __future__ import annotations

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

import numpy as np

from easycom_manifest import ALL_SESSIONS, TASK_FPS, TEST_SESSIONS, TRAIN_SESSIONS


FEATURE_SPECS = {
    "rgb_kinetics_resnet50": 2048,
    "flow_kinetics_resnet50": 2048,
    "audio_wav2vec2": 5120,
    "target_perframe": 3,
}


def load_raw_steps(path: Path | None) -> dict[str, int]:
    if path is None:
        return {}
    payload = json.loads(path.read_text())
    return {
        session["session"]: int(session["expected_5hz_steps"])
        for session in payload["sessions"]
        if "expected_5hz_steps" in session
    }


def audit(
    feature_root: Path,
    *,
    raw_manifest: Path | None,
) -> dict[str, Any]:
    feature_root = feature_root.resolve()
    raw_steps = load_raw_steps(raw_manifest)
    sessions: list[dict[str, Any]] = []
    issues: list[str] = []
    all_targets: list[np.ndarray] = []

    for session in ALL_SESSIONS:
        arrays: dict[str, np.ndarray] = {}
        shapes: dict[str, list[int]] = {}
        dtypes: dict[str, str] = {}
        for modality, expected_dim in FEATURE_SPECS.items():
            path = feature_root / modality / f"{session}.npy"
            if not path.is_file():
                issues.append(f"{session}: missing {path}")
                continue
            array = np.load(path, mmap_mode="r")
            arrays[modality] = array
            shapes[modality] = list(array.shape)
            dtypes[modality] = str(array.dtype)
            if array.ndim != 2 or array.shape[1] != expected_dim:
                issues.append(
                    f"{session}: {modality} expected [T,{expected_dim}], got {array.shape}"
                )

        if len(arrays) != len(FEATURE_SPECS):
            continue
        lengths = {name: int(array.shape[0]) for name, array in arrays.items()}
        if len(set(lengths.values())) != 1:
            issues.append(f"{session}: unequal modality lengths {lengths}")

        target = np.asarray(arrays["target_perframe"])
        if not np.all((target == 0) | (target == 1)):
            issues.append(f"{session}: target contains values outside {{0,1}}")
        allowed_rows = {(1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0), (0.0, 1.0, 1.0)}
        observed_rows = {tuple(row) for row in np.unique(target, axis=0)}
        unexpected = observed_rows - allowed_rows
        if unexpected:
            issues.append(f"{session}: unexpected target rows {sorted(unexpected)}")
        all_targets.append(target)

        num_steps = lengths["target_perframe"]
        record: dict[str, Any] = {
            "session": session,
            "split": "train" if session in TRAIN_SESSIONS else "test",
            "num_steps": num_steps,
            "duration_seconds_at_5hz": num_steps / TASK_FPS,
            "shapes": shapes,
            "dtypes": dtypes,
            "class_positive_steps": target.sum(axis=0).astype(int).tolist(),
            "overlap_wearer_other_steps": int(
                np.sum((target[:, 1] == 1) & (target[:, 2] == 1))
            ),
        }
        if session in raw_steps:
            record["raw_manifest_expected_steps"] = raw_steps[session]
            record["feature_minus_raw_steps"] = num_steps - raw_steps[session]
            if abs(record["feature_minus_raw_steps"]) > 1:
                issues.append(
                    f"{session}: feature/raw mismatch "
                    f"{record['feature_minus_raw_steps']:+d} steps"
                )
        sessions.append(record)

    combined = np.concatenate(all_targets, axis=0)
    unique_rows, row_counts = np.unique(combined, axis=0, return_counts=True)
    split_steps = {
        split: sum(item["num_steps"] for item in sessions if item["split"] == split)
        for split in ("train", "test")
    }
    return {
        "feature_root": str(feature_root),
        "raw_manifest": str(raw_manifest.resolve()) if raw_manifest is not None else None,
        "task_fps": TASK_FPS,
        "step_seconds": 1.0 / TASK_FPS,
        "sessions": sessions,
        "total_steps": int(combined.shape[0]),
        "total_duration_hours": float(combined.shape[0] / TASK_FPS / 3600),
        "split_steps": split_steps,
        "class_positive_steps": combined.sum(axis=0).astype(int).tolist(),
        "target_row_counts": {
            str(row.astype(int).tolist()): int(count)
            for row, count in zip(unique_rows, row_counts)
        },
        "issues": issues,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--feature-root", type=Path, required=True)
    parser.add_argument("--raw-manifest", type=Path)
    parser.add_argument("--output", type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    report = audit(args.feature_root, raw_manifest=args.raw_manifest)
    if args.output is not None:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(report, indent=2) + "\n")
        print(f"Saved report to {args.output}")

    print(
        f"{report['total_steps']} steps at {report['task_fps']} Hz = "
        f"{report['total_duration_hours']:.6f} hours"
    )
    print(f"split steps: {report['split_steps']}")
    print(f"class-positive steps [background, wearer, other]: {report['class_positive_steps']}")
    print(f"target rows: {report['target_row_counts']}")
    if report["issues"]:
        print(f"Found {len(report['issues'])} issue(s):")
        for issue in report["issues"]:
            print(f"  - {issue}")
    else:
        if report["raw_manifest"] is not None:
            print("All modality lengths, dimensions, labels, and raw steps match.")
        else:
            print("All feature-modality lengths, dimensions, and labels are internally consistent.")


if __name__ == "__main__":
    main()