"""Fixed-stratum generation and plotting for v6 protocol validation."""

from __future__ import annotations

import hashlib
from dataclasses import fields, replace
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch

from pff_pk.tasks.batch import FlowBatch
from pff_pk.tasks.batching import inverse_concentration

CONTEXT_COLOR = "#ff7f0e"
GENERATED_COLOR = "#1f77b4"
TRUTH_COLOR = "#202020"

PROTOCOL_CATEGORIES = {
    "identity": (
        "reference_identity",
        "global_identity",
        "add_on_identity",
    ),
    "dose_scaling": (
        "reference_to_lower_scale",
        "reference_to_higher_scale",
        "lower_scale_to_reference",
        "higher_scale_to_reference",
        "global_scale_cross",
    ),
    "add_on_dose": (
        "reference_to_early_add_on",
        "reference_to_late_add_on",
        "add_on_to_reference",
        "add_on_cross",
    ),
    "cross_intervention": (
        "global_to_add_on",
        "add_on_to_global",
    ),
}
ACQUISITIONS = ("unscheduled", "exact", "pseudo_scheduled")


def select_representatives(bank: dict) -> list[tuple[str, int, int]]:
    """Select one deterministic bank row for every protocol/acquisition stratum."""
    candidates: dict[str, list[tuple[int, int]]] = {}
    for record_index, record in enumerate(bank["batches"]):
        candidates.setdefault(record["stratum"], []).extend(
            (record_index, row) for row in range(len(record["study_id"]))
        )
    expected = {
        f"{case}__{acquisition}"
        for cases in PROTOCOL_CATEGORIES.values()
        for case in cases
        for acquisition in ACQUISITIONS
    }
    if set(candidates) != expected:
        raise ValueError("validation bank does not contain the expected 42 v6 strata")
    selected = []
    for stratum in sorted(candidates):
        digest = hashlib.blake2b(stratum.encode(), digest_size=8).digest()
        choices = candidates[stratum]
        record_index, row = choices[int.from_bytes(digest, "little") % len(choices)]
        selected.append((stratum, record_index, row))
    return selected


def protocol_seed(stratum: str) -> int:
    """Return the fixed source-process seed for a protocol stratum."""
    digest = hashlib.blake2b(stratum.encode(), digest_size=4).digest()
    return 2026 + int.from_bytes(digest, "little")


def _select(batch: FlowBatch, row: int, device: torch.device) -> FlowBatch:
    index = torch.tensor([row], dtype=torch.long)
    values = {}
    for field in fields(batch):
        value = getattr(batch, field.name)
        if value is None:
            values[field.name] = None
        elif isinstance(value, torch.Tensor):
            values[field.name] = value.index_select(0, index).to(device)
        else:
            values[field.name] = [value[row]]
    return replace(batch, **values)


def _event_arrays(batch: FlowBatch) -> tuple[np.ndarray, np.ndarray]:
    assert batch.context_protocol_events is not None
    assert batch.context_protocol_event_mask is not None
    assert batch.target_protocol_events is not None
    assert batch.target_protocol_event_mask is not None
    context_mask = batch.context_protocol_event_mask[0, 0].cpu().numpy()
    target_mask = batch.target_protocol_event_mask[0].cpu().numpy()
    context = batch.context_protocol_events[0, 0].cpu().numpy()[context_mask]
    target = batch.target_protocol_events[0].cpu().numpy()[target_mask]
    return context, target


def generate_protocol_draws(
    model, cpu_batch: FlowBatch, row: int, samples: int, seed: int
) -> dict:
    """Generate persistent conditional draws for one fixed bank row."""
    device = next(model.parameters()).device
    batch = _select(cpu_batch, row, device)
    repeated_values = {}
    for field in fields(batch):
        value = getattr(batch, field.name)
        if value is None:
            repeated_values[field.name] = None
        elif isinstance(value, torch.Tensor):
            repeated_values[field.name] = value.expand(samples, *value.shape[1:]).clone()
        else:
            repeated_values[field.name] = value * samples
    repeated = replace(batch, **repeated_values)
    cuda_devices = [device.index] if device.type == "cuda" else []
    with torch.random.fork_rng(devices=cuda_devices), torch.inference_mode():
        torch.manual_seed(seed)
        generated = inverse_concentration(model.sample(repeated), repeated)[..., 0]
    context = inverse_concentration(batch.context_value, batch)[0, ..., 0]
    target = inverse_concentration(batch.target_value, batch)[0, ..., 0]
    context_events, target_events = _event_arrays(batch)
    return {
        "study_id": batch.study_id[0],
        "context_time": batch.context_time[0, ..., 0].cpu(),
        "context_value": context.cpu(),
        "context_mask": batch.context_mask[0].cpu(),
        "target_time": batch.target_time[0, :, 0].cpu(),
        "target_value": target.cpu(),
        "target_mask": batch.target_mask[0].cpu(),
        "generated": generated.cpu(),
        "context_events": torch.from_numpy(context_events),
        "target_events": torch.from_numpy(target_events),
    }


