# from wappy import *
from earsim import *
from atmocal import *

import matplotlib.pyplot as plt
from scipy.signal import hilbert
from scipy import signal
from scipy.interpolate import interp1d
from scipy.optimize import curve_fit,minimize
#import pandas as pd
import os

try:
    from tqdm import tqdm
except:
    tqdm = lambda x: x

try:
    from joblib import Parallel, delayed
except:
    Parallel = None
    delayed = lambda x: x

plt.rcParams.update({'font.size': 16})

atm = AtmoCal()

from matplotlib import cm

def location_to_shower_plane(loc, u=None, ev=None):
    if ev is not None:
        uAxB = ev.uAxB
        uAxAxB = ev.uAxAxB
        uA = ev.uA
    else:
        uAxB, uAxAxB, uA = u

    return np.dot(loc, uAxB), np.dot(loc, uAxAxB), np.dot(loc, uA)

def shower_plane_to_location( x, dXref=0, u=None, ev=None):
    if len(x) == 2:
        x, y = x
    else:
        x, y, dXref = x

    if ev is not None:
        uAxB = ev.uAxB
        uAxAxB = ev.uAxAxB
        uA = ev.uA
    else:
        uAxB, uAxAxB, uA = u

    return x * uAxB + y * uAxAxB + dXref * uA


def set_pol_and_bp(e,low=0.03,high=0.08):
    for ant in e.antennas:
        E = [np.dot(e.uAxB,[ex,ey,ez]) for ex,ey,ez in zip(ant.Ex,ant.Ey,ant.Ez)]
        dt =  ant.t[1] -ant.t[0]
        E = block_filter(E,dt,low,high)
        ant.E_AxB = E
        ant.t_AxB = ant.t


def pow_and_time(test_loc,ev,dt=1.0):
    t_ = []
    a_ = []
    t_min = 1e9
    t_max = -1e9
    for ant in ev.antennas:
        #propagate to test location
        aloc = [ant.x,ant.y,ant.z]
        delta,dist = atm.light_travel_time(test_loc,aloc)
        delta = delta*1e9
        t__ = np.subtract(ant.t_AxB,delta)
        t_.append(t__)
        a_.append(ant.E_AxB)
        if t__[0] < t_min:
            t_min = t__[0]
        if t__[-1] > t_max:
            t_max = t__[-1]

    t_sum = np.arange(t_min+1,t_max-1,dt)
    a_sum = np.zeros(len(t_sum))
    #interpolation
    for t_r,E_ in zip (t_,a_):
        f = interp1d(t_r,E_,assume_sorted=True,bounds_error=False,fill_value=0.)
        a_int = f(t_sum)
        a_sum = np.add(a_sum,a_int)
    if len(a_sum) != 0:
        P = np.sum(np.square(np.absolute(np.fft.fft(a_sum))))
        # normalise P with the length of the traces
        P = P/( t_sum[-1] - t_sum[0])
    else:
        print("ERROR, a_sum lenght = 0",
              "tmin ",t_min,
              "t_max ",t_max,
              "dt",dt)
        P = 0
    return P,t_,a_,a_sum,t_sum

def shower_axis_slice(e,Xb=200,Xe=1200,dX=2,zgr=0):
    zgr = zgr + e.core[2]
    N = int((Xe-Xb)/dX)
    Xs = np.array(np.linspace(Xb,Xe,N+1))
    ds = np.array([atm.distance_to_slant_depth(np.deg2rad(e.zenith),X,zgr) for X in Xs])

    locs = []
    for d_ in ds:
        xc = np.sin(np.deg2rad(e.zenith))*np.cos(np.deg2rad(e.azimuth))* d_
        yc = np.sin(np.deg2rad(e.zenith))*np.sin(np.deg2rad(e.azimuth))* d_
        zc = np.cos(np.deg2rad(e.zenith))* d_
        locs.append([xc,yc,zc])
    p = []
    for loc in locs:
            P,t_,pulses_,wav,twav = pow_and_time(loc,e)
            p.append(P)
    p = np.asanyarray(p)
    return ds,Xs,locs,p

