"""This is a module for calcuating groups and integrations with JWST on
the fly. The main function (``perform_calculation``) takes a dictionary
of inputs (modeled around how the web tool takes inputs) which must
include: ``observation_time``, ``num_groups``, ``magnitude``,
``model``, ``band``, ``filt``, ``filt_ta``, ``instrument``,
``subarray``, ``subarray_ta``, ``saturation_mode``, ``max_saturation``,
and ``infile``. It produces and dictionary of outputs that includes all
of the original information as well as groups, integrations, saturation
levels, and observation time estimates for target acquisition and
science observations with JWST.

Authors
-------

    Jules Fowler, April 2017
    Matthew Bourque, February 2021

Use
---

    This is mostly a module to be used by the ExoCTK web application,
    but the main function can be run standalone as such:
    ::

        from exoctk.groups_integrations.groups_integrations import perform_calculation
        perform_calculation()

Dependencies
------------

    - ``astropy``
    - ``numpy``
    - ``scipy``
"""

import json
import math

import numpy as np
from scipy import interpolate

def calc_duration_time(num_groups, num_integrations, num_reset_frames, frame_time, frames_per_group=1):
    """Calculates duration time (or exposure duration as told by APT)

    Parameters
    ----------
    num_groups : int
        Groups per integration
    num_integrations : int
        Integrations per exposure
    num_reset_frames : int
        Reset frames per integration
    frame_time : float
        Frame time (in seconds)
    frames_per_group : int, optional
        Frames per group -- always one except brown dwarves

    Returns
    -------
    duration_time : float
        Duration time (in seconds).
    """

    duration_time = frame_time * (num_groups * frames_per_group + num_reset_frames) * num_integrations

    return duration_time

def calc_exposure_time(num_integrations, ramp_time):
    """Calculates exposure time (or photon collection duration as told
    by APT.)

    Parameters
    ----------
    num_integrations : int
        Integrations per exposure.
    ramp_time : float
        Ramp time (in seconds).

    Returns
    -------
    exposure_time : float
        Exposure time (in seconds).
    """

    exposure_time = num_integrations * ramp_time

    return exposure_time

def calc_frame_time(num_columns, num_rows, num_amps, instrument):
    """Calculates the frame time for a given
    instrument/readmode/subarray.

    Parameters
    ----------
    num_columns : int
        Number of columns
    num_rows : int
        Number of rows
    num_amps : int
        Amplifiers reading data
    instrument : str
        The instrument

    Returns
    -------
    frame_time : float
        The frame time (in seconds)
    """

    num_columns, num_amps, num_rows = int(num_columns), int(num_amps), int(num_rows)

    if instrument == 'nirspec':
        n = 2
    if instrument in ['nircam', 'niriss']:
        n = 1

    frame_time = (num_columns / num_amps + 12) * (num_rows + n) * (1e-5)

    return frame_time

def calc_groups_from_exp_time(max_exptime_per_int, frame_time):
    """Given the maximum saturation time, calculates the number of
    frames per group.

    Parameters
    ----------
    max_exptime_per_int : float
        The maximum number of seconds an integration can last before
        it's oversaturated
    frame_time : float
        The time per frame

    Returns
    -------
    groups : int
        The required number of groups.
    """

    groups = np.floor(max_exptime_per_int / frame_time)

    return groups

def calc_integration_time(num_groups, frame_time, frames_per_group, num_skips):
    """Calculates the integration time.

    Parameters
    ----------
    num_groups : int
        Groups per integration.]
    frame_time : float
        Frame time (in seconds)
    frames_per_group : int
        Frames per group -- always 1 except maybe brown dwarves
    num_skips : int
        Skips per integration -- always 0 except maybe brown dwarves

    Returns
    -------
    integration_time : float
        Integration time (in seconds)
    """

    integration_time = (num_groups * (frames_per_group + num_skips) - num_skips) * frame_time

    return integration_time

def calc_num_integrations(transit_time, num_groups, num_reset_frames, frame_time, frames_per_group):
    """Calculates number of integrations required.

    Parameters
    ----------
    transit_time : float
        The time of the transit (in hours)
    num_groups : int
        Groups per integration
    num_reset_frames : int
        Number of reset frames per integration
    frame_time : float
        The frame time (in seconds)
    frames_per_group : int
        Frames per group -- always 1 except maybe for brown dwarves

    Returns
    -------
    num_integrations : float
        The required number of integraitions.
    """

    num_integrations = math.ceil((float(transit_time) * 3600) / (frame_time * (num_groups * frames_per_group + num_reset_frames)))

    return num_integrations

def calc_observation_efficiency(exposure_time, duration_time):
    """Calculates the observation efficiency.

    Parameters
    ----------
    exposure_time : float
        Exposure time (in seconds).
    duration_time : float
        Duration time (in seconds).

    Returns
    -------
    observation_efficiency : float
        Observation efficiency.
    """

    observation_efficiency = exposure_time / duration_time

    return observation_efficiency

def calc_ramp_time(integration_time, num_reset_frames, frame_time):
    """Calculates the ramp time -- or the integration time plus overhead
    for resets.

    Parameters
    ----------
    integration_time : float
        Integration time (in seconds)
    num_reset_frames : int
        Rest frames per integration
    frame_time : float
        Frame time (in seconds)

    Returns
    -------
    ramp_time : float
        Ramp time (in seconds).
    """

    ramp_time = integration_time + (num_reset_frames - 1) * frame_time

    return ramp_time

def convert_saturation(max_saturation, saturation_mode, instrument, infile, target_acq_mode=False):
    """Converts full well fraction to a saturation in counts OR
    provides the max fullwell for TA mode.

    Parameters
    ----------
    max_saturation : float
        Either a full well fraction or counts
    saturation_mode : str
        ``well`` or ``counts``
    instrument : str
        The instrument
    infile : str
        The path to the data file
    target_acq_mode : bool, optional
        Whether or not it's TA mode

    Returns
    -------
    max_saturation : float
        The fullwell to use in counts.
    """

    with open(infile) as f:
        data = json.load(f)

    instrument_dict = data['fullwell']

    if saturation_mode == 'well':
        max_saturation = float(max_saturation) * float(instrument_dict[instrument])

    if target_acq_mode:
        max_saturation = instrument_dict[instrument]

    return max_saturation

def interpolate_from_pandeia(magnitude, instrument, filt, subarray, model, band, frame_time, saturation_level, infile, target_acq_mode=False):
    """Interpolates the precalculated ``pandeia`` data to estimate the
    saturation limit.

    Parameters
    ----------
    magnitude : float
        The magnitude of the source. (Takes between 4.5-12.5)
    instrument : str
       The instrument, allowable ``miri``, ``niriss``, ``nirspec``,
       ``nircam``
    filt : str
        The filter
    subarray : str
        The subarray
    model : str
        Phoenix model key
    band : str
        Magnitude band
    frame_time : float
        Frame time
    saturation_level : float
        The maximum fullwell saturation we'll allow
    infile : str
        The data file to use
    target_acq_mode : bool, optional
        Whether or not we're running this for TA

    Returns
    -------
    num_groups : int
        The number of groups that won't oversaturate the detector
    max_sat : int
        The maximum saturation level reached by that number of groups
    """

    # Create the dictionaries for each filter and select out the prerun data
    with open(infile) as f:
        data = json.load(f)

    ta_or_sci = 'sci_sat'

    if target_acq_mode:
        ta_or_sci = 'ta_sat'

    # The data
    magnitudes = np.array(data['mags'])
    saturation = data[ta_or_sci][instrument][filt][subarray][model]
    log_saturation = np.log10(saturation)

    # Interpolate the given magnitude
    func_log = interpolate.interp1d(magnitudes, log_saturation)
    max_log_saturation = func_log(float(magnitude))
    max_saturation = 10**(max_log_saturation)

    # Figure out what it means in wake of the given saturation lvl
    max_exptime = float(saturation_level) / max_saturation

    # Calculate the nearest number of groups
    num_groups = calc_groups_from_exp_time(max_exptime, frame_time)

    # Can't have zero groups
    num_groups = num_groups or 1

    return num_groups, max_saturation

def map_to_ta_modes(instrument, max_num_groups, min_num_groups):
    """Turns the min/max groups into the closest allowable TA group
    mode.

    Parameters
    ----------
    instrument : str
        The instrument
    max_num_groups : int
        The maximum number of groups without oversaturating
    min_num_groups : int
        The groups needed to hit the target SNR

    Returns
    -------
    min_ta_groups : int
        The min possible groups to hit target SNR
    max_ta_groups : int
        The max possible groups before saturation
    """

    # Allowable group modes for each instrument
    groups = {'miri': [3, 5, 9, 15, 23, 33, 45, 59, 75, 93, 113, 135, 159, 185, 243, 275, 513],
              'niriss': [3, 5, 7, 9, 1, 13, 15, 17, 19],
              'nirspec': [3],
              'nircam': [3, 5, 9, 17, 33, 65]}

    # Match the literal min and max groups to the nearest mode.
    allowable_groups = groups[instrument]
    min_ta_groups = min(allowable_groups, key=lambda x: abs(x - min_num_groups))
    max_ta_groups = min(allowable_groups, key=lambda x: abs(x - max_num_groups))

    # Unless it was oversaturated from the get-go OR there aren't enough groups
    # for SNR
    if min_num_groups == 0:
        min_ta_groups = 0
        max_ta_groups = 0
    if min_num_groups > max(allowable_groups):
        min_ta_groups = -1
        max_ta_groups = 0

    return max_ta_groups, min_ta_groups

def min_num_groups_for_sat(magnitude, instrument, filt, subarray, model, band, infile):
    """Estimates the minimum number of groups to reach target acq
    saturation requirements.

    Parameters
    ----------
    magnitude : float
        Magnitude of star.
    instrument : str
        Instrument.
    filt : str
        Filter.
    subarray : str
        Subarray.
    model : str
        Phoenix model key.
    band : str, currently unused?
        The band -- right now only k sooo?
    infile : str
        The file with the pandeia data.

    Returns
    -------
    min_num_groups_for_sat : int
        The minimum number of groups to reach target snr.
    """

    with open(infile) as f:
        data = json.load(f)

    # Match to closest magnitude
    magnitudes = [float(i) for i in data['mags']]
    closest_magnitude = min(magnitudes, key=lambda x: abs(x - float(magnitude)))
    index = magnitudes.index(closest_magnitude)

    # Match to data
    minimum_num_groups_for_sat = data['ta_snr'][instrument][filt][subarray][model][index]

    return minimum_num_groups_for_sat

