# -*- coding: utf-8 -*-
#
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# Copyright 2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
#
# Ported into VLTK from expression_classifier/flame_src/flame_file.py.
# Changes: resolved asset paths via sltk weight registry, removed numpy
# deprecation monkey-patches, added type hints.

from __future__ import annotations

import io
import logging
import pickle
import types
from pathlib import Path

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

from sltk.extraction.teaser_deps.flame.lbs import (
    batch_rodrigues,
    lbs,
    rot_mat_to_euler,
    vertices2landmarks,
)
from sltk.extraction.weights import WEIGHTS_DIR

logger = logging.getLogger(__name__)

# Base directory for FLAME assets bundled in VLTK
_FLAME_ASSETS = WEIGHTS_DIR / "teaser" / "flame_assets"


def _ensure_chumpy_stub() -> None:
    """Install a minimal chumpy stub so FLAME pickle can be loaded without chumpy.

    The FLAME generic_model.pkl was serialized with chumpy ``Ch`` arrays.
    We stub just enough for pickle to reconstruct them as numpy arrays.
    """
    import sys

    if "chumpy" in sys.modules:
        return

    # Create a fake chumpy.ch.Ch that unpickles into an ndarray
    class _ChStub(np.ndarray):
        """Minimal chumpy.ch.Ch replacement for unpickling."""

        def __new__(cls, *args: object, **kwargs: object) -> _ChStub:
            return np.ndarray.__new__(cls, shape=(0,))

        def __setstate__(self, state: object) -> None:
            # chumpy Ch serializes with __getstate__; the actual array
            # data is usually in a dict. Try to recover it.
            if isinstance(state, dict) and "x" in state:
                arr = np.array(state["x"])
                self.resize(arr.shape, refcheck=False)
                np.copyto(self, arr)
            elif isinstance(state, tuple):
                # Standard ndarray setstate
                super().__setstate__(state)

        def __reduce_ex__(self, protocol: int) -> object:
            return np.array, (np.asarray(self),)

    # Build stub module hierarchy
    chumpy_mod = types.ModuleType("chumpy")
    chumpy_ch = types.ModuleType("chumpy.ch")
    chumpy_ch.Ch = _ChStub  # type: ignore[attr-defined]
    chumpy_mod.Ch = _ChStub  # type: ignore[attr-defined]
    chumpy_mod.ch = chumpy_ch  # type: ignore[attr-defined]

    sys.modules["chumpy"] = chumpy_mod
    sys.modules["chumpy.ch"] = chumpy_ch


def _to_tensor(
    array: np.ndarray | torch.Tensor, dtype: torch.dtype = torch.float32
) -> torch.Tensor:
    if not isinstance(array, torch.Tensor):
        return torch.tensor(np.asarray(array), dtype=dtype)
    return array.to(dtype)


def _to_np(array: object, dtype: type = np.float32) -> np.ndarray:
    if "scipy.sparse" in str(type(array)):
        array = array.todense()  # type: ignore[union-attr]
    return np.array(array, dtype=dtype)


class _Struct:
    """Simple attribute container loaded from FLAME pickle."""

    def __init__(self, **kwargs: object) -> None:
        for key, val in kwargs.items():
            setattr(self, key, val)


