#import yt.mods as yt
import yt
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import glob


__all__ = ['future_snr', 'snr']


_core_collapse_labels = ["SNII", "II", "2", "SN_II", "TypeII", "Type 2",
                         "Type II", "type II", "typeII", 'core collapse']

_snia_labels = ["SN1a", "SNIa", "Type1a", "TypeIa", "Type Ia", "Type 1a",
                 "type 1a", "type Ia", "type ia", "type1a", "typeIa"]

_agb_labels = ['AGB', 'agb']


def future_snr(ds, data, times = None, sn_type = 'II'):
    """
    Looks forward from current time to compute future (projected) SN
    rate
    """

    current_time  = ds.current_time.convert_to_units('Myr').value

    if times is None:
        bin_spacing = 2.0* yt.units.Myr
        times = np.arange(np.min(creation_time) - bin_spacing*2.0, currentTime, bin_spacing)*yt.units.Myr
    elif np.size(times) == 1:
        bin_spacing = times
        if not hasattr(bin_spacing, 'value'):
            bin_spacing = bin_spacing * yt.units.Myr

        times = np.linspace(current_time, current_time + 2000.0, bin_spacing)
        times = times * yt.units.Myr

    birth_mass    = data['birth_mass'].value
    mass          = data['particle_mass'].convert_to_units('Msun').value
    creation_time = data['creation_time'].convert_to_units('Myr').value
    lifetimes     = data['dynamical_time'].convert_to_units('Myr').value
    pt            = data['particle_type']


    if  any( [sn_type in x for x in _core_collapse_labels]):

        collapse_threshold = ds.parameters['IndividualStarDirectCollapseThreshold']
        agb_threshold      = ds.parameters['IndividualStarSNIIMassCutoff']

        pcut = (pt == 11) * (birth_mass <= collapse_threshold) *\
                            (birth_mass  > agb_threshold)

    elif any( [sn_type in x for x in _snia_labels]):

        pcut = (pt == 12) * (mass > 0.0)

    elif any( [sn_type in x for x in _agb_labels]):

        pcut = (pt == 11) * (birth_mass < agb_threshold)

    explosion_times = creation_time[pcut] + lifetimes[pcut]
    explosion_times = explosion_times * yt.units.Myr

    times = times.convert_to_units('yr')
    snr   = np.zeros(np.size(times.value) - 1)

    # compute SNR
    for i in np.arange(np.size(times) - 1):
        dt = times[i+1] - times[i]
        dN = np.size( explosion_times[explosion_times <= times[i+1]]) -\
             np.size( explosion_times[explosion_times <= times[i]])

        snr[i] = dN / dt

    return times, snr


def snr(ds, data, times = None, sn_type = 'II'):
    """
    Computes the supernova rate of the desired time for a given dataset
    as a function of time. The way the particle types and particle lifetimes
    are handled, this can be done for the entire galaxy history using a single
    snapshot, rather than having to sort through each dump.

    One can provide sample times using "times" argument, or leave it alone for
    a 10 Myr sample spacing from t = 0 to t = current_time. If a single value
    is provided, this is taken to be the sample spacing (dt), sampled over
    t = 0 to t = current_time. Units are assumed to be Myr if not provided.

    Accounts for direct collapse model in computing SNII rates using
    parameter file.
    """

    current_time  = ds.current_time.convert_to_units('Myr').value

    if times is None:
        bin_spacing = 10.0 * yt.units.Myr
        times = np.linspace(np.min(creation_time), current_time, bin_spacing)*yt.units.Myr
    elif np.size(times) == 1:
        bin_spacing = times
        if not hasattr(bin_spacing, 'value'):
            bin_spacing = bin_spacing * yt.units.Myr

        times = np.linspace(np.min(creation_time), current_time, bin_spacing)
        times = times *yt.units.Myr


    # load particle properties
    birth_mass    = data['birth_mass'].value
    mass          = data['particle_mass'].convert_to_units("Msun").value
    creation_time = data['creation_time'].convert_to_units('Myr').value
    metallicity   = data['metallicity_fraction'].value
#    lifetimes     = data['dynamical_time'].convert_to_units('Myr').value
    lifetimes     = data[('io','particle_model_lifetime')].convert_to_units('Myr').value
    pt            = data['particle_type'].value

    # check to see if there are any SN candidates in the first place