def perform_calculation(params, frames_per_group=1, num_skips=0):
    """Calculates all of the outputs and puts them in a dictionary for
    easy access.

    Parameters
    ----------
    params : dict
        Dictionary of all the needed parameters. Must include:
        ``obs_time``, ``n_group``, ``mag``, ``mod``, ``band``,
        ``filt``, ``filt_ta``, ``ins``, ``subarray``, ``subarray_ta``,
        ``sat_mode``, ``max_sat``, ``infile``
    frames_per_group : int, optional
        The number of frames -- almost always 1
    num_skips: int, optional
        Number of skips -- almost always 0

    Returns
    -------
    params : dict, str
        Dictionary of outputs and inputs. If the calculation throws an
        error it will return a string error message instead
    """

    # TARGET ACQ
    ta_frame_time = set_frame_time(
        params['infile'], params['ins'], params['subarray_ta'], target_acq_mode=True)

    max_saturation_ta_level = convert_saturation(
        params['sat_max'], params['sat_mode'], params['ins'], params['infile'], target_acq_mode=True)

    max_num_groups, saturation_rate_ta = interpolate_from_pandeia(
        params['mag'], params['ins'], params['filt_ta'], params['subarray_ta'], params['mod'], params['band'],
        ta_frame_time, max_saturation_ta_level, params['infile'], target_acq_mode=True)

    min_num_groups = min_num_groups_for_sat(
        params['mag'], params['ins'], params['filt_ta'], params['subarray_ta'], params['mod'], params['band'], params['infile'])

    min_ta_groups, max_ta_groups = map_to_ta_modes(params['ins'], max_num_groups, min_num_groups)

    duration_time_ta_min = calc_duration_time(min_ta_groups, 1, 1, ta_frame_time)

    duration_time_ta_max = calc_duration_time(max_ta_groups, 1, 1, ta_frame_time)

    # Science obs
    # Figure out the rows/cols/amps/pixel_size and saturation
    instrument_params = set_params_from_instrument(params['ins'], params['subarray'])
    frame_time = set_frame_time(params['infile'], params['ins'], params['subarray'])

    # Run all the calculations
    num_rows, num_columns, num_amps, pixel_size, frame_time, num_reset_frames = instrument_params
    params['sat_max'] = convert_saturation(params['sat_max'], params['sat_mode'], params['ins'], params['infile'])

    # Calculate countrate and n_groups if it isn't supplied
    num_groups, saturation_rate = interpolate_from_pandeia(
        params['mag'], params['ins'], params['filt'], params['subarray'], params['mod'], params['band'],
        frame_time, params['sat_max'], params['infile'])
    if str(params['n_group']) == 'optimize':
        params['n_group'] = int(num_groups)
    else:
        params['n_group'] = int(float(params['n_group']))

    # Aditional helpful params
    # Calculate times/ramps/etc
    integration_time = calc_integration_time(params['n_group'], frame_time, frames_per_group, num_skips)
    ramp_time = calc_ramp_time(integration_time, num_reset_frames, frame_time)

    # Calculate nubmer of integrations (THE MEAT)
    num_integrations = calc_num_integrations(params['obs_time'], params['n_group'], num_reset_frames, frame_time, frames_per_group)

    # Other things that may come in handy who knows?
    exposure_time = calc_exposure_time(num_integrations, ramp_time)
    duration_time = calc_duration_time(params['n_group'], num_integrations, num_reset_frames, frame_time, frames_per_group)
    observation_efficiency = calc_observation_efficiency(exposure_time, duration_time)

    # Update params with new information
    params['duration_time'] = round(duration_time / 3600, 3)
    params['duration_time_ta_max'] = duration_time_ta_max
    params['duration_time_ta_min'] = duration_time_ta_min
    params['exposure_time'] = round(exposure_time / 3600, 3)
    params['frames_per_group'] = frames_per_group
    params['frame_time'] = round(frame_time, 3)
    params['integration_time'] = round(integration_time, 3)
    params['max_saturation_prediction'] = round(saturation_rate * frame_time * params['n_group'], 3)
    params['max_saturation_ta'] = round(saturation_rate_ta * ta_frame_time * max_ta_groups, 3)
    params['min_saturation_ta'] = round(saturation_rate_ta * ta_frame_time * min_ta_groups, 3)
    params['max_ta_groups'] = int(max_ta_groups)
    params['min_ta_groups'] = int(min_ta_groups)
    params['num_amps'] = num_amps
    params['num_columns'] = num_columns
    params['num_integrations'] = num_integrations
    params['num_reset_frames'] = num_reset_frames
    params['num_rows'] = num_rows
    params['num_skips'] = num_skips
    params['observation_efficiency'] = round(observation_efficiency, 3)
    params['ramp_time'] = round(ramp_time, 3)
    params['ta_frame_time'] = ta_frame_time

    return params

def set_params_from_instrument(instrument, subarray):
    """Sets/collects the running parameters from the instrument.

    Parameters
    ----------
    instrument : str
        Instrument, options are ``nircam``, ``niriss``, ``nirpec``, and
        ``miri``
    subarray : str
        Subarray mode

    Returns
    -------
    rows : int
        The number of pixels per row.
    cols : int
        The number of columns per row.
    amps : int
        The number of amplifiers.
    pixel_size : int
        The pixel size.
    frame_time : float
        The frame time.
    num_reset_frames : int
        The number of reset frames.
    """

    num_reset_frames = 1

    if instrument == 'nirspec':

        pixel_size = (40e-4)**2
        if subarray == 'sub2048':
            rows, cols = 2048, 32
        elif subarray in ['sub1024a', 'sub1024b']:
            rows, cols = 1024, 32
        elif subarray == 'sub512':
            rows, cols = 512, 32
        amps = 1  # 4 if not NRSRAPID????
        frame_time = calc_frame_time(cols, rows, amps, instrument)

    elif instrument == 'nircam':

        pixel_size = (18e-4)**2
        if subarray == 'full':
            rows, cols, amps = 2048, 2048, 4
        elif subarray == 'subgrism256':
            rows, cols, amps = 256, 256, 1
        elif subarray == 'subgrism128':
            rows, cols, amps = 128, 2048, 1
        elif subarray == 'subgrism64':
            rows, cols, amps = 64, 2048, 1
        frame_time = calc_frame_time(cols, rows, amps, instrument)

    elif instrument == 'miri':

        pixel_size = (25e-4)**2
        if subarray == 'slitlessprism':
            rows, cols, frame_time = 416, 72, .159
        amps = 4
        num_reset_frames = 0

    elif instrument == 'niriss':
        pixel_size = (40e-4)**2
        if subarray == 'substrip96':
            rows, cols = 2048, 96
        elif subarray == 'substrip256':
            rows, cols = 2048, 256
        amps = 1
        frame_time = calc_frame_time(cols, rows, amps, instrument)

    return rows, cols, amps, pixel_size, frame_time, num_reset_frames

def set_frame_time(infile, instrument, subarray, target_acq_mode=False):
    """Assign the appropriate frame time based on the instrument and
    subarray. For now, modes are implied.

    Parameters
    ----------
    infile: str
        The path to the data file.
    instrument : str
        The instrument : ``miri``, ``niriss``, ``nirspec``, or
        ``nircam``
    subarray : str
        The subarray
    target_acq_mode : bool
        Whether this is for TA or not.

    Returns
    -------
    frame_time : float
        The frame time for this instrument/subarray combo.
    """

    # Read in dict with frame times
    with open(infile) as f:
        frame_time = json.load(f)['frame_time']

    if target_acq_mode:
        frame_time = frame_time[instrument]['ta'][subarray]
    else:
        frame_time = frame_time[instrument][subarray]

    return frame_time


#cython: profile=True,boundscheck=False

from lattice_extension import validate_rxn as validate_rxn
from lattice_extension import propensity as rxn_propensity
from lattice_extension import update_propensities as agents_propensity
from lattice_extension import update_agents_react_prop
from lattice_extension import choose_biased_index as cbindex
from lattice_extension import pick_agent_to_act as pick_agent

from math import log
import numpy as np
import itertools as it
import time as ti

import random as rm

import modular_core.libprofile as lprf

import pdb

def set_state(state, surface_state, time, latt):
    state[0] = time
    state[1] = latt.agent_count
    state[2] = np.mean(surface_state[0])
    #state[1] = latt.total_population()
    #state[2:] = spectrum_population(
    #    latt.spec_count, latt.agent_count,
    #    latt.species_index_dict, latt.agents)
    state[3:] = latt.spectrum_population()

def set_surface_state(surf_state, time, latt):
    surf_state[0] = latt.population_grid()
    surf_state[1] = latt.identity_grid()
    if latt.has_resources:
        surf_state[2:] = latt.spectrum_resource_grid()

def capture(data, state, capture_dex, target_dexes):
    data[:,capture_dex] = [state[dex] for dex in target_dexes]

class agent(object):

    def __init__(self, variety, motion_prop = 0.0, react_prop = 0.0):
        self.variety = variety
        self.motion_prop = motion_prop
        self.react_prop = react_prop
        self.total_prop = self.react_prop + self.motion_prop
        self.behaviors = [self.move, self.react]

    def move_to(self, pos, lattice):
        site = lattice._lattice_[self.pos]
        site.agents.remove(self)
        newpopulace = site.populace.replace(self.variety, '', 1)
        site.populace = newpopulace
        site.population = len(newpopulace)
        site.who = site.identity()
        lattice.flagged_locs.append(self.pos)

        self.pos = pos
        newsite = lattice._lattice_[pos]
        newsite.agents.append(self)
        newsite.populace += self.variety
        newsite.population = len(newsite.populace)
        newsite.who = newsite.identity()
        lattice.flagged_locs.append(pos)

    def act(self, lattice):
        prtot = self.total_prop
        if prtot > 0.0:
            #normed = [self.motion_prop/prtot, self.react_prop/prtot]
            #behvs = [normed[0], normed[0] + normed[1]]
            #rd = rm.random()
            #if behvs[0] > rd:
            #    self.move(lattice)
            #elif behvs[1] > rd:
            #    self.react(lattice)
            #else: pdb.set_trace()
            behvs = [self.motion_prop, self.react_prop]
            behv = cbindex(behvs,2)
            self.behaviors[behv](lattice)
        else: print 'PASSING!'

    def move(self, lattice):
        new_pos = lattice._lattice_[self.pos].diffuse_pressure()
        #new_pos = lattice._lattice_[self.pos].diffuse_random()
        self.move_to(new_pos, lattice)

    def react(self, lattice):
        site = lattice._lattice_[self.pos]
        populace = site.populace
        resources = site.resources
        valid = [rx for rx in lattice.reactions
            if validate_rxn(rx.used, populace, resources)]
        rxn_props = [rxn_propensity(
            rx.used, populace, rx.rate)
                for rx in valid]
        rxdex = pick_rxn(rxn_props)
        rxn = valid[rxdex]
        rxn.react(site, lattice)
        lattice.flagged_locs.append(self.pos)

    def update_propensity(self, lattice):
        lattice.total_agent_propensity -= self.total_prop
        self.update_react_prop(lattice)
        self.update_motion_prop(lattice)
        self.total_prop = self.react_prop + self.motion_prop
        lattice.total_agent_propensity += self.total_prop

    def update_react_prop(self, lattice):
        site = lattice._lattice_[self.pos]
        rxns = lattice.reactions
        self.react_prop = update_agents_react_prop(
            site, rxns, lattice.reaction_count)

    def update_motion_prop(self, lattice):
        self.motion_prop = lattice._lattice_[
            self.pos].population**lattice.dimensions

class agent_immobile(agent):

    def move_to(self, pos, lattice):
        print 'immobile should not be asked to move!'

    def act(self, lattice):
        self.react(lattice)

    def update_propensity(self, lattice):
        lattice.total_agent_propensity -= self.total_prop
        self.update_react_prop(lattice)
        self.total_prop = self.react_prop
        lattice.total_agent_propensity += self.total_prop

    def update_motion_prop(self, lattice):
        print 'immobile agent should not be asked up update motion_prop!!'
        self.motion_prop = 0.0

