import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
import numpy as np
import operator as op
from utils import *
import datas
import functools as ft
import maxflow
from itertools import combinations
from itertools import product
import time

params = {'num_epochs': 300,
            'batch_size': 100,
            'learning_rate': 0.001,
            'gumbels' : 1,
            'N_K': (15,2),
            'eps_0':1.0,
            'anneal_rate':1e-5,
            'min_eps':0.1,
            'structure': True,
            'dataset':'omniglot',
            'split_valid':False,
            'binarize':True,
            'random_seed':777,
            'print_model':True,
            'print_result':True}


is_cuda = torch.cuda.is_available()
def ncr(n, r):
    r = min(r, n-r)
    numer = ft.reduce(op.mul, range(n, n-r, -1), 1)
    denom = ft.reduce(op.mul, range(1, r+1), 1)
    return numer//denom
    
def get_argmax_and_max(phi_i, phi_ij):
    #n=3
    #phi_i = [phi_0, phi_1, phi_2]
    #phi_ij = [phi_{0,1}, phi_{0,2}, phi_{1,2}]
    phi_i, phi_ij = -phi_i, -phi_ij
    g = maxflow.Graph[float]()
    nids = g.add_nodes(len(phi_i))
    obj_idx = [i for i,_ in enumerate(phi_i)]

    mapping = {ij:ind for ind,ij in enumerate(combinations(range(len(phi_i)),2))}
    for i,j in list(mapping):
        mapping[(j,i)] = mapping[(i,j)]
    i =0
    while i < len(phi_i):
        j = i+1
        s = phi_i[i]
        p = 0.5
        t = sum(p*phi_ij[mapping[(i,j_)]] for j_ in range(len(phi_i)) if i != j_ ) + s
        g.add_tedge(i,t,0)

        while j < len(phi_i):
            c = phi_ij[mapping[(i,j)]]
            g.add_edge(i,j,-p*c,-p*c)

            j +=1
        i+=1
    
    f = g.maxflow()

    return 1.0*g.get_grid_segments(nids),-f

def argmax_maxflow(phi_list):
    k_batch = []
    for phi_i,phi_ij in phi_list:
        phi = phi_i[:,1] - phi_i[:,0]
        k,_ = get_argmax_and_max(phi.detach().cpu().numpy().astype(float),phi_ij.reshape(-1).detach().cpu().numpy().astype(float))
        #k = get_argmax_and_max(phi.detach().cpu().numpy().astype(float))
        k_batch.append(k)

    k =to_var(torch.tensor(k_batch,dtype=torch.long))
    
    return k


    #print 
class Encoder(nn.Module):
    def __init__(self, image_size=784, N=6,K=6,M=20):
        super(Encoder, self).__init__()
        self.comb_num = ncr(N,K)
        self.encoder = nn.Sequential(
            nn.Linear(image_size,300),
            nn.ReLU(),
            nn.Linear(300,N*K+self.comb_num))
        
        #self.encode_phi_ij = nn.Sequential(
        #    nn.Linear(image_size,self.comb_num))

        self.N = N
        self.K = K
        self.M = M
    
        #print self.comb_num
    def forward(self, x,structure=True):
        phi_x = self.encoder(x)
        #phi_i_j = self.encode_phi_ij(x)
        phi_i, phi_i_j = phi_x.split([self.N*self.K,self.comb_num],dim=1)
        #phi_i_j = torch.exp(phi_i_j)
        phi_i_j = F.softplus(phi_i_j)
        z,phi_x_g = self.gumbel_perturbation(phi_i,phi_i_j,structure)
        return z,phi_x_g,phi_x

    def sample_gumbel(self,shape, eps=1e-20):
        #Sample from Gumbel(0, 1)
        U = torch.rand(shape).float()
        return -torch.log(eps - torch.log(U + eps))
    
    def gumbel_perturbation(self,phi_x,phi_i_j,structure = True, eps=1e-10):
        M,K,N = self.M,self.K,self.N
        phi_x=phi_x.contiguous().view(-1,K)
        
        phi_x = phi_x.repeat(M,1)
        shape = phi_x.size()
        gumbel_noise = to_var(self.sample_gumbel(shape, eps=eps))
        phi_x_gamma = phi_x + gumbel_noise
        batch_phi =list( zip(phi_x_gamma.view(-1,N,K),phi_i_j))
        #global self.structure
        if structure:
            k = argmax_maxflow(batch_phi)
        else:
            _, k = phi_x_gamma.data.max(-1)
        if is_cuda:
            z = torch.cuda.FloatTensor(*shape).zero_().scatter_(-1, k.view(-1, 1), 1.0)
        else:
            z = torch.FloatTensor(*shape).zero_().scatter_(-1, k.view(-1, 1), 1.0)
        z_phi_gamma = to_var(z)

        return z_phi_gamma,(batch_phi,k) #phi_x_gamma


