"""Offline temporal proposals that preserve StarVLA's dataset/episode marginals."""

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
from typing import Any, Sequence

import numpy as np

from retime.sampler import (
    LogicalDraw,
    SampleTrace,
    ValidAnchorManifest,
    _stateless_seed,
)


class TemporalProposal:
    schema_version = 1
    valid_score_kinds = {
        "synthetic",
        "random",
        "shuffle",
        "action_change",
        "residual",
        "likelihood",
    }

    def __init__(
        self,
        manifest: ValidAnchorManifest,
        alpha: float,
        anchor_probabilities: Sequence[Sequence[Sequence[float]]],
        score_kind: str,
        scorer_sha256: str,
        validation_certificate: dict[str, Any] | None = None,
    ) -> None:
        if not 0.0 <= alpha < 1.0:
            raise ValueError("alpha must satisfy 0 <= alpha < 1")
        if score_kind not in self.valid_score_kinds:
            raise ValueError(f"unsupported score_kind: {score_kind}")
        if len(anchor_probabilities) != len(manifest.datasets):
            raise ValueError("proposal dataset count does not match the manifest")

        normalized_probabilities = []
        importance_weights = []
        normalized_masses = []
        for dataset_index, (dataset, dataset_probabilities) in enumerate(
            zip(manifest.datasets, anchor_probabilities)
        ):
            if len(dataset_probabilities) != len(dataset.episodes):
                raise ValueError(
                    f"proposal episode count differs for dataset {dataset_index}"
                )
            normalized_dataset = []
            dataset_weights = []
            normalized_dataset_masses = []
            for episode_index, (episode, proposal_mass) in enumerate(
                zip(dataset.episodes, dataset_probabilities)
            ):
                mass = np.asarray(proposal_mass, dtype=np.float64)
                if mass.shape != (len(episode.anchors),):
                    raise ValueError(
                        f"proposal anchor count differs for dataset {dataset_index}, "
                        f"episode {episode_index}"
                    )
                if not np.all(np.isfinite(mass)) or np.any(mass < 0) or mass.sum() <= 0:
                    raise ValueError("proposal masses must be finite, non-negative, and non-zero")
                residual_probability = mass / mass.sum(dtype=np.float64)
                baseline_probability = np.full(len(mass), 1.0 / len(mass), dtype=np.float64)
                safe_probability = (1.0 - alpha) * baseline_probability + alpha * residual_probability
                weights = baseline_probability / safe_probability
                if not np.isclose(safe_probability.sum(dtype=np.float64), 1.0, atol=1e-12):
                    raise ValueError("safe temporal proposal does not sum to one")
                if np.max(weights) > 1.0 / (1.0 - alpha) + 1e-12:
                    raise ValueError("importance weight exceeds the analytical safe-mixture bound")
                normalized_dataset.append(safe_probability)
                dataset_weights.append(weights)
                normalized_dataset_masses.append(mass.copy())
            normalized_probabilities.append(tuple(normalized_dataset))
            importance_weights.append(tuple(dataset_weights))
            normalized_masses.append(tuple(normalized_dataset_masses))

        self.manifest_sha256 = manifest.sha256
        self.manifest_support = manifest.to_dict(include_hash=False)["datasets"]
        self.alpha = float(alpha)
        self.score_kind = score_kind
        self.scorer_sha256 = scorer_sha256
        self.validation_certificate = validation_certificate
        self.proposal_masses = tuple(normalized_masses)
        self.anchor_probabilities = tuple(normalized_probabilities)
        self.importance_weights = tuple(importance_weights)
        canonical = json.dumps(self.to_dict(include_hash=False), sort_keys=True, separators=(",", ":"))
        self.sha256 = hashlib.sha256(canonical.encode("utf-8")).hexdigest()

    @classmethod
    def from_masses(
        cls,
        manifest: ValidAnchorManifest,
        alpha: float,
        masses: Sequence[Sequence[Sequence[float]]],
        score_kind: str = "synthetic",
        scorer_sha256: str = "synthetic-fixture",
        validation_certificate: dict[str, Any] | None = None,
    ) -> "TemporalProposal":
        return cls(
            manifest=manifest,
            alpha=alpha,
            anchor_probabilities=masses,
            score_kind=score_kind,
            scorer_sha256=scorer_sha256,
            validation_certificate=validation_certificate,
        )

    @classmethod
    def from_dict(
        cls, manifest: ValidAnchorManifest, payload: dict[str, Any]
    ) -> "TemporalProposal":
        if payload.get("schema_version") != cls.schema_version:
            raise ValueError(
                f"unsupported temporal proposal schema: {payload.get('schema_version')}"
            )
        expected_support = payload.get("manifest_support")
        current_support = json.loads(
            json.dumps(manifest.to_dict(include_hash=False)["datasets"])
        )
        if expected_support is not None and expected_support != current_support:
            raise ValueError("temporal proposal changes dataset/episode/anchor support")
        proposal = cls(
            manifest=manifest,
            alpha=float(payload["alpha"]),
            anchor_probabilities=payload["proposal_masses"],
            score_kind=str(payload["score_kind"]),
            scorer_sha256=str(payload["scorer_sha256"]),
            validation_certificate=payload.get("validation_certificate"),
        )
        return proposal

    @classmethod
    def load(
        cls, path: str | Path, manifest: ValidAnchorManifest
    ) -> "TemporalProposal":
        with open(path, "r", encoding="utf-8") as proposal_file:
            return cls.from_dict(manifest, json.load(proposal_file))

    def save(self, path: str | Path) -> None:
        destination = Path(path)
        destination.parent.mkdir(parents=True, exist_ok=True)
        temporary = destination.with_suffix(destination.suffix + ".tmp")
        with open(temporary, "w", encoding="utf-8") as proposal_file:
            json.dump(self.to_dict(), proposal_file, indent=2)
        os.replace(temporary, destination)

    def sample(
        self, manifest: ValidAnchorManifest, seed: int, draw: LogicalDraw
    ) -> SampleTrace:
        rng = np.random.default_rng(_stateless_seed(seed, draw.epoch, draw.global_draw_id))
        dataset_index = int(rng.choice(len(manifest.datasets), p=manifest._dataset_probabilities))
        dataset = manifest.datasets[dataset_index]
        trajectory_index = int(
            rng.choice(
                len(dataset.episodes), p=manifest._episode_probabilities[dataset_index]
            )
        )
        episode = dataset.episodes[trajectory_index]
        if self.alpha == 0.0:
            # Preserve the exact baseline RNG path, not merely its distribution.
            anchor_index = int(rng.integers(len(episode.anchors)))
        else:
            anchor_index = int(
                rng.choice(
                    len(episode.anchors),
                    p=self.anchor_probabilities[dataset_index][trajectory_index],
                )
            )
        baseline_probability = (
            dataset.probability
            * dataset.episode_probabilities[trajectory_index]
            / len(episode.anchors)
        )
        proposal_probability = (
            dataset.probability
            * dataset.episode_probabilities[trajectory_index]
            * self.anchor_probabilities[dataset_index][trajectory_index][anchor_index]
        )
        return SampleTrace(
            manifest_sha256=manifest.sha256,
            epoch=draw.epoch,
            global_draw_id=draw.global_draw_id,
            dataset_index=dataset_index,
            dataset_name=dataset.name,
            trajectory_index=episode.trajectory_index,
            trajectory_id=episode.trajectory_id,
            anchor_index=anchor_index,
            base_index=episode.anchors[anchor_index],
            baseline_probability=float(baseline_probability),
            proposal_sha256=self.sha256,
            proposal_probability=float(proposal_probability),
            importance_weight=float(
                self.importance_weights[dataset_index][trajectory_index][anchor_index]
            ),
        )

    def to_dict(self, include_hash: bool = True) -> dict[str, Any]:
        payload = {
            "schema_version": self.schema_version,
            "manifest_sha256": self.manifest_sha256,
            "manifest_support": self.manifest_support,
            "alpha": self.alpha,
            "score_kind": self.score_kind,
            "scorer_sha256": self.scorer_sha256,
            "validation_certificate": self.validation_certificate,
            "proposal_masses": [
                [probabilities.tolist() for probabilities in dataset]
                for dataset in self.proposal_masses
            ],
        }
        if include_hash:
            payload["sha256"] = self.sha256
        return payload