import pickle
import os
import numpy as np
import matplotlib.pyplot as plt
import torch.utils.data as Data
from tqdm import tqdm
from config import config
from datetime import datetime

import torch
import torch.nn as nn
import torch.nn.functional as F
# BATCH_SIZE=64
# LR=0.01
# EPOCH=100
init_seed = 1222
torch.manual_seed(init_seed)
# torch.cuda.manual_seed(init_seed)
np.random.seed(init_seed) # 用于numpy的随机数

data_dir='./cifar-10-batches-py'
meta_path=os.path.join(data_dir,'batches.meta')
test_path=os.path.join(data_dir,'test_batch')
train_pathname='data_batch_'
log_path='./log/log'
train_path = os.path.join(data_dir, train_pathname )
now=0
def metric(y_true,pred):
    # print(y_true)
    # print(pred)
    TP=torch.tensor([0.]*2)
    FP=torch.tensor([0.]*2)
    FN=torch.tensor([0.]*2)
    for y,p in zip(y_true,pred):
        if(y==p):
            TP[y]+=1
        else:
            FP[p]+=1
            FN[y]+=1
    # print(TP,FP,FN)
    # print(TP+FP)
    precision=TP/(TP+FP+1e-4)
    recall=TP/(TP+FN+1e-4)
    # print(precision,recall)
    F1=2*precision*recall/(precision+recall+1e-4)
    mean_F1=torch.mean(F1)
    return F1,mean_F1

def unpickle(file):
    with open(file,'rb') as f:
        dict=pickle.load(f,encoding='bytes')
    return dict

def pickle_save(file,t):
    with open(file,'wb') as f:
        dict=pickle.dump(t,f)
    return dict

# print(labellist)
def get_img(data,label):
    img = np.reshape(data, (3, 32, 32))
    img = img.transpose((1, 2, 0))
    plt.imshow(img)
    if(label==0):
        plt.title('cat' )
    else:
        plt.title('dog')
    plt.show()
def get_catdog():
    train=[]
    cat_id=3
    dog_id=5
    for i in range(5):
        # train_path=os.path.join(data_dir,train_pathname+str(i+1))
        t=unpickle(train_path+str(i + 1)+'_catdog')
        # print(t[b'labels'])
        t[b'labels']=np.array(t[b'labels'])
        index=np.where((t[b'labels']==cat_id) | (t[b'labels']==dog_id))
        # print(index)
        # t[b'labels'][(t[b'labels']==cat_id) | (t[b'labels']==dog_id)])
        t[b'labels']=t[b'labels'][index]
        t[b'labels'][t[b'labels'] == cat_id] = 0
        t[b'labels'][t[b'labels'] == dog_id] = 1
        t[b'labels']= t[b'labels'].tolist()
        t[b'data']=t[b'data'][index]
        # print(t[b'labels'].shape, t[b'data'].shape)
        # t[b'labels']=[1 if x==cat_id else 2 if x==dog_id else 0 for x in t[b'labels']]
        # print(t[b'labels'])
        # print(np.where(t[b'labels']>3 ))
        new_t = {}
        new_t['labels'] = t[b'labels']
        new_t['data'] = t[b'data']
        pickle_save(train_path+'_catdog',new_t)
        train.append(t)
    meta=unpickle(meta_path)
    test=unpickle(test_path)
    # print(t)
    # print(test)
    # labellist=meta[b'label_names']
    # get_img(test['data'].reshape(-1,3,32,32)[3],test['labels'][3])
    test[b'labels']=np.array(test[b'labels'])
    index=np.where((test[b'labels']==cat_id) | (test[b'labels']==dog_id))
    test[b'labels']=test[b'labels'][index]
    test[b'labels'][test[b'labels']==cat_id]=0
    test[b'labels'][test[b'labels']==dog_id]=1
    test[b'labels']= test[b'labels'].tolist()
    test[b'data']=test[b'data'][index]

    new_test={}
    new_test['labels']=test[b'labels']
    new_test['data']=test[b'data']
    pickle_save(test_path+'_catdog',new_test)
    print(len(test[b'labels']))
# print(test)

# print(test[b'labels'].shape,test[b'data'].shape)
# print(train[0][b'data'].shape)
# print(type(train[0][b'labels']))
# get_img(test[0]['data'].reshape(-1,3,32,32)[3],test[0]['labels'][3])
# print(Data.TensorDataset(torch.from_numpy(i[b'data'].reshape(10000,3,32,32)),torch.tensor(i[b'labels'])))
# train_loaders=[Data.DataLoader(dataset= Data.TensorDataset(torch.from_numpy(i[b'data'].reshape(-1,3,32,32)),torch.FloatTensor(i[b'labels'])),batch_size=BATCH_SIZE,shuffle=True,num_workers=0) for i in train[:-1]]
# val_loader=Data.DataLoader(dataset= Data.TensorDataset(torch.from_numpy(train[-1][b'data'].reshape(-1,3,32,32)),torch.FloatTensor(train[-1][b'labels'])),batch_size=train[-1][b'data'].shape[0],shuffle=True,num_workers=0)
# test_loader=Data.DataLoader(dataset=Data.TensorDataset(torch.from_numpy(test[b'data'].reshape(-1,3,32,32)),torch.FloatTensor(test[b'labels'])),batch_size=test[b'data'].shape[0],shuffle=True,num_workers=0)
def predict(model):
    with torch.no_grad():
        accuracy = 0.0
        l = 0.0
        test = unpickle(test_path + '_catdog')
        test_loader = Data.DataLoader(
            dataset=Data.TensorDataset(torch.from_numpy(test['data'].reshape(-1, 3, 32, 32)),
                                       torch.FloatTensor(test['labels'])), batch_size=test['data'].shape[0], shuffle=False,
            num_workers=0)
        for x, y in test_loader:
            test_output = model(x)
            # p_y = torch.argmax(test_output, -1)
            test_output = test_output.squeeze(-1)
            test_output[test_output > 0] = 1
            test_output[test_output <= 0] = 0
            p_y=test_output
            # print(pred_y.shape)
            accuracy = accuracy + float((p_y == y).sum())
            l += float(y.size(0))
        accuracy /= l
        print("test accuracy：%.4f" % accuracy)

# def get_label_data(train,label):
#     data={"data":[],"label":[]}
#     for t in train:
#         for img,l in zip(t[b'data'],t[b'labels']):
#             if(labellist[l]==label):
#
# # get_img(train[-1][b'data'][0],train[-1][b'labels'][0])
#
# get_label_data(train, b'cat')
class cifar10_classify(nn.Module):
    def __init__(self,config):
        super(cifar10_classify,self).__init__()

        def pool(x):
            if (x == 'M'):
                return nn.MaxPool2d(kernel_size=2)
            elif (x == 'A'):
                return nn.AvgPool2d(kernel_size=2)
            else:
                return nn.MaxPool2d(kernel_size=config.dim*2)

        def activate(x):
            if x == 'relu':
                return nn.ReLU()
            elif x == 'tanh':
                return nn.Tanh()

        self.convs2 = \
            nn.ModuleList([
                nn.ModuleList([
                    nn.Sequential(
                        nn.Conv2d(
                            in_channels=config.conv_struct[i - 1][-1] * len(config.kernel_size[i - 1]),
                            out_channels=config.conv_struct[i][-1],
                            kernel_size=k,
                            stride=1,
                            padding=(k - 1) // 2  # padding=(kernel_size-stride)/2
                        ),
                        activate(config.conv_struct[i][1]),
                        nn.BatchNorm2d(config.conv_struct[i][-1]),
                        pool(config.conv_struct[i][0])
                    )  # 2*2的窗口
                    for k in config.kernel_size[i]
                ])
                for i in range(1,len(config.conv_struct))
            ] )#输出为[batch,32，8,8]
        self.dropout=nn.ModuleList([nn.Dropout(i) for i in config.conv_dropout])
        # print(config.fc)
        self.fc=nn.ModuleList(
            [
                    nn.Sequential(
                        nn.Linear(config.fc[i-1],config.fc[i]),
                        activate(config.fc_act[i]),
                        nn.Dropout(config.fc_dropout[i]))
                    if i+2<len(config.fc)
                    else
                    nn.Linear(config.fc[i - 1], config.fc[i])
                for i in range(1,len(config.fc))
            ]
        )

    def forward(self,x):
        # print(x.shape)
        # print(x)

        i=0
        x=x.float()/255.
        for i,modulse in enumerate(self.convs2):
            x=[conv(x) for conv in modulse]
            # print(x.shape)
            # xs=[]
            # for conv in modulse:
            #     # print(conv)
            #     xs +=[conv(x) ]
                # print(xs[-1].shape)

            x=torch.cat(x, 1)
            x = self.dropout[i](x)

        x = x.view(x.size(0), -1)
        # print(x.shape)
        for f in self.fc:
            x=f(x)
        # x=self.dropout1(x)
        # x=self.bn1(x)
        # print(x.shape)
        # x = self.conv1(x.float()/255.)
        # x=self.conv2(x)
        # x = self.dropout2(x)
        # x = self.conv3(x)
        # x = self.conv4(x)
        # x = self.dropout3(x)

        # print(x.shape)
        # x=self.fc1(x)
        # output=self.fc2(x)
        # output=self.prediction(x)
        return x

def one_fold(opt):
    opt.time=str(now)+'_'+str(datetime.now())
    print("---------------------------------------------------------------------------------------------------------------------------")
    print(opt.string())
    print("---------------------------------------------------------------------------------------------------------------------------")
    BATCH_SIZE=opt.BATCH_SIZE
    model=cifar10_classify(opt)
    #大数据常用Adam优化器，参数需要model的参数，以及学习率
    optimizer = torch.optim.Adam(model.parameters(), opt.LR)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=opt.gamma)
    #定义损失函数，交叉熵
    loss_func = nn.BCEWithLogitsLoss()
    val_id=opt.val_id
    max_val_acc=0.
    max_train_acc=0.
    early_stopping=opt.early_stopping
    # print(opt.EPOCH)
    for epoch in range(opt.EPOCH):
        scheduler.step()
        i=0
        for i in range(0,5):
            if(i+1==val_id):
                continue
            train = unpickle(train_path+str(i + 1)+'_catdog')
            train_loader=Data.DataLoader(dataset= Data.TensorDataset(torch.from_numpy(train['data'].reshape(-1,3,32,32)),torch.FloatTensor(train['labels'])),batch_size=BATCH_SIZE,shuffle=False,num_workers=0)
            i+=1
            total_loss=0.0
            train_accuracy=0.0
            step=0
            total_mean_F1=0.0
            for step,(batch_x,batch_y) in enumerate(train_loader):

                output = model(batch_x)
                # print(output)

                output=output.squeeze(-1)
                # print(output.shape)
                # print(batch_y.shape)
                loss = loss_func(output, batch_y)
                total_loss+=loss.data
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                # print(output)
                output[output>0]=1
                output[output <= 0] = 0
                pred_y=output
                # print(pred_y.shape)
                # print(pred_y)
                train_accuracy=train_accuracy+ float((pred_y == batch_y).sum()) / float(batch_y.size(0))
                F1, mean_F1 = metric(batch_y.long(), pred_y.long())
                # print(F1,mean_F1)
                total_mean_F1= total_mean_F1+mean_F1
            with torch.no_grad():
                accuracy = 0.0
                l = 0.0
                val_loss=0.0
                # val_path = os.path.join(data_dir, train_pathname + str(5) + '_catdog')
                val = unpickle(train_path+str(val_id)+'_catdog')
                val_loader = Data.DataLoader(
                    dataset=Data.TensorDataset(torch.from_numpy(val['data'].reshape(-1, 3, 32, 32)),
                                               torch.FloatTensor(val['labels'])), batch_size=val['data'].shape[0], shuffle=False,
                    num_workers=0)

                for x, y in val_loader:
                    val_output = model(x)
                    val_output = val_output.squeeze(-1)
                    val_output[val_output > 0] = 1
                    val_output[val_output <= 0] = 0
                    p_y = val_output

                    # print(pred_y.shape)
                    t_loss = loss_func(val_output, y)
                    val_loss=val_loss+t_loss.data
                    accuracy = accuracy + float((p_y == y).sum())
                    F1,mean_F1=metric(y.long(), p_y.long())
                    l += float(y.size(0))
                accuracy /= l
                # print("epoch:", epoch, "| step:", step, "|val accuracy：%.4f" % accuracy)
                train_accuracy/=step
                total_mean_F1/=step
                if(train_accuracy>max_train_acc):
                    max_train_acc=train_accuracy
                if(accuracy>max_val_acc):
                    max_val_acc=accuracy
                    early_stopping = opt.early_stopping
                else:
                    early_stopping-=1
                print("| epoch:", epoch,  "| filename:",i,"| train loss:%.4f" % total_loss,
                      "|train accuracy：%.4f" % train_accuracy,"|train mean_F1：%.4f" % total_mean_F1.data, "|max val accuracy：%.4f" %  max_val_acc,"| val loss:%.4f" % val_loss, "|val accuracy：%.4f" % accuracy,"|val F1：" , F1.data,"|val mean_F1：%.4f" % mean_F1.data," |")
        if(early_stopping<=0):
            print("early_stopping")
            break
    return max_val_acc,max_train_acc