def pick_rxn(rxn_props):
    table = []
    ptot = sum(rxn_props)
    rd = rm.uniform(0.0, ptot)
    for p in rxn_props: table.append(sum(table) + p)
    for tx, t in enumerate(table):
        if rd < t: return tx

class reaction(object):

    def __init__(self, rxn, specs):
        def convert(ag):
            stoch = int(ag[ag.find('(')+1:ag.find(')')])
            spec = ag[ag.find(')')+1:]
            isresource = not spec in specs
            return (stoch, spec, isresource)
        self.rxnstr = rxn
        rxnspl = rxn.split('->')
        front = rxnspl[0]
        self.rate = float(front[front.find('[')+1:front.find(']')])
        used, prod = front[:front.find('[')], rxnspl[1]
        if not used == '': used = used.split('+')
        else: used = []
        if not prod == '': prod = prod.split('+')
        else: prod = []
        self.used = [convert(us) for us in used]
        self.prod = [convert(pr) for pr in prod]
        self.flux = len(self.prod) - len(self.used)

    def react(self, site, lattice):
        pop = site.populace[:]
        resources = site.resources
        res_varis = [res.variety for res in resources]
        for u in self.used:
            if u[2]:
                rdex = res_varis.index(u[1])
                resource = resources[rdex]
                resource.quantity -= u[0]
            else: pop = pop.replace(u[1], '', u[0])
        for p in self.prod:
            if p[2]:
                rdex = res_varis.index(p[1])
                resource = resources[rdex]
                resource.quantity += p[0]
            else: pop += p[1]*p[0]
        site.resolve_agents(lattice, pop)

class site_resource(object):
    def __init__(self, pos, quantity = 1, variety = 'food'):
        self.variety = variety
        self.quantity = quantity

#site_topology instances describe the connectivity of an N
# dimensional lattice site to its neighboring sites
class site_topology(object):
    def __init__(self, pos, lshape, max_occupancy = None):
        self.pos = pos
        self.lshape = lshape
        self.dims = len(pos)
        self.max_occupancy = max_occupancy
        self.set_step_positions()

    def set_step_positions(self):
        pos = self.pos
        self.dn_steps = [list(pos[:]) for x in range(self.dims)]
        self.up_steps = [list(pos[:]) for x in range(self.dims)]
        axlen = self.lshape[0]
        for d in range(self.dims):
            if pos[d] == 0:
                dstep = 1 - axlen
                ustep = 1
            elif pos[d] == axlen - 1:
                dstep = 1
                ustep = 1 - axlen
            else:
                dstep = 1
                ustep = 1

            self.dn_steps[d][d] -= dstep
            self.up_steps[d][d] += ustep

        self.up_steps = [tuple(x) for x in self.up_steps]
        self.dn_steps = [tuple(x) for x in self.dn_steps]
        self.all_steps = self.up_steps + self.dn_steps
        self.neighbor_count = len(self.all_steps)

    def step(self, axdex, dir_):
        if dir_: new = self.up_steps[axdex]
        else: new = self.dn_steps[axdex]
        return new

class lattice_site(object):

    def __init__(self, args, pop = ''):
        self.position = args[0]
        self.lshape = args[1]
        self.species = args[2].species
        self.resources = args[3]
        #self.resources = [site_resource(self.position)]
        self.resource_names = [r.variety for r in self.resources]
        self.dims = len(self.position)
        self.agents = []
        self.populace = pop
        self.population = len(pop)
        self.who = self.identity()
        self.lattice = args[2]
        max_occupancy = args[2].max_occupancy
        self.topology = site_topology(
            self.position, self.lshape, max_occupancy)

    def resolve_agents(self, lattice, pop):
        local_ags = [ag for ag in self.agents if not ag.variety == '_']
        #local_ags = [ag for ag in lattice.agents if ag.pos == self.position]
        local_pop = [ag.variety for ag in local_ags]
        left = []
        for ag in local_ags:
            if ag.variety in pop: pop = pop.replace(ag.variety, '', 1)
            else: left.append(ag)

        for p in pop: lattice.put_onto_lattice_at_site(self.position, agent(p))
        for p in left: lattice.remove_from_lattice_site(self.position, p)

    def diffuse_pressure(self):
        adjacents = self.topology.all_steps
        ncnt = self.topology.neighbor_count
        latt = self.lattice._lattice_
        adjac_pops = [latt[pos].population for pos in adjacents]
        max_po = max(adjac_pops)
        adjac_pops = [-1.0*po + max_po for po in adjac_pops]
        lookup = cbindex(adjac_pops, ncnt)
        step = adjacents[lookup]
        return step

    def diffuse_random(self):
        step_dim = rm.randrange(self.dims)
        dir_ = rm.random() > 0.5
        step = self.topology.step(step_dim, dir_)
        return step

    def identity(self):
        if not self.populace: return 0
        for sdx in range(len(self.species)):
            sp = self.species[sdx]
            if sp in self.populace: return sdx + 1
        return 0

    def what(self, resource_dex):
        res = self.resources[resource_dex]
        return res.quantity

class lattice(object):

    def __init__(self, *args, **kwargs):
        if 'dims' in kwargs.keys(): dims = kwargs['dims']
        else: dims = 2
        if 'ax_length' in kwargs.keys(): ax_length = kwargs['ax_length']
        else: ax_length = 5
        if 'species' in kwargs.keys(): specs = kwargs['species']
        else: specs = []
        if 'reactions' in kwargs.keys(): rxns = kwargs['reactions']
        else: rxns = []
        if 'birthing_flag' in kwargs.keys(): bflag = kwargs['birthing_flag']
        else: bflag = False
        if 'max_occupancy' in kwargs.keys():
            max_occupancy = kwargs['max_occupancy']
        else: max_occupancy = None
        if 'resources' in kwargs.keys(): ress = kwargs['resources']
        else: ress = []
        self.agents = []
        self.agent_count = 0
        self.species = specs
        self.spec_count = len(specs)
        self.total_agent_propensity = 0
        self.agents_propensities = []
        #spdxdict = {}
        #for sdx in xrange(self.spec_count):
        #    spec = self.species[sdx]
        #    spdxdict[spec] = sdx
        #self.species_index_dict = spdxdict
        self.species_spectrum = {}
        for ke in self.species: self.species_spectrum[ke] = 0
        self.reactions = [reaction(rx, specs) for rx in rxns]
        self.reaction_count = len(rxns)
        #self.birth_reactions =\
        #    [rx for rdx, rx in enumerate(self.reactions) if rdx in brxns]
        self.birth_flag = bflag
        self.max_occupancy = max_occupancy
        self.resources = ress
        if len(ress) == 0: self.has_resources = False
        else: self.has_resources = True
        self.make_lattice(dims, ax_length, ress)

    def make_lattice(self, dims, maxax, ress):
        #must add ress contents here; right now its hardcoded
        self.dimensions = dims
        self.max_axis_size = maxax
        v_site = np.vectorize(lattice_site)
        init_arry = np.empty((maxax,)*dims,dtype = object)
        self.locs = [x for x in it.product(*[range(maxax)]*dims)]
        lattice = np.empty((maxax,)*dims,dtype = object)
        shp = lattice.shape
        for pos in self.locs:
            res = []
            for re in ress:
                if re[0] == 'all' or re[0].count(pos) > 0:
                    res.append(site_resource(pos, re[2], re[1]))
            init_arry[pos] = (pos, shp, self, res)
        lattice[:] = v_site(init_arry)
        self._lattice_ = lattice
        self.flagged_locs = []
        if self.birth_flag:
            for pos in self.locs:
                bagent = agent_immobile('_')
                self.put_onto_lattice_at_site(pos, bagent)
                self.flagged_locs.append(pos)

    def put_onto_lattice_at_site(self, pos, agent):
        self.agents.append(agent)
        if pos == 'random':pos = rm.choice(self.locs)
        agent.pos = pos
        site = self._lattice_[pos]
        site.agents.append(agent)
        if not agent.variety == '_':
            self.species_spectrum[agent.variety] += 1
            site.populace += agent.variety
            self.agent_count += 1
        site.population = len(site.populace)
        site.who = site.identity()

    def remove_from_lattice_site(self, pos, agent):
        self.species_spectrum[agent.variety] -= 1
        site = self._lattice_[pos]
        site.populace = site.populace.replace(agent.variety, '', 1)
        site.population = len(site.populace)
        site.who = site.identity()
        site.agents.remove(agent)
        self.agents.remove(agent)
        self.agent_count -= 1

    def total_propensity(self):
        return self.total_agent_propensity

    def act_agent(self, agdex):
        agent = self.agents[agdex]
        agent.act(self)

    def get_clean_fixed_surf(self):
        def fix(d):
            if d in relev: return slice(None)
            else: return fixins[d]
        fixins = [0 for d in range(self.dimensions)]
        relev = [0,1]
        fixed = [fix(d) for d in range(len(fixins))]
        maxax = self.max_axis_size
        surf = np.zeros((maxax, maxax), dtype = np.int)
        return surf, fixed

    def population_grid(self):
        surf, fixed = self.get_clean_fixed_surf()
        surf[:,:] = [[sit.population for sit in ax]
            for ax in self._lattice_[fixed]]
        return surf

    def identity_grid(self):
        surf, fixed = self.get_clean_fixed_surf()
        surf[:,:] = [[sit.who for sit in ax]
            for ax in self._lattice_[fixed]]
        return surf

    def spectrum_population(self):
        spectrum = [0]*self.spec_count
        for sdx in xrange(self.spec_count):
            sp = self.species[sdx]
            spectrum[sdx] = self.species_spectrum[sp]
        return spectrum

    def spectrum_resource_grid(self):
        def resource_grid(res):
            surf, fixed = self.get_clean_fixed_surf()
            surf[:,:] = [[sit.what(res) for sit in ax]
                for ax in self._lattice_[fixed]]
            return surf
        grids = []
        #for res in self.resources:
        for rsdx in range(len(self.resources)):
            grids.append(resource_grid(rsdx))
        return grids

