from Models.MNIST_Model import MNIST_Model as Model
from Dataloaders.Mnist import Mnist as Dataloader
from Clients.SGDClient import SGDClient as Client
from Servers.FedKpServer import FedKpServer
from Servers.FedAvgServer import FedAvgServer
from Algorithms.FedAvg import FedAvg as Alg
import matplotlib.pyplot as plt
from sklearn.neighbors import KernelDensity
import matplotlib.gridspec as grid_spec
import seaborn as sns
import pandas as pd
import torch
import numpy as np
import os
from Models.Callbacks.Callbacks import Callbacks
import json

out_path  = os.path.join('data','Results','MNIST_exploration')
result_path = os.path.join(out_path,'results')
plot_path = os.path.join(out_path,'plots')
if not os.path.exists(out_path): os.mkdir(out_path)
if not os.path.exists(result_path): os.mkdir(result_path)
if not os.path.exists(plot_path): os.mkdir(plot_path)
# Parameters
alphas = [100,50,10,5,1,0.5,0.1,0.05,0.01,0.005,0.001]
n_clients = 10
seed = 0
batch_size = 16
device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
dl = Dataloader(n_clients)
test_data = dl.get_test_dataloader(batch_size)

##################### Plot distributions based on alpha ######################
plt.style.use('seaborn-whitegrid')
torch.manual_seed(seed)
np.random.seed(seed)

