import glob
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import pdb
import argparse

params = {
    'text.latex.preamble': ['\\usepackage{gensymb}'],
    'image.origin': 'lower',
    'image.interpolation': 'nearest',
    'image.cmap': 'gray',
    'axes.grid': False,
    'savefig.dpi': 150,  # to adjust notebook inline plot size
    'axes.labelsize': 10, # fontsize for x and y labels (was 10)
    'axes.titlesize': 10,
    'font.size': 8, # was 10
    'legend.fontsize': 8, # was 10
    'xtick.labelsize': 10,
    'ytick.labelsize': 10,
    'text.usetex': True,
    'figure.figsize': [12, 6],
    'font.family': 'serif',
}
matplotlib.rcParams.update(params)

parser = argparse.ArgumentParser()

parser.add_argument("--noISAT", help="Is ISAT on")

args = parser.parse_args()

fig,axes = plt.subplots(nrows=2,ncols=3)
if (args.noISAT !="True"):

    # Select number of elements you want to read from arrays
    # This is much more easier
    #nElem = 450

    firstRank = True
    for rank in sorted(glob.glob('processor*')):
        # Plot add, retrieve, grow and solve
    
        solve_time = []
        add_time = []
        grow_time = []
        retrieve_time = []
        solve_cpu = []
        add_cpu = []
        retrieve_cpu = []
        grow_cpu = []
        ISAT_time = []
        ISAT_size = []


        cpu_solve = open(rank+'/TDAC/cpu_solve.out')
        cpu_retrieve = open(rank+'/TDAC/cpu_retrieve.out')
        cpu_grow = open(rank+'/TDAC/cpu_grow.out')
        cpu_add = open(rank+'/TDAC/cpu_add.out')
        size_ISAT = open(rank+'/TDAC/size_isat.out')

        lines_solve = cpu_solve.readlines()
        lines_add = cpu_add.readlines()
        lines_grow = cpu_grow.readlines()
        lines_retrieve = cpu_retrieve.readlines()
        lines_size = size_ISAT.readlines()
        for x in lines_solve:
            solve_time.append(x.split('    ')[0])
            solve_cpu.append(x.split('    ')[1])




        for x in lines_add:
            add_time.append(x.split('    ')[0])
            add_cpu.append(x.split('    ')[1])
        



        for x in lines_retrieve:
            retrieve_time.append(x.split('    ')[0])
            retrieve_cpu.append(x.split('    ')[1])
        



        for x in lines_grow:
            grow_time.append(x.split('    ')[0])
            grow_cpu.append(x.split('    ')[1]) 


        for x in lines_size:
            ISAT_time.append(x.split('    ')[0])
            ISAT_size.append(x.split('    ')[1]) 

        solve_cpu = [float(i) for i in solve_cpu]
        add_cpu = [float(i) for i in add_cpu]
        retrieve_cpu = [float(i) for i in retrieve_cpu]
        grow_cpu = [float(i) for i in grow_cpu]
        ISAT_size = [float(i) for i in ISAT_size]
        solve_time = [float(i) for i in solve_time]
        add_time = [float(i) for i in add_time]
        retrieve_time = [float(i) for i in retrieve_time]
        grow_time = [float(i) for i in grow_time]
        ISAT_time = [float(i) for i in ISAT_time]

        if(firstRank):
            nElem = np.size(solve_cpu)-5
            
        solve_cpu = solve_cpu[:nElem]
        add_cpu = add_cpu[:nElem]
        retrieve_cpu = retrieve_cpu[:nElem]
        grow_cpu = grow_cpu[:nElem]
        ISAT_size = ISAT_size[:nElem]
        solve_time = solve_time[:nElem]
        add_time = add_time[:nElem]
        retrieve_time = retrieve_time[:nElem]
        grow_time = grow_time[:nElem]
        ISAT_time = ISAT_time[:nElem]
        

        if(firstRank):
            solve_cpu_ref = solve_cpu
            add_cpu_ref = add_cpu
            retrieve_cpu_ref = retrieve_cpu
            grow_cpu_ref = grow_cpu
            solve_mean = np.zeros(len(solve_cpu_ref))
            add_mean = np.zeros(len(add_cpu_ref))
            retrieve_mean = np.zeros(len(retrieve_cpu_ref))
            grow_mean = np.zeros(len(grow_cpu_ref))
            firstRank = False
            total_cpu_ref  = [sum(x) for x in zip(add_cpu_ref,solve_cpu_ref,grow_cpu_ref,retrieve_cpu_ref)]
            solve_min = solve_cpu
            add_min = add_cpu
            retrieve_min = retrieve_cpu
            grow_min = grow_cpu
            solve_max = solve_cpu
            add_max = add_cpu
            retrieve_max = retrieve_cpu
            grow_max = grow_cpu
            size_ref = ISAT_size
            size_mean =  np.zeros(len(size_ref))
            size_min = ISAT_size
            size_max = ISAT_size
            total_cpu_max = [sum(x) for x in zip(add_cpu,solve_cpu,grow_cpu,retrieve_cpu)]
        
        print(rank)
        total_cpu =  [sum(x) for x in zip(add_cpu,solve_cpu,grow_cpu,retrieve_cpu)]
        total_cpu_max = np.maximum(total_cpu_max,total_cpu)

        # Get the minimum
        solve_min = np.minimum(solve_min,solve_cpu)   
        add_min = np.minimum(add_min,add_cpu)   
        retrieve_min = np.minimum(retrieve_min,retrieve_cpu)   
        grow_min = np.minimum(grow_min,grow_cpu)   
        size_min = np.minimum(size_min,ISAT_size)   
        
        # Get the maximum
        solve_max = np.maximum(solve_max,solve_cpu)   
        add_max = np.maximum(add_max,add_cpu)   
        retrieve_max = np.maximum(retrieve_max,retrieve_cpu)   
        grow_max = np.maximum(grow_max,grow_cpu)   
        size_max = np.maximum(size_max,ISAT_size)
           
        # Get the mean
        solve_mean = solve_mean + solve_cpu
        add_mean = add_mean + add_cpu
        retrieve_mean = retrieve_mean + retrieve_cpu
        grow_mean = grow_mean + grow_cpu
        print(rank)
        
    nRank  = len(sorted(glob.glob('processor*')))
    solve_mean /= nRank 
    add_mean /= nRank 
    retrieve_mean /= nRank 
    grow_mean /= nRank 
    n = 10
    print("Starting the plotting")
    axes[0,0].plot(solve_time[::n],solve_max[::n],label='Max')
    axes[0,0].plot(solve_time[::n],solve_min[::n],label='Min')
    axes[0,0].plot(solve_time[::n],solve_mean[::n],label='Mean')
    axes[0,0].set_ylabel('solve [s]')
    axes[0,0].set_xlabel('Sim time [s]')
    axes[0,0].set_xticks(np.linspace(solve_time[0],solve_time[-1],3))
    axes[0,0].legend(loc=2)
    
    
    axes[0,1].plot(add_time[::n],add_max[::n],label='Max')
    axes[0,1].plot(add_time[::n],add_min[::n],label='Min')
    axes[0,1].plot(add_time[::n],add_mean[::n],label='Mean')
    axes[0,1].set_ylabel('add [s]')
    axes[0,1].set_xlabel('Sim time [s]')
    axes[0,1].set_xticks(np.linspace(add_time[0],add_time[-1],3))

    axes[0,2].plot(grow_time[::n],grow_max[::n],label='Max')
    axes[0,2].plot(grow_time[::n],grow_min[::n],label='Min')
    axes[0,2].plot(grow_time[::n],grow_mean[::n],label='Mean')
    axes[0,2].set_ylabel('grow [s]')
    axes[0,2].set_xlabel('Sim time [s]')
    axes[0,2].set_xticks(np.linspace(grow_time[0],grow_time[-1],3))
    
    axes[1,0].plot(retrieve_time[::n],retrieve_max[::n],label='Max')
    axes[1,0].plot(retrieve_time[::n],retrieve_min[::n],label='Min')
    axes[1,0].plot(retrieve_time[::n],retrieve_mean[::n],label='Mean')
    axes[1,0].set_ylabel('retrieve [s]')
    axes[1,0].set_xlabel('Sim time [s]')
    axes[1,0].set_xticks(np.linspace(retrieve_time[0],retrieve_time[-1],3))
    
    
    axes[1,1].plot(solve_time[::n],solve_max[::n],label='Solve')
    axes[1,1].plot(retrieve_time[::n],retrieve_max[::n],label='Retrieve')
    axes[1,1].plot(grow_time[::n],grow_max[::n],label='Grow')
    axes[1,1].plot(add_time[::n],add_max[::n],label='Add')
    axes[1,1].plot(add_time[::n],total_cpu_max[::n],label='Total')
    axes[1,1].set_ylabel('max [s]')
    axes[1,1].set_xlabel('Sim time [s]')
    axes[1,1].set_xticks(np.linspace(retrieve_time[0],retrieve_time[-1],3))
    axes[1,1].legend(loc=1,frameon=False)
    
    axes[1,2].plot(ISAT_time[::n],size_max[::n],label='Max')
    axes[1,2].plot(ISAT_time[::n],size_min[::n],label='Min')
    axes[1,2].set_ylabel('Table size [-]')
    axes[1,2].set_xlabel('Sim time [s]')
    axes[1,2].set_xticks(np.linspace(ISAT_time[0],ISAT_time[-1],3))
    axes[1,2].legend(loc=2)
        
    
    fig.tight_layout()
    fig.savefig('ISAT_statistics.png',format='png',bbox_inches='tight',pad_inches=0,dpi=700)  
 
    print("Total CPU Times are:\n Retrieve: {} ({} %) \n Solve: {} ({} %) \n Add: {} ({} %)  \n Grow: {} ({} %) \n Total:{} \n".format(np.sum(retrieve_max),100*(np.sum(retrieve_max)/np.sum(total_cpu_max)),np.sum(solve_max),100*(np.sum(solve_max)/np.sum(total_cpu_max)),np.sum(add_max),100*(np.sum(add_max)/np.sum(total_cpu_max)),np.sum(grow_max),100*(np.sum(grow_max)/np.sum(total_cpu_max)),np.sum(total_cpu_max)))
    print("Total number of iterations: {} ".format(np.size(retrieve_max)))
   
