#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
schedmaster.py: This code identifies candidates for followup with 12 hours of
further VLA time on FERMI sources, then helps to create a schedule, testing
that it's possible. If everything looks good, we send the schedule to the
to be converted to OPT format and write it to disk. Note that this code serves
also as a prototype for the pythonic OPT project.
"""

__author__ = "Seth Bruzewski"
__email__ = "bruzewskis@unm.edu"
__license__ = "GPL"

# Third Party Modules
import numpy as np
import matplotlib.pyplot as plt
from astropy.coordinates import SkyCoord,Angle,EarthLocation
from astropy.table import Table
import astropy.units as u
from scipy.optimize import brentq

# Custom Modules
from SimEVLA import EquitorialToHorizontal as E2H, MotionSimulator
from SimEVLA import PointToPointTime as P2P
from optscheduler import write_OPT

def block_argsort(ra, dec):
    '''
    Sort the sources into blocks of RA, moving up and down in declination.
    With the name 'argsort' borrowed from numpy, this function returns the
    indices in the sorted order, allowing for arrays or tables to be resorted
    easily.

    Args
    =======
    ra (array) - The RA of the sources to be sorted
    dec (array) - The dec of the sources to be sorted

    Returns
    =======
    new_order (array) - The sorted indices, to be used in sorting objects

    Raises
    =======
    None
    '''

    # Divide into hours
    min_block = min(ra)//30*30
    max_block = max(ra)//30*30+30
    num_hours = (max_block - min_block) / 30 # 30 = 2 hr blocks
    div = np.linspace(min_block, max_block, num_hours+1)

    # Generate index
    index = np.arange(len(dec))

    new_order = np.array([], dtype=int)
    for i in range(len(div)-1):
        # Find indices inside RA range
        inrange = np.logical_and(ra>div[i], ra<=div[i+1])
        ind = index[inrange]

        # Sort RA filtered indices by DEC
        inorder = ind[np.argsort(dec[ind])]

        # Flip every other one
        if i%2!=0:
            inorder = np.flip(inorder)

        # Append
        new_order = np.append(new_order, inorder)

    return new_order

def plot_sky(sky, show_tour=False):
    '''
    From input coordinate pairs expressed in degrees, convert to the
    appropriate system and plot in a Mollweide projection friendly way.

    Args
    =======
    sky (array) - A (2,N) array of points in RA and DEC to be plotted.
    show_tour (bool) - If True, plot lines connecting points in order.

    Returns
    =======
    None

    Raises
    =======
    None
    '''

    ra = Angle(sky[:,0], unit='deg').wrap_at('180d').rad
    dec = np.deg2rad(sky[:,1])

    plt.scatter(ra,dec,s=10)

    if show_tour:
        for i in range(len(sky)-1):
            plt.plot([ra[i],ra[i+1]], [dec[i], dec[i+1]])

    return None

def make_block(data, ramin=None, ramax=None, dcmin=None, dcmax=None):
    '''
    Returns a subset of `data` fitting just inside the bounds provided as
    inputs. Note that because ra is a wrapping coordinate system, some
    considerations must be made. Namely when we say "angle a is between amin
    and amax," we mean more specifically "angle a is counterclockwise of amin
    and clockwise of amax." The sources are downselected using this method,
    then sorted by their RA and DEC into convenient blocks

    Args
    =======
    data (Table) - The data to be filtered down.
    ramin (float) - The minimum RA, specified in degrees.
    ramax (float) - The maximum RA, specified in degrees.
    dcmin (float) - The minimum Dec, specified in degrees.
    dcmax (float) - The maximum Dec, specified in degrees.

    Returns
    =======
    outdata (Table) - The input `data` table, filtered and sorted

    Raises
    =======
    None
    '''

    # Filter for Valid RA
    ra = data['RAJ2000']
    ra_minus = (ra-ramin)%360
    ra_max_minus = (ramax-ramin)%360
    ra_inbounds = np.logical_and(ra_minus>0, ra_minus<ra_max_minus)

    # Filter for Valid Dec
    dec = data['DEJ2000']
    dec_inbounds = np.logical_and(dec>dcmin, dec<dcmax)

    # All sources in bounds
    inblock = np.logical_and(ra_inbounds, dec_inbounds)

    # Find coordinates inside requested area
    inblock_ind = np.nonzero(inblock)[0]
    ra_inblock = data['RAJ2000'][inblock_ind]
    dec_inblock = data['DEJ2000'][inblock_ind]

    # Figure out best place to wrap
    radmin = np.deg2rad(ramin)
    radmax = np.deg2rad(ramax)
    dist = np.arctan2(np.sin(radmax-radmin), np.cos(radmax-radmin))
    wrap = radmin + dist%(2*np.pi)/2 + np.pi

    # Wrap at best spot then sort into RA-DEC blocks
    ra_wrapped = Angle(ra_inblock).wrap_at(wrap*u.rad).deg
    blocksort_ind = block_argsort(ra_wrapped, dec_inblock)
    master_ind = inblock_ind[blocksort_ind]

    # Generate data subset
    outdata = data[master_ind]

    return outdata

def find_nearest_cal(sc, calibrators, config):
    '''
    For an input coordinate, returns the nearest high quality VLA calibrator
    for the specified band and array configuration.

    Args
    =======
    sc (SkyCoord) - The sky coordinate to search around for a calibrator.
    band (string) - The observational band.
    config (string) - The observational array configuration.

    Returns
    =======
    best_sc (SkyCoord) - The skycoordinate of the nearest calibrator.
    name (string) - The J2000 epoch name for the returned calibrator.

    Raises
    =======
    None
    '''

    # Narrow down to good quality
    config_qual = calibrators[config+'_qual']
    good_qual = np.logical_or(config_qual=='P', config_qual=='S')
    cals = calibrators[good_qual]

    # Convert to SkyCoords
    sc_cals = SkyCoord(cals['ra'], cals['dec'])

    # Find minimum distance calibrator
    seps = sc.separation(sc_cals).deg
    near = np.argmin(seps)
    best_sc = sc_cals[near]
    name = cals['name'][near]

    return best_sc, name

def worst_setup_slew(az, el):
    '''
    This function finds the worst possible setup slew, assuming the telescope
    begins the observation at the worst possible wrap end.

    Args
    =======
    az (Angle) - The target azimuth of the first source
    el (Angle) - The target elevation of the first source

    Returns
    =======
    worst_time (Quantity) - The time the worst possible slew will take

    Raises
    =======
    None
    '''

    # HARD CODING SLEW SPEEDS
    vel_az = 40/60 * u.Unit('deg/s')
    vel_el = 20/60 * u.Unit('deg/s')

    # HARD CODING WRAP LIMITS
    az_ends = np.array([-85, 445])
    el_ends = np.array([0,90])

    # Figure out max in each direction
    az_dist = max(abs(az_ends-az.deg)) * u.deg
    el_dist = max(abs(el_ends-el.deg)) * u.deg

    time = max([ az_dist/vel_az, el_dist/vel_el])

    return time

def to_dur_fmt(t):
    '''
    Convert to format duration likes, input is time object
    '''
    time_hour = t.to('hour').value
    time_str = Angle(time_hour, unit=u.hourangle).to_string('h', pad=True)

    return time_str

def skyAz(lst, ra, dec):
    '''
    Helper function to predict_target. Extrapolates the azimuthal position for
    a source as a function of LST
    '''

    # Let's just assume we're at the VLA
    vla = EarthLocation.of_site('VLA')
    ha = Angle(lst*u.hourangle).rad - np.deg2rad(ra)
    lat = vla.lat.rad
    dec = np.deg2rad(dec)

    # Do Trigonometry
    cosAlt_sinAz = -np.cos(dec)*np.sin(ha)
    cosAlt_cosAz = np.sin(dec)*np.cos(lat) - np.cos(dec)*np.cos(ha)*np.sin(lat)

    Az = np.rad2deg(np.arctan2(cosAlt_sinAz,cosAlt_cosAz))%360

    return Az

def skyEl(lst, ra, dec):
    '''
    Helper function to predict_target. Extrapolates the elevation position for
    a source as a function of LST
    '''

    # Let's just assume we're at the VLA
    vla = EarthLocation.of_site('VLA')
    ha = Angle(lst*u.hourangle).rad - np.deg2rad(ra)
    lat = vla.lat.rad
    dec = np.deg2rad(dec)

    # Do Trigonometry
    sinAlt = np.sin(dec)*np.sin(lat) + np.cos(dec)*np.cos(ha)*np.cos(lat)

    Alt = np.rad2deg(np.arcsin(sinAlt))

    return Alt

def predict_target(tgt_ra, tgt_dec, tel_az, tel_el, lst0):
    '''
    Predict the location of the target after the slew, such that we intersect
    the objects path of motion. This effect will be most severe near zenith.
    '''
    # Convert to floats
    ra = tgt_ra.deg
    dec = tgt_dec.deg
    tel_az = tel_az.deg
    tel_el = tel_el.deg
    lst0 = lst0.hourangle

    # Current telescope position and capabilities
    vel_az = 40/60 * u.Unit('deg/s').to('deg/hour')
    vel_el = 20/60 * u.Unit('deg/s').to('deg/hour')

    # Azimuth calculations for intercept time
    az_dir = np.sign( skyAz(lst0, ra, dec) - tel_az)
    azt = lambda lst_t : tel_az + az_dir * vel_az * (lst_t-lst0)
    az_func = lambda lst, rr, dd : azt(lst) - skyAz(lst, rr, dd)
    az_time = brentq(az_func, lst0, lst0+0.5, args=(ra, dec))

    # Elevation calculations for intercept time
    el_dir = np.sign( skyEl(lst0, ra, dec) - tel_el)
    elt = lambda lst_t : tel_el + el_dir * vel_el * (lst_t-lst0)
    el_func = lambda lst, rr, dd : elt(lst) - skyEl(lst, rr, dd)
    el_time = brentq(el_func, lst0, lst0+0.5, args=(ra, dec))

    # Final coordinates
    worst_time = max([az_time, el_time])
    azf = skyAz(worst_time, ra, dec)
    elf = skyEl(worst_time, ra, dec)
    azf_ang = Angle(azf*u.deg)
    elf_ang = Angle(elf*u.deg)

    return azf_ang, elf_ang

def schedule_block(block, mosaics, source_scan=280, cal_scan=300,
                   ha_offset=0, verbose=True, config='A', wrap_pref=None):
    '''
    For input points and start LST, run a simulated observation to see if
    the particular schedule is possible. This will produce some figures
    displaying the hour angle and elevation of the telescope over time,
    highlighting in particular any areas where the telescope would go beyond
    certian limits. The function returns the total time the observation would
    take

    Args
    =======
    block (Table) - Target list in the FERMI format. The needed columns are
        essentially just RAJ2000, DEJ2000, and Source_Name.
    mosaics (dict) - The mosaic dictionary maping each source name to its
        mosaic coordinates
    source_scan (float/Quantity) - The amount of time to spend on each
        pointing. Default value is 30 and default units are seconds.
    cal_scan (float/Quantity) - The amount of time to spend on phase
        calibrators. Default value is 60 and default units are seconds.
    ha_offset (Quantity) - The hour angle of the first source when the
        observation will begin. For instance, if RA[0]=1h and ha_offset=-1h,
        then the observation will start at LST=0h. Defaults to 0h, such that
        things start when the first source is on the meridian. Default units
        are hourangle
    verbose (bool) - Whether to generate HA and El plots of the projected
        schedule and produce explanatory text.
    band (string) - The band which the observation will be performed in. This
        is used to search for calibrators.
    config (string) - The array configuration when the observation will be
        performed. This is used to search for calibrators.

    Returns
    =======
    sched_table (Table) - A table summarizing the observation. If the user
        likes how it went, this table can be passed to OPTSCHEDULER for
        formatting.

    Raises
    =======
    None
    '''

    ############### SET UP STUFF ###############
    # Check units
    if not isinstance(source_scan, u.Quantity):
        source_scan *= u.s
    if not isinstance(cal_scan, u.Quantity):
        cal_scan *= u.s
    if not isinstance(ha_offset, u.Quantity):
        ha_offset *= u.hourangle

    if wrap_pref is None:
        wrap_pref = ''

    # Define the coherence time manually

    # Convenience function for sc->str
    sc_to_ra = lambda sc:sc.ra.to_string('h', sep=':', pad=True)
    sc_to_dc = lambda sc:sc.dec.to_string('deg', sep=':', pad=True,
                                          alwayssign=True)

    # Convenience function for time->ang
    time_to_ang = lambda t : Angle(t.to('s').value/3600, unit=u.hourangle)

    # Convenience function for formatting LST angle
    lst_to_str = lambda lsti : lsti.to_string('h',sep=':',pad=True)[:8]

    ############### INITIALIZE TIMEKEEPING ###############
    # Figure out when to start
    tgt0 = SkyCoord(block['RAJ2000'][0]*u.deg, block['DEJ2000'][0]*u.deg)
    lst0 = tgt0.ra + ha_offset
    lst0_str = lst0.wrap_at('360d').to_string(u.hourangle, sep=':') #for print

    # This will keep a running track of time
    lst = lst0.wrap_at('360d') # initialize
    total_time = 0*u.s

    ############### DETERMINE BEST CALIBRATOR ###############
    # Values for best flux cal
    best_flux_name = 'Cygnus A'
    best_flux_sc = SkyCoord('19h59m28.34s', '40d44\'02.12"')
    best_flux_az, best_flux_el = E2H(best_flux_sc.ra, best_flux_sc.dec, lst)

    # What is the worst possible slew we might have to do
    worst_slew = worst_setup_slew(best_flux_az, best_flux_el)
    setup_atten = 1*u.min
    setup_req = 30*u.s
    slew_target = P2P(best_flux_sc, tgt0, lst0)
    scan_tgt = 10*u.min

    # Calculate length of setup slew
    setup_slew = worst_slew - setup_atten - setup_req

    # Figure out delays to apply to lst0 for flux cal
    delay_slew = worst_slew + scan_tgt + slew_target
    delay_atten = setup_atten + setup_req + scan_tgt + slew_target
    delay_req = setup_req + scan_tgt + slew_target
    delay_tgt = scan_tgt + slew_target

    ############### SETUP SCANS ###############
    # Define scheduler keeper
    sched = []

    # Slew scan
    start_lst = lst0 - time_to_ang(delay_slew)
    end_lst = lst0 - time_to_ang(delay_atten)
    start_az, start_el = E2H(best_flux_sc.ra, best_flux_sc.dec, start_lst)
    end_az, end_el = E2H(best_flux_sc.ra, best_flux_sc.dec, end_lst)
    start_ha = (start_lst-best_flux_sc.ra).wrap_at('180d')
    end_ha = (end_lst-best_flux_sc.ra).wrap_at('180d')
    slew_line = {'scanName': 'slew',
                 'sourceName': best_flux_name,
                 'resourceName': 'X band pointing',
                 'timeType': 'DUR',
                 'time': to_dur_fmt(setup_slew),
                 'antennaWrap': wrap_pref,
                 'applyRefPtg': 'N',
                 'applyPhase': 'N',
                 'recordOnMark5': 'N',
                 'allowOverTop': 'N',
                 'use10HzNoise': 'Y',
                 'scanIntents': 'SetAtnGain,',
                 'comments': '',
                 'ra': sc_to_ra(best_flux_sc),
                 'dec': sc_to_dc(best_flux_sc),
                 'start_lst': lst_to_str(start_lst.wrap_at('360d')),
                 'end_lst': lst_to_str(end_lst.wrap_at('360d')),
                 'start_az': start_az.deg,
                 'end_az': end_az.deg,
                 'start_el': start_el.deg,
                 'end_el': end_el.deg,
                 'start_ha': start_ha.hourangle,
                 'end_ha': end_ha.hourangle,
                 'errors': 0}
    sched.append(slew_line)

    # Atten scan
    start_lst = end_lst
    end_lst = lst0 - time_to_ang(delay_req)
    start_az, start_el = end_az, end_el
    end_az, end_el = E2H(best_flux_sc.ra, best_flux_sc.dec, end_lst)
    start_ha = (start_lst-best_flux_sc.ra).wrap_at('180d')
    end_ha = (end_lst-best_flux_sc.ra).wrap_at('180d')
    atten_line = {'scanName': 'atten',
                  'sourceName': best_flux_name,
                  'resourceName': 'L16f3B',
                  'timeType': 'DUR',
                  'time': to_dur_fmt(setup_atten),
                  'antennaWrap': wrap_pref,
                  'applyRefPtg': 'N',
                  'applyPhase': 'N',
                  'recordOnMark5': 'N',
                  'allowOverTop': 'N',
                  'use10HzNoise': 'Y',
                  'scanIntents': 'SetAtnGain,',
                  'comments': '',
                  'ra': sc_to_ra(best_flux_sc),
                  'dec': sc_to_dc(best_flux_sc),
                  'start_lst': lst_to_str(start_lst.wrap_at('360d')),
                  'end_lst': lst_to_str(end_lst.wrap_at('360d')),
                  'start_az': start_az.deg,
                  'end_az': end_az.deg,
                  'start_el': start_el.deg,
                  'end_el': end_el.deg,
                  'start_ha': start_ha.hourangle,
                  'end_ha': end_ha.hourangle,
                  'errors': 0}
    sched.append(atten_line)

    # Req scan
    start_lst = end_lst
    end_lst = lst0 - time_to_ang(delay_tgt)
    start_az, start_el = end_az, end_el
    end_az, end_el = E2H(best_flux_sc.ra, best_flux_sc.dec, end_lst)
    start_ha = (start_lst-best_flux_sc.ra).wrap_at('180d')
    end_ha = (end_lst-best_flux_sc.ra).wrap_at('180d')
    req_line = {'scanName': 'req',
                'sourceName': best_flux_name,
                'resourceName': 'L16f3B',
                'timeType': 'DUR',
                'time': to_dur_fmt(setup_req),
                'antennaWrap': wrap_pref,
                'applyRefPtg': 'N',
                'applyPhase': 'N',
                'recordOnMark5': 'N',
                'allowOverTop': 'N',
                'use10HzNoise': 'Y',
                'scanIntents': 'SetAtnGain,',
                'comments': '',
                'ra': sc_to_ra(best_flux_sc),
                'dec': sc_to_dc(best_flux_sc),
                'start_lst': lst_to_str(start_lst.wrap_at('360d')),
                'end_lst': lst_to_str(end_lst.wrap_at('360d')),
                'start_az': start_az.deg,
                'end_az': end_az.deg,
                'start_el': start_el.deg,
                'end_el': end_el.deg,
                'start_ha': start_ha.hourangle,
                'end_ha': end_ha.hourangle,
                'errors': 0}
    sched.append(req_line)

    # Flux target scan
    start_lst = end_lst
    end_lst = lst0 - time_to_ang(slew_target)
    start_az, start_el = end_az, end_el
    end_az, end_el = E2H(best_flux_sc.ra, best_flux_sc.dec, end_lst)
    start_ha = (start_lst-best_flux_sc.ra).wrap_at('180d')
    end_ha = (end_lst-best_flux_sc.ra).wrap_at('180d')
    flux_tgt_line = {'scanName': '',
                     'sourceName': best_flux_name,
                     'resourceName': 'L16f3B',
                     'timeType': 'DUR',
                     'time': to_dur_fmt(scan_tgt),
                     'antennaWrap': wrap_pref,
                     'applyRefPtg': 'N',
                     'applyPhase': 'N',
                     'recordOnMark5': 'N',
                     'allowOverTop': 'N',
                     'use10HzNoise': 'Y',
                     'scanIntents': 'CalBP,CalFlux,',
                     'comments': '',
                     'ra': sc_to_ra(best_flux_sc),
                     'dec': sc_to_dc(best_flux_sc),
                     'start_lst': lst_to_str(start_lst.wrap_at('360d')),
                     'end_lst': lst_to_str(end_lst.wrap_at('360d')),
                     'start_az': start_az.deg,
                     'end_az': end_az.deg,
                     'start_el': start_el.deg,
                     'end_el': end_el.deg,
                     'start_ha': start_ha.hourangle,
                     'end_ha': end_ha.hourangle,
                     'errors': 0}
    sched.append(flux_tgt_line)

    ############### RUN SIMULATION ###############
    # Initialize at defaults
    if wrap_pref == 'CCW':
        w = 'COUNTERCLOCKWISE'
    elif wrap_pref == 'CW':
        w = 'CLOCKWISE'
    else:
        w = None
    sim = MotionSimulator(end_az, end_el, wrap=w)

    # Run Simulation
    for i in range(len(block)):
        # Grab mosaic coordinates
        fname = block['Source_Name'][i]
        mosi = mosaics[fname]

        # Generate some names and notes
        b26 = lambda n : chr(97+n//26) + chr(97+n%26)
        names = [ fname+b26(n) for n in range(len(mosi)) ]
        intents = ['ObsTgt']*len(mosi)
        dur = [source_scan]*len(mosi)

        # Prepend a calibrator
        csc = SkyCoord('21h04m06.937s', '+76d33\'10.41"')
        cname = '3C427.1'
        cal_arr = np.array([[csc.ra.deg, csc.dec.deg]])

        # Prepend calibrator
        mosi = np.concatenate((cal_arr, mosi), axis=0)
        names.insert(0, cname)
        intents.insert(0, 'CalGain')
        dur.insert(0, cal_scan)

        if i==len(block)-1:
            mosi = np.concatenate((mosi, cal_arr), axis=0)
            names.append(cname)
            intents.append('CalGain')
            dur.append(cal_scan)

        for j in range(len(mosi)):
            # Move to a new spot on sky
            scj = SkyCoord(mosi[j][0]*u.deg, mosi[j][1]*u.deg)

            # Trying something new
            current_az = sim.getCurrentAntennaAzimuth()
            current_el = sim.getCurrentAntennaElevation()
            azj, elj = predict_target(scj.ra, scj.dec, current_az,
                                        current_el, lst)

            # Move to predicted position
            slew = sim.moveTo(azj, elj, wrap=w)

            # Convert slew time to an angle
            slew_ang = time_to_ang(slew)
            scan = dur[j] + 1*u.s
            scan_ang = time_to_ang(scan)

            # Record new time
            start_lst = lst
            lst += slew_ang + scan_ang
            total_time += slew + scan

            # Telescope tracked to new AltAz
            azj_new, elj_new = E2H(scj.ra, scj.dec, lst)
            sim.moveTo(azj_new, elj_new, wrap=w)

            # Num errors
            num_errors = len(sim.getErrors())

            # Calculate hour angles for recording
            start_ha = (start_lst - scj.ra).wrap_at('180d')
            end_ha = (lst - scj.ra).wrap_at('180d')

            # Keep track of scan
            source_line = {'scanName': '',
                           'sourceName': names[j],
                           'resourceName': 'L16f3B',
                           'timeType': 'DUR',
                           'time': to_dur_fmt(slew+scan),
                           'antennaWrap': wrap_pref,
                           'applyRefPtg': 'N',
                           'applyPhase': 'N',
                           'recordOnMark5': 'N',
                           'allowOverTop': 'N',
                           'use10HzNoise': 'Y',
                           'scanIntents': intents[j],
                           'comments': '',
                           'ra': sc_to_ra(scj),
                           'dec': sc_to_dc(scj),
                           'start_lst': lst_to_str(start_lst.wrap_at('360d')),
                           'end_lst': lst_to_str(lst.wrap_at('360d')),
                           'start_az': azj.deg,
                           'end_az': azj_new.deg,
                           'start_el': elj.deg,
                           'end_el': elj_new.deg,
                           'start_ha': start_ha.hourangle,
                           'end_ha': end_ha.hourangle,
                           'errors': num_errors}
            sched.append(source_line)

    sched_tab = Table(sched)

    # Display how run went
    if verbose:
        # Skip setup scans
        plot_tab = sched_tab[3:]

        # Some nice text
        print('\n'+'='*25, 'SCHEDULING BLOCK', '='*25+'\n')
        verb_str = 'Observing {} sources, starting at LST {}'
        print(verb_str.format(len(block), lst0_str))
        print('Flux cal:', best_flux_name)

        #DELETE ME
        print('worst slew', worst_slew)

        times = Angle(plot_tab['time'], unit=u.hourangle).hourangle
        start_times = np.cumsum(times)

        # Get errors and calibrators
        cc = ['C0' if e==0 else 'C3' for e in plot_tab['errors']]
        calind = ['CalGain' in intn for intn in plot_tab['scanIntents']]
        cal_times = start_times[calind]

        # Plot sky
        plt.figure(figsize=(8,8))
        ax1 = plt.subplot2grid((5,5), (0,0), rowspan=3, colspan=5,
                               projection='mollweide')
        ra = Angle(block['RAJ2000']).wrap_at('180d').rad
        dec = Angle(block['DEJ2000']).rad
        ax1.scatter(ra, dec, s=10)
        for i in range(len(ra)-1):
            p1 = [ra[i], dec[i]]
            p2 = [ra[i+1], dec[i+1]]
            plt.plot([p1[0], p2[0]], [p1[1], p2[1]])
        ax1.grid(True)

        # Plot az vs time
        ax2 = plt.subplot2grid((5,5), (3,0), colspan=3)
        ax2.scatter(start_times, plot_tab['start_az'], c=cc, s=5)
        ax2.axhline(85, ls='--', c='g')
        ax2.axhline(275, ls='--', c='r')
        for ct in cal_times:
            ax2.axvline(ct, c='C2', ls='--')
        ax2.set_ylabel('Hour Angle [hr]')
        ax2.set_xticklabels([])
        ax2.set_ylim(-85,445)
        ax2.grid(True)

        # Plot El vs time
        ax3 = plt.subplot2grid((5,5), (4,0), colspan=3)
        ax3.scatter(start_times, plot_tab['start_el'], c=cc, s=5)
        ax3.axhline(8, ls='--', c='k')
        for ct in cal_times:
            ax3.axvline(ct, c='C2', ls='--')
        ax3.set_xlabel('Time [hr]')
        ax3.set_ylabel('Elevation [deg]')
        ax3.set_ylim(0,90)
        ax3.grid(True)

        # Alt-Az map stuff
        a_az = np.deg2rad(plot_tab['start_az'])
        r_el = np.cos(np.deg2rad(plot_tab['start_el']))
        circle = np.linspace(0, 2*np.pi, len(r_el))
        r_min = np.full_like(r_el, np.cos(np.deg2rad(8)))

        # Plot alt-az map
        ax4 = plt.subplot2grid((5,5), (3,3), colspan=2, rowspan=2, polar=True)
        ax4.scatter(a_az, r_el, s=2)
        ax4.plot(circle, r_min, 'k--')
        ax4.plot(np.deg2rad([85,85]), [0,1], ls='--', c='g')
        ax4.plot(np.deg2rad([-85,-85]), [0,1], ls='--', c='r')
        ax4.set_ylim(0,1)
        ax4.set_xticklabels(['N', '', 'E', '', 'S', '', 'W'])
        ax4.set_yticklabels([])
        ax4.set_theta_direction(-1)
        ax4.set_theta_zero_location('N')
        ax4.grid(True)
        plt.show()

    # Just to be safe, clear the sim from memory
    del sim

    # Return scheduler
    return sched_tab

def time_block(block):
    tangs = Angle(block['time'], unit=u.hourangle).hourangle
    diff = sum(tangs)
    return diff

def make_best_block(block, mosaics, haoff, wrap, verbose,
                    doplots=True):

    dummy = schedule_block(block, mosaics, ha_offset=haoff,
                           wrap_pref=wrap, verbose=True)

    # Figure out what schedules to test
    width = -1
    steps = 10
    ha_range = np.linspace(haoff-width, haoff+width, steps)

    # Generate schedules
    sub_blocks = []
    for ha in ha_range:
        sblock = schedule_block(block, mosaics, ha_offset=ha,
                                wrap_pref=wrap, verbose=verbose)
        sub_blocks.append(sblock)

    # Some helpful plots, for now
    if doplots:
        el_min = [ min(sb['end_el']) for sb in sub_blocks]
        az_min = [ min(sb['end_az']) for sb in sub_blocks]
        az_max = [ max(sb['end_az']) for sb in sub_blocks]
        bl_tme = [ time_block(sb) for sb in sub_blocks]

        ax1 = plt.subplot2grid((2,2), (0,0))
        ax1.plot(ha_range, el_min)
        ax1.axhline(8)
        ax1.set_ylabel('Min El')

        ax2 = plt.subplot2grid((2,2), (0,1))
        ax2.plot(ha_range, az_max)
        ax2.set_ylabel('Max Az')

        ax3 = plt.subplot2grid((2,2), (1,1))
        ax3.plot(ha_range, az_min)
        ax3.set_ylabel('Min Az')

        ax4 = plt.subplot2grid((2,2), (1,0))
        ax4.plot(ha_range, bl_tme)
        ax4.set_ylabel('Run Time')
        plt.tight_layout()
        plt.show()

    # Copy middle block
    mid_ind = len(sub_blocks)//2
    master_block = sub_blocks[mid_ind].copy()

    # Check that all schedules are similar
    same_lengths = np.all([ len(s)==len(sub_blocks[0]) for s in sub_blocks])
    if same_lengths:
        # Extract max times
        print('Building out worst time scenario')
        tangs = Angle( [ sb['time'] for sb in sub_blocks ], unit=u.hourangle)
        tmaxs = np.max(tangs, axis=0).to_string('h', pad=True)
        master_block['time'] = tmaxs
    else:
        # Warn user if schedules are dissimilar
        print('ERROR: Could not build worst time scenario')

    # Time and print for user
    master_time = time_block(master_block)
    print('Master block time:', master_time)

    return master_block

def mosaic_new(radius, beamsize=0.0585):
    '''
    From input parameters of a FERMI positional uncertianty ellipse, generate
    a mosaic of beams to completely cover it. The coordinates of the mosaic
    pointing centers are returned.

    Args
    =======
    ra (float) - The right ascension of the FERMI source
    dec (float) - The declination of the FERMI source
    semimaj (float) - The semimajor axis of the FERMI source
    semimin (float) - The semiminor axis of the FERMI source
    angle (float) - The angle of the FERMI source, note that by default this
        is measured clockwise of North, and for our purposes we need it
        counterclockwise of East.

    Returns
    =======
    mosaic_coords (array) - An numpy array of coordinate points with shape
        (n,2). Note that the order of these points should amount to a fairly
        optimized slew path, alternating up and down rows.

    Raises
    =======
    None
    '''

    # Figure out how many rows/columns we need
    ver_bm = np.sqrt(3)/2 * beamsize # scaled to spacing
    nh = np.ceil(radius/beamsize)
    nv = np.ceil(radius/ver_bm)//2*2+1 # oddify

    # Build ranges
    rangev = np.linspace(-(nv-1)/2, (nv-1)/2, int(nv))

    coords = []
    for i in range(int(nv)):
        # For any row the y will be the same
        y = np.sqrt(3) * beamsize * rangev[i]

        # Flips center-like rows
        # Non-center-like rows shift left, add point on right
        tnh = nh
        if i%2==((nv-1)/2)%2:
            tnh = 1-nh
        rangeh = np.linspace(-tnh/2, tnh/2, int(abs(tnh)+1))

        # Loop through points in row
        for j in range(len(rangeh)):
            # Calculated x for each point
            x = 2 * beamsize * rangeh[j]

            # Record coordinate pair
            coords.append([x,y])

    # Arrayify
    coords = np.array(coords)

    # Remove any points not inside ellipse + fraction of beam radius
    fr = 0.8 # Derived -> (arccos(fr) - fr*sqrt(1-fr**2)) = 0.05*pi
    ewp = radius + fr*beamsize
    ehp = radius + fr*beamsize
    inside = (coords[:,0]/ewp)**2+(coords[:,1]/ehp)**2 <= 1
    coords_in = coords[inside]

    return coords_in

def rotate(sc_pnt, sc_shift):
    if sc_pnt.ndim == 0:
        sc_pnt = SkyCoord([sc_pnt.ra], [sc_pnt.dec])

    new_ra = np.zeros(len(sc_pnt))
    new_dec = np.zeros(len(sc_pnt))
    for i in range(len(sc_pnt)):
        pnt_ra = sc_pnt[i].ra
        pnt_dec = sc_pnt[i].dec

        ra_shift = sc_shift.ra
        dec_shift = sc_shift.dec

        v = np.array([[np.cos(pnt_dec)*np.cos(pnt_ra+ra_shift)],
                      [np.cos(pnt_dec)*np.sin(pnt_ra+ra_shift)],
                      [np.sin(pnt_dec)]])

        uv = np.array([[-np.sin(ra_shift)],
                      [np.cos(ra_shift)],
                      [0]])

        cross_u = np.array([[0, -uv[2][0], uv[1][0]],
                            [uv[2][0], 0, -uv[0][0]],
                            [-uv[1][0], uv[0][0], 0]])

        t1 = np.cos(dec_shift)*np.identity(3)
        t2 = -np.sin(dec_shift)*cross_u
        t3 = (1-np.cos(dec_shift))*np.outer(uv,uv)
        R = t1+t2+t3

        v_pr = R.dot(v)

        nra = np.arctan2(v_pr[1][0], v_pr[0][0])
        ndc = np.arcsin(v_pr[2][0])

        new_ra[i] = nra.to('deg').value
        new_dec[i] = ndc.to('deg').value

    new_sc = SkyCoord(new_ra*u.deg, new_dec*u.deg)

    if len(new_sc)==1:
        return new_sc[0]
    else:
        return new_sc

def main():
    beamsize = 1.1
    moc = mosaic_new(6, beamsize)
    scm = SkyCoord(moc[:,0]*u.deg, moc[:,1]*u.deg)
    scmr = rotate(scm, SkyCoord('0d', '90d'))
    moc = {'NCP': np.array([scmr.ra.deg, scmr.dec.deg]).T}

    tgt = Table({'Source_Name': ['NCP']*2, 'RAJ2000':[0]*2 *u.deg,
                 'DEJ2000':[89.99]*2 *u.deg})

    plt.scatter(scmr.ra.deg, scmr.dec.deg)
    plt.xlabel('RA [deg]')
    plt.ylabel('DEC [deg]')
    plt.grid()
    plt.show()

    # Block list
    blocks = []
    has = []

    # Make block
    ha1 = -2
    bestblock1 = make_best_block(tgt, moc, haoff=ha1, wrap='CCW', verbose=False)
    blocks.append(bestblock1)
    has.append(ha1)

    # Write
    write_OPT(blocks, has)

main()

import numpy,from numpy,from scipy,import astropy,from astropy


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

""" Functions for dealing with the Stagger model photospheres. """

from __future__ import division, absolute_import, print_function

__author__ = "Andy Casey <arc@ast.cam.ac.uk>"

import logging
import numpy as np

from .interpolator import BaseInterpolator

logger = logging.getLogger(__name__)

class Interpolator(BaseInterpolator):

    opacity_scale = "logtau"

    def __init__(self, filename, **kwargs):
        return super(self.__class__, self).__init__(filename, **kwargs)

def pickle_from_tsv_file(filename, depth_scale="optical", skiprows=72,
    delimiter=";"):
    """
    Pickle the Stagger-grid models from TSV-formatted filename.

    :param filename:
        The path of the TSV-formatted file.

    :type filename:
        str

    :param depth_scale: [optional, optical assumed]
        Which horizontal averaging method to use. Available options are:
        optical, mass density, Rosseland, or geometric height

    :type depth_scale:
        str

    :param skiprows: [optional]
        The number of rows at the top of the file before the header information.

    :type skiprows:
        int

    :param delimiter: [optional]
        The delimiting character between columns.

    :type delimiter:
        str
    """

    depth_scale_hint = depth_scale.lower()[0] # work it out from first letter
    if depth_scale_hint not in ("o", "m", "r", "z", "g", "h"): # zgh are same
        raise ValueError(
            "depth scale expected to be 'optical', 'mass density', "
            "Rosseland, or geometric height")
    if depth_scale_hint in ("g", "h"):
        depth_scale_hint = "z"
    elif depth_scale_hint == "r":
        depth_scale_hint = "R"

    depth_scale = {
        "o": "optical",
        "m": "mass density",
        "R": "Rossland opacity",
        "z": "geometric height",
    }[depth_scale_hint]

    with open(filename, "r") as fp:
        contents = fp.readlines()[skiprows + 1:]
    if contents[-1] == "\n": contents.pop(-1)

    # Number of extra columns in each row.
    n = 4

    # First three lines are for headers
    names = contents[0].strip().split(delimiter)
    units = contents[1].strip().split(delimiter)
    contents = contents[3:]

    num_models = len(set([row.split(delimiter)[n - 1] for row in contents]))
    parameters = np.nan * np.ones((num_models, n - 1))

    # Assume they all have the same number of depth points.
    assert (len(contents) % num_models) == 0
    num_depth_points = int(len(contents) / num_models)
    num_photospheric_quantitites = len(names) - n
    photospheres = np.nan * np.ones(
        (num_models, num_depth_points, num_photospheric_quantitites))

    for i in range(num_models):
        # The '4:' arises from the first four columns being the model parameters
        parameters[i, :] = \
            map(float, contents[i*num_depth_points].split(delimiter)[:n-1])
        photospheres[i, :, :] = np.array(
            [map(float, map(str.strip, _.split(delimiter)[n:])) \
                for _ in contents[i*num_depth_points:(i + 1)*num_depth_points]])

    names, units = names[n:], units[n:]
    # Replace dimensionless columns with "" for astropy.

    # Which depth scale do we want?
    indices = np.array([0] + [i for i, name in enumerate(names) \
        if name.endswith("({})".format(depth_scale_hint))])
    names = [names[0]] + [names[i][:-3] for i in indices[1:]]

    units = [units[i].replace("[-]", "") for i in indices]
    photospheres = photospheres[:, :, indices]

    meta = {
        "kind": "Stagger",
        "source_path": filename,
        "horizontal_averaging": depth_scale,
        "photospheric_units": units
    }
    assert np.all(np.isfinite(parameters))
    assert np.all(np.isfinite(photospheres))

    parameters = np.core.records.fromarrays(parameters.T,
        names=("effective_temperature", "surface_gravity", "metallicity"))

    return (parameters, photospheres, names, meta)


"""
Author: Joseph(Junzhe) Zhu, 2021/5. Email: josefzhu@stanford.edu / junzhe.joseph.zhu@gmail.com
For the original code for the paper[1], please refer to https://github.com/JunzheJosephZhu/MultiDecoder-DPRNN
Demo Page: https://junzhejosephzhu.github.io/Multi-Decoder-DPRNN/
Multi-Decoder DPRNN is a method for source separation when the number of speakers is unknown.
Our contribution is using multiple output heads, with each head modelling a distinct number of source outputs.
In addition, we design a selector network which determines which output head to use, i.e. estimates the number of sources.
The "DPRNN" part of the architecture is orthogonal to our contribution, and can be replaced with any other separator, e.g. Conv/LSTM-TasNet.
References:
    [1] "Multi-Decoder DPRNN: High Accuracy Source Counting and Separation",
        Junzhe Zhu, Raymond Yeh, Mark Hasegawa-Johnson. https://arxiv.org/abs/2011.12022