def cross_validate(opt):
    mean_acc=0.
    for i in range(0,5):
        opt.val_id=i+1
        max_acc,train_acc=one_fold(opt)
        mean_acc=mean_acc+max_acc
        opt.write_string(log_path,train_acc,max_acc,mean_acc/(i+1.))
    return mean_acc

def fc_grid(opt1):
    global now
    for i in range(1, 3):  # fc
        ll = []
        # print(opt1.string())
        opt1.change_conv()
        l = (opt1.fc[0] - opt1.fc[-1]) // i
        # print(l)
        # print(opt1.fc[0],opt1.fc[1])
        for j in range(i):
            ll.append(opt1.fc[0] - j * l)
            # print(ll)
        ll.append(opt1.fc[-1])
        opt1.fc = ll
        if(now<13):#4
            now += 1
            continue
        cross_validate(opt1)
        now += 1

def main():
    global now
    now = 38
    opt = config()
    for lr in [0.01,0.001,1e-4,1e-5]:
        opt.LR=lr
        for dropout in [0,0.2,0.5]:
            opt.conv_dropout=[dropout]*4
            for gamma in [0,0.5,0.9,1]:
                opt.gamma=gamma

    cross_validate(opt)
    # opt=config()
    # # with open(log_path, 'w', encoding='utf-8') as f:
    # #     f.write(opt.get_title())
    # kernel = [3, 5, 7]
    # pool=['M','A','G']
    #
    # # cross_validate(opt)
    # print('###########################################################################################')
    # print('实验1：前馈层的层数+多尺寸+卷积层数+不同类型的池化')
    # print('###########################################################################################')
    #
    # # opt1.fc=[opt.fc[0],opt.fc[-1]]
    # # [feature,1]
    # # cross_validate(opt1)
    # # [feature,64,1]
    # global now
    # now=0
    #
    # for k in range(2, 4):#卷积深度
    #     opt1 = config()
    #     opt1.kernel_size = opt.kernel_size[:k]
    #     opt1.conv_struct = opt.conv_struct[:k]
    #     for z in range(1, 3):  # 第一层卷积的宽度
    #         opt1.kernel_size[1] = kernel[:z]  #
    #         # print(opt1.string())
    #         if (len(opt1.conv_struct) == 3):  # pool
    #             for a in range(2):
    #                 opt1.conv_struct[1][0] = pool[a]
    #                 for b in range(3):
    #                     opt1.conv_struct[2][0] = pool[b]
    #                     fc_grid(opt1)
    #         elif (len(opt1.conv_struct) == 4):
    #             for a in range(2):
    #                 opt1.conv_struct[1][0] = pool[a]
    #                 for b in range(2):
    #                     opt1.conv_struct[2][0] = pool[b]
    #                     for c in range(3):
    #                         opt1.conv_struct[3][0] = pool[c]
    #                         fc_grid(opt1)
    #         else:
    #             for a in range(3):
    #                 opt1.conv_struct[1][0] = pool[a]
    #                 fc_grid(opt1)
    #
    #
    # print('###########################################################################################')
    # print('实验2：前馈层的层数+多尺寸+卷积层数')
    # print('###########################################################################################')
    # opt1 = config()
    # # opt1.fc=[opt.fc[0],opt.fc[-1]]
    # # [feature,1]
    # # cross_validate(opt1)
    # # [feature,64,1]
    # for i in range(1, 3):
    #     ll = []
    #     l = (opt.fc[-1] - opt.fc[0]) // i
    #     for j in range(i):
    #         ll.append(opt.fc[0] + i * l)
    #     ll.append(opt.fc[-1])
    #     opt1.fc = ll
    #     for k in range(1, 3):
    #         opt1.kernel_size = opt.kernel_size[:k]
    #         opt1.conv_struct = opt.conv_struct[:k]
    #         for z in range(1,3):
    #             opt1.kernel_size[0] = kernel[:z]
    #             cross_validate(opt1)
    #
    # print('###########################################################################################')
    # print('实验3：前馈层的层数+多尺寸+卷积层数')
    # print('###########################################################################################')
    # opt1 = config()
    # # opt1.fc=[opt.fc[0],opt.fc[-1]]
    # # [feature,1]
    # # cross_validate(opt1)
    # # [feature,64,1]
    # for i in range(1, 3):
    #     ll = []
    #     l = (opt.fc[-1] - opt.fc[0]) // i
    #     for j in range(i):
    #         ll.append(opt.fc[0] + j * l)
    #     ll.append(opt.fc[-1])
    #     opt1.fc = ll
    #     for k in range(1, 3):
    #         opt1.kernel_size[0] = kernel[:k]
    #         cross_validate(opt1)
    # print('###########################################################################################')
    # print('实验4：前馈层的层数+卷积层数')
    # print('###########################################################################################')
    # opt1 = config()
    # # opt1.fc=[opt.fc[0],opt.fc[-1]]
    # # [feature,1]
    # # cross_validate(opt1)
    # # [feature,64,1]
    # for i in range(1, 3):
    #     ll = []
    #     l = (opt.fc[-1] - opt.fc[0]) // i
    #     for j in range(i):
    #         ll.append(opt.fc[0] + j * l)
    #     ll.append(opt.fc[-1])
    #     opt1.fc = ll
    #     for k in range(1, 3):
    #         opt1.kernel_size = opt.kernel_size[:k]
    #         opt1.conv_struct = opt.conv_struct[:k]
    #         cross_validate(opt1)
    #     # cross_validate(opt1)
    # print('###########################################################################################')
    # print('实验5：前馈层的层数')
    # print('###########################################################################################')
    # opt1=config()
    # # opt1.fc=[opt.fc[0],opt.fc[-1]]
    # # [feature,1]
    # # cross_validate(opt1)
    # #[feature,64,1]
    # for i in range(1,3):
    #     ll=[]
    #     l=(opt.fc[-1]-opt.fc[0])//i
    #     for j in range(i):
    #         ll.append(opt.fc[0]+i*l)
    #     ll.append(opt.fc[-1])
    #     opt1.fc=ll
    #     cross_validate(opt1)
    # print('###########################################################################################')
    # print('实验6：卷积核层+深')
    # print('###########################################################################################')
    # # kernel = [3, 5, 7]
    # for i in range(1, 3):
    #     for j in range(1,3):
    #         opt1 = config()
    #         opt1.kernel_size = opt.kernel_size[:i]
    #         opt1.conv_struct = opt.conv_struct[:i]
    #         opt1.kernel_size[0] = kernel[:i]
    #         cross_validate(opt1)
    # ##第一个实验：卷积层数的影响
    # print('###########################################################################################')
    # print('实验7：卷积层数的影响')
    # print('###########################################################################################')
    # for i in range(1,3):
    #     opt1 = config()
    #     opt1.kernel_size=opt.kernel_size[:i]
    #     opt1.conv_struct=opt.conv_struct[:i]
    #     cross_validate(opt1)
    # print('###########################################################################################')
    # print('实验8：多尺寸卷积核的影响')
    # print('###########################################################################################')
    # for i in range(1,3):
    #     opt1 = config()
    #     opt1.kernel_size[0]=kernel[:i]
    #     cross_validate(opt1)
    # print('###########################################################################################')
    # print('实验9：不同激活层的影响')
    # print('###########################################################################################')
    # for i in range(1, 3):
    #     opt1 = config()
    #     opt1.kernel_size[0] = kernel[:i]
    #     cross_validate(opt1)

if __name__=='__main__':

    main()


##############################################
#
# Solving Re-ionisation
# By DS 2017, updated by DS 2018
#
##############################################
from pylab import *
from scipy.integrate import odeint
import matplotlib
#matplotlib.rc('text', usetex=True)   # Commented out for iMAC
from math import *
from scipy.interpolate import interp1d
import pyfits as py
import pyfits
from scipy.integrate import quadpack
from scipy.integrate import quadrature
from scipy import special
import os

# Quick function to open fits catalogue and deliver it back
def fits_open(file_in):
    Open_C=pyfits.open(file_in)
    Fits_catalogue=Open_C[1].data
    return Fits_catalogue

#### LOAD SC4K RESULTS
def Load_SC4K(sample='SSC4K'):

    CAT='../Santos2017_SC4K/Analysis_Global_Lya_LF/MCMC_Results_11122017_vfinal2.fits'

    # Open the catalogue
    data=fits_open(CAT)
    # FIELDS with relevant information
    SUB_SET = data.field('Sub_set_condition') # Name of sub-set
    Alpha = data.field('Alpha') # Name of sub-set
    Upper_lum_fit = data.field('Fit_upp_luminosity')
    ALPHA_fixed = (Alpha==-1.8)

    #LOAD ALL info
    L_dens, Ldens_up, Ldens_down, SFRD_, SFRD_up, SFRD_down, Ldens_PL, Ldens_PL_up, Ldens_PL_down = data.field('Ldens'), data.field('Ldens_error_up'), data.field('Ldens_error_down'), data.field('Ldens'), data.field('Ldens_error_up'), data.field('Ldens_error_down'), data.field('Ldens_PL2'), data.field('Ldens_error_up_PL2'), data.field('Ldens_error_down_PL2')
    # SETUP ERRORS
    ####
    Ldens_up_log   =  np.log10(L_dens+Ldens_up)-np.log10(L_dens)
    Ldens_down_log = np.log10(L_dens)-np.log10(L_dens-Ldens_down)
    #
    Ldens_PL_up_log =  np.log10(Ldens_PL+Ldens_PL_up)-np.log10(Ldens_PL)
    Ldens_PL_down_log =  np.log10(Ldens_PL)-np.log10(Ldens_PL-Ldens_PL_down)

    # GET REDSHIFTS
    Redshift_c, Delta_redshift = (data.field('Redshift_min')+data.field('Redshift_max'))/2.,(data.field('Redshift_max')-data.field('Redshift_min'))/2.
    # Copy to each filter
    MASK_each_filter = Alpha>100000.0
    MASK_each_fil_DRAKE= Alpha>100000.0
    # Now get the conditions properly:
    FILTERS=['NB392','IA427','IA464','IA484','IA505','IA527','IA574','IA624','IA679','IA709','IA738','IA767','IA827']
    INTEGRAL_con = [43.3,43.3,43.3,43.3,43.3,43.3,44.5,44.5,44.5,44.5,44.5,44.5,44.5]
    #
    for i in range(len(FILTERS)):
        MASK_each_filter = MASK_each_filter+ ((SUB_SET == '%s'%FILTERS[i])*ALPHA_fixed*(Upper_lum_fit==INTEGRAL_con[i]))
        MASK_each_fil_DRAKE = MASK_each_fil_DRAKE+ ((SUB_SET == '%s_with_All_lit'%FILTERS[i])*(Upper_lum_fit==INTEGRAL_con[i]))

    # NOW for redshift bins
    FILTERS = ['z22_ALL','z25_ALL', 'z31_ALL', 'z39_ALL', 'z47_ALL', 'z54_ALL']
    INTEGRAL_con = [43.3,43.3,43.3,44.5,44.5,44.5]
    MASK_stacks = Alpha>100000.0
    MASK_stacks_DRAKE= Alpha>100000.0
    #
    for i in range(len(FILTERS)):
        MASK_stacks = MASK_stacks+ ((SUB_SET == '%s'%FILTERS[i])*ALPHA_fixed*(Upper_lum_fit==INTEGRAL_con[i]))
        MASK_stacks_DRAKE = MASK_stacks_DRAKE+ ((SUB_SET == '%s_with_All_lit'%FILTERS[i])*(Upper_lum_fit==INTEGRAL_con[i]))

    # Now for full stack
    # NOW for redshift bins
    FILTERS = ['FULL_SC4K','FULL_SC4K']
    INTEGRAL_con = [43.3,43.5]
    FULL_SC4K = Alpha>100000.0
    FULL_SSC4K = Alpha>100000.0
    #
    for i in range(len(FILTERS)):
        FULL_SC4K = FULL_SC4K+ ((SUB_SET == '%s'%FILTERS[i])*ALPHA_fixed*(Upper_lum_fit==INTEGRAL_con[i]))
        FULL_SSC4K = FULL_SSC4K+ ((SUB_SET == '%s_with_All_lit'%FILTERS[i])*(Upper_lum_fit==INTEGRAL_con[i]))

    if sample=='SSC4K':
        return Redshift_c[MASK_each_fil_DRAKE],np.log10(L_dens[MASK_each_fil_DRAKE]), Ldens_down_log[MASK_each_fil_DRAKE],Ldens_up_log[MASK_each_fil_DRAKE]

############################################
# Return condicence levels
def plot_confidence_Lya(x_in,ysample,colour='g'):
    lower = percentile(ysample,16,axis=0)  # Equivalent to 1 sigma
    upper = percentile(ysample,84,axis=0) # Equivalent to 1 sigma
    lower2 = percentile(ysample,100-97.72,axis=0)  #, Equivalent to 2 sigma
    upper2 = percentile(ysample,97.72,axis=0) # Equivalent to 2 sigma
    lower3 = percentile(ysample,0.15,axis=0)  #, Equivalent to 3 sigma
    upper3 = percentile(ysample,99.85,axis=0) # Equivalent to 3 sigma
    return x_in,lower,upper, lower2, upper2, lower3, upper3

