import numpy as np
from .weighted_yield import SSP, lifetime_Argast, lifetime_Raiteri
from .imf import IMF
from .yields import SN2_feedback, AGB_feedback, SN1a_feedback, Hypernova_feedback

class SSP_wrap():
    '''
    This is the wrapper around the SSP function. It preloads the needed classes and calls all nucleosynthetic enrichment processes when the enrichment is calculated.
    '''
    def __init__(self, a):
        '''
        Upon initialization the default IMF, CC-SN yields, SN Ia yields and AGB yields is loaded.

        INPUT:

        a = Modelparameter class. So the default IMF etc are loaded. If we want other yield sets etc. loaded we need to specify that in paramter.py
        '''

        ## loading the IMF and the yieldsets prescribed in a (containing all the model parameters)
        basic_imf = IMF(a.mmin,a.mmax,a.mass_steps)
        getattr(basic_imf, a.imf_type_name)(a.imf_parameter)
        basic_sn2 = SN2_feedback()
        getattr(basic_sn2, a.yield_table_name_sn2)()
        basic_1a = SN1a_feedback()
        getattr(basic_1a, a.yield_table_name_1a)()
        basic_agb = AGB_feedback()
        getattr(basic_agb, a.yield_table_name_agb)()
        ## mixing of Nomoto CC-SN and HN yields
        if a.yield_table_name_sn2 == 'Nomoto2013':
            basic_hn = Hypernova_feedback()
            getattr(basic_hn, a.yield_table_name_hn)()
            for item in basic_sn2.metallicities:
                x = np.copy(basic_sn2.table[item])
                y = np.copy(basic_hn.table[item])
                for jtem in basic_hn.masses:
                    basic_sn2.table[item]['mass_in_remnants'][np.where(basic_sn2.table[item]['Mass']==jtem)] = a.sn2_to_hn * (x['mass_in_remnants'][np.where(x['Mass']==jtem)]) + (1-a.sn2_to_hn) * (y['mass_in_remnants'][np.where(y['Mass']==jtem)])
                    basic_sn2.table[item]['unprocessed_mass_in_winds'][np.where(basic_sn2.table[item]['Mass']==jtem)] = a.sn2_to_hn * (x['unprocessed_mass_in_winds'][np.where(x['Mass']==jtem)]) + (1-a.sn2_to_hn) * (y['unprocessed_mass_in_winds'][np.where(y['Mass']==jtem)])
                    hn_mass = []
                    sn_mass = []
                    for stem in basic_sn2.elements:
                        sn_mass.append(x[stem][np.where(x['Mass']==jtem)])
                        hn_mass.append(y[stem][np.where(y['Mass']==jtem)])
                        basic_sn2.table[item][stem][np.where(basic_sn2.table[item]['Mass']==jtem)]= a.sn2_to_hn * (x[stem][np.where(x['Mass']==jtem)]) + (1-a.sn2_to_hn) * (y[stem][np.where(y['Mass']==jtem)])
        ## to pass the information on to the feedback calculation
        self.a = a
        self.imf = basic_imf
        self.sn2 = basic_sn2
        self.sn1a = basic_1a
        self.agb = basic_agb

    def calculate_feedback(self, z, elements, element_fractions, time_steps):
        '''
        The feedback is calculated for the initializes SSP.

        INPUT:

        z = metallicity of the SSP in mass fraction (not normed to solar!)

        elements = which elements to follow

        element_fractions = the birth material of the SSP in the same order as 'elements'

        time_steps = the time-steps for which the enrichment of the SSP should be calculated (usually the time-steps until the end of the chempy simulation)
        '''
        basic_ssp = SSP(False, float(z), np.copy(self.imf.x), np.copy(self.imf.dm), np.copy(self.imf.dn), np.copy(time_steps), list(elements), str(self.a.stellar_lifetimes), str(self.a.interpolation_scheme), bool(self.a.only_net_yields_in_process_tables))
        basic_ssp.sn2_feedback(list(self.sn2.elements), dict(self.sn2.table), np.copy(self.sn2.metallicities), float(self.a.sn2mmin), float(self.a.sn2mmax),list(element_fractions))
        basic_ssp.agb_feedback(list(self.agb.elements), dict(self.agb.table), list(self.agb.metallicities), float(self.a.agbmmin), float(self.a.agbmmax),np.hstack(element_fractions))
        basic_ssp.sn1a_feedback(list(self.sn1a.elements), list(self.sn1a.metallicities), dict(self.sn1a.table), str(self.a.time_delay_functional_form), float(self.a.sn1ammin), float(self.a.sn1ammax), self.a.sn1a_parameter, float(self.a.total_mass), bool(self.a.stochastic_IMF))
        basic_ssp.bh_feedback(float(self.a.bhmmin),float(self.a.bhmmax),list(elements), np.hstack(element_fractions) , float(self.a.percentage_of_bh_mass))

        # exposing these tables to the outside wrapper
        self.table = basic_ssp.table
        self.sn2_table = basic_ssp.sn2_table
        self.agb_table = basic_ssp.agb_table
        self.sn1a_table = basic_ssp.sn1a_table
        self.bh_table = basic_ssp.bh_table
        self.inverse_imf = basic_ssp.inverse_imf

def initialise_stuff(a):
    '''
    Convenience function initialising the solar abundance, SFR and infall with the default values provided in parameter.py as a
    '''
    from .solar_abundance import solar_abundances
    from .sfr import SFR
    from .infall import INFALL

    basic_solar = solar_abundances()
    getattr(basic_solar, a.solar_abundance_name)()

    basic_sfr = SFR(a.start,a.end,a.time_steps)
    if a.basic_sfr_name == 'gamma_function':
        getattr(basic_sfr, a.basic_sfr_name)(S0 = a.S_0 * a.mass_factor,a_parameter = a.a_parameter, loc = a.sfr_beginning, scale = a.sfr_scale)
    elif a.basic_sfr_name == 'model_A':
        basic_sfr.model_A(a.mass_factor*a.S_0,a.t_0,a.t_1)
    elif a.basic_sfr_name == 'prescribed':
        basic_sfr.prescribed(a.mass_factor, a.name_of_file)
    elif a.basic_sfr_name == 'doubly_peaked':
        basic_sfr.doubly_peaked(S0 = a.mass_factor*a.S_0, peak_ratio = a.peak_ratio, decay = a.sfr_decay, t0 = a.sfr_t0, peak1t0 = a.peak1t0, peak1sigma = a.peak1sigma)

    basic_sfr.sfr = np.divide(basic_sfr.sfr,sum(basic_sfr.sfr))

    basic_infall = INFALL(np.copy(basic_sfr.t),np.copy(basic_sfr.sfr))
    if a.basic_infall_name == 'exponential':
        getattr(basic_infall, a.basic_infall_name)((a.infall_amplitude,a.tau_infall,a.infall_time_offset,a.c_infall,a.norm_infall))
    elif a.basic_infall_name == 'gamma_function':
        getattr(basic_infall, a.basic_infall_name)(mass_factor = a.norm_infall, a_parameter = a.infall_a_parameter, loc = a.infall_beginning, scale = a.infall_scale)
    elif a.basic_infall_name == 'sfr_related':
        getattr(basic_infall, a.basic_infall_name)()

    return basic_solar, basic_sfr, basic_infall