class Decoder(nn.Module):
    def __init__(self, image_size=784,N=6,K=6,M=3):
        super(Decoder, self).__init__()

        self.decoder = nn.Sequential(
            nn.Linear(N*K, 300),
            nn.ReLU(),
            nn.Linear(300,image_size))
        self.N = N
        self.K = K
        self.M = M
    
    def forward(self, y):
        out = self.decoder(y.view(-1,self.N*self.K))
        return out

class Direct_VAE:
    def __init__(self,params):

        self.N,self.K  = params['N_K']
        self.M = params['gumbels']
        self.encoder = Encoder(N=self.N,K=self.K,M=self.M)
        self.decoder = Decoder(N=self.N,K=self.K,M=self.M)
        self.eps = params['eps_0']
        self.annealing_rate = params['anneal_rate']
        self.structure = params['structure']
        self.params = params
        if params['print_model']:
            
            params['print_model'] = False
            print ('encoder: ',self.encoder)
            print ('decoder: ',self.decoder)

        if torch.cuda.is_available():
            self.encoder.cuda()
            self.decoder.cuda()
        
        lr = params['learning_rate']
        self.optimizer_e = torch.optim.Adam(self.encoder.parameters(), lr=lr)
        self.optimizer_d = torch.optim.Adam(self.decoder.parameters(), lr=lr)
        self.bce_loss = nn.BCEWithLogitsLoss(reduction = 'none') #size_average=False, reduce=False
        self.training_iter = 0

    def train(self,train_loader,return_time=None):
        kl_sum,bce_sum = 0,0
        eps_0,ANNEAL_RATE,min_eps = params['eps_0'],params['anneal_rate'],params['min_eps']
        training_time = []
        for i, (im, _) in enumerate(train_loader):
            start = time.time()
            images = to_var(im.view(im.size(0), -1))
            print images[0]
            bs = im.size(0)
            ground_truth = images.repeat(self.M,1)
            # forward
            z_hard,phi_x_g,phi_x = self.encoder(images,structure=self.structure)
            out = self.decoder(z_hard)

            # backward

            encoder_loss = self.compute_encoder_gradients(z_hard,phi_x_g,ground_truth,self.eps)

            if self.M !=0:
                decoder_loss  = self.bce_loss(out,ground_truth).view(self.M,bs,-1).mean(0).sum()
            else:
                decoder_loss  = self.bce_loss(out,ground_truth).sum()
            
            kl = kl_multinomial(phi_x[:,:self.N*self.K].contiguous().view(-1,self.K))

            #encoder_loss += kl
            self.optimizer_e.zero_grad()
            self.optimizer_d.zero_grad()

            decoder_loss.backward()
            encoder_loss.backward()

            self.optimizer_d.step()
            self.optimizer_e.step()
            bce = (decoder_loss+kl)/bs
            bce_sum += bce.detach().item()
            if self.training_iter % 500 == 0:
                a = eps_0*math.exp(-ANNEAL_RATE*self.training_iter)
                if a > min_eps:
                    self.eps = a
                else:
                    self.eps=min_eps
            
            self.training_iter += 1
            end = time.time()
            training_time.append(end-start)

        nll_bce = bce_sum/len(train_loader)
        avg_time = np.mean(training_time)
        to_return =nll_bce
        
        if return_time:
            to_return = avg_time
        return to_return

    def evaluate(self,test_loader):
        self.encoder.eval()
        self.decoder.eval()
        #self.encoder.M = 100
        #self.decoder.M = 100
        #self.M = 100
        bce_sum =0
        kl_div = 0
        with torch.no_grad():
            for images, _ in test_loader:
                images = to_var(images.view(images.size(0), -1))
                
                ground_truth = images.repeat(self.M,1)
                bs = images.size(0)
                hards,_,phi_x = self.encoder(images,structure=self.structure)

                out = self.decoder(hards)

                #print self.bce_loss(out,ground_truth).sum().item()
                decoder_loss  = self.bce_loss(out,ground_truth).view(self.M,bs,-1).mean(0).sum()
                kl = kl_multinomial(phi_x[:,:self.N*self.K].contiguous().view(-1,self.K))
                bce = (decoder_loss+kl)/images.size(0)
                bce_sum += bce.detach().item()
    

        self.encoder.train()
        self.decoder.train()
        self.encoder.M = params['gumbels']
        self.decoder.M = params['gumbels']
        self.M = params['gumbels']
        nll_bce = bce_sum/len(test_loader)
        return nll_bce
    def compute_encoder_gradients(self,z_hard,phi_x_g,ground_truth,epsilon=1.0):
        N = self.N
        K = self.K
        phi_x_g,z_opt = phi_x_g
        #phi_x_g = list(phi_x_g)
        soft_copy,phi_ij = zip(*phi_x_g)
        soft_copy = torch.cat(soft_copy)
        soft_copy = soft_copy.view(-1,K)
        hard_copy = z_hard.data.view(-1,N,K)
        self.decoder.eval()
        new_batch = []
        gt_batch = []
        for n in range(N):
            a_clone = hard_copy.clone()
            idx = n*torch.ones(hard_copy.size(0),1,hard_copy.size(2)).long()
            if is_cuda:
                idx = idx.cuda()
            a_clone.scatter_(1,idx, 0)

            for k in range(K):
                clone2 = a_clone.clone()
                clone2[:,n,k]=1
                new_batch.append(clone2)
                gt_batch.append(ground_truth)
        new_batch = torch.cat(new_batch,1)
        gt_batch = torch.cat(gt_batch,1).view(-1,ground_truth.size(-1))
        
        out = self.decoder(to_var(new_batch))
        losses = self.bce_loss(out,gt_batch).sum(dim = 1) #ground_truth.repeat(K*N,1)
        
        hard_copy = hard_copy.view(-1,K)
        losses = epsilon*losses.view(-1,K).data
        soft_copy = soft_copy - losses
        
        shape = soft_copy.size()
        if not self.structure:
            _, k = soft_copy.max(-1)
        else:
            batch_phi = zip(soft_copy.view(-1,N,K),phi_ij)
            k = argmax_maxflow(batch_phi)#.view(-1,K)
        z_direct = k.tolist()
        z_opt =  z_opt.tolist()
        
        if is_cuda:
            change = torch.cuda.FloatTensor(*shape).zero_().scatter_(-1, k.view(-1, 1),1.0)
        else:
            change = torch.FloatTensor(*shape).zero_().scatter_(-1, k.view(-1, 1),1.0)
        gradients_sign_phi_i = (hard_copy - change).view(-1,N,K)
        phi_i = soft_copy.view(-1,N,K)
        grad_phi_i = (gradients_sign_phi_i*phi_i).view(-1,N*K)
        ij_sign_batch = []
        if self.structure:
            for zopt, zdirect in zip(z_opt,z_direct):
                
                zij_opt = torch.tensor([l*r for l,r in combinations(zopt,2)])
                zij_direct = torch.tensor([l*r for l,r in combinations(zdirect,2)])
                phi_ij_sign_grad =  zij_opt - zij_direct
                ij_sign_batch.append(phi_ij_sign_grad.unsqueeze(0))

            phi_ij = [p.unsqueeze(0) for p in phi_ij]
            phi_ij = torch.cat(phi_ij)
            ij_sign_batch = to_var(torch.cat(ij_sign_batch).type(torch.float))
        else:
            phi_ij = torch.cat(phi_ij).view(grad_phi_i.size(0),-1)
            ij_sign_batch = torch.zeros_like(phi_ij)
        
        grad_phi_ij = ij_sign_batch*phi_ij
        gradients = torch.cat((grad_phi_i, grad_phi_ij),-1)
        self.decoder.train()
        gradients = gradients*(1.0/epsilon)
        return torch.sum(gradients)

