from __future__ import print_function
import numpy as np
import powderday.config as cfg
import pdb

import astropy.units as u
import astropy.constants as constants
from astropy import cosmology as cosmo

import fsps
from datetime import datetime
from powderday.grid_construction import stars_coordinate_boost

from multiprocessing import Pool
from functools import partial
from itertools import repeat
from scipy.integrate import simps

from powderday.nebular_emission.cloudy_tools import calc_LogQ, cmdf, get_nearest,convert_metals
from powderday.analytics import logu_diagnostic,dump_emlines
from powderday.nebular_emission.cloudy_model import get_nebular

#this is required to keep the reg as a strong reference.  for some
#reason in the star_list.append in star_list_gen, reg otherwise gets
#garbage collected.
import gc
gc.set_threshold(0)

# Lazily initialize FSPS
sp = None

class Stars:
    def __init__(self,mass,metals,positions,age,sed_bin=[-1,-1,-1],lum=-1,fsps_zmet=20,all_metals=[-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1]):
        self.mass = mass
        self.metals = metals
        self.positions = positions
        self.age = age
        self.sed_bin = sed_bin
        self.lum = lum
        self.fsps_zmet = fsps_zmet
        self.all_metals = all_metals

    def info(self):
        return(self.mass,self.metals,self.positions,self.age,self.sed_bin,self.lum,self.fsps_zmet)

def star_list_gen(boost,dx,dy,dz,reg,ds):
    print ('[SED_gen/star_list_gen]: reading in stars particles for SPS calculation')
    mass = reg["starmasses"].value
    positions = reg["starcoordinates"].value
    age = reg["stellarages"].value
    nstars = len(reg["stellarages"].value)
    el = ['He', 'C', 'N', 'O', 'Ne', 'Mg', 'Si', 'S', 'Ca', 'Fe' ]

    try:
        metals = np.zeros((nstars,11))-10.0
        for i in range(11):
            if i == 0:
                el_str = ""
            else:
                el_str = "_"+el[i-1]
            metals[:, i] = reg["starmetals"+el_str].value
    except:
        metals = reg["starmetals"].value
    print ('number of new stars =',nstars)

    #calculate the fsps interpolated metallicity

    #if the metallicity has many fields, and not just global
    #metallicity then just extract the global metallicity
    if metals.ndim > 1:
        metals_tot = metals[:,0]
    else:
        metals_tot = metals

    print ('[SED_gen/star_list_gen:] Manually increasing the newstar metallicities by: ',cfg.par.Z_init)
    metals_tot += cfg.par.Z_init

    #ADVANCED FEATURE - if force_stellar_metallcities or force_stellar_ages are set, then we set to those values
    if cfg.par.FORCE_STELLAR_AGES:
        print ("[SED_GEN/stars_list_gen:]  FORCE_STELLAR_AGES is set to True: setting all stars to age: %e Gyr"%cfg.par.FORCE_STELLAR_AGES_VALUE)
        age = np.repeat(cfg.par.FORCE_STELLAR_AGES_VALUE,nstars)

    if cfg.par.FORCE_STELLAR_METALLICITIES:
        print ("[SED_GEN/stars_list_gen:]  FORCE_STELLAR_METALLICITIES is set to True: setting all stars to metallicity: %e "%cfg.par.FORCE_STELLAR_METALLICITIES_VALUE)
        metals_tot = np.repeat(cfg.par.FORCE_STELLAR_METALLICITIES_VALUE,nstars)

    zmet = fsps_metallicity_interpolate(metals_tot)
    #mwd(zmet,mass,'zmet_distribution.png')

    #print '[SED_gen/star_list_gen: ] fsps zmet codes:',zmet

    #create the stars_list full of Stars objects
    stars_list = []

    if metals.ndim > 1:
        for i in range(nstars):
            stars_list.append(Stars(mass[i],metals_tot[i],positions[i],age[i],fsps_zmet=zmet[i],all_metals = metals[i]))
    else:
        for i in range(nstars):
            stars_list.append(Stars(mass[i],metals_tot[i],positions[i],age[i],fsps_zmet=zmet[i]))

    #boost stellar positions to grid center
    print ('boosting new stars to coordinate center')
    stars_list = stars_coordinate_boost(stars_list,boost)

    #orig_stars_list_len = len(stars_list)

    #ASSIGN DISK AND BULGE STARS - note, if these don't exist, it will
    #just make empty lists

    bulgestars_list = []
    diskstars_list = []

    #in principle, we should just be able to do the following blocks
    #if the particle types exist. the issue is that different groups
    #use PartType2 and 3 as 'filler' particle types, so they may exist
    #even if they don't correspond to disk/bulge stars.

    if ds.cosmological_simulation == False:

        #Disk Stars

        if ("diskstarcoordinates") in ds.derived_field_list:

            disk_positions = reg[("diskstarcoordinates")].value
            disk_masses =  reg[("diskstarmasses")].value
            nstars_disk = len(disk_masses)

            #create the disk_list full of DiskStars objects
            for i in range(nstars_disk):
                diskstars_list.append(Stars(disk_masses[i],cfg.par.solar,disk_positions[i],cfg.par.disk_stars_age))

            print ('boosting disk stars to coordinate center')
            diskstars_list = stars_coordinate_boost(diskstars_list,boost)

        #orig_disk_stars_list_len = nstars_disk

        #Bulge Stars

        if ("bulgestarcoordinates") in ds.derived_field_list:
            bulge_positions = reg[("bulgestarcoordinates")].value
            bulge_masses =  reg[("bulgestarmasses")].value
            nstars_bulge = len(bulge_masses)

            #create the bulge_list full of BulgeStars objects

            for i in range(nstars_bulge):
                bulgestars_list.append(Stars(bulge_masses[i],cfg.par.solar,bulge_positions[i],cfg.par.bulge_stars_age))

            print ('boosting bulge stars to coordinate center')
            bulgestars_list = stars_coordinate_boost(bulgestars_list,boost)

    #EXPERIMENTAL FEATURES
    if cfg.par.SOURCES_IN_CENTER == True:
        for i in range(nstars):
            stars_list[i].positions[:] =  np.array([0,0,0])
        if ("bulgestarcoordinates") in ds.derived_field_list:
            for i in range(nstars_bulge):
                bulgestars_list[i].positions[:] =  np.array([0,0,0])
            for i in range(nstars_disk):
                diskstars_list[i].positions[:] = np.array([0,0,0])

    if cfg.par.SOURCES_RANDOM_POSITIONS == True:
        print ("================================")
        print ("SETTING SOURCES TO RANDOM POSITIONS")
        print ("================================")
        for i in range(nstars):
            xpos,ypos,zpos = np.random.uniform(-0.9*dx/2.,0.9*dx/2.),np.random.uniform(-0.9*dy/2.,0.9*dy/2.),np.random.uniform(-0.9*dz/2.,0.9*dz/2.)
            stars_list[i].positions[:] = np.array([xpos,ypos,zpos])

        if ("bulgestarcoordinates") in ds.derived_field_list:
            for i in range(nstars_bulge):
                xpos,ypos,zpos = np.random.uniform(-0.9*dx/2.,0.9*dx/2.),np.random.uniform(-0.9*dy/2.,0.9*dy/2.),np.random.uniform(-0.9*dz/2.,0.9*dz/2.)
                bulgestars_list[i].positions[:] = np.array([xpos,ypos,zpos])
            for i in range(nstars_disk):
                xpos,ypos,zpos = np.random.uniform(-0.9*dx/2.,0.9*dx/2.),np.random.uniform(-0.9*dy/2.,0.9*dy/2.),np.random.uniform(-0.9*dz/2.,0.9*dz/2.)
                diskstars_list[i].positions[:] = np.array([xpos,ypos,zpos])

    return stars_list,diskstars_list,bulgestars_list,reg

def allstars_sed_gen(stars_list,cosmoflag,sp):

    #NOTE this part is just for the gadget simulations - this will
    #eventually become obviated as it gets passed into a function to
    #populate the stars_list with objects as we start to feed in new
    #types of simulation results.

    nstars = len(stars_list)

    #get just the wavelength array
    sp.params["tage"] = stars_list[0].age
    sp.params["imf_type"] = cfg.par.imf_type
    sp.params["pagb"] = cfg.par.pagb
    sp.params["sfh"] = 0
    sp.params["zmet"] = stars_list[0].fsps_zmet
    sp.params["add_neb_emission"] = cfg.par.add_neb_emission
    sp.params["add_agb_dust_model"] = cfg.par.add_agb_dust_model
    sp.params['gas_logu'] = cfg.par.gas_logu
    if cfg.par.FORCE_gas_logz == False:
        sp.params['gas_logz'] = np.log10(stars_list[0].metals/cfg.par.solar)
    else:
        sp.params['gas_logz'] = cfg.par.gas_logz

        '''
    sp = fsps.StellarPopulation(tage=stars_list[0].age,imf_type=cfg.par.imf_type,pagb = cfg.par.pagb,sfh=0,zmet=stars_list[0].fsps_zmet,
                                add_neb_emission = cfg.par.add_neb_emission, add_agb_dust_model=cfg.par.add_agb_dust_model)
                                '''
    spec = sp.get_spectrum(tage=stars_list[0].age,zmet=stars_list[0].fsps_zmet)
    nu = 1.e8*constants.c.cgs.value/spec[0]
    nlam = len(nu)

    nprocesses = np.min([cfg.par.n_processes,len(stars_list)]) #the pool.map will barf if there are less star bins than process threads

    #initialize the process pool and build the chunks
    p = Pool(processes = nprocesses)
    nchunks = nprocesses

    chunk_start_indices = []
    chunk_start_indices.append(0) #the start index is obviously 0

    #this should just be int(nstars/nchunks) but in case nstars < nchunks, we need to ensure that this is at least  1
    delta_chunk_indices = np.max([int(nstars / nchunks),1])
    print ('delta_chunk_indices = ',delta_chunk_indices)

    for n in range(1,nchunks):
        chunk_start_indices.append(chunk_start_indices[n-1]+delta_chunk_indices)

    '''
    chunk_start_indices = list(np.fix(np.arange(0,nstars,np.fix(nstars/nchunks))))
    #because this can result in too many chunks sometimes given the number of processors:
    chunk_start_indices = chunk_start_indices[0:nchunks]
    '''
    print ('Entering Pool.map multiprocessing for Stellar SED generation')
    list_of_chunks = []
    for n in range(nchunks):
        stars_list_chunk = stars_list[chunk_start_indices[n]:chunk_start_indices[n]+delta_chunk_indices]
        #if we're on the last chunk, we might not have the full list included, so need to make sure that we have that here
        if n == nchunks-1:
            stars_list_chunk = stars_list[chunk_start_indices[n]::]

        list_of_chunks.append(stars_list_chunk)

    t1=datetime.now()
    chunk_sol = p.map(newstars_gen, [arg for arg in list_of_chunks])

    t2=datetime.now()
    print ('Execution time for SED generation in Pool.map multiprocessing = '+str(t2-t1))

    stellar_fnu = np.zeros([nstars,nlam])
    star_counter=0
    for i in range(nchunks):
        fnu_list = chunk_sol[i] #this is a list of the stellar_fnu's returned by that chunk
        for j in range(len(fnu_list)):
            stellar_fnu[star_counter,:] = fnu_list[j,:]
            star_counter+=1

    p.close()
    p.terminate()
    p.join()

    stellar_nu = nu

    if cosmoflag == False:

        #calculate the SED for disk stars; note, this gets calculated
        #whether or not disk stars actually exist.  if they don't exist,
        #bogus values for the disk age and metallicity are assigned based
        #on whatever par.disk_stars_age and metallicity are.  it's no big
        #deal since these SEDs don't end up getting added to the model in
        #source_creation.

        #note, even if there are no disk/bulge stars, these are still
        #created since they're completely based on input parameters in
        #parameters_master.  they just won't get used at a later point
        #as there will be no disk/bulge star positions to add them to.

        #dust_tesc is an absolute value (not relative to min star age) as the ages of these stars are input by the user

        # Load in the metallicity legend
        fsps_metals = np.loadtxt(cfg.par.metallicity_legend)

        sp.params["tage"] = cfg.par.disk_stars_age
        sp.params["imf_type"] = cfg.par.imf_type
        sp.params["imf1"] = cfg.par.imf1
        sp.params["imf2"] = cfg.par.imf2
        sp.params["imf3"] = cfg.par.imf3
        sp.params["pagb"] = cfg.par.pagb
        sp.params["sfh"] = 0
        sp.params["zmet"] = cfg.par.disk_stars_metals
        sp.params["add_neb_emission"] = cfg.par.add_neb_emission
        sp.params["add_agb_dust_model"] = cfg.par.add_agb_dust_model
        sp.params['gas_logu'] = cfg.par.gas_logu
        if cfg.par.FORCE_gas_logz == False:
            sp.params['gas_logz'] = np.log10(fsps_metals[cfg.par.disk_stars_metals]/cfg.par.solar)
        else:
            sp.params['gas_logz'] = cfg.par.gas_logz

        spec = sp.get_spectrum(tage=cfg.par.disk_stars_age,zmet=cfg.par.disk_stars_metals)
        disk_fnu = spec[1]

        #calculate the SED for bulge stars
        sp.params["tage"] = cfg.par.bulge_stars_age
        sp.params["imf_type"] = cfg.par.imf_type
        sp.params["imf1"] = cfg.par.imf1
        sp.params["imf2"] = cfg.par.imf2
        sp.params["imf3"] = cfg.par.imf3
        sp.params["pagb"] = cfg.par.pagb
        sp.params["sfh"] = 0
        sp.params["zmet"] = cfg.par.bulge_stars_metals
        sp.params["add_neb_emission"] = cfg.par.add_neb_emission
        sp.params["add_agb_dust_model"] = cfg.par.add_agb_dust_model
        sp.params['gas_logu'] = cfg.par.gas_logu
        if cfg.par.FORCE_gas_logz == False:
            sp.params['gas_logz'] = np.log10(fsps_metals[cfg.par.bulge_stars_metals]/cfg.par.solar)
        else:
            sp.params['gas_logz'] = cfg.par.gas_logz

        spec = sp.get_spectrum(tage=cfg.par.bulge_stars_age,zmet=cfg.par.bulge_stars_metals)
        bulge_fnu = spec[1]

    else: #we have a cosmological simulation

        disk_fnu = []
        bulge_fnu = []

    total_lum_in_sed_gen = 0.
    for i in range(stellar_fnu.shape[0]):
        total_lum_in_sed_gen += np.absolute(np.trapz(stellar_fnu[i,:],x=nu))

    print ('[SED_gen: ] total_lum_in_sed_gen = ',total_lum_in_sed_gen)

    #return positions,disk_positions,bulge_positions,mass,stellar_nu,stellar_fnu,disk_masses,disk_fnu,bulge_masses,bulge_fnu
    return stellar_nu,stellar_fnu,disk_fnu,bulge_fnu