def_string = ''
def simulate(sys_string = def_string):
    species = ['A', 'B', 'C']
    #agents = [
    #    ((1,3),'A'),((4,3),'B'),((8,7),'C'),
    #    ((95,7),'A'),((84,9),'B'),((88,2),'C'),
    #    ((93,83),'A'),((84,93),'B'),((88,87),'C'),
    #    ((2,93),'A'),((6,83),'B'),((4,97),'C'),
    #        ]
    agents =\
        [('random', 'A')]*3333 +\
        [('random', 'B')]*3333 +\
        [('random', 'C')]*3333
    reactions = [
        #'(1)A+(1)food[5.0]->(2)A',
        #'(1)B+(1)food[5.0]->(2)B',
        #'(1)C+(1)food[5.0]->(2)C',
        '(1)A+(1)B[5.0]->(2)A',
        '(1)B+(1)C[5.0]->(2)B',
        '(1)C+(1)A[5.0]->(2)C',
        #'(1)A[0.01]->',
        #'(1)B[0.01]->',
        #'(1)C[0.01]->',
        #'[0.01]->(1)A',
        #'[0.01]->(1)B',
        #'[0.01]->(1)C',
                ]
    #birthing_rxns = [6]
    birthing_flag = False#True
    resources = []#[('all', 'food', 3)]
    lattice_dims = 2
    lattice_size = 100
    max_occupancy = 2
    timed_out = False
    timed_out_limit = 3600.0

    # I TREAT AGENTS AT EACH LATTICE SITE AS INDISTINGUISHABLE
    #  BUT THE LOCAL AGENTS ARE REMOVED IN REVERSE ORDER OF ENTERING THE SITE
    ### to properly include birth reactions - treat lattice sites as agents
    ###  with zero motion propensity
    # still want occupancy handling
    # species can only have a single letter to identify
    # voxel visualization pipeline for 3-d subspaces (perfectly analogous to
    #  surface_data pipeline)
    # statistics on average agent (average agent trajectory really....)
    # add full function/variable pipeline for reactions
    # add full mcfg/sys_string pipeline
    # introduce concept of nations (groups of agents)
    # i double count propensities since which agent initiates a reaction
    #  is never relevant - only the variety of agents post reaction

    time = 0.0
    last_time = 0.0
    end_time = 50.0
    incr_time = 0.1

    total_captures = end_time/incr_time
    capture_count = 0
    res_names = [r[1] for r in resources]
    targets = ['time', 'population', 'mean_occupancy'] + species[:]# + res_names
    plot_targets = ['time', 'population', 'mean_occupancy'] + species[:]# + res_names
    surf_targets = [
        'population_surfaces',
        'identity_surfaces'] + ['resource_surface_' + r for r in res_names]
    plot_surf_targets = [
        'population_surfaces',
        'identity_surfaces'] + ['resource_surface_' + r for r in res_names]
    #voxel_targets = ['identity_voxels']
    #plot_voxel_targets = ['identity_voxels']
    target_dexes = [targets.index(ta) for ta in plot_targets]
    surf_target_dexes = [surf_targets.index(ta) for ta in plot_surf_targets]
    #voxel_target_dexes = [voxel_targets.index(ta) for ta in plot_voxel_targets]
    _lattice_ = lattice(
        dims = lattice_dims, ax_length = lattice_size, species = species,
        reactions = reactions, max_occupancy = max_occupancy,
        birthing_flag = birthing_flag, resources = resources)
    for ag in agents: _lattice_.put_onto_lattice_at_site(ag[0],agent(ag[1]))
    for ag in _lattice_.agents: ag.update_propensity(_lattice_)
    maxax = _lattice_.max_axis_size
    state = np.zeros(shape = (len(targets)), dtype = np.float)
    surface_state = np.zeros(shape=(len(surf_targets),maxax,maxax),dtype=np.float)
    #voxel_state = np.zeros(shape=(len(voxel_targets),maxax,maxax,maxax),dtype=np.float)
    data = np.zeros(shape=(len(plot_targets),total_captures),dtype=np.float)
    surface_data = np.zeros(shape = (len(surf_targets),
    	total_captures, maxax, maxax), dtype = np.float)
    #voxel_data = np.zeros(shape = (len(voxel_targets),
    #	total_captures, maxax, maxax, maxax), dtype = np.float)
    #set_voxel_state(voxel_state, time, _lattice_)
    set_surface_state(surface_state, time, _lattice_)
    set_state(state, surface_state, time, _lattice_)

    start_time = ti.time()
    if birthing_flag: dead_eco_check = lambda : False
    else: dead_eco_check = lambda : not _lattice_.agent_count > 0
    dead_eco = dead_eco_check()
    while capture_count < total_captures and not dead_eco and not timed_out:
        propensity_total = _lattice_.total_propensity()
        if propensity_total > 0.0 and _lattice_.agent_count > 0:
            propensity_total_inv = 1.0/propensity_total
            time_step = -1.0 * log(rm.random()) * propensity_total_inv
            acnt = len(_lattice_.agents)
            agrand = rm.uniform(0.0, propensity_total)
            agent_dex = pick_agent(_lattice_.agents, acnt, agrand)
        else:
            time_step = incr_time
            agent_dex = -1

        time += time_step
        #set_surface_state(surface_state, time, _lattice_)
        set_state(state, surface_state, time, _lattice_)

        real_time = state[0]
        if last_time < real_time and capture_count < total_captures:
            #set_voxel_state(voxel_state, time, _lattice_)
            set_surface_state(surface_state, time, _lattice_)
            set_state(state, surface_state, time, _lattice_)
        while last_time < real_time and capture_count < total_captures:
            state[0] = last_time
            last_time += incr_time
            capture(data, state, capture_count, target_dexes)
            capture(surface_data, surface_state, capture_count, surf_target_dexes)
            #capture(voxel_data, voxel_state, capture_count, voxel_target_dexes)
            capture_count += 1
            print 'time', last_time, 'pop', len(_lattice_.agents)
        state[0] = real_time

        #agent activity cant depend on state/surface_state...
        if agent_dex >= 0: _lattice_.act_agent(agent_dex)
        agents_propensity(_lattice_)

        if ti.time() - start_time > timed_out_limit: timed_out = True
        dead_eco = dead_eco_check()

    if timed_out or dead_eco:
        #set_surface_state(surface_state, time, _lattice_)
        state[0] = last_time
        last_time += incr_time
        capture(data, state, capture_count, target_dexes)
        capture(surface_data, surface_state, capture_count, surf_target_dexes)
        #capture(voxel_data, voxel_state, capture_count, voxel_target_dexes)
        capture_count += 1
        print 'time', last_time, 'pop', len(_lattice_.agents)

        toss = capture_count
        #voxel_data = voxel_data[:,:toss,:,:]
        surface_data = surface_data[:,:toss,:,:]
        data = data[:,:toss]
        return data,surface_data,targets + surf_targets
    return data, surface_data, targets + surf_targets

if __name__ == '__main__':
    import sys
    if len(sys.argv) > 1: use_pyx = sys.argv[1]
    else: use_pyx = False
    if use_pyx:
        import lattice_simulator as lsim
        print 'USING CYTHON!'
        lprf.profile_function(lsim.simulate)
    else:
        lprf.profile_function(simulate)
    #lprf.profile_function(simulate)


""" Functions used to assign restframe g-r and r-i SDSS colors to mock galaxies.
"""
import numpy as np
from .sawtooth_binning import sawtooth_bin_indices
from astropy.utils.misc import NumpyRNGContext
from halotools.empirical_models import polynomial_from_table
from ..sdss_colors.sdss_completeness_model import retrieve_sdss_sample_mask

__all__ = ('mc_sdss_gr_ri', )

default_seed = 43

def assign_data_source(mock_logsm, table_abscissa=np.array([8.5, 9, 9.5, 10]),
            table_ordinates=np.array([1, 0.8, 0.35, 0]), seed=default_seed):
    """
    Determine the source of observational data that will be used
    to map colors onto mock galaxies.

    For entries in the returned ndarray equal to zero, real SDSS objects
    will be used to map colors onto mock galaxies.

    For entries in the returned ndarray equal to one, fake SDSS objects will be used.

    Parameters
    ----------
    mock_logsm : ndarray
        Numpy array of shape (ngals, ) storing the stellar mass of mock galaxies

    table_abscissa : ndarray, optional
        Control points in log10 stellar mass at which the data source probability is defined

    table_ordinates : ndarray, optional
        Control points defining the probability that
        a mock galaxy will be assigned to data source one.

    seed : int, optional
        Random number seed. Default is default_seed, set at the top of
        the module where the function is defined.

    Returns
    -------
    data_source : ndarray
        Numpy array of shape (ngals, ) storing an integer indicating the data source
        from which colors will be drawn
    """
    input_abscissa = np.linspace(table_abscissa.min(), table_abscissa.max(), 100)
    ordinates = polynomial_from_table(table_abscissa, table_ordinates, input_abscissa)

    prob_fake = np.interp(mock_logsm, input_abscissa, ordinates)
    with NumpyRNGContext(seed):
        fake_mask = np.random.rand(len(mock_logsm)) < prob_fake
    data_source = np.zeros_like(fake_mask).astype(int)
    data_source[fake_mask] = 1

    return data_source

def fuzzy_sawtooth_magr_binning(mock_rmag, data_source, magr_bins=None, seed=default_seed):
    """ Assign galaxies to overlapping bins based on their restframe absolute r-band magnitude.

    Binning will be done separately for mock galaxies above and below the SDSS completeness limit.

    Parameters
    ----------
    mock_rmag : ndarray
        Numpy array of shape (ngals, ) storing the
        restframe absolute r-band magnitude of mock galaxies

    data_source : ndarray
        Numpy array of shape (ngals, ) storing an integer indicating the data source
        from which colors will be drawn (zero for real data, one for fake data).

    magr_bins : ndarray, optional
        Numpy array of shape (nbins, ) storing the bin boundaries.
        Must strictly encompass the range spanned by the mock.
        Default is to use 25 bins linearly spaced in Magr covering a range that
        just beyond the boundaries of the data.

    seed : int, optional
        Random number seed. Default is default_seed, set at the top of
        the module where the function is defined.

    Returns
    -------
    rmag_bin_number : ndarray
        Numpy integer array of shape (ngals, ) storing the bin number of each mock galaxy
    """
    if magr_bins is None:
        epsilon = 0.01
        rmin, rmax, dr = mock_rmag.min()-epsilon, mock_rmag.max()+epsilon, 0.2
        magr_bins = np.arange(rmin, rmax+dr, dr)

    source0_mask = data_source == 0
    rmag_bin_number = -np.ones_like(mock_rmag).astype('i4')
    rmag_bin_number[source0_mask] = sawtooth_bin_indices(mock_rmag[source0_mask], magr_bins, seed=seed)
    rmag_bin_number[~source0_mask] = sawtooth_bin_indices(mock_rmag[~source0_mask], magr_bins, seed=seed)
    return rmag_bin_number

