from typing import Dict

import torch
from torch.nn import functional as F
import torchvision

from ..lap_inception import LAPInception


class RSNALAPInception(LAPInception):

    def __init__(self, aux_weight: float, pool_factory, adaptive_pool_factory):
        super().__init__(aux_weight, n_classes=1, pool_factory=pool_factory, adaptive_pool_factory=adaptive_pool_factory)
    
    def forward(self, x: torch.Tensor, y: torch.Tensor = None) -> Dict[str, torch.Tensor]:

        x = x.repeat_interleave(3, dim=1)  # B 3   224 224
        if self.training:
            out, aux = torchvision.models.Inception3.forward(self, x)  # B 1
            out, aux = out.flatten(), aux.flatten()  # B
        else:
            out = torchvision.models.Inception3.forward(self, x).flatten()  # B
            aux = None

        if y is not None:
            main_loss = F.binary_cross_entropy(out, y)
            if aux is not None:
                aux_loss = F.binary_cross_entropy(aux, y)
                loss = (main_loss + self.aux_weight * aux_loss) / (1 + self.aux_weight)
            else:
                loss = main_loss

            return {
                'positive_class_probability': out,
                'loss': loss
            }

        return {
            'positive_class_probability': out
        }