def newstars_gen(stars_list):
    global sp
    if sp is None:
        sp = fsps.StellarPopulation()

    #the newstars (particle type 4; so, for cosmological runs, this is all
    #stars) are calculated in a separate function with just one argument so that it is can be fed
    #into pool.map for multithreading.
    #sp = fsps.StellarPopulation()
    sp.params["tage"] = stars_list[0].age
    sp.params["imf_type"] = cfg.par.imf_type
    sp.params["pagb"] = cfg.par.pagb
    sp.params["sfh"] = 0
    sp.params["zmet"] = stars_list[0].fsps_zmet
    sp.params["add_neb_emission"] = False
    sp.params["add_agb_dust_model"] = cfg.par.add_agb_dust_model

    #first figure out how many wavelengths there are

    spec = sp.get_spectrum(tage=stars_list[0].age,zmet=stars_list[0].fsps_zmet)
    nu = 1.e8*constants.c.cgs.value/spec[0]

    nlam = len(nu)

    stellar_nu = np.zeros([nlam])
    stellar_fnu = np.zeros([len(stars_list),nlam])

    minage = 13 #Gyr
    for i in range(len(stars_list)):
        if stars_list[i].age < minage:
            minage = stars_list[i].age

    tesc_age = np.log10((minage+cfg.par.birth_cloud_clearing_age)*1.e9)

    # Get the number of ionizing photons from SED

    #calculate the SEDs for new stars
    for i in range(len(stars_list)):

        sp.params["tage"] = stars_list[i].age
        sp.params["imf_type"] = cfg.par.imf_type
        sp.params["imf1"] = cfg.par.imf1
        sp.params["imf2"] = cfg.par.imf2
        sp.params["imf3"] = cfg.par.imf3
        sp.params["pagb"] = cfg.par.pagb
        sp.params["sfh"] = 0
        sp.params["zmet"] = stars_list[i].fsps_zmet
        sp.params["add_neb_emission"] = False
        sp.params["add_agb_dust_model"] = cfg.par.add_agb_dust_model

        if cfg.par.CF_on == True:
            sp.params["dust_type"] = 0
            sp.params["dust1"] = 1
            sp.params["dust2"] = 0
            sp.params["dust_tesc"] = tesc_age

        spec_noneb = sp.get_spectrum(tage=stars_list[i].age,zmet=stars_list[i].fsps_zmet)
        f = spec_noneb[1]

        pagb = cfg.par.add_pagb_stars and cfg.par.PAGB_min_age <= stars_list[i].age <= cfg.par.PAGB_max_age
        young_star = cfg.par.add_young_stars and cfg.par.HII_min_age <= stars_list[i].age <= cfg.par.HII_max_age

        if (cfg.par.add_neb_emission or cfg.par.use_cmdf) and (young_star or pagb):

            # Cluster Mass Distribution Funtion is used only when the star particle's mass is gretaer than the maximum cluster mass and use_cmdf is True.

            if stars_list[i].mass/constants.M_sun.cgs.value > 10**cfg.par.cmdf_max_mass and cfg.par.use_cmdf:
                cluster_mass, num_clusters = cmdf(stars_list[i].mass/constants.M_sun.cgs.value,int(cfg.par.cmdf_bins),cfg.par.cmdf_min_mass,
                        cfg.par.cmdf_max_mass, cfg.par.cmdf_beta)

            else:
                cluster_mass = [np.log10(stars_list[i].mass/constants.M_sun.cgs.value)]
                num_clusters = [1]

            f = np.zeros(nlam)
            cloudy_nlam = len(np.genfromtxt(cfg.par.pd_source_dir + "/powderday/nebular_emission/data/refLines.dat", delimiter=','))
            line_em = np.zeros([cloudy_nlam])

            for j in range(len(cluster_mass)):
                num_HII_clusters = num_clusters[j]
                neb_file_output = cfg.par.NEB_DEBUG

                sp.params["add_neb_emission"] = False
                if cfg.par.add_neb_emission:
                    # id_val = 0, 1, 2 for young stars, Post-AGB star and AGNs respectively.
                    if young_star:
                        id_val = 0
                        Rinner_per_Rs = cfg.par.HII_Rinner_per_Rs
                        nh = cfg.par.HII_nh
                        escape_fraction  = cfg.par.HII_escape_fraction

                    elif pagb:
                        id_val = 1
                        Rinner_per_Rs = cfg.par.PAGB_Rinner_per_Rs
                        nh = cfg.par.PAGB_nh
                        escape_fraction  = cfg.par.PAGB_escape_fraction

                    if cfg.par.HII_alpha_enhance: #Setting Zstar based on Fe/H
                        Fe = stars_list[i].all_metals[-1]

                        # Gizmo metallicity structure, photospheric abundances from Asplund et al. 2009:
                        # Photospheric mass fraction of H = 0.7381
                        # Photospheric mass fraction of Fe = 1.31e-3

                        # Converting from mass fraction to atomic fraction of Fe
                        # Taking atmoic mass of H = 1.008u
                        # Taking atmoic mass of Fe = 55.845u
                        FeH = (Fe/0.7381)*(1.008/55.845)

                        # Solar atomic fraction of Fe. Calculated by substituting Fe = 1.31e-3 in the previous equation
                        FeH_sol = 3.22580645e-5

                        Logzsol = np.log10(FeH/FeH_sol)

                        sp1 = fsps.StellarPopulation(zcontinuous=1)
                        sp1.params["tage"] = stars_list[i].age
                        sp1.params["imf_type"] = cfg.par.imf_type
                        sp1.params["imf1"] = cfg.par.imf1
                        sp1.params["imf2"] = cfg.par.imf2
                        sp1.params["imf3"] = cfg.par.imf3
                        sp1.params["pagb"] = cfg.par.pagb
                        sp1.params["sfh"] = 0
                        sp1.params["zmet"] = stars_list[i].fsps_zmet
                        sp1.params["add_neb_emission"] = False
                        sp1.params["add_agb_dust_model"] = cfg.par.add_agb_dust_model
                        sp1.params["logzsol"] = Logzsol

                        if cfg.par.CF_on == True:
                            sp1.params["dust_type"] = 0
                            sp1.params["dust1"] = 1
                            sp1.params["dust2"] = 0
                            sp1.params["dust_tesc"] = tesc_age

                        spec = sp1.get_spectrum(tage=stars_list[i].age)

                    else:
                        spec = sp.get_spectrum(tage=stars_list[i].age,zmet=stars_list[i].fsps_zmet)

                    alpha = 2.5e-13 # Recombination Rate (assuming T = 10^4 K)

                    if cfg.par.FORCE_gas_logu[id_val]:
                        LogU = cfg.par.gas_logu[id_val]
                        LogQ = np.log10((10 ** (3*LogU))*(36*np.pi*(constants.c.cgs.value**3))/((alpha**2)*nh))
                        Rs = ((3*(10 ** LogQ))/(4*np.pi*(nh**2)*alpha))**(1./3.)

                    elif cfg.par.FORCE_logq[id_val]:
                        LogQ = cfg.par.source_logq[id_val]
                        Rs = ((3*(10 ** LogQ))/(4*np.pi*(nh**2)*alpha))**(1./3.)
                        LogU = np.log10((10**LogQ)/(4*np.pi*Rs*Rs*nh*constants.c.cgs.value))

                    else:
                        LogQ = calc_LogQ(1.e8*constants.c.cgs.value/spec[0], spec[1]*constants.L_sun.cgs.value
                                , efrac=escape_fraction, mstar=10**cluster_mass[j])
                        Rs = ((3*(10 ** LogQ))/(4*np.pi*(nh**2)*alpha))**(1./3.)
                        LogU = np.log10((10**LogQ)/(4*np.pi*Rs*Rs*nh*constants.c.cgs.value))+cfg.par.gas_logu_init[id_val]
                        LogQ = np.log10((10 ** (3*LogU))*(36*np.pi*(constants.c.cgs.value**3))/((alpha**2)*nh))
                        Rs = ((3*(10 ** LogQ))/(4*np.pi*(nh**2)*alpha))**(1./3.)

                    if cfg.par.FORCE_inner_radius[id_val]:
                        Rin = cfg.par.inner_radius[id_val]

                    else:
                        Rin = Rinner_per_Rs*Rs

                    if cfg.par.FORCE_gas_logz[id_val]:
                        LogZ = cfg.par.gas_logz[id_val]

                    else:
                        LogZ = np.log10(stars_list[i].metals/cfg.par.solar)

                    if neb_file_output:
                        if cfg.par.use_cloudy_tables:
                            Rin = 1.e19   # Rinner is fixed at 1.e19 cm for lookup tables

                        if cfg.par.FORCE_inner_radius[id_val]:
                            Rin = cfg.par.inner_radius[id_val]

                        LogU = np.log10((10**LogQ)/(4*np.pi*Rin*Rin*nh*constants.c.cgs.value))

                        logu_diagnostic(LogQ, LogU, LogZ, Rs, 10**cluster_mass[j], num_HII_clusters, stars_list[i].age, append=True)
                        neb_file_output = False

                    sp.params['gas_logu'] = LogU
                    sp.params['gas_logz'] = LogZ
                    sp.params["add_neb_emission"] = True
                    if cfg.par.use_cloudy_tables:
                        lam_neb, spec_neb = sp.get_spectrum(tage=stars_list[i].age, zmet=stars_list[i].fsps_zmet)
                        line_lum = sp.emline_luminosity
                        wave_line = sp.emline_wavelengths
                    else:
                        try:
                            # Calculating ionizing photons again but for 1 Msun in order to scale the output for FSPS
                            LogQ_1 = calc_LogQ(1.e8 * constants.c.cgs.value / spec[0], spec[1] * constants.L_sun.cgs.value,
                                    efrac=escape_fraction)
                            #LogQ_1 = LogQ_1 + cfg.par.gas_logu_init[id_val]

                            spec_neb, wave_line, line_lum = get_nebular(spec[0], spec[1], nh, LogQ, Rin, LogU, LogZ, LogQ_1, stars_list[i].all_metals,
                                                    Dust=False, abund=cfg.par.neb_abund[id_val], clean_up = cfg.par.cloudy_cleanup, index=id_val)
                        except ValueError as err:
                            # If the CLOUDY run crashes we switch to using lookup tables for young stars but throw an error for post-AGB stars.
                            if  young_star:
                                print ("WARNING: Switching to using lookup tables pre-packed with FSPS to calculate nebular emission for this particle.")
                                print ("WARNING: The emission line fluxes repoted may not be accurate if the particle lies outside the range of the lookup table paramters.")
                                lam_neb, spec_neb = sp.get_spectrum(tage=stars_list[i].age, zmet=stars_list[i].fsps_zmet)
                                line_lum = sp.emline_luminosity
                                wave_line = sp.emline_wavelengths
                            else:
                                print ("ERROR: Can't switch to using lookup tables.")
                                print ("ERROR: Please check the CLOUDY output file to figure out why the run was unsuccessful" )
                                raise ValueError('CLOUDY run was unsucessful')

                else:
                    lam_neb, spec_neb = sp.get_spectrum(tage=stars_list[i].age, zmet=stars_list[i].fsps_zmet)

                weight = num_HII_clusters*(10**cluster_mass[j])/(stars_list[i].mass/constants.M_sun.cgs.value)
                f = f + spec_neb*weight
                if cfg.par.add_neb_emission and cfg.par.dump_emlines:
                    line_em = line_em + line_lum*weight

            if cfg.par.add_neb_emission and cfg.par.dump_emlines:
                #the stellar population returns the calculation in units of Lsun/1 Msun: https://github.com/dfm/python-fsps/issues/117#issuecomment-546513619
                line_em = line_em * ((stars_list[i].mass*u.g).to(u.Msun).value) * (3.839e33)  # Units: ergs/s
                line_em = np.append(line_em, stars_list[i].age)
                dump_emlines(wave_line, line_em)

        stellar_nu[:] = 1.e8*constants.c.cgs.value/spec[0]
        stellar_fnu[i,:] = f

    return stellar_fnu

