"""PEPG-style optimizer for policy search over control parameters."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from pathlib import Path
import json

import numpy as np


@dataclass
class PEPGState:
    """Serializable optimizer state."""

    mean: list[float]
    sigma: list[float]
    learning_rate: float
    sigma_learning_rate: float
    iteration: int
    seed: int
    rng_state: dict


class PEPGOptimizer:
    """Parameter-Exploring Policy Gradients (PEPG) with diagonal Gaussian."""

    def __init__(
        self,
        parameter_dim: int,
        *,
        seed: int = 0,
        init_sigma: float = 0.1,
        learning_rate: float = 0.05,
        sigma_learning_rate: float = 0.02,
        sigma_min: float = 1e-4,
    ):
        if parameter_dim <= 0:
            raise ValueError("parameter_dim must be positive")
        self.parameter_dim = parameter_dim
        self.mean = np.zeros(parameter_dim, dtype=np.float64)
        self.sigma = np.full(parameter_dim, init_sigma, dtype=np.float64)
        self.learning_rate = float(learning_rate)
        self.sigma_learning_rate = float(sigma_learning_rate)
        self.sigma_min = float(sigma_min)
        self.iteration = 0
        self.seed = int(seed)
        self._rng = np.random.default_rng(self.seed)

    def ask(self, population_size: int) -> tuple[np.ndarray, np.ndarray]:
        """Sample antithetic population and return (candidates, perturbations)."""
        if population_size <= 0 or population_size % 2 != 0:
            raise ValueError("population_size must be a positive even integer")
        half = population_size // 2
        eps = self._rng.standard_normal((half, self.parameter_dim))
        perturb = np.vstack([eps, -eps]) * self.sigma
        candidates = self.mean + perturb
        return candidates, perturb

    def tell(self, perturbations: np.ndarray, rewards: np.ndarray) -> None:
        """Update mean and sigma from population rewards."""
        perturbations = np.asarray(perturbations, dtype=np.float64)
        rewards = np.asarray(rewards, dtype=np.float64)
        if perturbations.shape[0] != rewards.shape[0]:
            raise ValueError("perturbations and rewards length mismatch")

        centered = rewards - np.mean(rewards)
        scale = np.std(rewards)
        if scale > 0:
            centered = centered / (scale + 1e-8)

        grad_mean = perturbations.T @ centered / len(rewards)
        self.mean += self.learning_rate * grad_mean

        sigma_grad = (
            ((perturbations**2 - self.sigma**2) / np.maximum(self.sigma, 1e-8)).T
            @ centered
            / len(rewards)
        )
        self.sigma = np.maximum(self.sigma + self.sigma_learning_rate * sigma_grad, self.sigma_min)
        self.iteration += 1

    def state_dict(self) -> PEPGState:
        """Return serializable state."""
        return PEPGState(
            mean=self.mean.tolist(),
            sigma=self.sigma.tolist(),
            learning_rate=self.learning_rate,
            sigma_learning_rate=self.sigma_learning_rate,
            iteration=self.iteration,
            seed=self.seed,
            rng_state=self._rng.bit_generator.state,
        )

    def load_state_dict(self, state: PEPGState) -> None:
        """Load optimizer from state."""
        self.mean = np.asarray(state.mean, dtype=np.float64)
        self.sigma = np.asarray(state.sigma, dtype=np.float64)
        self.learning_rate = state.learning_rate
        self.sigma_learning_rate = state.sigma_learning_rate
        self.iteration = state.iteration
        self.seed = state.seed
        self._rng = np.random.default_rng()
        self._rng.bit_generator.state = state.rng_state

    def save_checkpoint(self, path: str | Path) -> None:
        """Persist optimizer checkpoint as JSON."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as f:
            json.dump(asdict(self.state_dict()), f, indent=2)

    @classmethod
    def load_checkpoint(cls, path: str | Path) -> "PEPGOptimizer":
        """Restore optimizer from checkpoint."""
        path = Path(path)
        with path.open("r", encoding="utf-8") as f:
            payload = json.load(f)
        state = PEPGState(**payload)
        opt = cls(parameter_dim=len(state.mean), seed=state.seed)
        opt.load_state_dict(state)
        return opt