def training_procedure(params):
    """Trains over the MNIST standard spilt (50K/10K/10K)
    Saves the best model on validation set
    Evaluates over test set every epoch for plots"""

    train_loader,test_loader = datas.load_data(params)
    
    N,K = params['N_K']
    direct_vae = Direct_VAE(params)
    
    best_state_dicts = None
    print( 'hyper parameters: ' ,params)

    train_results,test_results = [],[]

    print ('len(train_loader)', len(train_loader))
    print ('len(test_loader)', len(test_loader))
    for epoch in range(params['num_epochs']):
        epoch_results = [0,0]
        train_nll = direct_vae.train(train_loader)
        train_results.append(train_nll)
        epoch_results[0] = train_nll

        test_nll  = direct_vae.evaluate(test_loader)
        test_results.append(test_nll)

        epoch_results[1] = test_nll
        if params['print_result']:
            print_results(epoch_results,epoch,params['num_epochs'])

    return train_results,test_results

def check_running_time(params):
    torch.manual_seed(params['random_seed'])
    params['batch_size']=1
    train_loader,test_loader = datas.load_data(params)
    print ('len(train_loader)', len(train_loader))
    print ('len(test_loader)', len(test_loader))
    print ('hyper parameters: ' ,params)
    time_to_plot = []
    nk = [(i,2) for i in range(5,16)]
    
    for n_k in nk:
        params['N_K']=n_k
        direct_vae = Direct_VAE(params)
        time = direct_vae.train(train_loader,return_time=True)
        time_to_plot.append((n_k[0],time))
        print (n_k,time)
    return time_to_plot
    
results = training_procedure(params)