def get_gas_metals(ngas):
    # This function outputs the metallicity (total as well as all the 10 elements tracked by the simulation)
    # for all the gas particles
    reg = reg_gl
    el = ['He', 'C', 'N', 'O', 'Ne', 'Mg', 'Si', 'S', 'Ca', 'Fe']
    metals = np.zeros((ngas,11))-10.0
    try:
        for i in range(11):
            if i == 0:
                el_str = ""
            else:
                el_str = "_"+el[i-1]
            metals[:, i] = reg["gasmetals"+el_str].value
    except:
        metals[:,0] = reg["gasmetals"].value

    return metals

def get_nearest_gas_metals(all_gas_coordinates, particle_coordinates):

    all_gas_metals = get_gas_metals(len(all_gas_coordinates))

    # Getting N nearest gas particles to the AGN where N is defined by cfg.par.AGN_num_gas
    nearest_gas_dist, nearest_gas_id = get_nearest(all_gas_coordinates, particle_coordinates, num=cfg.par.AGN_num_gas)

    metals_avg = []

    # We take the distance weighted avearge of the metallicity of nearest N gas particles.
    # This is used as input to the CLOUDY model.
    for q in range(11):
        nearest_gas_metals = np.array(all_gas_metals[:,q][nearest_gas_id])
        metals_avg.append(np.sum(nearest_gas_metals*nearest_gas_dist)/np.sum(nearest_gas_dist))

    return metals_avg

def get_agn_seds(agn_ids, reg):

    print ('Starting AGN SED generation')

    t1 = datetime.now()
    nprocesses = np.min([cfg.par.n_processes,len(agn_ids)])
    p = Pool(processes = nprocesses)

    # Pre-calculating average metallicity and fluxes for all the BH particles
    nu = []
    fnu_in = []
    metals_avg = []
    for agn_id in agn_ids:
        fnu_in.append(reg["bhsed"][agn_id,:].in_units("Lsun").value/reg["bhnu"].in_units("Hz").value)
        nu.append(reg["bhnu"].in_units("Hz").value)
        all_gas_coordinates = reg["gascoordinates"].in_units('kpc').value
        agn_coordinates = reg["bhcoordinates"][agn_id].in_units('kpc').value
        metals_avg.append(get_nearest_gas_metals(all_gas_coordinates, agn_coordinates))

    z = zip(agn_ids, nu, fnu_in, metals_avg)
    fnu_out = p.starmap(agn_sed, z)
    fnu_out = np.atleast_2d(fnu_out)
    p.close()
    p.terminate()
    p.join()
    t2 = datetime.now()

    print ('Execution time for AGN SED generation = '+str(t2-t1))

    return fnu_out

def agn_sed(agn_id, nu, fnu, metals_avg):
    agn_id = int(agn_id)

    # Hopkins model returns the nu and fnu in reversed order.
    # Thus, we have to reverse it back to make it compatible.
    if cfg.par.BH_model != 'Nenkova':
        nu = nu[::-1]
        fnu = fnu[::-1]

    if cfg.par.add_neb_emission and cfg.par.add_AGN_neb:

        id_val = 2

        tot_metals = metals_avg[0]
        metals = metals_avg

        if cfg.par.FORCE_gas_logz[id_val]:
            LogZ = gas_logz[id_val]

        else:
            LogZ = tot_metals/cfg.par.solar

        Rin = cfg.par.inner_radius[id_val]

        if cfg.par.FORCE_gas_logu[id_val]:
            LogU = cfg.par.gas_logu[id_val]
            LogQ = np.log10((10**LogU)*(4*np.pi*Rin*Rin*cfg.par.AGN_nh*constants.c.cgs.value))

        elif cfg.par.FORCE_logq[id_val]:
            LogQ = cfg.par.source_logq[id_val]
            LogU = np.log10((10**LogQ)/(4*np.pi*Rin*Rin*cfg.par.AGN_nh*constants.c.cgs.value))

        else:
            LogQ = calc_LogQ(nu, fnu*constants.L_sun.cgs.value)
            LogU = np.log10((10**LogQ)/(4*np.pi*Rin*Rin*cfg.par.AGN_nh*constants.c.cgs.value))+cfg.par.gas_logu_init[id_val]
            LogQ = np.log10((10**LogU)*(4*np.pi*Rin*Rin*cfg.par.AGN_nh*constants.c.cgs.value))

        spec, wave_line, line_lum = get_nebular(1.e8 * constants.c.cgs.value / nu, fnu, cfg.par.AGN_nh, LogQ, Rin, LogU, LogZ, LogQ, metals,
                                        Dust=False, abund=cfg.par.neb_abund[id_val],clean_up = cfg.par.cloudy_cleanup, index=id_val)

        if cfg.par.dump_emlines:
        # The stellar population returns the calculation in units of Lsun
            line_em = line_lum * 3.839e33  # Units: ergs/s
            # The last column in dump_emlines is reserved for age of the star particle. For AGN we just set it to -1 as a place holder.
            line_em = np.append(line_em, -1.0)
            dump_emlines(wave_line, line_em)

    else:
        spec = fnu

    return spec

def fsps_metallicity_interpolate(metals):

    # takes a list of metallicities for star particles, and returns a
    # list of interpolated metallicities

    fsps_metals = np.loadtxt(cfg.par.metallicity_legend)
    nstars = len(metals)

    zmet = []

    for i in range(nstars):
        zmet.append(find_nearest_zmet(fsps_metals,metals[i]))

    return zmet

def find_nearest_zmet(array,value):
    # this is modified from the normal find_nearest in that it forces
    # the output to be 1 index higher than the true value since the
    # minimum zmet value fsps will take is 1 (not 0)

    idx = (np.abs(array-value)).argmin()

    return idx+1

import numpy,from scipy,import astropy,from astropy


# Contents in this file are specified for RASCIL

import ctypes
import os
import warnings

import numpy as np
from ctypes import c_double
from ctypes import c_int
from ctypes import c_float
from ctypes import c_void_p

c_int_p = ctypes.POINTER(c_int)
c_float_p = ctypes.POINTER(c_float)
c_double_p = ctypes.POINTER(c_double)

# TODO: See if there is a way to improve this so it is less hacky.
lib = None
# Try to load a local library directly.
lib_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),"libcurafft.so")
try:
    lib = ctypes.cdll.LoadLibrary(lib_path)
except Exception:
    raise RuntimeError('Failed to find curagridder library')

ms2dirty_1 = lib.ms2dirty_1
# the last two parameters have default value
ms2dirty_1.argtypes = [c_int, c_int, c_int, c_double, c_double, np.ctypeslib.ndpointer(np.double, flags='C'),
                     np.ctypeslib.ndpointer(np.complex128, flags='C'), np.ctypeslib.ndpointer(np.complex128, flags='C'), c_double, c_double, c_int]
ms2dirty_1.restype = c_int

ms2dirty_2 = lib.ms2dirty_2
ms2dirty_2.argtypes = [c_int, c_int, c_int, c_double, c_double, np.ctypeslib.ndpointer(np.double, flags='C'),
                     np.ctypeslib.ndpointer(np.complex128, flags='C'), np.ctypeslib.ndpointer(np.double, flags='C'), np.ctypeslib.ndpointer(np.complex128, flags='C'), c_double, c_double, c_int]
ms2dirty_2.restype = c_int

dirty2ms_1 = lib.dirty2ms_1
# the last two parameters have default value
dirty2ms_1.argtypes = [c_int, c_int, c_int, c_double, c_double, np.ctypeslib.ndpointer(np.double, flags='C'),
                     np.ctypeslib.ndpointer(np.complex128, flags='C'), np.ctypeslib.ndpointer(np.complex128, flags='C'), c_double, c_double, c_int]
dirty2ms_1.restype = c_int

dirty2ms_2 = lib.dirty2ms_2
dirty2ms_2.argtypes = [c_int, c_int, c_int, c_double, c_double, np.ctypeslib.ndpointer(np.double, flags='C'),
                     np.ctypeslib.ndpointer(np.complex128, flags='C'), np.ctypeslib.ndpointer(np.double, flags='C'), np.ctypeslib.ndpointer(np.complex128, flags='C'), c_double, c_double, c_int]
dirty2ms_2.restype = c_int

#----------------------------------------
# the interfaces below are idential to NIFTY
#-----------------------------------------

def ms2dirty(uvw, freq, ms, wgt, nxdirty, nydirty, rad_pix_x, rad_pix_y, nx, ny, epsilon, do_wstacking, *args):
    """
    Generate an image from visibility by non-uniform fourier transform
    Arguments:
        uvw - 3D coordinates, numpy array, shape - (nrow,3)
        freq - frequencies
        ms - visibility, shape - (nrow,)
        wgt - weight
        nxdirty, nydirty - image size
        deg_pix_ - degree per pixel
        epsilon - tolerance of relative error (expect, default 1e-6)
        do_wstacking - True, improved w stacking.

    Return:
        dirty image - shape-[nxdirty,nydirty]
    """
    nrow = uvw.shape[0]
    sigma = 2
    fov = rad_pix_x * nxdirty * 180 / np.pi
    dirty = np.zeros((nxdirty,nydirty),dtype=np.complex128)
    sign = -1
    # u = np.ctypeslib.as_ctypes(uvw[:,0])
    # v = np.ctypeslib.as_ctypes(uvw[:,1])
    # w = np.ctypeslib.as_ctypes(uvw[:,2])
    if(wgt is None):
        ms2dirty_1(nrow,nxdirty,nydirty,fov,freq[0],uvw
            ,ms,dirty,epsilon,sigma,sign)
    else:
        ms2dirty_2(nrow,nxdirty,nydirty,fov,freq[0],uvw
                ,ms,wgt,dirty,epsilon,sigma,sign)
    dirty = np.reshape(dirty,[nxdirty,nydirty])
    return dirty.real