else:
        # Select number of elements you want to read from arrays
    # This is much more easier
    #nElem = 450

    firstRank = True
    for rank in sorted(glob.glob('processor*')):
        # Plot add, retrieve, grow and solve
    
        solve_time = []
   
        solve_cpu = []



        cpu_solve = open(rank+'/TDAC/cpu_solve.out')
      
        lines_solve = cpu_solve.readlines()
        
        for x in lines_solve:
            solve_time.append(x.split('    ')[0])
            solve_cpu.append(x.split('    ')[1])




      

        solve_cpu = [float(i) for i in solve_cpu]
    
        solve_time = [float(i) for i in solve_time]
       
        if(firstRank):
            nElem = np.size(solve_cpu)-5
            
        solve_cpu = solve_cpu[:nElem]
       
        solve_time = solve_time[:nElem]
      
        if(firstRank):
            solve_cpu_ref = solve_cpu
          
            solve_mean = np.zeros(len(solve_cpu_ref))
          
            firstRank = False
           
            solve_min = solve_cpu
        
            solve_max = solve_cpu
          
              
        # Get the minimum
        solve_min = np.minimum(solve_min,solve_cpu)   
      
        # Get the maximum
        solve_max = np.maximum(solve_max,solve_cpu)   
      
        # Get the mean
        solve_mean = solve_mean + solve_cpu
   
        print(rank)
        
    nRank  = len(sorted(glob.glob('processor*')))
    solve_mean /= nRank 

    n = 10
    print("Starting the plotting")
    plt.plot(solve_time[::n],solve_max[::n],label='Max')
    plt.plot(solve_time[::n],solve_min[::n],label='Min')
    plt.plot(solve_time[::n],solve_mean[::n],label='Mean')
    plt.ylabel('solve [s]')
    plt.xlabel('Sim time [ss]')
    plt.xticks(np.linspace(solve_time[0],solve_time[-1],3))
    plt.legend(loc=1,frameon=False)
    plt.tight_layout()
    ax = plt.gca()
    print("Total CPU Times are:\n Solve: {} ".format(np.sum(solve_max)))
    print("Total number of iterations: {} ".format(np.size(solve_max)))
    plt.savefig('ISAT_statistics.png',dpi=500,bbox_inches='tight',pad_inches=0)