def Chempy(a):
    '''
    Chemical evolution run with the default parameters using the net yields.

    INPUT:

    a = ModelParameters() from parameter.py

    OUTPUT:

    cube = The ISM evolution class

    abundances = The abundances of the ISM
    '''
    from .infall import PRIMORDIAL_INFALL
    from .time_integration import ABUNDANCE_MATRIX
    from .making_abundances import mass_fraction_to_abundances
    from numpy.lib.recfunctions import append_fields
    basic_solar, basic_sfr, basic_infall = initialise_stuff(a)
    elements_to_trace = a.elements_to_trace
    basic_primordial = PRIMORDIAL_INFALL(list(elements_to_trace),np.copy(basic_solar.table))
    basic_primordial.primordial()

    # Needed a rescaling for the shortened sfr
    gas_reservoir_mass_factor = a.gas_reservoir_mass_factor / a.shortened_sfr_rescaling
    #sfr_factor_for_cosmic_accretion    = a.sfr_factor_for_cosmic_accretion / a.shortened_sfr_rescaling
    gas_at_start = a.gas_at_start / a.shortened_sfr_rescaling

    cube = ABUNDANCE_MATRIX(np.copy(basic_sfr.t),np.copy(basic_sfr.sfr),np.copy(basic_infall.infall),list(elements_to_trace),
        list(basic_primordial.symbols),list(basic_primordial.fractions),float(gas_at_start),list(basic_primordial.symbols),list(basic_primordial.fractions),
        float(gas_reservoir_mass_factor),float(a.outflow_feedback_fraction),bool(a.check_processes),float(a.starformation_efficiency),float(a.gas_power),
        float(a.sfr_factor_for_cosmic_accretion), list(basic_primordial.symbols), list(basic_primordial.fractions))

    basic_ssp = SSP_wrap(a)
    for i in range(len(basic_sfr.t)-1):
        j = len(basic_sfr.t)-i
        element_fractions = []
        for item in elements_to_trace:
            element_fractions.append(float(np.copy(cube.cube[item][max(i-1,0)]/cube.cube['gas'][max(i-1,0)])))## gas element fractions from one time step before
            if element_fractions[-1]<0:
                print('-ve Error')
                #raise Exception('-ve Error')
        metallicity = float(cube.cube['Z'][i])
        #print(metallicity)

        time_steps = np.copy(basic_sfr.t[:j])
        basic_ssp.calculate_feedback(float(metallicity), list(elements_to_trace), list(element_fractions), np.copy(time_steps))
        cube.advance_one_step(i+1,np.copy(basic_ssp.table),np.copy(basic_ssp.sn2_table),np.copy(basic_ssp.agb_table),np.copy(basic_ssp.sn1a_table),np.copy(basic_ssp.bh_table))
        if cube.cube['gas'][i] < 0:
            print(i, basic_sfr.t[i])
            print('gas became negative. returning -inf')
            return -np.inf, [0]
        if cube.gas_reservoir['gas'][i] < 0:
            print('gas_reservoir became negative. returning -inf')
            return -np.inf, [0]

    abundances,elements,numbers = mass_fraction_to_abundances(np.copy(cube.cube),np.copy(basic_solar.table))
    weights = cube.cube['sfr']
    abundances = append_fields(abundances,'weights',weights)
    abundances = append_fields(abundances,'time', cube.cube['time'])
    abundances = np.array(abundances)

    for element in elements:
        if element != 'Fe':
            try:
                abundances[element] -= abundances['Fe']
            except RuntimeWarning: # Remove error from first Fe abundance = -inf
                pass
    #TEST output
    #print('Chempy output')
    #print(abundances[:][-1])

    return cube, abundances

def Chempy_all_times(a):
    '''
    Chemical evolution run with the default parameters using the net yields.

    INPUT:

    a = ModelParameters() from parameter.py

    OUTPUT:

    cube = The ISM evolution class

    abundances = The abundances of the ISM
    '''
    from .infall import PRIMORDIAL_INFALL
    from .time_integration import ABUNDANCE_MATRIX
    from .making_abundances import mass_fraction_to_abundances
    from numpy.lib.recfunctions import append_fields
    basic_solar, basic_sfr, basic_infall = initialise_stuff(a)
    elements_to_trace = a.elements_to_trace
    basic_primordial = PRIMORDIAL_INFALL(list(elements_to_trace),np.copy(basic_solar.table))
    basic_primordial.primordial()
    gas_reservoir_mass_factor = a.gas_reservoir_mass_factor / a.shortened_sfr_rescaling
    gas_at_start = a.gas_at_start / a.shortened_sfr_rescaling # unlikely to be needed unless a.gas_at_start is non-zero

    cube = ABUNDANCE_MATRIX(np.copy(basic_sfr.t),np.copy(basic_sfr.sfr),np.copy(basic_infall.infall),list(elements_to_trace),list(basic_primordial.symbols),
        list(basic_primordial.fractions),float(gas_at_start),list(basic_primordial.symbols),list(basic_primordial.fractions),float(gas_reservoir_mass_factor),
        float(a.outflow_feedback_fraction),bool(a.check_processes),float(a.starformation_efficiency),float(a.gas_power),
        float(a.sfr_factor_for_cosmic_accretion),list(basic_primordial.symbols), list(basic_primordial.fractions))

    basic_ssp = SSP_wrap(a)

    for i in range(len(basic_sfr.t)-1):
        j = len(basic_sfr.t)-i
        element_fractions = []
        for item in elements_to_trace:
            element_fractions.append(float(np.copy(cube.cube[item][max(i-1,0)]/cube.cube['gas'][max(i-1,0)])))## gas element fractions from one time step before
            if element_fractions[-1]<0:
                print('-ve Error')
                #raise Exception('-ve Error')
        metallicity = float(cube.cube['Z'][i])
        #print(metallicity)

        time_steps = np.copy(basic_sfr.t[:j])
        basic_ssp.calculate_feedback(float(metallicity), list(elements_to_trace), list(element_fractions), np.copy(time_steps))
        cube.advance_one_step(i+1,np.copy(basic_ssp.table),np.copy(basic_ssp.sn2_table),np.copy(basic_ssp.agb_table),np.copy(basic_ssp.sn1a_table),np.copy(basic_ssp.bh_table))

        for item in elements_to_trace:
            if cube.cube[item][i]<0:
                print(i,item)
                print('element %s became negative. returning -inf'%item)
                return -np.inf,[0]
        if cube.cube['gas'][i] < 0:
            print(i, basic_sfr.t[i])
            print('gas became negative. returning -inf')
            return -np.inf, [0]
        if cube.gas_reservoir['gas'][i] < 0:
            print('gas_reservoir became negative. returning -inf')
            return -np.inf, [0]

    abundances,elements,numbers = mass_fraction_to_abundances(np.copy(cube.cube),np.copy(basic_solar.table))
    weights = cube.cube['sfr']
    abundances = append_fields(abundances,'weights',weights)
    abundances = append_fields(abundances,'time', cube.cube['time'])
    abundances = np.array(abundances)

    for element in elements:
        if element != 'Fe':
            filt = np.where(np.logical_not(np.isfinite(abundances['Fe'])))
            abundances[element][filt]=np.inf
            filt2 = np.where(np.isfinite(abundances['Fe']))
            abundances[element][filt2]-=abundances['Fe'][filt2]
    #TEST output
    #print('Chempy output')
    #print(abundances[:][-1])

    return cube, abundances

def Chempy_gross(a):
    '''
    Chemical evolution run with the default parameters but now using solar scaled material (testing the worse case when total yields provided).

    INPUT:

    a = ModelParameters() from parameter.py

    OUTPUT:

    cube = The ISM evolution class

    abundances = The abundances of the ISM
    '''
    from infall import PRIMORDIAL_INFALL
    from time_integration import ABUNDANCE_MATRIX
    from making_abundances import mass_fraction_to_abundances
    from numpy.lib.recfunctions import append_fields
    basic_solar, basic_sfr, basic_infall = initialise_stuff(a)
    elements_to_trace = a.elements_to_trace
    basic_primordial = PRIMORDIAL_INFALL(list(elements_to_trace),np.copy(basic_solar.table))
    basic_primordial.primordial(0)
    gas_reservoir_mass_factor = a.gas_reservoir_mass_factor / a.shortened_sfr_rescaling
    gas_at_start = a.gas_at_start / a.shortened_sfr_rescaling
    cube = ABUNDANCE_MATRIX(np.copy(basic_sfr.t),np.copy(basic_sfr.sfr),np.copy(basic_infall.infall),list(elements_to_trace),list(basic_primordial.symbols),
        list(basic_primordial.fractions),float(gas_at_start),list(basic_primordial.symbols),list(basic_primordial.fractions),float(gas_reservoir_mass_factor),
        float(a.outflow_feedback_fraction),bool(a.check_processes),float(a.starformation_efficiency),float(a.gas_power), float(a.sfr_factor_for_cosmic_accretion),
        list(basic_primordial.symbols), list(basic_primordial.fractions))
    basic_ssp = SSP_wrap(a)
    for i in range(len(basic_sfr.t)-1):
        j = len(basic_sfr.t)-i
        metallicity = float(cube.cube['Z'][i])
        solar_scaled_material = PRIMORDIAL_INFALL(list(elements_to_trace),np.copy(basic_solar.table))
        solar_scaled_material.solar(np.log10(metallicity/basic_solar.z))
        element_fractions = list(solar_scaled_material.fractions)
        for item in elements_to_trace:
            element_fractions.append(float(np.copy(cube.cube[item][max(i-1,0)]/cube.cube['gas'][max(i-1,0)])))## gas element fractions from one time step before
        time_steps = np.copy(basic_sfr.t[:j])
        basic_ssp.calculate_feedback(float(metallicity), list(elements_to_trace), list(element_fractions), np.copy(time_steps))
        cube.advance_one_step(i+1,np.copy(basic_ssp.table),np.copy(basic_ssp.sn2_table),np.copy(basic_ssp.agb_table),np.copy(basic_ssp.sn1a_table))
    abundances,elements,numbers = mass_fraction_to_abundances(np.copy(cube.cube),np.copy(basic_solar.table))
    weights = cube.cube['sfr']
    abundances = append_fields(abundances,'weights',weights)
    abundances = np.array(abundances)

    return cube, abundances