"""
import json
import os
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.functional import fold, unfold

from asteroid import torch_utils
from asteroid.models import BaseModel
from asteroid_filterbanks import make_enc_dec
from asteroid.engine.optimizers import make_optimizer
from asteroid.masknn import activations, norms
from asteroid.masknn.recurrent import DPRNNBlock
from asteroid.models.base_models import _shape_reconstructed, _unsqueeze_to_3d
from asteroid.utils.generic_utils import has_arg
from asteroid.utils.torch_utils import pad_x_to_y, script_if_tracing, jitable_shape
from asteroid.losses import PITLossWrapper, pairwise_neg_sisdr

def make_model_and_optimizer(conf, sample_rate):
    """Function to define the model and optimizer for a config dictionary.
    Args:
        conf: Dictionary containing the output of hierachical argparse.
    Returns:
        model, optimizer.
    The main goal of this function is to make reloading for resuming
    and evaluation very simple.
    """
    model = MultiDecoderDPRNN(**conf["masknet"], **conf["filterbank"], sample_rate=sample_rate)
    optimizer = make_optimizer(model.parameters(), **conf["optim"])
    return model, optimizer

class MultiDecoderDPRNN(BaseModel):
    """Multi-Decoder Dual-Path RNN as proposed in [1].

    Args:
        n_srcs (list of int): range of possible number of sources
        bn_chan (int): Number of channels after the bottleneck.
            Defaults to 128.
        hid_size (int): Number of neurons in the RNNs cell state.
            Defaults to 128.
        chunk_size (int): window size of overlap and add processing.
            Defaults to 100.
        hop_size (int or None): hop size (stride) of overlap and add processing.
            Default to `chunk_size // 2` (50% overlap).
        n_repeats (int): Number of repeats. Defaults to 6.
        norm_type (str, optional): Type of normalization to use. To choose from
            - ``'gLN'``: global Layernorm
            - ``'cLN'``: channelwise Layernorm
        mask_act (str, optional): Which non-linear function to generate mask.
        bidirectional (bool, optional): True for bidirectional Inter-Chunk RNN
            (Intra-Chunk is always bidirectional).
        rnn_type (str, optional): Type of RNN used. Choose between ``'RNN'``,
            ``'LSTM'`` and ``'GRU'``.
        num_layers (int, optional): Number of layers in each RNN.
        dropout (float, optional): Dropout ratio, must be in [0,1].
        kernel_size (int): Length of the filters.
        n_filters (int): Number of filters / Input dimension of the masker net.
        stride (int, optional): Stride of the convolution.
            If None (default), set to ``kernel_size // 2``.

    References
        [1] "Multi-Decoder DPRNN: High Accuracy Source Counting and Separation",
            Junzhe Zhu, Raymond Yeh, Mark Hasegawa-Johnson. https://arxiv.org/abs/2011.12022
    """

    def __init__(
        self,
        n_srcs,
        bn_chan=128,
        hid_size=128,
        chunk_size=100,
        hop_size=None,
        n_repeats=6,
        norm_type="gLN",
        mask_act="sigmoid",
        bidirectional=True,
        rnn_type="LSTM",
        num_layers=1,
        dropout=0,
        kernel_size=16,
        n_filters=64,
        stride=8,
        encoder_activation=None,
        use_mulcat=False,
        sample_rate=8000,
    ):
        super().__init__(sample_rate=sample_rate)
        self.encoder_activation = encoder_activation
        self.enc_activation = activations.get(encoder_activation or "linear")()
        hop_size = hop_size if hop_size is not None else chunk_size // 2
        self.encoder, _ = make_enc_dec(
            "free",
            kernel_size=kernel_size,
            n_filters=n_filters,
            stride=stride,
        )
        # Update in_chan
        self.masker = DPRNN_MultiStage(
            in_chan=n_filters,
            bn_chan=bn_chan,
            hid_size=hid_size,
            chunk_size=chunk_size,
            hop_size=hop_size,
            n_repeats=n_repeats,
            norm_type=norm_type,
            bidirectional=bidirectional,
            rnn_type=rnn_type,
            use_mulcat=use_mulcat,
            num_layers=num_layers,
            dropout=dropout,
        )
        self.decoder_select = Decoder_Select(
            kernel_size=kernel_size,
            stride=stride,
            in_chan=n_filters,
            n_srcs=n_srcs,
            bn_chan=bn_chan,
            chunk_size=chunk_size,
            hop_size=hop_size,
            mask_act=mask_act,
        )

        """
        Args:
            wav: 2D or 3D Tensor, Tensor of shape $(batch, T)$
            ground_truth: oracle number of speakers, None or list of $(batch)$ ints
        Return:
            reconstructed: torch.Tensor, $(batch, num_stages, max_spks, T)$
                where max_spks is the maximum possible number of speakers.
                if training, num_stages=n_repeats; otherwise num_stages=0
            Speaker dimension is zero-padded for examples with num_spks < max_spks
        """

    def forward(self, wav, ground_truth=None):
        shape = jitable_shape(wav)
        # [batch, 1, T]
        wav = _unsqueeze_to_3d(wav)
        tf_rep = self.enc_activation(self.encoder(wav))
        est_masks_list = self.masker(tf_rep)
        decoded, selector_output = self.decoder_select(
            est_masks_list, tf_rep, ground_truth=ground_truth
        )
        reconstructed = pad_x_to_y(decoded, wav)
        return _shape_reconstructed(reconstructed, shape), _shape_reconstructed(
            selector_output, shape
        )

    def forward_wav(self, wav, slice_size=32000, *args, **kwargs):
        """Separation method for waveforms.
        Unfolds a full audio into slices, estimate
        Args:
            wav (torch.Tensor): waveform array/tensor.
                Shape: 1D, 2D or 3D tensor, time last.
        Return:
            output_cat (torch.Tensor): concatenated output tensor.
                [num_spks, T]
        """
        assert not self.training, "forward_wav is only used for test mode"
        T = wav.size(-1)
        if wav.ndim == 1:
            wav = wav.reshape(1, wav.size(0))
        assert wav.ndim == 2  # [1, T]
        slice_stride = slice_size // 2
        # pad wav to integer multiple of slice_stride
        T_padded = max(int(np.ceil(T / slice_stride)), 2) * slice_stride
        wav = F.pad(wav, (0, T_padded - T))
        slices = wav.unfold(
            dimension=-1, size=slice_size, step=slice_stride
        )  # [1, slice_nb, slice_size]
        slice_nb = slices.size(1)
        slices = slices.squeeze(0).unsqueeze(1)
        tf_rep = self.enc_activation(self.encoder(slices))
        est_masks_list = self.masker(tf_rep)
        selector_input = est_masks_list[-1]  # [slice_nb, bn_chan, chunk_size, n_chunks]
        selector_output = self.decoder_select.selector(selector_input).reshape(
            slice_nb, -1
        )  # [slice_nb, num_decs]
        est_idx, _ = selector_output.argmax(-1).mode()
        est_spks = self.decoder_select.n_srcs[est_idx]
        output_wavs, _ = self.decoder_select(
            est_masks_list, tf_rep, ground_truth=[est_spks] * slice_nb
        )  # [slice_nb, 1, n_spks, slice_size]
        output_wavs = output_wavs.squeeze(1)[:, :est_spks, :]
        # TODO: overlap and add (with division)
        output_cat = output_wavs.new_zeros(est_spks, slice_nb * slice_size)
        output_cat[:, :slice_size] = output_wavs[0]
        start = slice_stride
        for i in range(1, slice_nb):
            end = start + slice_size
            overlap_prev = output_cat[:, start : start + slice_stride].unsqueeze(0)
            overlap_next = output_wavs[i : i + 1, :, :slice_stride]
            pw_losses = pairwise_neg_sisdr(overlap_next, overlap_prev)
            _, best_indices = PITLossWrapper.find_best_perm(pw_losses)
            reordered = PITLossWrapper.reorder_source(output_wavs[i : i + 1, :, :], best_indices)
            output_cat[:, start : start + slice_size] += reordered.squeeze(0)
            output_cat[:, start : start + slice_stride] /= 2
            start += slice_stride
        return output_cat[:, :T]

class DPRNN_MultiStage(nn.Module):
    """Implementation of the Dual-Path-RNN model,
    with multi-stage output, without Conv2D projection
    """

    def __init__(
        self,
        in_chan,
        bn_chan,
        hid_size,
        chunk_size,
        hop_size,
        n_repeats,
        norm_type,
        bidirectional,
        rnn_type,
        use_mulcat,
        num_layers,
        dropout,
    ):
        super(DPRNN_MultiStage, self).__init__()
        self.in_chan = in_chan
        self.bn_chan = bn_chan
        self.hid_size = hid_size
        self.chunk_size = chunk_size
        self.hop_size = hop_size
        self.n_repeats = n_repeats
        self.norm_type = norm_type
        self.bidirectional = bidirectional
        self.rnn_type = rnn_type
        self.num_layers = num_layers
        self.dropout = dropout
        self.use_mulcat = use_mulcat

        layer_norm = norms.get(norm_type)(in_chan)
        bottleneck_conv = nn.Conv1d(in_chan, bn_chan, 1)
        self.bottleneck = nn.Sequential(layer_norm, bottleneck_conv)

        # Succession of DPRNNBlocks.
        self.net = nn.ModuleList([])
        for i in range(self.n_repeats):
            self.net.append(
                DPRNNBlock(
                    bn_chan,
                    hid_size,
                    norm_type=norm_type,
                    bidirectional=bidirectional,
                    rnn_type=rnn_type,
                    use_mulcat=use_mulcat,
                    num_layers=num_layers,
                    dropout=dropout,
                )
            )

    def forward(self, mixture_w):
        """Forward.
        Args:
            mixture_w (:class:`torch.Tensor`): Tensor of shape $(batch, nfilters, nframes)$
        Returns:
            list of (:class:`torch.Tensor`): Tensor of shape $(batch, bn_chan, chunk_size, n_chunks)
        """
        batch, n_filters, n_frames = mixture_w.size()
        output = self.bottleneck(mixture_w)  # [batch, bn_chan, n_frames]
        output = unfold(
            output.unsqueeze(-1),
            kernel_size=(self.chunk_size, 1),
            padding=(self.chunk_size, 0),
            stride=(self.hop_size, 1),
        )
        n_chunks = output.shape[-1]
        output = output.reshape(batch, self.bn_chan, self.chunk_size, n_chunks)
        # Apply stacked DPRNN Blocks sequentially
        output_list = []
        for i in range(self.n_repeats):
            output = self.net[i](output)
            output_list.append(output)
        return output_list

class SingleDecoder(nn.Module):
    """
    Base decoder module, including the projection layer from (bn_chan) to (n_src * bn_chan).
    Takes a single example mask and encoding, outputs waveform
    """

    def __init__(
        self, kernel_size, stride, in_chan, n_src, bn_chan, chunk_size, hop_size, mask_act
    ):
        super(SingleDecoder, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.in_chan = in_chan
        self.bn_chan = bn_chan
        self.chunk_size = chunk_size
        self.hop_size = hop_size
        self.n_src = n_src
        self.mask_act = mask_act

        # Masking in 3D space
        net_out_conv = nn.Conv2d(bn_chan, n_src * bn_chan, 1)
        self.first_out = nn.Sequential(nn.PReLU(), net_out_conv)
        # Gating and masking in 2D space (after fold)
        self.net_out = nn.Sequential(nn.Conv1d(bn_chan, bn_chan, 1), nn.Tanh())
        self.net_gate = nn.Sequential(nn.Conv1d(bn_chan, bn_chan, 1), nn.Sigmoid())
        self.mask_net = nn.Conv1d(bn_chan, in_chan, 1, bias=False)

        # Get activation function.
        mask_nl_class = activations.get(mask_act)
        # For softmax, feed the source dimension.
        if has_arg(mask_nl_class, "dim"):
            self.output_act = mask_nl_class(dim=1)
        else:
            self.output_act = mask_nl_class()

        _, self.trans_conv = make_enc_dec(
            "free", kernel_size=kernel_size, stride=stride, n_filters=in_chan
        )

    def forward(self, output, mixture_w):
        """
        Args:
            output: LSTM output, Tensor of shape $(num_stages, bn_chan, chunk_size, n_chunks)$
            mixture_w: Encoder output, Tensor of shape $(num_stages, in_chan, nframes)
        outputs:
            est_wavs: Signal, Tensor of shape $(num_stages, n_src, T)
        """
        batch, bn_chan, chunk_size, n_chunks = output.size()
        _, in_chan, n_frames = mixture_w.size()
        assert self.bn_chan == bn_chan
        assert self.in_chan == in_chan
        assert self.chunk_size == chunk_size
        output = self.first_out(output)
        output = output.reshape(batch * self.n_src, self.bn_chan, self.chunk_size, n_chunks)
        # Overlap and add:
        # [batch, out_chan, chunk_size, n_chunks] -> [batch, out_chan, n_frames]
        to_unfold = self.bn_chan * self.chunk_size
        output = fold(
            output.reshape(batch * self.n_src, to_unfold, n_chunks),
            (n_frames, 1),
            kernel_size=(self.chunk_size, 1),
            padding=(self.chunk_size, 0),
            stride=(self.hop_size, 1),
        )
        # Apply gating
        output = output.reshape(batch * self.n_src, self.bn_chan, -1)
        output = self.net_out(output) * self.net_gate(output)
        # Compute mask
        score = self.mask_net(output)
        est_mask = self.output_act(score)
        est_mask = est_mask.reshape(batch, self.n_src, self.in_chan, n_frames)
        mixture_w = mixture_w.unsqueeze(1)
        source_w = est_mask * mixture_w
        source_w = source_w.reshape(batch * self.n_src, self.in_chan, n_frames)
        est_wavs = self.trans_conv(source_w)
        est_wavs = est_wavs.reshape(batch, self.n_src, -1)
        return est_wavs

class Decoder_Select(nn.Module):
    """Selects which SingleDecoder to use, as well as whether to use multiloss, as proposed in [1]
    References
        [1] "Multi-Decoder DPRNN: High Accuracy Source Counting and Separation",
            Junzhe Zhu, Raymond Yeh, Mark Hasegawa-Johnson. https://arxiv.org/abs/2011.12022
    """

    def __init__(
        self, kernel_size, stride, in_chan, n_srcs, bn_chan, chunk_size, hop_size, mask_act
    ):
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.in_chan = in_chan
        self.n_srcs = n_srcs
        self.bn_chan = bn_chan
        self.chunk_size = chunk_size
        self.hop_size = hop_size
        self.mask_act = mask_act

        self.n_src2idx = {n_src: i for i, n_src in enumerate(n_srcs)}
        self.decoders = torch.nn.ModuleList()
        for n_src in n_srcs:
            self.decoders.append(
                SingleDecoder(
                    kernel_size=kernel_size,
                    stride=stride,
                    in_chan=in_chan,
                    n_src=n_src,
                    bn_chan=bn_chan,
                    chunk_size=chunk_size,
                    hop_size=hop_size,
                    mask_act=mask_act,
                )
            )
        self.selector = nn.Sequential(
            nn.Conv2d(bn_chan, in_chan, 1),
            nn.AdaptiveAvgPool2d(1),
            nn.ReLU(),
            nn.Conv2d(in_chan, len(n_srcs), 1),
        )

    def forward(self, output_list, mixture_w, ground_truth):
        """Forward
        Args:
            output_list: list of $(batch, bn_chan, chunk_size, n_chunks)$
            mixture_w: torch.Tensor, $(batch, in_chan, n_frames)$
            ground_truth: None, or list of [B] ints, or Long Tensor of $(B)
                if None, use inferred number of speakers to determine output shape
        Output:
            output_wavs: torch.Tensor, $(batch, num_stages, max_spks, T)$
                where the speaker dimension is padded for examples with num_spks < max_spks
                if training, num_stages=n_repeats; otherwise, num_stages=1
            selector_output: output logits from selector module. torch.Tensor, $(batch, num_stages, num_decoders)$
        """
        batch, bn_chan, chunk_size, n_chunks = output_list[0].size()
        _, in_chan, n_frames = mixture_w.size()
        assert self.chunk_size == chunk_size
        if not self.training:
            output_list = output_list[-1:]
        num_stages = len(output_list)
        # [batch, num_stages, bn_chan, chunk_size, n_chunks]
        output = torch.stack(output_list, 1).reshape(
            batch * num_stages, bn_chan, chunk_size, n_chunks
        )
        selector_output = self.selector(output).reshape(batch, num_stages, -1)
        output = output.reshape(batch, num_stages, bn_chan, chunk_size, n_chunks)
        # [batch, num_stages, in_chan, n_frames]
        mixture_w = mixture_w.unsqueeze(1).repeat(1, num_stages, 1, 1)
        if ground_truth is not None:  # oracle
            decoder_selected = torch.LongTensor([self.n_src2idx[truth] for truth in ground_truth])
        else:
            assert num_stages == 1  # can't use select with multistage
            decoder_selected = selector_output.reshape(batch, -1).argmax(1)
        T = self.kernel_size + self.stride * (n_frames - 1)
        output_wavs = torch.zeros(batch, num_stages, max(self.n_srcs), T).to(output.device)
        for i in range(batch):
            output_wavs[i, :, : self.n_srcs[decoder_selected[i]], :] = self.decoders[
                decoder_selected[i]
            ](output[i], mixture_w[i])
        return output_wavs, selector_output

def load_best_model(train_conf, exp_dir, sample_rate):
    """Load best model after training.

    Args:
        train_conf (dict): dictionary as expected by `make_model_and_optimizer`
        exp_dir(str): Experiment directory. Expects to find
            `'best_k_models.json'` of `checkpoints` directory in it.

    Returns:
        nn.Module the best (or last) pretrained model according to the val_loss.
    """
    # Create the model from recipe-local function
    model, _ = make_model_and_optimizer(train_conf, sample_rate=sample_rate)
    try:
        # Last best model summary
        with open(os.path.join(exp_dir, "best_k_models.json"), "r") as f:
            best_k = json.load(f)
        best_model_path = min(best_k, key=best_k.get)
    except FileNotFoundError:
        # Get last checkpoint
        all_ckpt = os.listdir(os.path.join(exp_dir, "checkpoints/"))
        all_ckpt = [
            (ckpt, int("".join(filter(str.isdigit, os.path.basename(ckpt)))))
            for ckpt in all_ckpt
            if ckpt.find("ckpt") >= 0
        ]
        all_ckpt.sort(key=lambda x: x[1])
        best_model_path = os.path.join(exp_dir, "checkpoints", all_ckpt[-1][0])
    # Load checkpoint
    checkpoint = torch.load(best_model_path, map_location="cpu")
    # Load state_dict into model.
    model = torch_utils.load_state_dict_in(checkpoint["state_dict"], model)
    model.eval()
    return model

# Training notes:
# Weight different stages in accordance with facebook code
if __name__ == "__main__":
    network = MultiDecoderDPRNN(n_srcs=[2, 3], bn_chan=32, hid_size=32, n_filters=16)
    # training
    input = torch.rand(2, 3200)
    wavs, selector_output = network(input, [3, 2])
    print(wavs.shape)
    assert (wavs[1, :, 2] == 0).all()
    # validation
    network.eval()
    wavs, selector_output = network(input)
    print(wavs.shape)
    # test
    input_wav = torch.rand(64351)
    output_wavs = network.forward_wav(input_wav)
    print(output_wavs.shape)


from numba import prange, njit, jit, objmode
import logging
import numpy as np

from tardis.montecarlo.montecarlo_numba.r_packet import (
    RPacket,
    PacketStatus,
)
from tardis.montecarlo.montecarlo_numba.utils import MonteCarloException

from tardis.montecarlo.montecarlo_numba.numba_interface import (
    PacketCollection,
    VPacketCollection,
    NumbaModel,
    numba_plasma_initialize,
    Estimators,
    configuration_initialize,
)

from tardis.montecarlo import (
    montecarlo_configuration as montecarlo_configuration,
)

from tardis.montecarlo.montecarlo_numba.single_packet_loop import (
    single_packet_loop,
)
from tardis.montecarlo.montecarlo_numba import njit_dict
from numba.typed import List
from tardis.util.base import update_iterations_pbar, update_packet_pbar

def montecarlo_radial1d(
    model,
    plasma,
    iteration,
    no_of_packets,
    total_iterations,
    show_progress_bars,
    runner,
):
    packet_collection = PacketCollection(
        runner.input_r,
        runner.input_nu,
        runner.input_mu,
        runner.input_energy,
        runner._output_nu,
        runner._output_energy,
    )

    numba_model = NumbaModel(
        runner.r_inner_cgs,
        runner.r_outer_cgs,
        model.time_explosion.to("s").value,
    )
    numba_plasma = numba_plasma_initialize(plasma, runner.line_interaction_type)
    estimators = Estimators(
        runner.j_estimator,
        runner.nu_bar_estimator,
        runner.j_blue_estimator,
        runner.Edotlu_estimator,
    )
    packet_seeds = montecarlo_configuration.packet_seeds

    number_of_vpackets = montecarlo_configuration.number_of_vpackets

    (
        v_packets_energy_hist,
        last_interaction_type,
        last_interaction_in_nu,
        last_line_interaction_in_id,
        last_line_interaction_out_id,
        virt_packet_nus,
        virt_packet_energies,
        virt_packet_initial_mus,
        virt_packet_initial_rs,
        virt_packet_last_interaction_in_nu,
        virt_packet_last_interaction_type,
        virt_packet_last_line_interaction_in_id,
        virt_packet_last_line_interaction_out_id,
    ) = montecarlo_main_loop(
        packet_collection,
        numba_model,
        numba_plasma,
        estimators,
        runner.spectrum_frequency.value,
        number_of_vpackets,
        packet_seeds,
        iteration=iteration,
        show_progress_bars=show_progress_bars,
        no_of_packets=no_of_packets,
        total_iterations=total_iterations,
    )

    runner._montecarlo_virtual_luminosity.value[:] = v_packets_energy_hist
    runner.last_interaction_type = last_interaction_type
    runner.last_interaction_in_nu = last_interaction_in_nu
    runner.last_line_interaction_in_id = last_line_interaction_in_id
    runner.last_line_interaction_out_id = last_line_interaction_out_id

    if montecarlo_configuration.VPACKET_LOGGING and number_of_vpackets > 0:
        runner.virt_packet_nus = np.concatenate(
            np.array(virt_packet_nus)
        ).ravel()
        runner.virt_packet_energies = np.concatenate(
            np.array(virt_packet_energies)
        ).ravel()
        runner.virt_packet_initial_mus = np.concatenate(
            np.array(virt_packet_initial_mus)
        ).ravel()
        runner.virt_packet_initial_rs = np.concatenate(
            np.array(virt_packet_initial_rs)
        ).ravel()
        runner.virt_packet_last_interaction_in_nu = np.concatenate(
            np.array(virt_packet_last_interaction_in_nu)
        ).ravel()
        runner.virt_packet_last_interaction_type = np.concatenate(
            np.array(virt_packet_last_interaction_type)
        ).ravel()
        runner.virt_packet_last_line_interaction_in_id = np.concatenate(
            np.array(virt_packet_last_line_interaction_in_id)
        ).ravel()
        runner.virt_packet_last_line_interaction_out_id = np.concatenate(
            np.array(virt_packet_last_line_interaction_out_id)
        ).ravel()
    update_iterations_pbar(1)

@njit(**njit_dict)
def montecarlo_main_loop(
    packet_collection,
    numba_model,
    numba_plasma,
    estimators,
    spectrum_frequency,
    number_of_vpackets,
    packet_seeds,
    iteration,
    show_progress_bars,
    no_of_packets,
    total_iterations,
):
    """
    This is the main loop of the MonteCarlo routine that generates packets
    and sends them through the ejecta.

    Parameters
    ----------
    packet_collection : PacketCollection
    numba_model : NumbaModel
    estimators : NumbaEstimators
    spectrum_frequency : astropy.units.Quantity
        frequency bins
    number_of_vpackets : int
        VPackets released per interaction
    packet_seeds : numpy.array
    """
    output_nus = np.empty_like(packet_collection.packets_output_nu)
    last_interaction_types = (
        np.ones_like(packet_collection.packets_output_nu, dtype=np.int64) * -1
    )
    output_energies = np.empty_like(packet_collection.packets_output_nu)

    last_interaction_in_nus = np.empty_like(packet_collection.packets_output_nu)
    last_line_interaction_in_ids = (
        np.ones_like(packet_collection.packets_output_nu, dtype=np.int64) * -1
    )
    last_line_interaction_out_ids = (
        np.ones_like(packet_collection.packets_output_nu, dtype=np.int64) * -1
    )

    v_packets_energy_hist = np.zeros_like(spectrum_frequency)
    delta_nu = spectrum_frequency[1] - spectrum_frequency[0]

    # Pre-allocate a list of vpacket collections for later storage
    vpacket_collections = List()
    for i in range(len(output_nus)):
        vpacket_collections.append(
            VPacketCollection(
                i,
                spectrum_frequency,
                montecarlo_configuration.v_packet_spawn_start_frequency,
                montecarlo_configuration.v_packet_spawn_end_frequency,
                number_of_vpackets,
                montecarlo_configuration.temporary_v_packet_bins,
            )
        )

    # Arrays for vpacket logging
    virt_packet_nus = []
    virt_packet_energies = []
    virt_packet_initial_mus = []
    virt_packet_initial_rs = []
    virt_packet_last_interaction_in_nu = []
    virt_packet_last_interaction_type = []
    virt_packet_last_line_interaction_in_id = []
    virt_packet_last_line_interaction_out_id = []

    for i in prange(len(output_nus)):
        if show_progress_bars:
            with objmode:
                update_amount  = 1
                update_packet_pbar(
                    update_amount,
                    current_iteration=iteration,
                    no_of_packets=no_of_packets,
                    total_iterations=total_iterations,
                )

        if montecarlo_configuration.single_packet_seed != -1:
            seed = packet_seeds[montecarlo_configuration.single_packet_seed]
            np.random.seed(seed)
        else:
            seed = packet_seeds[i]
            np.random.seed(seed)
        r_packet = RPacket(
            numba_model.r_inner[0],
            packet_collection.packets_input_mu[i],
            packet_collection.packets_input_nu[i],
            packet_collection.packets_input_energy[i],
            seed,
            i,
        )
        vpacket_collection = vpacket_collections[i]

        loop = single_packet_loop(
            r_packet, numba_model, numba_plasma, estimators, vpacket_collection
        )
        # if loop and 'stop' in loop:
        #     raise MonteCarloException

        output_nus[i] = r_packet.nu
        last_interaction_in_nus[i] = r_packet.last_interaction_in_nu
        last_line_interaction_in_ids[i] = r_packet.last_line_interaction_in_id
        last_line_interaction_out_ids[i] = r_packet.last_line_interaction_out_id

        if r_packet.status == PacketStatus.REABSORBED:
            output_energies[i] = -r_packet.energy
            last_interaction_types[i] = r_packet.last_interaction_type
        elif r_packet.status == PacketStatus.EMITTED:
            output_energies[i] = r_packet.energy
            last_interaction_types[i] = r_packet.last_interaction_type

        vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx]
        vpackets_energy = vpacket_collection.energies[: vpacket_collection.idx]
        vpackets_initial_mu = vpacket_collection.initial_mus[
            : vpacket_collection.idx
        ]
        vpackets_initial_r = vpacket_collection.initial_rs[
            : vpacket_collection.idx
        ]

        v_packets_idx = np.floor(
            (vpackets_nu - spectrum_frequency[0]) / delta_nu
        ).astype(np.int64)
        # if we're only in a single-packet mode
        # if montecarlo_configuration.single_packet_seed == -1:
        #    break
        for j, idx in enumerate(v_packets_idx):
            if (vpackets_nu[j] < spectrum_frequency[0]) or (
                vpackets_nu[j] > spectrum_frequency[-1]
            ):
                continue
            v_packets_energy_hist[idx] += vpackets_energy[j]

    if montecarlo_configuration.VPACKET_LOGGING:
        for vpacket_collection in vpacket_collections:
            vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx]
            vpackets_energy = vpacket_collection.energies[
                : vpacket_collection.idx
            ]
            vpackets_initial_mu = vpacket_collection.initial_mus[
                : vpacket_collection.idx
            ]
            vpackets_initial_r = vpacket_collection.initial_rs[
                : vpacket_collection.idx
            ]
            virt_packet_nus.append(np.ascontiguousarray(vpackets_nu))
            virt_packet_energies.append(np.ascontiguousarray(vpackets_energy))
            virt_packet_initial_mus.append(
                np.ascontiguousarray(vpackets_initial_mu)
            )
            virt_packet_initial_rs.append(
                np.ascontiguousarray(vpackets_initial_r)
            )
            virt_packet_last_interaction_in_nu.append(
                np.ascontiguousarray(
                    vpacket_collection.last_interaction_in_nu[
                        : vpacket_collection.idx
                    ]
                )
            )
            virt_packet_last_interaction_type.append(
                np.ascontiguousarray(
                    vpacket_collection.last_interaction_type[
                        : vpacket_collection.idx
                    ]
                )
            )
            virt_packet_last_line_interaction_in_id.append(
                np.ascontiguousarray(
                    vpacket_collection.last_interaction_in_id[
                        : vpacket_collection.idx
                    ]
                )
            )
            virt_packet_last_line_interaction_out_id.append(
                np.ascontiguousarray(
                    vpacket_collection.last_interaction_out_id[
                        : vpacket_collection.idx
                    ]
                )
            )

    packet_collection.packets_output_energy[:] = output_energies[:]
    packet_collection.packets_output_nu[:] = output_nus[:]

    return (
        v_packets_energy_hist,
        last_interaction_types,
        last_interaction_in_nus,
        last_line_interaction_in_ids,
        last_line_interaction_out_ids,
        virt_packet_nus,
        virt_packet_energies,
        virt_packet_initial_mus,
        virt_packet_initial_rs,
        virt_packet_last_interaction_in_nu,
        virt_packet_last_interaction_type,
        virt_packet_last_line_interaction_in_id,
        virt_packet_last_line_interaction_out_id,
    )


"""
Wrapper for openquake models.
Can import without openquake but using openquake models will raise ImportError.
"""
from math import exp

import numpy as np

from empirical.util.classdef import TectType, GMM

try:
    # openquake constants and models
    from openquake.hazardlib import const, imt, gsim
    from openquake.hazardlib.site import Site, SiteCollection
    from openquake.hazardlib.geo import Point

    OQ = True
except ImportError:
    # fail silently, only an issue if openquake models wanted
    OQ = False

SITE_PROPERTIES = [
    ("vs30", "vs30"),
    ("vs30measured", "vs30measured"),
    ("z1pt0", "z1p0"),
    ("z2pt5", "z2p5"),
    ("fpeak", "fpeak"),
]
RUPTURE_PROPERTIES = [
    ("mag", "Mw"),
    ("rake", "rake"),
    ("width", "width"),
    ("ztor", "ztor"),
    ("hypo_depth", "hdepth"),
]
DISTANCE_PROPERTIES = [
    ("rrup", "Rrup"),
    ("rjb", "Rjb"),
    ("rx", "Rx"),
    ("ry0", "Ry"),
    ("rvolc", "Rtvz"),
]

OQ_GMM_LIST = [
    GMM.P_20,
    GMM.HA_20,
    GMM.G_17,
    GMM.BC_16,
    GMM.S_16,
    GMM.Ph_20,
    GMM.Ch_20,
    GMM.AG_20,
    GMM.AG_20_NZ,
    GMM.K_20,
    GMM.K_20_NZ,
    GMM.Si_20,
    GMM.Z_16,
]
if OQ:
    oq_models = {
        GMM.P_20: {
            TectType.SUBDUCTION_SLAB: gsim.parker_2020.ParkerEtAl2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.parker_2020.ParkerEtAl2020SInter,
        },
        GMM.HA_20: {
            TectType.ACTIVE_SHALLOW: gsim.hassani_atkinson_2020.HassaniAtkinson2020Asc,
            TectType.SUBDUCTION_SLAB: gsim.hassani_atkinson_2020.HassaniAtkinson2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.hassani_atkinson_2020.HassaniAtkinson2020SInter,
        },
        GMM.G_17: {TectType.ACTIVE_SHALLOW: gsim.gulerce_2017.GulerceEtAl2017},
        GMM.BC_16: {
            TectType.ACTIVE_SHALLOW: gsim.bozorgnia_campbell_2016.BozorgniaCampbell2016
        },
        GMM.S_16: {TectType.ACTIVE_SHALLOW: gsim.stewart_2016_vh.StewartEtAl2016VH},
        GMM.Ph_20: {
            TectType.ACTIVE_SHALLOW: gsim.phung_2020.PhungEtAl2020Asc,
            TectType.SUBDUCTION_SLAB: gsim.phung_2020.PhungEtAl2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.phung_2020.PhungEtAl2020SInter,
        },
        GMM.Ch_20: {
            TectType.ACTIVE_SHALLOW: gsim.chao_2020.ChaoEtAl2020Asc,
            TectType.SUBDUCTION_SLAB: gsim.chao_2020.ChaoEtAl2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.chao_2020.ChaoEtAl2020SInter,
        },
        GMM.AG_20: {
            TectType.SUBDUCTION_SLAB: gsim.abrahamson_gulerce_2020.AbrahamsonGulerce2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.abrahamson_gulerce_2020.AbrahamsonGulerce2020SInter,
        },
        GMM.AG_20_NZ: {
            TectType.SUBDUCTION_SLAB: gsim.abrahamson_gulerce_2020.AbrahamsonGulerce2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.abrahamson_gulerce_2020.AbrahamsonGulerce2020SInter,
        },
        GMM.K_20: {
            TectType.SUBDUCTION_SLAB: gsim.kuehn_2020.KuehnEtAl2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.kuehn_2020.KuehnEtAl2020SInter,
        },
        GMM.K_20_NZ: {
            TectType.SUBDUCTION_SLAB: gsim.kuehn_2020.KuehnEtAl2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.kuehn_2020.KuehnEtAl2020SInter,
        },
        GMM.Si_20: {
            TectType.SUBDUCTION_SLAB: gsim.si_2020.SiEtAl2020SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.si_2020.SiEtAl2020SInter,
        },
        GMM.Z_16: {
            TectType.ACTIVE_SHALLOW: gsim.zhao_2016.ZhaoEtAl2016Asc,
            TectType.SUBDUCTION_SLAB: gsim.zhao_2016.ZhaoEtAl2016SSlab,
            TectType.SUBDUCTION_INTERFACE: gsim.zhao_2016.ZhaoEtAl2016SInter,
        },
    }

class Properties(object):
    """
    Stores values for sites, rup and dists.
    """

    def __init__(self):
        # this allows attaching arbitrary attributes to self later
        pass

def oq_mean_stddevs(model, sites, rup, dists, imr, stddev_types):
    """
    Calculate mean and standard deviations given openquake input structures.
    """
    mean, stddevs = model.get_mean_and_stddevs(sites, rup, dists, imr, stddev_types)
    mean = exp(mean[0]) if hasattr(mean, "__len__") else exp(mean)
    stddevs = [s[0] if hasattr(s, "__len__") else s for s in stddevs]

    return mean, stddevs

def oq_run(model, site, fault, im, period=None, **kwargs):
    """
    Run an openquake model using Empirical_Engine input structures.
    model: model or value from empirical.util.classdef.GMM or openquake class:
           GMM.P_20 gsim.parker_2020.ParkerEtAl2020SInter
    site / fault: instances from empirical.classdef -- A tect_type must be able to be set to retrieve the correct model
    im: intensity measure name
    period: for spectral acceleration, openquake tables automatically
            interpolate values between specified values, fails if outside range
    kwargs: pass extra (model specific) parameters to models
    """
    if not OQ:
        raise ImportError("openquake is not installed, models not available")

    # model can be given multiple ways
    if type(model).__name__ == "GMM":
        model = oq_models[model][fault.tect_type](**kwargs)
    elif type(model).__name__ == "MetaGSIM":
        model = model(**kwargs)

    trt = model.DEFINED_FOR_TECTONIC_REGION_TYPE
    if trt == const.TRT.SUBDUCTION_INTERFACE:
        assert fault.tect_type == TectType.SUBDUCTION_INTERFACE
    elif trt == const.TRT.SUBDUCTION_INTRASLAB:
        assert fault.tect_type == TectType.SUBDUCTION_SLAB
    elif trt == const.TRT.ACTIVE_SHALLOW_CRUST:
        assert fault.tect_type == TectType.ACTIVE_SHALLOW
    else:
        raise ValueError("unknown tectonic region: " + trt)

    stddev_types = []
    for st in [const.StdDev.TOTAL, const.StdDev.INTER_EVENT, const.StdDev.INTRA_EVENT]:
        if st in model.DEFINED_FOR_STANDARD_DEVIATION_TYPES:
            stddev_types.append(st)

    location = Point(
        0.0, 0.0, 0.0
    )  # Create a dummy location as OQ calculation doesn't use a location
    oq_site = Site(location)
    extra_site_parameters = set(model.REQUIRES_SITES_PARAMETERS).difference(
        list(zip(*SITE_PROPERTIES))[0]
    )
    if len(extra_site_parameters) > 0:
        raise ValueError("unknown site property: " + extra_site_parameters)
    oq_site = check_properties(site, model, SITE_PROPERTIES, oq_site, np_array=True)

    sites = SiteCollection([oq_site])

    extra_rup_properties = set(model.REQUIRES_RUPTURE_PARAMETERS).difference(
        list(zip(*RUPTURE_PROPERTIES))[0]
    )
    if len(extra_rup_properties) > 0:
        raise ValueError("unknown rupture property: " + " ".join(extra_rup_properties))
    rupture = check_properties(fault, model, RUPTURE_PROPERTIES, Properties())

    extra_dist_properties = set(model.REQUIRES_DISTANCES).difference(
        list(zip(*DISTANCE_PROPERTIES))[0]
    )
    if len(extra_dist_properties) > 0:
        raise ValueError(
            "unknown distance property: " + " ".join(extra_dist_properties)
        )
    dists = check_properties(
        site, model, DISTANCE_PROPERTIES, Properties(), np_array=True
    )

    if period is not None:
        assert imt.SA in model.DEFINED_FOR_INTENSITY_MEASURE_TYPES
        # use sorted instead of max for full list
        max_period = max([i.period for i in model.COEFFS.sa_coeffs.keys()])
        single = False
        if not hasattr(period, "__len__"):
            single = True
            period = [period]
        results = []
        for p in period:
            imr = imt.SA(period=min(p, max_period))
            m, s = oq_mean_stddevs(model, sites, rupture, dists, imr, stddev_types)
            # interpolate pSA value up based on maximum available period
            if p > max_period:
                m = m * (max_period / p) ** 2
            results.append((m, s))
        if single:
            return results[0]
        return results
    else:
        imc = getattr(imt, im)
        assert imc in model.DEFINED_FOR_INTENSITY_MEASURE_TYPES
        return oq_mean_stddevs(model, sites, rupture, dists, imc(), stddev_types)

def check_properties(ee_object, model, properties, properties_obj, np_array=False):
    for oq_property_name, ee_property_name in properties:
        ee_property = getattr(ee_object, ee_property_name)
        if ee_property:
            setattr(
                properties_obj,
                oq_property_name,
                np.array([ee_property]) if np_array else ee_property,
            )
        else:
            check_param(model, oq_property_name)
    return properties_obj

def check_param(model, rp):
    if rp in model.REQUIRES_RUPTURE_PARAMETERS:
        raise ValueError(f"{rp} is a required parameter for {model}")


"""
The classes in this file do preprocessing on data and monte carlo to be used
to do a point source analysis.
"""

__author__ = 'John Evans'
__copyright__ = 'Copyright 2020 John Evans'
__credits__ = ['John Evans', 'Jason Fan', 'Michael Larson']
__license__ = 'Apache License 2.0'
__version__ = '0.0.1'
__maintainer__ = 'John Evans'
__email__ = 'john.evans@icecube.wisc.edu'
__status__ = 'Development'

from typing import List, Optional, Tuple, Union

import numpy as np
from scipy.interpolate import UnivariateSpline as Spline

from dataclasses import dataclass
from dataclasses import field
from dataclasses import InitVar

from . import sources
from . import _models

@dataclass
class _I3EventModelBase(_models.EventModelBase):
    """Docstring

    Attributes:
        sin_dec_bins (np.array): An array of sin(dec) bin edges for the energy
            maps.
        log_energy_bins (np.array): An array of log(energy) bin edges for the
            energy maps.
        log_sob_gamma_splines (List[List[scipy.interpolate.UnivariateSpline]]):
            A 2D list of spline fits of the log(signal-over-background) vs.
            gamma at a binned energy and sin(dec).
        """
    _sin_dec_bins: np.array = field(init=False)
    _log_energy_bins: np.array = field(init=False)
    _log_sob_gamma_splines: List[List[Spline]] = field(init=False)

@dataclass
class _I3EventModelDefaultsBase(_models.TdEventModelDefaultsBase):
    """Docstring"""
    signal_sin_dec_bins: InitVar[Union[np.array, int]] = field(default=50)
    log_energy_bins: InitVar[Union[np.array, int]] = field(default=50)
    gamma_bins: InitVar[Union[np.array, int]] = field(default=50)
    verbose: InitVar[bool] = field(default=False)

@dataclass
class I3EventModel(
    _models.TdEventModel,
    _I3EventModelDefaultsBase,
    _I3EventModelBase,
):
    """Docstring"""
    def __post_init__(
        self,
        source: sources.Source,
        data: np.ndarray,
        sim: np.ndarray,
        grl: np.ndarray,
        gamma: float,
        sampling_width: Optional[float],
        background_sin_dec_bins: Union[np.array, int],
        background_window: float,
        withinwindow: bool,
        signal_sin_dec_bins: Union[np.array, int],
        log_energy_bins: Union[np.array, int],
        gamma_bins: Union[np.array, int],
        verbose: bool,
    ) -> None:
        """Docstring"""
        super().__post_init__(
            source,
            data,
            sim,
            grl,
            gamma,
            sampling_width,
            background_sin_dec_bins,
            background_window,
            withinwindow,
        )

        if isinstance(signal_sin_dec_bins, int):
            signal_sin_dec_bins = np.linspace(-1, 1, 1 + signal_sin_dec_bins)
        self._sin_dec_bins = signal_sin_dec_bins

        if isinstance(log_energy_bins, int):
            log_energy_bins = np.linspace(1, 8, 1 + log_energy_bins)
        self._log_energy_bins = log_energy_bins

        if isinstance(gamma_bins, int):
            gamma_bins = np.linspace(-4.25, -0.5, 1 + gamma_bins)

        self._log_sob_gamma_splines = self._init_log_sob_gamma_splines(
            gamma_bins, verbose=verbose)

    def _init_sob_map(self, gamma: float, *args, verbose: bool = False,
                      **kwargs) -> np.array:
        """Creates sob histogram for a given spectral index (gamma).

        The UnivariateSpline function call uses these default arguments:
        k=1, s=0, ext=3. To replace any of these defaults, or to pass any other
        args/kwargs to UnivariateSpline, just pass them to this function.
        Spline is used here to smooth over the energies to get the values for
        each bin.

        Args:
            gamma: The gamma value to use to weight the signal.
            verbose: A flag to print progress.

        Returns:
            An array of signal-over-background values binned in sin(dec) and
            log(energy) for a given gamma.
        """
        if verbose:
            print(f'Building map for gamma = {gamma}...', end='')
        bins = np.array([self._sin_dec_bins, self._log_energy_bins])
        bin_centers = bins[1, :-1] + np.diff(bins[1]) / 2

        # background
        bg_h, _, _ = np.histogram2d(self._data['sindec'], self._data['logE'],
                                    bins=bins, density=True)

        # signal
        sig_w = self._sim['ow'] * self._sim['trueE']**gamma
        sig_h, _, _ = np.histogram2d(self._sim['sindec'], self._sim['logE'],
                                     bins=bins, weights=sig_w, density=True)

        # Normalize histograms by dec band
        bg_h /= np.sum(bg_h, axis=1)[:, None]
        sig_h /= np.sum(sig_h, axis=1)[:, None]

        # div-0 okay here
        with np.errstate(divide='ignore', invalid='ignore'):
            ratio = sig_h / bg_h

        if 'k' not in kwargs:
            kwargs['k'] = 1
        if 's' not in kwargs:
            kwargs['s'] = 0
        if 'ext' not in kwargs:
            kwargs['ext'] = 3

        for i in range(ratio.shape[0]):
            # Pick out the values we want to use.
            # We explicitly want to avoid NaNs and infinities
            good = np.isfinite(ratio[i]) & (ratio[i] > 0)
            good_bins, good_vals = bin_centers[good], ratio[i][good]

            # Do a linear interpolation across the energy range
            spline = Spline(good_bins, good_vals, *args, **kwargs)

            # And store the interpolated values
            ratio[i] = spline(bin_centers)
        if verbose:
            print('done')
        return ratio

    def _init_log_sob_gamma_splines(self, gamma_bins: np.array, *args,
                                    verbose: bool = False,
                                    **kwargs) -> List[List[Spline]]:
        """Builds a 3D hist of sob vs. sin(dec), log(energy), and gamma, then
            returns splines of sob vs. gamma.

        The UnivariateSpline function call uses these default arguments:
        k=3, s=0, ext='raise'. To replace any of these defaults, or to pass any
        other args/kwargs to UnivariateSpline, just pass them to this function.
        Spline is used here to smooth over the energies to get the values for
        each bin.

        Args:
            gamma_bins: The spectral indicies at which to build the histograms.
            verbose: A flag to print progress.

        Returns: A Nested spline list of shape (sin_dec_bins, log_energy_bins).
        """
        if verbose:
            print('Building signal-over-background maps...')
        sob_maps = np.array([self._init_sob_map(gamma, verbose=verbose)
                             for gamma in gamma_bins])
        if verbose:
            print('done.')

        if 'k' not in kwargs:
            kwargs['k'] = 3
        if 's' not in kwargs:
            kwargs['s'] = 0
        if 'ext' not in kwargs:
            kwargs['ext'] = 'raise'

        transposed_log_sob_maps = np.log(sob_maps.transpose(1, 2, 0))

        if verbose:
            print('Fitting log(signal-over-background vs. gamma splines)...',
                  end='')

        splines = [[
            Spline(gamma_bins, log_ratios, *args, **kwargs)
            for log_ratios in dec_bin
        ] for dec_bin in transposed_log_sob_maps]

        if verbose:
            print('done')

        return splines

    def log_sob_spline_prepro(
        self,
        events: np.ndarray,
    ) -> Tuple[np.ndarray, List]:
        """Docstring"""
        # Get the bin that each event belongs to
        sin_dec_idx = np.searchsorted(self._sin_dec_bins[:-1],
                                      events['sindec'])

        log_energy_idx = np.searchsorted(self._log_energy_bins[:-1],
                                         events['logE'])

        spline_idxs, event_spline_idxs = np.unique(
            [sin_dec_idx - 1, log_energy_idx - 1],
            return_inverse=True,
            axis=1
        )

        splines = [
            self._log_sob_gamma_splines[i][j]
            for i, j in spline_idxs.T
        ]

        return np.array(event_spline_idxs, dtype=int), splines

    def get_sob_energy(
        self,
        gamma: float,
        splines: List[Spline],
        event_spline_idxs: np.ndarray,
    ) -> np.array:
        """Docstring"""
        spline_evals = np.exp([spline(gamma) for spline in splines])
        return spline_evals[event_spline_idxs]


"""Kernels for scikit-image label.

