from dataclasses import dataclass
from typing import Optional

import timm
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

try:  # relative import works as Hub remote code; absolute works for local scripts
    from .configuration_edgeface import EdgeFaceConfig
except ImportError:
    from configuration_edgeface import EdgeFaceConfig


# ---------------------------------------------------------------------------
# Static low-rank linear factorization (EdgeFace's "gamma" trick).
# This is baked into the pretrained weights and is unrelated to PEFT/LoRA
# adapters -- naming it LowRankLinear keeps "lora" free for real adapters.
# The submodule attribute names (linear1, linear2) are kept so the original
# published checkpoints load unchanged.
# ---------------------------------------------------------------------------
class LowRankLinear(nn.Module):
    def __init__(self, in_features, out_features, rank, bias=True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.linear1 = nn.Linear(in_features, rank, bias=False)
        self.linear2 = nn.Linear(rank, out_features, bias=bias)

    def forward(self, x):
        return self.linear2(self.linear1(x))


def _factorize_recursive(module, ratio):
    for name, child in module.named_children():
        if isinstance(child, nn.Linear) and "head" not in name:
            rank = max(2, int(min(child.in_features, child.out_features) * ratio))
            bias = child.bias is not None
            setattr(module, name, LowRankLinear(child.in_features, child.out_features, rank, bias))
        else:
            _factorize_recursive(child, ratio)


def factorize_linear_layers(module, ratio=0.2):
    """Replace eligible nn.Linear layers with LowRankLinear, in place."""
    _factorize_recursive(module, ratio)
    return module


@dataclass
class EdgeFaceOutput(ModelOutput):
    embeddings: Optional[torch.FloatTensor] = None


class EdgeFaceModel(PreTrainedModel):
    config_class = EdgeFaceConfig
    main_input_name = "pixel_values"
    input_modalities="image"

    def __init__(self, config: EdgeFaceConfig):
        super().__init__(config)

        # Keep the attribute named `self.model` so the original published
        # checkpoints (keys prefixed with "model.") load unchanged.
        self.model = timm.create_model(config.timm_model)
        self.model.reset_classifier(config.featdim)

        if config.use_low_rank:
            factorize_linear_layers(self.model, ratio=config.low_rank_ratio)

        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.trunc_normal_(module.weight, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(
        self,
        pixel_values: torch.FloatTensor,
        normalize: bool = False,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        embeddings = self.model(pixel_values)

        # Reference code does not normalize inside the model (it normalizes at
        # comparison time via cosine similarity). Off by default for parity.
        if normalize:
            embeddings = F.normalize(embeddings, dim=-1)

        if not return_dict:
            return (embeddings,)
        return EdgeFaceOutput(embeddings=embeddings)