# -*- coding: utf-8 -*-
"""
collatz_update_noyau.py

Met à jour le noyau en soustrayant les classes couvertes par un certificat de fusion.
Charge le noyau précédent, charge le certificat (classes couvertes),
soustrait et écrit le nouveau noyau.

CLI: --fusion CERT_JSON --previous NOYAU_JSON --output OUTPUT_JSON
"""

from __future__ import annotations
from pathlib import Path
import argparse
import csv
import json


def load_noyau(path: str) -> list[int]:
    """Load noyau from JSON: list of residues or dict with R*_after / noyau / residues."""
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    if isinstance(data, list):
        return [int(x) for x in data]
    if isinstance(data, dict):
        for key in ("R25_after", "R24_after", "noyau", "residues", "uncovered"):
            if key in data and isinstance(data[key], list):
                return [int(x) for x in data[key]]
        raise ValueError(f"Noyau JSON: no known key in {list(data.keys())}")
    raise ValueError("Noyau JSON must be a list or dict with residue list")


def load_covered_classes(path: str) -> set[int]:
    """
    Load covered classes from fusion certificate.
    Supports: JSON with 'covered', 'covered_classes', or list; CSV with classe_mod_2^m column.
    """
    p = Path(path)
    suffix = p.suffix.lower()

    if suffix == ".json":
        data = json.loads(p.read_text(encoding="utf-8"))
        if isinstance(data, list):
            return {int(x) for x in data}
        if isinstance(data, dict):
            for key in ("covered", "covered_classes", "classe_mod_2^m"):
                if key in data:
                    val = data[key]
                    if isinstance(val, list):
                        return {int(x) for x in val}
            # Try top-level keys like "11", "12" for per-horizon data
            covered: set[int] = set()
            for v in data.values():
                if isinstance(v, dict) and "covered" in v and isinstance(v["covered"], list):
                    covered.update(int(x) for x in v["covered"])
                elif isinstance(v, list):
                    covered.update(int(x) for x in v if isinstance(x, (int, float)))
            if covered:
                return covered
        raise ValueError(f"Fusion cert JSON: no covered classes found in {list(data.keys()) if isinstance(data, dict) else 'list'}")

    if suffix == ".csv":
        covered: set[int] = set()
        with p.open("r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            col = "classe_mod_2^m"
            for row in reader:
                if col in row:
                    covered.add(int(row[col]))
        return covered

    raise ValueError(f"Fusion cert must be .json or .csv, got {suffix}")


def _get_palier(path: str) -> int | None:
    """Extract palier from noyau JSON if present."""
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    if isinstance(data, dict) and "palier" in data:
        return int(data["palier"])
    return None


def run_update_noyau(fusion_cert: str, previous_noyau: str, output: str) -> None:
    noyau = set(load_noyau(previous_noyau))
    covered = load_covered_classes(fusion_cert)
    new_noyau = sorted(noyau - covered)
    palier = _get_palier(previous_noyau)

    out_path = Path(output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    payload: list[int] | dict = new_noyau
    if palier is not None:
        payload = {"noyau": new_noyau, "palier": palier}
    out_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")

    print(f"Previous noyau: {len(noyau)}, covered: {len(covered)}, new noyau: {len(new_noyau)}")
    print(f"Wrote: {out_path}")


def main() -> None:
    ap = argparse.ArgumentParser(description="Update noyau by subtracting fusion-covered classes")
    ap.add_argument("--fusion", required=True, help="Path to fusion certificate (JSON or CSV with covered classes)")
    ap.add_argument("--previous", required=True, help="Path to previous noyau JSON")
    ap.add_argument("--output", required=True, help="Path to output new noyau JSON")
    args = ap.parse_args()

    run_update_noyau(
        fusion_cert=args.fusion,
        previous_noyau=args.previous,
        output=args.output,
    )


if __name__ == "__main__":
    main()