"""Distributed fixed-bank diagnostics for v6 dose-intervention tasks."""

from __future__ import annotations

import json
import os
from pathlib import Path

import lightning.pytorch as pl
import torch
import torch.distributed as dist

from pff_pk.evaluation.bank import load_bank, record_to_batch
from pff_pk.inference.protocol_validation import (
    PROTOCOL_CATEGORIES,
    generate_protocol_draws,
    plot_protocol_category,
    protocol_seed,
    select_representatives,
)


class ProtocolValidationDiagnostics(pl.Callback):
    """Persist and plot the same 42 protocol strata at fixed step intervals.

    Every DDP rank evaluates a disjoint slice of the fixed validation bank. The
    resulting CPU records are gathered on rank zero, saved with full provenance,
    rendered into the four protocol panels, and logged to the existing Comet run.
    The bank and per-stratum random seeds remain fixed across checkpoints.
    """

    def __init__(
        self,
        bank_path: str | Path,
        output_directory: str | Path,
        *,
        interval: int,
        samples: int,
    ) -> None:
        super().__init__()
        self.bank_path = Path(bank_path)
        self.output_directory = Path(output_directory)
        self.interval = interval
        self.samples = samples
        self._logged: set[int] = set()
        self._bank: dict | None = None

    def _completed_milestones(self) -> set[int]:
        """Recover completed milestones so resumed jobs do not duplicate diagnostics."""
        completed: set[int] = set()
        if not self.output_directory.is_dir():
            return completed
        for manifest in self.output_directory.glob("step*/manifest.json"):
            try:
                payload = json.loads(manifest.read_text(encoding="utf-8"))
                step = int(payload["checkpoint_step"])
                milestone = int(
                    payload.get("trigger_milestone", step // self.interval * self.interval)
                )
            except (KeyError, TypeError, ValueError, json.JSONDecodeError, OSError):
                continue
            if milestone > 0:
                completed.add(milestone)
        return completed

    def _load_bank(self) -> dict:
        if self._bank is None:
            self._bank = load_bank(self.bank_path)
        return self._bank

    def on_validation_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
        step = int(trainer.global_step)
        milestone = step // self.interval * self.interval
        if milestone == 0:
            return
        if not self._logged:
            self._logged.update(self._completed_milestones())
        if milestone in self._logged:
            return

        bank = self._load_bank()
        rank = int(trainer.global_rank)
        world_size = int(trainer.world_size)
        representatives = select_representatives(bank)
        selected = representatives[rank::world_size]
        device = pl_module.device
        cuda_devices = [device.index] if device.type == "cuda" else []
        was_training = pl_module.training
        local_records = []
        pl_module.eval()
        try:
            with torch.random.fork_rng(devices=cuda_devices), torch.inference_mode():
                for stratum, record_index, row in selected:
                    batch = record_to_batch(bank["batches"][record_index], torch.device("cpu"))
                    result = generate_protocol_draws(
                        pl_module, batch, row, self.samples, protocol_seed(stratum)
                    )
                    local_records.append({"stratum": stratum, **result})
        finally:
            pl_module.train(was_training)

        gathered: list[list[dict] | None] | None
        if dist.is_available() and dist.is_initialized():
            gathered = [None] * world_size if trainer.is_global_zero else None
            dist.gather_object(local_records, gathered, dst=0)
        else:
            gathered = [local_records]
        if not trainer.is_global_zero:
            self._logged.add(milestone)
            return

        assert gathered is not None
        records = {
            record["stratum"]: record
            for rank_records in gathered
            for record in (rank_records or [])
        }
        if len(records) != 42:
            raise RuntimeError(f"expected 42 protocol strata, found {len(records)}")

        output = self.output_directory / f"step{step:06d}"
        output.mkdir(parents=True, exist_ok=True)
        draws_path = output / "protocol_draws.pt"
        torch.save(
            {
                "format": "v6-protocol-validation-draws-v1",
                "checkpoint_step": step,
                "trigger_milestone": milestone,
                "validation_bank": str(self.bank_path.resolve()),
                "validation_bank_sha256": bank["sha256"],
                "samples_per_task": self.samples,
                "records": [records[key] for key in sorted(records)],
            },
            draws_path,
        )
        figures = [
            plot_protocol_category(records, category, cases, output, step)
            for category, cases in PROTOCOL_CATEGORIES.items()
        ]
        manifest = {
            "format": "v6-protocol-validation-draws-v1",
            "checkpoint_step": step,
            "trigger_milestone": milestone,
            "validation_bank": str(self.bank_path.resolve()),
            "validation_bank_sha256": bank["sha256"],
            "samples_per_task": self.samples,
            "strata": sorted(records),
            "draws": draws_path.name,
            "figures": [path.name for pair in figures for path in pair],
        }
        manifest_path = output / "manifest.json"
        manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")

        comet_step = step + int(os.environ.get("PFFF_COMET_STEP_OFFSET", "0"))
        experiment = trainer.logger.experiment
        for png, pdf in figures:
            category = png.stem.removeprefix("protocol_validation_").rsplit("_step", 1)[0]
            experiment.log_image(
                image_data=str(png),
                name=f"validation/protocol_stratified/{category}.png",
                step=comet_step,
            )
            experiment.log_asset(
                str(pdf),
                file_name=f"validation/protocol_stratified/step_{step}/{pdf.name}",
                step=comet_step,
            )
        experiment.log_asset(
            str(draws_path),
            file_name=f"validation/protocol_stratified/step_{step}/protocol_draws.pt",
            step=comet_step,
        )
        experiment.log_asset(
            str(manifest_path),
            file_name=f"validation/protocol_stratified/step_{step}/manifest.json",
            step=comet_step,
        )
        self._logged.add(milestone)