def dirty2ms(uvw, freq, dirty, wgt, rad_pix_x, rad_pix_y, nx, ny, epsilon, do_wstacking, *args):
    """
    Generate Visibility from dirty image by non-uniform fourier transform
    Arguments:
        uvw - 3D coordinates, numpy array, shape - (nrow,3)
        freq - frequencies
        ms - visibility, shape - (nrow,)
        wgt - weight
        nxdirty, nydirty - image size
        fov - field of view
        epsilon - tolerance of relative error (expect, default 1e-6)
        sigma - upsampling factor for grid (default 1.25)
    Return:
        vis - shape-[M,]
    """
    nrow = uvw.shape[0]
    nxdirty = dirty.shape[0]
    nydirty = dirty.shape[1]
    sigma = 2
    fov = rad_pix_x * nxdirty * 180 / np.pi
    sign = -1
    ms = np.zeros((nrow,1),dtype=np.complex128)
    dirty1 = np.zeros(dirty.shape,dtype=np.complex128)
    dirty1.real = dirty

    if(wgt is None):
        dirty2ms_1(nrow,nxdirty,nydirty,fov,freq[0],uvw
            ,ms,dirty1,epsilon,sigma,sign)
    else:
        dirty2ms_2(nrow,nxdirty,nydirty,fov,freq[0],uvw
            ,ms,wgt,dirty1,epsilon,sigma,sign)
    return ms


# coding: utf-8
# Copyright (c) Tingzheng Hou.
# Distributed under the terms of the MIT License.

import numpy as np
from tqdm.notebook import tqdm
from MDAnalysis.analysis.distances import distance_array
from scipy.signal import savgol_filter
from mdgo.util import atom_vec

__author__ = "Tingzheng Hou"
__version__ = "1.0"
__maintainer__ = "Tingzheng Hou"
__email__ = "tingzheng_hou@berkeley.edu"
__date__ = "Feb 9, 2021"

def trajectory(nvt_run, li_atom, run_start, run_end, species, selection_dict, distance):
    dist_values = {}
    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    if species not in list(selection_dict):
        print("Invalid species selection")
        return None
    for ts in trj_analysis:
        selection = (
            "(" + selection_dict.get(species) + ") and (around " + str(distance) + " index " + str(li_atom.id - 1) + ")"
        )
        shell = nvt_run.select_atoms(selection, periodic=True)
        for atom in shell.atoms:
            if str(atom.id) not in dist_values:
                dist_values[str(atom.id)] = np.full(run_end - run_start, 100.0)
        time_count += 1
    time_count = 0
    for ts in trj_analysis:
        for atomid in dist_values.keys():
            dist = distance_array(ts[li_atom.id - 1], ts[(int(atomid) - 1)], ts.dimensions)
            dist_values[atomid][time_count] = dist
        time_count += 1
    return dist_values

def find_nearest(trj, time_step, distance, hopping_cutoff, smooth=51):
    """Returns an array of binding sites (unique on each timestep),
    the frequency of hopping between sites, and steps when each binding site
    exhibits the closest distance to the central atom.

    Args:
        trj (dict): A python dict of distances between central atom and selected atoms.
        time_step (int): The time step of the simulation.
        distance (int or float): Binding cutoff distance.
        hopping_cutoff: (int or float): Detaching cutoff distance.
        smooth (int): The length of the smooth filter window. Default to 51.
    """
    time_span = len(list(trj.values())[0])
    for kw in list(trj):
        trj[kw] = savgol_filter(trj.get(kw), smooth, 2)
    site_distance = [100 for _ in range(time_span)]
    sites = [0 for _ in range(time_span)]
    sites[0] = min(trj, key=lambda k: trj[k][0])
    site_distance[0] = trj.get(sites[0])[0]
    for time in range(1, time_span):
        if sites[time - 1] == 0:
            old_site_distance = 100
        else:
            old_site_distance = trj.get(sites[time - 1])[time]
        if old_site_distance > hopping_cutoff:
            new_site = min(trj, key=lambda k: trj[k][time])
            new_site_distance = trj.get(new_site)[time]
            if new_site_distance > distance:
                site_distance[time] = 100
            else:
                sites[time] = new_site
                site_distance[time] = new_site_distance
        else:
            sites[time] = sites[time - 1]
            site_distance[time] = old_site_distance
    sites = [int(i) for i in sites]
    sites_and_distance_array = np.array([[sites[i], site_distance[i]] for i in range(len(sites))])
    steps = []
    closest_step = 0
    previous_site = sites_and_distance_array[0][0]
    for i, step in enumerate(sites_and_distance_array):
        site = step[0]
        distance = step[1]
        if site == 0:
            pass
        else:
            if site == previous_site:
                if distance < sites_and_distance_array[closest_step][1]:
                    closest_step = i
                else:
                    pass
            else:
                steps.append(closest_step)
                closest_step = i
                previous_site = site
    if previous_site is not None:
        steps.append(closest_step)
    change = (np.diff([i for i in sites if i != 0]) != 0).sum()
    frequency = change / (time_span * time_step)
    return sites, frequency, steps

def find_in_n_out(trj, distance, hopping_cutoff, smooth=51, cool=20):
    """Returns two arrays of time step of hopping in and hopping out, respectively.

    Args:
        trj (dict): A python dict of distances between central atom and selected atoms.
        distance (int or float): Binding cutoff distance.
        hopping_cutoff: (int or float): Detaching cutoff distance.
        smooth (int): The length of the smooth filter window. Default to 51.
        cool (int): The cool down timesteps between hopping in and hopping out.
    """
    time_span = len(list(trj.values())[0])
    for kw in list(trj):
        trj[kw] = savgol_filter(trj.get(kw), smooth, 2)
    site_distance = [100 for _ in range(time_span)]
    sites = [0 for _ in range(time_span)]
    sites[0] = min(trj, key=lambda k: trj[k][0])
    site_distance[0] = trj.get(sites[0])[0]
    for time in range(1, time_span):
        if sites[time - 1] == 0:
            old_site_distance = 100
        else:
            old_site_distance = trj.get(sites[time - 1])[time]
        if old_site_distance > hopping_cutoff:
            new_site = min(trj, key=lambda k: trj[k][time])
            new_site_distance = trj.get(new_site)[time]
            if new_site_distance > distance:
                site_distance[time] = 100
            else:
                sites[time] = new_site
                site_distance[time] = new_site_distance
        else:
            sites[time] = sites[time - 1]
            site_distance[time] = old_site_distance
    sites = [int(i) for i in sites]

    last = sites[0]
    steps_in = list()
    steps_out = list()
    in_cool = cool
    out_cool = cool
    for i, s in enumerate(sites):
        if last == s:
            pass
        elif last == 0:
            in_cool = 0
            steps_in.append(i)
            if out_cool < cool:
                steps_out.pop()
        elif s == 0:
            out_cool = 0
            steps_out.append(i)
            if in_cool < cool:
                steps_in.pop()
        else:
            pass
        last = s
        in_cool += 1
        out_cool += 1
    return steps_in, steps_out

def check_contiguous_steps(nvt_run, li_atom, species_dict, select_dict, run_start, run_end, checkpoints, lag=20):
    """Returns two arrays of time step of hopping in and hopping out, respectively.

    Args:
        nvt_run (MDAnalysis.Universe): An Universe object of wrapped trajectory.
        li_atom (MDAnalysis.core.groups.Atom): the interested central atom object.
        species_dict (dict): Dict of Cutoff distance of neighbor for each species.
        select_dict (dict): A dictionary of selection language of atom species.
        run_start (int): Start time step.
        run_end (int): End time step.
        checkpoints (numpy.array): The time step of interest to check for contiguous steps
        lag (int): The range (+/- lag) of the contiguous steps
    """
    coord_num = {x: [[] for _ in range(lag * 2 + 1)] for x in species_dict.keys()}
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    has = False
    for i, ts in enumerate(trj_analysis):
        log = False
        checkpoint = None
        for j in checkpoints:
            if abs(i - j) <= lag:
                log = True
                has = True
                checkpoint = j
        if log:
            for kw in species_dict.keys():
                selection = (
                    "("
                    + select_dict[kw]
                    + ") and (around "
                    + str(species_dict[kw])
                    + " index "
                    + str(li_atom.id - 1)
                    + ")"
                )
                shell = nvt_run.select_atoms(selection, periodic=True)
                coord_num[kw][i - checkpoint + lag].append(len(shell))
    if has:
        for kw in coord_num:
            np_arrays = np.array([np.array(time).mean() for time in coord_num[kw]])
            coord_num[kw] = np_arrays
    return coord_num

def heat_map(
    nvt_run,
    li_atom,
    sites,
    dist_to_center,
    bind_atom_type,
    cartesian_by_ref,
    run_start,
    run_end,
):
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    coordinates = []
    for i, ts in enumerate(trj_analysis):
        if sites[i] == 0:
            pass
        else:
            center_atom = nvt_run.select_atoms("index " + str(sites[i] - 1))[0]
            selection = (
                "("
                + bind_atom_type
                + ") and "
                + "(around "
                + str(dist_to_center)
                + " index "
                + str(center_atom.id - 1)
                + ")"
            )
            bind_atoms = nvt_run.select_atoms(selection, periodic=True)
            distances = distance_array(ts[li_atom.id - 1], bind_atoms.positions, ts.dimensions)
            idx = np.argpartition(distances[0], 3)
            vertex_atoms = bind_atoms[idx[:3]]
            vector_li = atom_vec(li_atom, center_atom, ts.dimensions)
            vector_a = atom_vec(vertex_atoms[0], center_atom, ts.dimensions)
            vector_b = atom_vec(vertex_atoms[1], center_atom, ts.dimensions)
            vector_c = atom_vec(vertex_atoms[2], center_atom, ts.dimensions)
            basis_abc = np.transpose([vector_a, vector_b, vector_c])
            abc_li = np.linalg.solve(basis_abc, vector_li)
            unit_x = np.linalg.norm(
                cartesian_by_ref[0, 0] * vector_a
                + cartesian_by_ref[0, 1] * vector_b
                + cartesian_by_ref[0, 2] * vector_c
            )
            unit_y = np.linalg.norm(
                cartesian_by_ref[1, 0] * vector_a
                + cartesian_by_ref[1, 1] * vector_b
                + cartesian_by_ref[1, 2] * vector_c
            )
            unit_z = np.linalg.norm(
                cartesian_by_ref[2, 0] * vector_a
                + cartesian_by_ref[2, 1] * vector_b
                + cartesian_by_ref[2, 2] * vector_c
            )
            vector_x = cartesian_by_ref[0] / unit_x
            vector_y = cartesian_by_ref[1] / unit_y
            vector_z = cartesian_by_ref[2] / unit_z
            basis_xyz = np.transpose([vector_x, vector_y, vector_z])
            xyz_li = np.linalg.solve(basis_xyz, abc_li)
            coordinates.append(xyz_li)
    return np.array(coordinates)

def get_full_coords(coords, reflection=None, rotation=None, inversion=None, sample=None):
    coords_full = coords
    if reflection:
        for vec in reflection:
            coords_full = np.concatenate((coords, coords * vec), axis=0)
    if rotation:
        coords_copy = coords_full
        for mat in rotation:
            coords_rot = np.dot(coords_copy, mat)
            coords_full = np.concatenate((coords_full, coords_rot), axis=0)
    if inversion:
        coords_copy = coords_full
        for mat in inversion:
            coords_inv = np.dot(coords_copy, mat)
            coords_full = np.concatenate((coords_full, coords_inv), axis=0)
    if sample:
        index = np.random.choice(coords_full.shape[0], sample, replace=False)
        coords_full = coords_full[index]
    return coords_full

