import os
import numpy as np
import matplotlib.pyplot as plt
import sys


"""
    Plots for producing plots of the results, as generated by the training functions.
    
    Call with argument 'perm', 'trans' or 'rot' depending on which results you want to plot
    Second argument is number of experiments performed
"""

def get_perm_statistics(dir_name,nbr_exp,tp):
    
    dirr_name = os.path.join('results',dir_name)
    filename = os.path.join(dirr_name + '_00','eq')
    
    f = open(filename,'r')
    N = len(f.readlines())
    f.close()
    
    dirr_name = dir_name
    
    varnames = 'loss','lay0','lay1','shunt0','shunt1'
    res={}
    for var in varnames:
        res[var] = np.zeros((N+1,nbr_exp)) 
    
    dir_name = os.path.join('results', dirr_name)
    

    for n in range(nbr_exp):
        file = os.path.join(dir_name+'_0'+str(n),tp)

        f = open(file)    
        for k,line in enumerate(f.readlines()):
                    
                g=line.split()
                res['loss'][k+1,n] = float(g[1].replace(':',''))
                res['lay0'][k+1,n] = float(g[3].replace('[',''))
                res['lay1'][k+1,n] = float(g[4].replace(']',''))
                res['shunt0'][k+1,n] = float(g[6].replace('[','').replace(',',''))
                res['shunt1'][k+1,n] = float(g[7].replace(']',''))
        f.close()
            
        
        
        
        
    return res


def get_translation_statistics(dir_name,nbr_exp,tp):
    
    varnames = 'loss','lay0','lay1','lay2','shunt0','shunt1', 'shunt2'
    
    dirr_name = os.path.join('results',dir_name)
    filename = os.path.join(dirr_name + '_00','eq')
    
    f = open(filename,'r')
    N = len(f.readlines())
    f.close()
    res={}
    for var in varnames:
        res[var] = np.zeros((N+1,nbr_exp)) 
        

    for n in range(nbr_exp):
        file = os.path.join(dirr_name+'_0'+str(n),tp)

        f = open(file)    
        for k,line in enumerate(f.readlines()):
                    
               g=line.split()
               res['loss'][k+1,n] = float(g[1].replace(':',''))
               res['lay0'][k+1,n] = float(g[3].replace('[',''))
               res['lay1'][k+1,n]  = float(g[4].replace(',',''))
               res['lay2'][k+1,n]  = float(g[5].replace(']',''))
               res['shunt0'][k+1,n]  = float(g[7].replace('[','').replace(',',''))
               res['shunt1'][k+1,n]  = float(g[8].replace(',',''))
               res['shunt2'][k+1,n]  = float(g[9].replace(']',''))
        f.close()
        
    return res



def get_rotation_statistics(dir_name,nbr_exp,tp):
    
   dirr_name = os.path.join('results',dir_name)
   filename = os.path.join(dirr_name + '_01','eq')
    
   f = open(filename,'r')
   N = len(f.readlines())
   f.close()
    
   varnames = 'loss','lay0','lay1','lay2','lay3','shunt0','shunt1', 'shunt2','shunt3'
   res={}
   for var in varnames:
       res[var] = np.zeros((N+1,nbr_exp)) 
       res[var][:,0] =0.0
        
    
    # 1 3 4 6 7 9 10
   for n in range(nbr_exp):
        file = os.path.join(dirr_name+'_0'+str(n),tp)
        f = open(file)    
        for l,line in enumerate(f.readlines()):
            g=line.split()
           
            k = l+1
            
            
            res['loss'][k,n] = float(g[1].replace(':',''))
            res['lay0'][k,n] = float(g[3].replace('[',''))
            res['lay1'][k,n] = float(g[4].replace(',',''))
            res['lay2'][k,n]= float(g[5].replace(']',''))
            res['lay3'][k,n] = float(g[6].replace(']',''))
            res['shunt0'][k,n]= float(g[8].replace('[','').replace(',',''))
            res['shunt1'][k,n] = float(g[9].replace(',',''))
            res['shunt2'][k,n] = float(g[10].replace(']','').replace(',',''))
            res['shunt3'][k,n] = float(g[11].replace(']',''))
           

            f.close()
        
   return res

def plot_stat_to_stat(eqres,augres,non_augres,stat0,stat1,name = 'Unknown', y2 = 10e-6, non_aug=True):

    # plot two statistics against each other. Produces a plot similar (slightly less tidy) to the one in the paper
    

    fig = plt.figure()
    ax0 = fig.add_subplot(111)    # The big subplot
    ax1 = fig.add_subplot(121)
    ax2 = fig.add_subplot(122)

    ax = [ax1,ax2]
