#!/usr/bin/env python
# coding: utf-8

# In[1]:

print('here')
import numpy as np
import scipy as sp
import scipy.stats
import itertools
import logging
import matplotlib.pyplot as plt
import pandas as pd
import torch.utils.data as utils
import math
import time
import tqdm

import torch
import torch.optim as optim
import torch.nn.functional as F
from argparse import ArgumentParser
from torch.distributions import MultivariateNormal

import torch.nn as nn
import torch.nn.init as init
import sys

sys.path.append("../new_flows")
from flows import RealNVP, Planar, MAF
from models import NormalizingFlowModel


# ## Load and process the data

# In[2]:


testprior_data = pd.read_hdf("/data/t3home000/spark/QUASAR/preprocessing/delphes_output_4500_500_150.h5")
dt = testprior_data.values

# In[3]:


correct = (dt[:,3]>0) &(dt[:,19]>0) & (dt[:,1]>0) & (dt[:,2]>0)
dt = dt[correct]

for i in range(13,19):
    dt[:,i] = dt[:,i]/dt[:,3]

for i in range(29,35):
    dt[:,i] = dt[:,i]/(dt[:,19])


correct = (dt[:,0]>=2800)
dt = dt[correct]


#Y = dt[:,[3,4,5,6,11,12,19,20,21,22,27,28]]
sig_idx = np.where((dt[:,3]>200)  & (dt[:,19]>50) & (dt[:,0]>4200) & (dt[:,0]<4800))[0]

dt = dt[sig_idx]
#Y = dt[:,[3,4,5,6,11,12,13,14,15,16,17,18,19,20,21,22,27,28,29,30,31,32,33,34]]
#Y = dt[:,[3,4,5,6,11,12,13,14,15,16,17,18,19,20,21,22,27,28,29,30,31,32,33,34]]


# In[4]:


print(dt.shape)


# In[5]:


Y = dt[:,[3,4,5,6,11,12,19,20,21,22,27,28]]


# In[6]:


print(Y.shape)


# In[7]:



#if whichbkg == 'rndbkg':
#    dt = df_rndbkg.values
#    correct = (dt[:,3]>0) &(dt[:,19]>0) &(dt[:,1]>0) &(dt[:,2]>0)&(dt[:,16]>0)&(dt[:,32]>0)
#    dt = dt[correct]
#    for i in range(13,19):
#        dt[:,i] = dt[:,i]/dt[:,3]
#
#    for i in range(29,35):
#        dt[:,i] = dt[:,i]/(dt[:,19])
#
#    correct = (dt[:,16]>0) &(dt[:,29]>=0) &(dt[:,29]<=1)&(dt[:,30]>=0) &(dt[:,30]<=1)&(dt[:,31]>=0) &(dt[:,31]<=1)&(dt[:,32]>=0) &(dt[:,32]<=1)&(dt[:,33]>=0) &(dt[:,33]<=1)&(dt[:,34]>=-0.01) &(dt[:,34]<=1)
#    dt = dt[correct]
#    correct = (dt[:,0]>=2800)
#    dt = dt[correct]
#    #Y = dt[:,[3,4,5,6,11,12,19,20,21,22,27,28]]
#    idx = dt[:,-1]
#    bkg_idx = np.where(idx==0)[0]
#    signal_idx = np.where(idx==1)[0]
#    dt = dt[bkg_idx]
#    #Y = dt[:,[3,4,5,6,11,12,13,14,15,16,17,18,19,20,21,22,27,28,29,30,31,32,33,34]]
#    Y = dt[:,[3,4,5,6,11,12,13,14,15,16,17,18,19,20,21,22,27,28,29,30,31,32,33,34]]

sig_mean = []
sig_std = []
for i in range(12):
    mean = np.mean(Y[:,i])
    std = np.std(Y[:,i])
    sig_mean.append(mean)
    sig_std.append(std)
    Y[:,i] = (Y[:,i]-mean)/std




total_PureBkg = torch.tensor(Y)

total_PureBkg_selection = total_PureBkg




