"""TEASER encoder network for FLAME face parameter estimation.

Ported from TEASER (ICLR 2025) — extracts jaw pose, expression, shape,
eyelid, head pose, and camera parameters from cropped 224x224 face images.

Only the encoder is included; the generator, FLAME model, and renderer
are not needed for parameter extraction.
"""

from __future__ import annotations

import timm
import torch
import torch.nn.functional as F
from torch import nn


def create_backbone(backbone_name: str, pretrained: bool = False) -> tuple[nn.Module, int]:
    """Create a timm feature-extraction backbone.

    Args:
        backbone_name: timm model identifier.
        pretrained: Whether to load ImageNet weights (False when loading
            from a TEASER checkpoint that already has fine-tuned weights).

    Returns:
        Tuple of (backbone module, output feature dimension).
    """
    backbone = timm.create_model(backbone_name, pretrained=pretrained, features_only=True)
    feature_dim: int = backbone.feature_info[-1]["num_chs"]
    return backbone, feature_dim


class PoseEncoder(nn.Module):
    """Predicts 3-DoF head pose and 3-DoF camera parameters."""

    def __init__(self) -> None:
        super().__init__()
        self.encoder, feature_dim = create_backbone("tf_mobilenetv3_small_minimal_100")
        self.pose_cam_layers = nn.Sequential(nn.Linear(feature_dim, 6))
        self.init_weights()

    def init_weights(self) -> None:
        self.pose_cam_layers[-1].weight.data *= 0.001
        self.pose_cam_layers[-1].bias.data *= 0.001
        self.pose_cam_layers[-1].weight.data[3] = 0
        self.pose_cam_layers[-1].bias.data[3] = 7

    def forward(self, img: torch.Tensor) -> dict[str, torch.Tensor]:
        features = self.encoder(img)[-1]
        features = F.adaptive_avg_pool2d(features, (1, 1)).squeeze(-1).squeeze(-1)
        pose_cam = self.pose_cam_layers(features).reshape(img.size(0), -1)
        return {
            "pose_params": pose_cam[..., :3],
            "cam": pose_cam[..., 3:],
        }


class ShapeEncoder(nn.Module):
    """Predicts FLAME shape (identity) parameters."""

    def __init__(self, n_shape: int = 300) -> None:
        super().__init__()
        self.encoder, feature_dim = create_backbone("tf_mobilenetv3_large_minimal_100")
        self.shape_layers = nn.Sequential(nn.Linear(feature_dim, n_shape))
        self.init_weights()

    def init_weights(self) -> None:
        self.shape_layers[-1].weight.data *= 0
        self.shape_layers[-1].bias.data *= 0

    def forward(self, img: torch.Tensor) -> dict[str, torch.Tensor]:
        features = self.encoder(img)[-1]
        features = F.adaptive_avg_pool2d(features, (1, 1)).squeeze(-1).squeeze(-1)
        parameters = self.shape_layers(features).reshape(img.size(0), -1)
        return {"shape_params": parameters}


class ExpressionEncoder(nn.Module):
    """Predicts FLAME expression, jaw pose, and eyelid parameters."""

    def __init__(self, n_exp: int = 50) -> None:
        super().__init__()
        self.encoder, feature_dim = create_backbone("tf_mobilenetv3_large_minimal_100")
        # num expressions + jaw (3) + eyelid (2)
        self.expression_layers = nn.Sequential(nn.Linear(feature_dim, n_exp + 2 + 3))
        self.n_exp = n_exp
        self.init_weights()

    def init_weights(self) -> None:
        self.expression_layers[-1].weight.data *= 0.1
        self.expression_layers[-1].bias.data *= 0.1

    def forward(self, img: torch.Tensor) -> dict[str, torch.Tensor]:
        features = self.encoder(img)[-1]
        features = F.adaptive_avg_pool2d(features, (1, 1)).squeeze(-1).squeeze(-1)
        parameters = self.expression_layers(features).reshape(img.size(0), -1)
        n = self.n_exp
        return {
            "expression_params": parameters[..., :n],
            "eyelid_params": torch.clamp(parameters[..., n : n + 2], 0, 1),
            "jaw_params": torch.cat(
                [
                    F.relu(parameters[..., n + 2].unsqueeze(-1)),
                    torch.clamp(parameters[..., n + 3 : n + 5], -0.2, 0.2),
                ],
                dim=-1,
            ),
        }


class TokenEncoder(nn.Module):
    """Encodes multi-scale visual tokens (used by generator, not for param extraction)."""

    def __init__(self) -> None:
        super().__init__()
        self.encoder, _feature_dim = create_backbone("tf_mobilenetv3_small_minimal_100")
        feature_dim_list = [576, 192, 384, 1024]
        self.token_layers = nn.ModuleList(
            [nn.Sequential(nn.Linear(feature_dim_list[i], 256)) for i in range(4)]
        )
        self.init_weights()

    def init_weights(self) -> None:
        for layer in self.token_layers:
            layer[-1].weight.data *= 0.001
            layer[-1].bias.data *= 0.001
            layer[-1].weight.data[3] = 0
            layer[-1].bias.data[3] = 7

    def forward(self, img: torch.Tensor) -> torch.Tensor:
        token_list = []
        features = self.encoder(img)

        feature = features[-1]
        feature = F.adaptive_avg_pool2d(feature, (1, 1)).squeeze(-1).squeeze(-1)
        token_list.append(self.token_layers[0](feature).reshape(img.size(0), -1))

        for i in range(2, 5):
            feature = features[-i]
            feature = F.adaptive_avg_pool2d(feature, (2 ** (i - 1), 2 ** (i - 1))).flatten(
                start_dim=1
            )
            token_list.append(self.token_layers[i - 1](feature).reshape(img.size(0), -1))

        return torch.stack(token_list, dim=0)


class TeaserEncoder(nn.Module):
    """Full TEASER encoder: pose + shape + expression + tokens.

    Args:
        n_exp: Number of expression blend-shape coefficients.
        n_shape: Number of FLAME shape (identity) coefficients.
    """

    def __init__(self, n_exp: int = 50, n_shape: int = 300) -> None:
        super().__init__()
        self.pose_encoder = PoseEncoder()
        self.shape_encoder = ShapeEncoder(n_shape=n_shape)
        self.expression_encoder = ExpressionEncoder(n_exp=n_exp)
        self.token_encoder = TokenEncoder()

    def forward(self, img: torch.Tensor) -> dict[str, torch.Tensor]:
        outputs: dict[str, torch.Tensor] = {}
        outputs.update(self.pose_encoder(img))
        outputs.update(self.shape_encoder(img))
        outputs.update(self.expression_encoder(img))
        outputs["token"] = self.token_encoder(img)
        return outputs