def multi_star_optimization():
    '''
    This function will optimize the parameters of all stars in a hierachical manner (similar to gibbs sampling)

    INPUT:

    a = will be loaded from parameter.py (prepare all variables there)

    OUTPUT:

    log_list = a list of intermediate results (so far only for debugging)
    '''
    import time
    import multiprocessing as mp
    from .optimization import minimizer_initial, minimizer_global, minimizer_local
    from .cem_function import global_optimization_error_returned
    from .parameter import ModelParameters

    # For testing
    import warnings
    warnings.filterwarnings("ignore")

    a = ModelParameters()
    print(a.stellar_identifier_list)
    start_time = time.time()

    log_list = []
    # I: Minimization for each star seperately
    # 1: for each star make initial conditions (each star needs other model parameters)
    parameter_list = []
    for item in a.stellar_identifier_list:
        parameter_list.append(item)
    # 2: call posterior_function_for_minimization with scipy.optimize.minimize in multiprocess for each star and recover the found parameters
    p = mp.Pool(len(parameter_list))
    t = p.map(minimizer_initial, parameter_list)
    p.close()
    p.join()
    result = np.vstack(t)

    log_list.append(np.copy(result))
    log_list.append('initial minimization')
    initial = time.time()
    print('first minimization for each star separately took: %2.f seconds' %(initial - start_time))

    # IV: repeat II and III until posterior does not change much
    result[:,:len(a.SSP_parameters)] = np.mean(result[:,:len(a.SSP_parameters)], axis = 0)
    posteriors = []
    counter = 0
    while True:
        counter += 1
        if len(posteriors) > 1:
            if np.abs(posteriors[-1] - posteriors[-2]) < a.gibbs_sampler_tolerance:
                break
            if len(posteriors) > a.gibbs_sampler_maxiter:
                break

        initial = time.time()
        # II: Global parameter minimization:
        # 1: only SSP parameters free. Use mean SSP parameter values and individual (but fixed ISM parameter values)
        changing_parameter = result[0,:len(a.SSP_parameters)]
        # 2: Call each star in multiprocess but only return the predictions
        # 3: Calculate the likelihood for each star and optimize the common model error (is all done within minimizer global, which is calling 'global optimization')
        x = minimizer_global(changing_parameter,  a.tol_minimization, a.maxiter_minimization, a.verbose, result)

        # 4: return global SSP parameters and common model error
        posterior, error_list, elements = global_optimization_error_returned(x, result)
        posteriors.append(posterior)
        print(posteriors)

        global_iteration1 = time.time()
        print('step %d global minimization took: %2.f seconds' %(counter, global_iteration1 - initial))

        # III: Local parameter minimization:
        # 1: Use fixed global parameters and fixed common errors make initial conditions
        result[:,:len(a.SSP_parameters)] = x

        log_list.append((np.copy(x),posterior))
        log_list.append('step %d global minimization' %(counter))

        p0_list = []
        parameter_list = []
        x_list = []
        error_list_mp = []
        element_list_mp = []

        for i,item in enumerate(a.stellar_identifier_list):
            parameter_list.append(item)
            p0_list.append(result[i,len(a.SSP_parameters):])
            x_list.append(x)
            error_list_mp.append(error_list)
            element_list_mp.append(elements)

        args = zip(p0_list,parameter_list,x_list,error_list_mp,element_list_mp)

        # 2: Minimize each star ISM parameters in multiprocess
        p = mp.Pool(len(parameter_list))
        t = p.map(minimizer_local, args)
        p.close()
        p.join()
        local_parameters = np.vstack(t)
        result[:,len(a.SSP_parameters):] = local_parameters

        log_list.append(np.copy(result))
        log_list.append('step %d local minimization' %(counter))
        local_iteration1 = time.time()
        print('step %d local minimization took: %2.f seconds' %(counter, local_iteration1 - global_iteration1))

    log_list.append(posteriors)
    print(log_list)

    # V: MCMC run
    ## reshape the result to have global parameters in the front and the local parameters following
    changing_parameter = list(result[0,:len(a.SSP_parameters)])
    for i in range(result.shape[0]):
        changing_parameter.append(list(result[i,len(a.SSP_parameters):]))
    changing_parameter = np.hstack(changing_parameter)
    ## jitter the parameters to initialise the chain (add a validation later, i.e. testing that the particular parameters yield a result)
    mcmc_multi(changing_parameter, error_list, elements)
    # 1: Free all parameters and optimize common error (SSP should be the same for all stars)
    # 2: Plug everything into emcee and sample the posterior
    return log_list

def mcmc(a):
    '''
    Convenience function to use the MCMC. A subdirectory mcmc/ will be created in the current directory and intermediate chains will be stored there.

    The MCMC will sample the volume of best posterior for the likelihood functions that are declared in parameter.py. Default is ['sol_norm','gas_reservoir','sn_ratio'] which corresponds to 'Sun+' from the paper.
    '''
    import time
    import os
    import multiprocessing as mp
    from .optimization import creating_chain, posterior_probability
    import emcee

    start1 = time.time()
    directory = 'mcmc/'
    if os.path.exists(directory):
        if a.verbose:
            print('%s already existed. Content might be overwritten' %(directory))
    else:
        os.makedirs(directory)

    a.check_processes = False
    a.number_of_models_overplotted = 1
    a.only_net_yields_in_process_tables = False
    a.testing_output = False
    a.summary_pdf = False
    a.nthreads = mp.cpu_count()
    if a.nthreads == 4:
        a.nthreads = 2

    chain = creating_chain(a,np.copy(a.p0))
    sampler = emcee.EnsembleSampler(a.nwalkers,a.ndim,posterior_probability,threads=a.nthreads, args = [a])
    pos,prob,state,blobs = sampler.run_mcmc(chain,a.mburn)

    mean_prob = mean_prob_beginning = np.zeros((a.m))
    posterior_list = []
    posterior_std_list = []
    for i in range(a.m):
        print('step ', i+1 , 'of ',a.m)
        pos, prob, state, blobs = sampler.run_mcmc(pos, a.save_state_every, rstate0=state, lnprob0=prob, blobs0 = blobs, storechain = True)
        np.save('%s/flatchain' %(directory),sampler.chain)
        np.save('%s/flatlnprobability' %(directory),sampler.lnprobability)
        np.save('%s/flatblobs' %(directory),sampler.blobs)
        posterior = np.load('%s/flatlnprobability.npy' %(directory))
        posterior_list.append(np.mean(posterior, axis = 0)[-1])
        posterior_std_list.append(np.std(posterior, axis = 0)[-1])
        np.save('%s/flatmeanposterior' %(directory), posterior_list)
        np.save('%s/flatstdposterior' %(directory), posterior_std_list)
        print(np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1])

        if i>202:
            print('posterior -1, -100, -200',np.mean(posterior, axis = 0)[-1], np.mean(posterior, axis = 0)[-100], np.mean(posterior, axis = 0)[-200])
            print('posterior 0, 100, 200',np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[100], np.mean(posterior, axis = 0)[200])
        #print("Mean acceptance fraction:", sampler.acceptance_fraction)
        elapsed1 = (time.time() - start1)
        print('calculation so far took', elapsed1, ' seconds')
        if i>a.min_mcmc_iterations and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-100]) < a.mcmc_tolerance and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-200]) < a.mcmc_tolerance:
            break
    if a.send_email:
        send_email(a.nthreads, i, np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1], a, elapsed1)

