#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
from copy import deepcopy
from scipy.interpolate import interp1d
import pandas as pd
from pathos.multiprocessing import Pool
from scipy.signal import savgol_filter
from os import makedirs
from os.path import exists as path_exist
from ..ANTARESS_analysis.ANTARESS_inst_resp import calc_FWHM_inst
from ..ANTARESS_general.utils import init_parallel_func,np_where1D,is_odd,stop,dataload_npz,datasave_npz,check_data


def process_spectral_cont(vis_mode,data_type_gen,inst,data_dic,gen_dic,vis):
    r"""**Stellar continuum: main routine.**
    
    Returns generic stellar continuum from binned master spectrum.

    Args:
        TBD
    
    Returns:
        TBD
    
    """       
    print('   > Defining stellar continuum on '+gen_dic['type_name'][data_type_gen]+' master') 
    data_inst = data_dic[inst]   

    #Using master from several visits
    if vis_mode=='multivis':
        vis_det='binned'
        if data_inst['type']!='spec1D':stop('Spectra must be 1D') 

    #Using master from single visit
    elif vis_mode=='':
        vis_det=vis   
        if data_inst[vis]['type']!='spec1D':stop('Spectra must be 1D')    

    #Processing data    
    save_data_paths = gen_dic['save_data_dir']+'Stellar_cont_'+data_type_gen+'/'+inst+'_'+vis_det+'/'
    if not path_exist(save_data_paths):makedirs(save_data_paths)  
    if (gen_dic['calc_'+data_type_gen+'_stcont']):
        print('         Calculating data')
        prop_dic = deepcopy(data_dic[data_type_gen]) 
    
        #Retrieve binning information
        data_bin = dataload_npz(gen_dic['save_data_dir']+data_type_gen+'bin_data/'+gen_dic['add_txt_path'][data_type_gen]+inst+'_'+vis_det+'_'+prop_dic['dim_bin']+'_add')
    
        #Retrieve master spectrum
        if data_bin['n_exp']>1:stop('Bin data into a single master spectrum')
        data_mast = dataload_npz(gen_dic['save_data_dir']+data_type_gen+'bin_data/'+gen_dic['add_txt_path'][data_type_gen]+inst+'_'+vis_det+'_'+prop_dic['dim_bin']+str(0))

        #Check for alignment
        if (not gen_dic['align_'+data_type_gen]) or ((data_type_gen=='DI') and (not gen_dic['align_DI'] and data_bin['sysvel']==0.)):
            stop('Data must have been aligned in the stellar rest frame')
            
        #Check for in-transit contamination
        if (data_type_gen=='DI') and data_bin['in_inbin']:
            stop('Disk-integrated master contain in-transit profiles')    
    
        #Limit master to minimum definition range
        idx_def_mast = np_where1D(data_mast['cond_def'][0])
        flux_mast = data_mast['flux'][:,idx_def_mast[0]:idx_def_mast[-1]+1]
        cen_bins_mast = data_mast['cen_bins'][:,idx_def_mast[0]:idx_def_mast[-1]+1]
        cond_def_mast = data_mast['cond_def'][:,idx_def_mast[0]:idx_def_mast[-1]+1]
        edge_bins_mast = data_mast['edge_bins'][:,idx_def_mast[0]:idx_def_mast[-1]+2]
        
        #Stellar continuum
        min_edge_ord = cen_bins_mast[0,0]
        dic_sav = {}
        _,cont_func_dic,_ = calc_spectral_cont(1,[0],flux_mast,cen_bins_mast,edge_bins_mast,cond_def_mast,None,None,inst,gen_dic['contin_roll_win'][inst],gen_dic['contin_smooth_win'][inst],gen_dic['contin_locmax_win'][inst],\
                                                         gen_dic['contin_stretch'][inst],gen_dic['contin_pinR'][inst],min_edge_ord,dic_sav,1)
            
        #Saving data     
        datasave_npz(save_data_paths+'St_cont',{'cont_func_dic': cont_func_dic[0]})

    else:
        data_paths={'path':save_data_paths+'St_cont'}
        check_data(data_paths)          
        
    return None