def _schedule(events: torch.Tensor) -> str:
    labels = []
    for time, amount, duration, oral in events.tolist():
        route = "PO" if oral >= 0.5 else "IV"
        label = f"{route} {amount:.3g}× @ {time:.3g}"
        if duration > 0:
            label += f" over {duration:.3g}"
        labels.append(label)
    return " + ".join(labels) if labels else "no dose"


def _human_case(value: str) -> str:
    return value.replace("_", " ").replace("to", "→")


def plot_protocol_category(
    records: dict[str, dict],
    category: str,
    cases: tuple[str, ...],
    output: Path,
    step: int,
) -> tuple[Path, Path]:
    """Render one protocol category over the three acquisition designs."""
    figure, axes = plt.subplots(
        len(cases),
        len(ACQUISITIONS),
        figsize=(15.6, 3.25 * len(cases) + 1.4),
        squeeze=False,
        constrained_layout=True,
    )
    legend_handles = None
    for row, case in enumerate(cases):
        for column, acquisition in enumerate(ACQUISITIONS):
            axis = axes[row, column]
            record = records[f"{case}__{acquisition}"]
            context_time = record["context_time"].numpy()
            context_value = record["context_value"].numpy()
            context_mask = record["context_mask"].numpy()
            for times, values, mask in zip(
                context_time, context_value, context_mask, strict=True
            ):
                axis.plot(
                    times[mask],
                    values[mask],
                    color=CONTEXT_COLOR,
                    linewidth=0.5,
                    alpha=0.28,
                    zorder=1,
                )
            target_mask = record["target_mask"].numpy()
            target_time = record["target_time"].numpy()[target_mask]
            for generated in record["generated"].numpy():
                axis.plot(
                    target_time,
                    generated[target_mask],
                    color=GENERATED_COLOR,
                    linewidth=1.0,
                    alpha=0.72,
                    zorder=3,
                )
            truth = axis.plot(
                target_time,
                record["target_value"].numpy()[target_mask],
                color=TRUTH_COLOR,
                linestyle="--",
                linewidth=1.5,
                zorder=4,
                label="held-out truth",
            )[0]
            context_line = axis.plot([], [], color=CONTEXT_COLOR, label="context cohort")[0]
            generated_line = axis.plot([], [], color=GENERATED_COLOR, label="generated draws")[0]
            legend_handles = (context_line, generated_line, truth)
            axis.text(
                0.02,
                0.97,
                f"study {record['study_id'][-8:]}\n"
                f"context: {_schedule(record['context_events'])}\n"
                f"target: {_schedule(record['target_events'])}",
                transform=axis.transAxes,
                fontsize=7.2,
                ha="left",
                va="top",
                bbox={"facecolor": "white", "alpha": 0.78, "edgecolor": "none", "pad": 2},
                zorder=10,
            )
            if row == 0:
                axis.set_title(
                    acquisition.replace("_", " "),
                    fontsize=10.5,
                    fontweight="bold",
                    pad=8,
                )
            if column == 0:
                axis.set_ylabel(f"{_human_case(case)}\nC / Cref", fontsize=8)
            if row == len(cases) - 1:
                axis.set_xlabel("dimensionless time tau", fontsize=8)
            axis.set_yscale("log")
            axis.tick_params(labelsize=7)
            axis.grid(alpha=0.13, linewidth=0.4)
    figure.suptitle(
        f"v6 validation: {category.replace('_', ' ')} | checkpoint step {step}\n"
        "one held-out family per protocol/acquisition stratum; five conditional draws",
        fontsize=13,
    )
    if legend_handles is not None:
        figure.legend(
            legend_handles,
            [handle.get_label() for handle in legend_handles],
            loc="outside lower center",
            ncol=3,
            frameon=False,
        )
    png = output / f"protocol_validation_{category}_step{step:06d}.png"
    pdf = png.with_suffix(".pdf")
    figure.savefig(png, dpi=220, facecolor="white")
    figure.savefig(pdf, facecolor="white")
    plt.close(figure)
    return png, pdf