
import torch
import torch.nn as nn
import torch.optim as optim


def supported_hyperparameters():
    return {'lr', 'momentum', 'dropout'}


class UNet2DModel(nn.Module):
    def __init__(
            self,
            sample_size=64,  # Randomly changed from 32
            in_channels=1,    # Randomly changed from 3
            out_channels=128,
            layers_per_block=2,
            block_out_channels=(32, 64, 128),
            down_block_types=("DownBlock2D", "AttnDownBlock2D", "DownBlock2D"),
            up_block_types=("UpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
    ):
        super(UNet2DModel, self).__init__()
        self.sample_size = sample_size
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.layers_per_block = layers_per_block

        self.down_blocks = nn.ModuleList()
        in_ch = in_channels
        for out_ch, block_type in zip(block_out_channels, down_block_types):
            self.down_blocks.append(self._make_block(in_ch, out_ch, block_type))
            in_ch = out_ch

        self.up_blocks = nn.ModuleList()
        for out_ch, block_type in zip(block_out_channels[::-1], up_block_types):
            self.up_blocks.append(self._make_block(in_ch, out_ch, block_type))
            in_ch = out_ch

        self.final_conv = nn.Conv2d(in_ch, out_channels, kernel_size=1)

    def _make_block(self, in_channels, out_channels, block_type):
        if block_type.startswith("Down"):
            return nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),
                nn.ReLU(),
                nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
                nn.ReLU(),
            )
        elif block_type.startswith("Up"):
            return nn.Sequential(
                nn.ConvTranspose2d(
                    in_channels, out_channels, kernel_size=2, stride=2
                ),
                nn.ReLU(),
                nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
                nn.ReLU(),
            )
        else:
            return nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
                nn.ReLU(),
            )

    def forward(self, x, timesteps):
        down_features = []
        for block in self.down_blocks:
            x = block(x)
            down_features.append(x)

        for block in self.up_blocks:
            x = block(x)

        x = self.final_conv(x)
        return nn.Identity()(x)


class Net(nn.Module):
    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
        super(Net, self).__init__()
        self.device = device

        channel_number = in_shape[1]
        image_size = in_shape[2]
        class_number = out_shape[0]

        self.unet = UNet2DModel(
            sample_size=image_size,
            in_channels=channel_number,
            out_channels=128,
            layers_per_block=2,
            block_out_channels=(32, 64, 128),
            down_block_types=("DownBlock2D", "AttnDownBlock2D", "DownBlock2D"),
            up_block_types=("UpBlock2D", "AttnUpBlock2D", "UpBlock2D"), )
        self.classifier = nn.Sequential(nn.Dropout(prm['dropout']), nn.Linear(128, class_number))
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        batch_size = x.shape[0]
        timesteps = torch.full((batch_size,), 50, dtype=torch.long, device=x.device)
        unet_output = self.unet(x, timesteps)
        unet_output = unet_output.to(torch.float32)
        pooled_features = unet_output.mean(dim=(2, 3))
        logits = self.classifier(pooled_features)
        return logits

    def train_setup(self, prm):
        self.to(self.device)
        self.criteria = nn.CrossEntropyLoss().to(self.device)

        lr = prm['lr']
        momentum = prm['momentum']

        self.optimizer = optim.SGD(self.parameters(), lr=lr, momentum=momentum)

    def learn(self, train_data):
        self.train()
        for inputs, labels in train_data:
            inputs, labels = inputs.to(self.device), labels.to(self.device)
            self.optimizer.zero_grad()
            outputs = self(inputs)
            loss = self.criteria(outputs, labels)
            loss.backward()
            nn.utils.clip_grad_norm_(self.parameters(), 3)
            self.optimizer.step()
