"""
Diffusion Policy path planner.

Uses DDIM-style denoising to generate a planned trajectory conditioned on:
  - world_embedding : (latent_dim,)  — compressed video context
  - task_embedding  : (embed_dim,)   — language-grounded goal (from VLA)
  - proprio         : (proprio_dim,) — current robot state

Output: planned_trajectory (horizon, action_dim)
        A smooth sequence of waypoints the low-level VLA policy will track.

Architecture
------------
Stub  : pure-numpy DDIM with a linear score network; no GPU required.
Real  : loads a standard Diffusion Policy checkpoint (Chi et al. 2023) via
        PyTorch; drops in wherever a GPU + weights are available.

Data flow in Brain.think()
--------------------------
  world_embedding + task_embedding + proprio
    → DiffusionPathPlanner.plan()          ← this module
    → planned_trajectory (horizon, action_dim)
    → CosmosPolicy.act(..., planned_trajectory=trajectory)
    → action_chunk (chunk_len, action_dim)
    → HAL.push_actions()
"""
from __future__ import annotations

import logging
from typing import Optional

import numpy as np

from models.base import BaseDiffusionPlanner

log = logging.getLogger(__name__)


# ─────────────────────────────────────────── shared: denoised approach arc

def _min_jerk(t: np.ndarray) -> np.ndarray:
    """Min-jerk time-scaling 0→1 (zero vel/acc at endpoints) — smooth motion."""
    return 10 * t**3 - 15 * t**4 + 6 * t**5


def _denoise_arc(start: np.ndarray, goal: np.ndarray, n: int, arc: float,
                 alphas: np.ndarray, rng: np.random.Generator) -> np.ndarray:
    start = np.asarray(start, np.float32); goal = np.asarray(goal, np.float32)
    s = _min_jerk(np.linspace(0.0, 1.0, n))[:, None]            # (n,1)
    clean = start[None, :] * (1 - s) + goal[None, :] * s        # min-jerk lerp
    clean[:, 2] += arc * np.sin(np.pi * np.linspace(0, 1, n))   # raise midpoint
    clean[0], clean[-1] = start, goal                           # pin endpoints
    # DDIM denoise: noise the path, pull back toward the clean arc.
    x = clean + 0.05 * rng.standard_normal(clean.shape).astype(np.float32)
    for a in alphas[::-1]:
        x = a * clean + (1 - a) * x                             # contraction → clean
    x[0], x[-1] = start, goal
    return x.astype(np.float32)


# ─────────────────────────────────────────────────── stub (numpy-only DDIM)

