#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
import numpy as np
import matplotlib.pyplot as plt
c = 299792458
pi = np.pi

def read_data(dirname, distance, zshift):
    media = [1,2] # 1 - positive zshift, 2 - negative (need to add a minus sign for real shift).
    #min_mesh_step = 2.5 #nm
    data = []
    data.append([])
    for x in distance:
        data.append([])
        data[x].append([])
        for m in media:
            data[x].append([])
            for z in zshift:            
                monitor_name = "mon_x"+str(x)+"mkm_media"+str(m)+"_zshift"+z+"nm"
                data[x][m].append(
                    np.transpose(
                    np.genfromtxt(dirname+"/"+monitor_name+".txt", delimiter=", ",skip_header=1
                                  ,dtype=None, encoding = None
                                      , converters={0: lambda s: complex(s),
                                                    1: lambda s: complex(s),
                                                    2: lambda s: complex(s.replace('i', 'j')),
                                                    3: lambda s: complex(s.replace('i', 'j')),
                                                    4: lambda s: complex(s.replace('i', 'j')),
                                                    5: lambda s: complex(s.replace('i', 'j')),
                                                    6: lambda s: complex(s.replace('i', 'j')),
                                                    7: lambda s: complex(s.replace('i', 'j')),
                                                    8: lambda s: complex(s.replace('i', 'j'))
                                                    }
                                      )
                        )
                    )
    return data


def find_nearest(array,value):
    idx = (np.abs(array-value)).argmin()
    return array[idx],idx


def get_WLs_idx(WLs, data):
    dist = 1 #mkm
    mmedia = 1 # vacuum
    shift = 1 # one mesh step
    WLs_idx = []
    for wl in WLs:
        val, idx = find_nearest(data[dist][mmedia][shift][0,:],wl*1e-9)
        WLs_idx.append(idx)
    return WLs_idx


def analyze(data, dist, z_vec, wl_idx):
    #data = [dist][mmedia][shift] "lambda, dip.power, Ex, Ey, Ez, Hx, Hy, Hz, n_Au"
    #                              0     , 1        , 2 , 3 , 4 , 5 , 6 , 7 , 8   "
    data_in_air = np.array(data[dist][1])
    data_in_gold = np.array(data[dist][2])
    lambd = data_in_air[0][0,:]
    omega = 2*pi*c/lambd
    
    eps1 = complex(1)
    n_Au = data_in_air[0][8,:]
    eps2 = n_Au**2

    k_0 = omega/c #air
    k_spp = k_0*np.sqrt(eps1*eps2/(eps1+eps2))
    kappa1= np.sqrt(k_spp**2 - eps1*k_0**2)
    kappa2= np.sqrt(k_spp**2 - eps2*k_0**2)

    H1 = data_in_air[:,6,wl_idx]
    H2 = data_in_gold[:,6,wl_idx]
    for i in range(len(z_vec)):
        z = z_vec[i]*1e-9
        print("z =",z)
        H1_0 = H1[i]/np.exp(-kappa1[wl_idx]*z)
        H2_0 = H2[i]/np.exp(-kappa2[wl_idx]*z)
        print("H0 air ",H1_0," from H1",H1[i])
        print("H0 gold",H2_0," from H2",H2[i])
    # H1_0 = H1/np.exp(-kappa1* 
    # print(H1[0], H2[0],H1[0]- H2[0])
    # pl_data = (np.absolute(data_gold[:,2,wl_idx]*np.sqrt(dist)))
    # plt.semilogy(z_vec, pl_data,marker="o")
    # # legend = []
    # # legend.append(zshift[shift]+"@"+str(WLs[i])+" nm")
    # # plt.legend(legend)
    # # #plt.xlabel(r'THz')
    # plt.xlabel(r'Z shift, nm')
    # plt.ylabel(r'$Abs(E_x) \sqrt{R}$',labelpad=-5)
    # # plt.title(' r = '+str(core_r))
    # plt.savefig(dirname+"_z."+file_ext)
    # plt.clf()
    # plt.close()
    
file_ext="png"
dirname="template-dipole-on-sphere-on-surf-z.fsp.results"
def main ():
    distance = [1,2,3,4,5,6,7,8,9,10]
    zshift = ["5","20","200","400","600"]
    z_vec = [int(val) for val in zshift]

    data = read_data(dirname, distance, zshift)

    #WLs=[300,350,400,450,600,700,800]
    #WLs=[600,700, 800, 450]
    WLs=[800]#, 450]
    WLs_idx = get_WLs_idx(WLs, data)


    dist = 8 #mkm
    wl_idx = WLs_idx[0]
    
    analyze(data, dist, z_vec, wl_idx)


    # legend = []
    # for shift in range(len(zshift)):
    #     for i in range(len(WLs)):
    #         pl_data = []
    #         idx = WLs_idx[i]
    #         legend.append(zshift[shift]+"@"+str(WLs[i])+" nm")
    #         for dist in distance:
    #             pl_data.append(np.absolute(data[dist][mmedia][shift][2,idx]*np.sqrt(dist)))
    #         print(len(pl_data))
    #         plt.semilogy(distance, pl_data,marker="o")
    # plt.legend(legend)
    # # #plt.xlabel(r'THz')
    # plt.xlabel(r'Monitor R, $\mu$m')
    # plt.ylabel(r'$Abs(E_x) \sqrt{R}$',labelpad=-5)
    # # plt.title(' r = '+str(core_r))
    # plt.savefig(dirname+"_WLs."+file_ext)
    # plt.clf()
    # plt.close()
main()