import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from dspn import DSPN
from models import FSEncoder, F_reconstruct
import random



class ConvInputModel(nn.Module):
    def __init__(self):
        super(ConvInputModel, self).__init__()
        
        self.conv1 = nn.Conv2d(3, 24, 3, stride=2, padding=1)
        self.batchNorm1 = nn.BatchNorm2d(24)
        self.conv2 = nn.Conv2d(24, 24, 3, stride=2, padding=1)
        self.batchNorm2 = nn.BatchNorm2d(24)
        self.conv3 = nn.Conv2d(24, 24, 3, stride=2, padding=1)
        self.batchNorm3 = nn.BatchNorm2d(24)
        self.conv4 = nn.Conv2d(24, 24, 3, stride=2, padding=1)
        self.batchNorm4 = nn.BatchNorm2d(24)

        
    def forward(self, img):
        """convolution"""
        x = self.conv1(img)
        x = F.relu(x)
        x = self.batchNorm1(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.batchNorm2(x)
        x = self.conv3(x)
        x = F.relu(x)
        x = self.batchNorm3(x)
        x = self.conv4(x)
        x = F.relu(x)
        x = self.batchNorm4(x)
        return x

  
class FCOutputModel(nn.Module):
    def __init__(self):
        super(FCOutputModel, self).__init__()

        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.fc2(x)
        x = F.relu(x)
        x = F.dropout(x)
        x = self.fc3(x)
        return F.log_softmax(x)

  

class BasicModel(nn.Module):
    def __init__(self, args, name):
        super(BasicModel, self).__init__()
        self.name=name
        self.args=args

    def train_(self, input_img, input_qst, label):
        self.optimizer.zero_grad()
        output, objects = self(input_img, input_qst)
        if objects == None:
            sparse_loss = 0
        else:
            sparse_loss = objects.permute(0,2,1).reshape(-1, objects.shape[2]).norm(p=2, dim=1).sum()
        loss = F.nll_loss(output, label) + self.args.sparse_loss*sparse_loss
        if random.randint(0,100)==0:
            print(loss, self.args.sparse_loss*sparse_loss, sparse_loss)
        loss.backward()
        self.optimizer.step()
        pred = output.data.max(1)[1]
        correct = pred.eq(label.data).cpu().sum()
        accuracy = correct * 100. / len(label)
        return accuracy
        
    def test_(self, input_img, input_qst, label):
        output, objects = self(input_img, input_qst)
        pred = output.data.max(1)[1]
        correct = pred.eq(label.data).cpu().sum()
        accuracy = correct * 100. / len(label)
        return accuracy

    def save_model(self, epoch):
        torch.save(self.state_dict(), 'model/epoch_{}_{:02d}.pth'.format(self.name, epoch))


class RN(BasicModel):
    def __init__(self, args):
        super(RN, self).__init__(args, 'SRN' if args.use_srn else 'RN')
   
        self.conv = ConvInputModel()
        self.use_coord = True
        
        self.USE_SRN = args.use_srn
        if self.USE_SRN:
            iters = args.inner_iters
            lr = args.inner_lr
            CNN_ENCODER_SPACE = 25
            self.set_size = 25
            self.element_dims = 24 + 2*self.use_coord
            self.encoder = FSEncoder(self.element_dims, CNN_ENCODER_SPACE, 512)
            self.decoder = DSPN(
                        self.encoder, self.element_dims, iters=iters, lr=lr
            )
            self.proj_vector = nn.Linear((24+self.use_coord*2)*25, CNN_ENCODER_SPACE)
        
        ##(number of filters per object+coordinate of object)*2+question vector
        self.g_fc1 = nn.Linear((24+self.use_coord*2)*2+11, 256)

        self.g_fc2 = nn.Linear(256, 256)
        self.g_fc3 = nn.Linear(256, 256)
        self.g_fc4 = nn.Linear(256, 256)

        self.f_fc1 = nn.Linear(256, 256)

        self.coord_oi = torch.FloatTensor(args.batch_size, 2)
        self.coord_oj = torch.FloatTensor(args.batch_size, 2)
        if args.cuda:
            self.coord_oi = self.coord_oi.cuda()
            self.coord_oj = self.coord_oj.cuda()
        self.coord_oi = Variable(self.coord_oi)
        self.coord_oj = Variable(self.coord_oj)

        # prepare coord tensor
        def cvt_coord(i):
            return [(i/5-2)/2., (i%5-2)/2.]
        
        self.coord_tensor = torch.FloatTensor(args.batch_size, 25, 2)
        if args.cuda:
            self.coord_tensor = self.coord_tensor.cuda()
        self.coord_tensor = Variable(self.coord_tensor)
        np_coord_tensor = np.zeros((args.batch_size, 25, 2))
        for i in range(25):
            np_coord_tensor[:,i,:] = np.array( cvt_coord(i) )
        self.coord_tensor.data.copy_(torch.from_numpy(np_coord_tensor))

        self.fcout = FCOutputModel()
        
        self.optimizer = optim.Adam(self.parameters(), lr=args.lr)


    def forward(self, img, qst):
        x = self.conv(img) ## x = (64 x 24 x 5 x 5)
        """g"""
        mb = x.size()[0]
        n_channels = x.size()[1]
        d = x.size()[2]
        x_flat = x.view(mb,n_channels,d*d)
        if self.use_coord:
            x_flat = torch.cat([x_flat, self.coord_tensor.permute(0,2,1)], 1)

        if self.USE_SRN:
            vec = self.proj_vector(x_flat.reshape(x_flat.shape[0], -1))
            intermediate_sets, losses, grad_norms = self.decoder(vec, x_flat)
            if random.randint(0, 100) == 0:
                print('inner losses: ', [round(l.item()/x.shape[0], 2) for l in losses])
            x_flat = intermediate_sets[-1]
        x_flat = x_flat.permute(0,2,1)
 
        
        # add coordinates
        
        # add question everywhere
        qst = torch.unsqueeze(qst, 1)
        qst = qst.repeat(1,25,1)
        qst = torch.unsqueeze(qst, 2)
        
        # cast all pairs against each other
        x_i = torch.unsqueeze(x_flat,1) # (64x1x25x26+11)
        x_i = x_i.repeat(1,25,1,1) # (64x25x25x26+11)
        x_j = torch.unsqueeze(x_flat,2) # (64x25x1x26+11)
        x_j = torch.cat([x_j,qst],3)
        x_j = x_j.repeat(1,1,25,1) # (64x25x25x26+11)
        
        # concatenate all together
        x_full = torch.cat([x_i,x_j],3) # (64x25x25x2*26+11)
        
        # reshape for passing through network
        x_ = x_full.view(mb*d*d*d*d,(24+2*self.use_coord)*2 + 11)
        x_ = self.g_fc1(x_)
        x_ = F.relu(x_)
        x_ = self.g_fc2(x_)
        x_ = F.relu(x_)
        x_ = self.g_fc3(x_)
        x_ = F.relu(x_)
        x_ = self.g_fc4(x_)
        x_ = F.relu(x_)
        
        # reshape again and sum
        x_g = x_.view(mb,d*d*d*d,256)
        x_g = x_g.sum(1).squeeze()
        
        """f"""
        x_f = self.f_fc1(x_g)
        x_f = F.relu(x_f)
        
        return self.fcout(x_f), x_flat

class RN_reconstruct(nn.Module):
    def __init__(self, args):
        super(RN_reconstruct, self).__init__()
   
        self.conv = ConvInputModel()
        
        self.USE_SRN = args.use_srn
        if self.USE_SRN:
            iters = args.inner_iters
            lr = args.inner_lr
            CNN_ENCODER_SPACE = 100
            self.set_size = 25
            self.element_dims = 24
            self.encoder = FSEncoder(self.element_dims, CNN_ENCODER_SPACE, 512)
            self.decoder = DSPN(
                        self.encoder, self.element_dims, iters=iters, lr=lr
            )
            self.proj_vector = nn.Linear(24*25, CNN_ENCODER_SPACE)
        
        ##(number of filters per object+coordinate of object)*2+question vector
        self.g_fc1 = nn.Linear((24)*2+11, 256)

        self.g_fc2 = nn.Linear(256, 256)
        self.g_fc3 = nn.Linear(256, 256)
        self.g_fc4 = nn.Linear(256, 256)

        self.f_fc1 = nn.Linear(256, 256)
        
        self.fcout = FCOutputModel()
        self.reconstruct = F_reconstruct(24)
        

    def forward(self, img):
        x = self.conv(img) ## x = (64 x 24 x 5 x 5)
        """g"""
        mb = x.size()[0]
        n_channels = x.size()[1]
        d = x.size()[2]
        x_flat = x.view(mb,n_channels,d*d)
        if self.USE_SRN:
            vec = self.proj_vector(x_flat.reshape(x_flat.shape[0], -1))
            intermediate_sets, losses, grad_norms = self.decoder(vec, x_flat)
            x_flat = intermediate_sets[-1]
        x_flat = x_flat.permute(0,2,1)
        generated_f, generated_set = self.reconstruct(x_flat.permute(0,2,1))
        if self.USE_SRN:
            return generated_f, generated_set, losses
        else:
            return generated_f, generated_set, []
        
      

class CNN_MLP(BasicModel):
    def __init__(self, args):
        super(CNN_MLP, self).__init__(args, 'CNNMLP')

        self.conv  = ConvInputModel()
        self.fc1   = nn.Linear(5*5*24 + 11, 256)  # question concatenated to all
        self.fcout = FCOutputModel()

        self.optimizer = optim.Adam(self.parameters(), lr=args.lr)
        #print([ a for a in self.parameters() ] )
  
    def forward(self, img, qst):
        x = self.conv(img) ## x = (64 x 24 x 5 x 5)

        """fully connected layers"""
        x = x.view(x.size(0), -1)
        
        x_ = torch.cat((x, qst), 1)  # Concat question
        
        x_ = self.fc1(x_)
        x_ = F.relu(x_)
        
        return self.fcout(x_), None

