"""Two-branch CNN + azimuth metadata fusion model."""

import torch
from torch import nn
from torchvision.models import ResNet18_Weights, resnet18


class LunarFusionModel(nn.Module):
    def __init__(self, pretrained: bool = True, metadata_dropout: float = 0.1,
                 classifier_dropout: float = 0.3, channels: int = 1, use_metadata: bool = True):
        super().__init__()
        weights = ResNet18_Weights.DEFAULT if pretrained else None
        backbone = resnet18(weights=weights)

        old_conv = backbone.conv1
        if channels not in (1, 3):
            raise ValueError('channels must be 1 or 3')
        if channels == 1:
            backbone.conv1 = nn.Conv2d(1, old_conv.out_channels, kernel_size=old_conv.kernel_size,
                                       stride=old_conv.stride, padding=old_conv.padding, bias=False)
            with torch.no_grad():
                backbone.conv1.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
        self.use_metadata = use_metadata
        image_features = backbone.fc.in_features
        backbone.fc = nn.Identity()
        self.image_branch = backbone

        self.metadata_branch = nn.Sequential(
            nn.Linear(2, 32), nn.ReLU(inplace=True),
            nn.Dropout(metadata_dropout), nn.Linear(32, 16), nn.ReLU(inplace=True),
        )
        self.classifier = nn.Sequential(
            nn.Linear(image_features + 16, 128), nn.ReLU(inplace=True),
            nn.Dropout(classifier_dropout), nn.Linear(128, 1),
        )

    def forward(self, image: torch.Tensor, azimuth: torch.Tensor) -> torch.Tensor:
        image_features = self.image_branch(image)
        metadata_features = self.metadata_branch(azimuth) if self.use_metadata else image_features.new_zeros((len(image), 16))
        return self.classifier(torch.cat([image_features, metadata_features], dim=1)).squeeze(1)