#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

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

from dovla_cil.generation.tangent_memory import (  # noqa: E402
    evaluate_tangent_memory_generator,
    load_tangent_targets,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Evaluate a train-only positive-tangent memory generator against heldout "
            "same-state positive tangent support."
        )
    )
    parser.add_argument("--targets", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--k-values", default="1,2,4,8,16")
    parser.add_argument("--thresholds", default="0.05,0.1,0.2,0.4")
    parser.add_argument("--val-fraction", type=float, default=0.2)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--diversity-weight", type=float, default=0.25)
    parser.add_argument("--no-groups", action="store_true", help="Drop per-group rows from output.")
    args = parser.parse_args(argv)

    try:
        k_values = _parse_ints(args.k_values)
        thresholds = _parse_floats(args.thresholds)
    except ValueError as exc:
        parser.error(str(exc))

    examples = load_tangent_targets(args.targets)
    report = evaluate_tangent_memory_generator(
        examples,
        k_values=k_values,
        thresholds=thresholds,
        val_fraction=args.val_fraction,
        seed=args.seed,
        diversity_weight=args.diversity_weight,
    )
    report["targets"] = str(args.targets)
    if args.no_groups:
        report.pop("groups", None)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, indent=2) + "\n")
    summary = {
        key: value
        for key, value in report.items()
        if key not in {"groups"}
    }
    print(json.dumps(summary, indent=2))
    print(f"Wrote {args.out}")
    return 0


def _parse_ints(value: str) -> tuple[int, ...]:
    items = tuple(int(item.strip()) for item in value.split(",") if item.strip())
    if not items or any(item <= 0 for item in items):
        raise ValueError("integer list must contain positive values")
    return items


def _parse_floats(value: str) -> tuple[float, ...]:
    items = tuple(float(item.strip()) for item in value.split(",") if item.strip())
    if not items or any(item < 0 for item in items):
        raise ValueError("float list must contain non-negative values")
    return items


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