"""
2D to 3D pose uplift using inverse kinematics.

This module provides utilities for converting 2D pose sequences to 3D
using either:
1. Single-view lifting (MLP-based depth estimation)
2. Multi-view triangulation (when multiple camera views available)
3. IK optimization (with anatomical constraints)

Based on the ik_solver codebase.

Usage:
    from sltk.extraction.uplift import PoseUplifter, UpliftConfig

    uplifter = PoseUplifter()
    poses_3d = uplifter.uplift_sequence(poses_2d, method="mlp")
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import torch

import numpy as np

from sltk.data.core import PoseSequence
from sltk.exceptions import ConfigurationError, ValidationError


@dataclass
class UpliftConfig:
    """Configuration for 2D to 3D uplift."""

    # Method: "mlp", "triangulation", "ik"
    method: str = "mlp"

    # MLP model path (for method="mlp")
    mlp_model_path: str | None = None

    # Camera parameters (for triangulation)
    camera_matrices: np.ndarray | None = None  # (num_cams, 3, 4)

    # IK optimization parameters
    ik_iterations: int = 100
    ik_lr: float = 0.01
    bone_length_constraint: bool = True
    joint_angle_constraint: bool = True

    # Processing
    device: str = "cuda:0"
    verbose: bool = True


class PoseUplifter:
    """
    Convert 2D pose sequences to 3D.

    Supports multiple uplift methods:
    - MLP: Learning-based depth estimation
    - Triangulation: Multi-view geometry
    - IK: Inverse kinematics with anatomical constraints
    """

    def __init__(self, config: UpliftConfig | None = None):
        self.config = config or UpliftConfig()
        self._mlp_model: Any = None
        self._skeleton: Any = None

    def uplift_sequence(
        self,
        poses_2d: PoseSequence | np.ndarray,
        camera_matrices: np.ndarray | None = None,
        method: str | None = None,
    ) -> PoseSequence:
        """
        Convert 2D poses to 3D.

        Args:
            poses_2d: 2D pose sequence (T, N, 2) or PoseSequence
            camera_matrices: Optional camera projection matrices
            method: Uplift method override

        Returns:
            PoseSequence with 3D poses
        """
        method = method or self.config.method

        if isinstance(poses_2d, PoseSequence):
            data_2d = poses_2d.data[:, :, :2]
            fps = poses_2d.fps
            pose_format = poses_2d.format
            confidence = poses_2d.confidence
        else:
            data_2d = poses_2d
            fps = 25.0
            pose_format = "unknown"
            confidence = None

        if method == "mlp":
            data_3d = self._uplift_mlp(data_2d)
        elif method == "triangulation":
            if camera_matrices is None:
                camera_matrices = self.config.camera_matrices
            if camera_matrices is None:
                raise ConfigurationError("Camera matrices required for triangulation")
            data_3d = self._uplift_triangulation(data_2d, camera_matrices)
        elif method == "ik":
            data_3d = self._uplift_ik(data_2d)
        else:
            raise ConfigurationError(
                f"Unknown uplift method: {method}. Valid methods: mlp, triangulation, ik"
            )

        return PoseSequence(
            data=data_3d,
            fps=fps,
            format=f"{pose_format}_3d",
            confidence=confidence,
            metadata={
                "uplift_method": method,
                "source_format": pose_format,
            },
        )

    def _uplift_mlp(self, poses_2d: np.ndarray) -> np.ndarray:
        """Uplift using MLP depth estimator."""
        if self._mlp_model is None:
            self._load_mlp_model()

        T, N, _ = poses_2d.shape

        # Normalize 2D poses
        poses_normalized = self._normalize_poses_2d(poses_2d)

        # Predict depth
        import torch

        with torch.no_grad():
            poses_tensor = torch.from_numpy(poses_normalized).float()
            poses_tensor = poses_tensor.to(self.config.device)

            # Reshape for batch processing
            poses_flat = poses_tensor.reshape(T, -1)
            depth_pred = self._mlp_model(poses_flat)
            depth = depth_pred.reshape(T, N, 1).cpu().numpy()

        # Combine 2D + depth
        poses_3d = np.concatenate([poses_2d, depth], axis=-1)

        return poses_3d.astype(np.float32)

    def _load_mlp_model(self) -> None:
        """Load MLP depth estimation model."""
        if self.config.mlp_model_path is None:
            from sltk.config import get_ik_solver_root

            # Check environment variable first, then fallback to common locations
            ik_root = get_ik_solver_root()
            model_paths = []
            if ik_root:
                model_paths.append(ik_root / "models" / "depth_mlp.pt")
            # Fallback to common locations
            model_paths.extend(
                [
                    Path.home() / "Projects" / "ik_solver" / "models" / "depth_mlp.pt",
                ]
            )

            for default_path in model_paths:
                if default_path.exists():
                    self.config.mlp_model_path = str(default_path)
                    break

            if self.config.mlp_model_path is None:
                raise ConfigurationError(
                    "MLP model path required. Set via UpliftConfig.mlp_model_path "
                    "or SLTK_IK_SOLVER_ROOT environment variable"
                )

        import torch

        # Try loading from ik_solver
        try:
            import sys

            from sltk.config import get_ik_solver_root

            ik_root = get_ik_solver_root()
            ik_paths = []
            if ik_root:
                ik_paths.append(ik_root)
            # Fallback to common locations
            ik_paths.extend(
                [
                    Path.home() / "Projects" / "ik_solver",
                ]
            )

            for ik_path in ik_paths:
                if ik_path.exists():
                    sys.path.insert(0, str(ik_path))
                    break

            from mlp_model import DepthMLP

            self._mlp_model = DepthMLP.load(self.config.mlp_model_path)
        except ImportError:
            # Fallback: load as generic TorchScript
            self._mlp_model = torch.jit.load(self.config.mlp_model_path)

        self._mlp_model = self._mlp_model.to(self.config.device)
        self._mlp_model.eval()

    def _uplift_triangulation(
        self,
        poses_2d: np.ndarray,
        camera_matrices: np.ndarray,
    ) -> np.ndarray:
        """Uplift using multi-view triangulation."""
        T, N, _ = poses_2d.shape
        camera_matrices.shape[0]

        if poses_2d.ndim == 4:
            # Multiple camera views: (T, num_cams, N, 2)
            multi_view_poses = poses_2d
        else:
            raise ValidationError(
                "Triangulation requires multi-view 2D poses with shape (T, num_cams, N, 2)"
            )

        poses_3d = np.zeros((T, N, 3), dtype=np.float32)

        for t in range(T):
            for n in range(N):
                # Get 2D points from all cameras
                points_2d = multi_view_poses[t, :, n, :]  # (num_cams, 2)

                # Triangulate
                point_3d = self._triangulate_point(points_2d, camera_matrices)
                poses_3d[t, n] = point_3d

        return poses_3d

    def _triangulate_point(
        self,
        points_2d: np.ndarray,
        P: np.ndarray,
    ) -> np.ndarray:
        """
        Triangulate a single 3D point from multiple 2D observations.

        Uses Direct Linear Transform (DLT) with SVD.

        Args:
            points_2d: (num_cams, 2) 2D observations
            P: (num_cams, 3, 4) projection matrices

        Returns:
            (3,) 3D point coordinates
        """
        num_cams = P.shape[0]

        # Build A matrix for DLT
        A = np.zeros((2 * num_cams, 4))

        for i in range(num_cams):
            x, y = points_2d[i]
            A[2 * i] = x * P[i, 2] - P[i, 0]
            A[2 * i + 1] = y * P[i, 2] - P[i, 1]

        # SVD solution: X is last column of V
        _, _, Vt = np.linalg.svd(A)
        X = Vt[-1]

        # Convert from homogeneous
        return X[:3] / X[3]

    def _uplift_ik(self, poses_2d: np.ndarray) -> np.ndarray:
        """Uplift using inverse kinematics optimization."""
        T, N, _ = poses_2d.shape

        if self._skeleton is None:
            self._load_skeleton(N)

        import torch

        poses_2d_tensor = torch.from_numpy(poses_2d).float().to(self.config.device)

        # Initialize angles from heuristics
        angles = self._initialize_angles(poses_2d_tensor)

        # Optimization loop
        optimizer = torch.optim.Adam([angles], lr=self.config.ik_lr)

        for i in range(self.config.ik_iterations):
            optimizer.zero_grad()

            # Forward kinematics
            poses_3d = self._forward_kinematics(angles)

            # Project to 2D
            poses_2d_proj = self._project_to_2d(poses_3d)

            # Reprojection loss
            loss = torch.mean((poses_2d_proj - poses_2d_tensor) ** 2)

            # Constraints
            if self.config.bone_length_constraint:
                loss = loss + 0.1 * self._bone_length_loss(poses_3d)

            if self.config.joint_angle_constraint:
                loss = loss + 0.1 * self._angle_constraint_loss(angles)

            loss.backward()
            optimizer.step()

            if self.config.verbose and i % 20 == 0:
                print(f"  IK iteration {i}: loss = {loss.item():.6f}")

        # Get final 3D poses
        with torch.no_grad():
            poses_3d = self._forward_kinematics(angles)

        return poses_3d.cpu().numpy()

    def _load_skeleton(self, num_keypoints: int) -> None:
        """Load skeleton definition."""
        # Try to load from ik_solver
        try:
            import sys

            from sltk.config import get_ik_solver_root

            ik_root = get_ik_solver_root()
            ik_paths = []
            if ik_root:
                ik_paths.append(ik_root)
            # Fallback to common locations
            ik_paths.extend(
                [
                    Path.home() / "Projects" / "ik_solver",
                ]
            )

            for ik_path in ik_paths:
                if ik_path.exists():
                    sys.path.insert(0, str(ik_path))
                    break

            from skeleton import Skeleton

            self._skeleton = Skeleton.from_num_keypoints(num_keypoints)
            return
        except ImportError:
            pass

        # Fallback: create simple kinematic chain
        self._skeleton = self._create_simple_skeleton(num_keypoints)

    def _create_simple_skeleton(self, num_keypoints: int) -> dict[str, Any]:
        """Create a simple skeleton definition."""
        # Basic parent-child relationships for common formats
        if num_keypoints == 33:
            # MediaPipe body
            parents = [
                -1,
                0,
                0,
                1,
                2,
                3,
                4,
                5,
                6,
                7,
                8,
                5,
                6,
                11,
                12,
                13,
                14,
                15,
                16,
                15,
                16,
                15,
                16,
                11,
                12,
                23,
                24,
                25,
                26,
                27,
                28,
                27,
                28,
            ]
        elif num_keypoints == 21:
            # Hand keypoints
            parents = [-1, 0, 1, 2, 3, 0, 5, 6, 7, 0, 9, 10, 11, 0, 13, 14, 15, 0, 17, 18, 19]
        else:
            # Generic chain
            parents = [-1] + list(range(num_keypoints - 1))

        return {
            "num_keypoints": num_keypoints,
            "parents": parents,
        }

    def _normalize_poses_2d(self, poses_2d: np.ndarray) -> np.ndarray:
        """Normalize 2D poses to [-1, 1] range."""
        poses = poses_2d.copy()

        # Center on hip/root
        root_idx = 0
        root = poses[:, root_idx : root_idx + 1, :]
        poses = poses - root

        # Scale by torso height
        scale = np.std(poses, axis=(1, 2), keepdims=True) + 1e-6
        poses = poses / scale

        return poses

    def _initialize_angles(self, poses_2d: torch.Tensor) -> torch.Tensor:
        """Initialize joint angles from 2D poses."""
        import torch

        T, N, _ = poses_2d.shape
        num_angles = (N - 1) * 3  # 3 angles per joint

        # Initialize with small random values
        angles = torch.randn(T, num_angles, device=self.config.device) * 0.1
        angles.requires_grad_(True)

        return angles

    def _forward_kinematics(self, angles: torch.Tensor) -> torch.Tensor:
        """Compute 3D poses from joint angles."""
        import torch

        T = angles.shape[0]
        N = self._skeleton["num_keypoints"]
        parents = self._skeleton["parents"]

        # Initialize positions
        positions = torch.zeros(T, N, 3, device=angles.device)

        # Root position
        positions[:, 0, :] = 0

        # Forward pass through kinematic chain
        angle_idx = 0
        for i in range(1, N):
            parent = parents[i]
            if parent >= 0:
                # Get rotation angles
                rot = angles[:, angle_idx * 3 : (angle_idx + 1) * 3]
                angle_idx += 1

                # Simple rotation (could use proper rotation matrices)
                bone_length = 0.1  # Default bone length
                offset = torch.stack(
                    [
                        bone_length * torch.cos(rot[:, 0]) * torch.cos(rot[:, 1]),
                        bone_length * torch.sin(rot[:, 0]) * torch.cos(rot[:, 1]),
                        bone_length * torch.sin(rot[:, 1]),
                    ],
                    dim=-1,
                )

                positions[:, i] = positions[:, parent] + offset

        return positions

    def _project_to_2d(self, poses_3d: torch.Tensor) -> torch.Tensor:
        """Project 3D poses to 2D (orthographic)."""
        return poses_3d[:, :, :2]

    def _bone_length_loss(self, poses_3d: torch.Tensor) -> torch.Tensor:
        """Compute bone length consistency loss."""
        import torch

        parents = self._skeleton["parents"]
        loss = torch.tensor(0.0, device=poses_3d.device)

        for i, parent in enumerate(parents):
            if parent >= 0:
                bone = poses_3d[:, i] - poses_3d[:, parent]
                bone_length = torch.norm(bone, dim=-1)

                # Encourage consistent bone length across time
                target_length = bone_length.mean()
                loss = loss + torch.mean((bone_length - target_length) ** 2)

        return loss

    def _angle_constraint_loss(self, angles: torch.Tensor) -> torch.Tensor:
        """Compute joint angle constraint loss."""
        import torch

        # Simple constraint: keep angles within reasonable range
        max_angle = np.pi / 2
        violation = torch.relu(torch.abs(angles) - max_angle)
        return torch.mean(violation**2)


def uplift_poses(
    poses_2d: PoseSequence | np.ndarray,
    method: str = "mlp",
    **kwargs,
) -> PoseSequence:
    """
    Quick 2D to 3D pose conversion.

    Args:
        poses_2d: 2D pose sequence
        method: Uplift method ("mlp", "triangulation", "ik")
        **kwargs: Additional config parameters

    Returns:
        PoseSequence with 3D poses
    """
    config = UpliftConfig(method=method, **kwargs)
    uplifter = PoseUplifter(config)
    return uplifter.uplift_sequence(poses_2d)