def mc_true_sdss_gr_ri(sdss_redshift, sdss_magr, sdss_gr, sdss_ri,
        mock_magr_bin_number, mock_magr, mock_sfr_percentile, sigma=0., k=10, seed=default_seed):
    """ Given {Mr, SFR-percentile} for mock galaxies, search SDSS data
    to select galaxies with matching {Mr, g-r-percentile}, and use the colors
    {g-r, r-i} of the selected SDSS galaxies to paint on to the mock galaxies.
    """
    from scipy.spatial import cKDTree
    from scipy.stats import gaussian_kde

    mock_gr = np.zeros_like(mock_magr)
    mock_ri = np.zeros_like(mock_magr)

    #  We loop over probabilistically overlapping bins of r-band magnitude
    magr_bin_numbers = list(set(mock_magr_bin_number))
    for bin_number in magr_bin_numbers:
        mock_magr_mask = mock_magr_bin_number == bin_number
        npts_mock_bin = np.count_nonzero(mock_magr_mask)

        mock_magr_bin = mock_magr[mock_magr_mask]
        mock_sfr_percentile_bin = mock_sfr_percentile[mock_magr_mask]

        #  Make a cut to apply a color-completeness limiting redshift
        Mr_min, Mr_max = mock_magr_bin.min(), mock_magr_bin.max()
        sdss_mask = retrieve_sdss_sample_mask(sdss_redshift, sdss_magr, Mr_min, Mr_max)
        sdss_rmag_bin = sdss_magr[sdss_mask]
        sdss_gr_bin = sdss_gr[sdss_mask]
        sdss_ri_bin = sdss_ri[sdss_mask]

        #  Use Gaussian kernel density estimation on the SDSS data
        #  to generate a random sampling of colors for each mock galaxy
        #  This step helps smooth out artificial clustering of mock galaxies
        #  around outlier fluctuations in the observed data
        X_bin = np.vstack((sdss_rmag_bin, sdss_gr_bin, sdss_ri_bin))
        kde_bin = gaussian_kde(X_bin)
        with NumpyRNGContext(seed):
            resampled_bin = kde_bin.resample(npts_mock_bin)
        sdss_rmag_bin_resampled = resampled_bin[0, :]
        sdss_gr_bin_resampled = resampled_bin[1, :]
        sdss_ri_bin_resampled = resampled_bin[2, :]

        #  Sort the mock galaxies in the bin by their SFR-rank-order-percentile (at fixed M*)
        #  Sort the Monte Carlo realization of g-r
        #  Apply the implicit non-parametric map to paint g-r onto mock galaxies
        idx_mock_percentile_sorted = np.argsort(mock_sfr_percentile_bin)
        mock_gr_bin = np.zeros(npts_mock_bin).astype('f4')
        mock_gr_bin[idx_mock_percentile_sorted] = np.sort(sdss_gr_bin_resampled)[::-1]

        #  Now run a nearest-neighbor search on the KDE-resampled SDSS data
        #  to identify an r-i color of a galaxy with a closely matching {Mr, g-r}.
        #  Randomly selecting one of the 10 nearest neighbors helps smear out
        #  artificial amplification of hard edges in the observed data
        Y = np.vstack((sdss_rmag_bin_resampled, sdss_gr_bin_resampled)).T
        sdss_tree = cKDTree(Y)
        knn = min(k, len(sdss_rmag_bin_resampled))
        result = sdss_tree.query(np.vstack((mock_magr_bin, mock_gr_bin)).T, k=knn)
        nn_indices = result[1]
        with NumpyRNGContext(seed):
            a = np.random.randint(0, nn_indices.shape[1], nn_indices.shape[0])
        idx = nn_indices[np.arange(nn_indices.shape[0]), a]
        mock_ri_bin = sdss_ri_bin_resampled[idx]

        mock_gr[mock_magr_mask] = mock_gr_bin
        mock_ri[mock_magr_mask] = mock_ri_bin

    return mock_gr, mock_ri

def mc_fake_sdss_gr_ri(sdss_gr, sdss_ri, rmag_bin_number, mock_log10_mstar, seed=default_seed):
    """ Set up a simple multivariate Gaussian model to extrapolate SDSS colors
    {g-r, r-i} into the very faint end.
    """
    gr_center = np.median(sdss_gr)-0.35
    ri_center = np.median(sdss_ri)-0.15
    median_gr = np.interp(mock_log10_mstar, [6, 9], [gr_center-0.2, gr_center])
    median_ri = np.interp(mock_log10_mstar, [6, 9], [ri_center-0.3, ri_center])
    median_array = np.vstack((median_gr, median_ri)).T

    ngals_mock = len(mock_log10_mstar)
    X = np.vstack((sdss_gr, sdss_ri))
    cov = np.cov(X)/2.

    with NumpyRNGContext(seed):
        Z = np.random.multivariate_normal(
            mean=(0, 0), cov=cov, size=ngals_mock) + median_array
    mock_gr, mock_ri = Z[:, 0], Z[:, 1]
    return mock_gr, mock_ri

def mc_sdss_gr_ri(mock_rmag, mock_mstar, mock_sfr_percentile,
            sdss_redshift, sdss_magr, sdss_gr, sdss_ri, k=10, seed=default_seed):
    """ Divide mock galaxies into a sample for which colors from real SDSS galaxies
    will be drawn, and another sample for which colors from the extrapolation model
    will be drawn, creating a fuzzy boundary in stellar mass that stitches the
    two samples together. Then for each sample, call the appropriate
    Monte Carlo function to generate {g-r, r-i} colors.
    """
    mock_data_source = assign_data_source(np.log10(mock_mstar), seed=seed)
    mock_rmag_bin_number = fuzzy_sawtooth_magr_binning(mock_rmag, mock_data_source, seed=seed)

    source0_mask = mock_data_source == 0

    mock_source0_gr, mock_source0_ri = mc_true_sdss_gr_ri(
        sdss_redshift, sdss_magr, sdss_gr, sdss_ri,
        mock_rmag_bin_number[source0_mask], mock_rmag[source0_mask],
        mock_sfr_percentile[source0_mask], k, seed=seed)

    mock_source1_gr, mock_source1_ri = mc_fake_sdss_gr_ri(sdss_gr, sdss_ri,
            mock_rmag_bin_number[~source0_mask], np.log10(mock_mstar[~source0_mask]), seed=seed)

    output_mock_gr = np.zeros_like(mock_rmag) - 999.
    output_mock_ri = np.zeros_like(mock_rmag) - 999.
    output_mock_gr[source0_mask] = mock_source0_gr
    output_mock_ri[source0_mask] = mock_source0_ri
    output_mock_gr[~source0_mask] = mock_source1_gr
    output_mock_ri[~source0_mask] = mock_source1_ri

    return output_mock_gr, output_mock_ri


"""
This is to generate interaction energies and corresponding translational vectors,
given a fixed receptor and an ensemble of ligand coordinates (including rotations and/or configurations)
"""
from __future__ import print_function

import numpy as np
import netCDF4

try:
    from bpmfwfft.grids import RecGrid
    from bpmfwfft.grids import LigGrid

except:
    from grids import RecGrid
    from grids import LigGrid

KB = 0.001987204134799235

class Sampling(object):
    def __init__(self, rec_prmtop, lj_sigma_scal_fact, rec_inpcrd,
                        bsite_file, grid_nc_file,
                        lig_prmtop, lig_inpcrd,
                        lig_coord_ensemble,
                        energy_sample_size_per_ligand,
                        output_nc,
                        temperature=300.):
        """
        :param rec_prmtop: str, name of receptor prmtop file
        :param lj_sigma_scal_fact: float, used to check consitency when loading receptor and ligand grids
        :param rec_inpcrd: str, name of receptor inpcrd file
        :param bsite_file: None or str, name of file defining the box, the same as
        from AlGDock pipeline. "measured_binding_site.py"
        :param grid_nc_file: str, name of receptor precomputed grid netCDF file
        :param lig_prmtop: str, name of ligand prmtop file
        :param lig_inpcrd: str, name of ligand inpcrd file
        :param lig_coord_ensemble: list of 2d array, each array is an ligand coordinate
        :param energy_sample_size_per_ligand: int, number of energies and translational vectors to store for each ligand crd
        :param output_nc: str, name of nc file
        :param temperature: float
        """
        self._energy_sample_size_per_ligand = energy_sample_size_per_ligand
        self._beta = 1./ temperature / KB

        rec_grid = self._create_rec_grid(rec_prmtop, lj_sigma_scal_fact, rec_inpcrd,
                                        bsite_file, grid_nc_file)
        self._rec_crd = rec_grid.get_crd()

        self._lig_grid = self._create_lig_grid(lig_prmtop, lj_sigma_scal_fact, lig_inpcrd, rec_grid)

        self._lig_coord_ensemble = self._load_ligand_coor_ensemble(lig_coord_ensemble)

        self._nc_handle = self._initialize_nc(output_nc)

    def _create_rec_grid(self, rec_prmtop, lj_sigma_scal_fact, rec_inpcrd, bsite_file, grid_nc_file):
        rec_grid = RecGrid(rec_prmtop, lj_sigma_scal_fact, rec_inpcrd, bsite_file,
                            grid_nc_file, new_calculation=False)
        return rec_grid

    def _create_lig_grid(self, lig_prmtop, lj_sigma_scal_fact, lig_inpcrd, rec_grid):
        lig_grid = LigGrid(lig_prmtop, lj_sigma_scal_fact, lig_inpcrd, rec_grid)
        return lig_grid

    def _load_ligand_coor_ensemble(self, lig_coord_ensemble):
        assert len(lig_coord_ensemble.shape) == 3, "lig_coord_ensemble must be 3-D array."
        ensemble = lig_coord_ensemble
        natoms = self._lig_grid.get_natoms()

        for i in range(len(ensemble)):
            if (ensemble[i].shape[0] != natoms) or (ensemble[i].shape[1] != 3):
                raise RuntimeError("Ligand crd %d does not have correct shape"%i)
        return ensemble

    def _initialize_nc(self, output_nc):
        nc_handle = netCDF4.Dataset(output_nc, mode="w", format="NETCDF4")

        nc_handle.createDimension("three", 3)
        rec_natoms = self._rec_crd.shape[0]
        nc_handle.createDimension("rec_natoms", rec_natoms)

        lig_natoms = self._lig_grid.get_natoms()
        nc_handle.createDimension("lig_natoms", lig_natoms)
        nc_handle.createDimension("lig_sample_size", self._lig_coord_ensemble.shape[0])

        nc_handle.createDimension("energy_sample_size_per_ligand", self._energy_sample_size_per_ligand)

        nc_handle.createVariable("rec_positions", "f8", ("rec_natoms", "three"))
        nc_handle.variables["rec_positions"][:,:] = self._rec_crd

        nc_handle.createVariable("lig_positions", "f8", ("lig_sample_size", "lig_natoms", "three"))
        nc_handle.createVariable("lig_com", "f8", ("lig_sample_size", "three"))
        nc_handle.createVariable("volume", "f8", ("lig_sample_size"))
        nc_handle.createVariable("nr_grid_points", "i8", ("lig_sample_size"))

        nc_handle.createVariable("exponential_sums", "f8", ("lig_sample_size"))
        nc_handle.createVariable("log_of_divisors",  "f8", ("lig_sample_size"))

        nc_handle.createVariable("mean_energy",  "f8", ("lig_sample_size"))
        nc_handle.createVariable("min_energy",  "f8", ("lig_sample_size"))
        nc_handle.createVariable("energy_std",  "f8", ("lig_sample_size"))

        nc_handle.createVariable("resampled_energies", "f8", ("lig_sample_size", "energy_sample_size_per_ligand"))
        nc_handle.createVariable("resampled_trans_vectors", "i8", ("lig_sample_size", "energy_sample_size_per_ligand", "three"))

        nc_handle = self._write_grid_info(nc_handle)
        return nc_handle

    def _write_grid_info(self, nc_handle):
        """
        write grid info, "x", "y", "z" ...
        """
        data = self._lig_grid.get_grids()
        grid_func_names = self._lig_grid.get_grid_func_names()
        keys = [key for key in data.keys() if key not in grid_func_names]

        for key in keys:
            for dim in data[key].shape:
                dim_name = "%d"%dim
                if dim_name not in nc_handle.dimensions.keys():
                    nc_handle.createDimension(dim_name, dim)

        for key in keys:
            if data[key].dtype == int:
                store_format = "i8"
            elif data[key].dtype == float:
                store_format = "f8"
            else:
                raise RuntimeError( "Unsupported dtype %s"%data[key].dtype )
            dimensions = tuple([ "%d"%dim for dim in data[key].shape ])
            nc_handle.createVariable(key, store_format, dimensions)

        for key in keys:
            nc_handle.variables[key][:] = data[key]
        return nc_handle

    def _save_data_to_nc(self, step):
        self._nc_handle.variables["lig_positions"][step, :, :] = self._lig_grid.get_crd()

        self._nc_handle.variables["lig_com"][step, :] = self._lig_grid.get_initial_com()

        self._nc_handle.variables["volume"][step] = self._lig_grid.get_box_volume()

        self._nc_handle.variables["nr_grid_points"][step] = self._lig_grid.get_number_translations()

        self._nc_handle.variables["exponential_sums"][step] = self._exponential_sum

        self._nc_handle.variables["log_of_divisors"][step] = self._log_of_divisor

        self._nc_handle.variables["mean_energy"][step] = self._mean_energy
        self._nc_handle.variables["min_energy"][step] = self._min_energy
        self._nc_handle.variables["energy_std"][step] = self._energy_std

        self._nc_handle.variables["resampled_energies"][step,:] = self._resampled_energies

        self._nc_handle.variables["resampled_trans_vectors"][step,:,:] = self._resampled_trans_vectors
        return None

    def _do_fft(self, step):
        print("Doing FFT for step %d"%step, "test")
        lig_conf = self._lig_coord_ensemble[step]
        self._lig_grid.cal_grids(molecular_coord = lig_conf)

        energies = self._lig_grid.get_meaningful_energies()
        print("Energies shape:", energies.shape)
        self._mean_energy = energies.mean()
        self._min_energy  = energies.min()
        self._energy_std  = energies.std()
        print("Number of finite energy samples", energies.shape[0])

        exp_energies = -self._beta * energies
        print(f"Max exp energy {exp_energies.max()}, Min exp energy {exp_energies.min()}")
        self._log_of_divisor = exp_energies.max()
        exp_energies[exp_energies < 0] = 0
        exp_energies = np.exp(exp_energies - self._log_of_divisor)
        self._exponential_sum = exp_energies.sum()
        exp_energies /= self._exponential_sum
        print("Number of exponential energy samples", exp_energies.sum())
        # sel_ind = np.random.choice(exp_energies.shape[0], size=self._energy_sample_size_per_ligand, p=exp_energies, replace=True)
        try:
            sel_ind = np.random.choice(exp_energies.shape[0], size=self._energy_sample_size_per_ligand, p=exp_energies, replace=False)
        except:
            print(f"Only {np.count_nonzero(exp_energies)} non-zero entries in p, falling back to replacement")
            sel_ind = np.random.choice(exp_energies.shape[0], size=self._energy_sample_size_per_ligand, p=exp_energies, replace=True)

        del exp_energies

        self._resampled_energies = [energies[ind] for ind in sel_ind]
        del energies
        self._lig_grid.set_meaningful_energies_to_none()

        trans_vectors = self._lig_grid.get_meaningful_corners()
        self._resampled_trans_vectors = [trans_vectors[ind] for ind in sel_ind]
        del trans_vectors

        self._resampled_energies = np.array(self._resampled_energies, dtype=float)
        self._resampled_trans_vectors = np.array(self._resampled_trans_vectors, dtype=int)

        self._save_data_to_nc(step)
        return None

    def run_sampling(self):
        """
        """
        for step in range(self._lig_coord_ensemble.shape[0]):
            self._do_fft(step)

            print("Min energy", self._min_energy)
            print("Mean energy", self._mean_energy)
            print("STD energy", self._energy_std)
            print("Initial center of mass", self._lig_grid.get_initial_com())
            print("Grid volume", self._lig_grid.get_box_volume())
            print("Number of translations", self._lig_grid.get_number_translations())
            print("-------------------------------\n\n")

        self._nc_handle.close()
        return None