def mcmc_multi(changing_parameter, error_list, elements):
    '''
    Convenience function to use the MCMC for multiple zones (and therefore multiple observations). A subdirectory mcmc/ will be created in the current directory and intermediate chains will be stored there.
    The MCMC will sample the volume of best posterior for the likelihood functions that are declared in parameter.py.
    Default is a list of Proto-sun, Arcturus and B-stars. The MCMC uses many walkers and can use multiple threads. Each walker will evaluate a series of Chempy zones and add their posterior together which then will be returned.

    INPUT:

    changing_parameter = the parameter vector for initialization (will usually be found from minimization before). The initial chain will be created by jittering slightly the initial parameter guess

    error_list = the vector of element errors

    elements = the corresponding element symbols

    OUTPUT:

    The function will create a folder and store the chain as well as the predicted element values

    The MCMC stops when the convergence criteria is met, which is when the median posterior of all walkers does not change much inbetween 200 steps anymore.
    '''
    import time
    import os
    import multiprocessing as mp
    from .cem_function import  posterior_function_many_stars
    from .parameter import ModelParameters
    import emcee

    a = ModelParameters()
    start1 = time.time()
    directory = 'mcmc/'
    if os.path.exists(directory):
        if a.verbose:
            print('%s already existed. Content might be overwritten' %(directory))
    else:
        os.makedirs(directory)

    nthreads = mp.cpu_count()
    if nthreads == 4:
        nthreads = 2
    ndim = len(changing_parameter)
    a.nwalkers = max(a.nwalkers, int(ndim*2))
    chain = np.empty(shape = (a.nwalkers,ndim))

    for i in range(a.nwalkers):
        result = -np.inf
        while result == -np.inf:
            jitter = np.random.normal(loc = 0, scale = 0.001, size = ndim)
            result, dummy = posterior_function_many_stars(changing_parameter + jitter,error_list,elements)
        chain[i] = changing_parameter + jitter
    print('Chain created')
    sampler = emcee.EnsembleSampler(a.nwalkers,ndim,posterior_function_many_stars,threads=nthreads, args = [error_list,elements])
    pos,prob,state,blobs = sampler.run_mcmc(chain,a.mburn)

    mean_prob = mean_prob_beginning = np.zeros((a.m))
    posterior_list = []
    posterior_std_list = []
    for i in range(a.m):
        print('step ', i+1 , 'of ',a.m)
        pos, prob, state, blobs = sampler.run_mcmc(pos, a.save_state_every, rstate0=state, lnprob0=prob, blobs0 = blobs, storechain = True)
        np.save('%s/flatchain' %(directory),sampler.chain)
        np.save('%s/flatlnprobability' %(directory),sampler.lnprobability)
        np.save('%s/flatblobs' %(directory),sampler.blobs)
        posterior = np.load('%s/flatlnprobability.npy' %(directory))
        posterior_list.append(np.mean(posterior, axis = 0)[-1])
        posterior_std_list.append(np.std(posterior, axis = 0)[-1])
        np.save('%s/flatmeanposterior' %(directory), posterior_list)
        np.save('%s/flatstdposterior' %(directory), posterior_std_list)
        print(np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1])

        if i>202:
            print('posterior -1, -100, -200',np.mean(posterior, axis = 0)[-1], np.mean(posterior, axis = 0)[-100], np.mean(posterior, axis = 0)[-200])
            print('posterior 0, 100, 200',np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[100], np.mean(posterior, axis = 0)[200])
        #print("Mean acceptance fraction:", sampler.acceptance_fraction)
        elapsed1 = (time.time() - start1)
        print('calculation so far took', elapsed1, ' seconds')
        if i>a.min_mcmc_iterations and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-100]) < a.mcmc_tolerance and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-200]) < a.mcmc_tolerance:
            break
    if a.send_email:
        send_email(nthreads, i, np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1], a, elapsed1)

def send_email(thread_count, iteration_count, posterior_beginning, posterior_end, parameters, time):
    from email.MIMEMultipart import MIMEMultipart
    from email.MIMEText import MIMEText
    import smtplib

    fromaddr = "pythonspeaking@gmail.com"
    toaddr = "philcox@mpia.de"
    msg = MIMEMultipart()
    msg['From'] = fromaddr
    msg['To'] = toaddr
    msg['Subject'] = "Threads = %d, Run finished after %.2f hours" %(thread_count, time/3600.)
    body = "After %.1f hours %d threads produced %d iterations.\n The posterior at beginning was: %.2f. The posterior now is: %.2f.\n The stellar identifier list = %s.\n The error marginalization is %s \n  The yields are: %s %s %s \n " %(time/3600., thread_count, iteration_count, posterior_beginning, posterior_end, str(parameters.stellar_identifier_list), str(parameters.error_marginalization), parameters.yield_table_name_sn2, parameters.yield_table_name_agb, parameters.yield_table_name_1a)
    msg.attach(MIMEText(body, 'plain'))

    server = smtplib.SMTP('smtp.gmail.com', 587)
    server.ehlo()
    server.starttls()
    server.ehlo()
    server.login("pythonspeaking@gmail.com", "MPIA_Server_runs")
    text = msg.as_string()
    server.sendmail(fromaddr, toaddr, text)

def mcmc_quick(changing_parameter,elements,preload):
    '''
    Convenience function to use the MCMC for one zone. A subdirectory mcmc/ will be created in the current directory and intermediate chains will be stored there.
    The MCMC will sample the volume of best posterior for the likelihood functions that are declared in parameter.py.
    This is a cut down version to speed up MCMC for one star only
    INPUT:

    changing_parameter = the parameter vector for initialization (will usually be found from minimization before). The initial chain will be created by jittering slightly the initial parameter guess

    error_list = the vector of element errors

    elements = the corresponding element symbols

    OUTPUT:

    The function will create a folder and store the chain as well as the predicted element values

    The MCMC stops when the convergence criteria is met, which is when the median posterior of all walkers does not change much inbetween 200 steps anymore.
    '''
    import time
    import os
    import multiprocessing as mp
    from .cem_function import  posterior_function_mcmc_quick
    from .score_function import preload_params_mcmc
    from .parameter import ModelParameters
    import emcee

    a = ModelParameters()
    start1 = time.time()
    directory = 'mcmc/'
    if os.path.exists(directory):
        if a.verbose:
            print('%s already existed. Content might be overwritten' %(directory))
    else:
        os.makedirs(directory)

    nthreads = mp.cpu_count()
    if nthreads == 4:
        nthreads = 2
    ndim = len(changing_parameter)
    a.nwalkers = max(a.nwalkers, int(ndim*2))
    chain = np.empty(shape = (a.nwalkers,ndim))

    for i in range(a.nwalkers):
        result = -np.inf
        while result == -np.inf:
            jitter = np.random.normal(loc = 0, scale = 0.001, size = ndim)
            result = posterior_function_mcmc_quick(changing_parameter + jitter,elements,preload)
        chain[i] = changing_parameter + jitter

    pool=mp.Pool()
    sampler = emcee.EnsembleSampler(a.nwalkers,ndim,posterior_function_mcmc_quick,threads=nthreads, args = [elements,preload],pool=pool)
    pos,prob,state,blobs = sampler.run_mcmc(chain,a.mburn)

    mean_prob = mean_prob_beginning = np.zeros((a.m))
    posterior_list = []
    posterior_std_list = []
    for i in range(a.m):
        print('step ', i+1 , 'of ',a.m)
        pos, prob, state, blobs = sampler.run_mcmc(pos, a.save_state_every, rstate0=state, lnprob0=prob, blobs0 = blobs, storechain = True)
    #   np.save('%s/flatchain' %(directory),sampler.chain)
    #   np.save('%s/flatlnprobability' %(directory),sampler.lnprobability)
    #   np.save('%s/flatblobs' %(directory),sampler.blobs)
    #   posterior = np.load('%s/flatlnprobability.npy' %(directory))
        posterior = sampler.lnprobability
        posterior_list.append(np.mean(posterior, axis = 0)[-1])
        posterior_std_list.append(np.std(posterior, axis = 0)[-1])
    #   np.save('%s/flatmeanposterior' %(directory), posterior_list)
    #   np.save('%s/flatstdposterior' %(directory), posterior_std_list)
        print(np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1])

        #if i>202:
            #print('posterior -1, -100, -200',np.mean(posterior, axis = 0)[-1], np.mean(posterior, axis = 0)[-100], np.mean(posterior, axis = 0)[-200])
            #print('posterior 0, 100, 200',np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[100], np.mean(posterior, axis = 0)[200])
        #print("Mean acceptance fraction:", sampler.acceptance_fraction)
        elapsed1 = (time.time() - start1)
        print('calculation so far took', elapsed1, ' seconds')
        if i>a.min_mcmc_iterations and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-100]) < a.mcmc_tolerance and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-200]) < a.mcmc_tolerance:
            break
    np.save('%s/flatchain' %(directory),sampler.chain)
    np.save('%s/flatlnprobability' %(directory),sampler.lnprobability)
    np.save('%s/flatblobs' %(directory),sampler.blobs)
    posterior = sampler.lnprobability
    #posterior = np.load('%s/flatlnprobability.npy' %(directory))
    posterior_list.append(np.mean(posterior, axis = 0)[-1])
    posterior_std_list.append(np.std(posterior, axis = 0)[-1])
    np.save('%s/flatmeanposterior' %(directory), posterior_list)
    np.save('%s/flatstdposterior' %(directory), posterior_std_list)
    pool.close()
    if a.send_email:
        send_email(nthreads, i, np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1], a, elapsed1)

