import numpy as np
import matplotlib.pyplot as pl
import sys

termbg_black = '\033[107m'
termbg_red = '\033[41m'
termbg_green = '\033[42m'
termbg_orange = '\033[43m'
termbg_blue = '\033[44m'
termbg_purple = '\033[45m'
termbg_cyan = '\033[46m'
termbg_lightgrey = '\033[47m'

def colored_i(text):   return "\033[38;2;{};{};{}m{} \033[38;2;255;255;255m".format(100 ,  100,  100, text)
def colored_R(text):   return "\033[38;2;{};{};{}m{} \033[38;2;255;255;255m".format( 41,  84, 255, text)
def colored_T(text):   return "\033[38;2;{};{};{}m{} \033[38;2;255;255;255m".format(77 , 232, 159, text)
def colored_A(text):   return "\033[38;2;{};{};{}m{} \033[38;2;255;255;255m".format(250, 124, 112, text)
def colored_TOT(text): return "\033[1m\033[38;2;{};{};{}m{} \033[38;2;255;255;255m\033[0m".format(245, 10, 10  , text)


def dtrap_poy(fname_in,nx,ny):
    poy_data = np.loadtxt(fname_in)
    x_export2D = poy_data[:,2].reshape((nx,ny))
    y_export2D = poy_data[:,3].reshape((nx,ny))
    poy_y_grid_re  = poy_data[:,10].reshape((nx,ny))
    temp=np.trapz(poy_y_grid_re,x_export2D[:,0],axis=0)
    return np.trapz(temp,y_export2D[0,:]) #[x_export2D,y_export2D,poy_y_grid_re] #

myDir = sys.argv[1]
str_FLAG_TOTAL = myDir[myDir.find('FLAG_TOTAL'):]

intpoyz_tot = abs(dtrap_poy(myDir+'/Poy_tot_gd.pos',50,50))
intpoyz_ref = abs(dtrap_poy(myDir+'/Poy_ref_gd.pos',50,50))
intpoyz_inc = abs(dtrap_poy(myDir+'/Poy_inc_gd.pos',50,50))
R_poy = intpoyz_ref/intpoyz_inc
T_poy = intpoyz_tot/intpoyz_inc

R1nm = np.loadtxt(myDir+'/eff_r1.txt',ndmin=2)[:,1] + 1j*np.loadtxt(myDir+'/eff_r1.txt',ndmin=2)[:,2]
T1nm = np.loadtxt(myDir+'/eff_t1.txt',ndmin=2)[:,1] + 1j*np.loadtxt(myDir+'/eff_t1.txt',ndmin=2)[:,2]
R2nm = np.loadtxt(myDir+'/eff_r2.txt',ndmin=2)[:,1] + 1j*np.loadtxt(myDir+'/eff_r2.txt',ndmin=2)[:,2]
T2nm = np.loadtxt(myDir+'/eff_t2.txt',ndmin=2)[:,1] + 1j*np.loadtxt(myDir+'/eff_t2.txt',ndmin=2)[:,2]
Q = [np.loadtxt(myDir+'/temp-Q_L_%g.txt'%k,ndmin=2)[:,1] for k in range(2,7)]
Q.append(np.loadtxt(myDir+'/temp-Q_scat.txt',ndmin=2)[:,1])
Q=np.array(Q)
TOT1 = R1nm.real.sum()+T1nm.real.sum()+Q.sum()
TOT2 = R2nm.real.sum()+T2nm.real.sum()+Q.sum()

if 'nanowire_solarcell' in myDir:
    print('cf pdf')
    Nmax=2
    tab_lambdas=np.loadtxt(myDir+'/temp_lambda_step.txt',ndmin=2)[:,8]
    nb_lambdas=len(tab_lambdas)
    R1nm = R1nm.reshape((nb_lambdas,2*Nmax+1,2*Nmax+1))
    T1nm = T1nm.reshape((nb_lambdas,2*Nmax+1,2*Nmax+1))
    Rtot = [R1nm[i].real.sum() for i in range(nb_lambdas)]
    Ttot = [T1nm[i].real.sum() for i in range(nb_lambdas)]
    Abs_rods = Q[-1]
    Abs_ITO  = Q[0]
    Abs_subs = Q[2]+Q[3]+Q[4]+Ttot
    TOT = Rtot+Abs_rods+Abs_ITO+Abs_subs
    pl.figure()
    pl.plot(tab_lambdas,Abs_ITO,label='absorption ITO electrode')
    pl.plot(tab_lambdas,Abs_rods,label='absorption in Silicon rods')
    pl.plot(tab_lambdas,Abs_subs,label='absorption in Silicon subs')
    pl.plot(tab_lambdas,Rtot,label='reflection')
    pl.plot(tab_lambdas,TOT,label='total')
    pl.legend()
    pl.xlabel(r'$\lambda$ [nm]')
    pl.ylabel('fraction of incident energy')
    pl.savefig('fig_solar_balance.pdf')
elif 'convergence' in myDir:
    pl.figure()
    data_abs = np.loadtxt(myDir+'/temp-Q_scat.txt',ndmin=2)[:,1]
    pl.plot(data_abs[0:int(len(data_abs)/2)],label='linear elements')
    pl.plot(data_abs[int(len(data_abs)/2):] ,label='curved elements')
    pl.xlabel('N / mesh size=$\lambda_0/N$')
    pl.xlabel('Absorption$')
    pl.legend()
    fname = 'fig_convergence_absorption_%s.pdf'%str_FLAG_TOTAL
    pl.savefig(fname)
    print('cf %s'%fname)
else:
    print(colored_i('===> Computed from diffraction efficiencies with tangential components only'))
    print(colored_R('Rtot1  = %.9f'%R1nm.real.sum()))
    print(colored_T('Ttot1  = %.9f'%T1nm.real.sum()))
    print(colored_A('Atot   = %.9f'%Q.sum()))
    print(colored_TOT('TOTAL1 = %.9f'%TOT1))
    print(colored_i('===> Computed from diffraction efficiencies with normal component'))
    print(colored_R('Rtot2  = %.9f'%R2nm.real.sum()))
    print(colored_T('Ttot2  = %.9f'%T2nm.real.sum()))
    print(colored_A('Atot   = %.9f'%Q.sum()))
    print(colored_TOT('TOTAL2 = %.9f'%TOT2))
    print(colored_i('===> Computed manually (trapz) from normal component of Poynting vector'))
    print(colored_i('(expected to be less precise)'))
    print(colored_R('Rtot3  = %.9f'%R_poy))
    print(colored_T('Ttot3  = %.9f'%T_poy))
    print(colored_A('Atot   = %.9f'%Q.sum()))
    print(colored_TOT('TOTAL3 = %.9f'%(Q.sum()+R_poy+T_poy)))