#
#TODO   the class above assumes that the resample size is smaller than number of meaningful energies
#       in general, the number of meaningful energies can be very smaller or even zero (no energy)
#       when the number of meaningful energies is zero, that stratum contributes n_points zeros to the exponential mean
#
#       so when needs to consider separately 3 cases:
#           len(meaningful energies) == 0
#           0< len(meaningful energies) <= resample size
#           len(meaningful energies) > resample size
#

class Sampling_PL(Sampling):

    def _write_data_key_2_nc(self, data, key):
        if data.shape[0] == 0:
            return None

        for dim in data.shape:
            dim_name = "%d"%dim
            if dim_name not in self._nc_handle.dimensions.keys():
                self._nc_handle.createDimension(dim_name, dim)

        if data.dtype == int:
            store_format = "i8"
        elif data.dtype == float:
            store_format = "f8"
        else:
            raise RuntimeError("unsupported dtype %s"%data.dtype)
        dimensions = tuple(["%d"%dim for dim in data.shape])
        self._nc_handle.createVariable(key, store_format, dimensions)

        self._nc_handle.variables[key][:] = data
        return None

    def _initialize_nc(self, output_nc):
        """
        """
        nc_handle = netCDF4.Dataset(output_nc, mode="w", format="NETCDF4")

        nc_handle.createDimension("three", 3)
        rec_natoms = self._rec_crd.shape[0]
        nc_handle.createDimension("rec_natoms", rec_natoms)

        lig_natoms = self._lig_grid.get_natoms()
        nc_handle.createDimension("lig_natoms", lig_natoms)
        nc_handle.createDimension("lig_sample_size", self._lig_coord_ensemble.shape[0])

        #nc_handle.createDimension("energy_sample_size_per_ligand", self._energy_sample_size_per_ligand)

        nc_handle.createVariable("rec_positions", "f8", ("rec_natoms", "three"))
        nc_handle.variables["rec_positions"][:,:] = self._rec_crd

        nc_handle.createVariable("lig_positions", "f8", ("lig_sample_size", "lig_natoms", "three"))
        nc_handle.createVariable("lig_com", "f8", ("lig_sample_size", "three"))
        nc_handle.createVariable("volume", "f8", ("lig_sample_size"))
        nc_handle.createVariable("nr_grid_points", "i8", ("lig_sample_size"))
        nc_handle.createVariable("nr_finite_energy", "i8", ("lig_sample_size"))

        nc_handle.createVariable("exponential_sums", "f8", ("lig_sample_size"))
        nc_handle.createVariable("log_of_divisors",  "f8", ("lig_sample_size"))

        nc_handle.createVariable("mean_energy",  "f8", ("lig_sample_size"))
        nc_handle.createVariable("min_energy",  "f8", ("lig_sample_size"))
        nc_handle.createVariable("energy_std",  "f8", ("lig_sample_size"))

        #nc_handle.createVariable("resampled_energies", "f8", ("lig_sample_size", "energy_sample_size_per_ligand"))
        #nc_handle.createVariable("resampled_trans_vectors", "i8", ("lig_sample_size", "energy_sample_size_per_ligand", "three"))

        nc_handle = self._write_grid_info(nc_handle)
        return nc_handle

    def _save_data_to_nc(self, step):
        self._nc_handle.variables["lig_positions"][step, :, :] = self._lig_grid.get_crd()

        self._nc_handle.variables["lig_com"][step, :] = self._lig_grid.get_initial_com()

        self._nc_handle.variables["volume"][step] = self._lig_grid.get_box_volume()

        self._nc_handle.variables["nr_grid_points"][step] = self._lig_grid.get_number_translations()

        self._nc_handle.variables["nr_finite_energy"][step] = self._nr_finite_energy

        self._nc_handle.variables["exponential_sums"][step] = self._exponential_sum

        self._nc_handle.variables["log_of_divisors"][step] = self._log_of_divisor

        self._nc_handle.variables["mean_energy"][step] = self._mean_energy

        self._nc_handle.variables["min_energy"][step] = self._min_energy

        self._nc_handle.variables["energy_std"][step] = self._energy_std

        self._write_data_key_2_nc(self._resampled_energies, "resampled_energies_%d"%step)

        self._write_data_key_2_nc(self._resampled_trans_vectors, "resampled_trans_vectors_%d"%step)
        return None

    def _do_fft(self, step):
        print("Doing FFT for step %d"%step)
        lig_conf = self._lig_coord_ensemble[step]
        print(self._lig_grid["SASAr"])
        self._lig_grid.cal_grids(molecular_coord = lig_conf)

        energies = self._lig_grid.get_meaningful_energies()
        self._nr_finite_energy = energies.shape[0]
        print("Number of finite energy samples", self._nr_finite_energy)

        if energies.shape[0] > 0:

            self._mean_energy = energies.mean()
            self._min_energy  = energies.min()
            self._energy_std  = energies.std()

            exp_energies = -self._beta * energies
            self._log_of_divisor = exp_energies.max()
            exp_energies = np.exp(exp_energies - self._log_of_divisor)
            self._exponential_sum = exp_energies.sum()
            exp_energies /= self._exponential_sum

            sample_size = min(exp_energies.shape[0], self._energy_sample_size_per_ligand)
            sel_ind = np.random.choice(exp_energies.shape[0], size=sample_size, p=exp_energies, replace=True)

            del exp_energies

            self._resampled_energies = [energies[ind] for ind in sel_ind]
            del energies
            self._lig_grid.set_meaningful_energies_to_none()

            trans_vectors = self._lig_grid.get_meaningful_corners()
            self._resampled_trans_vectors = [trans_vectors[ind] for ind in sel_ind]
            del trans_vectors

            self._resampled_energies = np.array(self._resampled_energies, dtype=float)
            self._resampled_trans_vectors = np.array(self._resampled_trans_vectors, dtype=int)

        else:

            self._mean_energy = np.inf
            self._min_energy  = np.inf
            self._energy_std  = np.inf

            self._log_of_divisor  = 1.
            self._exponential_sum = 0.

            self._resampled_energies = np.array([], dtype=float)
            del energies
            self._lig_grid.set_meaningful_energies_to_none()

            self._resampled_trans_vectors = np.array([], dtype=float)

        self._save_data_to_nc(step)
        return None

if __name__ == "__main__":
    # test
    rec_prmtop = "../examples/amber/ubiquitin_ligase/receptor.prmtop"
    lj_sigma_scal_fact = 0.8
    rec_inpcrd = "../examples/amber/ubiquitin_ligase/receptor.inpcrd"

    # bsite_file = "../examples/amber/t4_lysozyme/measured_binding_site.py"
    bsite_file = None
    grid_nc_file = "../examples/grid/ubiquitin_ligase/grid.nc"

    lig_prmtop = "../examples/amber/ubiquitin/ligand.prmtop"
    lig_inpcrd = "../examples/amber/ubiquitin/ligand.inpcrd"

    energy_sample_size_per_ligand = 200
    output_nc = "../examples/fft_sampling/ubql_ubiquitin/fft_sampling.nc"

    ligand_md_trj_file = "../examples/ligand_md/ubiquitin/rotation.nc"
    lig_coord_ensemble = netCDF4.Dataset(ligand_md_trj_file, "r").variables["positions"][:]

    sampler = Sampling(rec_prmtop, lj_sigma_scal_fact, rec_inpcrd,
                        bsite_file, grid_nc_file,
                        lig_prmtop, lig_inpcrd,
                        lig_coord_ensemble,
                        energy_sample_size_per_ligand,
                        output_nc,
                        temperature=300.)
    sampler.run_sampling()


'''Collection of functions that generate figures from the fit outputs. Creates figures from fit results using both simulated and real datasets.'''

from fancy.analysis import results
import numpy as np
import os
import h5py
import matplotlib as mpl
from matplotlib import pyplot as plt
import seaborn as sns
from pandas import DataFrame
from astropy.coordinates import SkyCoord
from astropy import units as u