fig,axs = plt.subplots(2,2,figsize=(6,6))
for i,alpha in enumerate([10,1,0.1,0.01]):
    ax = axs[i//2,i%2]
    dl = Dataloader(n_clients,alpha=alpha)
    for j,c_data in enumerate(dl.get_training_raw_data()):
        labels = [l for x,l in c_data]
        counts = np.zeros(10)
        for l in labels:
            counts[l]+=1
        counts = counts/5420
        ax.scatter(np.repeat(j,10)+1,np.arange(10)+1,s = 100*counts,color='blue')
    ax.set_title('alpha = %.2f'%alpha)
    ax.set_xticks(np.arange(10)+1)
    ax.set_yticks(np.arange(10)+1)
    ax.grid(color='grey',alpha=0.2)
    if(i%2==0):
        ax.set_ylabel('Labels')
    if(i//2==1):
        ax.set_xlabel('Client')

plt.savefig(os.path.join(plot_path,'sampledist.jpg'))
plt.show()

################# Res 1 ###############################################

torch.manual_seed(seed)
np.random.seed(seed)
init_model = Model()
cbs = Callbacks(test_data,verbose=False,device=device)
callbacks = [cbs.server_loss,cbs.server_accuracy,cbs.server_training_loss,cbs.server_training_accuracy,
                cbs.skew,cbs.kurtosis,cbs.ks_test]
res = pd.DataFrame(columns=['alpha','train_acc','train_loss','val_loss','val_acc',
                            'ks_test','kurtosis','skew','rounds','epochs','n_clients'])

for alpha in alphas[1:]:
    dl = Dataloader(n_clients,alpha=alpha)
    alg = Alg(dl,Model,batch_size=batch_size)
    alg.server.model.set_weights(init_model.get_weights())
    alg.run(1,epochs=1,callbacks=callbacks,device=device)
    r = alg.get_callback_data()
    train_acc = np.mean(r['server_training_accuracy'][-1])
    train_loss = np.mean(r['server_training_loss'][-1])
    val_loss = r['server_loss'][-1]
    val_acc = r['server_accuracy'][-1]
    for skew,kurtosis,ks in zip(r['skew'][-1],r['kurtosis'][-1],r['ks_test'][-1]):
        res = res.append({'alpha':alpha,'train_acc':train_acc,'train_loss':train_loss,
                            'val_acc':val_acc,'val_loss':val_loss,
                            'ks_test':ks,'kurtosis':kurtosis,
                            'skew':skew,'rounds':1,'epochs':1,
                            'n_clients':n_clients},ignore_index = True)

import matplotlib.pyplot as plt
from matplotlib import gridspec  as grid_spec
from sklearn.neighbors import KernelDensity
import seaborn as sns
import os

res.to_csv(os.path.join(result_path,'res_alphas.csv'))

################# Res 5 ###############################################
seed=1
torch.manual_seed(seed)
np.random.seed(seed)
repeats = 5
init_model = Model()
cbs = Callbacks(test_data,verbose=False,device=device)
callbacks = [cbs.server_loss,cbs.server_accuracy,cbs.server_training_loss,cbs.server_training_accuracy]
res = pd.DataFrame(columns=['alpha','train_acc','train_loss','val_loss','val_acc'])
for _ in range(repeats):
    for alpha in alphas:
        dl = Dataloader(n_clients,alpha=alpha)
        alg = Alg(dl,Model,batch_size=batch_size)
        alg.server.model.set_weights(init_model.get_weights())
        alg.run(1,epochs=1,callbacks=callbacks,device=device)
        r = alg.get_callback_data()
        train_acc = np.mean(r['server_training_accuracy'][-1])
        train_loss = np.mean(r['server_training_loss'][-1])
        val_loss = r['server_loss'][-1]
        val_acc = r['server_accuracy'][-1]
        res = res.append({'alpha':alpha,'train_acc':train_acc,'train_loss':train_loss,
                            'val_acc':val_acc,'val_loss':val_loss},ignore_index = True)

res['val_acc'] = res['val_acc'].apply(lambda x: x.item())
res.to_csv(os.path.join(result_path,'res_lossacc.csv'))
res

########################################################################
############################ Plot Results ##############################
########################################################################

######## Res 1 ###############################
res = pd.read_csv(os.path.join(result_path,'res_alphas.csv'))
alphas = [100,50,10,5,1,0.5,0.1,0.05,0.01,0.005,0.001]
out_path = os.path.join('data','Results','Plots','Distributions')
metrics = ['ks_test','skew','kurtosis']
colors = sns.color_palette("ch:s=.25,rot=-.25", n_colors=len(alphas))
for metric in metrics:
    fig = plt.figure(figsize=(6,6))
    gs = grid_spec.GridSpec(len(alphas),1,1)
    ax_objs=[]
    for i,alpha in enumerate(alphas[::-1]):
        ax_objs.append(fig.add_subplot(gs[i:i+1, 0:]))
        ax = ax_objs[-1]
        res_tmp = res[res['alpha'] == alpha]
        x = np.array(res_tmp[metric])
        acc = np.array(res_tmp['val_acc'])[0]
        loss = np.array(res_tmp['val_loss'])[0]
        x_d = np.linspace(np.min(x),np.max(x), 1000)

        kde = KernelDensity(bandwidth=0.03, kernel='gaussian')
        kde.fit(x[:, None])

        logprob = kde.score_samples(x_d[:, None])

        # plotting the distribution
        ax.plot(x_d, np.exp(logprob),color="white",lw=1)
        ax.fill_between(x_d, np.exp(logprob), alpha=1,color=colors[i])

        # setting uniform x and y lims
        ax.set_xlim(np.min(x),np.max(x))

        # make background transparent
        rect = ax.patch
        rect.set_alpha(0)

        # remove borders, axis ticks, and labels
        ax.set_yticklabels([])
        ax.set_yticks([])

        if i == len(alphas)-1:
            if metric == 'ks_test':
                ax.set_xlabel('p-value', fontsize=14)
            else:
                ax.set_xlabel(metric, fontsize=14)
        else:
            ax.set_xticklabels([])
            ax.set_xticks([])

        spines = ["top","right","left","bottom"]
        for s in spines:
            ax.spines[s].set_visible(False)

        ax.text(np.min(x),0,r'$\alpha=$'+'%.3f'%(alpha),fontsize=10,ha="right")
        ax.grid(False)

    gs.update(hspace=-0.5)

    plt.tight_layout()

    name = 'alphas_%s'%(metric)
    plt.savefig(os.path.join(plot_path,name))

############## Res 5 #############################
res = pd.read_csv(os.path.join(result_path,'res_lossacc.csv')).drop('Unnamed: 0',axis=1)
stats = res.groupby('alpha').aggregate(['mean','std'])
cols = np.unique([c[0] for c in stats.columns])[[0,2,1,3]]
colors = sns.color_palette("Paired", n_colors=4)

fig,ax = plt.subplots(figsize=(6,6))
ax2 = ax.twinx()
x = np.array(stats.index)
for i,col in enumerate(cols):
    mu = np.array(stats[col]['mean'])
    std = np.array(stats[col]['std'])
    l = mu-std
    u = mu+std
    ax_tmp = ax2 if('loss' in col) else ax
    ax_tmp.plot(x,mu,label=col,color=colors[i])
    ax_tmp.fill_between(x,l,u,color=colors[i],alpha=0.3)
ax.legend(bbox_to_anchor=(1, -0.1),ncol=2,prop={'size': 10})
ax2.grid(False)
ax.grid(False)
ax2.legend(bbox_to_anchor=(0.5, -0.1),ncol=2,prop={'size': 10})
ax.set_xlabel(r'$\alpha$',fontsize=14)
ax.set_ylabel('Accuracy',fontsize=14)
ax2.set_ylabel('Loss',fontsize=14)
ax.set_xscale('log')

plt.savefig(os.path.join(plot_path,'lossandacc'))
plt.show()
