"""Stage-two selective MiniROADA adaptation from a trained Omni residual.

The released input features and Omni states remain frozen. This module only
enables gradients in selected MiniROADA temporal/forecast modules and the
existing residual. Stage-one checkpoints are loaded strictly, without resetting
the learned residual or requiring the original checkpoint to be available.
"""

from __future__ import annotations

from copy import deepcopy
from pathlib import Path
from typing import Any, Literal

import torch

from easycom_head import NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS
from egospeak_omni_residual import (
    EgoSpeakOmniResidual,
    MiniROADABase,
    load_torch_payload,
    normalize_session_name,
)

BaseTrainable = Literal["none", "heads", "gru-heads"]
STAGE_ONE_FORMAT = "egospeak-omni-zero-residual-v1"
STAGE_TWO_FORMAT = "egospeak-omni-selective-finetune-v1"


def validate_training_sessions(values: list[str]) -> None:
    normalized = [normalize_session_name(value) for value in values]
    if not normalized or len(normalized) != len(set(normalized)):
        raise ValueError("Training sessions must be nonempty and unique")
    if not set(normalized).issubset(OFFICIAL_TRAIN_SESSIONS):
        raise ValueError("Training may only use EasyCom Sessions 4--12")


def load_stage_one_checkpoint(
    path: Path,
) -> tuple[EgoSpeakOmniResidual, dict[str, Any]]:
    path = path.expanduser().resolve()
    if not path.is_file():
        raise FileNotFoundError(path)
    payload = load_torch_payload(path)
    if not isinstance(payload, dict) or payload.get("format") != STAGE_ONE_FORMAT:
        raise ValueError(
            "--stage-one-checkpoint must be a stage-one residual checkpoint"
        )
    config = payload.get("model_config")
    if not isinstance(config, dict):
        raise TypeError("Stage-one checkpoint lacks model_config")
    required = {
        "omni_dim",
        "fusion_dim",
        "tail_embedding_dim",
        "num_tail_bins",
        "dropout",
    }
    if set(config) != required or config["num_tail_bins"] != NUM_HORIZONS:
        raise ValueError("Incompatible stage-one residual model_config")
    training = payload.get("training_config")
    if not isinstance(training, dict) or int(training.get("epochs", 0)) <= 0:
        raise ValueError("Stage-one checkpoint lacks a positive training epoch budget")
    validate_training_sessions(payload.get("train_sessions", []))
    base_config = payload.get("base_model_config", {})
    if not isinstance(base_config, dict):
        raise TypeError("Invalid base_model_config")
    # Released v1 checkpoints omit base_model_config and use default MiniROADA.
    model = EgoSpeakOmniResidual(MiniROADABase(**base_config), **config)
    model.load_state_dict(payload["model_state_dict"], strict=True)
    if any(not torch.isfinite(value).all() for value in model.state_dict().values()):
        raise ValueError("Stage-one checkpoint contains nonfinite parameters")
    model.requires_grad_(False)
    model.eval()
    return model, payload


class EgoSpeakOmniFinetune(EgoSpeakOmniResidual):
    """Continuation with explicit base policy and intact stage-one logits."""

    def __init__(
        self, *args: Any, base_trainable: BaseTrainable = "gru-heads", **kwargs: Any
    ):
        super().__init__(*args, **kwargs)
        if base_trainable not in ("none", "heads", "gru-heads"):
            raise ValueError(f"Invalid base-trainable policy: {base_trainable}")
        self.base_trainable = base_trainable
        self.requires_grad_(True)
        self.base.requires_grad_(False)
        if base_trainable in ("heads", "gru-heads"):
            self.base.anticipation_layer.requires_grad_(True)
            self.base.f_classification.requires_grad_(True)
        if base_trainable == "gru-heads":
            self.base.gru.requires_grad_(True)
        self.train(False)

    @classmethod
    def from_stage_one(
        cls,
        reference: EgoSpeakOmniResidual,
        config: dict[str, Any],
        *,
        base_trainable: BaseTrainable = "gru-heads",
    ) -> EgoSpeakOmniFinetune:
        model = cls(deepcopy(reference.base), **config, base_trainable=base_trainable)
        # Copy rather than share weights with the immutable first-stage reference.
        model.load_state_dict(reference.state_dict(), strict=True)
        return model

    def train(self, mode: bool = True) -> EgoSpeakOmniFinetune:
        super().train(mode)
        # Keep frozen layer1's dropout off. Enabling GRU/head train mode does
        # not turn the original frozen RGB+audio projection back on.
        if self.base_trainable in ("heads", "gru-heads"):
            self.base.anticipation_layer.train(mode)
            self.base.f_classification.train(mode)
        if self.base_trainable == "gru-heads":
            self.base.gru.train(mode)
        return self

    def optimizer_groups(
        self, *, fusion_lr: float, base_lr: float
    ) -> list[dict[str, Any]]:
        if fusion_lr <= 0 or base_lr <= 0:
            raise ValueError("Both learning rates must be positive")
        fusion = [
            p
            for name, p in self.named_parameters()
            if p.requires_grad and not name.startswith("base.")
        ]
        base = [
            p
            for name, p in self.named_parameters()
            if p.requires_grad and name.startswith("base.")
        ]
        groups = [{"name": "fusion", "params": fusion, "lr": fusion_lr}]
        if base:
            groups.append({"name": "miniroada", "params": base, "lr": base_lr})
        return groups

    def forward_last(
        self,
        rgb: torch.Tensor,
        audio: torch.Tensor,
        omni: torch.Tensor,
        tail_bin: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        hidden, _ = self.base.encode(rgb, audio)
        ego_hidden = hidden[:, -1, :]
        base_logits = self.base.anticipation_logits(ego_hidden)
        # No no_grad/detach: gradients must reach the selected MiniROADA GRU.
        residual = self._residual_logits(ego_hidden, omni, tail_bin)
        return base_logits + residual, base_logits

    def forward_sequence_chunk(
        self,
        rgb: torch.Tensor,
        audio: torch.Tensor,
        omni: torch.Tensor,
        tail_bin: torch.Tensor,
        recurrent_state: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        hidden, next_state = self.base.encode(rgb, audio, recurrent_state)
        base_logits = self.base.anticipation_logits(hidden)
        residual = self._residual_logits(hidden, omni, tail_bin)
        return base_logits + residual, base_logits, next_state


def base_model_config(base: MiniROADABase) -> dict[str, Any]:
    return {
        "rgb_dim": base.rgb_dim,
        "audio_dim": base.audio_dim,
        "embedding_dim": base.embedding_dim,
        "hidden_dim": base.hidden_dim,
        "num_layers": base.num_layers,
        "num_horizons": base.num_horizons,
        "num_classes": base.num_classes,
        "dropout": base.layer1[-1].p,
    }


def load_stage_two_checkpoint(
    path: Path,
) -> tuple[EgoSpeakOmniFinetune, dict[str, Any]]:
    payload = load_torch_payload(path)
    if not isinstance(payload, dict) or payload.get("format") != STAGE_TWO_FORMAT:
        raise ValueError("Not a stage-two fine-tuning checkpoint")
    model = EgoSpeakOmniFinetune(
        MiniROADABase(**payload["base_model_config"]),
        **payload["model_config"],
        base_trainable=payload["base_trainable"],
    )
    model.load_state_dict(payload["model_state_dict"], strict=True)
    model.eval()
    return model, payload