def single_star_optimization():
    '''
    This function will optimize the parameters of a single zone quickly

    INPUT:

    a = will be loaded from parameter.py (prepare all variables there)

    OUTPUT:

    log_list = a list of intermediate results (so far only for debugging)
    '''
    import time
    #import multiprocessing as mp
    from .optimization import minimizer_initial_quick
    from .cem_function import global_optimization_error_returned
    from .parameter import ModelParameters
    from .score_function import preload_params_mcmc

    # For testing
    import warnings
    warnings.filterwarnings("ignore")

    a = ModelParameters()
    preload = preload_params_mcmc()

    print(a.stellar_identifier_list)
    start_time = time.time()

    log_list = []
    # I: Minimization for each star seperately
    # 1: for each star make initial conditions (each star needs other model parameters)
    parameter_list = []
    for item in a.stellar_identifier_list:
        parameter_list.append(item)
    # 2: call posterior_function_for_minimization with scipy.optimize.minimize in multiprocess for each star and recover the found parameters
    #p = mp.Pool(len(parameter_list))
    #t = p.map(minimizer_initial_quick, parameter_list)
    #p.close()
    #p.join()
    #result = np.vstack(t)

    result  = minimizer_initial_quick(parameter_list)

    log_list.append(np.copy(result))
    log_list.append('initial minimization')
    initial = time.time()
    print('first minimization for each star separately took: %2.f seconds' %(initial - start_time))

    # IV: repeat II and III until posterior does not change much
    #result[:,:len(a.SSP_parameters)] = np.mean(result[:,:len(a.SSP_parameters)], axis = 0)
    #posteriors = []
    #counter = 0
    #while True:
    #   counter += 1
    #   if len(posteriors) > 1:
    #     if np.abs(posteriors[-1] - posteriors[-2]) < a.gibbs_sampler_tolerance:
    #      break
    #     if len(posteriors) > a.gibbs_sampler_maxiter:
    #      break

    #   initial = time.time()
        # II: Global parameter minimization:
        # 1: only SSP parameters free. Use mean SSP parameter values and individual (but fixed ISM parameter values)
    #   changing_parameter = result[0,:len(a.SSP_parameters)]
        # 2: Call each star in multiprocess but only return the predictions
        # 3: Calculate the likelihood for each star and optimize the common model error (is all done within minimizer global, which is calling 'global optimization')
    #   x = minimizer_global(changing_parameter,  a.tol_minimization, a.maxiter_minimization, a.verbose, result)

        # 4: return global SSP parameters and common model error
    #   posterior, error_list, elements = global_optimization_error_returned(x, result)
    #   posteriors.append(posterior)
    #   print(posteriors)

    #   global_iteration1 = time.time()
    #   print('step %d global minimization took: %2.f seconds' %(counter, global_iteration1 - initial))

        # III: Local parameter minimization:
        # 1: Use fixed global parameters and fixed common errors make initial conditions
    #   result[:,:len(a.SSP_parameters)] = x

    #   log_list.append((np.copy(x),posterior))
    #   log_list.append('step %d global minimization' %(counter))

    #   p0_list = []
    #   parameter_list = []
    #   x_list = []
    #   error_list_mp = []
    #   element_list_mp = []

    #   for i,item in enumerate(a.stellar_identifier_list):
    #     parameter_list.append(item)
    #     p0_list.append(result[i,len(a.SSP_parameters):])
    #     x_list.append(x)
    #     error_list_mp.append(error_list)
    #     element_list_mp.append(elements)

    #   args = zip(p0_list,parameter_list,x_list,error_list_mp,element_list_mp)

        # 2: Minimize each star ISM parameters in multiprocess
    #   p = mp.Pool(len(parameter_list))
    #   t = p.map(minimizer_local, args)
    #   p.close()
    #   p.join()
    #   local_parameters = np.vstack(t)
    #   result[:,len(a.SSP_parameters):] = local_parameters

    #   log_list.append(np.copy(result))
    #   log_list.append('step %d local minimization' %(counter))
    #   local_iteration1 = time.time()
    #   print('step %d local minimization took: %2.f seconds' %(counter, local_iteration1 - global_iteration1))

    #log_list.append(posteriors)
    #print(log_list)

    # V: MCMC run
    ## reshape the result to have global parameters in the front and the local parameters following
    #changing_parameter = list(result[0,:len(a.SSP_parameters)])

    elements = np.unique(a.elements_to_trace,preload.wildcard.dtype.names)
    changing_parameter = list(result)

    #for i in range(result.shape[0]):
    #   changing_parameter.append(list(result[i,len(a.SSP_parameters):]))

    changing_parameter = np.hstack(changing_parameter)

    ## jitter the parameters to initialise the chain (add a validation later, i.e. testing that the particular parameters yield a result)

    mcmc_quick(changing_parameter, elements,preload)

    # 1: Free all parameters and optimize common error (SSP should be the same for all stars)
    # 2: Plug everything into emcee and sample the posterior
    return log_list

def scoring_wrapper():
    """
    NO LONGER USED
    This function will calculate Bayes and CV scores for yield set, using the code in score_function.py.

    The neural network must be trained beforehand using training_data and create_network

    Main outputs are labelled .npz files in the Scores/ file.

    MUST set a.UseNeural = True for this and select correct dataset.
    """
    from Chempy.neural import training_data,create_network
    import time
    from Chempy.parameter import ModelParameters
    from Chempy.score_function import CV_wrapper, Bayes_wrapper
    init_time = time.time()
    a = ModelParameters()

    print('Step 1 (at time %.2f s): Calculate Bayes score' %(time.time()-init_time))
    Bayes_wrapper()

    print('Step 2 (at time %.2f s): Calculate cross-validation score' %(time.time()-init_time))
    CV_wrapper()

    print('Process complete in time %.2f s' %(time.time()-init_time))
    return None


# this is out of date, and no longer generates data that works for the analysis

import numpy as np
import networkx as nx
import json
import random
import string
import multiprocessing