from fancy import Data, Results
from fancy.plotting import Corner
from fancy.plotting.allskymap_cartopy import AllSkyMapCartopy as AllSkyMap
from fancy.plotting.colours import *
from fancy.interfaces.stan import Direction

# use minimalist style
plt.style.use("minimalist")

# paths to important files
path_to_this_file = os.path.abspath(os.path.dirname(__file__))
stan_path = os.path.join(path_to_this_file, "..", "stan")
source_file = os.path.join(path_to_this_file, "..", "data", "sourcedata.h5")
uhecr_file = os.path.join(path_to_this_file, "..", "data", "UHECRdata.h5")
table_path = os.path.join(path_to_this_file, "..", "tables")
output_path = os.path.join(path_to_this_file, "..", "output")

class OutputFigures():
    def __init__(self,
                 fig_args,
                 sim_model=None,
                 sim_output_file=None,
                 tight_B=False,
                 skymap_label_fmt="default"):
        '''Collection of figures obtained from output of fitting.'''
        # parameters set from argparse
        self.source_type = fig_args["source"]
        self.detector_type = fig_args["detector"]
        self.model_type = fig_args["model"]
        self.data_type = fig_args["dtype"]
        self.ptype = fig_args["ptype"]
        self.seed = fig_args["seed"]
        self.sim_model_type = sim_model
        self.sim_output_file = sim_output_file

        # output file from simulation / fitting
        if tight_B:
            self.output_file = os.path.join(
                output_path,
                "tmp_{0}_fit_{5}_{1}_{2}_{3}_{4}_tightB.h5".format(
                    self.model_type, self.source_type, self.detector_type,
                    self.seed, self.ptype, self.data_type))
        else:
            self.output_file = os.path.join(
                output_path, "tmp_{0}_fit_{5}_{1}_{2}_{3}_{4}.h5".format(
                    self.model_type, self.source_type, self.detector_type,
                    self.seed, self.ptype, self.data_type))

        if sim_output_file is not None:
            self.output_file = os.path.join(
                output_path, "{0}_fit_{5}_{1}_{2}_{3}_{4}_{6}.h5".format(
                    self.model_type, self.source_type, self.detector_type,
                    self.seed, self.ptype, self.data_type,
                    self.sim_model_type))
            # if self.detector_type == "joint_gmf":
            #     self.output_file = os.path.join(
            #         output_path, "{0}_fit_{5}_{1}_{2}_{3}_{4}_{5}.h5".format(
            #             self.model_type, self.source_type, self.detector_type,
            #             self.seed, self.ptype, self.data_type,
            #             self.sim_model_type))

            # else:
            #     self.output_file = os.path.join(
            #         output_path, "{0}_fit_{5}_{1}_{2}_{3}_{4}.h5".format(
            #             self.model_type, self.source_type, self.detector_type,
            #             self.seed, self.ptype, self.data_type))

        # obtain detector properties / params from imports
        self.detector_properties, self.detector_params = self._get_detectorimports(
        )

        # initialize the data object and add relevant data into it
        self._initialize_data()

        # format style for skymap labels (from "mpl" or "TA")
        self.skymap_label_fmt = skymap_label_fmt

    def _get_detectorimports(self):
        '''Get variables imported by (detector_name).py'''
        if self.detector_type == "TA2015":
            from fancy.detector.TA2015 import detector_properties, detector_params
        elif self.detector_type == "auger2014":
            from fancy.detector.auger2014 import detector_properties, detector_params
        elif self.detector_type == "auger2010":
            from fancy.detector.auger2010 import detector_properties, detector_params
        else:
            raise Exception("Undefined detector type!")

        return detector_properties, detector_params

    def _initialize_data(self):
        '''Initialize data object based on sim / real data'''
        self.data = Data()

        if self.data_type == "sim":
            self.data.from_file(self.sim_output_file)

        elif self.data_type == "real":
            self.data.add_source(source_file, self.source_type)
            self.data.add_uhecr(uhecr_file, self.detector_type)
            self.data.add_detector(self.detector_properties)

    def src_uhecr_skymap(self, savefile=None, coord="G", exposure="map"):
        '''
        Plot skymap with sources + UHECR, obtained from Data object.
        Optionally add either exposure map or exposure limit.
        '''
        # labels
        src_label = self.data.source.label
        uhecr_label = self.data.uhecr.label
        # read in initializations
        omega_src = Direction(self.data.source.unit_vector)
        omega_arr = Direction(self.data.uhecr.unit_vector)
        energy_arr = self.data.uhecr.energy

        # convert source / arrival directions to equatorial / galactic
        if coord == "G":
            x_src, y_src = omega_src.glons, omega_src.glats
            x_arr, y_arr = omega_arr.glons, omega_arr.glats
        elif coord == "E":
            x_src, y_src = omega_src.ras, omega_src.decs
            x_arr, y_arr = omega_arr.ras, omega_arr.decs

        Eth = self.data.detector.Eth
        Emax = np.ceil(np.max(energy_arr) / 10.) * 10.

        uhecr_color = [lightblue, midblue, darkblue]
        uhecr_cmap = mpl.colors.ListedColormap(uhecr_color)
        energy_bins = np.logspace(np.log(Eth), np.log(Emax), 4, base=np.e)
        uhecr_norm = mpl.colors.BoundaryNorm(energy_bins, uhecr_cmap.N)

        # Legend
        legend_elements = [
            mpl.lines.Line2D([0], [0],
                             marker='o',
                             color='w',
                             label='sources',
                             markersize=10,
                             markerfacecolor='k'),
            mpl.lines.Line2D([0], [0],
                             marker='o',
                             color='w',
                             label='UHECRs',
                             markersize=15,
                             markerfacecolor=midblue,
                             alpha=0.8)
        ]

        # create skymap
        skymap = AllSkyMap(projection='moll', lon_0=180)
        skymap.set_gridlines(label_fmt=self.skymap_label_fmt)

        # sources
        skymap.scatter(x_src, y_src, s=10.0, color='k', alpha=1.0, zorder=5)

        # UHECRs
        for lon, lat, E in np.nditer([x_arr, y_arr, energy_arr]):
            i = np.digitize(E, energy_bins) - 1
            skymap.tissot(lon,
                          lat,
                          3.0 + (i * 2),
                          30,
                          facecolor=uhecr_cmap.colors[i],
                          alpha=0.8,
                          zorder=i + 2)

        # exposure
        if exposure == "map":
            skymap.exposure_map(self.data.detector.params, coord=coord)
        elif exposure == "limit":
            skymap.exposure_limit(self.data.detector.limiting_dec.deg,
                                  coord=coord,
                                  s=2,
                                  color=grey,
                                  alpha=1)
        else:
            raise ValueError(
                "Exposure plot type {0} not defined.".format(exposure))

        # Annotations
        skymap.legend(handles=legend_elements,
                      loc='upper right',
                      bbox_to_anchor=(1., 1.),
                      fontsize=16,
                      fancybox=True)
        skymap.title("{0} + {1}".format(src_label, uhecr_label))

        # Colorbar
        cb_ax = plt.axes([0.25, 0.07, .5, .05], frameon=False)
        bar = mpl.colorbar.ColorbarBase(cb_ax,
                                        norm=uhecr_norm,
                                        cmap=uhecr_cmap,
                                        orientation='horizontal',
                                        drawedges=True,
                                        alpha=1)
        bar.set_label('$\hat{E}$ / EeV', color='k', fontsize=16)
        bar.ax.tick_params(labelsize=16)

        skymap.save(savefile, bbox_inches='tight')

    def eval_association_probs(self):
        '''
        Evaluate the assocation probabilities between sources and UHECR
        from some output file
        '''

        # Log probability
        results = Results(self.output_file)
        keys = ['lp']
        chain = results.get_chain(keys)
        logprob = chain['lp'].transpose(1, 2, 0)
        N = np.shape(logprob)[0]

        # Account for background component
        Ns = np.shape(logprob)[1] - 1

        # Calculate association probabilities for each source-UHECR combo
        uhecr_p = []
        for lp in logprob:
            lps = []
            for src in range(Ns + 1):
                lps.append(np.mean(np.exp(lp[src])))

            norm = sum(lps)
            ps = []
            for src in range(Ns + 1):
                ps.append(lps[src] / norm)
            uhecr_p.append(ps)

        # Normalise line weights
        pmax = max(max(uhecr_p))

        # Find names of dominant sources
        threshold_probability = 0.1

        dominant = []
        for p in uhecr_p:
            # for i in range(data.source.N):
            for i in range(Ns):
                if p[i] > threshold_probability:
                    dominant.append(i)

        seen = set()
        inds = []
        for d in dominant:
            if d not in seen:
                inds.append(d)
                seen.add(d)

        # dominant_sources = [self.data.source.name[i] for i in inds]

        # sort so that those with largest assossations appear first
        N_assos = {}
        for i in inds:
            N_assos[self.data.source.name[i].decode("UTF-8")] = len(
                np.argwhere([d == i for d in dominant]))

        N_assos_sorted = {
            k: v
            for k, v in sorted(N_assos.items(), key=lambda item: item[1])[::-1]
        }

        # print("Dominant sources: ", [self.data.source.name[i] for i in inds])

        return uhecr_p, N_assos_sorted

    def association_skymap(self, savefile, coord="G"):
        '''Plot association skymap between sources and UHECRs'''

        uhecr_p, N_assos_sorted = self.eval_association_probs()

        # labels
        # src_label = self.data.source.label

        src_label = "Swift-BAT" if self.data.source.label == "swift_BAT_213" else self.data.source.label.split(
            "_")[0]

        # uhecr_label = self.data.uhecr.label
        # read in initializations
        omega_src = Direction(self.data.source.unit_vector)
        omega_arr = Direction(self.data.uhecr.unit_vector)
        energy_arr = self.data.uhecr.energy

        # convert source / arrival directions to equatorial / galactic
        if coord == "G":
            x_src, y_src = 180. - omega_src.glons, omega_src.glats
            x_arr, y_arr = 180. - omega_arr.glons, omega_arr.glats
        elif coord == "E":
            x_src, y_src = omega_src.ras, omega_src.decs
            x_arr, y_arr = omega_arr.ras, omega_arr.decs

        Eth = self.data.detector.Eth
        Emax = np.ceil(np.max(energy_arr) / 10.) * 10.

        uhecr_color = [lightblue, midblue, darkblue]
        uhecr_cmap = mpl.colors.ListedColormap(uhecr_color)
        energy_bins = np.logspace(np.log(Eth), np.log(Emax), 4, base=np.e)
        uhecr_norm = mpl.colors.BoundaryNorm(energy_bins, uhecr_cmap.N)

        # Legend
        legend_elements = [
            mpl.lines.Line2D([0], [0],
                             marker='o',
                             color='w',
                             label=src_label,
                             markersize=10,
                             markerfacecolor='k'),
            mpl.lines.Line2D([0], [0],
                             marker='o',
                             color='w',
                             label='UHECRs',
                             markersize=15,
                             markerfacecolor=midblue,
                             alpha=0.8)
        ]

        # plot
        skymap = AllSkyMap(projection='moll', lon_0=180)
        skymap.set_gridlines(label_fmt=self.skymap_label_fmt)

        # sources
        skymap.scatter(x_src, y_src, s=10.0, color='k', alpha=1.0, zorder=5)

        # UHECRs
        for lon, lat, E in np.nditer([x_arr, y_arr, energy_arr]):
            i = np.digitize(E, energy_bins) - 1
            skymap.tissot(lon,
                          lat,
                          3.0 + (i * 2),
                          30,
                          facecolor=uhecr_cmap.colors[i],
                          alpha=0.8,
                          zorder=i + 2)

        # Association
        # add some way to include dominant sources with this association plot later
        pmax = np.max(np.max(uhecr_p))

        for i, p in enumerate(uhecr_p):
            for j, psrc in enumerate(p[0:self.data.source.N]):
                if psrc > 0.001:
                    skymap.geodesic(x_arr[i],
                                    y_arr[i],
                                    x_src[j],
                                    y_src[j],
                                    color='k',
                                    lw=3,
                                    alpha=psrc / pmax,
                                    zorder=10)

        # Annotations
        skymap.legend(handles=legend_elements,
                      loc='upper right',
                      bbox_to_anchor=(1.1, 1.),
                      fontsize=16,
                      fancybox=True)

        # exposure limit
        skymap.exposure_limit(
            self.data.detector.limiting_dec.deg,
            coord=coord,
            s=0.5,
            marker="o",
            color=lightpurple,
            alpha=0.01,
        )

        # Colorbar
        cb_ax = plt.axes([0.25, 0.07, .5, .05], frameon=True)
        bar = mpl.colorbar.ColorbarBase(cb_ax,
                                        norm=uhecr_norm,
                                        cmap=uhecr_cmap,
                                        orientation='horizontal',
                                        drawedges=True,
                                        alpha=1)
        bar.set_label('$\hat{E}$ / EeV', color='k', fontsize=16)
        bar.ax.tick_params(labelsize=16)
        skymap.save(savefile)

        return N_assos_sorted

    def corner(self, savefile):
        '''Plot corner plot'''

        if self.dtype == "sim":
            self.corner_sim(savefile)
        elif self.dtype == "real":
            self.corner_data(savefile)

    def corner_data(self, savefile):
        '''Plot corner plot of data'''
        # Get chains from joint fit
        results_fit = Results(self.output_file)

        # get keys and corresponding labels
        labels = {}
        if self.model_type.find("join") != -1:
            keys = ['alpha', 'B', 'f']

            labels['B'] = r'$B$ / $\mathrm{nG}$'
            labels['alpha'] = r'$\alpha$'
            labels['f'] = r'$f$'

        elif self.model_type == "arrival_direction":
            keys = ['kappa', 'L', 'f']

            labels['L'] = r'$L$'
            labels['kappa'] = r'$\kappa$'
            labels['f'] = r'$f$'

        chain = results_fit.get_chain(keys)

        # Make nicely labelled dict
        chain_for_df = {}
        for key in keys:
            chain_for_df[labels[key]] = chain[key]

        # Make ordered dataframe
        df = DataFrame(data=chain_for_df)
        df = df[[labels[keys[0]], labels[keys[1]], labels[keys[2]]]]

        corner = Corner(df, color=midblue, contour_color=midblue_contour)
        corner.save(savefile)

    def corner_sim(self, savefile, cumul=False, output_files=None):

        if cumul:
            self.corner_sim_cumul(output_files, savefile)
        else:
            results_fit = Results(self.output_file)

            results_sim = Results(self.sim_output_file)
            truth_keys = ['F0', 'L', 'alpha', 'B', 'f']
            truth = results_sim.get_truths(truth_keys)
            info_keys = ['Eth', 'Eth_sim']
            info = results_sim.get_truths(info_keys)

            # Correct for different Eth in sim and fit
            # Also scale to plot units
            flux_scale = (info['Eth'] / info['Eth_sim'])**(1 - truth['alpha'])
            truth['F0'] = truth['F0'] * flux_scale  # km^-2 yr^-1
            truth[
                'L'] = truth['L'][0] * flux_scale / 1.0e39 * 10  # 10^-38 yr^-1

            if self.model_type == "arrival_direction":
                self.corner_sim_arrival(results_fit, truth, savefile)
            else:
                self.corner_sim_joint(results_fit, truth, savefile)

    def corner_sim_arrival(self, results_fit, truth, savefile):
        '''Corner plot for Arrival Direction model'''

        keys = ['kappa', 'L', 'f']
        chain = results_fit.get_chain(keys)

        chain['L'] = chain['L'] * 10  # 10^-38 yr^-1

        labels = {}
        labels['L'] = r'$L$ / $10^{38}$ $\mathrm{yr}^{-1}$'
        labels['kappa'] = r'$\kappa$'
        labels['f'] = r'$f$'

        truths = [truth["L"], truth["f"]]

        # Make nicely labelled dict
        chain_for_df = {}
        for key in keys:
            chain_for_df[labels[key]] = chain[key]

        # Make ordered dataframe
        df = DataFrame(data=chain_for_df)
        df = df[[labels['L'], labels['f'], labels['kappa']]]

        corner = Corner(df,
                        truths,
                        color=midblue,
                        contour_color=midblue_contour)

        corner.save(savefile)

    def corner_sim_joint(self, results_fit, truth, savefile):
        '''Plot corner plot resulting from simulations'''

        keys = ['F0', 'L', 'alpha', 'B', 'f']
        chain = results_fit.get_chain(keys)

        # Convert form Stan units to plot units
        chain['F0'] = chain['F0'] / 1.0e3  # km^-2 yr^-1
        chain['L'] = chain['L'] * 10  # 10^-38 yr^-1

        labels = {}
        labels['L'] = r'$L$ / $10^{38}$ $\mathrm{yr}^{-1}$'
        labels['F0'] = r'$F_0$ / $\mathrm{km}^{-2} \ \mathrm{yr}^{-1}$'
        labels['B'] = r'$B$ / $\mathrm{nG}$'
        labels['alpha'] = r'$\alpha$'
        labels['f'] = r'$f$'

        params = np.column_stack([chain[key] for key in keys])
        truths = [truth[key] for key in keys]

        # Make nicely labelled dict
        chain_for_df = {}
        for key in keys:
            chain_for_df[labels[key]] = chain[key]

        # Make ordered dataframe
        df = DataFrame(data=chain_for_df)
        df = df[[
            labels['F0'], labels['L'], labels['alpha'], labels['B'],
            labels['f']
        ]]

        corner = Corner(df, truths, color=purple, contour_color=purple_contour)
        corner.save(savefile)

    def corner_sim_cumul(self, output_files, savefile):
        '''Cumulative distribution of corner plots from simulation'''

        keys = ['F0', 'L', 'alpha', 'B', 'f']
        chain_avgs = {key: 0 for key in keys}
        chain_list = []
        Nseeds = len(output_files)

        # get chains for each output file
        # ignores self.outpuf_file
        for output_file in output_files:
            chain = Results(output_file).get_chain(keys)

            chain_list.append(chain)

        # evaluate averaged value for each seed
        for key in keys:
            chain_sum = 0
            for i in range(Nseeds):
                chain_sum += chain_list[i][key]

            chain_sum /= Nseeds

            chain_avgs[key] = chain_sum

        # Convert form Stan units to plot units
        chain_avgs['F0'] = chain_avgs['F0'] / 1.0e3  # km^-2 yr^-1
        chain_avgs['L'] = chain_avgs['L'] * 10  # 10^-38 yr^-1

        # Get truths from simulation
        results_sim = Results(self.sim_output_file)

        truth_keys = ['F0', 'L', 'alpha', 'B', 'f']
        truth = results_sim.get_truths(truth_keys)
        info_keys = ['Eth', 'Eth_sim']
        info = results_sim.get_truths(info_keys)

        # Correct for different Eth in sim and fit
        # Also scale to plot units
        flux_scale = (info['Eth'] / info['Eth_sim'])**(1 - truth['alpha'])
        truth['F0'] = truth['F0'] * flux_scale  # km^-2 yr^-1
        truth['L'] = truth['L'][0] * flux_scale / 1.0e39 * 10  # 10^-38 yr^-1

        labels = {}
        labels['L'] = r'$L$ / $10^{38}$ $\mathrm{yr}^{-1}$'
        labels['F0'] = r'$F_0$ / $\mathrm{km}^{-2} \ \mathrm{yr}^{-1}$'
        labels['B'] = r'$B$ / $\mathrm{nG}$'
        labels['alpha'] = r'$\alpha$'
        labels['f'] = r'$f$'

        params = np.column_stack([chain_avgs[key] for key in keys])
        truths = [truth[key] for key in keys]

        # Make nicely labelled dict
        chain_for_df = {}
        for key in keys:
            chain_for_df[labels[key]] = chain_avgs[key]

        # Make ordered dataframe
        df = DataFrame(data=chain_for_df)
        df = df[[
            labels['F0'], labels['L'], labels['alpha'], labels['B'],
            labels['f']
        ]]

        corner = Corner(df, truths, color=purple, contour_color=purple_contour)
        corner.save(savefile)