def cluster_coordinates(
    nvt_run,
    select_dict,
    run_start,
    run_end,
    species,
    distance,
    basis_vectors=None,
    cluster_center="center",
):
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    cluster_center = nvt_run.select_atoms(select_dict.get(cluster_center), periodic=True)[0]
    selection = (
        "("
        + " or ".join([s for s in species])
        + ") and (around "
        + str(distance)
        + " index "
        + str(cluster_center.id - 1)
        + ")"
    )
    print(selection)
    shell = nvt_run.select_atoms(selection, periodic=True)
    cluster = []
    for atom in shell:
        coord_list = []
        for ts in trj_analysis:
            coord_list.append(atom.position)
        cluster.append(np.mean(np.array(coord_list), axis=0))
    cluster = np.array(cluster)
    if basis_vectors:
        if len(basis_vectors) == 2:
            vec1 = basis_vectors[0]
            vec2 = basis_vectors[1]
            vec3 = np.cross(vec1, vec2)
            vec2 = np.cross(vec1, vec3)
        elif len(basis_vectors) == 3:
            vec1 = basis_vectors[0]
            vec2 = basis_vectors[1]
            vec3 = basis_vectors[2]
        else:
            raise ValueError("incorrect vector format")
        vec1 = vec1 / np.linalg.norm(vec1)
        vec2 = vec2 / np.linalg.norm(vec2)
        vec3 = vec3 / np.linalg.norm(vec3)
        basis_xyz = np.transpose([vec1, vec2, vec3])
        cluster_norm = np.linalg.solve(basis_xyz, cluster.T).T
        cluster_norm = cluster_norm - np.mean(cluster_norm, axis=0)
        return cluster_norm
    else:
        return cluster

def num_of_neighbor_one_li(
    nvt_run,
    li_atom,
    species_dict,
    select_dict,
    run_start,
    run_end,
    write=False,
    structure_code=None,
    write_freq=0,
    write_path=None,
    element_id_dict=None,
):

    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    cn_values = dict()
    species = list(species_dict.keys())
    for kw in species:
        if kw in select_dict.keys():
            cn_values[kw] = np.zeros(int(len(trj_analysis)))
        else:
            print("Invalid species selection")
            return None
    cn_values["total"] = np.zeros(int(len(trj_analysis)))
    for ts in trj_analysis:
        digit_of_species = len(species) - 1
        for kw in species:
            selection = (
                "("
                + select_dict.get(kw)
                + ") and (around "
                + str(species_dict.get(kw))
                + " index "
                + str(li_atom.id - 1)
                + ")"
            )
            shell = nvt_run.select_atoms(selection, periodic=True)
            # for each atom in shell, create/add to dictionary
            # (key = atom id, value = list of values for step function)
            for _ in shell.atoms:
                cn_values[kw][time_count] += 1
                cn_values["total"][time_count] += 10 ** digit_of_species
            digit_of_species = digit_of_species - 1
        if write and cn_values["total"][time_count] == structure_code:
            a = np.random.random()
            if a > 1 - write_freq:
                print("writing")
                selection_write = " or ".join(
                    "(same resid as ("
                    + select_dict.get(kw)
                    + " and around "
                    + str(species_dict.get(kw))
                    + " index "
                    + str(li_atom.id - 1)
                    + "))"
                    for kw in species
                )
                selection_write = "((" + selection_write + ")and not " + select_dict.get("cation") + ")"
                structure = nvt_run.select_atoms(selection_write, periodic=True)
                li_pos = ts[(int(li_atom.id) - 1)]
                path = write_path + str(li_atom.id) + "_" + str(int(ts.time)) + "_" + str(structure_code) + ".xyz"
                write_out(li_pos, structure, element_id_dict, path)
        time_count += 1
    return cn_values

def num_of_neighbor_one_li_simple(nvt_run, li_atom, species_dict, select_dict, run_start, run_end):

    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    species = list(species_dict.keys())[0]
    if species in select_dict.keys():
        cn_values = np.zeros(int(len(trj_analysis)))
    else:
        print("Invalid species selection")
        return None
    for ts in trj_analysis:
        selection = (
            "("
            + select_dict.get(species)
            + ") and (around "
            + str(species_dict.get(species))
            + " index "
            + str(li_atom.id - 1)
            + ")"
        )
        shell = nvt_run.select_atoms(selection, periodic=True)
        shell_len = len(shell)
        if shell_len == 0:
            cn_values[time_count] = 1
        elif shell_len == 1:
            selection_species = (
                "("
                + select_dict.get("cation")
                + " and around "
                + str(species_dict.get(species))
                + " index "
                + str(shell.atoms[0].id - 1)
                + ")"
            )
            shell_species = nvt_run.select_atoms(selection_species, periodic=True)
            shell_species_len = len(shell_species) - 1
            if shell_species_len == 0:
                cn_values[time_count] = 2
            else:
                cn_values[time_count] = 3
        else:
            cn_values[time_count] = 3
        time_count += 1
    cn_values = {"total": cn_values}
    return cn_values

def num_of_neighbor_one_li_simple_extra(nvt_run, li_atom, species, select_dict, distance, run_start, run_end):

    time_count = 0
    emc_angle = list()
    ec_angle = list()
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    if species in select_dict.keys():
        cn_values = np.zeros(int(len(trj_analysis)))
    else:
        print("Invalid species selection")
        return None
    for ts in trj_analysis:
        selection = (
            "(" + select_dict.get(species) + ") and (around " + str(distance) + " index " + str(li_atom.id - 1) + ")"
        )
        shell = nvt_run.select_atoms(selection, periodic=True)
        shell_len = len(shell)
        if shell_len == 0:
            cn_values[time_count] = 1
        elif shell_len == 1:
            selection_species = (
                "("
                + select_dict.get("cation")
                + " and around "
                + str(distance)
                + " index "
                + str(shell.atoms[0].id - 1)
                + ")"
            )
            shell_species = nvt_run.select_atoms(selection_species, periodic=True)
            shell_species_len = len(shell_species) - 1
            if shell_species_len == 0:
                cn_values[time_count] = 2
                li_pos = li_atom.position
                p_pos = shell.atoms[0].position
                ec_select = (
                    "(" + select_dict.get("EC") + ") and (around " + str(3) + " index " + str(li_atom.id - 1) + ")"
                )
                emc_select = (
                    "(" + select_dict.get("EMC") + ") and (around " + str(3) + " index " + str(li_atom.id - 1) + ")"
                )
                ec_group = nvt_run.select_atoms(ec_select, periodic=True)
                emc_group = nvt_run.select_atoms(emc_select, periodic=True)
                for atom in ec_group.atoms:
                    theta = angle(p_pos, li_pos, atom.position)
                    ec_angle.append(theta)
                for atom in emc_group.atoms:
                    theta = angle(p_pos, li_pos, atom.position)
                    emc_angle.append(theta)
            else:
                cn_values[time_count] = 3
        else:
            cn_values[time_count] = 3
        time_count += 1
    return cn_values, np.array(ec_angle), np.array(emc_angle)

def num_of_neighbor_one_li_simple_extra_two(nvt_run, li_atom, species_list, select_dict, distances, run_start, run_end):
    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    cip_step = list()
    ssip_step = list()
    agg_step = list()
    cn_values = dict()
    for kw in species_list:
        if kw in select_dict.keys():
            cn_values[kw] = np.zeros(int(len(trj_analysis)))
        else:
            print("Invalid species selection")
            return None
    cn_values["total"] = np.zeros(int(len(trj_analysis)))
    for ts in trj_analysis:
        digit_of_species = len(species_list) - 1
        for kw in species_list:
            selection = (
                "("
                + select_dict.get(kw)
                + ") and (around "
                + str(distances.get(kw))
                + " index "
                + str(li_atom.id - 1)
                + ")"
            )
            shell = nvt_run.select_atoms(selection, periodic=True)
            # for each atom in shell, create/add to dictionary
            # (key = atom id, value = list of values for step function)
            for _ in shell.atoms:
                cn_values[kw][time_count] += 1
                cn_values["total"][time_count] += 10 ** digit_of_species
            digit_of_species = digit_of_species - 1

        selection = (
            "("
            + select_dict.get("anion")
            + ") and (around "
            + str(distances.get("anion"))
            + " index "
            + str(li_atom.id - 1)
            + ")"
        )
        shell = nvt_run.select_atoms(selection, periodic=True)
        shell_len = len(shell)
        if shell_len == 0:
            ssip_step.append(time_count)
        elif shell_len == 1:
            selection_species = (
                "("
                + select_dict.get("cation")
                + " and around "
                + str(distances.get("anion"))
                + " index "
                + str(shell.atoms[0].id - 1)
                + ")"
            )
            shell_species = nvt_run.select_atoms(selection_species, periodic=True)
            shell_species_len = len(shell_species) - 1
            if shell_species_len == 0:
                cip_step.append(time_count)
            else:
                agg_step.append(time_count)
        else:
            agg_step.append(time_count)
        time_count += 1
    cn_ssip = dict()
    cn_cip = dict()
    cn_agg = dict()
    for kw in species_list:
        cn_ssip[kw] = np.mean(cn_values[kw][ssip_step])
        cn_cip[kw] = np.mean(cn_values[kw][cip_step])
        cn_agg[kw] = np.mean(cn_values[kw][agg_step])
    return cn_ssip, cn_cip, cn_agg

def angle(a, b, c):
    ba = a - b
    bc = c - b
    cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
    cosine_angle = np.clip(cosine_angle, -1.0, 1.0)
    angle_in_radian = np.arccos(cosine_angle)
    return np.degrees(angle_in_radian)

# Depth-first traversal
def num_of_neighbor_one_li_complex(nvt_run, li_atom, species, selection_dict, distance, run_start, run_end):
    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    cn_values = np.zeros((int(len(trj_analysis)), 4))
    for ts in trj_analysis:
        cation_list = [li_atom.id]
        anion_list = []
        shell = nvt_run.select_atoms(
            "(" + selection_dict.get(species) + " and around " + str(distance) + " index " + str(li_atom.id - 1) + ")",
            periodic=True,
        )
        for anion_1 in shell.atoms:
            if anion_1.resid not in anion_list:
                anion_list.append(anion_1.resid)
                cn_values[time_count][0] += 1
                shell_anion_1 = nvt_run.select_atoms(
                    "(type 17 and around 3 resid " + str(anion_1.resid) + ")",
                    periodic=True,
                )
                for cation_2 in shell_anion_1:
                    if cation_2.id not in cation_list:
                        cation_list.append(cation_2.id)
                        cn_values[time_count][1] += 1
                        shell_cation_2 = nvt_run.select_atoms(
                            "(type 15 and around 3 index " + str(cation_2.id - 1) + ")",
                            periodic=True,
                        )
                        for anion_3 in shell_cation_2.atoms:
                            if anion_3.resid not in anion_list:
                                anion_list.append(anion_3.resid)
                                cn_values[time_count][2] += 1
                                shell_anion_3 = nvt_run.select_atoms(
                                    "(type 17 and around 3 resid " + str(anion_3.resid) + ")",
                                    periodic=True,
                                )
                                for cation_4 in shell_anion_3:
                                    if cation_4.id not in cation_list:
                                        cation_list.append(cation_4.id)
                                        cn_values[time_count][3] += 1

def coord_shell_array(nvt_run, func, li_atoms, species_dict, select_dict, run_start, run_end):
    """
    Args:
        nvt_run: MDAnalysis Universe
        func: One of the neighbor statistical method (num_of_neighbor_one_li,
            num_of_neighbor_one_li_simple)
        li_atoms: Atom group of the Li atoms.
        species_dict (dict): A dict of coordination cutoff distance
            of the interested species.
        select_dict: A dictionary of species selection.
        run_start (int): Start time step.
        run_end (int): End time step.
    """
    num_array = func(nvt_run, li_atoms[0], species_dict, select_dict, run_start, run_end)
    for li in tqdm(li_atoms[1::]):
        this_li = func(nvt_run, li, species_dict, select_dict, run_start, run_end)
        for kw in num_array.keys():
            num_array[kw] = np.concatenate((num_array.get(kw), this_li.get(kw)), axis=0)
    return num_array

def write_out(li_pos, selection, element_id_dict, path):
    lines = list()
    lines.append(str(len(selection) + 1))
    lines.append("")
    lines.append("Li 0.0000000 0.0000000 0.0000000")
    box = selection.dimensions
    half_box = np.array([box[0], box[1], box[2]]) / 2
    for atom in selection:
        locs = list()
        for i in range(3):
            loc = atom.position[i] - li_pos[i]
            if loc > half_box[i]:
                loc = loc - box[i]
            elif loc < -half_box[i]:
                loc = loc + box[i]
            else:
                pass
            locs.append(loc)
        line = element_id_dict.get(int(atom.type)) + " " + " ".join(str(loc) for loc in locs)
        lines.append(line)
    with open(path, "w") as xyz_file:
        xyz_file.write("\n".join(lines))


#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
2/24/2021

This script takes outputs from a regional climate model (RCM) - e.g. MERRA,
MAR - for a particular site and puts that data into a pandas dataframe.