def calc_spectral_cont(nord,iord_proc_list,flux_mast_in,cen_bins_mast,edge_bins_mast,cond_def_mast,flux_Earth_all_in,cond_def_all_in,inst,roll_win,smooth_win,locmax_win,par_stretch,\
                       contin_pinR,min_edge_ord,dic_sav,nthreads):
    r"""**Stellar continuum: wrap-up for calculation.**
    
    Calls direct calculation or multi-threaded one.

    Args:
        TBD
    
    Returns:
        TBD
    
    """     
    mean_flux_mast = np.zeros(nord,dtype=float)*np.nan
    if flux_mast_in is None:
        flux_mast=np.repeat(None,len(iord_proc_list))
        flux_Earth_all = deepcopy(flux_Earth_all_in)   
        cond_def_all = deepcopy(cond_def_all_in)    
    else:
        flux_mast = deepcopy(flux_mast_in)
        flux_Earth_all=np.tile(None,[1,len(iord_proc_list)])
        cond_def_all=np.tile(None,[1,len(iord_proc_list)])
    common_args = (inst,roll_win,smooth_win,locmax_win,par_stretch,contin_pinR,min_edge_ord)
    if nthreads>1:mean_flux_mast[iord_proc_list],dic_sav_loc = para_sub_spectral_cont(sub_spectral_cont,nthreads,len(iord_proc_list),[iord_proc_list,cond_def_mast,cen_bins_mast,flux_mast,edge_bins_mast,flux_Earth_all,cond_def_all],common_args)                           
    else:mean_flux_mast[iord_proc_list],dic_sav_loc = sub_spectral_cont(iord_proc_list,cond_def_mast,cen_bins_mast,flux_mast,edge_bins_mast,flux_Earth_all,cond_def_all,*common_args)  
    dic_sav.update(dic_sav_loc)

    #----------------------------------------------------
    #Define continuum function for current order
    #    - performed outside of the parallelized loop because for some reason it raises issues with 'cont_func_dic'
    cont_func_dic={}
    for isub_ord,iord in enumerate(iord_proc_list):
        cont_func_dic[iord] = interp1d(dic_sav[iord]['wav_max'],dic_sav[iord]['flux_max'],fill_value='extrapolate')  

    return mean_flux_mast,cont_func_dic,dic_sav



