"""Predictor architecture metadata carried inside safetensors checkpoints.

Stage-1 trainers write ``predictor_config`` into the safetensors metadata so
Stage-2 DMD training and the evaluators can rebuild the exact module (source
layer, input variant, ATC hyper-parameters) without extra side files.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import torch
from safetensors import safe_open

from predictor_training.single_block import (
    SingleBlockPredictor,
    initialize_predictor_block,
)


SUPPORTED_INPUT_VARIANTS = ("self_forcing", "disca", "atc")
LEGACY_PREDICTOR_CONFIG = {"source_layer": 17, "input_variant": "self_forcing"}


def predictor_config_metadata(config: dict[str, Any]) -> dict[str, str]:
    """Serialize a Predictor config as safetensors metadata."""
    return {
        "predictor_config": json.dumps(
            config, sort_keys=True, separators=(",", ":")
        )
    }


def read_predictor_config(
    weights: str | Path, expected_input_variant: str = "auto"
) -> dict[str, Any]:
    """Recover the Predictor architecture from a checkpoint's metadata.

    Checkpoints written before metadata existed are the original concat
    Predictor on Teacher layer 17; they are accepted only when the caller does
    not demand a different variant.
    """
    with safe_open(str(weights), framework="pt", device="cpu") as handle:
        metadata = handle.metadata() or {}
    raw = metadata.get("predictor_config")
    if raw is None:
        if expected_input_variant not in ("auto", "self_forcing"):
            raise ValueError(
                f"{weights} has no predictor_config metadata; cannot assume "
                f"input_variant={expected_input_variant!r}"
            )
        config = dict(LEGACY_PREDICTOR_CONFIG)
    else:
        config = json.loads(raw)
    actual = str(config.get("input_variant", "self_forcing"))
    if actual not in SUPPORTED_INPUT_VARIANTS:
        raise ValueError(f"Unsupported input_variant={actual!r} in {weights}")
    if expected_input_variant != "auto" and actual != expected_input_variant:
        raise ValueError(
            "Predictor input variant mismatch: "
            f"expected={expected_input_variant} actual={actual}"
        )
    config["input_variant"] = actual
    config["source_layer"] = int(config.get("source_layer", 17))
    return config


def build_single_block_predictor(
    teacher: torch.nn.Module,
    config: dict[str, Any],
    *,
    gradient_checkpointing: bool,
    collect_diagnostics: bool = False,
) -> SingleBlockPredictor:
    """Instantiate a SingleBlockPredictor matching ``config`` (weights not loaded)."""
    source_layer = int(config.get("source_layer", 17))
    predictor = SingleBlockPredictor(
        block=initialize_predictor_block(teacher.blocks[source_layer], "teacher_full"),
        dim=teacher.dim,
        gradient_checkpointing=gradient_checkpointing,
        input_variant=str(config.get("input_variant", "self_forcing")),
        gate_mode=str(config.get("gate_mode", "baseline")),
        gate_hidden_dim=int(config.get("gate_hidden_dim", 128)),
        gate_initial_bias=float(config.get("gate_initial_bias", 4.6)),
        gate_floor=float(config.get("gate_floor", 0.0)),
        constant_gate=float(config.get("constant_gate", 1.0)),
        atc_previous_scope=str(config.get("atc_previous_scope", "chunk")),
        atc_freq_dim=int(config.get("atc_freq_dim", 256)),
        atc_mlp_hidden_dim=int(config.get("atc_mlp_hidden_dim", 3072)),
        atc_gate_hidden_dim=int(config.get("atc_gate_hidden_dim", 512)),
        atc_transport_residual_scale=float(
            config.get("atc_transport_residual_scale", 0.1)
        ),
        atc_gate_initial_probability=float(
            config.get("atc_gate_initial_probability", 0.3)
        ),
        atc_collect_diagnostics=collect_diagnostics,
    )
    predictor.source_layer = source_layer
    return predictor