import torch
import torch.nn as nn
import torch.nn.functional as F  # noqa: F401

from sltk.embedding.signrep_models.backbones.hiera_nodec_mae import mae_hiera_base_16x224
from sltk.embedding.signrep_models.heads.dict_head import ProjectionModel


class Model(nn.Module):
    def __init__(
        self,
        backbone_params,
        ckpt_dir,
        mask_ratio,
        out_dim,
        cls_head_params,
        init_weights=True,
    ):
        super().__init__()
        self.backbone = mae_hiera_base_16x224(num_classes=400, pretrained=False, **backbone_params)
        if ckpt_dir is not None:
            ckpt = torch.load(ckpt_dir, map_location="cpu", weights_only=True)

            res = self.backbone.load_state_dict(ckpt["model_state"], strict=False)
            print(res)
        self.mask_ratio = mask_ratio
        self.num_features = out_dim

        try:
            self.cls_head = ProjectionModel(**cls_head_params, in_dim=self.num_features)

            self.fc = torch.nn.Sequential(
                nn.LayerNorm(out_dim),
                nn.Linear(out_dim, out_dim, bias=False),
            )

            self.active_head = torch.nn.Sequential(
                torch.nn.Linear(out_dim, 128),
                torch.nn.GELU(),
                torch.nn.Linear(128, 2),
                # torch.nn.Sigmoid(),
            )

            if init_weights:
                self.cls_head.apply(self._init_weights)
        except Exception:
            self.cls_head = nn.Identity()
            print("SKIPPED CLASS HEAD")

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.xavier_uniform_(m.weight)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward(self, imgs, extract_features=False, apply_results=False):
        b, c, t, h, w = imgs.shape

        # if extract_features:
        #     return {"features": self.backbone(imgs, eval=True)}
        latent, intermediates_feature_map, intermediate_features_cls = self.backbone(
            imgs,
            mask_ratio=self.mask_ratio if self.training else 0.0,
            # eval=not self.training,
        )

        res = self.fc(latent)

        cls_y = self.cls_head(res)
        cls_active = self.active_head(res)

        return {
            "intermediates_feature_map": intermediates_feature_map,
            "intermediate_features_cls": intermediate_features_cls,
            "features": res,
            "latent": latent,
            # "post_features": res,
            "cls_head": cls_y,
            "cls_active": cls_active,
        }

    def inference_forward(self, imgs, need_active=True):
        """Fast inference path that skips intermediates and cls_head.

        Uses backbone eval=True path which calls Hiera.forward with
        return_intermediates=False, skipping all reroll operations at
        stage boundaries and the cls_head (pose/keypoint projections).

        Returns raw backbone latent features (768-dim). The fc projection
        head is only computed internally when need_active=True (the
        active_head takes fc output as input) but is never exposed.
        """
        latent = self.backbone(imgs, mask_ratio=0.0, eval=True)
        result = {"latent": latent}
        if need_active:
            projected = self.fc(latent)
            result["cls_active"] = self.active_head(projected)
        return result