def sub_spectral_cont(iord_proc_list,cond_def_mast,cen_bins_mast,flux_mast,edge_bins_mast,flux_Earth_all,cond_def_all,inst,roll_win,smooth_win,locmax_win,par_stretch,contin_pinR,min_edge_ord):
    r"""**Stellar continuum: direct calculation.**
    
    Estimates stellar continuum from master spectrum.

    Args:
        TBD
    
    Returns:
        TBD
    
    """  
    mean_flux_mast = np.zeros(len(iord_proc_list),dtype=float)*np.nan
    dic_sav={}
    low_edge_bins = edge_bins_mast[:,0:-1]
    high_edge_bins = edge_bins_mast[:,1::]
    dcen_bins = high_edge_bins - low_edge_bins  
    for isub_ord,iord in enumerate(iord_proc_list):

        #Master tables over order
        cond_def = cond_def_mast[isub_ord]
        cen_bins_ord = cen_bins_mast[isub_ord,cond_def]
        low_edge_bins_ord = low_edge_bins[isub_ord][cond_def]
        high_edge_bins_ord = high_edge_bins[isub_ord][cond_def]
        dcen_bins_ord = dcen_bins[isub_ord][cond_def]
        
        #Calculate master from median flux if undefined
        if flux_mast[isub_ord] is None:
            flux_mast_ord = np.zeros(np.sum(cond_def))*np.nan                 
            for ipix_sub,ipix in enumerate(np_where1D(cond_def)):flux_mast_ord[ipix_sub] = np.median(flux_Earth_all[cond_def_all[:,isub_ord,ipix],isub_ord,ipix])
        else:
            flux_mast_ord = flux_mast[isub_ord,cond_def]
       
        #Mean master flux in order
        mean_flux_mast[isub_ord] = np.sum(flux_mast_ord*dcen_bins_ord)/np.sum(dcen_bins_ord)

        #Mean bin size over order (A)                
        dwav_mean = np.mean(dcen_bins_ord)

        #-------------------------------------------------------
        #Preliminary peak exclusion
        
        #Half-integer range over which peaks are excluded
        #    - exclusion range set to 3*fwhm of the inst
        fwhm_wav = calc_FWHM_inst(inst,np.mean(cen_bins_ord))  
        hn_exc = int(np.round(0.5*(3*fwhm_wav /dwav_mean)))

        #Apply successive sigma-clipping to remove spurious peaks
        #    - this step should not be necessary if the cosmic correction module was applied
        pd_spectre = pd.DataFrame(flux_mast_ord)
        n_roll = int(np.round(roll_win/dwav_mean))
        for iteration in range(5): 
            
            #Envelopes of the spectrum at different quantile
            maxi_roll_fast = np.ravel(pd_spectre.rolling(5*n_roll,min_periods=1,center=True).quantile(0.99))
            Q3_fast = np.ravel(pd_spectre.rolling(n_roll,min_periods=1,center=True).quantile(0.75)) 
            median = np.ravel(pd_spectre.rolling(n_roll,min_periods=1,center=True).quantile(0.50))

            #Condition of peak removal
            sup_fast = Q3_fast+3*(Q3_fast-median)  
            mask = (flux_mast_ord>sup_fast) & (flux_mast_ord>maxi_roll_fast)
            
            #Replace peaks and the chosen range around them
            mask_range = deepcopy(mask)
            for j in range(1,hn_exc):mask_range  |= np.roll(mask,-j) | np.roll(mask,j) 
            flux_mast_ord[mask_range] = median[mask_range]

        #-------------------------------------------------------
        #Smoothing
        #    - using Savitzky-Golay filter   
        n_smooth_win = int(smooth_win/dwav_mean)
        if not is_odd(n_smooth_win):n_smooth_win+=1
        smooth_flux = savgol_filter(flux_mast_ord, n_smooth_win, 3)

        #Suppress the smoothing of sharp peaks that create sinc-like wiggle
        spectre_backup = flux_mast_ord.copy()
        median = np.median(abs(spectre_backup-smooth_flux))
        IQ = np.percentile(abs(spectre_backup-smooth_flux),75) - median
        mask_out = np.where(abs(spectre_backup-smooth_flux)>(median+20*IQ))[0]
        mask_out = np.unique(mask_out+np.arange(-n_smooth_win,n_smooth_win+1,1)[:,np.newaxis])
        mask_out = mask_out[(mask_out>=0)&(mask_out<len(cen_bins_ord))].astype('int')
        smooth_flux[mask_out] = spectre_backup[mask_out] 

        #Set negative flux values to 0
        smooth_flux[smooth_flux<0] = 0. 

        #-------------------------------------------------------
        #Identification of local maxima 
        hn_locmax = int(np.round(0.5*locmax_win/dwav_mean))
        smooth_flux_noedge = smooth_flux[hn_locmax:-hn_locmax]
        maxima = np.ones(len(smooth_flux_noedge))
        for k in range(1,hn_locmax):
            maxima *= 0.5*(1+np.sign(smooth_flux_noedge - smooth_flux[hn_locmax-k:-hn_locmax-k]))*0.5*(1+np.sign(smooth_flux_noedge - smooth_flux[hn_locmax+k:-hn_locmax+k]))
        index = np_where1D(maxima==1)+hn_locmax
        maxima_flux = smooth_flux[index]
        maxima_cen_bins = cen_bins_ord[index]

        #Dealing with edge issues
        if maxima_flux[0] < smooth_flux[0]:
            maxima_cen_bins = np.insert(maxima_cen_bins,0,cen_bins_ord[0])
            maxima_flux = np.insert(maxima_flux,0,smooth_flux[0])
        if maxima_flux[-1] < smooth_flux[-1]:
            maxima_cen_bins = np.hstack([maxima_cen_bins,cen_bins_ord[-1]])
            maxima_flux = np.hstack([maxima_flux,smooth_flux[-1]])

        #Removing cosmics
        pd_max_flux = pd.DataFrame(maxima_flux)
        median = np.ravel(pd_max_flux.rolling(10,center=True).quantile(0.50))
        IQ = np.ravel(pd_max_flux.rolling(10,center=True).quantile(0.75)) - median
        IQ[np.isnan(IQ)] = smooth_flux.max()
        median[np.isnan(median)] = smooth_flux.max()
        cond_nopeaks = (maxima_flux <= median + 20 * IQ)
        maxima_cen_bins = maxima_cen_bins[cond_nopeaks]
        maxima_flux = maxima_flux[cond_nopeaks]

        #-------------------------------------------------------
        #Stretching
        #    - Y-axis stretching value
        #    - to scale the x and y axis    
        normalisation = (np.max(flux_mast_ord) - np.min(flux_mast_ord))/(high_edge_bins_ord[-1]  - low_edge_bins_ord[0])
        maxima_flux = maxima_flux/par_stretch/normalisation

        #----------------------------------------------------
        #Rolling pin
        waves = maxima_cen_bins - maxima_cen_bins[:,np.newaxis]
        distance = np.sign(waves)*np.sqrt((waves)**2+(maxima_flux - maxima_flux[:,np.newaxis])**2)
        distance[distance<0] = 0

        #Chromatic radius, constant in velocity space
        radius = np.repeat(contin_pinR,len(maxima_cen_bins))*maxima_cen_bins/min_edge_ord
        radius[0] = radius[0]/1.5
        
        #Applying rolling pin
        numero = np.arange(len(distance)).astype('int')
        j = 0
        keep = [0]
        while (len(maxima_cen_bins)-j>3): 
            par_R = float(radius[j]) #take the radius from the penality law
            mask = (distance[j,:]>0)&(distance[j,:]<2.*par_R) #recompute the points closer than the diameter if radius changed with the penality      
            while np.sum(mask)==0:
                par_R *=1.5
                mask = (distance[j,:]>0)&(distance[j,:]<2.*par_R) #recompute the points closer than the diameter if radius changed with the penality      
            p1 = np.array([maxima_cen_bins[j],maxima_flux[j]]).T #vector of all the local maxima 
            p2 = np.array([maxima_cen_bins[mask],maxima_flux[mask]]).T #vector of all the maxima in the diameter zone
            delta = p2 - p1 # delta x delta y
            c = np.sqrt(delta[:,0]**2+delta[:,1]**2) # euclidian distance 
            h = np.sqrt(par_R**2-0.25*c**2) 
            cx = p1[0] + 0.5*delta[:,0] - h/c*delta[:,1] #x coordinate of the circles center
            cy = p1[1] + 0.5*delta[:,1] + h/c*delta[:,0] #y coordinates of the circles center
            cond1 = (cy-p1[1])>=0
            thetas = cond1*(-1*np.arccos((cx - p1[0])/par_R)+np.pi) + (1-1*cond1)*(-1*np.arcsin((cy - p1[1])/par_R) + np.pi)
            j2 = thetas.argmin()
            j = numero[mask][j2] #take the numero of the local maxima falling in the diameter zone
            keep.append(j)
        maxima_flux = maxima_flux[keep] #we only keep the local maxima with the rolling pin condition
        maxima_cen_bins = maxima_cen_bins[keep]

        #----------------------------------------------------
        #Removing outlying maxima
        diff_deri = abs(np.diff(np.diff(maxima_flux)/np.diff(maxima_cen_bins)))
        mask_out = diff_deri  > (np.percentile(diff_deri,99.5))
        mask_out = np.array([False]+mask_out.tolist()+[False])
        maxima_cen_bins = maxima_cen_bins[~mask_out]
        maxima_flux = maxima_flux[~mask_out]
        
        #Unstretch maxima spectrum
        norm_maxima_flux = maxima_flux*normalisation*par_stretch

        #Save
        dic_sav[iord]={'wav_max':maxima_cen_bins,'flux_max':norm_maxima_flux}  
        if flux_mast[isub_ord] is None:
            dic_sav[iord].update({'wav_mast':cen_bins_ord,'flux_mast':flux_mast_ord})       

    return mean_flux_mast,dic_sav