These are copied from CuPy, with modification to add a greyscale_mode
parameter as needed for scikit-image.

"""
import cupy
import numpy

def _label(x, structure, y, greyscale_mode=False):
    elems = numpy.where(structure != 0)
    vecs = [elems[dm] - 1 for dm in range(x.ndim)]
    offset = vecs[0]
    for dm in range(1, x.ndim):
        offset = offset * 3 + vecs[dm]
    indxs = numpy.where(offset < 0)[0]
    dirs = [[vecs[dm][dr] for dm in range(x.ndim)] for dr in indxs]
    dirs = cupy.array(dirs, dtype=numpy.int32)
    ndirs = indxs.shape[0]
    y_shape = cupy.array(y.shape, dtype=numpy.int32)
    count = cupy.zeros(2, dtype=numpy.int32)
    _kernel_init()(x, y)
    try:
        int_t = int_types[y.dtype.char]
    except KeyError:
        raise ValueError("y must have int32, uint16, uint32 or uint64 dtype")
    if int_t != "int":
        raise NotImplementedError(
            "Currently only 32-bit integer case is implemented"
        )
    if greyscale_mode:
        _kernel_connect(True, int_t)(
            x, y_shape, dirs, ndirs, x.ndim, y, size=y.size
        )
    else:
        _kernel_connect(False, int_t)(
            y_shape, dirs, ndirs, x.ndim, y, size=y.size
        )
    _kernel_count()(y, count, size=y.size)
    maxlabel = int(count[0])  # synchronize
    labels = cupy.empty(maxlabel, dtype=numpy.int32)
    _kernel_labels()(y, count, labels, size=y.size)
    _kernel_finalize()(maxlabel, cupy.sort(labels), y, size=y.size)
    return maxlabel

"""
Elementwise kernels for use by label
"""

def _kernel_init():
    return cupy.ElementwiseKernel(
        "X x",
        "Y y",
        "if (x == 0) { y = -1; } else { y = i; }",
        "cucim_nd_label_init",
    )

def _kernel_connect(greyscale_mode=False, int_t="int"):
    """
    Notes
    -----
    dirs is a (n_neig//2, ndim) of relative offsets to the neighboring voxels.
    For example, for structure = np.ones((3, 3)):
        dirs = array([[-1, -1],
                      [-1,  0],
                      [-1,  1],
                      [ 0, -1]], dtype=int32)
    (Implementation assumes a centro-symmetric structure)
    ndirs = dirs.shape[0]

    In the dirs loop below, there is a loop over the ndim neighbors:
        Here, index j corresponds to the current pixel and k is the current
        neighbor location.
    """
    in_params = "raw int32 shape, raw int32 dirs, int32 ndirs, int32 ndim"
    if greyscale_mode:
        # greyscale mode -> different values receive different labels
        x_condition = "if (x[k] != x[j]) continue;"
        in_params = "raw X x, " + in_params
    else:
        # binary mode -> all non-background voxels treated the same
        x_condition = ""

    # Note: atomicCAS is implemented for int, unsigned short, unsigned int, and
    # unsigned long long

    code = """
        if (y[i] < 0) continue;
        for (int dr = 0; dr < ndirs; dr++) {{
            {int_t} j = i;
            {int_t} rest = j;
            {int_t} stride = 1;
            {int_t} k = 0;
            for (int dm = ndim-1; dm >= 0; dm--) {{
                int pos = rest % shape[dm] + dirs[dm + dr * ndim];
                if (pos < 0 || pos >= shape[dm]) {{
                    k = -1;
                    break;
                }}
                k += pos * stride;
                rest /= shape[dm];
                stride *= shape[dm];
            }}
            if (k < 0) continue;
            if (y[k] < 0) continue;
            {x_condition}
            while (1) {{
                while (j != y[j]) {{ j = y[j]; }}
                while (k != y[k]) {{ k = y[k]; }}
                if (j == k) break;
                if (j < k) {{
                    {int_t} old = atomicCAS( &y[k], (Y)k, (Y)j );
                    if (old == k) break;
                    k = old;
                }}
                else {{
                    {int_t} old = atomicCAS( &y[j], (Y)j, (Y)k );
                    if (old == j) break;
                    j = old;
                }}
            }}
        }}
        """.format(
        x_condition=x_condition, int_t=int_t
    )

    return cupy.ElementwiseKernel(
        in_params, "raw Y y", code, "cucim_nd_label_connect",
    )

def _kernel_count():
    return cupy.ElementwiseKernel(
        "",
        "raw Y y, raw int32 count",
        """
        if (y[i] < 0) continue;
        int j = i;
        while (j != y[j]) { j = y[j]; }
        if (j != i) y[i] = j;
        else atomicAdd(&count[0], 1);
        """,
        "cucim_nd_label_count",
    )

def _kernel_labels():
    return cupy.ElementwiseKernel(
        "",
        "raw Y y, raw int32 count, raw int32 labels",
        """
        if (y[i] != i) continue;
        int j = atomicAdd(&count[1], 1);
        labels[j] = i;
        """,
        "cucim_nd_label_labels",
    )

def _kernel_finalize():
    return cupy.ElementwiseKernel(
        "int32 maxlabel",
        "raw int32 labels, raw Y y",
        """
        if (y[i] < 0) {
            y[i] = 0;
            continue;
        }
        int yi = y[i];
        int j_min = 0;
        int j_max = maxlabel - 1;
        int j = (j_min + j_max) / 2;
        while (j_min < j_max) {
            if (yi == labels[j]) break;
            if (yi < labels[j]) j_max = j - 1;
            else j_min = j + 1;
            j = (j_min + j_max) / 2;
        }
        y[i] = j + 1;
        """,
        "cucim_nd_label_finalize",
    )

int_types = {
    "i": "int",
    "H": "unsigned short",
    "I": "unsigned int",
    "L": "unsigned long long",
}