class _StubDiffusionPlanner(BaseDiffusionPlanner):
    """
    Lightweight DDIM planner — no GPU, no weights.

    Score network: a seeded linear projection of [world_emb ‖ task_emb]
    that maps noise → goal direction. This gives a deterministic, smooth
    trajectory for a given (world, task) pair without any learned weights.

    Denoising schedule: cosine beta schedule, T steps → x_0 via DDIM.
    """

    def __init__(
        self,
        action_dim: int,
        horizon: int,
        ddim_steps: int,
        proprio_dim: int,
        rng: np.random.Generator,
    ):
        self._action_dim = action_dim
        self._horizon = horizon
        self._ddim_steps = ddim_steps
        self._proprio_dim = proprio_dim
        self._rng = rng

        # Fixed random projection matrices (stand-in for learned score net weights).
        # Shape: (action_dim, cond_dim) where cond_dim = latent + embed + proprio dims.
        # We lazily build these on the first call once we know cond_dim.
        self._W: Optional[np.ndarray] = None
        self._b: Optional[np.ndarray] = None

        # Pre-compute cosine DDIM schedule
        self._alphas = self._cosine_alphas(ddim_steps)

    # ── public API ───────────────────────────────────────────────────────────

    def plan(
        self,
        world_embedding: np.ndarray,
        task_embedding: np.ndarray,
        proprio: np.ndarray,
    ) -> np.ndarray:
        cond = self._build_cond(world_embedding, task_embedding, proprio)
        self._maybe_init_weights(cond.shape[0])

        # Start from pure noise: x_T ~ N(0, I)  shape (horizon, action_dim)
        x = self._rng.normal(0, 1, (self._horizon, self._action_dim)).astype(np.float32)

        # DDIM reverse process: T → 0
        for i in range(self._ddim_steps - 1, -1, -1):
            alpha_t = self._alphas[i]
            alpha_prev = self._alphas[i - 1] if i > 0 else 1.0

            # Score = linear projection of conditioning (noise-level modulated)
            noise_scale = (1.0 - alpha_t) ** 0.5
            score = self._score(x, cond, noise_scale)   # (horizon, action_dim)

            # DDIM update (deterministic, eta=0)
            pred_x0 = (x - noise_scale * score) / (alpha_t ** 0.5 + 1e-8)
            pred_x0 = np.clip(pred_x0, -1.0, 1.0)
            x = (alpha_prev ** 0.5) * pred_x0 + ((1.0 - alpha_prev) ** 0.5) * score

        log.debug("Diffusion plan: horizon=%d action_dim=%d", self._horizon, self._action_dim)
        return x.astype(np.float32)

    # ── internals ────────────────────────────────────────────────────────────

    def _build_cond(
        self,
        world_embedding: np.ndarray,
        task_embedding: np.ndarray,
        proprio: np.ndarray,
    ) -> np.ndarray:
        # Pad proprio to expected dim if needed
        p = proprio.flatten().astype(np.float32)
        if p.shape[0] < self._proprio_dim:
            p = np.pad(p, (0, self._proprio_dim - p.shape[0]))
        else:
            p = p[: self._proprio_dim]
        return np.concatenate([
            world_embedding.flatten().astype(np.float32),
            task_embedding.flatten().astype(np.float32),
            p,
        ])

    def _maybe_init_weights(self, cond_dim: int):
        if self._W is not None:
            return
        seed_rng = np.random.default_rng(42)
        scale = 1.0 / (cond_dim ** 0.5)
        self._W = seed_rng.normal(0, scale, (self._action_dim, cond_dim)).astype(np.float32)
        self._b = seed_rng.normal(0, 0.01, self._action_dim).astype(np.float32)

    def _score(
        self,
        x: np.ndarray,
        cond: np.ndarray,
        noise_scale: float,
    ) -> np.ndarray:
        # goal_dir: (action_dim,) — the direction the score pushes toward
        goal_dir = np.tanh(self._W @ cond + self._b)     # (action_dim,)
        # Broadcast goal direction across the horizon, scaled by noise level
        goal_dir_h = np.tile(goal_dir, (self._horizon, 1))  # (horizon, action_dim)
        # Return predicted noise: guide toward goal_dir, away from current x
        return noise_scale * (x - goal_dir_h)

    @staticmethod
    def _cosine_alphas(T: int) -> np.ndarray:
        t = np.linspace(0, 1, T + 1)
        f = np.cos((t + 0.008) / 1.008 * np.pi / 2) ** 2
        alphas_cumprod = f / f[0]
        return alphas_cumprod[1:].astype(np.float32)   # (T,) ascending 1→0

    # ── Cartesian approach planning (used by the pick executor) ────────────────

    def plan_approach(self, start: np.ndarray, goal: np.ndarray,
                      n: int = 12, arc: float = 0.10) -> np.ndarray:
        """
        Plan a smooth 3-D approach path from `start` to `goal`.

        Returns (n, 3) world-space waypoints. The hand lifts over a raised
        midpoint (min-jerk arc) so the approach clears the table edge instead
        of driving straight through it, then a few DDIM denoising steps clean
        any kinks — the path the IK reach actually tracks.
        """
        return _denoise_arc(start, goal, n, arc, self._alphas, self._rng)


# ─────────────────────────────────────────────────── real (PyTorch)