####################################
# Gaussian Function
def gaussian(x,x0,sigma):
    return np.exp(-(x-x0)**2/(2*sigma**2))

####################################
# Double gaussian, asymmetric
# phi_dummy generates the space of values to pick from
# phi_probs generates the PDF for a 1sigma error around each phi
def double_normal(phi,err_down,err_up,size):
   phi_dummy = np.linspace(phi-5.*err_down,phi+5.*err_up,100000) #generate list from -5 to +5 error
   phi_probs = np.append(gaussian(phi_dummy[phi_dummy<phi],phi,err_down),gaussian(phi_dummy[phi_dummy>=phi],phi,err_up))
   return np.random.choice(phi_dummy,size,p=phi_probs/np.sum(phi_probs))

##################################
def dQ_dt (Q,t): #HII_fraction

    z = t_to_z(t)
    trec = 1./(C*alpha_B*n_e*(1.+z)**(3.)) # s
    nion = f_esc0*xi_ion*p(z)/ (2.938e+73) # (2.938e+73) converts from Mpc^-3 to cm^-3  - full units s^-1 Mpc^-3
    return ((nion)/n_H-Q/trec)*3.1536e+16 # Conversion to time in seconds

##################################
def dQ_dt_Lya (Q,t): #HII_fraction

    z = t_to_z(t)
    trec = 1./(C*alpha_B*n_e*(1.+z)**(3.)) # s
    cHa_K98 = 1.36*pow(10.,-12.)

    #Q_ionLya =  LLya/(cHa_K98*(1-f_esc0)(0.042*EW0)) # Equation 5 Sobral & Matthee 2018
    nion = f_esc0*(pLya(z)/(cHa_K98*(1.-f_esc0)*(0.042*EW(z))))/ (2.938e+73) # From Sobral & Matthee 2018
    #nion = f_esc0*xi_ion*p(z)/ (2.938e+73) # (2.938e+73) converts from Mpc^-3 to cm^-3  - full units s^-1 Mpc^-3
    return ((nion)/n_H-Q/trec)*3.1536e+16 # Conversion to time in seconds

##################################
def dQ_dt_direct (Q,t): #HII_fraction

    z = t_to_z(t)
    trec = 1./(C*alpha_B*n_e*(1.+z)**(3.)) # s
    cHa_K98 = 1.36*pow(10.,-12.)

    nion = pLyC(z)/ (2.938e+73) # (2.938e+73) converts from Mpc^-3 to cm^-3  - full units s^-1 Mpc^-3
    return ((nion)/n_H-Q/trec)*3.1536e+16 # Conversion to time in seconds

##################################
def tau (z):

    return sigma_T*n_e*(1.+z)*(1.+z)*Q(z)*c/(H)/z   # Units:
    # sigma_T  = cm^2
    # c km/s
    # H km/s/Mpc => now km/s/cm
    # Q(z) - unitless
    #n_e per cm^3
    # cm^2 * cm  / cm^3 = unitless

def compute_tau_z(Q,zmin=0.6,zmax=14.,dz=0.3):
    z_integrate = arange(zmin,zmax,dz)
    taus_plot = []
    for i in z_integrate:
        taus_plot.append(quadpack.quad(tau,0.3,i)[0])

    return z_integrate, taus_plot

###########################
# Main properties/constants
#####################################################################################
Omega_b = 0.04 #baryon_density
h = 0.7 # hubble constant / 100
X_p = 0.75 #hydrogen_mass_fraction
T = 20000. #IGM_temperature [K]
C = 3.  #clumping_factor
Y_p = (1.-X_p) # Helium mass
alpha_B = 2.6*pow(10,-13.)*(T/10.**(4.))**(-0.76)  #case_B_recombination_coefficient
n_H = 1.67*pow(10,-7.) * ((Omega_b * h* h)/0.02) * (X_p/0.75) # per cm^3
c = 299792.458  # Speed of light km/s

#### For getting Tau (z)
sigma_T = 6.65*pow(10,-25) # Thomson scattering cross section   cm^2  - from Mayer et al.
n_e = (1.+Y_p/(4.*X_p))*n_H   # comoving number density of electrons (assuming singly ionized He)
H = h*100. *(3.24077929*pow(10,-25.)) # Hubble parameter == hubble constant conversion per cm from per Mpc

#######################################################################################

##################################################################
# GET redshift - time conversion
######
# F737 Cosmology
F737 = 'Redshift_time_F737_cosmology.fits'
# PLanck cosmology
FPlanck = 'Redshift_time_FPlanck_cosmology.fits'

# Which Cosmology to use:
USING_NOW = FPlanck

#LOAD Redshift-time dependence
A=py.open(USING_NOW)
NOW=A[1].data
#LOAD GENERAL PROPERTIES - what will be plotted
Redshift = NOW.field('Redshift')
Time =  NOW.field('Time_Gyr')
# time to redshift relation
t_to_z = interp1d(Time,Redshift,kind='linear',bounds_error=None)

############################################################
# Integration limits and conversion z to time
#
ts = np.linspace(0.25,13.4,60000)  # Produce time grid  - 50000 steps seems to be a good minimum for convergence
zs= t_to_z(ts)
t0 = 0.0
###############################################################

#################################

#### START PLOT

# DEFINE the sub-plots to make them side by side and share Y axis
f, (ax1,ax2) = plt.subplots(2,1, sharex=True, figsize=(9,9))
plt.subplots_adjust(wspace = -0.04)  # adjust spacing between sub-figures
# setting up the font style throughout
matplotlib.rcParams.update({'font.size': 18, 'font.weight': 'bold'}) # change for diff style

###### QHII PLOT NOW ######
########################################
# Plot end of re-ionisation
za = [6.0,6.0000001]
zp = [-0.5,5.0]
Reionisation = plot(za,zp,'k--')

## PLOT REIONISATION EPOCH
za = [8.8,8.80001]
zp = [-0.5,5.0]
Reionisation = plot(za,zp,linestyle='-',color ='0.40', linewidth=8, zorder=-8)
Reionisation = ax1.plot(za,zp,linestyle='-',color ='0.60', linewidth=8, zorder=-18)

# Extra integration limit
extra_int = 0.45

########################################
# UV luminosity densities
B=py.open('RHO_UV.fits')
NOWB=B[1].data
Red_z = NOWB.field('Redshift')
rho_UV =  10.**(NOWB.field('p_UV_12')) # 'p_UV_12  p_UV_13', 'p_UV_17' - other integration limits
rho_UV_log = (NOWB.field('p_UV_12')) # 'p_UV_12  p_UV_13', 'p_UV_17' - other integration limits
p = interp1d(Red_z,rho_UV,kind='linear',bounds_error=None) # p is rho UV(z)

# LOAD Lyman-alpha luminosity density, from S-SC4K
z_Lya, logLumLya, err_downLya, err_upLya=Load_SC4K(sample='SSC4K')
# Convert to list so can be appended
z_Lya, logLumLya, err_downLya, err_upLya= list(z_Lya), list(logLumLya+extra_int), list(err_downLya), list(err_upLya)

##############################################################
# Add z=6.6, z=7.0, z=7.3 from Matthee+, Zheng+, Konno+
z_others = [6.6,7.0,7.3]
pLya_others = [39.4,38.92,38.55]
err_others = [0.3,0.4,0.3]
# Best linear fit
Bfit = [-0.21, 40.9+extra_int]
# Best linear fit is given by: A=-1.21351351,  B=47.41081081   in log10
# Add higher redshift data - underestimation likely due to reionisaiton effects
"""
for i in range(len(z_others)):
    z_Lya.append(z_others[i])
    logLumLya.append(pLya_others[i])
    err_downLya.append(err_others[i])
    err_upLya.append(err_others[i])
"""

# Linking Lya with UV assuming rho_Lya actually declines as UV - this should be a conservative assumption

# Add higher redshift extrapolated data
for i in [6.0,6.5,7.0,7.5,8.0,8.5,9.0,10.0,10.5,11.0,12.,13.0,14.,16.0,17.0,18.0]:
    z_Lya.append(i)
    logLumLya.append(i*Bfit[0]+Bfit[1])
    err_downLya.append(0.4)
    err_upLya.append(0.3)

z_Lya[0]=0.0

# Convert back to arrays
z_Lya, logLumLya, err_downLya, err_upLya= array(z_Lya), array(logLumLya), array(err_downLya), array(err_upLya)
# Fit higher redshift to extrapolate:

Q_storing, Tau_storing = [], []

iterations = 10000

xi_ions = [25.5,0.2,0.6] # central value, 1 sigma down, 1 sigma up
xi_ion_lya = [25.5,0.2,0.4]
f_esc0_lya = [0.15,0.05,0.05]

EW_0_dist = [80.0,5.,15.]

f_esc0s = [0.15,0.05,0.1]
pUV_uncert = [0.1,0.05]
##### TRY MCMC quick

for i in range(iterations):#### Draw from normal distribution

            rho_lya_new = logLumLya-logLumLya # create new array with zeros
            for j in range(len(logLumLya)): # Go through each one
                rho_lya_new[j]=pow(10.,double_normal(logLumLya[j],err_downLya[j],err_upLya[j],1)[0])

            #print rho_lya_new
            pLya = interp1d(z_Lya,rho_lya_new,kind='linear',bounds_error=None)

            #print pLya(2.22), z_Lya
            # for now assume EW fixed
            EW_pick_now = double_normal(EW_0_dist[0],EW_0_dist[1],EW_0_dist[2],1)[0]
            EW_fixed =rho_lya_new-rho_lya_new+EW_pick_now
            EW = interp1d(z_Lya,EW_fixed,kind='linear',bounds_error=None)

            # Now pick and f_esc LyC
            f_esc0 = double_normal(f_esc0s[0],f_esc0s[1],f_esc0s[2],1)[0]

            # Solve equation
            Q = odeint(dQ_dt_Lya,t0,ts)  # SOLVE NOW

            # SET boundary conditions
            Q[Q>1.0] = 1.0 # boundary condition up
            Q[Q<0.0] = 0.0 # boundary condition down

            # Make it an actual function - rder values
            Q_not = []
            z_not = []
            zs = array(zs)
            for i in range(len(Q)):
                Q_not.append(Q[len(Q)-i-1][0])
                z_not.append(zs[len(zs)-i-1])

            # now interpolate
            Q = interp1d(array(z_not),array(Q_not),kind='linear') # p is rho UV(z)
            Q_storing.append(Q(zs))

            # NOW TAU:
            z_plot, taus_plot = compute_tau_z(Q)

            Tau_storing.append(taus_plot)

            #plot(z_not,Q(z_not),'g-',linewidth=0.4,zorder=-1)

# FIND AND PLOT COUNTOURS FOR Fraction of Ionised Hydrogen
xi,lower,upper, lower2, upper2, lower3, upper3 = plot_confidence_Lya(zs,Q_storing,colour='#ff8080')
SHADE = fill_between(xi,lower,upper,lw=1,color='#0066ff',alpha=0.4,zorder = 90)

#SHADE = fill_between(xi,lower2,upper2,lw=1,color='#e6e6ff',alpha=0.8,zorder = 80)
#SHADE = fill_between(xi,lower3,upper3,lw=1,color='#cce6ff',alpha=0.4,zorder = -100)

# FIND AND PLOT CONTOUNRS FOR Tau - CMB optical depth to compare with Planck

xi,lower,upper, lower2, upper2, lower3, upper3 = plot_confidence_Lya(z_plot,Tau_storing,colour='#ff8080')
SHADE = ax1.fill_between(xi,lower,upper,lw=1,color='#0066ff',alpha=0.4,zorder = 90)

LABEL_contours = plot([-20,-10],[-20,-10],lw=8,alpha=0.4,color='#0066ff')

### FOR UV LBG formalism
"""
rho_UV_new = rho_UV_log-rho_UV_log # create array with zeros
for j in range(len(rho_UV_log)): # Go through each one
    rho_UV_new[j]=pow(10.,double_normal(rho_UV_log[j],pUV_uncert[0],pUV_uncert[1],1)[0])
p = interp1d(Red_z,rho_UV_new,kind='linear',bounds_error=None)

# Pick an x_ion - in log10
xi_ion=pow(10.,double_normal(xi_ions[0],xi_ions[1],xi_ions[2],1)[0]) ###ASYMETRIC
"""

p = interp1d(Red_z,rho_UV,kind='linear',bounds_error=None)

######

# Solve equation
Q = odeint(dQ_dt_Lya,t0,ts)  # SOLVE NOW

# SET boundary conditions
Q[Q>1.0] = 1.0 # boundary condition up
Q[Q<0.0] = 0.0 # boundary condition down

# Make it an actual function - rder values
Q_not = []
z_not = []
zs = array(zs)
for i in range(len(Q)):
    Q_not.append(Q[len(Q)-i-1][0])
    z_not.append(zs[len(zs)-i-1])