The output can be fed to RCMpkl_to_spin.py to generate a time series to force the CFM

YOU MAY HAVE TO EDIT THIS SCRIPT A LOT TO MAKE IT WORK WITH YOUR FILE STRUCTURE
AND WHAT CLIMATE FILES YOU HAVE.

And, for now there are little things you need to search out and change manually,
like the reference climate interval. Sorry!

@author: maxstev
'''

import netCDF4 as nc
import numpy as np
import scipy.io
import csv
import math
import sys
import decimal
import os
import sys
import matplotlib.pyplot as plt
from dateutil import rrule
from datetime import datetime, timedelta, date
import pandas as pd
import fnmatch
from scipy.spatial import cKDTree
from sklearn import datasets, linear_model
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.svm import SVR
import time
import xarray as xr
import glob
import hl_analytic as hla

def find_indices(points,lon,lat,tree=None):
    '''
    find the grid point nearest a given coordinate.
    '''
    if tree is None:
        # lon,lat = lon.T,lat.T
        lonlat = np.column_stack((lon.ravel(),lat.ravel()))
        tree = cKDTree(lonlat)
    dist,idx = tree.query(points,k=[1])
    ind = np.column_stack(np.unravel_index(idx,lon.shape))
    print(ind)
    for i,j in ind:
        ii=i
        jj=j

    return ii,jj #, [(i,j) for i,j in ind]

def read_netcdfs_merra(files, dim, ii, jj, vv, transform_func=None):
    '''
    Read merra files and concatenate into a pandas dataframe
    '''
    def process_one_path(path):
        with xr.open_dataset(path) as ds:
            # transform_func should do some sort of selection or
            # aggregation
            # if transform_func is not None:
            #     ds = transform_func(ds)
            # load all data from the transformed dataset, to ensure we can
            # use it after closing each original file
            ds = ds[vv].isel(lat=ii,lon=jj)
            ds.load()
            return ds
    datasets = [process_one_path(p) for p in files]
    combined = xr.concat(datasets, dim)
    df1 = combined.to_dataframe()
    return (df1.drop(labels=['lon','lat'],axis=1)).sort_index()

def read_netcdfs_mar(files, dim, ii, jj, vv):
    '''
    Read mar files and concatenate into a pandas dataframe
    '''
    def process_one_path(path):
        with xr.open_dataset(path) as ds:
            dsd = {}
            for v in vv:
                # print(v)
                if len(ds[v].dims)==4:
                    dsd[v] = ds[v][:,0,ii,jj].to_dataframe()
                else:
                    dsd[v] = ds[v][:,ii,jj].to_dataframe()
            df_list = [v for k,v in dsd.items()]
            df1 = pd.concat(df_list, axis=1)
            return df1[df1.columns.intersection(vv)]
    datasets = [process_one_path(p) for p in files]
    return (pd.concat(datasets)).sort_index()

def effectiveT(T):
    '''
    The Arrhenius mean temperature.
    '''
    Q   = -1 * 60.0e3
    R   = 8.314
    k   = np.exp(Q/(R*T))
    km  = np.mean(k)
    return Q/(R*np.log(km))

def getClimate(lat_int,lon_int,writer=True,datatype='MERRA',timeres='1D',melt=False,runtype='local',dsource = None):
    '''
    Load data from MERRA or MAR or whatever.
    Put it into a pandas dataframe, called df_CLIM. index must be datetimeindex for
    resampling.
    df_CLIM can have any number of columns: BDOT, TSKIN, SMELT, RAIN,
    SUBLIMATION (use capital letters. We use SMELT because melt is a pandas function)
    Hopefully this makes it easy to adapt for the different climate products.
    write df_CLIM into a pickle for future use.

    Reference for Summit, Greenland (my favorite test site):
    lat = 72.57972
    lon = -38.50454

    DYE-2 (my favorite wet test site):
    lat = 66.5
    lon = -46.2

    UNITS FOR MASS FLUXES IN THE DATAFRAMES ARE kg/m^2 PER TIME STEP SIZE IN
    THE DATA FRAME. e.g. if you have hourly data in the dataframe, the units
    for accumulation are kg/m^2/hour - the mass of precip that fell during that
    time interval.

    Parameters
    ----------
    lat_int: float
        the latitude of the site you want to build a climate history for
    lon_int: float
        the longitude of the site you want to build a climate history for
    writer: boolean
        Whether or not you want to write the pandas dataframe to a pickle
    datatype: string
        The type of RCM data you are using 'MERRA' or 'MAR' for now.
    melt: boolean
        Whether or not to put melt into the pandas dataframe
    Tinterp: 'mean', 'effective', or 'weighted'
        how to resample the temperature; mean is regular mean, 'effective' is
        Arrhenius mean; 'weighted' is accumulation-weighted mean
    runtype: 'local' or 'remote'
        Allows you easily switch between directory structures if you are testing
        code locally and running on a remote server
    dsource: 'ERA10k', 'ERA6k', or 'NCEP20k'
        MAR has several flavors; choose which one.

    Returns
    -------
    df_CLIM: pandas dataframe
        Dataframe containing the time series of each pertinent variable for
        the site, pulled from the RCM data. Index is a datetimeindex.

    '''

    if not writer:
        print('Files will not be written!')
    SPY = 365.25*24*3600

    todaystring = date.today().strftime("%Y%m%d")
    # write_out_dir = 'inputdata{}/'.format(todaystring) + datatype + 'input'
    write_out_dir = 'pickle'
    if writer:
        try:
            os.makedirs(write_out_dir)
        except:
            pass

    if datatype == 'MERRA':
        '''
        smb has dimensions of (time,lat,lon)
        smb has units of kg m^-2 s^-1 per day (because I sum the hourly values to get a value for each day, but do not divde by 24 after that) (pretty sure, at least!)
        temperature has dimensions of (time,lat,lon)
        temperature has units K

        '''

        ### Set directory to find climate files.
        if lat_int < 0: # Antarctica
            if runtype=='local':
                # ddir = 'PATH/TO/LOCAL/DATA/MERRA/Antarctica/Hourly'
                ddir = '/Volumes/Samsung_T1/MERRA/Antarctica/daily_melt'
            elif runtype=='remote':
                ddir = 'PATH/TO/REMOTE/DATA/MERRA/Antarctica/Hourly'
            elif runtype=='differentremote':
                ddir = 'PATH/TO/OTHER/REMOTE/DATA/CFM/MERRA/Antarctica/Hourly'

            # Adjust these as you see fit to set the Reference Climate Interval (RCI)
            spin_date_st = 1980
            spin_date_end = 2019

        else: # Greenland
            if runtype=='local':
                # ddir = 'PATH/TO/LOCAL/DATA/MERRA/Greenland/Hourly'
                # ddir = '/Volumes/Samsung_T1/MERRA/Greenland/daily_melt'
                ddir = '/Users/cdsteve2/RCMdata/MERRA2/Greenland/daily_melt'
            elif runtype=='remote':
                ddir = 'PATH/TO/REMOTE/DATA/MERRA/Greenland/Hourly'
            elif runtype == 'loki':
                ddir = '/home/maxstev/CFM_main/MERRA/Greenland/daily_melt'

            # Adjust these as you see fit to set the Reference Climate Interval (RCI)
            spin_date_st = 1980
            spin_date_end = 1995

        # input_datetimes = [dparser.parse((re.search(r'\d{8}',xx)).group()) for xx in ff] # this will extract the dates for each file
        # yy = np.array([float((re.search(r'\d{8}',xx)).group()[0:4]) for xx in glob.glob(ddir+'/TS/*.nc*')])
        # yrs = np.arange(min(yy),max(yy)+1)

        fn_ll = glob.glob(ddir + '/*.nc*')
        nc_ll = nc.Dataset(fn_ll[0],'r')
        lat_ll = nc_ll.variables['lat'][:]
        lon_ll = nc_ll.variables['lon'][:]
        ii, lat_val = min(enumerate(lat_ll), key=lambda x: abs(x[1]-lat_int))
        jj, lon_val = min(enumerate(lon_ll), key=lambda x: abs(x[1]-lon_int))
        nc_ll.close()
        print('lat_val: ', lat_val)
        print('lon_val: ', lon_val)

        if runtype=='local':
            # pickle_folder = '/PUT/PICKLES/HERE/MERRA/IDSpickle/pickle/'
            pickle_folder = 'example_pickle/'
        else:
            pickle_folder = 'IDS/pickle/'
        pickle_name = pickle_folder + 'MERRA2_CLIM_df_{}_{}.pkl'.format(lat_val,lon_val)
        if not os.path.exists(pickle_folder):
            os.makedirs(pickle_folder)

        if os.path.isfile(pickle_name):
            print('pickle found')
            writer = False
            loadnetcdf = False
            df_CLIM = pd.read_pickle(pickle_name)
            # try:
            #     df_BDOT = pd.DataFrame(df_CLIM['PRECTOT'])
            #     df_TS = pd.DataFrame(df_CLIM['TS'])
            #     df_CLIM.rename(columns={'PRECTOT':'BDOT','TS':'TSKIN'},inplace=True)
            # except Exception:
            #     df_BDOT = pd.DataFrame(xx['BDOT'])
            #     df_TS = pd.DataFrame(xx['TSKIN'])

            # if df_CLIM.BDOT.resample('1A').sum().mean()<1:
            #     df_CLIM.BDOT = df_CLIM.BDOT *3600 #get rid of seconds dimension - MERRA is hourly, so this gives precip per hour.

        else:
            vv=['TS','EVAP','SMELT','PRECTOT','PRECSNO']
            # flist_TS = glob.glob(ddir+'/TS/*.nc*')

            # df_TS = read_netcdfs_merra(flist_TS, dim='time',ii=ii,jj=jj,vv='TS')
            # df_TS.rename(columns={'TS':'TSKIN'},inplace=True)

            # flist_SMB = glob.glob(ddir+'/SMB/*.nc*')
            # df_BDOT = read_netcdfs_merra(flist_SMB, dim='time',ii=ii,jj=jj,vv='PRECTOT') # [kg m^-2 s^-1]
            # df_BDOT = (df_BDOT.rename(columns={'PRECTOT':'BDOT'}))*3600 # [kg m^-2 hour^-1] (this is amount of precip per MERRA time interval)

            df_merra = read_netcdfs_merra(fn_ll, dim='time',ii=ii,jj=jj,vv=vv)

            df_CLIM = df_merra
        # ACCVAR = 'PRECTOT'
        # TVAR = 'TS'

        # df_MELT = None
        # df_RAIN = None
        ####################
        #### end MERRA #####

    elif datatype == 'MAR':
        spin_date_st = 1980
        spin_date_end = 1995
        print('Using MAR')
        if lat_int < 0:
            print('no Antarctic MAR data')
            sys.exit()
        else:
            if runtype=='local':
                ddir = '/Volumes/Samsung_T1/MAR311/Greenland/Daily'

        if not dsource:
            dsource = 'ERA10k'
            print('using MAR ', dsource)

        if dsource == 'ERA10k':
            d2 = '/ERA_1958-2019-10km/'
            vv = ['ME','SF','ST2','RF','SU','TT']
        elif dsource == 'ERA6k':
            d2 = '/ERA_1979-2020-6km/'
            vv = ['ME','SF','ST2','RF','TT']
        elif dsource == 'NCEP20k':
            d2 = '/NCEP1_1948-2020_20km/'
            vv = ['ME','SF','ST2','RF','SU','TT']

        pickle_folder = ddir + '/pickles' + d2
        print(pickle_folder)
        if not os.path.exists(pickle_folder):
            os.makedirs(pickle_folder)
        # searchdir = ddir + d2 + '/*.nc'
        flist = glob.glob(ddir + d2 + '*.nc')
        rgr = nc.Dataset(flist[0],'r')
        lat = rgr['LAT'][:,:]
        lon = rgr['LON'][:,:]
        ii,jj = find_indices((lon_int,lat_int),lon,lat)
        lat_val = lat[ii,jj]
        lon_val = lon[ii,jj]
        print('lat_val: ', lat_val)
        print('lon_val: ', lon_val)
        rgr.close()

        PN = pickle_folder + 'MAR_{}_CLIM_df_{}_{}.pkl'.format(dsource,lat_val,lon_val)
        if os.path.isfile(PN):
            df_CLIM = pd.read_pickle(PN)
            print('Pickle found!')
            df_BDOT = pd.DataFrame(df_CLIM.BDOT)
            df_TS = pd.DataFrame(df_CLIM.TSKIN)

        # vv = ['ST2','SMB']
        else:
            df_CLIM = (read_netcdfs_mar(flist,'TIME',ii=ii,jj=jj,vv=vv))[str(spin_date_st):]

            if 'SMB' in df_CLIM.columns:
                df_BDOT = pd.DataFrame(df_CLIM['SMB']/1000*917).rename(columns = ['BDOT']) #put into units kg/m^2/day (i.e. per time resolution in the files))
                df_MELT = None
                df_RAIN = None
            else:
                if 'SU' in df_CLIM.columns:
                    df_BDOT = pd.DataFrame(((df_CLIM['SF']-df_CLIM['SU'])/1000*917),columns=['BDOT']) #put into units kg/m^2/day (i.e. per time resolution in the files))
                    df_CLIM['BDOT'] = df_BDOT.BDOT.values
                    df_CLIM.drop(['SF','SU'],axis=1,inplace=True)
                else:
                    df_BDOT = pd.DataFrame((df_CLIM['SF'])/1000*917).rename(columns={'SF':'BDOT'}) #put into units kg/m^2/day (i.e. per time resolution in the files))
                    df_CLIM['BDOT'] = df_BDOT.BDOT.values
                    df_CLIM.drop(['SF'],axis=1,inplace=True)
                df_CLIM['ME'] = df_CLIM['ME']/1000*917 #put into units kg/m^2/day (i.e. per time resolution in the files))
                df_CLIM['RF'] = df_CLIM['RF']/1000*917 #put into units kg/m^2/day (i.e. per time resolution in the files))
                # df_MELT = pd.DataFrame(df_CLIM['ME']/1000*917/3600).rename(columns={'ME':'MELT'}) #put into equivalent units to the merra data (kg/m^2/s)
                # df_RAIN = pd.DataFrame(df_CLIM['RF']/1000*917/3600).rename(columns={'RF':'RAIN'}) #put into equivalent units to the merra data (kg/m^2/s)
            df_TS = pd.DataFrame(df_CLIM[['ST2','TT']]).rename(columns = {'ST2':'TSKIN','TT':'T2M'}) + 273.15

            drn = {'ME':'SMELT','SU':'SUBLIMATION','SF':'BDOT','RF':'RAIN','ST2':'TSKIN','SMB':'BDOT','TT':'T2M'}
            df_CLIM.rename(mapper=drn,axis=1,inplace=True)
            df_CLIM.TSKIN = df_CLIM.TSKIN + 273.15
            df_CLIM.T2M = df_CLIM.T2M + 273.15
        ###############
        ### end MAR ###
        ###############

    elif datatype == 'RACMO':

        ### Set directory to find climate files.
        if lat_int < 0: # Antarctica
            if runtype=='local':
                ddir = '/Volumes/Samsung_T1/RACMO/Antarctica'
            elif runtype=='remote':
                ddir = 'PATH/TO/REMOTE/DATA/RACMO/Antarctica/Hourly'
            elif runtype=='differentremote':
                ddir = 'PATH/TO/OTHER/REMOTE/DATA/RACMO/Antarctica/Hourly'

            # Adjust these as you see fit to set the Reference Climate Interval (RCI)
            spin_date_st = 1980
            spin_date_end = 2019

        else: # Greenland
            if runtype=='local':
                # ddir = 'PATH/TO/LOCAL/DATA/MERRA/Greenland/Hourly'
                ddir = '/Volumes/Samsung_T1/RACMO/Greenland'
            elif runtype=='remote':
                ddir = 'PATH/TO/REMOTE/DATA/RACMO/Greenland/Hourly'
            elif runtype == 'differentremote':
                ddir = 'PATH/TO/OTHER/REMOTE/DATA/RACMO/Greenland/Hourly'

        spin_date_st = 1980
        spin_date_end = 1995

        flist = glob.glob(ddir + '/*1958-2016*.nc*')[0]
        rgr = nc.Dataset(flist[0],'r')
        lat = rgr['LAT'][:,:]
        lon = rgr['LON'][:,:]
        ii,jj = find_indices((lon_int,lat_int),lon,lat)
        lat_val = lat[ii,jj]
        lon_val = lon[ii,jj]
        print('lat_val: ', lat_val)
        print('lon_val: ', lon_val)
        rgr.close()

    if writer:
        if datatype =='MERRA':
            df_CLIM.to_pickle(pickle_folder + 'MERRA2_CLIM_df_{}_{}.pkl'.format(lat_val,lon_val))
        elif datatype == 'MAR':
            df_CLIM.to_pickle(pickle_folder + 'MAR_{}_CLIM_df_{}_{}.pkl'.format(dsource,lat_val,lon_val))

    return df_CLIM
    # return CD, stepsperyear, depth_S1, depth_S2, desired_depth

if __name__ == '__main__':
    tic = time.time()

    LLpair = sys.argv[1]
    nn = np.fromstring(LLpair,dtype =float, sep=' ')
    lat_int = nn[0]
    lon_int = nn[1]
    writer=True
    datatype='MERRA'
    runtype = 'local'

    df_CLIM = getClimate(lat_int,lon_int,writer = True, runtype = runtype)
    print(time.time()-tic)


# -*- coding: utf-8 -*-
"""
This script runs MCSED on objects and is organized in the same BLOCK method
that variable_house.py is.  A user should have read through variable_house.py
before interacting with main.py.  Heres an overview:

Block 1 - imports, ignore
Block 2 - for running multiple scripts. *INTERACTABLE*
Block 3 - set up how walkers obtain SEDs during random walk
Block 4 - defination of the stellar population/PRErams, user interactible
Block 5 - begin the for loop to run MCSED on each object
Block 6 - handle fluxes and errors, including error floors
Block 7 - begin actual random walk. pass lnprob() additional inputs here.
Block 8 - save results of walk in .npz
"""

VERSION = '1.4c' #see docs/change_log.txt

# = BLOCK 1 ========================================================= BLOCK 1 =
# Be polite and print a title.
print "\n ** MCSED ** \n"
print "written by Hunter Brooks, hzb5080@psu.edu"
print "Penn State, Dep't of Astronomy"
print "version:", VERSION, "\n"

# remove all old .pyc files to make sure no funny business is going on
print '... Cleaning old .pyc files'
import os
directory = os.listdir('.')
for filename in directory:
    if filename[-3:] == 'pyc':
        os.remove(filename)

# import everything that is needed
print '... Doing initial imports'
import time
start_time = time.time() # store start time
import numpy as np
import emcee_functions as mcF
import factory
import math_functions as mathF
# get home
import wheres_home
home = wheres_home.getHomeLocation()
import variable_house as vHouse
import sys

# = BLOCK 2 ========================================================= BLOCK 2 =
nScriptsRunning = 1 #up to 100, numbered with two digits each time
superIndexes = np.loadtxt(home+'data/active/superIndexes.dat')

# get name of this script, and which number it is
thisScriptName = sys.argv[0][-9:] # filename should be NNmain.py, eg 04main.py
thisScriptNumb = int(thisScriptName[0:2])
print thisScriptNumb
print

# split superIndexes into nScriptsRunning chunks
if (nScriptsRunning == 1):
    chunks = np.array_split(superIndexes, nScriptsRunning)[0]
else:
    chunks = np.array_split(superIndexes, nScriptsRunning)
print chunks

# = BLOCK 3 ========================================================= BLOCK 3 =
print '... Opening necesary MCSED data'
if vHouse.walker_SED_type == 'Interp':
    pSpaceGridSize = vHouse.pSpaceGridSize
    # being loading interpolation data... make a time stamp
    load_data_time = time.time()
    npzopendata = np.load(home+'data/pSpaces/size'+str(pSpaceGridSize)+'_sed.npz')
    sedGrid = npzopendata["sedGrid"]
    npzopendata.close()
    npzopendata = np.load(home+'data/pSpaces/size'+str(pSpaceGridSize)+'_other.npz')
    #valsGrid=npzopendata['valsGrid']  #BROKEN, see create_pSpace.py
    lambda_e=npzopendata['waves']
    stellarMassGrid=npzopendata['stellarMassGrid']
    npzopendata.close()
    import scipy.interpolate as spterp
    stellarmass_interp_f = spterp.RegularGridInterpolator(vHouse.paramRanges,
                                                  stellarMassGrid,
                                                  method='linear')

    walker_SED_notes = 'pSpace grid size: '+str(pSpaceGridSize)
    print '...      Took ', str(time.time() - load_data_time)
if vHouse.walker_SED_type == 'Direct':
    sp = vHouse.defSP()
    lambda_e = sp.wavelengths
    walker_SED_notes = 'Currently, direct does not work well.. be careful'

# = BLOCK 4 ========================================================= BLOCK 4 =
print '... Opening necesary observational data'
# find nFilters and load input data from mcsed/data/active
col_names = np.loadtxt(home+'data/active/filter_names.dat', dtype='str')
filter_lambdas = np.loadtxt(home+'data/active/filter_lambdas.dat')
nFilters = len(col_names)
# load input data, this is the data for the objects YOU want to match
superFluxes = np.loadtxt(home+'data/active/superFluxes.dat')
superErrors = np.loadtxt(home+'data/active/superErrors.dat')
superInfo = np.loadtxt(home+'data/active/superInfo.dat')
superRedshift = np.loadtxt(home+'data/active/superRedshift.dat')
# handle an irregularity that arises when there is only 1 input
if len(np.shape(superFluxes)) == 1: #there is only 1 row
    nSuperCols = np.shape(superFluxes)[0]
    # we need to reshape superdata because it looses this property when super
    # data contains only a single row.
    superFluxes = np.reshape(superFluxes, [1,nSuperCols])
    superErrors = np.reshape(superErrors, [1,nSuperCols])
    superInfo = np.reshape(superInfo, [1,np.shape(superInfo)[1]])

nObjects, nSuperCols = np.shape(superFluxes)
# make sure nSuperCols == nFilters
assert nFilters == nSuperCols, "All observations must have nFilters datum, even if some =0"
# make sure no negative errors of average fluxes
assert len(np.where(superErrors < 0 )[0]) == 0, "There was a negative average flux."

# = BLOCK 5 ========================================================= BLOCK 5 =
# being for loop to run MCSED on each object
print '... Starting runs, superIndexes: ', chunks[thisScriptNumb]

for run_i in chunks:
# for run_i in chunks[thisScriptNumb]: original
    # this keeps things simple, since run_i is just an index
    run_i = int(run_i)

    # check to see if this object is in /OUTPUTS
    if os.path.isfile(home+'OUTPUTS/output'+str(run_i)+'.npz') == True:
        print
        print '... '+str(run_i)+' has already been run.'
        continue

    # log starting time for this obj
    run_start_time = time.time()
    dataAvgFlux = np.zeros([nFilters])
    dataAvgFluxErr = np.zeros([nFilters])

    # = BLOCK 6 ===================================================== BLOCK 6 =
    z = superRedshift[run_i] #redshift for this object
    assert z >= 0.0, "Just a check to make sure Redshift is positive..."
    dataAvgFlux = np.zeros([nFilters])
    dataAvgFluxErr = np.zeros([nFilters])

    # feed all the data into the active variable (dataAvgFlux/Err)...
    for filt_i in range(0, nFilters):
        # pull data from superFluxes/Errors into working data
        dataAvgFlux[filt_i] = superFluxes[run_i, filt_i  ]
        dataAvgFluxErr[filt_i] = superErrors[run_i, filt_i  ]

    # ...then run the data/errors through filter machine
    dataAvgFlux, dataAvgFluxErr = mathF.filterMachine(    dataAvgFlux,
                                                          dataAvgFluxErr,
                                                          z,
                                                          filter_lambdas,
                                                          run_i)

    # make sure that dataAvgFlux is not entirely zeros!
    nNullObs = len(np.where(dataAvgFlux==0)[0])
    if nNullObs == len(dataAvgFlux):
        print
        print '... '+str(run_i)+" has only null observations (entire dataAvgFlux = 0)"
        continue

    # = BLOCK 7 ===================================================== BLOCK 7 =
    print
    print '... Customizing pSpace for super row index: ', run_i

    # this is used to move data to observed frame
    nu_obs = mathF.get_nu_obs(lambda_e, z)
    # create an R_nu array, dont need to do this each time but its fast and
    # for confusion's sake we just do it each time...
    R_nu, nFilters = factory.R_nu(home,
                                  nu_obs)
    # create D_nu array. D_nu is the space between nu's that corrispond to
    # sp.wavelengths
    D_nu = factory.D_nu(nFilters,
                        nu_obs)
    # create NU_obs array
    NU_obs = factory.NU_obs(nFilters,
                            nu_obs)
    if vHouse.walker_SED_type == 'Interp':
        # create flux_interp_f from luminosity SED grid
        fluxes, nu_obs = mathF.getFlux(sedGrid, lambda_e, z)
        # find the A value for each flux sed
        A_grid, A_compatible = factory.A_grid(fluxes,
                                              nFilters, NU_obs, D_nu, R_nu,
                                              dataAvgFlux, dataAvgFluxErr)
        # scale fluxes by each SEDs A value
        fluxes = fluxes*A_compatible
        # create flux interp function
        flux_interp_f = spterp.RegularGridInterpolator(vHouse.paramRanges,
                                                       fluxes,
                                                       method='linear')
        # create A interp function
        A_interp_f = spterp.RegularGridInterpolator(vHouse.paramRanges,
                                                    A_grid,
                                                    method='linear')
    if vHouse.walker_SED_type == 'Direct':
        flux_interp_f = sp
        A_interp_f = None
        stellarmass_interp_f = None

    print '... Conducting random walk'
    # BEGIN emcee RUN
    MCMCout = mcF.conductMCMC(      flux_interp_f,
                                    A_interp_f,
                                    stellarmass_interp_f,
                                    lambda_e,

                                    nFilters,
                                    z,
                                    NU_obs,
                                    D_nu,
                                    R_nu,

                                    dataAvgFlux,
                                    dataAvgFluxErr)

    print '...      Took ', str(time.time() - run_start_time)

    # = BLOCK 8 ===================================================== BLOCK 8 =
    # save data
    print '... Saving results of row index: ', run_i
    save_time = time.time()

    np.savez_compressed(home+'OUTPUTS/output'+str(run_i)+'.npz',

                        run_date=(time.strftime("%D :: %H:%M:%S")),
                        time_taken = time.time() - run_start_time,

                        dataAvgFlux = dataAvgFlux,
                        dataAvgFluxErr = dataAvgFluxErr,
                        z=z,

                        data_nu_obs = nu_obs,
                        D_nu = D_nu[0,:-1],
                        R_nu = R_nu,

                        nDim = vHouse.nDim,
                        nDram = vHouse.nDram,
                        nWalkers = vHouse.nWalkers,
                        nSteps = vHouse.nSteps,
                        nBurnInSteps = vHouse.nBurnInSteps,

                        paramNames = vHouse.nameL,
                        dramNames = vHouse.dRamNames,
                        paramRanges = vHouse.paramRanges,

                        walker_SED_type = vHouse.walker_SED_type,
                        walker_SED_notes = walker_SED_notes,

                        sampler_type = vHouse.sampler_type,
                        chain = MCMCout[0],
                        lnprobability = MCMCout[1],
                        auto_corr = MCMCout[2],
                        acc_frac = MCMCout[3],
                        flatchain = MCMCout[4],
                        flatlnprobability = MCMCout[5],
                        VERSION = VERSION)

    print '...      Took ', str(time.time() - save_time)


#!/usr/bin/python
'''
Program:
    This is a program for finding the minimum magnitude of source found in the image.