def play(game_data, n_steps=100, treatment_id="", adopt=None, forget=None):
    """ play a single round of the game """

    def setup_game(game_data):
        # set up agents
        g = nx.from_dict_of_lists(game_data['neighbors'])

        # give starting information
        nx.set_node_attributes(
            g,
            name='M',  # M for mind/memory
            values={i: nx.from_edgelist([game_data['clues'][bf]['nodes'] for bf in game_data['beliefs'][str(i)]])
                    for i in g}
        )

        nx.set_node_attributes(
            g,
            name='F',  # F for forgetory
            values={i: nx.Graph() for i in g}
        )

        ### this stuff is just for generating data that mirrors the experiment
        # save initial state
        nx.set_node_attributes(
            g,
            name="initialState",
            values={i: {"promising_leads": {"clueIDs": game_data['beliefs'][str(i)]},
                        "dead_ends": {"clueIDs": []}
                        }
                    for i in g}
        )

        # assign random identifiers
        nx.set_node_attributes(
            g,
            name="_id",
            values={i: ''.join(random.choices(string.ascii_letters + string.digits, k=16))
                    for i in g}
        )

        # initialize logs
        nx.set_node_attributes(
            g,
            name='data.log',  # M for mind
            values={i: list() for i in g}
        )

        return g

    def adopt(g, ego, edge, t, target):

        if g.node[ego]['M'].has_edge(*edge):
            return False

        M = g.node[ego]['M']

        exposers = [nb for nb in g[ego] if edge in g.node[nb]['M'].edges()]
        exposure = len(exposers)
        if exposure == 0:
            return False

        ends_present = len(set(edge).intersection(set(M.nodes())))

        if M.has_node(edge[0]) and M.has_node(edge[1]):
            #path_list = nx.all_simple_paths(M, *edge, cutoff=4)
            path_list = nx.all_simple_paths(M, *edge, cutoff=2)
        else:
            path_list = []

        path_counts = {i: 0 for i in range(6)}
        for path in path_list:
            path_counts[len(path)] += 1

        n_beliefs = len(M.edges())

        baseline = .01 / (100)
        c_exposure = np.log(3)  # for each exposure multiply the likelihood by 1.5
        c_ends_present = np.log(5)
        c_len2paths = np.log(10)
        c_len3paths = np.log(3)
        c_len4paths = np.log(2)
        c_offtarget = np.log(.6)

        likelihood = baseline * np.exp(
             c_exposure * exposure +
            c_ends_present * ends_present +
            c_len2paths * path_counts[2] +
            c_len3paths * path_counts[3] +
            c_len4paths * path_counts[4] +
            c_offtarget * np.sign(n_beliefs - target)*(n_beliefs-target)**2
        )

        #adoption = path_counts[2] > 0

        adoption = np.random.binomial(1, np.min([likelihood, .9999])) == 1

        if adoption:
            source = g.node[np.random.choice(exposers)]["_id"]
            g.node[ego]['data.log'].append({
                "event": "drop",
                "data": {
                    "clue": lookup_cluename(edge),
                    "source": source,
                    "dest": "promising_leads",
                    "destIndex": 0
                },
                "at": t  # str(datetime.datetime.now()) # t + np.random.rand()  # unique timestamp within the second
            })

        return adoption

    def forget(g, ego, edge, t, target):
        if not g.node[ego]['M'].has_edge(*edge):
            return False

        M = g.node[ego]['M']

        exposers = [nb for nb in g[ego] if edge in g.node[nb]['M'].edges()]
        exposure = len(exposers)
        if exposure == 0:
            return False

        ends_present = len(set(edge).intersection(set(M.nodes())))

        if M.has_node(edge[0]) and M.has_node(edge[1]):
            path_list = nx.all_simple_paths(M, *edge, cutoff=4)
        else:
            path_list = []

        path_counts = {i: 0 for i in range(6)}
        for path in path_list:
            path_counts[len(path)] += 1

        # path_counts = pd.Series(pd.Series([len(pth) - 1 for pth in path_list]).value_counts(),
        #                         index=range(2, 5)).fillna(0)

        n_beliefs = len(M.edges())

        baseline = 10 / 100
        c_exposure = np.log(1/3)  # for each exposure multiply the likelihood by 1.5
        c_ends_present = np.log(1/5)
        c_len2paths = np.log(1/10)
        c_len3paths = np.log(1/3)
        c_len4paths = np.log(1/2)
        c_offtarget = np.log(1/.6)

        likelihood = baseline * np.exp(
            c_exposure * exposure +
            c_ends_present * ends_present +
            c_len2paths * path_counts[2] +
            c_len3paths * path_counts[3] +
            c_len4paths * path_counts[4] +
            c_offtarget * np.sign(n_beliefs - target)*(n_beliefs-target)**2
        )

        forgetion = np.random.binomial(1, np.min([likelihood, .99])) == 1

        if forgetion:
            g.node[ego]['data.log'].append({
                "event": "drop",
                "data": {
                    "clue": lookup_cluename(edge),
                    "source": "promising_leads",
                    "dest": "dead_ends",
                    "destIndex": 0
                },
                "at": t  # + np.random.rand()  # unique timestamp within the second
            })

        return forgetion

    target = game_data['parameters']['target']
    g = setup_game(game_data)

    # useful helper functions
    beliefs = [cl["nodes"] for _, cl in game_data["clues"].items()]
    hashtable = {hash(tuple(sorted(clue["nodes"]))): key for key, clue in game_data["clues"].items()}

    def lookup_cluename(edge):
        return hashtable[hash(tuple(sorted(edge)))]

    # play the game
    for step in range(n_steps):
        substep = 0
        for ego in np.random.permutation(g):  # select ego in random order
            for edge in np.random.permutation(beliefs):  # select a belief in random order to propagate
                substep += .000001  # ensure proper ordering
                if adopt(g, ego, edge, step + substep, target):
                    g.node[ego]['M'].add_edges_from([edge])
                    if g.node[ego]['F'].has_edge(*edge):
                        g.node[ego]['F'].remove_edges_from([edge])
                substep += .000001  # so you can't have a forget event at the same time as an adopt event
                if forget(g, ego, edge, step + substep, target):
                    g.node[ego]['M'].remove_edges_from([edge])
                    g.node[ego]['F'].add_edges_from([edge])

    # save to json for postprocessing (keep as string here)
    player_json = []
    for n in g:
        player_object = {
            "_id": str(g.node[n]["_id"]),
            "initialState": g.node[n]["initialState"],
            "log": g.node[n]["log"],
            "alterIDs": [g.node[nb]["_id"] for nb in g.neighbors(n)],
            "notebooks": {
                "promising_leads": {"clueIDs": [lookup_cluename(edge) for edge in g.node[n]['M'].edges()]},
                "dead_ends": {"clueIDs": [lookup_cluename(edge) for edge in g.node[n]['F'].edges()]}
            }
        }
        json_string = json.dumps(player_object)
        player_json.append(json_string)

    game_object = {"_id": ''.join(random.choices(string.ascii_letters + string.digits, k=16)),
                   "playerIds": [g.node[n]['_id'] for n in g],
                   "treatmentId": treatment_id,
                   "gameSetupId": game_data["gameSetupId"],
                   "createdAt": 0,
                   "finishedAt": n_steps}

    game_json = [json.dumps(game_object)]

    return game_json, player_json

def simulate_experiment(experiment_filename):
    """Run all the games in the experiment"""
    with open(experiment_filename, 'r') as infile:
        experiment = json.load(infile)

    with multiprocessing.Pool() as p:
        game_results = p.map(play, experiment["games"].values())

    #game_results = [play(g) for g in experiment["games"].values()]

    game_json_list = []
    player_json_list = []
    for res in game_results:
        game_json_list += res[0]
        player_json_list += res[1]

    return game_json_list, player_json_list


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""This module define the Frames available for computation and their relations
to each other.

The relations may be circular, thanks to the use of the Node class.

.. code-block:: text

    ,---.          ,-------.        ,----.
    |G50|---bias---|EME2000|..bias..|GCRF|
    `---'          `-------'        `----'
                       |              |
                   Precession         |
                       |              |
                     ,---.        Precession
                     |MOD|            +
                     `---'         Nutation
                       |     + model corrections
                    Nutation          |
              + model corrections     |
                       |              |
    ,----.           ,---.         ,----.
    |TEME|--Equinox--|TOD|         |CIRF|
    `----'           `---'         `----'
                       |              |
                 Sideral time   Sideral time
                       |              |
                     ,---.         ,----.
                     |PEF|         |TIRF|
                     `---'         `----'
                        \\            /
                    IAU 1980       IAU 2010
           Earth Orientation       Earth Orientation
                  Parameters       Parameters
                           \\     /
         ,-----.            ,----.
         |WGS84|--identity--|ITRF|
         `-----'            `----'
