from __future__ import annotations

from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file


class ConvAct(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        groups: int = 1,
        bias: bool = True,
        activation: str | None = "hardswish",
    ) -> None:
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=kernel_size // 2,
            groups=groups,
            bias=bias,
        )
        self.activation = activation

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv(x)
        if self.activation == "hardswish":
            return F.hardswish(x)
        if self.activation == "relu":
            return F.relu(x)
        if self.activation is None:
            return x
        raise ValueError(f"Unsupported activation: {self.activation}")


class DepthwisePointwise(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
    ) -> None:
        super().__init__()
        self.depthwise = ConvAct(
            in_channels,
            in_channels,
            kernel_size,
            stride=stride,
            groups=in_channels,
        )
        self.pointwise = ConvAct(in_channels, out_channels, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.pointwise(self.depthwise(x))


class SqueezeExcite(nn.Module):
    def __init__(self, channels: int, hidden_channels: int) -> None:
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.reduce = nn.Conv2d(channels, hidden_channels, kernel_size=1, bias=True)
        self.expand = nn.Conv2d(hidden_channels, channels, kernel_size=1, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        scale = self.pool(x)
        scale = F.relu(self.reduce(scale))
        scale = F.hardsigmoid(self.expand(scale))
        return x * scale


class MinerUOrientationClassifier(nn.Module):
    def __init__(self, num_classes: int = 4) -> None:
        super().__init__()
        self.stem = ConvAct(3, 16, kernel_size=3, stride=2)
        self.blocks = nn.ModuleList(
            [
                DepthwisePointwise(16, 32, kernel_size=3),
                DepthwisePointwise(32, 64, kernel_size=3, stride=2),
                DepthwisePointwise(64, 64, kernel_size=3),
                DepthwisePointwise(64, 128, kernel_size=3, stride=2),
                DepthwisePointwise(128, 128, kernel_size=3),
                DepthwisePointwise(128, 256, kernel_size=3, stride=2),
                DepthwisePointwise(256, 256, kernel_size=5),
                DepthwisePointwise(256, 256, kernel_size=5),
                DepthwisePointwise(256, 256, kernel_size=5),
                DepthwisePointwise(256, 256, kernel_size=5),
                DepthwisePointwise(256, 256, kernel_size=5),
            ]
        )
        self.downsample_256 = ConvAct(
            256,
            256,
            kernel_size=5,
            stride=2,
            groups=256,
        )
        self.se_256 = SqueezeExcite(256, 64)
        self.expand_512 = ConvAct(256, 512, kernel_size=1)
        self.depthwise_512 = ConvAct(512, 512, kernel_size=5, groups=512)
        self.se_512 = SqueezeExcite(512, 128)
        self.project_512 = ConvAct(512, 512, kernel_size=1)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.head_conv = nn.Conv2d(512, 1280, kernel_size=1, bias=False)
        self.classifier = nn.Linear(1280, num_classes, bias=True)

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        for block in self.blocks:
            x = block(x)
        x = self.downsample_256(x)
        x = self.se_256(x)
        x = self.expand_512(x)
        x = self.depthwise_512(x)
        x = self.se_512(x)
        x = self.project_512(x)
        return x

    def forward_logits(self, x: torch.Tensor) -> torch.Tensor:
        x = self.forward_features(x)
        x = self.pool(x)
        x = F.hardswish(self.head_conv(x))
        x = x * 0.8
        x = x.reshape(x.shape[0], 1280)
        return self.classifier(x)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.softmax(self.forward_logits(x), dim=1)


def from_file(model_dir: str | Path, device: str = "cpu") -> MinerUOrientationClassifier:
    model_dir = Path(model_dir)
    model = MinerUOrientationClassifier()
    state = load_file(model_dir / "model.safetensors", device=device)
    model.load_state_dict(state, strict=True)
    return model.to(device).eval()