"""SMPLFX fitting wrapper for SLTK.

Runs SMPLFX body-model optimisation using outputs from NLF, WiLoR, TEASER,
RTMPose, and (optionally) Video Depth Anything.  Produces optimised SMPL-X
parameters and 3D keypoints that can be consumed by downstream classifiers.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import NamedTuple

import numpy as np
import torch

logger = logging.getLogger(__name__)


def _align_smpl_translation(
    transl: torch.Tensor,
    K_old: torch.Tensor,
    K_new: torch.Tensor,
) -> torch.Tensor:
    """Analytically re-align SMPL-X translation from K_old to K_new camera space.

    Ensures the projected body position remains visually consistent when the
    intrinsic matrix is replaced (e.g. by MoGe-estimated intrinsics).

    Args:
        transl: ``(N, 3)`` tensor of original translations.
        K_old: ``(3, 3)`` original intrinsic matrix.
        K_new: ``(3, 3)`` new intrinsic matrix.

    Returns:
        ``(N, 3)`` updated translation tensor.
    """
    fx1, fy1, cx1, cy1 = K_old[0, 0], K_old[1, 1], K_old[0, 2], K_old[1, 2]
    fx2, _fy2, cx2, cy2 = K_new[0, 0], K_new[1, 1], K_new[0, 2], K_new[1, 2]
    t1x, t1y, t1z = transl[:, 0], transl[:, 1], transl[:, 2]
    t2z = t1z * (fx2 / fx1)
    t2x = t1x + (t1z / fx1) * (cx1 - cx2)
    t2y = t1y + (t1z / fy1) * (cy1 - cy2)
    return torch.stack((t2x, t2y, t2z), dim=1)


@dataclass
class SMPLFXFittingConfig:
    """Configuration for SMPLFX optimisation."""

    device: str = "cuda:0"
    max_iters: int = 350
    depth_weight_loss: float = 20.0
    mesh_intersection_weight: float = 200.0
    model_path: str | None = None
    segmentation_path: str | None = None
    flame_masks_path: str | None = None
    flame_corr_path: str | None = None
    optimize_hands: bool = False
    optimize_shape: bool = False
    optimize_faces: bool = True
    optimize_body_pose: bool = True
    chunk_size: int = 700
    lr: float = 5e-3
    min_lr: float = 1e-5
    scheduler_patience: int = 10
    print_every: int = 15
    verbose: bool = True


class FittingResult(NamedTuple):
    """Result from SMPLFX fitting."""

    optimised_pt: Path
    smplfx_h5: Path
    num_frames: int
    location_features_pt: Path | None = None


class SMPLFXFitter:
    """Post-processing fitter that optimises SMPLFX from multi-source H5 inputs."""

    name = "smplfx_fitting"

    def __init__(self, config: SMPLFXFittingConfig | None = None) -> None:
        self.config = config or SMPLFXFittingConfig()
        self._resolve_weight_paths()

    def _build_smplx_kwargs_from_params(self, smplx_params: dict) -> dict:
        """Construct SMPLFX kwargs from stored SMPLX params."""
        from .smplfx_deps.utils_smplx import convertSMPLX2SMPLFXKwargs

        cfg = self.config
        betas = smplx_params.get("betas")
        shape = smplx_params.get("shape")
        expression = smplx_params.get("expression")

        num_betas = int(betas.shape[-1]) if isinstance(betas, torch.Tensor) else 10
        num_betas_flame = int(shape.shape[-1]) if isinstance(shape, torch.Tensor) else num_betas
        num_expression = int(expression.shape[-1]) if isinstance(expression, torch.Tensor) else 50

        return convertSMPLX2SMPLFXKwargs(
            {
                "model_path": cfg.model_path,
                "model_type": "smplx",
                "gender": "female",
                "use_face_contour": False,
                "num_betas": num_betas,
                "num_expression_coeffs": num_expression,
                "use_pca": False,
                "ext": "npz",
                "batch_size": 1,
                "flat_hand_mean": True,
                "use_euler_angles": True,
                "gender_flame": "generic",
                "ext_flame": "pkl",
            },
            num_betas_flame=num_betas_flame,
            v_mask_path_flame=cfg.flame_masks_path,
            corr_fname=cfg.flame_corr_path,
        )

    @staticmethod
    def _infer_num_frames(smplx_params: dict) -> int:
        """Infer temporal length from SMPLX parameter tensors."""
        for key in (
            "body_pose",
            "transl",
            "global_orient",
            "expression",
            "jaw_pose",
            "left_hand_pose",
            "right_hand_pose",
        ):
            value = smplx_params.get(key)
            if isinstance(value, torch.Tensor) and value.ndim >= 1:
                return int(value.shape[0])
        return 0

    # ------------------------------------------------------------------
    # Weight resolution
    # ------------------------------------------------------------------

    def _resolve_weight_paths(self) -> None:
        """Auto-resolve model weight paths from the SLTK weights dir."""
        from sltk.extraction.weights import get_weight_path

        cfg = self.config

        if cfg.model_path is None:
            resolved = get_weight_path("smplfx_model")
            if resolved is not None:
                # The optimizer expects the *directory* containing the model
                # SMPLFX body_models.py uses string concat (model_path+'smplx/')
                # so we need a trailing /
                cfg.model_path = str(resolved.parent.parent) + "/"
            else:
                raise FileNotFoundError(
                    "SMPLX model not found. Set config.model_path or place "
                    "SMPLX_NEUTRAL.npz at sltk/weights/smplfx/smplx/"
                )

        if cfg.segmentation_path is None:
            resolved = get_weight_path("smplfx_segmentation")
            if resolved is not None:
                cfg.segmentation_path = str(resolved)
            else:
                logger.warning(
                    "smplx_vert_segmentation.json not found – mesh intersection "
                    "loss will be disabled."
                )

        if cfg.flame_masks_path is None:
            resolved = get_weight_path("smplfx_flame_masks")
            if resolved is not None:
                cfg.flame_masks_path = str(resolved)

        if cfg.flame_corr_path is None:
            resolved = get_weight_path("smplfx_flame_corr")
            if resolved is not None:
                cfg.flame_corr_path = str(resolved)

    # ------------------------------------------------------------------
    # Core fitting
    # ------------------------------------------------------------------

    def fit(
        self,
        teaser_h5: str | Path,
        rtm_h5: str | Path,
        nlf_h5: str | Path,
        wilor_h5: str | Path,
        depth_h5: str | Path | None = None,
        output_dir: str | Path = ".",
    ) -> FittingResult:
        """Run SMPLFX optimisation.

        Args:
            teaser_h5: Path to TEASER face H5.
            rtm_h5: Path to RTMPose wholebody H5.
            nlf_h5: Path to NLF body H5.
            wilor_h5: Path to WiLoR hands H5.
            depth_h5: Path to per-keypoint depth H5 (optional).
            output_dir: Directory for output files.

        Returns:
            ``FittingResult`` with paths and frame count.
        """
        cfg = self.config
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        pt_path = output_dir / "smplfx_optimised.pt"
        h5_path = output_dir / "body_smplfx.h5"

        # -- 1. Load data through adapters --
        from .smplfx_deps.extraction_utils import (
            replaceKeypointsWeightsWithHandModelSuccess,
        )
        from .smplfx_deps.h5_adapters import (
            load_depth_for_fitting,
            load_nlf_for_fitting,
            load_rtm_for_fitting,
            load_teaser_for_fitting,
            load_wilor_for_fitting,
        )
        from .smplfx_deps.smplx_smoother_optimizer import SMPLXOptimizer
        from .smplfx_deps.utils import (
            convertMMPoseScores2WeightsOpenPose133,
            matrix_to_euler_angles,
        )
        from .smplfx_deps.utils_smplx import (
            SMPLX2DJointMapper,
            convertSMPLX2SMPLFXKwargs,
            remapIndependentHandToBody,
        )
        from .smplfx_deps.vision3d import intrinsic_matrix_from_field_of_view

        logger.info("Loading TEASER data from %s", teaser_h5)
        data_teaser = load_teaser_for_fitting(
            teaser_h5, median_shape=True, jaw_pose="euler", conv="ZYX"
        )

        logger.info("Loading RTMPose data from %s", rtm_h5)
        data_kpts = load_rtm_for_fitting(rtm_h5)

        logger.info("Loading NLF data from %s", nlf_h5)
        data_nlf = load_nlf_for_fitting(nlf_h5, angle_format="euler", conv="ZYX")

        logger.info("Loading WiLoR data from %s", wilor_h5)
        data_wilor = load_wilor_for_fitting(wilor_h5)

        num_frames = data_wilor.num_frames

        # -- 2. Prepare 2D keypoints --
        if len(data_kpts["people"]) > 1:
            logger.warning("Multiple people in RTM detections – using person_0 only.")
        kpts_2d = torch.from_numpy(data_kpts["people"]["person_0"]["keypoints"]).float()
        assert (
            kpts_2d.shape[0] == num_frames
        ), f"RTM frames ({kpts_2d.shape[0]}) != video frames ({num_frames})"

        mapper = SMPLX2DJointMapper(
            with_hips=False,
            with_legs=False,
            with_top_feet=False,
            with_mp_main_face=False,
            with_hand_palm=False,
            with_shoulders=True,
            detector_name="rtmp",
            use_face_contour=False,
            is_full_face=False,
        )
        kpts_2d = replaceKeypointsWeightsWithHandModelSuccess(
            kpts_2d,
            mapper.rhand_idxs,
            mapper.lhand_idxs,
            data_wilor.data["mano"]["success_r"],
            data_wilor.data["mano"]["success_l"],
        )

        # -- 3. Build SMPLX init --
        smplx_init = data_nlf.data["smplx"]
        smplx_init["expression"] = torch.zeros(
            num_frames, data_teaser["flame"]["expression"].shape[-1]
        )
        smplx_init["eyelid_params"] = torch.zeros(
            num_frames, data_teaser["flame"]["eyelid_params"].shape[-1]
        )
        for key in data_teaser["flame"]:
            if key != "shape":
                smplx_init[key][data_teaser["successful_frames"]] = data_teaser["flame"][key]

        smplx_init["body_pose"] = remapIndependentHandToBody(
            data_wilor.data["mano"]["global_orient_r"],
            data_wilor.data["mano"]["global_orient_l"],
            smplx_init["body_pose"],
            smplx_init["global_orient"],
            "rot",
            "euler",
            conv="ZYX",
        )
        smplx_init["right_hand_pose"] = matrix_to_euler_angles(
            data_wilor.data["mano"]["right_hand_pose"], "ZYX"
        ).flatten(1, -1)
        smplx_init["left_hand_pose"] = matrix_to_euler_angles(
            data_wilor.data["mano"]["left_hand_pose"], "ZYX"
        ).flatten(1, -1)

        smplx_init["shape"] = data_teaser["flame"]["shape"]
        for key in ("betas", "shape", "global_orient", "transl"):
            smplx_init[key] = smplx_init[key].median(0).values[None]

        # -- 4. Keypoint weights --
        poses_2d_weights_bin = convertMMPoseScores2WeightsOpenPose133(
            kpts_2d[:, None][..., 2].clone(),
            mapper.face_idxs,
            mapper.hand_idxs,
            mapper.body_idxs,
            low_q_face=0.05,
            high_q_face=0.95,
            low_q_body=None,
            high_q_body=0.95,
            low_q_hand=None,
            high_q_hand=0.95,
            normalize_min_body=4.0,
            normalize_min_hand=4.5,
        )[..., None]

        # -- 5. SMPLX kwargs --
        smplx_kwargs = convertSMPLX2SMPLFXKwargs(
            {
                "model_path": cfg.model_path,
                "model_type": "smplx",
                "gender": "female",
                "use_face_contour": False,
                "num_betas": smplx_init["betas"].shape[-1],
                "num_expression_coeffs": smplx_init["expression"].shape[-1],
                "use_pca": False,
                "ext": "npz",
                "batch_size": 1,
                "flat_hand_mean": True,
                "use_euler_angles": True,
                "gender_flame": "generic",
                "ext_flame": "pkl",
            },
            num_betas_flame=smplx_init["shape"].shape[-1],
            v_mask_path_flame=cfg.flame_masks_path,
            corr_fname=cfg.flame_corr_path,
        )

        # -- 6. Camera intrinsics --
        nlf_attrs = data_nlf.data.get("_attrs", data_nlf.meta)
        fov = nlf_attrs.get("fov", 55.0)
        # Use RTMPose/native frame resolution in (H, W) order when available.
        # This matches intrinsic_matrix_from_field_of_view() expectation.
        resolution_hw = data_kpts.get(
            "resolution",
            nlf_attrs.get("resolution", np.array([1080, 1920])),
        )
        if isinstance(resolution_hw, np.ndarray):
            resolution_hw = resolution_hw.tolist()
        else:
            resolution_hw = list(resolution_hw)

        ks = intrinsic_matrix_from_field_of_view(fov, resolution_hw, to_torch=True)[None]
        rts = torch.eye(3, 4)[None]

        # -- 7. Fitting kwargs --
        mesh_weight = cfg.mesh_intersection_weight
        if cfg.segmentation_path is None:
            logger.warning("No segmentation path – disabling mesh intersection loss.")
            mesh_weight = 0.0

        fitting_kwargs: dict = {
            "smplx_kwargs": smplx_kwargs,
            "smplx_params_init": smplx_init,
            "right_hand_constraints": None,
            "left_hand_constraints": None,
            "K": ks,
            "RT": rts,
            "mapper": mapper,
            "poses_2d": kpts_2d[:, None][:, :, mapper.idxs, :2],
            "poses_2d_weights": poses_2d_weights_bin[:, :, mapper.idxs],
            "device": cfg.device,
            "MAX_ITERS": cfg.max_iters,
            "min_lr": cfg.min_lr,
            "chunk_size": cfg.chunk_size,
            "loss_fn": "huber",
            "optimize_shape": cfg.optimize_shape,
            "optimize_hands": cfg.optimize_hands,
            "optimize_smplx_rt": True,
            "OPTIMIZE_RT": True,
            "OPTIMIZE_RT_FIX_CAM_0": False,
            "MESH_INTERSECTION_LOSS_WEIGHT": mesh_weight,
            "other_params_init": {},
            "lr": cfg.lr,
            "SMOOTHNESS_TERMS": [("vertices", "acceleration", 1e2)],
            "is_smplfx": True,
            "keys": np.arange(num_frames),
            "PRINT_EVERY": cfg.print_every,
            "SCHEDULER_PATIENCE": cfg.scheduler_patience,
            "GLOBAL_ORIENT_ZERO_X": True,
            "GLOBAL_ORIENT_FIX_Z": True,
            "body_zero_angles_idxs": [0, 1, 3, 4, 6, 7, 9, 10],
            "LRs": {
                "J_regressor": 0.0,
                "expression": 1e-6,
                "jaw_pose": 1e-5,
                "eyelid_params": 1e-6,
                "body_pose": 5e-3,
            },
            "are_control_points": False,
            "REG_NORM_EXPRESSION_W": 0,
            "BETA_REG_W": 1e-4,
            "USE_SIGMOID_ACTIVATION": False,
            "spacing": 3,
            "APPLY_SMOOTHNESS_ON_CP": True,
            "apply_constraints_before_decode": False,
            "ANGLES_REG_W": 0.0,
            "ANGLES_REG_IDXS": [2, 5, 8, 11, 14, 12, 13],
            "OPTIMIZE_FLAME_SHAPE": False,
            "body_constraints": None,
            "is_single_smplx_rt": True,
            "poses_3d": None,
            "OPTIMIZE_FACES": cfg.optimize_faces,
            "IS_BODY_POSE_6D": False,
            "OPTIMIZE_BODY_POSE": cfg.optimize_body_pose,
            "depths": None,
            "depths_mask": None,
            "DEPTH_LOSS_WEIGHT": 0.0,
            "z_3d": None,
            "Z_3D_LOSS_WEIGHT": 0.0,
            "segmentation_path": cfg.segmentation_path,
            "smplx_model_path": cfg.model_path,
        }

        # -- 8. Depth supervision --
        if depth_h5 is not None and cfg.depth_weight_loss > 0:
            logger.info("Loading depth data from %s", depth_h5)
            depths_data = load_depth_for_fitting(depth_h5)
            depths = depths_data["person_0"]["depths"]
            depths_mask = depths_data["person_0"]["depths_mask"]
            fitting_kwargs["depths"] = depths[:, mapper.idxs][:, None]
            fitting_kwargs["depths_mask"] = depths_mask[:, mapper.idxs][:, None]
            fitting_kwargs["DEPTH_LOSS_WEIGHT"] = cfg.depth_weight_loss

            # Override K with MoGe-estimated intrinsics when available
            attrs = depths_data.get("attrs", {})
            if "focal_x" in attrs:
                logger.info("Using MoGe-estimated camera intrinsics from depth H5.")
                ks_nlf = ks.clone()  # save original FOV-based K
                ks[0, 0, 0] = float(attrs["focal_x"])
                ks[0, 1, 1] = float(attrs["focal_y"])
                ks[0, 0, 2] = float(attrs["pp_x"])
                ks[0, 1, 2] = float(attrs["pp_y"])
                fitting_kwargs["K"] = ks
                smplx_init["transl"] = _align_smpl_translation(
                    smplx_init["transl"], ks_nlf[0], ks[0]
                )

        elif cfg.depth_weight_loss > 0:
            logger.warning("Depth weight > 0 but no depth H5 provided – depth loss disabled.")

        # -- 9. Run optimisation --
        logger.info("Starting SMPLFX optimisation (%d iters)...", cfg.max_iters)
        fitter = SMPLXOptimizer(**fitting_kwargs)
        optimized_smplx, optimized_others = fitter.optimize()

        # -- 10. Save optimised params --
        torch.save(
            {"smplx": optimized_smplx, "others": optimized_others},
            pt_path,
        )
        logger.info("Saved optimised params to %s", pt_path)

        # -- 11. Extract 3D keypoints and save as H5 --
        location_features_pt = self._save_keypoints_h5(
            smplx_kwargs, optimized_smplx, h5_path, cfg.device
        )
        logger.info("Saved 3D keypoints to %s", h5_path)
        if location_features_pt is not None:
            logger.info("Saved location features to %s", location_features_pt)

        return FittingResult(
            optimised_pt=pt_path,
            smplfx_h5=h5_path,
            num_frames=num_frames,
            location_features_pt=location_features_pt,
        )

    def export_from_optimised(
        self,
        optimised_pt: str | Path,
        output_h5: str | Path,
    ) -> FittingResult:
        """Export SMPLFX keypoints/features from an existing optimised .pt file.

        Args:
            optimised_pt: Path to ``smplfx_optimised.pt``.
            output_h5: Output path for ``body_smplfx.h5``.

        Returns:
            ``FittingResult`` with regenerated H5 and optional location features.
        """
        optimised_path = Path(optimised_pt)
        output_h5_path = Path(output_h5)
        output_h5_path.parent.mkdir(parents=True, exist_ok=True)

        if not optimised_path.exists():
            raise FileNotFoundError(f"Optimised SMPLFX file not found: {optimised_path}")

        payload = torch.load(
            str(optimised_path),
            map_location="cpu",
            weights_only=False,
        )
        if isinstance(payload, dict) and "smplx" in payload:
            smplx_params = payload["smplx"]
        elif isinstance(payload, dict):
            smplx_params = payload
        else:
            raise TypeError("Unsupported optimised SMPLFX payload type: " f"{type(payload)}")
        if not isinstance(smplx_params, dict):
            raise TypeError("Expected dict for smplx params, got " f"{type(smplx_params)}")

        smplx_kwargs = self._build_smplx_kwargs_from_params(smplx_params)
        location_features_pt = self._save_keypoints_h5(
            smplx_kwargs=smplx_kwargs,
            smplx_params=smplx_params,
            output_path=output_h5_path,
            device=self.config.device,
        )
        num_frames = self._infer_num_frames(smplx_params)
        return FittingResult(
            optimised_pt=optimised_path,
            smplfx_h5=output_h5_path,
            num_frames=num_frames,
            location_features_pt=location_features_pt,
        )

    # ------------------------------------------------------------------
    # Extract and save 3D joints
    # ------------------------------------------------------------------

    @staticmethod
    def _save_keypoints_h5(
        smplx_kwargs: dict,
        smplx_params: dict,
        output_path: Path,
        device: str = "cuda:0",
    ) -> Path | None:
        """Run SMPLFX forward passes and save 3D joints/features.

        Args:
            smplx_kwargs: Model constructor arguments.
            smplx_params: Optimised parameter dict.
            output_path: Where to write the H5 file.
            device: Torch device for the forward pass.

        Returns:
            Path to ``location_features.pt`` if vertex-based features were
            extracted successfully, otherwise ``None``.
        """
        import h5py

        from .smplfx_deps.utils_smplx import getSMPLXOut

        loc_pt_path: Path | None = None
        with torch.no_grad():
            out_camera, _model = getSMPLXOut(
                smplx_kwargs,
                device,
                smplx_params,
                BS=512,
                run_with_verts=False,
            )

        joints_3d = out_camera["joints"].cpu().numpy()  # (T, J, 3)

        with h5py.File(str(output_path), "w") as hf:
            hf.create_dataset(
                "joints3d",
                data=joints_3d.astype(np.float32),
                compression="gzip",
                compression_opts=4,
            )
            hf.attrs["num_frames"] = joints_3d.shape[0]
            hf.attrs["num_joints"] = joints_3d.shape[1]
            hf.attrs["source"] = "smplfx_fitting"

        body_centred_params = {
            key: value
            for key, value in smplx_params.items()
            if key not in {"transl", "global_orient"}
        }
        try:
            with torch.no_grad():
                out_body_centred, _model = getSMPLXOut(
                    smplx_kwargs,
                    device,
                    body_centred_params,
                    BS=512,
                    run_with_verts=True,
                )
        except Exception as exc:
            logger.warning(
                "SMPLFX body-centred forward with vertices failed (%s). "
                "location_features.pt not generated.",
                exc,
            )
            return None

        # Extract location-classifier features when vertices are available:
        # 127 SMPL-X joints + 1141 sampled vertices = (T, 1268, 3)
        try:
            verts = out_body_centred.get("vertices")
            joints = out_body_centred.get("joints")
            if verts is None or joints is None:
                logger.warning(
                    "SMPLFX output missing vertices/joints; " "location_features.pt not generated."
                )
                return None

            vertex_idx_path = (
                Path(__file__).parent
                / "smplfx_deps"
                / "assets"
                / "sampled_vertex_indices_seed42.json"
            )
            if not vertex_idx_path.exists():
                logger.warning(
                    "Sampled vertex index file not found at %s; "
                    "location_features.pt not generated.",
                    vertex_idx_path,
                )
                return None

            sampled_payload = json.loads(vertex_idx_path.read_text())
            sampled_indices = sampled_payload.get("vertex_indices")
            if sampled_indices is None:
                sampled_indices = sampled_payload.get("sampled_vertex_indices")
            if not isinstance(sampled_indices, list):
                logger.warning(
                    "Invalid sampled vertex index payload at %s; "
                    "location_features.pt not generated.",
                    vertex_idx_path,
                )
                return None

            num_frames = int(joints.shape[0])
            if int(joints.shape[1]) < 127:
                pad = torch.zeros(
                    num_frames,
                    127 - int(joints.shape[1]),
                    3,
                    device=joints.device,
                    dtype=joints.dtype,
                )
                joints_127 = torch.cat([joints, pad], dim=1)
            else:
                joints_127 = joints[:, :127, :]

            max_vert = int(verts.shape[1])
            sampled_indices = [int(i) for i in sampled_indices if 0 <= int(i) < max_vert]
            if not sampled_indices:
                logger.warning(
                    "No valid sampled vertex indices found; " "location_features.pt not generated."
                )
                return None

            sampled_verts = verts[:, sampled_indices, :]
            location_features = torch.cat([joints_127, sampled_verts], dim=1)
            loc_pt_path = output_path.parent / "location_features.pt"
            torch.save(location_features.float().cpu(), str(loc_pt_path))
            return loc_pt_path
        except Exception as exc:
            logger.warning(
                "Failed to export location_features.pt (%s). Continuing without it.",
                exc,
            )
            return None

    # ------------------------------------------------------------------
    # Cleanup
    # ------------------------------------------------------------------

    def close(self) -> None:
        """Release GPU resources."""
        torch.cuda.empty_cache()
