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

from .General.Empty import Empty
from .General.Resnet import *
from .General.MLP import *


class IDM(nn.Module):

    def __init__(self, action_size, net='vgg', pretrained=True, input=8):
        super(IDM, self).__init__()

        self.net = net

        if net == 'inception':
            self.model = models.inception_v3(pretrained=pretrained)

        elif net == 'vgg':
            self.model = models.vgg19_bn(pretrained=pretrained)
            self.model.classifier = Empty()
            self.fc_layers = nn.Sequential(
                nn.Linear((512 * 7 * 7) * 2, 4096),
                nn.LeakyReLU(),
                nn.Dropout(0.5),
                nn.Linear(4096, 4096),
                nn.LeakyReLU(),
                nn.Dropout(0.5),
                nn.Linear(4096, action_size)
            )

            if pretrained:
                print('Freezing weights')
                for params in self.model.parameters():
                    params.requires_grad = False

        elif net in ['attention', 'resnet']:
            if net == 'attention':
                self.model = ResnetFirst(normalize=False)
            else:
                self.model = Resnet(normalize=False)

            self.model.features.fc = Empty()
            self.fc_layers = nn.Sequential(
                nn.Linear(512 * 2, 512),
                nn.LeakyReLU(),
                nn.Dropout(0.5),
                nn.Linear(512, 512),
                nn.LeakyReLU(),
                nn.Dropout(0.5),
                nn.Linear(512, action_size)
            )

        elif net == 'vector':
            self.model = MlpWithAttention(input, action_size)

    def forward(self, state, nState):
        if self.net == 'vector':
            input = torch.cat((state, nState), 1)
            x = self.model(input)
        else:
            s = self.model(state)
            nS = self.model(nState)
            x = self.fc_layers(torch.cat((s, nS), 1))
        return x

    def get_params(self):
        return (x for x in torch.cat((self.reduction.parameters(), self.fc_layers.parameters())))


def train(model, data, criterion, optimizer, device, tensorboard=None):
    if model.training is False:
        model.train()

    s, nS, a = data

    s = s.to(device)
    nS = nS.to(device)
    a = a.to(device)

    optimizer.zero_grad()
    pred = model(s, nS)

    pred_argmax = torch.argmax(pred, 1)
    loss = criterion(pred, a.long())
    loss.backward()
    optimizer.step()

    acc = ((pred_argmax == a).sum().item() / a.shape[0]) * 100
    return loss, acc


def validation(model, data, device, tensorboard=None):
    if model.training is True:
        model.eval()

    s, nS, a = data

    s = s.to(device)
    nS = nS.to(device)
    a = a.to(device)

    pred = model(s, nS)

    pred_argmax = torch.argmax(pred, 1)

    acc = ((pred_argmax == a).sum().item() / a.shape[0]) * 100

    return acc
