import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel

from .configuration_vehicle_encoder import VehicleEncoderConfig


class ResBlock(nn.Module):
    def __init__(self, in_ch, out_ch, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride, 1, bias=False)
        self.bn1   = nn.BatchNorm2d(out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, 1, 1, bias=False)
        self.bn2   = nn.BatchNorm2d(out_ch)
        if stride != 1 or in_ch != out_ch:
            self.skip = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 1, stride, bias=False),
                nn.BatchNorm2d(out_ch),
            )
        else:
            self.skip = nn.Identity()

    def forward(self, x):
        identity = self.skip(x)
        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
        out = self.bn2(self.conv2(out))
        return F.relu(out + identity, inplace=True)


def _make_stage(in_ch, out_ch, blocks, stride):
    layers = [ResBlock(in_ch, out_ch, stride=stride)]
    for _ in range(blocks - 1):
        layers.append(ResBlock(out_ch, out_ch, stride=1))
    return nn.Sequential(*layers)


class VehicleEncoderModel(PreTrainedModel):
    config_class = VehicleEncoderConfig

    def __init__(self, config: VehicleEncoderConfig):
        super().__init__(config)
        bps = config.blocks_per_stage

        self.stem = nn.Sequential(
            nn.Conv2d(3, 64, 7, 2, 3, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2, 1),
        )
        self.stage1 = _make_stage( 64,  64, bps, stride=1)
        self.stage2 = _make_stage( 64, 128, bps, stride=2)
        self.stage3 = _make_stage(128, 256, bps, stride=2)
        self.stage4 = _make_stage(256, 512, bps, stride=2)
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc  = nn.Linear(512, config.latent_dim)

        self.post_init()

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        x = self.stem(pixel_values)
        x = self.stage1(x); x = self.stage2(x)
        x = self.stage3(x); x = self.stage4(x)
        x = self.gap(x).flatten(1)
        return self.fc(x)