# now interpolate
Q = interp1d(array(z_not),array(Q_not),kind='linear') # p is rho UV(z)

########################################
# UV luminosity densities
B=py.open('RHO_UV.fits')
NOWB=B[1].data
Red_z = NOWB.field('Redshift')
rho_UV =  10.**(NOWB.field('p_UV_13')) # 'p_UV_12  p_UV_13', 'p_UV_17' - other integration limits
rho_UV_log = (NOWB.field('p_UV_13')) # 'p_UV_12  p_UV_13', 'p_UV_17' - other integration limits
p = interp1d(Red_z,rho_UV,kind='linear',bounds_error=None) # p is rho UV(z)

#################################################################
# ALL - Robertson+ like
xi_ion = 10.**(25.2) #production_efficiency [Hz erg**(-1.)]
f_esc0 = 0.20

Q = odeint(dQ_dt,t0,ts)  # SOLVE NOW
Q_LAE=Q
Q[Q>1.0] = 1.0 # boundary condition up
Q[Q<0.0] = 0.0 # boundary condition down

# Make it an actual function - rder values
Q_not = []
z_not = []
zs = array(zs)
for i in range(len(Q)):
    Q_not.append(Q[len(Q)-i-1][0])
    z_not.append(zs[len(zs)-i-1])

# now interpolate
Q = interp1d(array(z_not),array(Q_not),kind='linear') # p is rho UV(z)

##################################################################################
### NOW TAU PLOT

# SOLVE TAU:
z_plot, taus_plot = compute_tau_z(Q)

TOTAL = plot(z_not,Q(z_not),'g-',linewidth=3,zorder=-1)
Prediction = ax1.plot(z_plot, taus_plot,'g-',linewidth=3,zorder=-1)

#################################################################
# LAEs
xi_ion = 10.**(25.6) #production_efficiency [Hz erg**(-1.)]
f_esc0 = 0.10

Q = odeint(dQ_dt,t0,ts)  # SOLVE NOW
Q_LAE=Q
Q[Q>1.0] = 1.0 # boundary condition up
Q[Q<0.0] = 0.0 # boundary condition down

# Make it an actual function - rder values
Q_not = []
z_not = []
zs = array(zs)
for i in range(len(Q)):
    Q_not.append(Q[len(Q)-i-1][0])
    z_not.append(zs[len(zs)-i-1])

# now interpolate
Q = interp1d(array(z_not),array(Q_not),kind='linear') # p is rho UV(z)
LAEs = plot(z_not,Q(z_not),'b-.',linewidth=6,zorder=-1)

##################################################################################
### NOW TAU PLOT

# SOLVE TAU:
z_plot, taus_plot = compute_tau_z(Q)
Prediction = ax1.plot(z_plot, taus_plot,'b-.',linewidth=6,zorder=-1)

#################################################################
# LBGs z~3
xi_ion = 10.**(25.3) #production_efficiency [Hz erg**(-1.)]
f_esc0 = 0.03

Q = odeint(dQ_dt,t0,ts)  # SOLVE NOW
Q_LAE=Q
Q[Q>1.0] = 1.0 # boundary condition up
Q[Q<0.0] = 0.0 # boundary condition down

# Make it an actual function - rder values
Q_not = []
z_not = []
zs = array(zs)
for i in range(len(Q)):
    Q_not.append(Q[len(Q)-i-1][0])
    z_not.append(zs[len(zs)-i-1])

# now interpolate
Q = interp1d(array(z_not),array(Q_not),kind='linear') # p is rho UV(z)
LBGs_real = plot(z_not,Q(z_not),'r--',linewidth=3,zorder=-1)

##################################################################################
### NOW TAU PLOT

# SOLVE TAU:
z_plot, taus_plot = compute_tau_z(Q)
Prediction = ax1.plot(z_plot, taus_plot,'r--',linewidth=3,zorder=-1)

#### PLOT Planck
########################################
# Plot Planck optical depth
za = [-2.0,19.0000001]
zp = [0.066,0.06600001]
Reionisation = ax1.plot(za,zp,'k--',linewidth=3)
Reionisation = ax1.plot(za,zp,'k-',linestyle='-',color ='0.90', linewidth=50, zorder=-20) # INCLUDES ERROR
#Reionisation = ax1.plot(za,array(zp)+0.012,'k-',linewidth=3)
#Reionisation = ax1.plot(za,array(zp)-0.012,'k-',linewidth=3)

# LEGEND
I1=legend((LABEL_contours[0],TOTAL[0],LAEs[0],LBGs_real[0]),(r'S-SC4K+SM18 MCMC',r'R+13,15',r'LAEs $z\sim2-3$',r'LBGs $z\sim2-3$'), shadow = True, loc=3,numpoints=1)
ltext = gca().get_legend().get_texts()
setp(ltext[0], fontsize = 19, color='0.1')
setp(ltext[1], fontsize = 19, color='k')
frame=I1.get_frame()
frame.set_linewidth(0)
frame.set_visible(False)

# Define the axis
axis([0,13.1,-0.05,1.05])
ax1.axis([0,13.1,-0.01,0.11])

# Label axis - quantities and units
ax2.set_xlabel(r'Redshift ($\bf z$)', {'color':'k', 'fontsize': 23})
ax2.set_ylabel(r'Ionised H Fraction (\%)', {'color':'k', 'fontsize': 23})
ax1.set_ylabel(r'$\bf \tau (z)$', {'color'    : 'k', 'fontsize'   :24 })

# SET AXIS - labelling
ax2.set_xticks((0,2,4,6,8,10,12),(0,2,4,6,8,10,12))
ax2.set_yticks((0.0,0.2,0.4,0.6,0.8,1.0),(0,20,40,60,80,100))
ax1.set_yticks((0.0,0.02,0.04,0.06,0.08,0.10),(0.0,0.02,0.04,0.06,0.08,0.10))

savefig('Q_HII_and_Tau_evolution.pdf')
show()


import numpy as np
import pandas as pd

import tensorflow as tf
from sklearn.preprocessing import MinMaxScaler

milk = pd.read_csv('monthly-milk-production.csv',index_col='Month')
milk.index = pd.to_datetime(milk.index)

train_set = milk.head(156)
test_set = milk.tail(12)

scaler = MinMaxScaler()
train_scaled = scaler.fit_transform(train_set)
test_scaled = scaler.transform(test_set)

# Just one feature, the time series
num_inputs = 1
# Num of steps in each batch
num_time_steps = 12
# 100 neuron layer, play with this
num_neurons = 100
# Just one output, predicted time series
num_outputs = 1

learning_rate = 0.03
# how many iterations to go through (training steps), you can play with this
num_train_iterations = 1000
# Size of the batch of data
batch_size = 1

X = tf.placeholder(tf.float32, [None, num_time_steps, num_inputs])
y = tf.placeholder(tf.float32, [None, num_time_steps, num_outputs])

# Also play around with GRUCell
cell = tf.contrib.rnn.OutputProjectionWrapper(
    tf.contrib.rnn.BasicLSTMCell(num_units=num_neurons, activation=tf.nn.relu),
    output_size=num_outputs)

outputs, states = tf.nn.dynamic_rnn(cell, X, dtype=tf.float32)

loss = tf.reduce_mean(tf.square(outputs - y)) # MSE
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train = optimizer.minimize(loss)

init = tf.global_variables_initializer()
saver = tf.train.Saver()
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.9)

def next_batch(training_data,batch_size,steps):
    # Grab a random starting point for each batch
    rand_start = np.random.randint(0,len(training_data)-steps)
    # Create Y data for time series in the batches
    y_batch = np.array(training_data[rand_start:rand_start+steps+1]).reshape(1,steps+1)
    return y_batch[:, :-1].reshape(-1, steps, 1), y_batch[:, 1:].reshape(-1, steps, 1)

sess = tf.InteractiveSession()
sess.run(init)

for iteration in range(num_train_iterations):
    X_batch, y_batch = next_batch(train_scaled,batch_size,num_time_steps)
    sess.run(train, feed_dict={X: X_batch, y: y_batch})
    if iteration % 100 == 0:
        mse = loss.eval(feed_dict={X: X_batch, y: y_batch})
        print(iteration, "MSE:", mse)

# Save Model for Later
model_dir = './time_series_model/'
saver.save(sess, model_dir)

serialized_tf_example = tf.placeholder(tf.string, name='x')
feature_configs = {
    'x': tf.FixedLenFeature(
        shape=[num_time_steps, num_inputs], dtype=tf.float32),
    }
tf_example = tf.parse_example(serialized_tf_example, feature_configs)
input_x = tf_example['x']

SavedModel_folder = "SavedModel1"
builder = tf.saved_model.builder.SavedModelBuilder(SavedModel_folder)

predict_tensor_inputs_info = tf.saved_model.utils.build_tensor_info(input_x)
predict_tensor_outputs_info = tf.saved_model.utils.build_tensor_info(outputs)

prediction_signature = (
    tf.saved_model.signature_def_utils.build_signature_def(
        inputs={'inputs': predict_tensor_inputs_info},
        outputs={'outputs': predict_tensor_outputs_info},
        method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME
    )
)

# save the model
legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
builder.add_meta_graph_and_variables(
    sess, [tf.saved_model.tag_constants.SERVING],
    signature_def_map={
        'predict_images': prediction_signature
    },
    legacy_init_op=legacy_init_op)

builder.save()


import pandas as pd
import sys
import os
from tslearn.clustering import KShape, TimeSeriesKMeans
from tslearn.preprocessing import TimeSeriesScalerMeanVariance
from sklearn import metrics

dataset = pd.read_csv(sys.argv[1])
path_output = sys.argv[2]

#take all elements in dataset based on points in encoding process
matrix_data = []

headers = [key for key in dataset.keys() if key[0] == "P"]

for i in range(len(dataset)):
	row = [dataset[key][i] for key in headers]
	matrix_data.append(row)

# For this method to operate properly, prior scaling is required
X_train = TimeSeriesScalerMeanVariance().fit_transform(matrix_data)

matrix_scaler = []

for row in X_train:
	row_values = []
	for element in row:
		row_values.append(element[0])
	matrix_scaler.append(row_values)

matrix_response = []

dataset_categories = pd.DataFrame()
dataset_categories['id_sequence_by_algorithm'] = dataset['id_sequence_by_algorithm']

for k in range(2, 50):

	try:
		print("Process K: ", k)
		# kShape clustering
		ks = KShape(n_clusters=k, verbose=True, random_state=0)
		y_pred = ks.fit_predict(matrix_scaler)

		siluetas = metrics.silhouette_score(matrix_scaler, y_pred, metric='euclidean')
		calinski = metrics.calinski_harabasz_score(matrix_scaler, y_pred)
		davies_index = metrics.davies_bouldin_score(matrix_scaler, y_pred)
		description = "K-"+str(k)

		row_summary = ["KShape_clustering", description, siluetas, calinski, davies_index]

		matrix_response.append(row_summary)
		dataset_categories[description] = y_pred

	except:
		pass

dataset_categories.to_csv(path_output+"categories_exploring.csv", index=False)
dataSummary = pd.DataFrame(matrix_response, columns=["Algorithm", "Hyperparams", "silhouette_score", "calinski_harabasz_score", "davies_bouldin_score"])
dataSummary.to_csv(path_output+"summary_report.csv", index=False)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on  9:49 2019/7/17