Usage:
    plot_d_mag.py [image_list]
Editor:
    Jacob975
20181219
#################################
update log

20181219 version alpha 1:
    1. The code works.
'''
from astropy.io import fits as pyfits
from astropy import wcs
from astropy.coordinates import SkyCoord
from astropy import units as u
import numpy as np
import time
from sys import argv
import TAT_env
from cataio_lib import get_catalog, get_distance, get_app_mag
from reduction_lib import get_rid_of_exotic
from matplotlib import pyplot as plt
import os
from photometry import take_data_within_duration
import collections

def show_cata_mag(mag, ra, dec, VERBOSE = 1):
    # Load the index of columes
    reduce_data = np.transpose(np.array([mag, ra, dec], dtype = float))
    reduce_data = reduce_data[mag.argsort()]
    reduce_data = reduce_data[~np.isnan(reduce_data[:,0])]
    # Pick 50 brightest stars from the data
    reduce_data = reduce_data[:20]
    world = reduce_data[:,1:]
    #--------------------------------------------------
    # Find the catalog magnitude.
    # Query data from vizier
    mag_delta_list = []
    filter_ = 'V'
    for i in xrange(len(reduce_data)):
        inst_mag = reduce_data[i,0]
        RA = float(world[i, 0])
        DEC = float(world[i, 1])
        failure, match_star = get_catalog(RA, DEC, TAT_env.URAT_1, TAT_env.index_URAT_1)
        if failure:
            continue
        failure, app_mag = get_app_mag(match_star, filter_)
        if failure:
            continue
        if np.isnan(inst_mag):
            continue
        mag_delta = app_mag - inst_mag
        if VERBOSE == 1: print "INST_MAG = {0}, CATA_MAG = {1}, delta = {2}".format(inst_mag, app_mag, mag_delta)
        mag_delta_list.append(mag_delta)
    # Find the average of delta_mag
    # Check if the number of source is enough or not.
    if len(mag_delta_list) == 0:
        print "No enough source found in catalogue for comparison"
        return 1
    mag_delta_list = get_rid_of_exotic(mag_delta_list)
    if len(mag_delta_list) < 3:
        print "No enough source found in catalogue for comparison"
        return 1
    # remove np.nan
    mag_delta_array = np.array(mag_delta_list)
    mag_delta_array = mag_delta_array[~np.isnan(mag_delta_array)]
    # Find the median of the delta of the magnitude, and apply the result on all sources.
    median_mag_delta = np.median(mag_delta_array)
    app_mag = mag + median_mag_delta
    return 0, app_mag

def pick_median(app_mag, err_app_mag):
    cata_mag = np.linspace(9, 20, 45)
    cata_mag = (cata_mag[1:] + cata_mag[:-1]) / 2.0
    lim_err_mag = np.zeros(len(cata_mag))
    for i, mag in enumerate(cata_mag):
        index = np.where(   (app_mag >= mag - 0.125 ) &
                            (app_mag <  mag + 0.125))
        selected_err_app_mag = err_app_mag[index]
        lim_err_mag[i] = np.median(selected_err_app_mag)
    index_min_err_mag = np.nanargmin(lim_err_mag)
    lim_err_mag[:index_min_err_mag] = lim_err_mag[index_min_err_mag]
    nans, x= nan_helper(lim_err_mag)
    lim_err_mag[nans]= np.interp(x(nans), x(~nans), lim_err_mag[~nans])
    return cata_mag, lim_err_mag

def nan_helper(y):
    """Helper to handle indices and logical indices of NaNs.

    Input:
        - y, 1d numpy array with possible NaNs
    Output:
        - nans, logical indices of NaNs
        - index, a function, with signature indices= index(logical_indices),
            to convert logical indices of NaNs to 'equivalent' indices
    Example:
        >>> # linear interpolation of NaNs
        >>> nans, x= nan_helper(y)
        >>> y[nans]= np.interp(x(nans), x(~nans), y[~nans])
    """
    return np.isnan(y), lambda z: z.nonzero()[0]

def observable_kepler(cata_mag, lim_err_mag, kepler_mag, kepler_depth):
    num_observable_sources = np.zeros(len(cata_mag))
    for i, mag in enumerate(cata_mag):
        index = np.where(   (kepler_mag >= mag - 0.125 ) &
                            (kepler_mag < mag + 0.125) &
                            (kepler_depth > lim_err_mag[i]))
        num_observable_sources[i] = len(index[0])
    return cata_mag, num_observable_sources
#--------------------------------------------
# main code
if __name__ == "__main__":
    # Measure time
    start_time = time.time()
    #----------------------------------------
    # Initialize
    if len(argv) != 4:
        print "Error!\n The number of arguments is wrong."
        print "Usage: plot_d_mag.py [start time] [end time]"
        exit(1)
    start_date = argv[1]
    end_date = argv[2]
    transit_table_name = argv[3]
    #----------------------------------------
    # Processing
    print ('---plot d mag ---')
    # Take the data in this duration
    data = take_data_within_duration(start_date, end_date)
    # Load the index
    index_BJD = TAT_env.obs_data_titles.index("BJD")
    index_INST_MAG = TAT_env.obs_data_titles.index('INST_MAG')
    index_E_INST_MAG = TAT_env.obs_data_titles.index('E_INST_MAG')
    index_EP_MAG = TAT_env.obs_data_titles.index('EP_MAG')
    index_E_EP_MAG = TAT_env.obs_data_titles.index('E_EP_MAG')
    index_RA = TAT_env.obs_data_titles.index('RA')
    index_DEC = TAT_env.obs_data_titles.index('`DEC`')
    # Find the first image
    all_bjd = data[:,index_BJD]
    bjd_s = [item for item, count in collections.Counter(all_bjd).items() if count > 1]
    first_bjd = np.amin(bjd_s)
    # Take and analyize the data of the first image.
    first_frame_data = data[data[:,index_BJD] == first_bjd]
    ep_mag = np.array(first_frame_data[:, index_EP_MAG], dtype = float)
    e_ep_mag = np.array(first_frame_data[:, index_E_EP_MAG], dtype = float)
    index_no_nan_in_ep_mag = np.where(~np.isnan(ep_mag) & ~np.isnan(e_ep_mag))
    ra = first_frame_data[:, index_RA]
    dec = first_frame_data[:, index_DEC]
    ep_mag = ep_mag[index_no_nan_in_ep_mag]
    e_ep_mag = e_ep_mag[index_no_nan_in_ep_mag]
    ra = ra[index_no_nan_in_ep_mag]
    dec = dec[index_no_nan_in_ep_mag]
    failure, app_mag = show_cata_mag(ep_mag, ra, dec)
    # Load transit table
    index_Vmag = 11
    index_depth = 10
    index_depth_measure = 9
    index_DEC = 5
    transit_table = np.loadtxt(transit_table_name, dtype = str, delimiter = '\t')
    transit_table = transit_table[1:]
    # Ignore all no-observed data
    transit_table = transit_table[transit_table[:,index_Vmag] != '']
    # Take data
    Vmag = np.array(transit_table[:,index_Vmag], dtype = float)
    depth = np.array(transit_table[:,index_depth], dtype = float)
    depth_measure = np.array(transit_table[:,index_depth_measure], dtype = str)
    DEC = np.array(transit_table[:, index_DEC], dtype = float)
    above_south30 = DEC > -30.

    # Take source above -30 dec
    Vmag  = Vmag[above_south30]
    depth = depth[above_south30]
    depth_measure = depth_measure[above_south30]

    depth_measure[depth_measure == ''] = '0.0'
    depth_measure = np.array(depth_measure, dtype = float)
    where_no_depth = depth == 0.0
    depth[where_no_depth] = depth_measure[where_no_depth]
    depth = depth / 100.
    # Calculate the histogram
    numbers, bin_edges = np.histogram(Vmag, bins = np.linspace(9, 20, 45))
    bins = (bin_edges[1:] + bin_edges[:-1]) / 2.0
    cata_mag, lim_err_mag = pick_median(app_mag, e_ep_mag)
    _, num_obs_source = observable_kepler(cata_mag, lim_err_mag, Vmag, depth)
    print np.sum(numbers)
    print np.sum(num_obs_source)
    # plot the histogram of magnitude
    fig, axs = plt.subplots(1, 1, figsize = (8,6))
    axs.set_title('The observable transits')
    axs.bar(bins, numbers, width = 0.25, color = 'b', label = 'Comfirm transits')
    axs.bar(cata_mag, num_obs_source, width = 0.20, color = 'r', label = 'Observable comfirm transits')
    axs.set_xlabel('V magnitude')
    axs.set_ylabel("# of sources")
    axs.set_xlim(9, 20)
    axs.grid()
    axs.legend()
    plt.savefig('{0}_{1}_hist_kepler.png'.format(start_date, end_date))
    #---------------------------------------
    # Measure time
    elapsed_time = time.time() - start_time
    print "Exiting Main Program, spending ", elapsed_time, "seconds."