import torch as pt
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import ImageClassifierOutput
from .configuration_nula import NulaConfig

class BlurPool2d(nn.Module):
    def __init__(self, channels, stride=2):
        super().__init__()
        if stride not in (1, 2):
            raise ValueError
        self.channels = channels
        self.stride = stride
        kernel_1d = pt.tensor([1.0, 2.0, 1.0], dtype=pt.float32)
        kernel_2d = kernel_1d[:, None] * kernel_1d[None, :]
        kernel_2d = kernel_2d / kernel_2d.sum()
        filt = kernel_2d[None, None, :, :].repeat(channels, 1, 1, 1)
        self.register_buffer("filt", filt)
        self.pad = nn.ReflectionPad2d(1)

    def forward(self, x):
        if x.ndim != 4:
            raise ValueError("x needs shape (B, C, H, W)")
        if x.shape[1] != self.channels:
            raise ValueError(f"expected : {self.channels} channels\ngot : {x.shape[1]}")
        x = self.pad(x)
        x = F.conv2d(
            x, 
            self.filt, 
            stride=self.stride, 
            padding=0, 
            groups=self.channels
        )
        return x

class SEBlock(nn.Module):
    def __init__(self, channels, se_reduction=16):
        super().__init__()
        hidden = max(channels // se_reduction, 1)
        self.pool = nn.AdaptiveAvgPool2d((1,1))
        self.fc1 = nn.Linear(channels, hidden)
        self.act = nn.ReLU(inplace=True)
        self.fc2 = nn.Linear(hidden, channels)
        self.gate = nn.Sigmoid()

    def forward(self, x):
        b, c, h, w = x.shape
        z = self.pool(x)
        z = z.view(b, c)
        z = self.act(self.fc1(z))
        z = self.gate(self.fc2(z))

        z = z.view(b, c, 1, 1)
        return x * z


class NulaConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch, downsample=False, use_se=False, se_reduction=16):
        super().__init__()
        stride = 1
        self.conv1 = nn.Conv2d(
            in_ch,
            out_ch,
            kernel_size=3,
            stride=stride,
            padding=1,
            bias=False
        )
        self.bn1 = nn.BatchNorm2d(out_ch)
        self.act1 = nn.SiLU(inplace=True)
        self.conv2 = nn.Conv2d(
            out_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_ch)
        
        self.se = SEBlock(out_ch, se_reduction=se_reduction) if use_se else nn.Identity()
        self.blurpool = BlurPool2d(out_ch, stride=2) if downsample else nn.Identity()
        if downsample or in_ch != out_ch:
            layers = []
            layers.append(nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=1, bias=False))
            layers.append(nn.BatchNorm2d(out_ch))
            if downsample:
                layers.append(BlurPool2d(out_ch, stride=2))
            self.shortcut = nn.Sequential(*layers)
        else:
            self.shortcut = nn.Identity()
        self.out_act = nn.SiLU(inplace=True)

    def forward(self, x):
        identity = self.shortcut(x)

        out = self.act1(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out = self.se(out)
        out = self.blurpool(out)
        out = out + identity
        out = self.out_act(out)
        return out

class NulaForImageClassification(PreTrainedModel):
    config_class = NulaConfig
    base_model_prefix = "nula"

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

        dims = list(config.block_channels)
        if len(dims) != 3:
            raise ValueError("For v0.1, block_channels must exactly have 3 entries :>")
        c1, c2, c3 = dims
        
        self.stem_conv = nn.Conv2d(
            in_channels=config.in_channels,
            out_channels=c1,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=False
        )
        self.stem_bn = nn.BatchNorm2d(c1)
        self.stem_act = nn.SiLU(inplace=True)

        self.s1 = NulaConvBlock(c1, c1, False, config.use_se, config.se_reduction)
        self.s2 = NulaConvBlock(c1, c2, True, config.use_se, config.se_reduction)
        self.s3 = NulaConvBlock(c2, c3, True, config.use_se, config.se_reduction)
        self.global_pool = nn.AdaptiveAvgPool2d((1,1))
        self.head = nn.Sequential(
            nn.Linear(c3, config.classifier_hidden_dim),
            nn.SiLU(inplace=True),
            nn.Dropout(p=0.30),
            nn.Linear(config.classifier_hidden_dim, config.num_classes)
        )
        self.criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Conv2d):
            nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
        elif isinstance(module, nn.BatchNorm2d):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            nn.init.zeros_(module.bias)
        if isinstance(module, NulaConvBlock) and hasattr(module, "bn2"):
            nn.init.zeros_(module.bn2.weight)

    def forward(self, pixel_values=None, labels=None, return_dict=None):
        self._val_invars(pixel_values)
        return_dict = return_dict if return_dict is not None else self.config.return_dict
        x = pixel_values

        x = self.stem_act(self.stem_bn(self.stem_conv(x)))
        x = self.s1(x)
        x = self.s2(x)
        x = self.s3(x)
        x = self.global_pool(x)
        x = pt.flatten(x, 1)
        logits = self.head(x)

        loss = None
        if labels is not None:
            loss = self.criterion(logits, labels)
        if not return_dict:
            output = (logits,)
            return ((loss,) + output) if loss is not None else output

        return ImageClassifierOutput(
            loss=loss,
            logits=logits
        )

    def _val_invars(self, pixel_values):
        if pixel_values is None:
            raise ValueError("pixel_values must be provided")
        if pixel_values.ndim != 4:
            raise ValueError(f"Expected (B, C, H, W), got {pixel_values.shape}")

        B, C, H, W = pixel_values.shape
        if C != self.config.in_channels:
            raise ValueError(f"Channel mismatch:\nexpected : {self.config.in_channels},\ngot : {C}")