def shower_plane_slice(e,X=750.,Nx=10,Ny=10,wx=1e3,wy=1e3,xoff=0,yoff=0,zgr=0,n_jobs=None, xs=None, ys=None):
    zgr = zgr + e.core[2]
    dX = atm.distance_to_slant_depth(np.deg2rad(e.zenith),X,zgr)
    x = xs if xs is not None else np.linspace(-wx,wx,Nx)
    y = ys if ys is not None else np.linspace(-wy,wy,Ny)

    def loop_func(x_, y_, xoff=xoff, yoff=yoff):
            loc = (x_+xoff)* e.uAxB + (y_+yoff)*e.uAxAxB + dX *e.uA
            P,t_,pulses_,wav,twav = pow_and_time(loc,e)

            return x_+xoff, y_+yoff, P, loc

    res = ( delayed(loop_func)(x_, y_) for x_ in x for y_ in y)

    if Parallel:
        #if n_jobs is None change with `with parallel_backend`
        res = Parallel(n_jobs=n_jobs)(res)

    # unpack loop results
    xx, yy, p, locs = zip(*res)

    xx = np.asarray(xx)
    yy = np.asarray(yy)
    p = np.asanyarray(p)

    return xx,yy,p,locs[np.argmax(p)]

def slice_figure(e,X,xx,yy,p,mode='horizontal', scatter_kwargs={}, colorbar_kwargs={'label':'Power'},suptitle=True, figsize=(10,8)):
    scatter_kwargs = {
            **dict(
                cmap='Spectral_r',
                alpha=0.9,
                s=30
                ),
            **scatter_kwargs
            }
    crosshair_kwargs = dict(ms=30, mew=3)


    fig, axs = plt.subplots(1,figsize=figsize)
    if suptitle:
        fig.suptitle(r'E = %.1g PeV, $\theta$ = %.1f$^\circ$, $\phi$ = %.1f$^\circ$ X = %.f'%(e.energy*10**3,e.zenith,e.azimuth,X))
    sc = axs.scatter(xx/1e3,yy/1e3,c=p,**scatter_kwargs)
    cbar = fig.colorbar(sc,ax=axs, **colorbar_kwargs)
    zgr = 0 + e.core[2]
    dX = atm.distance_to_slant_depth(np.deg2rad(e.zenith),X,zgr)
    xc = np.sin(np.deg2rad(e.zenith))*np.cos(np.deg2rad(e.azimuth))* dX
    yc = np.sin(np.deg2rad(e.zenith))*np.sin(np.deg2rad(e.azimuth))* dX
    if mode == 'horizontal':
        axs.plot(xc/1e3,yc/1e3,'r+', **crosshair_kwargs)
        axs.set_xlabel('x [km]')
        axs.set_ylabel('y [km]')
    elif mode == "sp":
        axs.plot(0,0, 'r+', **crosshair_kwargs)
        axs.set_xlabel('-v x B [km]')
        axs.set_ylabel(' vxvxB [km]')

    # indicate maximum power
    im = np.argmax(p)
    im_color = 'blue'
    im_norm = cbar.norm(p[im])
    if im_norm < 0.4:
        im_color = '#4488FF'
    if True:
        cbar.add_lines(levels=[p[im]], colors=[im_color], linewidths=[crosshair_kwargs['mew']])
        #cbar.lines[-1].set_linestyles('dotted')

    axs.plot(xx[im]/1e3,yy[im]/1e3, 'x', color=im_color, **crosshair_kwargs)
        #cbar.ax.axhline(cbar.norm(p[im]), color='b')
    fig.tight_layout()
    return fig,axs

def dist_to_line(xp,core,u):
    xp = np.array(xp)
    xp_core = xp-core
    c2 = np.dot(xp_core,xp_core)
    a2 = np.dot((np.dot(xp_core,u)*u),(np.dot(xp_core,u)*u))
    d = (np.abs(c2 - a2))**0.5
    return d

def dist_to_line_sum(param,data,weights):
    #distance line point: a = xp-core is D= | (a)^2-(a dot n)n  |
    #where ux is direction of line and x0 is a point in the line (like t = 0)
    x0 = param[0]
    y0 = param[1]
    theta = param[2]
    phi = param[3]
    core = np.array([x0, y0, 0.])
    u = np.array([np.cos(phi)*np.sin(theta),np.sin(phi)*np.sin(theta),np.cos(theta)])
    dsum = 0
    for xp,w in zip(data,weights):
        dsum += dist_to_line(xp,core,u)*w**2
    # print('%.2e %.2e %.2e %.2e %.2e'%(x0,y0,theta,phi,dsum))
    return dsum/len(data)