@author:    Liu Jinbao
@mail:      liu.jinbao@outlook.com
@project:   PlasmaChemistry
@IDE:       PyCharm
"""

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from plasmistry.electron import EEDF
from plasmistry import constants as const
from plasmistry.electron import get_maxwell_eedf
# from scipy.integrate import ode as ode_ivp
from plasmistry.solvers import ode_ivp
from plasmistry.molecule import get_ideal_gas_density

inelas_df = pd.read_pickle("plasmistry/electron/tests/e.g._inelas_colli_dataframe.pkl")
species = ['CO2', 'CO2(v0)', 'CO2(v1)', 'CO2(vc)', 'CO', 'O2']
# ----------------------------------------------------------------------------------------------- #
eedf = EEDF(max_energy_eV=10,
            grid_number=100)
eedf.initialize(rctn_with_crostn_df=inelas_df,
                total_species=species)

eedf.set_parameters(E=1, Tgas=1000, N=1e20)
total_species_density = get_ideal_gas_density(p_Pa=1e5, Tgas_K=2000)
eedf.set_flux(total_species_density=total_species_density * np.array([1, 0, 0, 0, 0, 0]))

# eedf.set_flux(total_species_density=)

def dndt(t, y, _eedf):
    _eedf.density_in_J = y
    return _eedf.get_deriv_total()

# y_0 = np.zeros_like(eedf.energy_point)
# y_0[(eedf.energy_point < 5.5 * const.eV2J) & (eedf.energy_point > 0.5 * const.eV2J)] = \
#     1e16 / 5.0 / const.eV2J
# eedf.density_in_J = y_0
# density_0 = eedf.electron_density
# Te_0 = eedf.electron_temperature
y_0 = 1e14 * get_maxwell_eedf(eedf.energy_point, Te_eV=0.026)
sol = ode_ivp(deriv_func=dndt,
              func_args=(eedf,),
              time_span=(0, 1e5),
              y_0=y_0,
              rtol=1e-6,
              atol=1e-12,
              show_time=False)
# plt.plot(eedf.energy_point, a.transpose())

# eedf.density_in_J = a[-1]
# eedf.set_density_in_J(sol.y[-1])
plt.plot(eedf.energy_point, (sol.y / np.sqrt(eedf.energy_point)).transpose())


import pytest
import dhash
from PIL import Image
from subprocess import check_output

size_sample = list(range(4, 17))

@pytest.fixture
def test_filename():
    return "dhash-test.jpg"

@pytest.fixture
def pil_image(test_filename):
    img = Image.open(test_filename)
    return img

# TODO
# @pytest.fixture
# def wand_image():
#    ...

@pytest.fixture
def known_hash():
    """The known hex format size 8 of dhash-test.jpg"""
    dhash = "c1c4e4a484a0809083fbffcc0040831e"
    return dhash

@pytest.mark.parametrize("size", size_sample)
def test_get_grays(pil_image, size):
    """get_grays(...) should convert an image to grayscale, downsize, and give integer values"""

    # GIVEN a sample image
    img = pil_image

    # WHEN calling get_grays(...)
    result = dhash.get_grays(img, size, size)

    # THEN return a list of grayscale ints 0-255
    assert isinstance(result, list), "get_grays(...) should return a list"

    for item in result:
        assert isinstance(item, int), "each item in the get_grays(...) list should be an int"
        assert item >= 0 and item <= 255, "each item in the get_grays(...) list should be between 0 and 255"

    expected_len = size * size
    assert len(result) == expected_len, "get_grays(...) should have be the expected_len"

@pytest.mark.parametrize("size", size_sample)
def test_dhash_row_col(pil_image, size):
    """dhash_row_col(...) should convert image to grayscale, downsize, and produce ints which will later be combined"""

    # GIVEN a sample image
    img = pil_image

    # WHEN calling dhash_row_col(...)
    result = dhash.dhash_row_col(img, size=size)

    # THEN return a tuple of two ints
    assert isinstance(result, tuple), "dhash_row_col(...) should return a tuple"
    assert len(result) == 2, "the dhash_row_col(...) tuple should have len of 2"

    for item in result:
        assert isinstance(item, int), "each item in the dhash_row_col(...) tuple should be an int"

@pytest.mark.parametrize("size", size_sample)
def test_dhash_int(pil_image, size):
    """dhash_int(...) should calculate combined hash and output an integer"""

    # GIVEN a sample image
    img = pil_image

    # WHEN calling dhash_int(...)
    result = dhash.dhash_int(img, size=size)

    # THEN return an int
    assert isinstance(result, int), "dhash_int(...) should return an int"

@pytest.mark.parametrize("size", size_sample)
def test_format_grays(pil_image, size):
    """format_grays(...) should a matrix graph of the gray values"""

    # GIVEN grays from a sample image
    img = pil_image
    grays = dhash.get_grays(pil_image, size, size)

    # WHEN calling format_grays(...)
    result = dhash.format_grays(grays, size=size - 1)

    # THEN
    assert isinstance(result, str), "dhash.format_grays(...) should return a str"

    lines = result.split("\n")

    assert len(lines) == size, "the matrix should be the specified size"

    for line in lines:

        values_in_line = list(filter(lambda x: x != "", line.lstrip().rstrip().split(" ")))

        assert len(values_in_line) == size, "all lines (non-whitespace) should be the specified size"

        for value in values_in_line:
            assert int(value) >= 0 and int(value) <= 255, "each value should be between 0 and 255"

@pytest.mark.parametrize("size", list(filter(lambda x: x % 2 == 0, size_sample)))
def test_format_bytes(pil_image, size):
    """format_bytes(...) should return a hash in bytes format"""

    # NOTE: This test fails on odd numbers

    # GIVEN the row and col hash from an image
    img = pil_image
    row_hash, col_hash = dhash.dhash_row_col(img, size=size)

    # WHEN calling format_bytes(...)
    result = dhash.format_bytes(row_hash, col_hash, size=size)

    # THEN return bytes
    assert isinstance(result, bytes), "format_bytes(...) should return bytes"

@pytest.mark.parametrize("size", size_sample)
def test_format_hex(pil_image, size):
    """format_hex(...) should return a hash in hex format"""

    # GIVEN the row and col hash from an image
    img = pil_image
    row_hash, col_hash = dhash.dhash_row_col(img, size=size)

    # WHEN calling format_hex(...)
    result = dhash.format_hex(row_hash, col_hash, size=size)

    # THEN return a hex str
    assert isinstance(result, str), "format_hex(...) should return a str"

    hex_characters = "0123456789abcdef"
    chars_are_hex_chars = list(map(lambda x: x in hex_characters, result))
    assert all(chars_are_hex_chars), "all character should be between 0-9 and a-f"

@pytest.mark.parametrize("size", size_sample)
def test_format_matrix(pil_image, size):
    """format_matrix(...) should a matrix graph of the pixel comparison"""

    # GIVEN the row and col hash from an image
    img = pil_image
    row_hash, col_hash = dhash.dhash_row_col(img, size=size)

    # WHEN calling dhash.format_matrix(...) on the row and col hashes
    bits = [". ", "* "]
    row_result = dhash.format_matrix(row_hash, bits=bits, size=size)
    col_result = dhash.format_matrix(col_hash, bits=bits, size=size)

    # THEN
    for result in [row_result, col_result]:

        assert isinstance(result, str), "dhash.format_matrix(...) should return a str"

        lines = result.split("\n")
        bit_chars = "".join(set("".join(bits)))

        assert len(lines) == size, "the matrix should be the specified size"

        for line in lines:

            assert len(line.replace(" ", "")) == size, "all lines (non-whitespace) should the specified size"

            chars_are_bit_chars = list(map(lambda x: x in bit_chars, line))
            assert all(chars_are_bit_chars), "all chars should be bit chars"

@pytest.mark.parametrize("hash1,hash2,expected", [(0, 0, 0), (0, 1, 1), (1, 0, 1), (1, 1, 0)])
def test_get_num_bits_different(hash1, hash2, expected):
    """get_num_bits_different(...) should perform binary XOR and count 1s"""

    # GIVEN mock hashes
    # WHEN calling get_num_bits_different(...)
    result = dhash.get_num_bits_different(hash1, hash2)

    # THEN produce the expected XOR result
    assert result == expected

def test_force_pil():
    """force_pil() should overide wand to None"""

    # GIVEN wand as something other than None
    dhash.wand = 1

    # WHEN calling force_pil()
    dhash.force_pil()

    # THEN change dhash.wand to None
    assert dhash.wand is None

def test_dhash(test_filename, known_hash):
    """dhash(...) should perform the default configuartion of dhash"""
    # GIVEN a sample image
    # WHEN performing the default dhash
    result = dhash.dhash(test_filename)

    # THEN return a hex str of the known hash
    assert isinstance(result, str), "dhash(...) should return a str"

    hex_characters = "0123456789abcdef"
    chars_are_hex_chars = list(map(lambda x: x in hex_characters, result))
    assert all(chars_are_hex_chars), "all character should be between 0-9 and a-f"

    assert result == known_hash, "The result should match the known hash"

def test_dhash_diff(test_filename):
    """dhash(...) should provide an int of difference between two images"""

    # GIVEN the same image
    filename = test_filename

    # WHEN performing dhash_diff(...)
    result = dhash.dhash_diff(filename, filename)

    # THEN return a diff of zero
    assert isinstance(result, int), "dhash_diff(...) should return an int"
    assert result == 0, "The result should be zero"

    # GIVEN two different images
    filename180 = "dhash-test-180.jpg"

    # WHEN performing dhash_diff()
    result = dhash.dhash_diff(filename, filename180)
    assert isinstance(result, int), "dhash_diff(...) should return an int"
    assert result > 10, "The result should be zero"

def test_cli(pil_image, test_filename, known_hash):
    """dhash cli should produce hex format hash from size 8"""

    # GIVEN command line input for
    # dhash on the test image, no size or format specified
    cmd = f"python -m dhash {test_filename}"

    # WHEN executing on command line
    result = check_output(cmd, shell=True).decode().rstrip()

    # THEN dash, default at size 8 and return hex_format
    assert result == known_hash, "stdout should match expected output"


import numpy as np
from keras import layers
from keras.layers import Input, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D
from keras.layers import AveragePooling2D, MaxPooling2D, Dropout, GlobalMaxPooling2D, GlobalAveragePooling2D
from keras.utils import np_utils, print_summary
from keras.models import Sequential
from keras.callbacks import ModelCheckpoint
import pandas as pd
from sklearn.model_selection import train_test_split

import keras.backend as K

def keras_model(image_x, image_y):
    num_of_classes = 2
    model = Sequential()
    model.add(Conv2D(32, (5, 5), input_shape=(image_x, image_y, 1), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same'))
    model.add(Conv2D(64, (5, 5), activation='sigmoid'))
    model.add(MaxPooling2D(pool_size=(5, 5), strides=(5, 5), padding='same'))
    model.add(Flatten())
    model.add(Dense(1024, activation='relu'))
    model.add(Dropout(0.6))
    model.add(Dense(num_of_classes, activation='softmax'))

    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    filepath = "emojinator.h5"
    checkpoint1 = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True, mode='max')
    callbacks_list = [checkpoint1]

    return model, callbacks_list

def main():
    dataset = pd.read_csv("train_foo.csv")
    # Normalizing the training dataset (train.csv file)# Norma
    dataset = dataset.sample(frac=1)
    X = dataset.iloc[:,1:]
    Y = np.array(dataset.iloc[:,0])
    X = X/255.                      # for normalizing the inputs
    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, random_state=123, train_size = 0.8)

    #X_train = X_train.T
    #X_test = X_test.T
    print(Y_train.shape[0])
    Y_train = Y_train.reshape(Y_train.shape[0],1)
    Y_test = Y_test.reshape(-1,1)
    # Reshape
    print("number of training examples = " + str(X_train.shape[0]))
    print("number of test examples = " + str(X_test.shape[0]))
    print("X_train shape: " + str(X_train.shape))
    print("Y_train shape: " + str(Y_train.shape))
    print("X_test shape: " + str(X_test.shape))
    print("Y_test shape: " + str(Y_test.shape))
    image_x = 50
    image_y = 50

    import ipdb
    #ipdb.set_trace()
    train_y = np_utils.to_categorical(Y_train)
    test_y = np_utils.to_categorical(Y_test)
    train_y = train_y.reshape(train_y.shape[0], train_y.shape[1])
    test_y = test_y.reshape(test_y.shape[0], test_y.shape[1])
    X_train = np.array(X_train)
    X_train = X_train.reshape(X_train.shape[0], 50, 50, 1)
    X_test = np.array(X_test)
    X_test = X_test.reshape(X_test.shape[0], 50, 50, 1)
    print("X_train shape: " + str(X_train.shape))
    print("X_test shape: " + str(X_test.shape))

    model, callbacks_list = keras_model(image_x, image_y)
    model.fit(X_train, train_y, validation_data=(X_test, test_y), epochs=10, batch_size=64,
              callbacks=callbacks_list)
    scores = model.evaluate(X_test, test_y, verbose=0)
    print("CNN Error: %.2f%%" % (100 - scores[1] * 100))
    print_summary(model)

    model.save('emojinator.h5')

main()


#-*- encoding:utf-8 -*-
from db import *
import datetime
from random import *

#Sauvegarde de photo
def savePath(id_compte, mission_id, name_img, path_img, latitude, longitude):
	'''
	Cette fonction permet d'enregistrer le chemin du fichier image.
	Valeur de retour:
		* "" if success
		* 'error' if error
	'''
	try:
		db = Db()
		db.execute("INSERT INTO photoPath(id_compte, id_mission, name_photo, path_photo, \
			latitude_photo, longitude_photo) VALUES \
		 (%(compteID)s, %(missionID)s, %(namePhoto)s, %(pathPhoto)s, %(platitude)s, \
		 	%(plongitude)s)", {
		 "compteID":id_compte,
		 "missionID":mission_id,
		 "namePhoto":name_img,
		 "pathPhoto":path_img,
		 "platitude":latitude,
		 "plongitude":longitude
		 })
		db.close()
		return ""
	except:
		return 'error'

#Sauvegarde mission terminée
def save_mission_finished(id_compte, mission_id):
	'''
	Cette fonction permet d'enregistrer dans la table 'missionEnded' la mission nouvellement
	terminée.
	Valeur de retour:
		* Success: 'saved'
		* Error : 'error'
	'''
	try:
		db = Db()
		db.execute("INSERT INTO missionEnded(id_compte, id_mission, date_missionEnded) VALUES \
			(%(compteID)s, %(missionID)s, %(DateTime)s)", {
			"compteID": id_account,
			"missionID":id_compte,
			"DateTime":str(datetime.datetime.now())
			})
		db.close()
		return 'saved'
	except:
		return 'error'

#Sauvegarde photo validée
def saved_validate_photo(photo_id, id_compte, mission_id):
	'''
	Cette fonction permet de sauvegarder une photo qui a été validée.
	Valeur de retour:
		* Success: 'saved'
		* Error : 'error'
	'''
	try:
		#On enregistre la photo
		db = Db()
		db.execute("INSERT INTO validatePhoto(id_compte, id_mission, id_photo, date_validatePhoto)\
			VALUES(%(compteID)s, %(missionID)s, %(photoID)s, %(DateTime)s)", {
			"compteID":id_compte,
			"missionID":mission_id,
			"photoID":photo_id,
			"DateTime":str(datetime.datetime.now())
			})
		db.close()
		return 'saved'
	except:
		return 'error'

#Recupération de l'id de la photo
def get_id_photoWihName(name):
	'''
	Cette fonction permet de récupérer l'identifiant de la photo a partir de son nom.
	Valeur de retour:
		* Success : Id
		* Error : 0
	'''
	db = Db()
	identifiant = db.select("SELECT id_photo as id FROM photo WHERE (name_photo = %(namePhoto)s)", {
		"namePhoto":name
		})
	db.close()
	if (len(identifiant) == 1):
		return identifiant[0]['id']
	return 0

#Récupération du nom de la photo
def get_photoName(photo_id):
	'''
	Cette fonction permet de récupérer le nom d'une photo a partir de son identifiant.
	Valeur de retour:
		* Success : name (de la photo)
		* Error : ''
	'''
	db = Db()
	name = db.select("SELECT name_photo as name FROM photo WHERE id_photo = %(photoID)s", {
		"photoID":photo_id
		})
	db.close()
	if (len(name) == 1):
		return name[0]['name']
	return ''

#Récupération des points que vaut la photo
def get_pointsPhoto(photo_id):
	'''
	Cette fonction permet de récupérer le nombre de points que vaut une photo.
	Valeur de retour:
		* Succes: Id
		* Error : -1
	'''
	db = Db()
	pointsPhoto = db.select("SELECT points_photos as points FROM photo WHERE (id_photo = %(photoID)s)",{
		"photoID":photo_id
		})
	db.close()
	print(pointsPhoto)
	if (len(pointsPhoto) == 1):
		return pointsPhoto[0]['points']
	return -1

#Récupération de l'id du compte joueur
def getAccountDbID(identifier):
	'''
	Cette fonction permet de récupérer en base de données l'identifiant d'un compte à partir de son uuid.
	Valeur de retour:
		* Success: ID.
		* Error: 0
	'''
	db = Db()
	account_id = db.select("SELECT id_compte FROM compte WHERE (identifier_compte = %(the_identifier)s)", {
		"the_identifier":identifier
		})
	db.close()
	print(account_id)
	if (len(account_id) == 1):
		return account_id[0]['id_compte']
	return 0

#Récupération du nombre de points du compte joueur (score)
def get_pointsAccount(id_compte):
	'''
	Cette fonction permet de récupérer le nombre de points actuel d'un compte.
	Valeur de retour:
		* Success : nbDePoints
		* Error: -1
	'''
	db = Db()
	points = db.select("SELECT points_compte as points FROM compte WHERE id_compte = %(compteID)s", {
		"compteID":id_compte
		})
	db.close()

	if (len(points) == 1):
		return points[0]['points']
	return -1

#Récupération de l'id de la mission associé au compte 'compte_id'
def getMissionIDFrom(compte_id):
	'''
	Cette fonction permet de définir l'identifiant de la dernière mission associé à un compte.
	Valeur de retour:
		* Success: Id
		* Error : 0
	'''
	db = Db()
	missionID = db.select("SELECT id_mission FROM participer WHERE (id_compte = %(compteID)s)", {
		"compteID":compte_id
		})
	db.close()
	if (len(missionID) == 1):
		return missionID[0]['id_mission']
	return 0

#Récupération du nom de la mission
def get_missionName(mission_id):
	'''
	Cette fonction permet de récupérer le nom d'une mission a partir de son identifiant.
	Valeur de retour:
		* Success: Name of the mission
		* Error : ''
	'''
	db = Db()
	name_mission = db.select("SELECT name_mission FROM mission WHERE (id_mission = %(missionID)s)",{
		"missionID":mission_id
		})
	db.close()

	if (len(name_mission) == 1):
		return name_mission[0]['name_mission']
	return ''

#Dénombrement du nombre de photo validée
def count_nbPhotoValidate(id_compte, mission_id):
	'''
	Cette fonction permet de compter le nombre de photos validée par le compte.
	Valeur de retour:
		* Success: le nb de photo
		* Error: -1
	'''
	db = Db()
	count = db.select("SELECT COUNT(id_photo) FROM validatePhoto WHERE (id_compte = %(compteID)s AND \
		id_mission = %(missionID)s)", {
	"compteID":id_compte,
	"missionID":mission_id
	})
	db.close()

	if (len(count) == 1):
		return count[0]['count']
	return 'error'

#Dénombrement du nombre de photo affilié à une mission
def count_nbPhotoInMission(mission_id):
	'''
	Cette fonction permet de compter le nombre de photo liées à une mission.
	Valeur de retour:
		* Success: le nb de photo
		* Error: -1
	'''
	db = Db()
	count = db.select("SELECT COUNT(id_photo) FROM avoir WHERE (id_mission = %(missionID)s)", {
		"missionID":mission_id
		})
	db.close()
	if (len(count) == 1):
		return count[0]['count']
	return -1

#Calcul du nombre de points d'un joueur
def calculate_player_points(photo_id, id_compte):
	'''
	Cette fonction peremt de calculer les points du joueur en fonction de la validité ou non de la photo
	Valeur de retour:
		* Success : Score (Int) - le score du joueur
		* Error : 'error'
	'''
	#Recupération des points du compte du joueur
	points_player = get_pointsAccount(id_compte)

	#Recupération du nombre de points de la photo
	points_photo = get_pointsPhoto(photo_id)

	#Calcul de la nouvelle valeur
	if (points_player > -1 and points_photo > -1):
		return int(points_player) + int(points_photo)
	return 'error'

#Verification existence d'une photo validée
def check_if_photoAlready_validated(photo_id, id_compte, mission_id):
	'''
	Cette fonction permet de vérifier si pour la mission actuelle, une même photo a déjà été validée.
	Valeur de retour:
		* Success : True/False
		* Error : -1
	'''
	db = Db()
	count = db.select("SELECT COUNT(*) FROM validatePhoto WHERE (id_compte = %(compteID)s and \
		id_mission = %(missionID)s and id_photo = %(photoID)s)", {
		"compteID":id_compte,
		"missionID":mission_id,
		"photoID":photo_id
	})
	db.close()

	if (len(count) == 1):
		if (count[0]['count'] == 1):
			return True
		else:
			return False
	return -1

#Vérification de la validation d'une mission
def check_mission_validate(id_compte, mission_id):
	'''
	Cette fonction permet de vérifier si une mission donnée est validée par un joueur donné.
	Valeur de retour:
		* Success: True / False
		* Error : -1
	'''
	nb_photos_mission = count_nbPhotoInMission(mission_id)
	nb_photos_validated = count_nbPhotoValidate(id_compte, mission_id)

	if (nb_photos_validated > -1 and nb_photos_mission > -1):
		if (nb_photos_validated == nb_photos_mission):
			return True
		else:
			return False
	return -1

#Recupération de liste de données unique (identifiant ou noms)
def get_missionPhotosIDs(mission_id):
	'''
	Cette fonction permet de récupérer l'ensemble des identifiants des photos dans une mission.
	Valeur de retour:
		* Success: [id]
		* Error: []
	'''
	db = Db()
	photosIDS = db.select("SELECT id_photo AS id FROM avoir WHERE (id_mission = %(missionID)s)", {
		"missionID":mission_id
		})
	db.close()

	ids_photos = []
	if (len(photosIDS) > 0):
		for photoID in photosIDS:
			ids_photos.append(photoID['id'])
	return ids_photos

def get_photosNames(id_compte, mission_id):
	'''
	Cette fonction permet d'obtenir l'ensemble des noms des photos affiliées à une mission.
	Valeur de retour:
		* Success: [name]
		* Error : []
	'''
	#Valeur a retourner
	photosNames = []
	#Recupération de l'ensemble des identifiants des photos qui sont liées à la mission
	ids_photo = get_missionPhotosIDs(mission_id)
	if (len(ids_photo) > 0):
		#Recupération dans un tableau de l'ensemble des nom de ses photos
		for an_id in ids_photo:
			#On récupère le nom de la photo à l'identifiant en question
			photo_name = get_photoName(an_id)
			if (photo_name != ''):
				photosNames.append(photo_name)
	return photosNames

def get_MissionIDS():
	'''
	Cette fonction permet de récupérer l'ensemble des identifiants de missions présentes en base de données.
	Valeur de retour:
		* Success: [Id]
		* Error : []
	'''
	db = Db()
	missions_ids = db.select("SELECT id_mission as id FROM mission")
	db.close()

	ids_missions = []
	if (len(missions_ids) > 0):
		for mission_id in missions_ids:
			ids_missions.append(mission_id['id'])
	return ids_missions

#Update du nombre de points du compte joueur
def update_player_points(id_compte, new_value):
	'''
	Cette fonction permet de faire un update de la valeur des points d'un compte.
	Valeur de retour:
		* Success: 'updated'
		* Error : 'error'
	'''
	try:
		db = Db()
		db.execute("UPDATE compte SET points_compte= %d WHERE id_compte = %d" %(int(new_value), int(id_compte)))
		db.close()
		return 'updated'
	except:
		return 'error'

#Update de l'url du compte joueur
def update_url(id_compte, url_compte):
	'''
	Cette route permet d'effectuer une mise à jour du champs url du compte ayant l'id id_compte
	Valeur de retour:
		* Success: updated
		* Error : 'error'
	'''
	try:
		db = Db()
		the_url = "'"+ str(url_compte) +"'"
		print(the_url)
		db.execute("UPDATE compte SET url_compte= %s WHERE id_compte = %d" %(the_url, int(id_compte)))
		db.close()

		return 'updated'
	except:
		return 'error'

#Update d'information concernant la mission lié à un compte
def update_player_mission(id_compte, mission_id):
	'''
	Cette fonction permet de réaliser une mise à jour du champs de la table de relation 'participer' liant
	la table 'compte' et 'mission'.
	Cela permet d'attribuer une nouvelle mission à l'utilisateur du compte.
	Valeur de retour:
		* Success: 'updated'
		* Error : 'error'
	'''
	try:
		db = Db()
		db.execute("UPDATE participer SET id_mission=%d WHERE id_compte=%d" %(int(mission_id), int(id_compte)))
		db.close()
		return 'updated'
	except:
		return 'error'

#Choix de la mission pour un joueur
def choice_random_mission():
	'''
	Cette fonction permet d'attribuer une mission.
	Valeur de retour:
		* Success: Id (celui de la mission a assigner au joueur)
		* Error: 0
	'''
	mission_ids = get_MissionIDS()
	mission_id = 0
	if (len(mission_ids) > 0):
		shuffle(mission_ids)
		mission_id = choice(mission_ids)
	#print(mission_id)
	return mission_id

#Assignation d'une mission à un compte joueur
def assign_mission_at_player(player_id, mission_id):
	'''
	Cette fonction permet d'assigner à un joueur une mission présente en base de données.
	Le paramètre "mission_id" est une fonction
	Aucune valeur de retour :
		* Success: OK.
		* Error : Error
	'''
	try:
		db = Db()
		db.execute("INSERT INTO participer(id_compte, id_mission) VALUES (%(compteID)s, %(missionID)s)", {
			"compteID":player_id,
			"missionID":mission_id
			})
		db.close()
		return "OK"
	except:
		print("l'instance existe surement deja en base de données")
		return "Error"

#Récupération du nom et de la description d'une photo
def get_infosPhoto(photo_id):
	'''
	Cette fonction permet d'obtenir les informations a propos d'une photo.
	Valeur de retour:
		Succes: [{name:the_name, description:the_description}]
		Error : []
	'''
	db = Db()
	infosPhoto = db.select("SELECT name_photo as name, description_photo as description FROM photo WHERE (id_photo = %(photoID)s)", {
		"photoID":photo_id
		})
	db.close()
	return infosPhoto

#Récupération de l'identifiant du dernier compte créé
def recup_last_created_compteID():
	'''
	Cette fonction permet de récupérer l'identifiant de la dernière ligne ajoutée à la table
	compte.
	Valeur de retour:
		* Success: Id
		* Error : 0
	'''
	db = Db()
	max_id = db.select("SELECT MAX(id_compte) FROM compte")
	db.close()
	print(max_id)

	if (max_id > 0):
		return max_id[0]['max']
	return 0


import torch
import torch.nn as nn
import torch.nn.functional as F
import sys

class SFCN(nn.Module):
    def __init__(self, channel_number=[32, 64, 128, 256, 256, 64], output_dim=40, dropout=0.5,batch_norm=True):
        '''
            channel_number - list of ints - number of channels for 3d convolutions
            output_dim - output dimension of model - e.g. for sex prediction 2 or age prediction the number of bins
            dropout - float - probability of a single node to dropout at final layer
            batch_norm - bool - indicates whether batch normalization is applied
        '''
        super(SFCN, self).__init__()
        self.n_layer = len(channel_number)
        self.output_dim=output_dim
        self.channel_number=channel_number

        #Define part 1: feature extractor
        self.feature_extractor = nn.Sequential()
        for i in range(self.n_layer):
            if i == 0:
                in_channel = 1
            else:
                in_channel = channel_number[i-1]
            out_channel = channel_number[i]
            if i < self.n_layer-1:
                self.feature_extractor.add_module('conv_%d' % i,
                                                  self.conv_layer(in_channel,
                                                                  out_channel,
                                                                  maxpool=True,
                                                                  kernel_size=3,
                                                                  padding=1,
                                                                  batch_norm=batch_norm))
            else:
                self.feature_extractor.add_module('conv_%d' % i,
                                                  self.conv_layer(in_channel,
                                                                  out_channel,
                                                                  maxpool=False,
                                                                  kernel_size=1,
                                                                  padding=0,
                                                                  batch_norm=batch_norm))
        #Define part 2: classifier
        self.classifier = nn.Sequential()
        avg_shape = [5, 6, 5]
        self.classifier.add_module('average_pool', nn.AvgPool3d(avg_shape))

        if dropout is not None:
            self.classifier.add_module('dropout', nn.Dropout(dropout))

        in_channel = channel_number[-1]
        out_channel = self.output_dim
        self.classifier.add_module('conv_%d' % self.n_layer,
                                   nn.Conv3d(in_channel, out_channel, padding=0, kernel_size=1))

    @staticmethod
    def conv_layer(in_channel, out_channel, maxpool=True, kernel_size=3, padding=0, maxpool_stride=2,batch_norm=True):
        '''
        A function to build a convolutional layer including batch normalization, pooling and ReLU activation
        Input:
            in_channel, out_channel - ints - number of in and out channels
            maxpool - bool - indicates whether max-pooling is applied
            kernel_size.padding,maxpool_stride - int - parameters for model
            batch_norm - bool - indicates whether batch norm is applied
        Output:
            torch.nn.Sequential module - layer
        '''

        if batch_norm:
            batch_norm_layer=nn.BatchNorm3d(out_channel)
        else:
            batch_norm_layer=nn.Identity()

        if maxpool is True:
            layer = nn.Sequential(
                nn.Conv3d(in_channel, out_channel, padding=padding, kernel_size=kernel_size),
                batch_norm_layer,
                nn.MaxPool3d(2, stride=maxpool_stride),
                nn.ReLU(),
            )
        else:
            layer = nn.Sequential(
                nn.Conv3d(in_channel, out_channel, padding=padding, kernel_size=kernel_size),
                batch_norm_layer,
                nn.ReLU()
            )
        return layer

    def forward(self, x):
        '''
        Input: x - torch.Tensor - shape (batch_size,1,160,192,160)
        Output: torch.Tensor - shape (batch_size,self.output_dim)
                    - Every row consists of log-probabilities of a probability distribution.
        '''
        batch_size=x.size(0)
        x_f = self.feature_extractor(x)
        x = self.classifier(x_f)
        x=x.view(batch_size,-1)
        x = F.log_softmax(x, dim=1)
        return(x)

    def set_grad_classifier(self, val):
        '''
        Function to set the gradient parameters of classifier.
        Input: val - bool - value of gradient computation of classifier
        '''
        for p in self.classifier.parameters():
            p.requires_grad=val

    def set_grad_featext(self,vals,layers):
        '''
        Function to set the gradient parameters of feature extractor.
        Input:  val - list of bools - gradient parameters
                layers - list of ints - index of layer to set gradient parameter
        '''
        if len(layers)!=len(vals):
            sys.exit("Error in set_feature_extractor: Different lengths of val and layers.")

        for it in range(len(layers)):
            for p in self.feature_extractor[layers[it]].parameters():
                p.requires_grad=vals[it]

    def train_final_layer(self):
        '''
        Function to train classifier.
        '''
        self.set_grad_classifier(True)

    def train_full_model(self):
        '''
        After calling this function, all parameters
        of the model will be trained.
        '''
        for p in self.parameters():
            p.requires_grad=True

    def train_nothing(self):
        '''
        After calling this function, no parameters
        of the model will be trained.
        '''
        for p in self.parameters():
            p.requires_grad=False

    def train_only_final_layers(self,n_train_layers:int):
        '''
        Function to train (only!) the final layers of a network.
        Input: n_train_layers - int - number of final layers to train
                                (classifier is last layer)
        '''
        #Set gradient to zero everywhere:
        self.train_nothing()

        #Set gradient to true for the final layers:
        if n_train_layers>0:
            self.set_grad_classifier(True)

        if n_train_layers>1:
            n_feature_layers=n_train_layers-1

            #Count down from the end:
            feat_layers=[self.n_layer-it-1 for it in range(n_feature_layers)]
            #Set all true:
            feat_vals=[True for _ in feat_layers]

            self.set_grad_featext(feat_vals,feat_layers)

    def reinit_classifier_pres_scale(self):
        '''
        Function to reinitialize the parameters of classifier while preserving the scaling.
        '''
        for p in self.classifier.parameters():
            p_=p.data.detach()
            p.data=p_.std()*torch.randn(size=p_.shape).to(p_.device)+p_.mean()

    def reinit_featext_pres_scale(self,layers=None):
        '''
        Function to reinitialize the parameters of feature extractor while preserving the scaling.
        Input: layers - list of ints - indices of layers to reinitalize - if None everything is reinitialized
        '''
        if layers is None:
            layers=[it for it in range(self.n_layer)]
        for layer in layers:
            for p in self.feature_extractor[layer].parameters():
                p_=p.data.detach()
                p.data=p_.std()*torch.randn(size=p_.shape).to(p_.device)+p_.mean()

    def reinit_final_layers_pres_scale(self,n_train_layers:int):
        '''
        Function to reinitialize the final layers of the network while preserving the scaling.
        Input: n_train_layers - int - number of final layers to train
        '''
        n_train_layers=min(n_train_layers,self.n_layer+1)
        #Set gradient to true for the final layers:
        if n_train_layers>0:
            print("Initialize classifier layer.")
            self.reinit_classifier_pres_scale()
        if n_train_layers>1:
            n_feature_layers=n_train_layers-1
            feat_layers=[self.n_layer-it-1 for it in range(n_feature_layers)]
            print("Reinitialize feature layers: ", feat_layers)
            self.reinit_featext_pres_scale(feat_layers)

    def reinit_full_model_pres_scale(self):
        '''
        Function to reinitialize the full model while preserving scaling.
        '''
        self.reinit_classifier_pres_scale()
        self.reinit_featext_pres_scale()

    def change_output_dim(self, new_output_dim: int):
        '''
        new_output_dim - new output dimension
        '''
        c_in = self.classifier.conv_6.in_channels
        conv_last = nn.Conv3d(c_in, new_output_dim, kernel_size=1)
        self.classifier.conv_6 = conv_last
        self.output_dim=new_output_dim

#print(sfcn.feature_extractor)
# model=SFCN(dropout=0.5)
# #for name, module in model.named_modules():
# #    print(name)
# print(list(model.feature_extractor[0].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[1].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[2].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[3].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[4].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[5].parameters())[0].flatten()[0])
# print(list(model.classifier.parameters())[0].flatten()[0])
# # print()
# model.module.train_only_final_layers(2)
# model.reinit_full_model_pres_scale()
# print(list(model.feature_extractor[0].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[1].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[2].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[3].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[4].parameters())[0].flatten()[0])
# print(list(model.feature_extractor[5].parameters())[0].flatten()[0])
# print(list(model.classifier.parameters())[0].flatten()[0])


import numpy as np
import matplotlib.pyplot as plt

#percentiles = np.arange(100)
p = 20

# Read in IPSM fit result

centrality_class =  str(p) + '-' + str(p+1)
filename_fit = "../IPSM-Fit/output/" + centrality_class + ".txt"
One_sw2_N, sn2_N_N2 = np.loadtxt(filename_fit, unpack=True)

modes = [0, 1, 2, 3, 4]
lMax = 6

# Set up figure and image grid
figs, axs = plt.subplots(4, len(modes), figsize=(8.0,6.0),sharex="col", sharey="row", constrained_layout=True)

# Find maximal correlator value
maximal_value_Trento = 0
maximal_value_CGC = 0
maximal_value_Nc = 0
maximal_value_IPSM = 0
for mode in modes:
	#import two-point functions
	source = 'output/two_point_random_' + str(p) + '-' + str(p+1) + '_m' + str(mode) + '_real' +'.txt'
	profile = np.loadtxt(source)
	# import one-point functions
	source_one = 'output/one_point_' + str(p) + '-' + str(p+1)  +'.txt'
	profile_one_ml = np.loadtxt(source_one)
	# subtract disconnected part from m=0 mode
	if (mode == 0):
		profile[0,0] -= profile_one_ml[mode, 0] * profile_one_ml[mode, 0]

	current_max_Trento = max(np.amax(profile[:lMax,:lMax]), -np.amin(profile[:lMax,:lMax]))
	if (current_max_Trento > maximal_value_Trento):
		maximal_value_Trento = current_max_Trento
	# import CGC simple profiles
	source = 'Saclay_simplified/output/'+centrality_class+'/two_point_random_connected' + '_m_' + str(mode)  +'.txt'
	profile = np.loadtxt(source)
	# # add geometry part
	# source_one_point = 'output/one_point_'+centrality_class+'.txt'
	# one_points = np.loadtxt(source_one_point)
	# for i in range(0, lMax):
	# 	for j in range(0, lMax):
	# 		if ((i==0)&(j==0)&(mode==0)):
	# 			profile[i][j] += 1./np.pi/np.pi*(sn2_N_N2)
	# 		else:
	# 			profile[i][j] += (1.+sn2_N_N2)*one_points[mode, i]*one_points[mode, j]

	current_max_CGC = max(np.amax(profile[:lMax,:lMax]), -np.amin(profile[:lMax,:lMax]))
	if (current_max_CGC > maximal_value_CGC):
		maximal_value_CGC = current_max_CGC
	#import Large_Nc profile
	source = 'Saclay/output/'+centrality_class+'/Nr41/Nm64/m5.0e-3/two_point_random_connected' + '_m_' + str(mode)  +'.txt'
	profile = np.loadtxt(source)
	# # add geometry part
	# source_one_point = 'output/one_point_'+centrality_class+'.txt'
	# one_points = np.loadtxt(source_one_point)
	# for i in range(0, lMax):
	# 	for j in range(0, lMax):
	# 		if ((i==0)&(j==0)&(mode==0)):
	# 			profile[i][j] += 1./np.pi/np.pi*(sn2_N_N2)
	# 		else:
	# 			profile[i][j] += (1.+sn2_N_N2)*one_points[mode, i]*one_points[mode, j]

	current_max_Nc = max(np.amax(profile[:lMax,:lMax]), -np.amin(profile[:lMax,:lMax]))
	if (current_max_Nc > maximal_value_Nc):
		maximal_value_Nc = current_max_Nc
	#import IPSM profile
	source_one_point = 'output/one_point_'+centrality_class+'.txt'
	one_points = np.loadtxt(source_one_point)
	profile = np.zeros((lMax, lMax))
	clm_old = np.loadtxt("output/clm.txt")
	clm = np.zeros((lMax, len(modes)))
	for l in range(0, lMax):
		for m in range(0, len(modes)):
			if ((m==0)):
				if (l==0):
					clm[l][m] = 1./2
				else:
					clm[l][m] = clm_old[l-1][m]
			else:
				clm[l][m] = clm_old[l][m]
	for i in range(0, lMax):
		for j in range(0, lMax):
			if ((i == 0)&(j==0)&(mode == 0)):
				profile[i][j] = 1./np.pi/np.pi*(One_sw2_N+sn2_N_N2)
			elif ((i == j)):
				profile[i][j] = (-1.0)**mode/2./np.pi**2/clm[i, mode]*One_sw2_N + (1.+sn2_N_N2)*one_points[mode, i]*one_points[mode, j]
			else:
				profile[i][j] = (1.+sn2_N_N2)*one_points[mode, i]*one_points[mode, j]
	current_max_IPSM = max(np.amax(profile[:lMax,:lMax]), -np.amin(profile[:lMax,:lMax]))
	if (current_max_IPSM > maximal_value_IPSM):
		maximal_value_IPSM = current_max_IPSM
print("Trento: ", maximal_value_Trento)
print("CGC: ", maximal_value_CGC)
print("Nc: ", maximal_value_Nc)
print("IPSM: ", maximal_value_IPSM)

maximal_TN = max(maximal_value_Trento, maximal_value_Nc)
maximal_IC = max(maximal_value_CGC, maximal_value_IPSM)
maximal_total = max(maximal_TN, maximal_IC)
maximal_T = maximal_total
maximal_N = maximal_total
maximal_I = maximal_total
maximal_C = maximal_total

# maximal_T = maximal_value_Trento
# maximal_N = maximal_value_Nc
# maximal_I = maximal_value_IPSM
# maximal_C = maximal_value_CGC

cmap_Trento = plt.cm.seismic
cmap_LargeNc = plt.cm.seismic
cmap_IPSM = plt.cm.seismic
cmap_magma = plt.cm.seismic

from matplotlib.ticker import MaxNLocator
#for ax in axs:
# axs.xaxis.set_major_locator(MaxNLocator(integer=True))
# axs.yaxis.set_major_locator(MaxNLocator(integer=True))

# counter_fig = 0
# plot Trento diagrams
for mode in modes:
	#import two-point functions
	source = 'output/two_point_random_' + str(p) + '-' + str(p+1) + '_m' + str(mode) + '_real' +'.txt'
	profile = np.loadtxt(source)
	# import one-point functions
	source_one = 'output/one_point_' + str(p) + '-' + str(p+1)  +'.txt'
	profile_one_ml = np.loadtxt(source_one)
	# subtract disconnected part from m=0 mode
	if (mode == 0):
		profile[0,0] -= profile_one_ml[mode, 0] * profile_one_ml[mode, 0]

	im1 = axs[0][mode].imshow(profile[0:(lMax),0:(lMax)], interpolation=None, cmap=cmap_Trento, vmin = -maximal_T, vmax = maximal_T, extent = (-0.5+1, len(profile[0,0:(lMax)])-0.5+1, len(profile[0:(lMax),0])-0.5+1, -0.5+1))
	#centrality_class =  str(p) + '-' + str(p+1) + '%'
	#if (mode == 0):
		#axs[mode][0].set_title("TRENTo")

	axs[0][mode].set_title("$m={0}$".format(mode))
	axs[3][mode].set_xlabel("$l_2$")
	axs[0][mode].set_xticks(np.arange(lMax)+1)
	axs[0][mode].set_yticks(np.arange(lMax)+1)

	# # colorbar
	# if (mode == modes[-1]):
	# 	short = axs[mode][0].cax
	# 	short.colorbar(im)
	# 	tick_T = float('%.1g' % (maximal_T*0.6))
	# 	short.set_xticks([-tick_T, tick_T])
	# 	short.toggle_label(True)

	# counter_fig = counter_fig + 1

# plot Saclay diagrams
for mode in modes:
	#import two-point functions
	source = 'Saclay/output/'+centrality_class+'/Nr41/Nm64/m5.0e-3/two_point_random_connected' + '_m_' + str(mode)  +'.txt'
	profile = np.loadtxt(source)
	# # add geometry part
	# source_one_point = 'output/one_point_'+centrality_class+'.txt'
	# one_points = np.loadtxt(source_one_point)
	# for i in range(0, lMax):
	# 	for j in range(0, lMax):
	# 		if ((i==0)&(j==0)&(mode==0)):
	# 			profile[i][j] += 1./np.pi/np.pi*(sn2_N_N2)
	# 		else:
	# 			profile[i][j] += (1.+sn2_N_N2)*one_points[mode, i]*one_points[mode, j]

	axs[1][mode].imshow(profile[0:(lMax),0:(lMax)], interpolation=None, cmap=cmap_LargeNc, vmin = -maximal_N, vmax = maximal_N, extent = (-0.5+1, len(profile[0,0:(lMax)])-0.5+1, len(profile[0:(lMax),0])-0.5+1, -0.5+1))
	#centrality_class =  str(p) + '-' + str(p+1) + '%'
	#if (mode == 0):
		#axs[1][mode].set_title("Large $N_c$")
	#if (mode == modes[-1]):
		#axs[mode][1].set_xlabel("$l_2$")
	axs[1][mode].set_xticks(np.arange(lMax)+1)
	axs[1][mode].set_yticks(np.arange(lMax)+1)
	#ax.set_ylabel("$m={0}$\n$l_1$".format(mode))

	# # colorbar
	# if (mode == modes[-1]):
	# 	short = axs[mode][1].cax
	# 	short.colorbar(im)
	# 	tick_N = float('%.1g' % (maximal_N*0.6))
	# 	short.set_xticks([-tick_N, tick_N])
	# 	short.toggle_label(True)

# plot IPSM diagrams
for mode in modes:
	#import one-point functions
	source_one_point = 'output/one_point_'+centrality_class+'.txt'
	one_points = np.loadtxt(source_one_point)
	profile = np.zeros((lMax, lMax))

	clm_old = np.loadtxt("output/clm.txt")
	clm = np.zeros((lMax, len(modes)))
	for l in range(0, lMax):
		for m in range(0, len(modes)):
			if ((m==0)):
				if (l==0):
					clm[l][m] = 1./2
				else:
					clm[l][m] = clm_old[l-1][m]
			else:
				clm[l][m] = clm_old[l][m]

	for i in range(0, lMax):
		for j in range(0, lMax):
			if ((i == 0)&(j==0)&(mode == 0)):
				profile[i][j] = 1./np.pi/np.pi*(One_sw2_N+sn2_N_N2)
			elif ((i == j)):
				profile[i][j] = (-1.0)**mode/2./np.pi**2/clm[i, mode]*One_sw2_N + (1.+sn2_N_N2)*one_points[mode, i]*one_points[mode, j]
			else:
				profile[i][j] = (1.+sn2_N_N2)*one_points[mode, i]*one_points[mode, j]

	im2 = axs[2][mode].imshow(profile[0:(lMax),0:(lMax)], interpolation=None, cmap=cmap_IPSM, vmin = -maximal_I, vmax = maximal_I, extent = (-0.5+1, len(profile[0,0:(lMax)])-0.5+1, len(profile[0:(lMax),0])-0.5+1, -0.5+1))
	#centrality_class =  str(p) + '-' + str(p+1) + '%'
	#if (mode == 0):
		#axs[mode][2].set_title("IPSM")
	#if (mode == modes[-1]):
		#axs[mode][2].set_xlabel("$l_2$")
		#figs.colorbar(im2, ax = axs[2:4][mode], location = "right", shrink = 1, pad=0.0, aspect = 40)
	axs[2][mode].set_xticks(np.arange(lMax)+1)
	axs[2][mode].set_yticks(np.arange(lMax)+1)
	#ax.set_ylabel("$m={0}$\n$l_1$".format(mode))

	# # colorbar
	# if (mode == modes[-1]):
	# 	short = axs[mode][2].cax
	# 	short.colorbar(im)
	# 	tick_I = float('%.1g' % (maximal_I*0.6))
	# 	short.set_xticks([-tick_I, tick_I])
	# 	short.toggle_label(True)

# plot Saclay simple diagrams
for mode in modes:
	#import two-point functions
	source = 'Saclay_simplified/output/'+centrality_class+'/two_point_random_connected' + '_m_' + str(mode)  +'.txt'
	profile = np.loadtxt(source)
	# # add geometry part
	# source_one_point = 'output/one_point_'+centrality_class+'.txt'
	# one_points = np.loadtxt(source_one_point)
	# for i in range(0, lMax):
	# 	for j in range(0, lMax):
	# 		if ((i==0)&(j==0)&(mode==0)):
	# 			profile[i][j] += 1./np.pi/np.pi*(sn2_N_N2)
	# 		else:
	# 			profile[i][j] += (1.+sn2_N_N2)*one_points[mode, i]*one_points[mode, j]

	axs[3][mode].imshow(profile[0:(lMax),0:(lMax)], interpolation=None, cmap=cmap_magma, vmin = -maximal_C, vmax = maximal_C, extent = (-0.5+1, len(profile[0,0:(lMax)])-0.5+1, len(profile[0:(lMax),0])-0.5+1, -0.5+1))
	#centrality_class =  str(p) + '-' + str(p+1) + '%'
	#if (mode == 0):
		#axs[mode][3].set_title("magma")
	#if (mode == modes[-1]):
		#axs[mode][3].set_xlabel("$l_2$")
	axs[3][mode].set_xticks(np.arange(lMax)+1)
	axs[3][mode].set_yticks(np.arange(lMax)+1)
	#axs[mode][3].set_ylabel("$m={0}$\n$l_1$".format(mode))

	# # colorbar
	# if (mode == modes[-1]):
	# 	short = axs[mode][3].cax
	# 	short.colorbar(im)
	# 	tick_C = float('%.1g' % (maximal_C*0.6))
	# 	short.set_xticks([-tick_C, tick_C])
	# 	short.toggle_label(True)

	figs.colorbar(im1, ax = axs[0:4,4], shrink = 1, pad=0.0, aspect = 80, location = "right")
	axs[0][0].set_ylabel("TrENTo\n$l_1$")
	axs[1][0].set_ylabel("CGC large-$N_c$\n$l_1$")
	axs[2][0].set_ylabel("IPSM\n$l_1$")
	axs[3][0].set_ylabel("magma\n$l_1$")

	figs.suptitle("$G_{l_1, l_2}^{(m,-m)}$, "+centrality_class+"% class", x=0.5, y=1.05, fontsize=14)

	filename = "plots/two_point_comparison_horizontal_"+centrality_class+".pdf"
	#plt.subplots_adjust(wspace=0.05, hspace=0.05)
	plt.savefig(filename, format='pdf', bbox_inches = "tight")


import os
import torch
import torchvision
from dataset import CarvanaDataset
from torch.utils.data import DataLoader

CHECKPOINTS = 'checkpoints/'

def check_env():
    """
    Check PyTorch version and CUDA availability
    """
    print(f'CUDA: {torch.cuda.is_available()} - PyTorch: {torch.__version__}')

def save_checkpoint(state, epoch, verbose=0):
    """
    Save checkpoints
    """
    try:
        os.mkdir(CHECKPOINTS)
        print('Created checkpoint directory')
    except OSError:
        pass
    torch.save(state, CHECKPOINTS + f'BEST_CP_epoch{epoch + 1}.pth')
    if verbose == 1:
        print(f'Best checkpoint {epoch + 1} saved !')

def load_checkpoint(filename, model, verbose=0):
    """
    Load checkpoints
    """
    checkpoint = torch.load(CHECKPOINTS + filename)
    model.load_state_dict(checkpoint["state_dict"])
    if verbose == 1:
        print(f'Checkpoint {epoch + 1} loaded !')

def get_loaders(
    train_dir,
    train_maskdir,
    val_dir,
    val_maskdir,
    batch_size,
    train_transform,
    val_transform,
    num_workers=4,
    pin_memory=True,
):
    train_ds = CarvanaDataset(
        image_dir=train_dir,
        mask_dir=train_maskdir,
        transform=train_transform,
    )

    train_loader = DataLoader(
        train_ds,
        batch_size=batch_size,
        num_workers=num_workers,
        pin_memory=pin_memory,
        shuffle=True,
    )

    val_ds = CarvanaDataset(
        image_dir=val_dir,
        mask_dir=val_maskdir,
        transform=val_transform,
    )

    val_loader = DataLoader(
        val_ds,
        batch_size=batch_size,
        num_workers=num_workers,
        pin_memory=pin_memory,
        shuffle=False,
    )

    return train_loader, val_loader

def check_dice_score(loader, model, device="cuda"):
    """
    Compute the dice score
    """
    dice_score = 0
    model.eval()

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device)
            y = y.to(device).unsqueeze(1)
            preds = torch.sigmoid(model(x))
            preds = (preds > 0.5).float()
            dice_score += (2 * (preds * y).sum()) / (
                (preds + y).sum() + 1e-8
            )
    dice_score = dice_score / len(loader)
    model.train()

    return dice_score

def evaluate(loader, model, loss_fn, device="cuda"):
    loss = 0
    model.eval()
    for idx, (x, y) in enumerate(loader):
        x = x.to(device=device)
        y = y.float().unsqueeze(1).to(device=device)
        with torch.no_grad():
            predictions = model(x)
            loss = loss + loss_fn(predictions, y).item()

    model.train()

    return loss / len(loader)

def save_predictions_to_tensorboard(loader, model, epoch, writer, device="cuda"):
    model.eval() #https://stackoverflow.com/questions/60018578/what-does-model-eval-do-in-pytorch
    for idx, (x, y) in enumerate(loader):
        x = x.to(device=device)
        y = y.to(device=device)
        with torch.no_grad():
            preds = torch.sigmoid(model(x))
            preds = (preds > 0.5).float()

        writer.add_images('images', x, epoch)
        writer.add_images('masks/true', y.float().unsqueeze(1), epoch)
        writer.add_images('masks/pred', preds, epoch)

    model.train()

def save_results(loader, jeu, model, epoch, writer, device="cuda"):
    """
    Save results (Dice score and prdictions)
    """
    dice_score = 0
    model.eval()

    for idx, (x, y) in enumerate(loader):
        x = x.to(device)
        y = y.to(device).unsqueeze(1)
        with torch.no_grad():
            preds = torch.sigmoid(model(x))
            preds = (preds > 0.5).float()

        dice_score += (2 * (preds * y).sum()) / (
            (preds + y).sum() + 1e-8)
        writer.add_images(f'images/{jeu}', x, idx)
        writer.add_images(f'masks/{jeu}/true', y.float(), idx)
        writer.add_images(f'masks/{jeu}/pred', preds, idx)

    dice_score = dice_score / len(loader)
    writer.add_scalar(f'dice/{jeu}', dice_score, epoch)

    model.train()