bs = 800
bkgAE_train_iterator = utils.DataLoader(total_PureBkg_selection, batch_size=bs, shuffle=True)
bkgAE_test_iterator = utils.DataLoader(total_PureBkg_selection, batch_size=bs)


# ## Build the model

# In[17]:


####MAF
class VAE_NF(nn.Module):
    def __init__(self, K, D):
        super().__init__()
        self.dim = D
        self.K = K
        self.encoder = nn.Sequential(
            nn.Linear(12, 50),
            nn.LeakyReLU(True),
            nn.Linear(50, 30),
            nn.LeakyReLU(True),
            nn.Linear(30, 20),
            nn.LeakyReLU(True),
            nn.Linear(20, D * 2)
        )

        self.decoder = nn.Sequential(
            nn.Linear(D, 20),
            nn.LeakyReLU(True),
            nn.Linear(20, 30),
            nn.LeakyReLU(True),
            nn.Linear(30, 50),
            nn.LeakyReLU(True),
            nn.Linear(50, 12)
        )

        flow_init = MAF(dim=D)
        flows_init = [flow_init for _ in range(K)]
        prior = MultivariateNormal(torch.zeros(D).cuda(), torch.eye(D).cuda())
        self.flows = NormalizingFlowModel(prior, flows_init)

    def forward(self, x):
        # Run Encoder and get NF params
        enc = self.encoder(x)
        mu = enc[:, :self.dim]
        log_var = enc[:, self.dim: self.dim * 2]

        # Re-parametrize
        sigma = (log_var * .5).exp()
        z = mu + sigma * torch.randn_like(sigma)
        kl_div = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
        # Construct more expressive posterior with NF

        z_k, _, sum_ladj = self.flows(z)

        kl_div = kl_div / x.size(0) - sum_ladj.mean()  # mean over batch

        # Run Decoder
        x_prime = self.decoder(z_k)
        return x_prime, kl_div


# ## Creating Instance¶

# In[19]:


def train():
    global n_steps
    train_loss = []
    model.train()

    for batch_idx, x in enumerate(bkgAE_train_iterator):
        start_time = time.time()

        x = x.float().cuda()

        x_tilde, kl_div = model(x)



        mseloss = nn.MSELoss(size_average=False)

        huberloss = nn.SmoothL1Loss(size_average=False)


        #loss_recons = F.binary_cross_entropy(x_tilde, x, size_average=False) / x.size(0)
        loss_recons = mseloss(x_tilde,x ) / x.size(0)

        #loss_recons = huberloss(x_tilde,x ) / x.size(0)
        loss = loss_recons + beta * kl_div

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss.append([loss_recons.item(), kl_div.item()])

        if (batch_idx + 1) % PRINT_INTERVAL == 0:
            print('\tIter [{}/{} ({:.0f}%)]\tLoss: {} Time: {:5.3f} ms/batch beta:{}'.format(
                batch_idx * len(x), 50000,
                PRINT_INTERVAL * batch_idx / 50000,
                np.asarray(train_loss)[-PRINT_INTERVAL:].mean(0),
                1000 * (time.time() - start_time),
                beta

            ))

        n_steps += 1


# In[20]:


def evaluate(split='valid'):
    global n_steps
    start_time = time.time()
    val_loss = []
    model.eval()

    with torch.no_grad():
        for batch_idx, x in enumerate(bkgAE_test_iterator):

            x = x.float().cuda()

            x_tilde, kl_div = model(x)
            mseloss = nn.MSELoss(size_average=False)
            #loss_recons = F.binary_cross_entropy(x_tilde, x, size_average=False) / x.size(0)
            huberloss = nn.SmoothL1Loss(size_average=False)


            #loss_recons = F.binary_cross_entropy(x_tilde, x, size_average=False) / x.size(0)
            loss_recons = mseloss(x_tilde,x ) / x.size(0)
            #loss_recons = huberloss(x_tilde,x ) / x.size(0)
            loss = loss_recons + beta * kl_div

            val_loss.append(loss.item())
            #writer.add_scalar('loss/{}/ELBO'.format(split), loss.item(), n_steps)
            #writer.add_scalar('loss/{}/reconstruction'.format(split), loss_recons.item(), n_steps)
            #writer.add_scalar('loss/{}/KL'.format(split), kl_div.item(), n_steps)

    print('\nEvaluation Completed ({})!\tLoss: {:5.4f} Time: {:5.3f} s'.format(
        split,
        np.asarray(val_loss).mean(0),
        time.time() - start_time
    ))
    return np.asarray(val_loss).mean(0)