def get_axis_points(e,savefig=True,path="",zgr=0,Xlow=300, Xhigh=1000, N_X=15, n_jobs=None, tqdm=tqdm):
    Xsteps = np.linspace(Xlow, Xhigh, N_X)
    zgr=zgr+e.core[2] #not exact
    dXref = atm.distance_to_slant_depth(np.deg2rad(e.zenith),750,zgr)
    scale2d = dXref*np.tan(np.deg2rad(2.))
    scale4d = dXref*np.tan(np.deg2rad(4.))
    scale0_2d=dXref*np.tan(np.deg2rad(0.2))

    def loop_func(X):
        print("Starting", X)
        x,y,p,loc_max = shower_plane_slice(e,X,21,21,scale2d,scale2d)
        if savefig:
            fig,axs = slice_figure(e,X,x,y,p,'sp')
            fig.savefig(path+'X%d_a.pdf'%(X))
            plt.close(fig)
        im = np.argmax(p)
        if np.abs(x[im]) == np.max(x) or np.abs(y[im]) == (np.max(y)):
            x,y,p,loc_max = shower_plane_slice(e,X,21,21,scale4d,scale4d)
            if savefig:
                fig,axs = slice_figure(e,X,x,y,p,'sp')
                fig.savefig(path+'X%d_c.pdf'%(X))
                plt.close(fig)
        im = np.argmax(p)
        x,y,p,loc_max = shower_plane_slice(e,X,21,21,scale0_2d,scale0_2d,x[im],y[im])
        if savefig:
            fig,axs = slice_figure(e,X,x,y,p,'sp')
            fig.savefig(path+'X%d_b.pdf'%(X))
            plt.close(fig)

        print("Finished", X)
        return np.max(p), loc_max

    res = tqdm((delayed(loop_func)(X) for X in Xsteps), total=len(Xsteps))

    if Parallel:
        #if n_jobs is None change with `with parallel_backend`
        res = Parallel(n_jobs=n_jobs)(res)

    # unpack loop results
    max_vals, axis_points = zip(*res)

    return Xsteps,axis_points,max_vals

def fit_track(e,axis_points,vals,nscale=1e0):
    weights = vals/np.max(vals)
    data=axis_points[:]
    data = [d/nscale for d in data] #km, to have more comparable step sizes
    x0=[0,0,np.deg2rad(e.zenith),np.deg2rad(e.azimuth)]
    res = minimize(dist_to_line_sum,args=(data,weights),x0=x0)
    zen_r = np.rad2deg(res.x[2])
    azi_r = np.rad2deg(res.x[3])
    print(res,zen_r,e.zenith,azi_r,e.azimuth)
    return zen_r,azi_r,[res.x[0]*nscale,res.x[1]*nscale,0]


def update_event(e,core,theta,phi,axp=None):
    #recalculate
    e.zenith = theta
    e.azimuth = phi
    theta = np.deg2rad(theta)
    phi = np.deg2rad(phi)
    e.core = e.core+core
    e.uA = np.array([np.cos(phi)*np.sin(theta),np.sin(phi)*np.sin(theta),np.cos(theta)])
    e.uAxB = np.cross(e.uA,e.uB)
    e.uAxB = e.uAxB/(np.dot(e.uAxB,e.uAxB))**0.5
    e.uAxAxB = np.cross(e.uA,e.uAxB)
    #antenna position
    for a in e.antennas:
        a.x -= core[0]
        a.y -= core[1]
        a.z -= core[2]
    if axp != None:
        for ap in axp:
            ap[0] -= core[0]
            ap[1] -= core[1]
            ap[2] -= core[2]

def longitudinal_figure(dist,Xs,p,mode='grammage'):
    fig, axs = plt.subplots(1,figsize=(6,5))
    if mode=='grammage':
        axs.plot(Xs,p/np.max(p),'o-')
        axs.set_xlabel('X (g/cm$^2$)')
    if mode=='distance':
        axs.plot(dist/1e3,p/np.max(p),'o-')
        axs.set_xlabel('distance from ground (km)')
    axs.grid()
    fig.tight_layout()
    return fig

def time_residuals(e,tlable=True):
    ds,tp,nsum,ssum,swidth,azi,x,y,sid = lateral_parameters(e,True,[0,0,0])
    fig, axs = plt.subplots(1,figsize=(6,5),sharex=True)
    tp = tp-np.min(tp)
    cut_outlier = ~((ds<200)&(tp > 10))
    axs.plot(ds,tp,'o')
    if tlable:
        for d,t,s in zip(ds,tp,sid):
            plt.text(d,t,s)
        # axs.text(ds,tp,sid)
    axs.set_xlabel('distance (m)')
    axs.set_ylabel('$\Delta t (ns)$')
    axs.grid()
    z = np.polyfit(ds[cut_outlier],tp[cut_outlier],3)
    pfit = np.poly1d(z)
    xfit = np.linspace(np.min(ds),np.max(ds),100)
    yfit = pfit(xfit)
    tres = tp - pfit(ds)
    sigma =np.std(tres[cut_outlier])
    axs.plot(xfit,yfit,label=r'pol3 fit, $\sigma=%.2f$ (ns)'%(sigma))
    axs.legend()
    fig.tight_layout()
    return fig,tres