class FLAME(nn.Module):
    """FLAME face model — differentiable mesh from shape/expression/pose params.

    Ported from https://github.com/soubhiksanyal/FLAME_PyTorch.
    Asset paths are resolved from the VLTK bundled weights directory.

    Args:
        flame_model_path: Path to generic_model.pkl.
        flame_lmk_embedding_path: Path to landmark_embedding.npy.
        n_shape: Number of shape PCA components.
        n_exp: Number of expression PCA components.
    """

    def __init__(
        self,
        flame_model_path: str | Path | None = None,
        flame_lmk_embedding_path: str | Path | None = None,
        n_shape: int = 300,
        n_exp: int = 50,
    ) -> None:
        super().__init__()

        if flame_model_path is None:
            flame_model_path = _FLAME_ASSETS / "generic_model.pkl"
        if flame_lmk_embedding_path is None:
            flame_lmk_embedding_path = _FLAME_ASSETS / "landmark_embedding.npy"

        _ensure_chumpy_stub()
        with open(flame_model_path, "rb") as fp:
            ss = pickle.load(fp, encoding="latin1")
            flame_model = _Struct(**ss)

        masks_path = _FLAME_ASSETS / "FLAME_masks.pkl"
        with open(masks_path, "rb") as fp:
            self.flame_masks = pickle.load(fp, encoding="latin1")

        self.n_shape = n_shape
        self.n_exp = n_exp
        self.dtype = torch.float32

        self.register_buffer(
            "faces_tensor",
            _to_tensor(_to_np(flame_model.f, dtype=np.int64), dtype=torch.long),
        )
        logger.debug("Using generic FLAME model")
        self.register_buffer(
            "v_template", _to_tensor(_to_np(flame_model.v_template), dtype=self.dtype)
        )

        # Shape + expression directions
        shapedirs = _to_tensor(_to_np(flame_model.shapedirs), dtype=self.dtype)
        shapedirs = torch.cat(
            [shapedirs[:, :, :n_shape], shapedirs[:, :, 300 : 300 + n_exp]], 2
        )
        self.register_buffer("shapedirs", shapedirs)

        # Pose directions
        num_pose_basis = flame_model.posedirs.shape[-1]
        posedirs = np.reshape(flame_model.posedirs, [-1, num_pose_basis]).T
        self.register_buffer(
            "posedirs", _to_tensor(_to_np(posedirs), dtype=self.dtype)
        )

        self.register_buffer(
            "J_regressor",
            _to_tensor(_to_np(flame_model.J_regressor), dtype=self.dtype),
        )
        parents = _to_tensor(_to_np(flame_model.kintree_table[0])).long()
        parents[0] = -1
        self.register_buffer("parents", parents)
        self.register_buffer(
            "lbs_weights", _to_tensor(_to_np(flame_model.weights), dtype=self.dtype)
        )

        # Eyelid offsets
        self.register_buffer(
            "l_eyelid",
            torch.from_numpy(np.load(str(_FLAME_ASSETS / "l_eyelid.npy"))).to(
                self.dtype
            )[None],
        )
        self.register_buffer(
            "r_eyelid",
            torch.from_numpy(np.load(str(_FLAME_ASSETS / "r_eyelid.npy"))).to(
                self.dtype
            )[None],
        )

        # Fixed eye and neck pose
        default_eyball_pose = torch.zeros([1, 6], dtype=self.dtype, requires_grad=False)
        self.register_parameter(
            "eye_pose", nn.Parameter(default_eyball_pose, requires_grad=False)
        )
        default_neck_pose = torch.zeros([1, 3], dtype=self.dtype, requires_grad=False)
        self.register_parameter(
            "neck_pose", nn.Parameter(default_neck_pose, requires_grad=False)
        )

        # Static and dynamic landmark embeddings
        lmk_embeddings = np.load(
            str(flame_lmk_embedding_path), allow_pickle=True, encoding="latin1"
        )
        lmk_embeddings = lmk_embeddings[()]
        self.register_buffer(
            "lmk_faces_idx",
            torch.from_numpy(lmk_embeddings["static_lmk_faces_idx"]).long(),
        )
        self.register_buffer(
            "lmk_bary_coords",
            torch.from_numpy(lmk_embeddings["static_lmk_bary_coords"]).to(self.dtype),
        )
        self.register_buffer(
            "dynamic_lmk_faces_idx",
            lmk_embeddings["dynamic_lmk_faces_idx"].long(),
        )
        self.register_buffer(
            "dynamic_lmk_bary_coords",
            lmk_embeddings["dynamic_lmk_bary_coords"].to(self.dtype),
        )
        self.register_buffer(
            "full_lmk_faces_idx",
            torch.from_numpy(lmk_embeddings["full_lmk_faces_idx"]).long(),
        )
        self.register_buffer(
            "full_lmk_bary_coords",
            torch.from_numpy(lmk_embeddings["full_lmk_bary_coords"]).to(self.dtype),
        )

        # Neck kinematic chain
        neck_kin_chain: list[torch.Tensor] = []
        NECK_IDX = 1
        curr_idx = torch.tensor(NECK_IDX, dtype=torch.long)
        while curr_idx != -1:
            neck_kin_chain.append(curr_idx)
            curr_idx = self.parents[curr_idx]
        self.register_buffer("neck_kin_chain", torch.stack(neck_kin_chain))

        # MediaPipe landmark embeddings
        mp_path = _FLAME_ASSETS / "mediapipe_landmark_embedding.npz"
        lmk_embeddings_mp = np.load(str(mp_path))
        self.register_buffer(
            "mp_lmk_faces_idx",
            torch.from_numpy(lmk_embeddings_mp["lmk_face_idx"].astype("int32")).long(),
        )
        self.register_buffer(
            "mp_lmk_bary_coords",
            torch.from_numpy(lmk_embeddings_mp["lmk_b_coords"]).to(self.dtype),
        )

    def _find_dynamic_lmk_idx_and_bcoords(
        self,
        pose: torch.Tensor,
        dynamic_lmk_faces_idx: torch.Tensor,
        dynamic_lmk_b_coords: torch.Tensor,
        neck_kin_chain: torch.Tensor,
        dtype: torch.dtype = torch.float32,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Select face contour landmarks based on head rotation."""
        batch_size = pose.shape[0]

        aa_pose = torch.index_select(
            pose.view(batch_size, -1, 3), 1, neck_kin_chain
        )
        rot_mats = batch_rodrigues(aa_pose.view(-1, 3), dtype=dtype).view(
            batch_size, -1, 3, 3
        )

        rel_rot_mat = (
            torch.eye(3, device=pose.device, dtype=dtype)
            .unsqueeze_(dim=0)
            .expand(batch_size, -1, -1)
        )
        for idx in range(len(neck_kin_chain)):
            rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)

        y_rot_angle = torch.round(
            torch.clamp(rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi, max=39)
        ).to(dtype=torch.long)

        neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
        mask = y_rot_angle.lt(-39).to(dtype=torch.long)
        neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
        y_rot_angle = neg_mask * neg_vals + (1 - neg_mask) * y_rot_angle

        dyn_lmk_faces_idx = torch.index_select(
            dynamic_lmk_faces_idx, 0, y_rot_angle
        )
        dyn_lmk_b_coords = torch.index_select(
            dynamic_lmk_b_coords, 0, y_rot_angle
        )
        return dyn_lmk_faces_idx, dyn_lmk_b_coords

    def seletec_3d68(self, vertices: torch.Tensor) -> torch.Tensor:
        """Get 68-point 3D FAN landmarks from vertices.

        Args:
            vertices: Mesh vertices (B, 5023, 3).

        Returns:
            3D landmarks (B, 68, 3).
        """
        return vertices2landmarks(
            vertices,
            self.faces_tensor,
            self.full_lmk_faces_idx.repeat(vertices.shape[0], 1),
            self.full_lmk_bary_coords.repeat(vertices.shape[0], 1, 1),
        )

    def forward(
        self,
        param_dictionary: dict[str, torch.Tensor | None],
        zero_expression: bool = False,
        zero_shape: bool = False,
        zero_pose: bool = False,
    ) -> dict[str, torch.Tensor]:
        """Forward pass: FLAME parameters -> mesh + landmarks.

        Args:
            param_dictionary: Dict with keys: shape_params, expression, pose_params,
                jaw_pose, eye_pose_params, neck_pose_params, eyelid_params.
            zero_expression: Zero out expression params.
            zero_shape: Zero out shape params.
            zero_pose: Zero out pose params.

        Returns:
            Dict with vertices, landmarks_fan, landmarks_fan_3d, landmarks_mp.
        """
        shape_params = param_dictionary.get("shape_params")
        expression_params = param_dictionary.get("expression")
        pose_params = param_dictionary.get("pose_params")
        jaw_params = param_dictionary.get("jaw_pose")
        eye_pose_params = param_dictionary.get("eye_pose_params")
        neck_pose_params = param_dictionary.get("neck_pose_params")
        eyelid_params = param_dictionary.get("eyelid_params")

        if expression_params is not None:
            batch_size, device = expression_params.shape[0], expression_params.device
        elif jaw_params is not None:
            batch_size, device = jaw_params.shape[0], jaw_params.device
        elif shape_params is not None:
            batch_size, device = shape_params.shape[0], shape_params.device
        elif pose_params is not None:
            batch_size, device = pose_params.shape[0], pose_params.device
        else:
            raise ValueError("At least one param tensor must be provided")

        if expression_params is None:
            expression_params = torch.zeros(batch_size, self.n_exp).to(device)
        if expression_params.shape[1] < self.n_exp:
            pad = torch.zeros(
                expression_params.shape[0],
                self.n_exp - expression_params.shape[1],
            ).to(device)
            expression_params = torch.cat([expression_params, pad], dim=1)

        if shape_params is None:
            shape_params = torch.zeros(batch_size, self.n_shape).to(device)
        if shape_params.shape[1] < self.n_shape:
            pad = torch.zeros(
                shape_params.shape[0], self.n_shape - shape_params.shape[1]
            ).to(device)
            shape_params = torch.cat([shape_params, pad], dim=1)

        if pose_params is None:
            pose_params = torch.zeros(batch_size, 3).to(device)

        if zero_expression:
            expression_params = torch.zeros_like(expression_params).to(device)
            jaw_params = torch.zeros_like(jaw_params).to(device) if jaw_params is not None else None

        if zero_shape:
            shape_params = torch.zeros_like(shape_params).to(device)

        if zero_pose:
            pose_params = torch.zeros_like(pose_params).to(device)
            pose_params[..., 0] = 0.2
            pose_params[..., 1] = -0.7

        if eye_pose_params is None:
            eye_pose_params = self.eye_pose.expand(batch_size, -1)

        if neck_pose_params is None:
            neck_pose_params = self.neck_pose.expand(batch_size, -1)

        if jaw_params is None:
            jaw_params = torch.zeros(batch_size, 3).to(device)

        betas = torch.cat([shape_params, expression_params], dim=1)
        full_pose = torch.cat(
            [pose_params, neck_pose_params, jaw_params, eye_pose_params], dim=1
        )

        template_vertices = self.v_template.unsqueeze(0).expand(batch_size, -1, -1)

        vertices, _ = lbs(
            betas,
            full_pose,
            template_vertices,
            self.shapedirs,
            self.posedirs,
            self.J_regressor,
            self.parents,
            self.lbs_weights,
            dtype=self.dtype,
        )

        if eyelid_params is not None:
            vertices = (
                vertices
                + self.r_eyelid.expand(batch_size, -1, -1) * eyelid_params[:, 1:2, None]
            )
            vertices = (
                vertices
                + self.l_eyelid.expand(batch_size, -1, -1) * eyelid_params[:, 0:1, None]
            )

        # 2D landmarks (dynamic contour + static)
        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(batch_size, -1)
        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).expand(
            batch_size, -1, -1
        )

        dyn_lmk_faces_idx, dyn_lmk_bary_coords = (
            self._find_dynamic_lmk_idx_and_bcoords(
                full_pose,
                self.dynamic_lmk_faces_idx,
                self.dynamic_lmk_bary_coords,
                self.neck_kin_chain,
                dtype=self.dtype,
            )
        )
        lmk_faces_idx = torch.cat([dyn_lmk_faces_idx, lmk_faces_idx], 1)
        lmk_bary_coords = torch.cat([dyn_lmk_bary_coords, lmk_bary_coords], 1)

        landmarks2d = vertices2landmarks(
            vertices, self.faces_tensor, lmk_faces_idx, lmk_bary_coords
        )
        bz = vertices.shape[0]
        landmarks3d = vertices2landmarks(
            vertices,
            self.faces_tensor,
            self.full_lmk_faces_idx.repeat(bz, 1),
            self.full_lmk_bary_coords.repeat(bz, 1, 1),
        )
        landmarksmp = vertices2landmarks(
            vertices,
            self.faces_tensor,
            self.mp_lmk_faces_idx.repeat(bz, 1),
            self.mp_lmk_bary_coords.repeat(bz, 1, 1),
        )

        return {
            "vertices": vertices,
            "landmarks_fan": landmarks2d,
            "landmarks_fan_3d": landmarks3d,
            "landmarks_mp": landmarksmp,
        }
