"""
NeonVLA — The core Vision-Language-Action model for G1.

Combines:
1. Video foundation model backbone (Qwen2.5-VL or Cosmos-Reason2)
2. Proprioception encoder (current joint states)
3. Multi-head action decoder (arms, locomotion, grippers)
4. Action chunking for temporal smoothness

The model takes (video_frames, language_instruction, proprioception)
and outputs a chunk of future actions for the G1 humanoid.
"""

from __future__ import annotations

import json
import logging
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn

from neon.data.action_space import ControlMode, G1ActionSpace
from neon.model.action_heads import ActionChunkingHead, ActionHeadConfig, G1ActionHead
from neon.model.video_backbone import BackboneConfig, VideoBackbone

logger = logging.getLogger(__name__)


@dataclass
class NeonConfig:
    """Full configuration for the Neon VLA model."""

    # Backbone
    backbone: BackboneConfig = field(default_factory=BackboneConfig)
    # Action space
    control_mode: str = "arms_only"  # "arms_only", "upper_body", "whole_body"
    include_locomotion: bool = True
    # Action head
    num_action_steps: int = 16  # Action chunk size
    mlp_hidden: int = 512
    action_head_layers: int = 3
    action_head_dropout: float = 0.1
    use_separate_heads: bool = True
    use_action_chunking: bool = True
    # Proprioception
    use_proprioception: bool = True
    proprioception_hidden: int = 128
    # Training
    learning_rate: float = 2e-4
    weight_decay: float = 0.01
    warmup_ratio: float = 0.1
    # Audio
    audio: Optional[Dict] = None  # AudioConfig as dict, or None to disable

    def to_dict(self) -> Dict:
        return asdict(self)

    def save(self, path: str):
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)

    @classmethod
    def load(cls, path: str) -> "NeonConfig":
        with open(path) as f:
            d = json.load(f)
        # Reconstruct nested configs
        if "backbone" in d and isinstance(d["backbone"], dict):
            d["backbone"] = BackboneConfig(**d["backbone"])
        return cls(**d)


@dataclass
class NeonOutput:
    """Output of a NeonVLA forward pass or prediction."""

    actions: np.ndarray  # (num_steps, action_dim) normalized to [-1, 1]
    raw_actions: Optional[np.ndarray] = None  # Denormalized actions
    upper_body: Optional[np.ndarray] = None  # (num_steps, 14) arm joint positions
    locomotion: Optional[np.ndarray] = None  # (num_steps, 3) vx, vy, omega
    gripper: Optional[np.ndarray] = None  # (num_steps, 2) gripper positions
    speech_path: Optional[str] = None  # Path to generated speech audio (from PersonaPlex)


class ProprioceptionEncoder(nn.Module):
    """Encodes current joint states into a feature vector.

    Proprioception is crucial for closed-loop control — the model needs
    to know where the robot currently is to predict good next actions.
    """

    def __init__(self, input_dim: int, hidden_dim: int = 128, output_dim: int = 256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, output_dim),
            nn.GELU(),
        )

    def forward(self, proprio: torch.Tensor) -> torch.Tensor:
        """Encode proprioception.

        Args:
            proprio: (batch, joint_dim) current joint positions

        Returns:
            (batch, output_dim) proprioception features
        """
        return self.net(proprio)