def para_sub_spectral_cont(func_input,nthreads,n_elem,y_inputs,common_args):  
    r"""**Multithreading routine for sub_spectral_cont().**

    Args:
        func_input (function): multi-threaded function
        nthreads (int): number of threads
        n_elem (int): number of elements to thread
        y_inputs (list): threadable function inputs 
        common_args (tuple): common function inputs
    
    Returns:
        y_output (None or specific to func_input): function outputs 
    
    """
    pool_proc = Pool(processes=nthreads)   #cannot be passed through lmfit
    ind_chunk_list=init_parallel_func(nthreads,n_elem)
    chunked_args=[tuple(y_inputs[i][ind_chunk[0]:ind_chunk[1]] for i in range(5))+(y_inputs[5][:,ind_chunk[0]:ind_chunk[1]],y_inputs[6][:,ind_chunk[0]:ind_chunk[1]],)+common_args for ind_chunk in ind_chunk_list]	
    all_results=tuple(tab for tab in pool_proc.starmap(func_input,chunked_args))	

    mean_flux_mast = np.concatenate(tuple(all_results[i][0] for i in range(nthreads)),axis=0)
    dic_sav = {}
    for dic_save_proc in tuple(all_results[i][1] for i in range(nthreads)):dic_sav.update(dic_save_proc)
    y_output=(mean_flux_mast,dic_sav) 
    
    pool_proc.close()
    pool_proc.join() 				
    return y_output