class SourceUHECRDist():
    def __init__(self, figsize=(6, 4), sim_output_file=None):
        '''Class that organizes plotting of source-UHECR association fraction distribution'''
        self.fig, self.ax = plt.subplots(figsize=figsize)

        self.color_list = [grey, lightpurple, lightblue]

        self.sim_output_file = sim_output_file

    def get_fs(self, fname_list, cumul=False):
        '''Get source-UHECR association fraction from output files'''

        self.f_list = []

        if cumul:  # evaluate cumulative distribution
            # then each element in fname_list must be a list of fnames with different models
            # each of this list will contain fnames with different seeds
            for fname_models_list in fname_list:
                f_list = []
                for output_file in fname_models_list:
                    f_i = Results(output_file).get_chain(['f'])['f']
                    f_list.append(f_i)

                f_avg = np.mean(np.array(f_list), axis=0)
                self.f_list.append(f_avg)

        else:
            for output_file in fname_list:
                try:
                    f_i = Results(output_file).get_chain(['f'])['f']
                except KeyError:
                    f_i = np.zeros(100)
                self.f_list.append(f_i)

        return self.f_list

    def plotdist(self, labels, extend=False, title=None):
        '''Plot the distribution'''
        for i, label in enumerate(labels):
            # for arrival_direction -> arrival direction
            label = label.replace("_", " ") if "_" in label else label

            sns.distplot(self.f_list[i],
                         hist=False,
                         kde_kws={
                             'shade': True,
                             'lw': 2,
                             'zorder': i
                         },
                         color=self.color_list[i],
                         label=label)

        if self.sim_output_file is not None:
            f_true = Results(self.sim_output_file).get_truths(['f'])['f']
            self.ax.axvline(f_true,
                            0,
                            10,
                            color='k',
                            zorder=3,
                            lw=2.,
                            alpha=0.7)

        self._annotate(title=title, extend=extend)

    def _annotate(self, title=None, extend=False):
        '''Plot the annotations'''
        # annotations
        self.ax.set_xlim(0, 1)
        if title:
            self.ax.set_title(title, fontsize=24)
        self.ax.set_xlabel('$f$')
        self.ax.set_ylabel('$P(f | \hat{E}, \hat{\omega})$')

        legend_loc = (1.6, 1) if extend else (1.3, 1)
        self.ax.legend(fontsize=22, bbox_to_anchor=legend_loc)

    def save(self, savefile):
        '''Save the distribution'''
        # save
        self.fig.savefig(savefile, bbox_inches='tight')