#!/usr/bin/env python3
# llmindex.io — fit CLI: response matrices JSON → fitted 2PL parameters JSON
# Author:  Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# License: Proprietary — © Simon-Pierre Boucher, all rights reserved
#
# Pure compute: reads the matrix file written by the worker (no DB access),
# fits a 2PL per domain, writes thetas/SEs and item parameters back to JSON.

from __future__ import annotations

import argparse
import json
import sys

import numpy as np

from llmindex_psycho.bt import fit_bradley_terry
from llmindex_psycho.irt import fit_2pl


def main() -> int:
    parser = argparse.ArgumentParser(description="llmindex 2PL fit")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    with open(args.input) as f:
        payload = json.load(f)

    hyper = payload.get("hyperparams", {})
    max_iter = int(hyper.get("maxIterations", 500))
    tol = float(hyper.get("tolerance", 1e-6))

    out: dict = {"domains": {}}
    for domain, data in payload.get("domains", {}).items():
        models = data["models"]
        items = data["items"]
        matrix = np.array(
            [[np.nan if v is None else float(v) for v in row] for row in data["responses"]],
            dtype=float,
        )
        result = fit_2pl(matrix, max_iterations=max_iter, tolerance=tol)
        out["domains"][domain] = {
            "models": [
                {"slug": slug, "theta": float(t), "se": float(se)}
                for slug, t, se in zip(models, result.theta, result.theta_se)
            ],
            "items": [
                {"id": item_id, "a": float(a), "b": float(b)}
                for item_id, a, b in zip(items, result.a, result.b)
            ],
            "diagnostics": {
                "iterations": result.iterations,
                "converged": result.converged,
                "final_loglik": result.final_loglik,
                **result.diagnostics,
            },
        }
        print(
            f"[fit] {domain}: {len(models)} models × {len(items)} items — "
            f"{'converged' if result.converged else 'max iterations'} @ {result.iterations}",
            file=sys.stderr,
        )

    # Judged duel domains: Bradley-Terry over wins matrices; log-strengths are
    # standardized to a theta-like scale so downstream scoring is uniform.
    out["duel_domains"] = {}
    for domain, data in payload.get("duel_domains", {}).items():
        models = data["models"]
        wins = np.array(data["wins"], dtype=float)
        result = fit_bradley_terry(wins)
        sd = result.log_strength.std() or 1.0
        theta = result.log_strength / sd
        theta_se = result.log_strength_se / sd
        out["duel_domains"][domain] = {
            "models": [
                {"slug": slug, "theta": float(t), "se": float(min(se, 3.0))}
                for slug, t, se in zip(models, theta, theta_se)
            ],
            "diagnostics": {
                "iterations": result.iterations,
                "converged": result.converged,
                "comparisons": float(wins.sum()),
            },
        }
        print(f"[fit] duel {domain}: {len(models)} models, {wins.sum():.0f} comparisons", file=sys.stderr)

    with open(args.output, "w") as f:
        json.dump(out, f, indent=2)
    return 0


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