# Turn off axis lines and ticks of the big subplot
    ax0.spines['top'].set_color('none')
    ax0.spines['bottom'].set_color('none')
    ax0.spines['left'].set_color('none')
    ax0.spines['right'].set_color('none')
    ax0.tick_params(labelcolor='w', top=False, bottom=False, left=False, right=False)
    for k in range(2):
       
        ax[k].ticklabel_format(axis='both',scilimits=(0,0))
        ax[k].plot(eqres[stat0], eqres[stat1],c='green', alpha=.1,linestyle ='solid')
        ax[k].plot(augres[stat0], augres[stat1],c='blue', alpha=.1, linestyle = 'dashed')
        if non_aug:
            ax[k].plot(non_augres[stat0], non_augres[stat1],c='red', alpha=.1,linestyle ='dotted')
        ax[k].plot(eqres[stat0].mean(1), eqres[stat1].mean(1),c='green', alpha=1, linestyle = 'solid', label = 'Equivariant')
        ax[k].plot(augres[stat0].mean(1), augres[stat1].mean(1),c='blue', alpha=1, linestyle = 'dashed', label = 'Augmented')
        if non_aug:
            ax[k].plot(non_augres[stat0].mean(1), non_augres[stat1].mean(1),c='red', alpha=1, linestyle = 'dotted', label = 'Normal')
            
    
    ax2.set_ylim([-y2/10,y2])
    ax0.set_xlabel('$\Vert A - A^0\Vert$')
    ax0.set_ylabel('$\Vert A_{\mathcal{E^\perp}}\Vert$')
    ax2.legend()
    plt.subplots_adjust(left=.1)
    
    
def collect_layer_stats(eqres,augres,non_augres,nlay):
    reses = [eqres,augres,non_augres]
    
    # sum up the total errors /shifts
    
    for res in reses:
        res['totlay'] = 0.0
        res['totshunt'] = 0.0
        
        for k in range(nlay):
            res['totlay'] += res['lay'+str(k)]**2
            res['totshunt'] += res['shunt'+str(k)]**2
        
        res['totlay']=res['totlay']**.5
        res['totshunt']= res['totshunt']**.5
    
    return eqres,augres,non_augres


def plot_perm_statistics(n):
    eqres = get_perm_statistics('res_perm',n,'eq')
    augres = get_perm_statistics('res_perm',n,'aug')
    non_augres = get_perm_statistics('res_perm',n,'nonaug')
    
    
    eqres,augres,non_augres = collect_layer_stats(eqres,augres,non_augres, 2)
    
    plt.rcParams['text.usetex'] = True
    plt.rcParams['font.size'] = 14
    
    plot_stat_to_stat(eqres,augres,non_augres,'totshunt','totlay','Permutations',1e-7)
    plt.show()
    
def plot_transl_statistics(n):
    
    eqres = get_translation_statistics('res_trans',n,'eq')
    augres = get_translation_statistics('res_trans',n,'aug')
    non_augres = get_translation_statistics('res_trans',n,'nonaug')
    
    eqres,augres,non_augres = collect_layer_stats(eqres,augres,non_augres, 3)
    
    plt.rcParams['text.usetex'] = True
    plt.rcParams['font.size'] = 14
    
    
    plot_stat_to_stat(eqres,augres,non_augres,'totshunt','totlay','Translations',5e-7)
    plt.show()
    
def plot_rotation_statistics(n):
    eqres = get_rotation_statistics('res_rot',n,'eq')
    augres = get_rotation_statistics('res_rot',n,'aug')
    non_augres = get_rotation_statistics('res_rot',n,'nonaug')
    
    eqres,augres,non_augres = collect_layer_stats(eqres,augres,non_augres, 4)
    
    plt.rcParams['text.usetex'] = True
    plt.rcParams['font.size'] = 14
    
      
    plot_stat_to_stat(eqres,augres,non_augres,'totshunt','totlay','Rotations',8e-8)
    plt.title('Rotation')
    plt.show()
    
if __name__ == "__main__":
    
    name = sys.argv[1]
    nbr_exp = int(sys.argv[2])

    if name == 'perm':
        plot_perm_statistics(nbr_exp)
    if name == 'trans':
        plot_transl_statistics(nbr_exp)
    if name == 'rot':
        plot_rotation_statistics(nbr_exp)
    
    