# In[21]:


#version = 1
#zdim = [1,3,5,7,9,10]
#nflow = [1,3,5,7,9,10]
#lrs = [1e-3,1e-4,1e-5,1e-6,1e-7]
#betas = [0.1,0.5,1.0,2.0,10.0]#

#version = 0
#zdim = [2,4,6,8,10]
#nflow = [2,4,6,8,10]
#lrs = [1e-3,1e-4,1e-5,1e-6,1e-7]
#betas = [0.1,0.5,1.0,2.0,10.0]


# In[22]:


version = 0
#zdim = [2,4,6,8,10]
#zdim = [4,6,8,10]
#nflow = [2,4,6,8,10]
#lrs = [5e-3, 1e-3,5e-4, 1e-4,1e-5,1e-6,1e-7]
#betas = [0.1,0.5,1.0,2.0,10.0]


version = 0
#zdim = [2,4,6,8,10]
zdim = [4]
nflow = [2,4,6,8,10]
lrs = [5e-3, 1e-3,5e-4, 1e-4,1e-5,1e-6,1e-7]
betas = [0.1,0.5,1.0,2.0,10.0]

# In[43]:


N_EPOCHS = 30
PRINT_INTERVAL = 400
NUM_WORKERS = 4
n_steps = 0


# In[44]:


import re


# In[ ]:


for Z_DIM in zdim:
    for N_FLOWS in nflow:
        for beta in betas:
            model = VAE_NF(N_FLOWS, Z_DIM).cuda()
            ae_def = {
                        "type":"sig",
                        "trainon":"2prong",
                        "features":"12features",
                        "architecture":"MAF",
                        "selection":"mjjcut",
                        "trainloss":"MSELoss",
                        "beta":f"beta{re.sub('[.,]', 'p', str(beta))}",
                        "zdimnflow":f"z{Z_DIM}f{N_FLOWS}",
                        "version":f"ver{version}"

                     }

            BEST_LOSS = float('inf')
            for LR in lrs:
                optimizer = optim.Adam(model.parameters(), lr=LR)
                LAST_SAVED = -1
                PATIENCE_COUNT = 0
                for epoch in range(1, 1000):
                    print("Epoch {}:".format(epoch))
                    train()
                    cur_loss = evaluate()
                    print(cur_loss)

                    if cur_loss <= BEST_LOSS:
                        PATIENCE_COUNT = 0
                        BEST_LOSS = cur_loss
                        LAST_SAVED = epoch
                        print("Saving model!")
                        torch.save(model.state_dict(),f"/data/t3home000/spark/QUASAR/weights/{ae_def['type']}_{ae_def['trainon']}_{ae_def['features']}_{ae_def['architecture']}_{ae_def['selection']}_{ae_def['trainloss']}_{ae_def['beta']}_{ae_def['zdimnflow']}_{ae_def['version']}.h5")

                    else:
                        PATIENCE_COUNT += 1
                        print("Not saving model! Last saved: {}".format(LAST_SAVED))
                        if PATIENCE_COUNT > 10:
                            print(f"############Patience Limit Reached with LR={LR}, Best Loss={BEST_LOSS}")
                            break

                model.load_state_dict(torch.load(f"/data/t3home000/spark/QUASAR/weights/{ae_def['type']}_{ae_def['trainon']}_{ae_def['features']}_{ae_def['architecture']}_{ae_def['selection']}_{ae_def['trainloss']}_{ae_def['beta']}_{ae_def['zdimnflow']}_{ae_def['version']}.h5"))


