import torch
import torch.nn as nn
from torchvision.models import convnext_tiny, ConvNeXt_Tiny_Weights

class ConvNeXtFeatureExtractor(nn.Module):
    def __init__(
        self,
        freeze_backbone: bool = True,
        unfreeze_last_stage: bool = False,
    ):
        super().__init__()
        weights = ConvNeXt_Tiny_Weights.DEFAULT
        self.backbone = convnext_tiny(weights=weights)

        if freeze_backbone:
            for p in self.backbone.parameters():
                p.requires_grad = False

            if unfreeze_last_stage:
                for p in self.backbone.features[-1].parameters():
                    p.requires_grad = True

        self.output_dim = self.backbone.classifier[2].in_features
        self.backbone.classifier = nn.Identity()

    def forward(self, x):
        x = self.backbone(x)
        x = x.view(x.size(0), -1)
        return x

class MLPHead(nn.Module):
    def __init__(
        self,
        input_dim: int,
        num_classes: int,
        head_depth: int = 2,
        hidden_dim_1: int = 512,
        hidden_dim_2: int = 256,
        dropout: float = 0.1,
        head_style: str = "standard",
    ):
        super().__init__()

        if head_style == "rakyan":
            if head_depth == 2:
                self.net = nn.Sequential(
                    nn.Linear(input_dim, hidden_dim_1),
                    nn.BatchNorm1d(hidden_dim_1),
                    nn.ReLU(inplace=True),
                    nn.Dropout(p=0.3),
                    nn.Linear(hidden_dim_1, num_classes),
                )
            elif head_depth == 3:
                self.net = nn.Sequential(
                    nn.Linear(input_dim, hidden_dim_1),
                    nn.BatchNorm1d(hidden_dim_1),
                    nn.ReLU(inplace=True),
                    nn.Dropout(p=0.3),
                    nn.Linear(hidden_dim_1, hidden_dim_2),
                    nn.BatchNorm1d(hidden_dim_2),
                    nn.ReLU(inplace=True),
                    nn.Dropout(p=0.3),
                    nn.Linear(hidden_dim_2, num_classes),
                )
            else:
                raise ValueError("head_depth must be 2 or 3")

        elif head_style == "standard":
            if head_depth == 2:
                self.net = nn.Sequential(
                    nn.LayerNorm(input_dim),
                    nn.Linear(input_dim, hidden_dim_1),
                    nn.LayerNorm(hidden_dim_1),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(hidden_dim_1, num_classes),
                )
            elif head_depth == 3:
                self.net = nn.Sequential(
                    nn.LayerNorm(input_dim),
                    nn.Linear(input_dim, hidden_dim_1),
                    nn.LayerNorm(hidden_dim_1),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(hidden_dim_1, hidden_dim_2),
                    nn.LayerNorm(hidden_dim_2),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(hidden_dim_2, num_classes),
                )
            else:
                raise ValueError("head_depth must be 2 or 3")
        else:
            raise ValueError("head_style must be 'standard' or 'rakyan'")

    def forward(self, x):
        return self.net(x)

class ConvNextMLP(nn.Module):
    def __init__(
        self, 
        num_classes: int = 10, 
        head_depth: int = 2, 
        hidden_dim_1: int = 512, 
        hidden_dim_2: int = 256, 
        dropout: float = 0.1, 
        freeze_backbone: bool = True, 
        unfreeze_last_stage: bool = False,
        head_style: str = "standard"
    ):
        super().__init__()
        self.feature_extractor = ConvNeXtFeatureExtractor(
            freeze_backbone=freeze_backbone, 
            unfreeze_last_stage=unfreeze_last_stage
        )
        
        self.head = MLPHead(
            input_dim=self.feature_extractor.output_dim,
            num_classes=num_classes,
            head_depth=head_depth,
            hidden_dim_1=hidden_dim_1,
            hidden_dim_2=hidden_dim_2,
            dropout=dropout,
            head_style=head_style
        )

    def forward(self, x):
        features = self.feature_extractor(x)
        out = self.head(features)
        return out