#    if not any([ == x for x in np.unique(pt)]):
#        print "no supernova of type " + sn_type + " found"
#        return times, np.zeros(np.size(times.value) - 1)

    # looking for core collapse supernova rate
    if  any( [sn_type in x for x in _core_collapse_labels]):

        pcut = (pt == 13)

        # ignore stars that did not actually go supernova
        collapse_threshold = ds.parameters['IndividualStarDirectCollapseThreshold']
        agb_threshold      = ds.parameters['IndividualStarSNIIMassCutoff']
        if not any([(x <= collapse_threshold)*(x > agb_threshold) for x in birth_mass[pcut]]):
            print("no core collapse supernova present, only direct collapse")
            return times, np.zeros(np.size(times.value) - 1)

        # slice!
        pcut *= (birth_mass <= collapse_threshold)*(birth_mass > agb_threshold)

    elif any( [sn_type in x for x in _snia_labels]):

        pcut = (pt == 12)

        if np.size(mass[pcut]) < 1:
            return times, np.zeros(np.size(times))
        
        # SNIa are the ones that are just masless tracers, rest are WD
        if not any(mass[pcut] == 0.0):
            print("no Type Ia supernova, only white dwarfs")
            print("N_WD = %i -- Lowest mass = %.3f Msun"%(np.size(mass[pcut]), np.min(mass[pcut])))
            print("Current time = %.2E Myr - Next to explode at t = %.2E Myr"%(current_time, np.min(lifetimes[pcut] + creation_time[pcut])))
            return times, np.zeros(np.size(times.value) - 1)

        # slice!
        pcut *= (mass == 0.0)

    elif any( [sn_type in x for x in _agb_labels]):
        agb_threshold	   = ds.parameters['IndividualStarSNIIMassCutoff']

        pcut = (pt > 11)  # all dead stars
        pcut = pcut * (birth_mass <= agb_threshold)

 #       pcut  = (pt == 12)
 #       pcut *= (mass > 0.0)

  #      pcut = pcut + ( (pt == 13) * (birth_mass <= agb_threshold))

    else:
        print("sn_type :" + sn_type + " not a valid option - check spelling")
        return -1


    #
    # now get the explosion times for all supernova
    # when stars go SN, lifetime is set to be lifetime*huge_number
    # therefore, explosion time can be backed out as:
    #
    explosion_times = creation_time[pcut] + lifetimes[pcut]/ds.parameters['huge_number']
    explosion_times = explosion_times * yt.units.Myr

    times = times.convert_to_units('yr')
    snr   = np.zeros(np.size(times.value) - 1)

    # compute SNR
    for i in np.arange(np.size(times) - 1):
        dt = times[i+1] - times[i]
        dN = np.size( explosion_times[explosion_times <= times[i+1]]) -\
             np.size( explosion_times[explosion_times <= times[i]])

        snr[i] = dN / dt

    return times, snr


if __name__ == '__main__':
    # example usage - uses most recent data file

    log = False

    ds_list = np.sort( glob.glob('./DD????/DD????'))

    ds   = yt.load(ds_list[-1])
    data = ds.all_data()

    dt = 25.0
    times = np.arange(0.0, ds.current_time.convert_to_units('Myr').value + dt, dt)
    times = times*yt.units.Myr

    times, snrII = snr(ds, data, times = times, sn_type = 'TypeII')
    times, snrIa = snr(ds, data, times = times, sn_type = "TypeIa")

    center = 0.5 * (times[1:] + times[:-1])

    fig, ax = plt.subplots(figsize=(8,8))

    snialabel = 'Type Ia x 10'
    sniilabel = 'Core Collapse'

    ftimes = np.arange(ds.current_time.convert_to_units('Myr').value,
                             ds.current_time.convert_to_units('Myr').value + 800.0 + 10, 10)
    ftimes = ftimes * yt.units.Myr
    ftimes, fsnrII = future_snr(ds, data, times = ftimes, sn_type = 'TypeII')
    ftimes, fsnrIa = future_snr(ds, data, times = ftimes, sn_type = 'TypeIa')

    if log:
        ax.plot(center/1.0E6, snrII*1.0E6, color = 'black', lw = 3, ls = '-', label = sniilabel)
        ax.plot(center/1.0E6, snrIa*1.0E6*10, color = 'black', lw = 3, ls = '--', label = snialabel)
        x.semilogy()
    else:
        ax.step(times[:-1]/1.0E6, snrII*1.0E6, color ='black', lw = 3, ls = '-', label = sniilabel)
        ax.step(times[:-1]/1.0E6, snrIa*1.0E6 * 10, color ='orange', lw = 3, ls = '-', label = snialabel)
        ax.step(ftimes[:-1]/1.0E6, fsnrII*1.0E6, color = 'black', lw = 3, ls = ':')
        ax.step(ftimes[:-1]/1.0E6, fsnrIa*1.0E6 * 10, color = 'orange', lw = 3, ls = ':')

    ax.set_xlabel('Time (Myr)')

    ax.set_ylabel(r'SNR (Myr$^{-1}$)')
    ax.set_ylim( np.min( [np.min(snrIa), np.min(snrII)])*1.0E6,
                 np.max( [np.max(snrIa), np.max(snrII)])*1.25*1.0E6)
    ax.plot( [ds.current_time.convert_to_units('Myr').value]*2, ax.get_ylim(), ls = '--', lw = 3, color = 'black')

    ax.legend(loc ='best')
    plt.tight_layout()
    ax.minorticks_on()
    plt.savefig('snr.png')
