"""PyTorch implementation of body models and fitters."""

from __future__ import annotations

from typing import Optional

import functools
import torch
import warnings
from .bodymodel import BodyModel
from .bodyfitter import BodyFitter
from .bodyconverter import BodyConverter
from .bodyflipper import BodyFlipper
from smplfitter.common import _set_module_for_docs


__all__ = [
    'BodyModel',
    'BodyFitter',
    'BodyConverter',
    'BodyFlipper',
    'get_cached_body_model',
    'get_cached_fit_fn',
]
_set_module_for_docs(__name__, globals(), __all__)


@functools.lru_cache()
def get_cached_body_model(model_name='smpl', gender='neutral', model_root=None):
    """Return a cached BodyModel instance, creating it on first call.

    Subsequent calls with the same arguments return the same object,
    avoiding redundant file I/O and initialization.

    Parameters:
        model_name: Body model type (``'smpl'``, ``'smplx'``, ``'smplh'``, etc.).
        gender: Gender (``'neutral'``, ``'female'``, ``'male'``).
        model_root: Path to model directory. See :class:`BodyModel` for defaults.

    Returns:
        A :class:`BodyModel` instance (shared, do not modify in place).
    """
    return _get_body_model(model_name, gender, model_root)


def _get_body_model(model_name, gender, model_root=None):
    return BodyModel(model_root=model_root, gender=gender, model_name=model_name)


@functools.lru_cache()
def get_cached_fit_fn(
    body_model_name='smpl',
    gender='neutral',
    num_betas=10,
    enable_kid=False,
    requested_keys=('pose_rotvecs', 'shape_betas', 'trans'),
    beta_regularizer=1.0,
    beta_regularizer2=0.0,
    num_iter=3,
    vertex_subset=None,
    joint_regressor_post_lbs=None,
    share_beta=False,
    final_adjust_rots=True,
    scale_target=False,
    scale_fit=False,
    scale_regularizer=0.0,
    kid_regularizer=None,
    device='cuda',
):
    with torch.device(device):
        body_model = BodyModel(
            gender=gender,
            model_name=body_model_name,
            num_betas=num_betas,
            vertex_subset=torch.as_tensor(vertex_subset) if vertex_subset is not None else None,
            joint_regressor_post_lbs=joint_regressor_post_lbs,
        )
        fitter = BodyFitter(body_model, enable_kid=enable_kid)

        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            fitter = torch.jit.script(fitter)

        def fit_fn(
            verts: torch.Tensor,
            joints: Optional[torch.Tensor] = None,
            vertex_weights: Optional[torch.Tensor] = None,
            joint_weights: Optional[torch.Tensor] = None,
        ) -> dict[str, torch.Tensor]:
            res = fitter.fit(
                verts,
                target_joints=joints,
                vertex_weights=vertex_weights,
                joint_weights=joint_weights,
                num_iter=num_iter,
                beta_regularizer=beta_regularizer,
                beta_regularizer2=beta_regularizer2,
                scale_regularizer=scale_regularizer,
                kid_regularizer=kid_regularizer,
                share_beta=share_beta,
                final_adjust_rots=final_adjust_rots,
                scale_target=scale_target,
                scale_fit=scale_fit,
                requested_keys=list(requested_keys),
            )
            return {k: v for k, v in res.items()}

        def wrapped(verts, joints, vertex_weights, joint_weights):
            verts_resh = verts.view(-1, body_model.num_vertices, 3)
            joints_resh = joints.view(-1, body_model.num_joints, 3) if joints is not None else None
            vertex_weights_resh = (
                vertex_weights.view(-1, body_model.num_vertices)
                if vertex_weights is not None
                else None
            )
            joint_weights_resh = (
                joint_weights.view(-1, body_model.num_joints)
                if joint_weights is not None
                else None
            )
            res = fit_fn(verts_resh, joints_resh, vertex_weights_resh, joint_weights_resh)
            return {k: v.view(*verts.shape[:-2], *v.shape[1:]) for k, v in res.items()}

        return wrapped