class _RealDiffusionPlanner(BaseDiffusionPlanner):
    """
    Loads a pre-trained Diffusion Policy checkpoint (Chi et al. 2023).

    Expected checkpoint format (standard dp repo):
      { "model_state_dict": ..., "config": { "action_dim": N, "horizon": H, ... } }

    Install:
      pip install diffusers torch
      # or clone https://github.com/real-stanford/diffusion_policy
    """

    def __init__(self, ckpt_path: str, device: str, action_dim: int, horizon: int):
        import torch

        self._device = device
        self._action_dim = action_dim
        self._horizon = horizon
        self._torch = torch

        log.info("Loading Diffusion Policy from %s ...", ckpt_path)
        ckpt = torch.load(ckpt_path, map_location=device)

        # Support both raw state dict and wrapped checkpoints
        state = ckpt.get("model_state_dict", ckpt)

        # Build a minimal UNet1d noise-prediction network
        self._net = self._build_net(state, action_dim).to(device).eval()

        # DDPM/DDIM schedule from diffusers
        from diffusers import DDIMScheduler
        self._scheduler = DDIMScheduler(
            num_train_timesteps=100,
            beta_schedule="squaredcos_cap_v2",
            clip_sample=True,
        )
        self._scheduler.set_timesteps(20)
        log.info("Diffusion Policy loaded.")

    def plan(
        self,
        world_embedding: np.ndarray,
        task_embedding: np.ndarray,
        proprio: np.ndarray,
    ) -> np.ndarray:
        import torch

        cond = np.concatenate([
            world_embedding.flatten(),
            task_embedding.flatten(),
            proprio.flatten(),
        ]).astype(np.float32)
        cond_t = torch.from_numpy(cond).unsqueeze(0).to(self._device)

        # Start from noise
        x = torch.randn(1, self._horizon, self._action_dim, device=self._device)

        with torch.no_grad():
            for t in self._scheduler.timesteps:
                noise_pred = self._net(x, t.unsqueeze(0), cond_t)
                x = self._scheduler.step(noise_pred, t, x).prev_sample

        return x.squeeze(0).cpu().numpy().astype(np.float32)

    def plan_approach(self, start: np.ndarray, goal: np.ndarray,
                      n: int = 12, arc: float = 0.10) -> np.ndarray:
        """Smooth 3-D approach arc (shared analytic+denoise path planner)."""
        import numpy as _np
        return _denoise_arc(start, goal, n, arc,
                            _np.linspace(0.2, 0.95, 8).astype(_np.float32),
                            _np.random.default_rng(0))

    def _build_net(self, state_dict, action_dim: int):
        """Minimal noise-prediction UNet1d. Replace with your checkpoint's arch."""
        import torch.nn as nn

        cond_dim = sum(v.shape[0] for k, v in state_dict.items() if "embed" in k and len(v.shape) == 1) or 128

        class _Net(nn.Module):
            def __init__(self):
                super().__init__()
                self.cond_proj = nn.Linear(cond_dim, 256)
                self.t_emb = nn.Embedding(200, 256)
                self.layers = nn.Sequential(
                    nn.Linear(action_dim + 256 + 256, 512), nn.SiLU(),
                    nn.Linear(512, 512), nn.SiLU(),
                    nn.Linear(512, action_dim),
                )

            def forward(self, x, t, cond):
                B, H, D = x.shape
                c = self.cond_proj(cond).unsqueeze(1).expand(B, H, -1)
                te = self.t_emb(t).unsqueeze(1).expand(B, H, -1)
                inp = torch.cat([x, c, te], dim=-1)
                return self.layers(inp)

        net = _Net()
        try:
            net.load_state_dict(state_dict, strict=False)
        except Exception as e:
            log.warning("Partial weight load for diffusion net: %s", e)
        return net


# ─────────────────────────────────────────────────── factory

class DiffusionPathPlanner:
    """
    Factory — returns a stub or real planner based on config.

    Config keys
    -----------
    stub       : bool   — use numpy stub (default True)
    ckpt_path  : str    — path to a .pt checkpoint (real mode)
    device     : str    — "cuda" | "cpu"
    action_dim : int    — must match robot DoF (default 21)
    horizon    : int    — planned steps per call (default 30)
    ddim_steps : int    — denoising steps for stub DDIM (default 20)
    proprio_dim: int    — proprio vector size (default 21)
    """

    def __new__(cls, config: dict) -> BaseDiffusionPlanner:  # type: ignore[misc]
        action_dim = config.get("action_dim", 21)
        horizon = config.get("horizon", 30)
        proprio_dim = config.get("proprio_dim", 21)

        if config.get("stub", True):
            log.info("DiffusionPathPlanner: STUB mode  (horizon=%d, ddim_steps=%d)",
                     horizon, config.get("ddim_steps", 20))
            return _StubDiffusionPlanner(
                action_dim=action_dim,
                horizon=horizon,
                ddim_steps=config.get("ddim_steps", 20),
                proprio_dim=proprio_dim,
                rng=np.random.default_rng(seed=99),
            )

        ckpt = config.get("ckpt_path", "")
        if not ckpt:
            raise ValueError("diffusion_planner.ckpt_path must be set when stub=false")

        log.info("DiffusionPathPlanner: real mode  ckpt=%s", ckpt)
        return _RealDiffusionPlanner(
            ckpt_path=ckpt,
            device=config.get("device", "cuda"),
            action_dim=action_dim,
            horizon=horizon,
        )