"""실제 P Formula v1만 사용해 0.6 formula-domain residual adapter를 GPU 학습한다."""

from __future__ import annotations

import argparse
from collections import Counter
from copy import deepcopy
from datetime import datetime, timezone
from hashlib import sha256
import json
from pathlib import Path
import random
import sys
from typing import Any

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler

PROJECT_ROOT = Path(__file__).parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
for path in (PROJECT_ROOT, SOURCE_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from math_grid_drawer.research.external_corpus import read_jsonl
from math_grid_drawer.research.p_formula_dataset06 import (
    PFormulaTensorBatch06,
    materialize_p_formula_split06,
    p_formula_release_metrics06,
    p_formula_seed_gate06,
)
from math_grid_drawer.research.p_formula_gate06 import audit_p_formula_records06
from math_grid_drawer.research.skeleton_adapter06 import SkeletonTrajectoryAdapter06
from scripts.audit_math_ink_06_case_context import _load_model06
from scripts.train_math_ink_06_formula_adapter import (
    _forward06,
    _metrics06,
    _targets06,
)


def _parse_args() -> argparse.Namespace:
    """필요 변수: P Formula data·제품 adapter·학습/gate 설정. 작동 원리: 재현 가능한 P-only CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Train Math Ink 0.6 P formula adapter")
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--adapter", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, required=True)
    parser.add_argument("--epochs", type=int, default=16)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--learning-rate", type=float, default=4e-4)
    parser.add_argument("--weight-decay", type=float, default=2e-3)
    parser.add_argument("--exact-loss-weight", type=float, default=0.10)
    parser.add_argument("--context-dropout", type=float, default=0.30)
    parser.add_argument("--hidden-size", type=int, default=64)
    parser.add_argument("--patience", type=int, default=4)
    parser.add_argument("--minimum-independent-sources", type=int, default=2)
    parser.add_argument("--top1-minimum", type=float, default=0.92)
    parser.add_argument("--top5-minimum", type=float, default=0.99)
    parser.add_argument("--macro-f1-minimum", type=float, default=0.90)
    parser.add_argument("--writer-floor-minimum", type=float, default=0.75)
    parser.add_argument("--missing-drop-maximum-pp", type=float, default=3.0)
    parser.add_argument("--skip-test", action="store_true")
    parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
    return parser.parse_args()


def _seed06(seed: int) -> None:
    """필요 변수: seed. 작동 원리: Python·NumPy·PyTorch 난수를 함께 고정한다."""

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _file_sha25606(path: Path) -> str:
    """필요 변수: P Formula JSONL. 작동 원리: seed 간 동일 corpus를 증명할 byte-level SHA-256을 계산한다."""

    digest = sha256()
    with path.open("rb") as file:
        for chunk in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _source_label_loader06(
    batch: PFormulaTensorBatch06,
    exact_targets: Tensor,
    family_targets: Tensor,
    *,
    batch_size: int,
    seed: int,
) -> DataLoader:
    """필요 변수: 학습 batch·label target·seed. 작동 원리: source와 exact label 빈도를 함께 완화한 sampler를 만든다."""

    label_counts = Counter(int(value) for value in exact_targets.tolist())
    source_counts = Counter(batch.source_ids)
    weights = torch.tensor([
        1.0
        / (
            max(label_counts[int(label)], 1) ** 0.5
            * max(source_counts[source], 1) ** 0.5
        )
        for label, source in zip(
            exact_targets.tolist(),
            batch.source_ids,
            strict=True,
        )
    ], dtype=torch.float32)
    weights /= weights.mean().clamp_min(1e-8)
    sampler = WeightedRandomSampler(
        weights,
        num_samples=len(weights),
        replacement=True,
        generator=torch.Generator().manual_seed(seed),
    )
    return DataLoader(
        TensorDataset(batch.features, exact_targets, family_targets),
        batch_size=batch_size,
        sampler=sampler,
    )


def _release_metrics06(
    exact_logits: Tensor,
    batch: PFormulaTensorBatch06,
    exact_targets: Tensor,
    labels: tuple[str, ...],
) -> dict[str, Any]:
    """필요 변수: exact logits·P batch·targets. 작동 원리: 공통 release metric 호출의 identity 인자를 고정한다."""

    return p_formula_release_metrics06(
        exact_logits,
        exact_targets,
        labels=labels,
        writer_ids=batch.writer_ids,
        source_ids=batch.source_ids,
        timestamp_missing=batch.timestamp_missing,
        pressure_missing=batch.pressure_missing,
    )


def main() -> None:
    """필요 변수: P-only split corpus·seed별 product adapter. 작동 원리: validation 선택 후 test를 한 번 평가하고 seed gate를 기록한다."""

    args = _parse_args()
    if not 0.0 <= args.context_dropout <= 1.0:
        raise ValueError("context dropout은 0~1 범위여야 합니다.")
    device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA 학습을 요청했지만 사용할 수 없습니다.")
    _seed06(args.seed)

    records = list(read_jsonl(args.data))
    data_sha256 = _file_sha25606(args.data)
    audit = audit_p_formula_records06(
        records,
        minimum_independent_sources=args.minimum_independent_sources,
    )
    if not audit["eligible_for_product_evaluation"]:
        raise ValueError(
            "P Formula preflight 실패: "
            + json.dumps(audit["issues"][:10], ensure_ascii=False),
        )
    adapter_payload = torch.load(args.adapter, map_location="cpu", weights_only=False)
    base_checkpoint = Path(str(adapter_payload["base_checkpoint"]))
    if not base_checkpoint.is_absolute():
        base_checkpoint = PROJECT_ROOT / base_checkpoint
    engine, online_adapter = _load_model06(base_checkpoint, args.adapter, device)
    for parameter in engine.model.parameters():
        parameter.requires_grad_(False)
    for parameter in online_adapter.parameters():
        parameter.requires_grad_(False)
    labels = tuple(str(label) for label in engine.labels)
    by_split = {
        split: [record for record in records if str(record["split"]) == split]
        for split in ("training", "validation", "test")
    }
    train_batch = materialize_p_formula_split06(
        by_split["training"],
        allowed_labels=labels,
    )
    validation_batch = materialize_p_formula_split06(
        by_split["validation"],
        allowed_labels=labels,
    )
    test_batch = (
        None
        if args.skip_test
        else materialize_p_formula_split06(by_split["test"], allowed_labels=labels)
    )
    train_exact, train_family = _targets06(
        train_batch.truths,
        engine.labels,
        engine.family_labels,
    )
    validation_exact, validation_family = _targets06(
        validation_batch.truths,
        engine.labels,
        engine.family_labels,
    )
    test_exact, test_family = (
        _targets06(test_batch.truths, engine.labels, engine.family_labels)
        if test_batch is not None
        else (None, None)
    )
    loader = _source_label_loader06(
        train_batch,
        train_exact,
        train_family,
        batch_size=args.batch_size,
        seed=args.seed,
    )

    formula_adapter = SkeletonTrajectoryAdapter06(
        hidden_size=args.hidden_size,
    ).to(device)
    optimizer = torch.optim.AdamW(
        formula_adapter.parameters(),
        lr=args.learning_rate,
        weight_decay=args.weight_decay,
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=max(args.epochs, 1),
        eta_min=args.learning_rate * 0.1,
    )
    baseline_logits = _forward06(
        engine.model,
        online_adapter,
        torch.nn.Identity().to(device),
        validation_batch.features,
        device=device,
        batch_size=args.batch_size,
    )
    baseline_validation = _metrics06(
        *baseline_logits,
        validation_exact,
        validation_family,
        engine.labels,
    )
    best_key = (-1.0, -1.0)
    best_state: dict[str, Tensor] | None = None
    best_epoch = 0
    stale = 0
    history = []
    for epoch in range(1, args.epochs + 1):
        formula_adapter.train()
        losses = []
        for features, exact_target, family_target in loader:
            features = features.to(device)
            exact_target = exact_target.to(device)
            family_target = family_target.to(device)
            if args.context_dropout:
                drop = torch.rand(len(features), device=device) < args.context_dropout
                features = features.clone()
                features[drop, :, 10:15] = 0.0
            optimizer.zero_grad(set_to_none=True)
            with torch.no_grad():
                online = online_adapter(features)
            exact_logits, family_logits = engine.model.classify_trajectory(
                formula_adapter(online),
            )
            loss = (
                torch.nn.functional.cross_entropy(family_logits, family_target)
                + args.exact_loss_weight
                * torch.nn.functional.cross_entropy(exact_logits, exact_target)
            )
            loss.backward()
            torch.nn.utils.clip_grad_norm_(formula_adapter.parameters(), 2.0)
            optimizer.step()
            losses.append(float(loss.detach()))
        scheduler.step()
        validation_logits = _forward06(
            engine.model,
            online_adapter,
            formula_adapter,
            validation_batch.features,
            device=device,
            batch_size=args.batch_size,
        )
        validation_metrics = _metrics06(
            *validation_logits,
            validation_exact,
            validation_family,
            engine.labels,
        )
        row = {
            "epoch": epoch,
            "loss": sum(losses) / max(len(losses), 1),
            "validation": validation_metrics,
        }
        history.append(row)
        print(json.dumps(row, ensure_ascii=False), flush=True)
        key = (
            float(validation_metrics["family_head_top1"]),
            float(validation_metrics["visual_family_top1"]),
        )
        if key > best_key:
            best_key = key
            best_epoch = epoch
            best_state = deepcopy({
                name: value.detach().cpu()
                for name, value in formula_adapter.state_dict().items()
            })
            stale = 0
        else:
            stale += 1
        if stale >= args.patience:
            break
    if best_state is None:
        raise RuntimeError("P formula adapter checkpoint가 선택되지 않았습니다.")
    formula_adapter.load_state_dict(best_state)
    selected_validation_logits = _forward06(
        engine.model,
        online_adapter,
        formula_adapter,
        validation_batch.features,
        device=device,
        batch_size=args.batch_size,
    )
    selected_validation = _release_metrics06(
        selected_validation_logits[0],
        validation_batch,
        validation_exact,
        labels,
    )
    if test_batch is not None and test_exact is not None and test_family is not None:
        test_logits = _forward06(
            engine.model,
            online_adapter,
            formula_adapter,
            test_batch.features,
            device=device,
            batch_size=args.batch_size,
        )
        official_test = _release_metrics06(
            test_logits[0],
            test_batch,
            test_exact,
            labels,
        )
        seed_gate = p_formula_seed_gate06(
            official_test,
            top1_minimum=args.top1_minimum,
            top5_minimum=args.top5_minimum,
            macro_f1_minimum=args.macro_f1_minimum,
            writer_floor_minimum=args.writer_floor_minimum,
            missing_drop_maximum_pp=args.missing_drop_maximum_pp,
        )
    else:
        official_test = None
        seed_gate = None

    args.output.mkdir(parents=True, exist_ok=True)
    checkpoint = args.output / "p_formula_adapter.pt"
    torch.save({
        "schema": "aiflow-math-ink-06-p-formula-adapter-v1",
        "state_dict": best_state,
        "hidden_size": args.hidden_size,
        "base_checkpoint": str(base_checkpoint),
        "online_adapter": str(args.adapter),
        "selected_epoch": best_epoch,
        "context_dropout": args.context_dropout,
        "exact_loss_weight": args.exact_loss_weight,
        "track": "P_approved_formula_only",
        "seed_gate_passed": bool(seed_gate and seed_gate["passed"]),
        "product_validation": False,
        "distillation_allowed": False,
        "data_sha256": data_sha256,
    }, checkpoint)
    report = {
        "experiment": "P-MATH-INK-06-FORMULA-ADAPTER-001",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "seed": args.seed,
        "device": str(device),
        "cuda_device": (
            torch.cuda.get_device_name(device) if device.type == "cuda" else None
        ),
        "data": str(args.data),
        "data_sha256": data_sha256,
        "preflight": audit,
        "samples": {
            "training": len(train_batch.truths),
            "validation": len(validation_batch.truths),
            "test": 0 if test_batch is None else len(test_batch.truths),
        },
        "label_support": {
            "training": len(set(train_batch.truths)),
            "validation": len(set(validation_batch.truths)),
            "test": 0 if test_batch is None else len(set(test_batch.truths)),
        },
        "sampler": "inverse_sqrt_source_x_exact_label",
        "baseline_validation": baseline_validation,
        "selected_epoch": best_epoch,
        "selected_validation": selected_validation,
        "official_test": official_test,
        "seed_gate": seed_gate,
        "official_test_skipped": args.skip_test,
        "history": history,
        "checkpoint": checkpoint.name,
        "checkpoint_bytes": checkpoint.stat().st_size,
        "track": "P_approved_formula_only",
        "product_validation": False,
        "distillation_allowed": False,
        "next_gate": "seeds 17/31/47 individual pass, then single student distillation and Android LiteRT validation",
    }
    (args.output / "report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps({
        "seed": args.seed,
        "selected_epoch": best_epoch,
        "validation": selected_validation,
        "official_test": official_test,
        "seed_gate": seed_gate,
        "checkpoint": str(checkpoint),
        "product_validation": False,
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()