"""

import sys
import logging
import numpy as np

from ..errors import UnknownFrameError
from ..constants import Earth
from ..utils.matrix import rot3
from ..utils.node import Node
from . import iau1980, iau2010
from .local import to_qsw, to_tnw

CIO = ["ITRF", "TIRF", "CIRF", "GCRF"]
IAU1980 = ["TOD", "MOD"]
OTHER = ["EME2000", "TEME", "WGS84", "PEF", "G50"]

__all__ = CIO + IAU1980 + OTHER + ["get_frame"]

log = logging.getLogger(__name__)

class FrameCache(dict):
    """This class is here to emulate module behavior for dynamically
    created frames.

    It's useful when pickle is involved (e.g. multiprocessing)
    """

    def __getattr__(self, name):
        if name not in self:
            raise AttributeError(name)
        return self[name]

dynamic = FrameCache()
"""This dictionary contains all the frames. Those defined here, and those created on the fly
by the developer.
"""

sys.modules[__name__ + ".dynamic"] = dynamic

def get_frame(frame):
    """Frame factory

    Args:
        frame (str): name of the desired frame
    Return:
        ~beyond.frames.frames.Frame
    """
    if frame not in dynamic.keys():
        raise UnknownFrameError(frame)

    return dynamic[frame]

class _MetaFrame(type, Node):
    """This MetaClass is here to join the behaviors of ``type`` and ``Node``
    """

    def __init__(cls, name, bases, dct):

        bypass = dct.pop("bypass", False)

        super(_MetaFrame, cls).__init__(name, bases, dct)
        super(type, cls).__init__(name)

        if not bypass and cls.__name__ in dynamic:
            log.warning(
                "A frame with the name '%s' is already registered. Overriding"
                % cls.__name__
            )

        cls.__module__ = __name__ + ".dynamic"

        # Making the frame available to the get_frame function
        dynamic[cls.__name__] = cls

    def __repr__(cls):  # pragma: no cover
        return "<Frame '{}'>".format(cls.name)

class Frame(metaclass=_MetaFrame):
    """Frame base class
    """

    center = Earth

    def __init__(self, date, orbit):
        """
        Args:
            date (~beyond.utils.Date)
            orbit (numpy.ndarray)
        """
        self.date = date
        self.orbit = orbit

    def __str__(self):  # pragma: no cover
        return self.name

    def __repr__(self):  # pragma: no cover
        return "<Frame obj '{}'>".format(self.__class__.__name__)

    @classmethod
    def _convert(cls, x=None, y=None):
        m = np.identity(6)

        if x is not None:
            m[:3, :3] = x
        if y is not None:
            m[3:, 3:] = y

        return m

    def transform(self, new_frame):
        """Change the frame of the orbit

        Args:
            new_frame (str)
        Return:
            numpy.ndarray
        """

        steps = self.__class__.steps(new_frame)

        orbit = self.orbit

        for _from, _to in steps:

            from_obj = _from(self.date, orbit)
            direct = "_to_%s" % _to

            if hasattr(from_obj, direct):
                rotation, offset = getattr(from_obj, direct)()
            else:
                to_obj = _to(self.date, orbit)
                inverse = "_to_%s" % _from
                if hasattr(to_obj, inverse):
                    rotation, offset = getattr(to_obj, inverse)()
                    rotation = rotation.T
                    offset = -offset
                else:
                    raise NotImplementedError(
                        "Unknown transformation {} to {}".format(_from, _to)
                    )

            if getattr(_from, "_rotation_before_translation", False):
                # In case of topocentric frame, the rotation is done before the translation
                orbit = offset + (rotation @ orbit)
            else:
                orbit = rotation @ (offset + orbit)

        return orbit

class TEME(Frame):
    """True Equator Mean Equinox"""

    orientation = "TEME"

    def _to_TOD(self):
        equin = iau1980.equinox(
            self.date, eop_correction=False, terms=4, kinematic=False
        )
        m = rot3(-np.deg2rad(equin))
        return self._convert(m, m), np.zeros(6)

class GTOD(Frame):
    """Greenwich True Of Date"""

    orientation = "GTOD"

class WGS84(Frame):
    """World Geodetic System 1984"""

    orientation = "WGS84"

    def _to_ITRF(self):
        return np.identity(6), np.zeros(6)

class PEF(Frame):
    """Pseudo Earth Fixed"""

    orientation = "PEF"

    def _to_TOD(self):
        m = iau1980.sideral(self.date, model="apparent", eop_correction=False)
        offset = np.zeros(6)
        offset[3:] = np.cross(iau1980.rate(self.date), self.orbit[:3])
        return self._convert(m, m), offset

class TOD(Frame):
    """True (Equator) Of Date"""

    orientation = "TOD"

    def _to_MOD(self):
        m = iau1980.nutation(self.date, eop_correction=False)
        return self._convert(m, m), np.zeros(6)

class MOD(Frame):
    """Mean (Equator) Of Date"""

    orientation = "MOD"

    def _to_EME2000(self):
        m = iau1980.precesion(self.date)
        return self._convert(m, m), np.zeros(6)

class EME2000(Frame):
    """EME2000 inertial frame (also known as J2000)"""

    orientation = "EME2000"

class ITRF(Frame):
    """International Terrestrial Reference Frame"""

    orientation = "ITRF"

    def _to_PEF(self):
        m = iau1980.earth_orientation(self.date)
        return self._convert(m, m), np.zeros(6)

    def _to_TIRF(self):
        m = iau2010.earth_orientation(self.date)
        return self._convert(m, m), np.zeros(6)

class TIRF(Frame):
    """Terrestrial Intermediate Reference Frame"""

    orientation = "TIRF"

    def _to_CIRF(self):
        m = iau2010.sideral(self.date)
        offset = np.zeros(6)
        offset[3:] = np.cross(iau2010.rate(self.date), self.orbit[:3])
        return self._convert(m, m), offset

class CIRF(Frame):
    """Celestial Intermediate Reference Frame"""

    orientation = "CIRF"

    def _to_GCRF(self):
        m = iau2010.precesion_nutation(self.date)
        return self._convert(m, m), np.zeros(6)

class GCRF(Frame):
    """Geocentric Celestial Reference Frame"""

    orientation = "GCRF"

class G50(Frame):
    """Gamma50 Reference Frame
    """

    orientation = "G50"

    def _to_EME2000(self):

        m = [
            [0.9999256794956877, -0.0111814832204662, -0.0048590038153592],
            [0.0111814832391717, 0.9999374848933135, -0.0000271625947142],
            [0.0048590037723143, -0.0000271702937440, 0.9999881946023742],
        ]

        return self._convert(m, m), np.zeros(6)

def orbit2frame(name, ref_orbit, orientation=None, center=None, bypass=False):
    """Create a frame based on a Orbit or Ephem object.

    Args:
        name (str): Name to give the created frame
        ref_orbit (Orbit or Ephem):
        orientation (str): Orientation of the created frame
        bypass (bool): By-pass the warning when creating a frame with an already
            taken name
    Return:
        Frame:

    If orientation is ``None``, the new frame will keep the orientation of the
    reference frame of the Orbit and move along with the orbit.
    Other acceptable values are ``"QSW"`` (and its aliases "LVLH" and "RSW") or ``"TNW"``.

    See :py:func:`~beyond.frames.local.to_qsw` and :py:func:`~beyond.frames.local.to_tnw`
    for informations regarding these orientations.
    """

    if orientation is None:
        orientation = ref_orbit.frame.orientation
    elif orientation.upper() in ("RSW", "LVLH"):
        orientation = "QSW"
    elif orientation.upper() not in ("QSW", "TNW"):
        raise ValueError("Unknown orientation '%s'" % orientation)

    if center is None:
        center = Earth

    def _to_parent_frame(self):
        """Conversion from orbit frame to parent frame
        """
        offset = ref_orbit.propagate(self.date).base.copy()

        if orientation.upper() in ("QSW", "TNW"):

            # propagation of the reference orbit to the date of the
            # converted orbit
            orb = ref_orbit.propagate(self.date)

            m = to_qsw(orb) if orientation.upper() == "QSW" else to_tnw(orb)

            # we transpose the matrix because it represents the conversion
            # from inertial to local frame, and we'd like the other way around
            rotation = Frame._convert(m, m).T
        else:
            # The orientation is the same as the parent reference frame
            rotation = np.identity(6)

        return rotation, offset

    # define the name of the method of conversion
    mtd = "_to_%s" % ref_orbit.frame.__name__

    # dictionary which defines attributes of the created class
    dct = {
        mtd: _to_parent_frame,
        "orientation": orientation,
        "center": center,
        "bypass": bypass,
    }

    # Creation of the class
    cls = _MetaFrame(name, (Frame,), dct)

    # Link to the parent
    cls + ref_orbit.frame
    return cls

WGS84 + ITRF + PEF + TOD + MOD + EME2000
TOD + TEME
# EME2000 + GCRF
ITRF + TIRF + CIRF + GCRF
EME2000 + G50


import matplotlib.pyplot as plt
import os
import numpy as np
import yt

from pygrackle import \
    FluidContainer, \
    chemistry_data, \
    evolve_constant_density

from pygrackle.utilities.physical_constants import \
    mass_hydrogen_cgs, \
    sec_per_Myr, \
    cm_per_mpc

import sys

from multiprocessing import Pool
from contextlib import closing
import itertools

tiny_number = 1e-20

class NoStdStreams(object):
    def __init__(self,stdout = None, stderr = None):
        self.devnull = open(os.devnull,'w')
        self._stdout = stdout or self.devnull or sys.stdout
        self._stderr = stderr or self.devnull or sys.stderr

    def __enter__(self):
        self.old_stdout, self.old_stderr = sys.stdout, sys.stderr
        self.old_stdout.flush(); self.old_stderr.flush()
        sys.stdout, sys.stderr = self._stdout, self._stderr

    def __exit__(self, exc_type, exc_value, traceback):
        self._stdout.flush(); self._stderr.flush()
        sys.stdout = self.old_stdout
        sys.stderr = self.old_stderr
        self.devnull.close()

def cooling_cell(density = 12.2,
                 initial_temperature = 2.0E4,
                 final_time = 30.0,
                 metal_fraction = 4.0E-4,
                 make_plot = False,
                 save_output = False, primordial_chemistry = 2,
                 outname = None, save_H2_fraction = False,
                 return_result = False,
                 verbose = False, H2_converge = None,
                 *args, **kwargs):

    current_redshift = 0.

    # Set solver parameters
    my_chemistry = chemistry_data()
    my_chemistry.use_grackle = 1
    my_chemistry.with_radiative_cooling = 1
    my_chemistry.primordial_chemistry = primordial_chemistry
    my_chemistry.metal_cooling = 1
    my_chemistry.UVbackground = 1
    my_chemistry.self_shielding_method = 3

    if primordial_chemistry > 1:
        my_chemistry.H2_self_shielding = 2
        my_chemistry.h2_on_dust = 1
        my_chemistry.three_body_rate = 4

    grackle_dir = "/home/aemerick/code/grackle-emerick/"
    my_chemistry.grackle_data_file = os.sep.join( #['/home/aemerick/code/grackle-emerick/input/CloudyData_UVB=HM2012.h5'])
        [grackle_dir, "input","CloudyData_UVB=HM2012_shielded.h5"])

    # set the factors
    my_chemistry.LW_factor = kwargs.get("LW_factor", 1.0)
    my_chemistry.k27_factor = kwargs.get("k27_factor", 1.0)
    #if 'LW_factor' in kwargs.keys():
    #    my_chemistry.LW_factor = kwargs['LW_factor']
    #else:
    #    my_chemistry.LW_factor = 1.0

    #if 'k27_factor' in kwargs.keys():
    #    my_chemistry.k27_factor = kwargs['k27_factor']
    #else:
    #    my_chemistry.k27_factor = 1.0

    # Set units
    my_chemistry.comoving_coordinates = 0 # proper units
    my_chemistry.a_units = 1.0
    my_chemistry.a_value = 1. / (1. + current_redshift) / \
        my_chemistry.a_units
    my_chemistry.density_units = mass_hydrogen_cgs # rho = 1.0 is 1.67e-24 g
    my_chemistry.length_units = cm_per_mpc         # 1 Mpc in cm
    my_chemistry.time_units = sec_per_Myr          # 1 Myr in s
    my_chemistry.velocity_units = my_chemistry.a_units * \
        (my_chemistry.length_units / my_chemistry.a_value) / \
        my_chemistry.time_units

    rval = my_chemistry.initialize()

    fc = FluidContainer(my_chemistry, 1)
    fc["density"][:] = density
    if my_chemistry.primordial_chemistry > 0:
        fc["HI"][:] = 0.76 * fc["density"]
        fc["HII"][:] = tiny_number * fc["density"]
        fc["HeI"][:] = (1.0 - 0.76) * fc["density"]
        fc["HeII"][:] = tiny_number * fc["density"]
        fc["HeIII"][:] = tiny_number * fc["density"]
    if my_chemistry.primordial_chemistry > 1:
        fc["H2I"][:] = tiny_number * fc["density"]
        fc["H2II"][:] = tiny_number * fc["density"]
        fc["HM"][:] = tiny_number * fc["density"]
        fc["de"][:] = tiny_number * fc["density"]
        fc['H2_self_shielding_length'][:] = 1.8E-6
    if my_chemistry.primordial_chemistry > 2:
        fc["DI"][:] = 2.0 * 3.4e-5 * fc["density"]
        fc["DII"][:] = tiny_number * fc["density"]
        fc["HDI"][:] = tiny_number * fc["density"]
    if my_chemistry.metal_cooling == 1:
        fc["metal"][:] = metal_fraction * fc["density"] * \
          my_chemistry.SolarMetalFractionByMass

    fc["x-velocity"][:] = 0.0
    fc["y-velocity"][:] = 0.0
    fc["z-velocity"][:] = 0.0

    fc["energy"][:] = initial_temperature / \
        fc.chemistry_data.temperature_units
    fc.calculate_temperature()
    fc["energy"][:] *= initial_temperature / fc["temperature"]

    # timestepping safety factor
    safety_factor = 0.001

    # let gas cool at constant density

    #if verbose:
    print("Beginning Run")
    data = evolve_constant_density(
            fc, final_time=final_time, H2_converge = H2_converge,
            safety_factor=safety_factor, verbose = verbose)
    #else:
    #    print "Beginning Run"

    #    with NoStdStreams():
    #        data = evolve_constant_density(
    #            fc, final_time=final_time, H2_converge = 1.0E-6,
    #            safety_factor=safety_factor)
    #    print "Ending Run"

    if make_plot:
        p1, = plt.loglog(data["time"].to("Myr"), data["temperature"],
                            color="black", label="T")
        plt.xlabel("Time [Myr]")
        plt.ylabel("T [K]")

        data["mu"] = data["temperature"] / \
            (data["energy"] * (my_chemistry.Gamma - 1.) *
             fc.chemistry_data.temperature_units)
        plt.twinx()
        p2, = plt.semilogx(data["time"].to("Myr"), data["mu"],
                              color="red", label="$\\mu$")
        plt.ylabel("$\\mu$")
        plt.legend([p1,p2],["T","$\\mu$"], fancybox=True,
                      loc="center left")
        plt.savefig("cooling_cell.png")

    # save data arrays as a yt dataset
    if outname is None:
        outname = 'cooling_cell_%.2f_%.2f'%(my_chemistry.k27_factor,
                                           my_chemistry.LW_factor)

    if save_output:

        yt.save_as_dataset({}, outname + '.h5', data)

    if my_chemistry.primordial_chemistry > 1:
        H2_fraction = (data['H2I'] + data['H2II']) / data['density']
    else:
        H2_fraction = np.zeros(np.size(data['density']))

    if save_H2_fraction:
        #np.savetxt(outname + ".dat", [data['time'], H2_fraction])

        f = open("all_runs_d_%.2f.dat"%(density),"a")
#        f.write("# k27 LW f_H2 T time\n")
        f.write("%8.8E %8.8E %8.8E %8.8E %8.8E \n"%(my_chemistry.k27_factor,
                                              my_chemistry.LW_factor,
                                              H2_fraction[-1], data['temperature'][-1],
                                              data['time'][-1] ))
        f.close()

    if return_result:
        return data
    else:
        return

def _parallel_loop(i, k27, LW):

    primordial_chemistry = 1

    data = cooling_cell(k27_factor = k27, LW_factor = LW, save_output = False,
                        save_H2_fraction = False, primordial_chemistry = primordial_chemistry,
                        return_result = True)

    if primordial_chemistry > 1:
        H2_fraction = (data['H2I'] + data['H2II']) / data['density']
    else:
        H2_fraction = np.zeros(np.size(data['density']))

    T           = (data['temperature'])

    str_i = "%00005i"%(i)

    result = { str_i : {}}
    result[str_i]['k27'] = k27
    result[str_i]['LW']  = LW
    result[str_i]['H2_fraction'] = H2_fraction[-1]
    result[str_i]['T'] = T[-1]

    return result

def _parallel_loop_star(args):
    return _parallel_loop(*args)

def cooling_cell_grid(k27_factors = None, LW_factors = None,
                      fmin = 0.1, fmax = 10000.0, npoints = 100,
                      nproc = 1, outname = None):

    if outname is None:
        outname = "all_parallel_runs.dat"

    if k27_factors is None:
        k27_factors = np.logspace(np.log10(fmin),
                                  np.log10(fmax), npoints)
    if LW_factors is None:
        LW_factors  = 1.0 * k27_factors

    if nproc == 1:
        call_cell = lambda x, y : cooling_cell(k27_factor = x,
                                               LW_factor  = y, save_H2_fraction = True)

        for i,k27 in enumerate(k27_factors):
            print((i)*np.size(LW_factors))

            temp_cell = lambda y : call_cell(k27,y)
            list(map(temp_cell, LW_factors)) # this may not work anymore - AE python 2 to 3
    else:

        LW_mesh, k27_mesh = np.meshgrid(LW_factors, k27_factors)

        k27_mesh = k27_mesh.flatten()
        LW_mesh  = LW_mesh.flatten()

        for sub_list in itertools.zip_longest(*(iter( np.arange(np.size(k27_mesh))),) * nproc):
            sub_list = list(sub_list)
            sub_list = [s for s in sub_list if s is not None]
            reduced_nproc = np.min( [len(sub_list), nproc])

            print("running for ", sub_list)

            imin,imax = sub_list[0], (sub_list[-1] + 1)

            pool = Pool(reduced_nproc)
            results = pool.map_async(_parallel_loop_star,
                                     zip(sub_list,
                                                    k27_mesh[imin:imax], LW_mesh[imin:imax]))
            pool.close()
            pool.join()

            for r in results.get():
                str_i = list(r.keys())[0]

                f = open(outname,"a")
                f.write("%8.8E %8.8E %8.8E %8.8E\n"%(  r[str_i]['k27'],
                                                       r[str_i]['LW'], r[str_i]['H2_fraction'],
                                                       r[str_i]['T']))
                f.close()

            del(results)

    return

if __name__ == "__main__":

    # test this
    #cooling_cell(k27_factor = 0.99, LW_factor = 0.99,
    #             save_output = False, save_H2_fraction=True)
    import time

    npoints = 16
    nproc   = 4

    start = time.time()
    cooling_cell_grid(npoints = npoints, nproc = nproc, outname = str(sys.argv[1]))
    end = time.time()

    dt = end - start
    eff = dt / (1.0*nproc)
    print("This run of %i models on %i processors took %.3E s - Eff = %.1E"%(npoints*npoints, nproc, dt, eff))