def figure_3D(axis_points,max_vals,zen,azi,core,res = 0):
    fig = plt.figure(figsize=(5,9))
    # fig, axs = plt.subplots(1,2,figsize=(12,8))
    ax = fig.add_subplot(2,1,1,projection='3d')
    xp = [ap[0]/1e3 for ap in axis_points]
    yp = [ap[1]/1e3 for ap in axis_points]
    zp = [ap[2]/1e3 for ap in axis_points]
    max_vals = np.asarray(max_vals)
    ax.scatter(xp, yp, zp,c=max_vals,s=150*(max_vals/np.max(max_vals))**2,cmap='Spectral_r')
    ax = fig.add_subplot(2,1,2)
    core = np.array(core)
    theta = np.deg2rad(zen)
    phi = np.deg2rad(azi)
    u = np.array([np.cos(phi)*np.sin(theta),np.sin(phi)*np.sin(theta),np.cos(theta)])
    residuals = [dist_to_line(ap,core,u) for ap in axis_points]
    dist = [np.sum((ap-core)**2)**0.5 for ap in axis_points]
    ax.scatter(dist,residuals,c=max_vals,cmap='Spectral_r')
    ax.grid()
    # ax.plot(xl,yl,zl,'-')
    # ax.set_zlim(0,18)
    # ax.view_init(15, 10)
    fig.tight_layout()
    if res != 0:
        res.track_dis.append(dist)
        res.track_res.append(residuals)
        res.track_val.append(max_vals)
    return fig

class RITResult():
    """docstring for RITResult."""
    def __init__(self):
        super(RITResult, self).__init__()
        self.xmax_rit = []
        self.xmax = []
        self.profile_rit = []
        self.dX = []
        self.dl = []
        self.zenith_ini = []
        self.azimuth_ini = []
        self.core_ini = []
        self.dcore_rec = []
        self.zenith_rec = []
        self.azimuth_rec = []
        self.index = []
        self.isMC = []
        self.track_dis = []
        self.track_res =[]
        self.track_val =[]
        self.station_ids =[]
        self.station_x =[]
        self.station_y =[]
        self.station_z =[]
        self.station_maxE = []
        self.has_pulse = []

def fill_stations_propeties(e,res):
    x = np.array([a.x for a in e.antennas])
    y = np.array([a.y for a in e.antennas])
    z = np.array([a.z for a in e.antennas])
    ids = [a.name for a in e.antennas]
    maxE = np.array([np.max(a.E_AxB) for a in e.antennas])
    #has_pulse = np.array([np.max(a.has_pulse) for a in e.antennas])
    res.station_x.append(x)
    res.station_y.append(y)
    res.station_z.append(z)
    res.station_ids.append(ids)
    #res.has_pulse.append(has_pulse)

def reconstruction(e,outfile='', slice_outdir=None, Xlow=300, Xhigh=1000, N_X=15, disable_pol_and_bp=False, tqdm=tqdm):
    res = RITResult()
    res.isMC.append(True)
    res.zenith_ini.append(e.zenith)
    res.azimuth_ini.append(e.azimuth)
    res.core_ini.append(e.core)

    if not disable_pol_and_bp:
        set_pol_and_bp(e, 0.03, 0.08)

    #only use signal that have a signal in data
    fill_stations_propeties(e,res)
    Xs,axis_points,max_vals = get_axis_points(e,savefig=(slice_outdir is not None), path=slice_outdir, Xlow=Xlow, Xhigh=Xhigh, N_X=N_X, tqdm=tqdm)
    zen,azi,core = fit_track(e,axis_points,max_vals,1e2)
    fig = figure_3D(axis_points,max_vals,zen,azi,core,res)
    fig.savefig(outfile)
    update_event(e,core,zen,azi)
    ds,Xs,locs,p = shower_axis_slice(e)
    #result
    res.dX.append(Xs)
    res.dl.append(ds)
    res.profile_rit.append(p)
    res.xmax_rit.append(Xs[np.argmax(p)])
    res.azimuth_rec.append(e.azimuth)
    res.zenith_rec.append(e.zenith)
    res.dcore_rec.append(core)
    return res

if __name__ == "__main__":
    file = '../ZH_airshower/mysim.sry'
    ev = REvent(file)
    set_pol_and_bp(ev)
    X = 750

    dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),X,0)
    scale2d = dXref*np.tan(np.deg2rad(2.))
    xx,yy,p,km= shower_plane_slice(ev,X,21,21,scale2d,scale2d)
    
    slice_figure(ev,X,xx,yy,p,mode='sp')
    #plt.scatter(xx,yy,c=p)
    #plt.colorbar()
    plt.show()