class NeonVLA(nn.Module):
    """Neon Vision-Language-Action model for the G1 humanoid.

    Pipeline:
        1. Video backbone encodes (frames, language) → visual-language features
        2. Proprioception encoder encodes current joint states → proprio features
        3. Feature fusion: concat + project
        4. Action heads decode → action chunks

    Usage:
        model = NeonVLA(NeonConfig())
        model.load_backbone()  # Load video foundation model

        # Training
        loss = model.compute_loss(images, text, proprio, target_actions)

        # Inference
        output = model.predict(image, instruction, joint_states)
    """

    def __init__(self, config: NeonConfig):
        super().__init__()
        self.config = config

        # Build action space
        self.action_space = G1ActionSpace(
            mode=ControlMode(config.control_mode),
            include_locomotion=config.include_locomotion,
        )

        # Build video backbone (lazy-loaded)
        self.backbone = VideoBackbone(config.backbone)

        # Build audio encoder (optional)
        self.audio_encoder = None
        self.speech_head = None
        self.speaker = None
        self._audio_config = None
        if config.audio is not None:
            from neon.model.audio import AudioConfig, AudioEncoder, SpeechResponseHead, PersonaPlexSpeaker

            self._audio_config = AudioConfig(**config.audio)
            if self._audio_config.enabled:
                self.audio_encoder = AudioEncoder(self._audio_config, self.backbone.hidden_size)
                if self._audio_config.enable_speech_output:
                    self.speech_head = SpeechResponseHead(self.backbone.hidden_size)
                    self.speaker = PersonaPlexSpeaker(self._audio_config)
                logger.info(f"Audio encoder enabled: {self._audio_config.encoder_type}")

        # Build proprioception encoder
        if config.use_proprioception:
            self.proprio_encoder = ProprioceptionEncoder(
                input_dim=self.action_space.action_dim,
                hidden_dim=config.proprioception_hidden,
                output_dim=config.proprioception_hidden * 2,
            )
            fusion_dim = self.backbone.hidden_size + config.proprioception_hidden * 2
        else:
            self.proprio_encoder = None
            fusion_dim = self.backbone.hidden_size

        # Add audio features to fusion if audio encoder is present
        if self.audio_encoder is not None:
            fusion_dim += self.audio_encoder.output_dim

        # Feature fusion projection
        self.fusion = nn.Sequential(
            nn.Linear(fusion_dim, self.backbone.hidden_size),
            nn.GELU(),
            nn.Dropout(config.action_head_dropout),
        )

        # Build action heads
        head_config = ActionHeadConfig(
            hidden_size=self.backbone.hidden_size,
            action_dim=self.action_space.action_dim,
            num_action_steps=config.num_action_steps,
            mlp_hidden=config.mlp_hidden,
            num_layers=config.action_head_layers,
            dropout=config.action_head_dropout,
            use_separate_heads=config.use_separate_heads,
            use_action_chunking=config.use_action_chunking,
        )
        self.action_head = G1ActionHead(head_config, self.action_space)

        logger.info(f"NeonVLA initialized: {self.action_space}")

    def load_backbone(self):
        """Load the video backbone model weights. Call before training/inference."""
        self.backbone.load()

    def forward(
        self,
        images: Optional[List] = None,
        video_frames: Optional[List] = None,
        text: Optional[str] = None,
        proprioception: Optional[torch.Tensor] = None,
        audio: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Forward pass — encode inputs and decode actions.

        Args:
            images: List of PIL images (single-frame input)
            video_frames: List of PIL images as video sequence
            text: Language instruction (text)
            proprioception: (batch, action_dim) current joint states
            audio: (batch, samples) raw audio waveform at 16kHz

        Returns:
            If action_chunking: (batch, num_steps, action_dim)
            Else: (batch, action_dim)
        """
        # Encode visual + language
        backbone_out = self.backbone.encode(images=images, video_frames=video_frames, text=text)
        features = backbone_out["pooled"]  # (batch, hidden_size)

        # Fuse with proprioception
        if self.proprio_encoder is not None and proprioception is not None:
            proprio_feat = self.proprio_encoder(proprioception)
            features = torch.cat([features, proprio_feat], dim=-1)

        # Fuse with audio features
        if self.audio_encoder is not None and audio is not None:
            audio_feat = self.audio_encoder(audio)  # (batch, audio_output_dim)
            features = torch.cat([features, audio_feat], dim=-1)

        features = self.fusion(features)

        # Decode actions
        actions = self.action_head(features)
        return actions

    def compute_loss(
        self,
        images: Optional[List] = None,
        video_frames: Optional[List] = None,
        text: Optional[str] = None,
        proprioception: Optional[torch.Tensor] = None,
        target_actions: Optional[torch.Tensor] = None,
        action_mask: Optional[torch.Tensor] = None,
        audio: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor]:
        """Compute training loss.

        Args:
            images, video_frames, text, proprioception, audio: Model inputs
            target_actions: Ground truth action chunks (batch, num_steps, action_dim)
            action_mask: Valid timestep mask (batch, num_steps)

        Returns:
            Dict with "loss" and optional per-group losses
        """
        predicted = self.forward(images, video_frames, text, proprioception, audio=audio)
        loss = self.action_head.compute_loss(predicted, target_actions, action_mask)
        return {"loss": loss}

    @torch.no_grad()
    def predict(
        self,
        image: Any = None,
        video_frames: Optional[List] = None,
        instruction: str = "",
        proprioception: Optional[np.ndarray] = None,
        audio: Optional[np.ndarray] = None,
        speak: bool = False,
    ) -> NeonOutput:
        """Run inference — predict actions from observations.

        Args:
            image: Single PIL image (current camera frame)
            video_frames: List of recent camera frames
            instruction: Natural language instruction (text)
            proprioception: Current joint positions (action_dim,)
            audio: Raw audio waveform at 16kHz (samples,) — spoken command
            speak: If True, robot narrates its action via PersonaPlex

        Returns:
            NeonOutput with predicted actions (and optional speech_path)
        """
        self.eval()

        # Prepare inputs
        images = [image] if image is not None else None

        # Proprioception
        proprio_tensor = None
        if proprioception is not None and self.proprio_encoder is not None:
            proprio_tensor = torch.from_numpy(proprioception).float().unsqueeze(0)
            device = next(self.action_head.parameters()).device
            proprio_tensor = proprio_tensor.to(device)

        # Audio
        audio_tensor = None
        if audio is not None and self.audio_encoder is not None:
            audio_tensor = torch.from_numpy(audio).float().unsqueeze(0)
            device = next(self.action_head.parameters()).device
            audio_tensor = audio_tensor.to(device)

        # Forward
        actions_tensor = self.forward(
            images=images,
            video_frames=video_frames,
            text=instruction,
            proprioception=proprio_tensor,
            audio=audio_tensor,
        )

        # To numpy
        actions_np = actions_tensor.cpu().numpy().squeeze(0)  # (num_steps, action_dim) or (action_dim,)
        if actions_np.ndim == 1:
            actions_np = actions_np[np.newaxis, :]  # (1, action_dim)

        # Denormalize
        raw_actions = np.stack([self.action_space.denormalize(a) for a in actions_np])

        # Split into groups
        split = self.action_space.split_action(raw_actions)

        output = NeonOutput(
            actions=actions_np,
            raw_actions=raw_actions,
            upper_body=np.concatenate(
                [split.get("left_arm", np.array([])), split.get("right_arm", np.array([]))], axis=-1
            ) if "left_arm" in split else None,
            locomotion=split.get("locomotion"),
            gripper=None,  # Extracted from arms if needed
        )

        # Speech output
        if speak and self.speaker is not None:
            speech_path = self.speaker.speak(f"Executing: {instruction}")
            output.speech_path = speech_path

        return output

    def save_pretrained(self, path: str):
        """Save model config + action head weights.

        Note: Backbone weights are loaded from HuggingFace separately.
        We only save the action heads + fusion layers + proprio encoder.
        """
        save_dir = Path(path)
        save_dir.mkdir(parents=True, exist_ok=True)

        # Save config
        self.config.save(str(save_dir / "config.json"))

        # Save action space
        with open(save_dir / "action_space.json", "w") as f:
            json.dump(self.action_space.to_dict(), f, indent=2)

        # Save trainable weights (action heads, fusion, proprio encoder)
        state = {}
        for name, param in self.named_parameters():
            if param.requires_grad or "action_head" in name or "fusion" in name or "proprio" in name:
                state[name] = param.cpu()

        torch.save(state, save_dir / "neon_weights.pt")
        logger.info(f"Saved Neon model to {save_dir} ({len(state)} tensors)")

    @classmethod
    def from_pretrained(cls, path: str, load_backbone: bool = False) -> "NeonVLA":
        """Load a saved Neon model.

        Args:
            path: Path to saved model directory or HuggingFace model ID
            load_backbone: If True, also load the video backbone

        Returns:
            NeonVLA instance with loaded weights
        """
        model_dir = Path(path)

        # Check if it's a HuggingFace model ID
        if not model_dir.exists():
            from huggingface_hub import snapshot_download

            model_dir = Path(snapshot_download(path))

        # Load config
        config = NeonConfig.load(str(model_dir / "config.json"))

        # Create model
        model = cls(config)

        # Load trainable weights
        weights_path = model_dir / "neon_weights.pt"
        if weights_path.exists():
            state = torch.load(weights_path, map_location="cpu", weights_only=True)
            missing, unexpected = model.load_state_dict(state, strict=False)
            logger.info(f"Loaded Neon weights: {len(state)} tensors, {len(missing)} missing, {len(unexpected)} unexpected")

        # Optionally load backbone
        if load_backbone:
            model.load_backbone()

        return model

    def __repr__(self) -> str:
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        total = sum(p.numel() for p in self.parameters())
        return (
            f"NeonVLA(\n"
            f"  backbone={self.config.backbone.model_id},\n"
            f"  action_space={self.action_space},\n"
            f"  action_chunking={self.config.use_action_chunking} (steps={self.config.num_action_steps}),\n"
            f"  trainable={trainable:,} / {total:,} params\n"
            f")"
        )
