""" Module storing the primary driver script used for the v1 release of cosmoDC2.
"""
import os
import psutil
import numpy as np
import h5py
import re
import healpy as hp
from time import time
from astropy.table import Table, vstack
from astropy.cosmology import FlatLambdaCDM
from astropy.utils.misc import NumpyRNGContext
from cosmodc2.load_gio_halos import load_gio_halo_snapshot
from cosmodc2.sdss_colors import assign_restframe_sdss_gri
from cosmodc2.sdss_colors.sigmoid_magr_model import magr_monte_carlo

from cosmodc2.stellar_mass_remapping import remap_stellar_mass_in_snapshot
from galsampler import halo_bin_indices, source_halo_index_selection
from galsampler.cython_kernels import galaxy_selection_kernel
from halotools.utils import crossmatch

from cosmodc2.synthetic_subhalos import map_mstar_onto_lowmass_extension
from cosmodc2.synthetic_subhalos import model_synthetic_cluster_satellites
#from cosmodc2.synthetic_subhalos import synthetic_logmpeak

fof_halo_mass = 'fof_halo_mass'
mass = 'mass'
fof_max = 14.5
H0 = 71.0
OmegaM = 0.2648
OmegaB = 0.0448

#unique galaxy_id
galaxy_id_factor = int(1e4)  #  factor to guarantee unique galaxy_id across blocks in snapshot

desired_logm_completeness=9.8

def write_umachine_snapshot_mock_to_disk(
        umachine_mock_fname, umachine_halo_fname,
        target_halo_catalog_fname, snapshot, blocks, output_snapshot_mock_fname_list,
        redshift, commit_hash, Lbox=3000.):
    """
    Main driver function used to paint SDSS fluxes onto UniverseMachine,
    GalSample the mock into the halo snapshot and write the snapshot mock to disk.

    Parameters
    ----------
    umachine_mock_fname : str
        the absolute path to the value-added UniverseMachine snapshot mock

    umachine_halo_fname : str
        the absolute path to the
        value-added host halo catalog hosting the UniverseMachine snapshot mock

    target_halo_catalog_fname : str
        the absolute path to the gio file(s) of
        source halos into which UniverseMachine will be GalSampled

    snapshot : str
        snapshot being processed

    blocks : list
        list of blocks (str format) of halo-catalog file being processed

    output_snapshot_mock_fname_list : list
        list of absolute paths to the output snapshot mock filenames (1 per block)

    redshift : str
        value of the redshift for the target halo catalog

    commit_hash : string
        Commit hash of the version of the cosmodc2 repo used when
        calling this function.

        After updating the cosmodc2 repo to the desired version,
        the commit_hash can be determined by navigating to the root
        directory and typing ``git log --pretty=format:'%h' -n 1``

    """

    start_time = time()
    process = psutil.Process(os.getpid())

    #  determine factor for number of synthetic galaxies for this snapshot
    #  synthetic_number_factor = (OR_size/MDPL2_size)**3

    #  initialize book-keeping variables
    fof_halo_mass_max = 0.
    Ngals_total = 0

    print('\nStarting snapshot processing')
    print('Redshift for halo catalog = {}'.format(redshift))

    #  Get galaxy properties from UM catalogs
    print("\n...loading z = {0:.2f} galaxy catalog into memory".format(redshift))
    um_mock = Table.read(umachine_mock_fname, path='data')
    print('.....{} galaxies read in'.format(len(um_mock)))

    ### Get source halos
    print("\n...loading z = {0:.2f} source-halo catalogs into memory".format(redshift))
    source_halos = Table.read(umachine_halo_fname, path='data')

    #  Bin the halos in source simulation by mass
    dlogM = 0.15
    mass_bins = 10.**np.arange(10.5, 14.5+dlogM, dlogM)
    source_halos['mass_bin'] = halo_bin_indices(
        mass=(source_halos['mvir'], mass_bins))

    for block, output_snap_fname  in zip(blocks, output_snapshot_mock_fname_list):
        new_time_stamp = time()
        #  determine seed from output filename (includes snapshot and block)
        seed = get_random_seed(os.path.basename(output_snap_fname))
        print('\n.Processing block {}'.format(block))
        print('.Using galaxy_id factor {} + galaxy_id offset {}'.format(galaxy_id_factor, block))
        print('.Using seed = {} (for block {})'.format(seed, block))

        # copy um_mock to new table for this block
        mock = um_mock.copy()

        print("\n...loading step {} fof target-halo catalogs into memory".format(snapshot))
        target_halos = load_gio_halo_snapshot(target_halo_catalog_fname, block=block)
        target_halos.rename_column('fof_halo_tag', 'fof_halo_id')
        target_halos.rename_column('fof_halo_center_x', 'x')
        target_halos.rename_column('fof_halo_center_y', 'y')
        target_halos.rename_column('fof_halo_center_z', 'z')
        target_halos.rename_column('fof_halo_mean_vx', 'vx')
        target_halos.rename_column('fof_halo_mean_vy', 'vy')
        target_halos.rename_column('fof_halo_mean_vz', 'vz')
        max_fof_halo_mass = np.max(target_halos[fof_halo_mass].quantity.value)
        fof_halo_mass_max = max(max_fof_halo_mass, fof_halo_mass_max)
        print('.....Maximum fof halo mass = {:.3e}'.format(max_fof_halo_mass))

        print("\n...Finding halo--halo correspondence with GalSampler")
        #  Bin the halos in target simulation by mass
        target_halos['mass_bin'] = halo_bin_indices(
            mass=(target_halos[fof_halo_mass], mass_bins))

        #  Randomly draw halos from corresponding mass bins
        nhalo_min = 10
        source_halo_bin_numbers = source_halos['mass_bin']
        target_halo_bin_numbers = target_halos['mass_bin']
        target_halo_ids = target_halos['fof_halo_id']
        _result = source_halo_index_selection(source_halo_bin_numbers,
                      target_halo_bin_numbers, target_halo_ids, nhalo_min, mass_bins, seed=seed)
        source_halo_indx, matching_target_halo_ids = _result

        #  Transfer quantities from the source halos to the corresponding target halo
        target_halos['source_halo_id'] = source_halos['halo_id'][source_halo_indx]
        target_halos['matching_mvir'] = source_halos['mvir'][source_halo_indx]
        target_halos['richness'] = source_halos['richness'][source_halo_indx]
        target_halos['first_galaxy_index'] = source_halos['first_galaxy_index'][source_halo_indx]

        ################################################################################
        #  Use GalSampler to calculate the indices of the galaxies that will be selected
        ################################################################################
        print("\n...GalSampling z={0:.2f} galaxies to OuterRim halos".format(redshift))

        source_galaxy_indx = np.array(galaxy_selection_kernel(
            target_halos['first_galaxy_index'].astype('i8'),
            target_halos['richness'].astype('i4'), target_halos['richness'].sum()))

        ########################################################################
        #  Correct stellar mass for low-mass subhalos and create synthetic mpeak
        ########################################################################
        #print("...correcting low mass mpeak and assigning synthetic mpeak values")
        #  First generate the appropriate number of synthetic galaxies for the snapshot
        #mpeak_synthetic_snapshot = 10**synthetic_logmpeak(
        #    mock['mpeak'], seed=seed, desired_logm_completeness=synthetic_halo_minimum_mass)
        #print('...assembling {} synthetic galaxies'.format(len(mpeak_synthetic_snapshot)))

        ########################################################################
        #  Assign stellar mass
        ########################################################################
        print("...re-assigning high-mass mstar values")

        #  Map stellar mass onto mock using target halo mass instead of UM Mpeak for cluster BCGs
        new_mstar = remap_stellar_mass_in_snapshot(redshift, mock['mpeak'], mock['obs_sm'])
        mock.rename_column('obs_sm', '_obs_sm_orig_um_snap')
        mock['obs_sm'] = new_mstar

        #  Add call to map_mstar_onto_lowmass_extension function after pre-determining low-mass slope
        print("...re-assigning low-mass mstar values")
        min_obs_sm = np.min(mock['obs_sm'])
        mpeak_synthetic_snapshot = np.asarray([])
        new_mstar_real, mstar_synthetic_snapshot = map_mstar_onto_lowmass_extension(
            mock['mpeak'], mock['obs_sm'], mpeak_synthetic_snapshot,
            desired_logm_completeness=desired_logm_completeness)
        mock['obs_sm'] = new_mstar_real
        new_min_obs_sm = np.min(mock['obs_sm'])
        print('.....New min(obs_sm) = {:.3e}; old min(obs_sm) = {:.3e}'.format(new_min_obs_sm, min_obs_sm))
        print('.....Number of shifted values = {}'.format(np.count_nonzero(mock['obs_sm'] < min_obs_sm)))

        ###################################################
        #  Map restframe Mr, g-r, r-i onto mock
        ###################################################
        #  use the redshift of the snapshot of the target simulation
        print("...assigning rest-frame Mr and colors")
        check_time = time()
        redshift_mock = np.zeros(len(mock)) + redshift
        msg = (".....using snapshot redshift to assign restframe colors")
        print(msg)

        magr, gr_mock, ri_mock, is_red_gr, is_red_ri = assign_restframe_sdss_gri(
            mock['upid'], mock['obs_sm'], mock['sfr_percentile'],
            mock['host_halo_mvir'], redshift_mock, seed=seed, use_substeps=False)
        #  check for bad values
        for m_id, m in zip(['magr', 'gr', 'ri'], [magr, gr_mock, ri_mock]):
            num_infinite = np.sum(~np.isfinite(m))
            if num_infinite > 0:
                print('.....Warning: {} infinite values in mock {}'.format(num_infinite, m_id))

        mock['restframe_extincted_sdss_abs_magr'] = magr
        mock['restframe_extincted_sdss_gr'] = gr_mock
        mock['restframe_extincted_sdss_ri'] = ri_mock
        mock['is_on_red_sequence_gr'] = is_red_gr
        mock['is_on_red_sequence_ri'] = is_red_ri
        print('.....time to assign_restframe_sdss_gri = {:.2f} secs'.format(time()-check_time))

        ########################################################################
        #  Assemble the output mock by snapshot
        ########################################################################

        print("\n...building output snapshot mock for snapshot {}".format(snapshot))
        output_mock = build_output_snapshot_mock(redshift, mock, target_halos,
                                                          source_galaxy_indx, galaxy_id_factor,
                                                          int(block), Lbox=Lbox)
        Ngals_total += len(output_mock['galaxy_id'])
        print('...saved {} galaxies to dict'.format(len(output_mock['galaxy_id'])))

        ########################################################################
        #  Write the output mock to disk
        ########################################################################
        if len(output_mock) > 0:
            check_time = time()
            write_output_mock_to_disk(output_snap_fname, output_mock, commit_hash, seed,
                                      redshift, snapshot, block, Lbox)
            print('...time to write mock to disk = {:.2f} minutes'.format((time()-check_time)/60.))

        time_stamp = time()
        msg = "\n.Block runtime = {0:.2f} minutes"
        print(msg.format((time_stamp-new_time_stamp)/60.))
        mem = ".Memory usage =  {0:.2f} GB"
        print(mem.format(process.memory_info().rss/1.e9))

    file_info = 'snapshot {}, blocks {}'.format(snapshot, ', '.join(blocks))
    print('\n.Maximum halo mass in {} = {}\n'.format(file_info, fof_halo_mass_max))
    print('.Number of galaxies in {} = {}\n'.format(file_info, Ngals_total))

    time_stamp = time()
    msg = "\nEnd-to-end runtime = {0:.2f} minutes\n"
    print(msg.format((time_stamp-start_time)/60.))

def get_random_seed(filename, seed_max=4294967095):  #reduce max seed by 200 to allow for 60 z shells
    import hashlib
    s = hashlib.md5(filename).hexdigest()
    seed = int(s, 16)

    #  enforce seed is below seed_max and odd
    seed = seed%seed_max
    if seed%2 == 0:
        seed = seed + 1
    return seed

def build_output_snapshot_mock(
            snapshot_redshift, umachine, target_halos, galaxy_indices, galaxy_id_factor,
            galaxy_id_offset, Lbox=0.):
    """
    Collect the GalSampled snapshot mock into an astropy table

    Parameters
    ----------
    snapshot_redshift : float
        Float of the snapshot redshift

    umachine : astropy.table.Table
        Astropy Table of shape (num_source_gals, )
        storing the UniverseMachine snapshot mock

    target_halos : astropy.table.Table
        Astropy Table of shape (num_target_halos, )
        storing the target halo catalog

    galaxy_indices: ndarray
        Numpy indexing array of shape (num_target_gals, )
        storing integers valued between [0, num_source_gals)

    galaxy_id_factor: integer
        Multiplicative factor to ensure unique galaxy id's in a snapshot

    galaxy_id_offset: integer
        Offset to ensure unique galaxy id's in a snapshot

    Returns
    -------
    dc2 : astropy.table.Table
        Astropy Table of shape (num_target_gals, )
        storing the GalSampled galaxy catalog
    """
    dc2 = Table()
    dc2['source_halo_id'] = umachine['hostid'][galaxy_indices]
    dc2['target_halo_id'] = np.repeat(
        target_halos['fof_halo_id'], target_halos['richness'])
    #needed for synthetic cluster satellite assignment
    dc2['target_halo_redshift'] = np.repeat(snapshot_redshift, len(dc2['target_halo_id']))
    dc2['target_halo_fof_halo_id'] = dc2['target_halo_id']

    #  copy target halo information
    dc2['source_halo_mvir'] = np.repeat(
        target_halos['matching_mvir'], target_halos['richness'])

    idxA, idxB = crossmatch(dc2['target_halo_id'], target_halos['fof_halo_id'])

    msg = "target IDs do not match!"
    assert np.all(dc2['source_halo_id'][idxA] == target_halos['source_halo_id'][idxB]), msg

    dc2['target_halo_x'] = 0.
    dc2['target_halo_y'] = 0.
    dc2['target_halo_z'] = 0.
    dc2['target_halo_vx'] = 0.
    dc2['target_halo_vy'] = 0.
    dc2['target_halo_vz'] = 0.

    dc2['target_halo_x'][idxA] = target_halos['x'][idxB]
    dc2['target_halo_y'][idxA] = target_halos['y'][idxB]
    dc2['target_halo_z'][idxA] = target_halos['z'][idxB]

    dc2['target_halo_vx'][idxA] = target_halos['vx'][idxB]
    dc2['target_halo_vy'][idxA] = target_halos['vy'][idxB]
    dc2['target_halo_vz'][idxA] = target_halos['vz'][idxB]

    dc2['target_halo_mass'] = 0.
    dc2['target_halo_mass'][idxA] = target_halos['fof_halo_mass'][idxB]

    source_galaxy_keys = ('host_halo_mvir', 'upid', 'mpeak',
            'host_centric_x', 'host_centric_y', 'host_centric_z',
            'host_centric_vx', 'host_centric_vy', 'host_centric_vz',
            'obs_sm', 'obs_sfr', 'sfr_percentile',
            'restframe_extincted_sdss_abs_magr',
            'restframe_extincted_sdss_gr', 'restframe_extincted_sdss_ri',
            'is_on_red_sequence_gr', 'is_on_red_sequence_ri',
            '_obs_sm_orig_um_snap', 'halo_id')
    for key in source_galaxy_keys:
        try:
            dc2[key] = umachine[key][galaxy_indices]
        except KeyError:
            msg = ("The build_output_snapshot_mock function was passed a umachine mock\n"
                "that does not contain the ``{0}`` key")
            raise KeyError(msg.format(key))

    max_umachine_halo_mass = np.max(umachine['mpeak'])
    ultra_high_mvir_halo_mask = (dc2['upid'] == -1) & (dc2['target_halo_mass'] > max_umachine_halo_mass)
    num_to_remap = np.count_nonzero(ultra_high_mvir_halo_mask)
    if num_to_remap > 0:
        print("...remapping stellar mass of {0} BCGs in ultra-massive halos".format(num_to_remap))

        halo_mass_array = dc2['target_halo_mass'][ultra_high_mvir_halo_mask]
        mpeak_array = dc2['mpeak'][ultra_high_mvir_halo_mask]
        mhalo_ratio = halo_mass_array/mpeak_array
        mstar_array = dc2['obs_sm'][ultra_high_mvir_halo_mask]
        redshift_array = dc2['target_halo_redshift'][ultra_high_mvir_halo_mask]
        upid_array = dc2['upid'][ultra_high_mvir_halo_mask]

        assert np.shape(halo_mass_array) == (num_to_remap, ), "halo_mass_array has shape = {0}".format(np.shape(halo_mass_array))
        assert np.shape(mstar_array) == (num_to_remap, ), "mstar_array has shape = {0}".format(np.shape(mstar_array))
        assert np.shape(redshift_array) == (num_to_remap, ), "redshift_array has shape = {0}".format(np.shape(redshift_array))
        assert np.shape(upid_array) == (num_to_remap, ), "upid_array has shape = {0}".format(np.shape(upid_array))
        assert np.all(mhalo_ratio >= 1), "Bookkeeping error: all values of mhalo_ratio ={0} should be >= 1".format(mhalo_ratio)

        dc2['obs_sm'][ultra_high_mvir_halo_mask] = mstar_array*(mhalo_ratio**0.5)
        dc2['restframe_extincted_sdss_abs_magr'][ultra_high_mvir_halo_mask] = magr_monte_carlo(
            dc2['obs_sm'][ultra_high_mvir_halo_mask], upid_array, redshift_array)
        idx = np.argmax(dc2['obs_sm'])
        halo_id_most_massive = dc2['halo_id'][idx]
        assert dc2['obs_sm'][idx] < 10**13.5, "halo_id = {0} has stellar mass {1:.3e}".format(
            halo_id_most_massive, dc2['obs_sm'][idx])

    dc2['x'] = dc2['target_halo_x'] + dc2['host_centric_x']
    dc2['vx'] = dc2['target_halo_vx'] + dc2['host_centric_vx']

    dc2['y'] = dc2['target_halo_y'] + dc2['host_centric_y']
    dc2['vy'] = dc2['target_halo_vy'] + dc2['host_centric_vy']

    dc2['z'] = dc2['target_halo_z'] + dc2['host_centric_z']
    dc2['vz'] = dc2['target_halo_vz'] + dc2['host_centric_vz']

    print('...number of galaxies before adding synthetic satellites = {}'.format(len(dc2['halo_id'])))
    print("...generating and stacking any synthetic cluster satellites")
    fake_cluster_sats = model_synthetic_cluster_satellites(dc2, Lbox=Lbox, snapshot=True)
    if len(fake_cluster_sats) > 0:
        check_time = time()
        dc2 = vstack((dc2, fake_cluster_sats))
        print('...time to create {} galaxies in fake_cluster_sats = {:.2f} secs'.format(len(fake_cluster_sats['target_halo_id']), time()-check_time))

    # delete duplicate/unnecessary column after fake cluster satellites are added
    dc2.remove_column('target_halo_fof_halo_id')
    dc2.remove_column('target_halo_redshift')

    dc2['galaxy_id'] = np.arange(len(dc2['target_halo_id'])).astype(int)*galaxy_id_factor + galaxy_id_offset
    print('...Min and max galaxy_id = {} -> {}'.format(np.min(dc2['galaxy_id']), np.max(dc2['galaxy_id'])))

    #  Use gr and ri color to compute gi flux
    dc2['restframe_extincted_sdss_abs_magg'] = (
        dc2['restframe_extincted_sdss_gr'] +
        dc2['restframe_extincted_sdss_abs_magr'])
    dc2['restframe_extincted_sdss_abs_magi'] = (
        -dc2['restframe_extincted_sdss_ri'] +
        dc2['restframe_extincted_sdss_abs_magr'])

    #convert table to dict
    check_time = time()
    output_dc2 = {}
    for k in dc2.keys():
        output_dc2[k] = dc2[k].quantity.value

    print('...time to new dict = {:.4f} secs'.format(time()-check_time))

    return output_dc2

def write_output_mock_to_disk(output_snapshot_mock_fname, output_mock, commit_hash, seed,
                              redshift, snapshot, block, Lbox,
                              versionMajor=0, versionMinor=1, versionMinorMinor=0):
    """
    """

    print("\n...writing to file {} using commit hash {}".format(output_snapshot_mock_fname, commit_hash))
    hdfFile = h5py.File(output_snapshot_mock_fname, 'w')
    hdfFile.create_group('metaData')
    gkey = 'galaxyProperties'

    hdfFile['metaData']['commit_hash'] = commit_hash
    hdfFile['metaData']['seed'] = seed
    hdfFile['metaData']['versionMajor'] = versionMajor
    hdfFile['metaData']['versionMinor'] = versionMinor
    hdfFile['metaData']['versionMinorMinor'] = versionMinorMinor
    hdfFile['metaData']['H_0'] = H0
    hdfFile['metaData']['Omega_matter'] = OmegaM
    hdfFile['metaData']['Omega_b'] = OmegaB
    hdfFile['metaData']['box_size'] = Lbox
    hdfFile['metaData']['redshift'] = redshift
    hdfFile['metaData']['timestep'] = snapshot
    hdfFile['metaData']['block_number'] = block
    for d in ['x', 'y', 'z']:
        hdfFile['metaData'][d + '_max'] = np.max(output_mock[d])
        hdfFile['metaData'][d + '_min'] = np.min(output_mock[d])

    gGroup = hdfFile.create_group(gkey)
    check_time = time()
    for k, v in output_mock.items():
        gGroup[k] = v

    print('.....time to write group {} = {:.4f} secs'.format(gkey, time()-check_time))

    check_time = time()
    hdfFile.close()
    print('.....time to close file {:.4f} secs'.format(time()-check_time))


"""
Classes for generating lists of photons
"""
import numpy as np
from yt.funcs import get_pbar
from pyxsim.lib.sky_functions import pixel_to_cel, \
    scatter_events, doppler_shift
from yt.utilities.physical_constants import clight
from yt.utilities.cosmology import Cosmology
from yt.utilities.orientation import Orientation
from yt.utilities.parallel_tools.parallel_analysis_interface import \
    communication_system
from yt.units.yt_array import YTArray
import h5py
from pyxsim.spectral_models import absorb_models
from pyxsim.utils import parse_value, mylog
from soxs.utils import parse_prng

comm = communication_system.communicators[-1]

init_chunk = 100000

def determine_fields(ds, source_type, point_sources):
    from yt.geometry.particle_geometry_handler import ParticleIndex
    # Figure out if this is a particle field or otherwise
    ptype = (source_type in ds.particle_types) | \
            (source_type in ds.known_filters) | \
            ((source_type == "gas") & isinstance(ds.index, ParticleIndex))
    if ptype:
        ppos = [f"particle_position_{ax}" for ax in "xyz"]
        pvel = [f"particle_velocity_{ax}" for ax in "xyz"]
        if source_type in ds.known_filters:
            if ds.known_filters[source_type].filtered_type == "gas":
                ppos = ["x", "y", "z"]
                pvel = [f"velocity_{ax}" for ax in "xyz"]
        elif source_type == "gas":
            source_type = ds._sph_ptypes[0]
        position_fields = [(source_type, ppos[i]) for i in range(3)]
        velocity_fields = [(source_type, pvel[i]) for i in range(3)]
        if source_type in ds._sph_ptypes:
            width_field = (source_type, "smoothing_length")
        else:
            width_field = None
    else:
        position_fields = [("index", ax) for ax in "xyz"]
        velocity_fields = [(source_type, f"velocity_{ax}") for ax in "xyz"]
        width_field = ("index", "dx")
    if point_sources:
        width_field = None
    return position_fields, velocity_fields, width_field

def find_object_bounds(data_source):
    """
    This logic is required to determine the bounds of the object, which is
    solely for fixing coordinates at periodic boundaries
    """
    if hasattr(data_source, "base_object"):
        # This is a cut region so we'll figure out
        # its bounds from its parent object
        data_src = data_source.base_object
    else:
        data_src = data_source

    if hasattr(data_src, "left_edge"):
        # Region or grid
        c = 0.5 * (data_src.left_edge + data_src.right_edge)
        w = data_src.right_edge - data_src.left_edge
        le = -0.5 * w + c
        re = 0.5 * w + c
    elif hasattr(data_src, "radius") and not hasattr(data_src, "height"):
        # Sphere
        le = -data_src.radius + data_src.center
        re = data_src.radius + data_src.center
    else:
        # Not sure what to do with any other object yet, so just
        # return the domain edges and punt.
        mylog.warning("You are using a region that is not currently "
                      "supported for straddling periodic boundaries. "
                      "Check to make sure that your region does not "
                      "do so.")
        le = data_source.ds.domain_left_edge
        re = data_source.ds.domain_right_edge

    return le.to_value("kpc"), re.to_value("kpc")

def make_photons(photon_prefix, data_source, redshift, area,
                 exp_time, source_model, point_sources=False,
                 parameters=None, center=None, dist=None,
                 cosmology=None, velocity_fields=None):
    r"""
    Write a photon list dataset to disk from a yt data source and assuming a
    source model for the photons. The redshift, collecting area, exposure time,
    and cosmology are stored in the *parameters* dictionary which is passed to
    the *source_model* function.

    Parameters
    ----------
    photon_prefix : string
        The prefix of the filename(s) to contain the photon list. If run in
        serial, the filename will be "{photon_prefix}.h5", if run in
        parallel, the filenames will be "{photon_prefix}.{mpi_rank}.h5".
    data_source : :class:`~yt.data_objects.data_containers.YTSelectionContainer`
        The data source from which the photons will be generated.
    redshift : float
        The cosmological redshift for the photons.
    area : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
        The collecting area to determine the number of photons. If units are
        not specified, it is assumed to be in cm^2.
    exp_time : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
        The exposure time to determine the number of photons. If units are
        not specified, it is assumed to be in seconds.
    source_model : :class:`~pyxsim.source_models.SourceModel`
        A source model used to generate the photons.
    point_sources : boolean, optional
        If True, the photons will be assumed to be generated from the exact
        positions of the cells or particles and not smeared around within
        a volume. Default: False
    parameters : dict, optional
        A dictionary of parameters to be passed for the source model to use,
        if necessary.
    center : string or array_like, optional
        The origin of the photon spatial coordinates. Accepts "c", "max", or
        a coordinate. If array-like and without units, it is assumed to be in
        units of kpc. If not specified, pyxsim attempts to use the "center"
        field parameter of the data_source.
    dist : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`, optional
        The angular diameter distance, used for nearby sources. This may be
        optionally supplied instead of it being determined from the
        *redshift* and given *cosmology*. If units are not specified, it is
        assumed to be in kpc. To use this, the redshift must be set to zero.
    cosmology : :class:`~yt.utilities.cosmology.Cosmology`, optional
        Cosmological information. If not supplied, we try to get the
        cosmology from the dataset. Otherwise, LCDM with the default yt
        parameters is assumed.
    velocity_fields : list of fields
        The yt fields to use for the velocity. If not specified, the
        following will be assumed:
        ['velocity_x', 'velocity_y', 'velocity_z'] for grid datasets
        ['particle_velocity_x', 'particle_velocity_y', 'particle_velocity_z'] for particle datasets

    Returns
    -------
    A tuple of two integers, the number of photons and the number of cells with
    photons

    Examples
    --------
    >>> thermal_model = pyxsim.ThermalSourceModel("apec",  0.1, 10.0,
    ...                                           1000, Zmet=0.3)
    >>> redshift = 0.05
    >>> area = 6000.0 # assumed here in cm**2
    >>> time = 2.0e5 # assumed here in seconds
    >>> sp = ds.sphere("c", (500., "kpc"))
    >>> n_photons, n_cells = pyxsim.make_photons(sp, redshift, area,
    ...                                          time, thermal_model)
    """
    ds = data_source.ds

    if photon_prefix.endswith(".h5"):
        photon_prefix = photon_prefix[:-3]

    if comm.size > 1:
        photon_file = f"{photon_prefix}.{comm.rank:04d}.h5"
    else:
        photon_file = f"{photon_prefix}.h5"

    if parameters is None:
        parameters = {}
    if cosmology is None:
        if hasattr(ds, 'cosmology'):
            cosmo = ds.cosmology
        else:
            cosmo = Cosmology()
    else:
        cosmo = cosmology
    if dist is None:
        if redshift <= 0.0:
            msg = "If redshift <= 0.0, you must specify a distance to the " \
                  "source using the 'dist' argument!"
            mylog.error(msg)
            raise ValueError(msg)
        D_A = cosmo.angular_diameter_distance(0.0, redshift).to("Mpc")
    else:
        D_A = parse_value(dist, "kpc")
        if redshift > 0.0:
            mylog.warning("Redshift must be zero for nearby sources. "
                          "Resetting redshift to 0.0.")
            redshift = 0.0

    if isinstance(center, str):
        if center == "center" or center == "c":
            parameters["center"] = ds.domain_center
        elif center == "max" or center == "m":
            parameters["center"] = ds.find_max("density")[-1]
    elif isinstance(center, YTArray):
        parameters["center"] = center.in_units("code_length")
    elif isinstance(center, tuple):
        if center[0] == "min":
            parameters["center"] = ds.find_min(center[1])[-1]
        elif center[0] == "max":
            parameters["center"] = ds.find_max(center[1])[-1]
        else:
            raise RuntimeError
    elif isinstance(center, (list, np.ndarray)):
        parameters["center"] = ds.arr(center, "code_length")
    elif center is None:
        if hasattr(data_source, "left_edge"):
            parameters["center"] = 0.5*(data_source.left_edge +
                                        data_source.right_edge)
        else:
            parameters["center"] = data_source.get_field_parameter("center")

    parameters["fid_exp_time"] = parse_value(exp_time, "s")
    parameters["fid_area"] = parse_value(area, "cm**2")
    parameters["fid_redshift"] = redshift
    parameters["fid_d_a"] = D_A
    parameters["hubble"] = cosmo.hubble_constant
    parameters["omega_matter"] = cosmo.omega_matter
    parameters["omega_lambda"] = cosmo.omega_lambda
    parameters["center"].convert_to_units("kpc")

    if redshift > 0.0:
        mylog.info(f"Cosmology: h = {cosmo.hubble_constant}, "
                   f"omega_matter = {cosmo.omega_matter}, "
                   f"omega_lambda = {cosmo.omega_lambda}")
    else:
        mylog.info(f"Observing local source at distance {D_A}.")

    local_exp_time = parameters["fid_exp_time"].v/comm.size
    D_A = parameters["fid_d_a"].to_value("cm")
    dist_fac = 1.0/(4.*np.pi*D_A*D_A*(1.+redshift)**2)
    spectral_norm = parameters["fid_area"].v*local_exp_time*dist_fac

    source_model.setup_model(data_source, redshift, spectral_norm)

    p_fields, v_fields, w_field = determine_fields(ds,
                                                   source_model.ftype,
                                                   point_sources)

    if velocity_fields is not None:
        v_fields = velocity_fields

    if p_fields[0] == ("index", "x"):
        parameters["data_type"] = "cells"
    else:
        parameters["data_type"] = "particles"

    dw = ds.domain_width.to_value("kpc")
    le, re = find_object_bounds(data_source)
    c = parameters["center"].to_value("kpc")

    f = h5py.File(photon_file, "w")

    # Parameters

    p = f.create_group("parameters")
    p.create_dataset("fid_area", data=float(parameters["fid_area"]))
    p.create_dataset("fid_exp_time", data=float(parameters["fid_exp_time"]))
    p.create_dataset("fid_redshift", data=parameters["fid_redshift"])
    p.create_dataset("hubble", data=parameters["hubble"])
    p.create_dataset("omega_matter", data=parameters["omega_matter"])
    p.create_dataset("omega_lambda", data=parameters["omega_lambda"])
    p.create_dataset("fid_d_a", data=float(parameters["fid_d_a"]))
    p.create_dataset("data_type", data=parameters["data_type"])

    n_cells = 0
    n_photons = 0
    c_offset = 0
    p_offset = 0
    c_size = init_chunk
    p_size = init_chunk

    cell_fields = ["x", "y", "z", "vx", "vy", "vz", "num_photons", "dx"]

    d = f.create_group("data")
    for field in cell_fields + ["energy"]:
        if field == "num_photons":
            dtype = "int64"
        else:
            dtype = "float64"
        d.create_dataset(field, data=np.zeros(init_chunk, dtype=dtype),
                         maxshape=(None,), dtype=dtype, chunks=True)

    f.flush()

    for chunk in data_source.chunks([], "io"):

        chunk_data = source_model(chunk)

        if chunk_data is not None:

            chunk_nc, number_of_photons, idxs, energies = chunk_data

            chunk_nph = np.sum(number_of_photons)

            if chunk_nph == 0:
                continue

            if c_size < n_cells + chunk_nc:
                while chunk_nc + n_cells > c_size:
                    c_size *= 2
                for field in cell_fields:
                    d[field].resize((c_size,))

            if p_size < n_photons + chunk_nph:
                while chunk_nph + n_photons > p_size:
                    p_size *= 2
                d["energy"].resize((p_size,))

            for i, ax in enumerate("xyz"):
                pos = chunk[p_fields[i]][idxs].to_value("kpc")
                # Fix photon coordinates for regions crossing a periodic boundary
                if ds.periodicity[i]:
                    tfl = pos < le[i]
                    tfr = pos > re[i]
                    pos[tfl] += dw[i]
                    pos[tfr] -= dw[i]

                vel = chunk[v_fields[i]][idxs].to_value("km/s")
                # Coordinates are centered
                d[ax][c_offset:c_offset+chunk_nc] = pos-c[i]
                d[f"v{ax}"][c_offset:c_offset+chunk_nc] = vel

            d["num_photons"][c_offset:c_offset+chunk_nc] = number_of_photons
            d["energy"][p_offset:p_offset+chunk_nph] = energies

            if w_field is None:
                d["dx"][c_offset:c_offset+chunk_nc] = 0.0
            else:
                d["dx"][c_offset:c_offset+chunk_nc] = \
                    chunk[w_field][idxs].to_value("kpc")

            n_cells += chunk_nc
            n_photons += chunk_nph
            c_offset = n_cells
            p_offset = n_photons

        f.flush()

    if c_size > n_cells:
        for field in cell_fields:
            d[field].resize((n_cells,))

    if p_size > n_photons:
        d["energy"].resize((n_photons,))

    f.close()

    source_model.cleanup_model()

    all_nphotons = comm.mpi_allreduce(n_photons)
    all_ncells = comm.mpi_allreduce(n_cells)

    mylog.info("Finished generating photons.")
    mylog.info(f"Number of photons generated: {all_nphotons}")
    mylog.info(f"Number of cells with photons: {all_ncells}")

    return all_nphotons, all_ncells

def project_photons(photon_prefix, event_prefix, normal, sky_center,
                    absorb_model=None, nH=None, no_shifting=False,
                    north_vector=None, sigma_pos=None,
                    kernel="top_hat", prng=None):
    r"""
    Projects photons onto an image plane given a line of sight, and
    stores them in an HDF5 dataset which contains an event list.

    Parameters
    ----------
    photon_prefix : string
        The prefix of the filename(s) containing the photon list. If run in
        serial, the filename will be "{photon_prefix}.h5", if run in
        parallel, the filenames will be "{photon_prefix}.{mpi_rank}.h5".
    event_prefix : string
        The prefix of the filename(s) which will be written to contain the
        event list. If run in serial, the filename will be "{event_prefix}.h5",
        if run in parallel, the filename will be "{event_prefix}.{mpi_rank}.h5".
    normal : character or array-like
        Normal vector to the plane of projection. If "x", "y", or "z", will
        assume to be along that axis (and will probably be faster). Otherwise,
        should be an off-axis normal vector, e.g [1.0, 2.0, -3.0]
    sky_center : array-like
        Center RA, Dec of the events in degrees.
    absorb_model : string or :class:`~pyxsim.spectral_models.AbsorptionModel`
        A model for foreground galactic absorption, to simulate the
        absorption of events before being detected. This cannot be applied
        here if you already did this step previously in the creation of the
        :class:`~pyxsim.photon_list.PhotonList` instance. Known options for
        strings are "wabs" and "tbabs".
    nH : float, optional
        The foreground column density in units of 10^22 cm^{-2}. Only used
        if absorption is applied.
    no_shifting : boolean, optional
        If set, the photon energies will not be Doppler shifted.
    north_vector : a sequence of floats
        A vector defining the "up" direction. This option sets the
        orientation of the plane of projection. If not set, an arbitrary
        grid-aligned north_vector is chosen. Ignored in the case where a
        particular axis (e.g., "x", "y", or "z") is explicitly specified.
    sigma_pos : float, optional
        Apply a gaussian smoothing operation to the sky positions of the
        events. This may be useful when the binned events appear blocky due
        to their uniform distribution within simulation cells. However, this
        will move the events away from their originating position on the
        sky, and so may distort surface brightness profiles and/or spectra.
        Should probably only be used for visualization purposes. Supply a
        float here to smooth with a standard deviation with this fraction
        of the cell size. Default: None
    kernel : string, optional
        The kernel used when smoothing positions of X-rays originating from
        SPH particles, "gaussian" or "top_hat". Default: "top_hat".
    prng : integer or :class:`~numpy.random.RandomState` object
        A pseudo-random number generator. Typically will only be specified
        if you have a reason to generate the same set of random numbers,
        such as for a test. Default is to use the :mod:`numpy.random`
        module.

    Returns
    -------
    A integer for the number of events created

    Examples
    --------
    >>> L = np.array([0.1,-0.2,0.3])
    >>> n_events = pyxsim.project_photons("my_photons.h5", "my_events.h5", L,
    ...                                   [30., 45.], absorb_model='tbabs',
    ...                                   nH=0.04)
    """
    from yt.funcs import ensure_numpy_array
    prng = parse_prng(prng)

    if photon_prefix.endswith(".h5"):
        photon_prefix = photon_prefix[:-3]

    if event_prefix.endswith(".h5"):
        event_prefix = event_prefix[:-3]

    if not isinstance(normal, str):
        L = np.array(normal)
        orient = Orientation(L, north_vector=north_vector)
        x_hat = orient.unit_vectors[0]
        y_hat = orient.unit_vectors[1]
        z_hat = orient.unit_vectors[2]
    else:
        x_hat = np.zeros(3)
        y_hat = np.zeros(3)
        z_hat = np.zeros(3)

    if comm.size > 1:
        photon_file = f"{photon_prefix}.{comm.rank:04d}.h5"
        event_file = f"{event_prefix}.{comm.rank:04d}.h5"
    else:
        photon_file = f"{photon_prefix}.h5"
        event_file = f"{event_prefix}.h5"

    sky_center = ensure_numpy_array(sky_center)

    scale_shift = -1.0/clight.to_value("km/s")

    if isinstance(absorb_model, str):
        if absorb_model not in absorb_models:
            raise KeyError(f"{absorb_model} is not a known absorption model!")
        absorb_model = absorb_models[absorb_model]
    if absorb_model is not None:
        if nH is None:
            raise RuntimeError("You specified an absorption model, but didn't "
                               "specify a value for nH!")
        absorb_model = absorb_model(nH)
        if comm.rank == 0:
            mylog.info(f"Foreground galactic absorption: using the "
                       f"{absorb_model._name} model and nH = {nH}.")

    f = h5py.File(photon_file, "r")

    p = f["parameters"]

    data_type = p["data_type"].asstr()[()]

    if sigma_pos is not None and data_type == "particles":
        raise RuntimeError("The 'smooth_positions' argument should "
                           "not be used with particle-based datasets!")

    d = f["data"]

    D_A = p["fid_d_a"][()]*1.0e3

    if d["energy"].size == 0:

        mylog.warning(f"No photons are in file {photon_file}, so "
                      f"I am done.")
        n_events = 0

    else:

        fe = h5py.File(event_file, "w")

        pe = fe.create_group("parameters")
        pe.create_dataset("exp_time", data=float(p["fid_exp_time"][()]))
        pe.create_dataset("area", data=float(p["fid_area"][()]))
        pe.create_dataset("sky_center", data=sky_center)

        event_fields = ["xsky", "ysky", "eobs"]

        n_events = 0
        e_offset = 0
        e_size = init_chunk
        cell_chunk = init_chunk
        start_e = 0

        de = fe.create_group("data")
        for field in event_fields:
            de.create_dataset(field, data=np.zeros(init_chunk),
                              maxshape=(None,), chunks=True)

        if isinstance(normal, str):
            norm = "xyz".index(normal)
        else:
            norm = normal

        n_cells = d["num_photons"].size

        pbar = get_pbar("Projecting photons from cells/particles", n_cells)

        for start_c in range(0, n_cells, cell_chunk):

            end_c = min(start_c + cell_chunk, n_cells)

            n_ph = d["num_photons"][start_c:end_c]
            dx = d["dx"][start_c:end_c]
            end_e = start_e + n_ph.sum()
            eobs = d["energy"][start_e:end_e]

            if not no_shifting:
                if isinstance(normal, str):
                    shift = d[f"v{normal}"][start_c:end_c]*scale_shift
                else:
                    shift = (d["vx"][start_c:end_c]*z_hat[0] +
                             d["vy"][start_c:end_c]*z_hat[1] +
                             d["vz"][start_c:end_c]*z_hat[2])*scale_shift
                doppler_shift(shift, n_ph, eobs)

            if absorb_model is None:
                det = np.ones(eobs.size, dtype='bool')
                num_det = eobs.size
            else:
                det = absorb_model.absorb_photons(eobs, prng=prng)
                num_det = det.sum()

            if num_det > 0:

                xsky, ysky = scatter_events(norm, prng, kernel,
                                            data_type, num_det, det, n_ph,
                                            d["x"][start_c:end_c],
                                            d["y"][start_c:end_c],
                                            d["z"][start_c:end_c],
                                            dx, x_hat, y_hat)

                if data_type == "cells" and sigma_pos is not None:
                    sigma = sigma_pos*np.repeat(dx, n_ph)[det]
                    xsky += sigma*prng.normal(loc=0.0, scale=1.0, size=num_det)
                    ysky += sigma*prng.normal(loc=0.0, scale=1.0, size=num_det)

                xsky /= D_A
                ysky /= D_A

                pixel_to_cel(xsky, ysky, sky_center)

                if e_size < n_events + num_det:
                    while n_events + num_det > e_size:
                        e_size *= 2
                    for field in event_fields:
                        de[field].resize((e_size,))

                de["xsky"][e_offset:e_offset+num_det] = xsky
                de["ysky"][e_offset:e_offset+num_det] = ysky
                de["eobs"][e_offset:e_offset+num_det] = eobs[det]

                n_events += num_det
                e_offset = n_events

                f.flush()

            start_e = end_e

            pbar.update(end_c)

        pbar.finish()

        if e_size > n_events:
            for field in event_fields:
                de[field].resize((n_events,))

        fe.close()

    f.close()

    all_nevents = comm.mpi_allreduce(n_events)

    mylog.info(f"Detected {all_nevents} events.")

    return all_nevents


# -*- coding: utf-8 -*-
"""
"""
from __future__ import division, print_function, unicode_literals
from ..utilities.future_from_2 import str, object

#from phasor.utilities.print import print
#import declarative
import collections
#import copy
import numpy as np

from . import ports
from ..utilities.print import pprint

from ..system.matrix_injections import (
    FactorCouplingBase,
)

def pk_prefs(*preflist):
    def key(pk):
        p, k = pk
        pk = p | k
        ksort = []
        for ktype in preflist:
            ksort.append(
                str(pk.get(ktype, None))
            )
        pk.purge_keys(*preflist)
        ksort.extend(sorted(pk.items()))
        return tuple(ksort)
    return key

def lt_mult(lt, lt2):
    def lt_mult_inner(lt, lt2):
        if isinstance(lt2, list):
            sublist = []
            for sublt2 in lt2:
                sublist = sublist + lt_mult_inner(lt, sublt2)
            return sublist
        elif isinstance(lt2, tuple):
            #multiply the gains and merge the indices
            return [(lt2[0] * lt[0],) + lt[1:] + lt2[1:]]
        else:
            raise RuntimeError("BOO")
    if isinstance(lt, list):
        sublist = []
        for sublt in lt:
            sublist = sublist + lt_mult(sublt, lt2)
        return sublist
    elif isinstance(lt, tuple):
        return lt_mult_inner(lt, lt2)
    else:
        raise RuntimeError("BOO")
    return

def lt_sort_collect(new_list):
    if not new_list:
        return new_list
    sNL = [(tuple(sorted(L[1:])), idx) for idx, L in enumerate(new_list)]
    sNL.sort()
    list_gen = []

    NL, idx = sNL[0]
    prev_NL = NL
    prev_gain = new_list[idx][0]
    for NL, idx in sNL[1:]:
        if prev_NL == NL:
            prev_gain += new_list[idx][0]
        else:
            list_gen.append(
                (prev_gain,) + tuple(prev_NL)
            )
            prev_NL = NL
            prev_gain = new_list[idx][0]
    list_gen.append(
        (prev_gain,) + tuple(prev_NL)
    )
    return list_gen

class ExpMatCoupling(FactorCouplingBase):
    """
    Generated by a dict-dict mapping ddlt[out][in] = list-tup-expr
    where list-tup-expr are a means of expressing addition and multiplication as lists and tups.

    lists represent series addition

    tups represent series multiplication, the first term is a raw number coefficient and any further are keys in the source
    vector
    """
    floating_req_set = None
    #holds a list of 3-tuples, each one carrying an in-set, out-set, and coupling-func
    #it behaves as though all in-set is connected to out-set for requirements analysis
    #but then during coupling matrix construction, only the edges returned from coupling-func
    #will be used
    floating_in_out_func_pairs = None

    #must redefine since it was a property
    edges_req_pkset_dict = None
    def __init__(
        self,
        dLt,
        in_map,
        out_map,
        N_ode = 1,
    ):
        self.N_ode     = N_ode
        self.dLt       = dLt
        self.out_map   = out_map
        self.in_map    = in_map
        self.solution  = dict()
        self.vals_prev = dict()
        self.vals_inj  = dict()

        #make the internal solution have no action initially (so that cavity feedback converges faster)
        for pk_internal, pk_in in self.in_map.items():
            pk_out = self.out_map.get(pk_internal, None)
            if pk_out is not None:
                self.solution[pk_in, pk_out] = 1

        #all edges are generated immediately. Currently assumes full density
        self.edges_NZ_pkset_dict = {}
        self.edges_pkpk_dict = {}
        self.edges_req_pkset_dict = {}

        def gen_edge_func(pk_in, pk_out):
            return lambda sV, sB: self.edge_func(pk_in, pk_out, sV, sB)

        self.floating_req_set = frozenset(self.in_map.values())
        ins = set(self.in_map.values())
        outs = set([pkt for pkt in self.out_map.values() if pkt is not None])
        self.floating_in_out_func_pairs = [(ins, outs, self.edge_mat_func)]

        def gen_src_func_out(pk_out):
            return lambda sV, sB: self.source_func_out(pk_out, sV, sB)
        self.sources_pk_dict = {}
        self.sources_NZ_pkset_dict = {}
        for pk_out in self.out_map.values():
            if pk_out is None:
                continue
            self.sources_pk_dict[pk_out] = gen_src_func_out(pk_out)
            self.sources_NZ_pkset_dict[pk_out] = frozenset()

        pks = set()
        def pks_grab(lt):
            if isinstance(lt, list):
                for sublt in lt:
                    pks_grab(sublt)
            elif isinstance(lt, tuple):
                for pk in lt[1:]:
                    pks.add(pk)
            else:
                raise RuntimeError("BOO")
            return
        for pk_out, lt in self.dLt.items():
            pks.add(pk_out)
            pks_grab(lt)

        #print(ins_p, sol_vector)
        self.pks = list(pks)
        #TODO: debug config
        print("Number of states: ", len(pks))
        #pprint(pks)
        self.pks.sort(key = pk_prefs(
            ports.QuantumKey,
            ports.ElementKey,
            ports.PortKey,
            ports.OpticalFreqKey,
            ports.ClassicalFreqKey,
            ports.PolKEY,
        ))
        self.pks_inv = dict()
        for idx, pk in enumerate(self.pks):
            self.pks_inv[pk] = idx

        h = 1 / self.N_ode
        #remap the index keys into integer indexes for speed
        dLt_accel = dict()
        def lt_remap(lt):
            if isinstance(lt, list):
                sublist = []
                for sublt in lt:
                    sublist.append(lt_remap(sublt))
                return sublist
            elif isinstance(lt, tuple):
                #get the gain
                newtup = [lt[0] * h]
                for pk in lt[1:]:
                    newtup.append(self.pks_inv[pk])
                return tuple(newtup)
            else:
                raise RuntimeError("BOO")
            return
        for pk_out, lt in self.dLt.items():
                dLt_accel[self.pks_inv[pk_out]] = lt_sort_collect(lt_remap(lt))
        self.dLt_accel = dLt_accel

        #pprint("PKS:")
        #pprint(pks)

    _prev_sol_vector = None

    def update_solution(self, sol_vector):
        if self._prev_sol_vector is None or sol_vector != self._prev_sol_vector:
            self._prev_sol_vector = sol_vector
            self.generate_solution(sol_vector)
        return

    def edge_mat_func(self, sol_vector, sB):
        self.update_solution(sol_vector)
        return self.solution

    def source_func_out(self, pk_out, sol_vector, sB):
        self.update_solution(sol_vector)
        return self.vals_inj.get(pk_out, 0)

    def generate_solution(self, sol_vector):
        pks = self.pks
        pkv = np.empty(len(pks), dtype=object)

        all_zeros = True
        for idx, pk in enumerate(pks):
            pk_in = self.in_map[pk]
            #print("PK_G: ", (ins_p, pk))
            sol_val = sol_vector.get(pk_in, 0)
            prev_val = self.vals_prev.get(pk_in, 0)
            tot_val = sol_val - prev_val
            new_val = np.copy(sol_val)
            if np.all(new_val == 0):
                new_val = 0
            pkv[idx] = new_val
            if np.any(abs(tot_val) > 1e-8 * abs(sol_val)):
                all_zeros = False
        if all_zeros:
            #TODO debug statements
            #print("ALL ZEROS")
            #no need to update the solution vector since nothing changed
            return
            pass
        else:
            #TODO debug statements
            #print("NOT ZEROS")
            pass
        #print("SOLVING ", self.pks)

        pk_original = pkv.copy()

        def lt_val(lt):
            assert(isinstance(lt, list))
            val = 0
            for sublt in lt:
                assert(isinstance(sublt, tuple))
                gain = np.copy(sublt[0])
                for pk_idx in sublt[1:]:
                    gain = gain * pkv[pk_idx]
                val += gain
            return val

        def lt_val_matrix_d(vec_d, lt_M_d):
            for pk_idx_from, lt in lt_M_d.items():
                assert(isinstance(lt, list))
                for sublt in lt:
                    assert(isinstance(sublt, tuple))
                    local_gain = np.copy(sublt[0])
                    for pk_idx in sublt[1:]:
                        val = pkv[pk_idx]
                        local_gain = local_gain * val
                    vec_d[pk_idx_from] += local_gain

        def lt_val_matrix_d_generate(pkv_nz_set, lt):
            lt_M_d = dict()
            full = 0
            reduced = 0
            assert(isinstance(lt, list))
            for sublt in lt:
                assert(isinstance(sublt, tuple))
                gain = sublt[0]
                for idx_idx, pk_idx_from in enumerate(sublt[1:]):
                    newtup = [gain]
                    full += 1
                    for idx_idx2, pk_idx in enumerate(sublt[1:]):
                        if idx_idx == idx_idx2:
                            continue
                        if pk_idx not in pkv_nz_set:
                            break
                        newtup.append(pk_idx)
                    else:
                        #only occurs if break was NOT called
                        reduced += 1
                        lt_inj = lt_M_d.setdefault(pk_idx_from, [])
                        lt_inj.append(
                            tuple(newtup)
                        )
            #print("REDUCED LT_d: ", reduced / full)
            return lt_M_d

        def lt_reduced_generate(pkv_nz_set, lt):
            newlt = []
            assert(isinstance(lt, list))
            for sublt in lt:
                assert(isinstance(sublt, tuple))
                newtup = [sublt[0]]
                for pk_idx in sublt[1:]:
                    if pk_idx not in pkv_nz_set:
                        break
                    newtup.append(pk_idx)
                else:
                    #only occurs if break was NOT called
                    newlt.append(
                        tuple(newtup)
                    )
            #print("REDUCED LT: ", len(newlt) / len(lt))
            return newlt

        #a double map matrix dMexp_s1[idx_out][idx_in]
        ##this dMexp_s1 does not have the one. For speed that is applied separately
        dMexp_s1 = collections.defaultdict(lambda : collections.defaultdict(lambda : 0))

        pkv_nz_set = set()
        for idx, val in enumerate(pkv):
            if np.any(val != 0):
                pkv_nz_set.add(idx)

        dLt_base = sorted(self.dLt_accel.items())
        dLt_idx_list = [T[0] for T in dLt_base]
        dLt_lt_list  = [lt_reduced_generate(pkv_nz_set, T[1]) for T in dLt_base]
        dLt_lt_M_list  = [lt_val_matrix_d_generate(pkv_nz_set, T[1]) for T in dLt_base]

        for idx_N in range(self.N_ode):
            idx_in_list = 0
            fullskip = 0
            while idx_in_list < len(dLt_idx_list):
                idx_pk = dLt_idx_list[idx_in_list]
                lt     = dLt_lt_list[idx_in_list]
                lt_M_d = dLt_lt_M_list[idx_in_list]
                idx_in_list += 1
                if not lt and not lt_M_d:
                    fullskip += 1

                dPK = lt_val(lt)
                Mexp_vec_d = collections.defaultdict(lambda : 0)
                lt_val_matrix_d(Mexp_vec_d, lt_M_d)

                if pkv[idx_pk] is 0:
                    if np.any(dPK != 0):
                        #print("STATUS CHANGE: ", idx_pk)
                        pkv_nz_set.add(idx_pk)
                        dLt_idx_list = [T[0] for T in dLt_base]
                        dLt_lt_list  = [lt_reduced_generate(pkv_nz_set, T[1]) for T in dLt_base]
                        dLt_lt_M_list  = [lt_val_matrix_d_generate(pkv_nz_set, T[1]) for T in dLt_base]
                if np.any(dPK != 0):
                    pkv[idx_pk] = pkv[idx_pk] + dPK

                dMexp_update = collections.defaultdict(lambda : 0)
                #print("LEN: ", len(dMexp_s1))
                #print("VEC: ", len(Mexp_vec_d))
                N_zero = 0
                N_skip = 0
                for idx_out, vec_val in Mexp_vec_d.items():
                    for idx_in, edge in dMexp_s1[idx_out].items():
                        if np.all(edge == 0):
                            N_zero += 1
                            continue
                        dMexp_update[idx_in] += (edge * vec_val)

                #print("NZERO: ", N_zero, N_zero / (.001+len(dMexp_s1)))
                #print("NSKIP: ", N_skip, N_skip / (.001+len(dMexp_s1)))

                for idx_in, vec_val in dMexp_update.items():
                    if np.any(vec_val != 0):
                        dMexp_s1[idx_pk][idx_in] += vec_val

                #applied the initial vector again since the dMexp_s1 didn't start with diagonal ones
                for idx_in, vec_val in Mexp_vec_d.items():
                    if np.any(vec_val != 0):
                        dMexp_s1[idx_pk][idx_in] += vec_val
            #print("FULLSKIP: ", fullskip, fullskip / len(dLt_idx_list))

        self.vals_prev = dict()
        #inject subtracted values at input to remove the DC values an only get the derivative
        for idx_in in range(len(pks)):
            pkin = self.in_map[pks[idx_in]]

            val = pk_original[idx_in]
            if np.any(val != 0):
                self.vals_prev[pkin] = val

        #print("START EDGES")
        #for idx_out in range(len(pks)):
        #    for idx_in in range(len(pks)):
        #        edge = dMexp_s1.get(idx_out, dict()).get(idx_in, None)
        #        if edge is not None:
        #            pkin = self.in_map[pks[idx_in]]
        #            pkout = self.out_map[pks[idx_out]]
        #            if idx_in == idx_out:
        #                edge = edge + 1

        #            print("IN: ", pkin[1])
        #            print("  OUT: ", pkout[1])
        #            print("EDGE: ", edge)
        #print("DONE EDGES")
        #dval_out holds the product of the input through the derivative matrix
        #this way it can cancel the forward propagation so that the output is correct assuming
        #the inputs do not change
        solution = dict()
        dval_out = collections.defaultdict(lambda : 0)

        #alter the derivative matrix to pass through the direct values
        for idx in range(len(pks)):
            in_map = dMexp_s1.get(idx, None)
            if in_map is None:
                dMexp_s1[idx] = {idx : 1}
                continue
            in_map[idx] = in_map.get(idx, 0) + 1

        for idx_out, in_map in dMexp_s1.items():
            for idx_in, edge in in_map.items():
                pkin = self.in_map[pks[idx_in]]
                pkout = self.out_map[pks[idx_out]]

                if pkin is not None and pkout is not None:
                    solution[pkin, pkout] = edge

                #also compute dval_out
                val_orig = pk_original[idx_in]
                prod = (edge * val_orig)
                if np.any(prod != 0):
                    dval_out[idx_out] += prod
        self.solution = solution

        vals_inj = dict()
        for idx_out in range(len(pks)):
            pkout = self.out_map[pks[idx_out]]
            if pkout is None:
                continue
            val = pkv[idx_out]
            altered_val = val - dval_out.get(idx_out, 0)
            #print("ALT: ", pkout, val, dval_out.get(idx, 0), altered_val)
            vals_inj[pkout] = altered_val
        self.vals_inj = vals_inj

        #print("DONE SOLVING ")
        return

#Old version
"""
class ExpMatCoupling(FactorCouplingBase):

    #must redefine since it was a property
    edges_req_pkset_dict = None
    def __init__(
        self,
        ddlt,
        in_map,
        out_map,
        N_ode = 1,
        order = 2,
        symplectify = True,
    ):
        self.N_ode   = N_ode
        self.order   = order
        self.ddlt    = ddlt
        self.out_map = out_map
        self.in_map  = in_map

        #all edges are generated immediately. Currently assumes full density
        self.edges_NZ_pkset_dict = {}
        self.edges_pkpk_dict = {}
        self.edges_req_pkset_dict = {}
        def gen_edge_func(pk_in, pk_out):
            return lambda sV, sB: self.edge_func(pk_in, pk_out, sV, sB)
        for pk_out in self.out_map.values():
            if pk_out is None:
                continue
            for pk_in in self.in_map.values():
                self.edges_NZ_pkset_dict[(pk_in, pk_out)] = frozenset()
                self.edges_pkpk_dict[(pk_in, pk_out)] = gen_edge_func(pk_in, pk_out)
                self.edges_req_pkset_dict[(pk_in, pk_out)] = frozenset(self.in_map.values())

        #Currently, nonlinear doesn't need to make any sources. It may in the future as that may be a more stable way to converge
        self.sources_pk_dict = {}
        self.sources_NZ_pkset_dict = {}

        pks = set()
        def pks_grab(lt):
            if isinstance(lt, list):
                for sublt in lt:
                    pks_grab(sublt)
            elif isinstance(lt, tuple):
                for pk in lt[1:]:
                    pks.add(pk)
            else:
                raise RuntimeError("BOO")
            return
        for pk_out, din in self.ddlt.items():
            pks.add(pk_out)
            for pk_in, lt in din.items():
                pks.add(pk_in)
                pks_grab(lt)

        #print(ins_p, sol_vector)
        self.pks = list(pks)
        print("OMG: ", len(pks))
        pprint(pks)
        self.pks.sort()
        self.pks_inv = dict()
        for idx, pk in enumerate(self.pks):
            self.pks_inv[pk] = idx

        h = 1 / self.N_ode
        #remap the index keys into integer indexes for speed
        ddlt_accel = dict()
        def ddlt_remap(lt):
            if isinstance(lt, list):
                sublist = []
                for sublt in lt:
                    sublist.append(ddlt_remap(sublt))
                return sublist
            elif isinstance(lt, tuple):
                #get the gain
                newtup = [lt[0] * h]
                for pk in lt[1:]:
                    newtup.append(self.pks_inv[pk])
                return tuple(newtup)
            else:
                raise RuntimeError("BOO")
            return
        for pk_out, din in self.ddlt.items():
            for pk_in, lt in din.items():
                ddlt_accel[self.pks_inv[pk_out], self.pks_inv[pk_in]] = ddlt_remap(lt)
        self.ddlt_accel = ddlt_accel

        def ddlt_mult(lt, lt2):
            def ddlt_mult2(lt, lt2):
                if isinstance(lt2, list):
                    sublist = []
                    for sublt2 in lt2:
                        sublist = sublist + ddlt_mult2(lt, sublt2)
                    return sublist
                elif isinstance(lt2, tuple):
                    #multiply the gains and merge the indices
                    return [(lt2[0] * lt[0],) + lt[1:] + lt2[1:]]
                else:
                    raise RuntimeError("BOO")
            if isinstance(lt, list):
                sublist = []
                for sublt in lt:
                    sublist = sublist + ddlt_mult(sublt, lt2)
                return sublist
            elif isinstance(lt, tuple):
                return ddlt_mult2(lt, lt2)
            else:
                raise RuntimeError("BOO")
            return
        def sort_collect(new_list):
            if not new_list:
                return new_list
            sNL = [(tuple(sorted(L[1:])), idx) for idx, L in enumerate(new_list)]
            sNL.sort()
            list_gen = []

            NL, idx = sNL[0]
            prev_NL = NL
            prev_gain = new_list[idx][0]
            for NL, idx in sNL[1:]:
                if prev_NL == NL:
                    prev_gain += new_list[idx][0]
                else:
                    list_gen.append(
                        (prev_gain,) + tuple(prev_NL)
                    )
                    prev_NL = NL
                    prev_gain = new_list[idx][0]
            list_gen.append(
                (prev_gain,) + tuple(prev_NL)
            )
            return list_gen

        self.ddlt_accel = ddlt_accel

        if symplectify:
            #print("IS SYMPLECTIC!")
            ddlt_accel_SE = dict()
            #put in the diagonals first
            for idx, pk in enumerate(self.pks):
                ddlt_accel_SE[idx, idx] = [(1,)]

            for pk_out, din in self.ddlt.items():
                pk_out_idx = self.pks_inv[pk_out]
                for pk_in, lt in din.items():
                    pk_in_idx = self.pks_inv[pk_in]
                    assert(pk_out_idx != pk_in_idx)
                    for col_idx, pk in enumerate(self.pks):
                        #needs to multiply everything
                        lt_rm = ddlt_remap(lt)
                        lt_keep = ddlt_accel_SE.get((pk_out_idx, col_idx), [])
                        lt_mult = ddlt_accel_SE.get((pk_in_idx, col_idx), [])
                        new_list = lt_keep + ddlt_mult(lt_mult, lt_rm)
                        #clear out now couplings for speed
                        new_list = [L for L in new_list if len(L[1:]) <= 3]
                        if not new_list or new_list == [(1,)]:
                            ddlt_accel_SE[pk_out_idx, col_idx] = new_list
                            continue

                        new_list = sort_collect(new_list)
                        #if new_list:
                            #print(pk_out_idx, col_idx)
                            #print([("{0:.2f}".format(np.log10(abs(L[0]) / h)), len(L[1:])) for L in new_list])
                        ddlt_accel_SE[pk_out_idx, col_idx] = new_list

            ddlt_accel_SE_use = dict()
            for (idx_out, idx_in), lt in ddlt_accel_SE.items():
                #TODO use a search to remove this since they have been sorted
                if idx_out == idx_in:
                    if lt and lt[0] == (1,):
                        lt = lt[1:]
                    else:
                        lt = lt + [(-1,)]
                if lt:
                    ddlt_accel_SE_use[idx_out, idx_in] = lt
            #pprint(ddlt_accel_SE_use)

            self.ddlt_accel = ddlt_accel_SE_use

            #add in linear term
            new_mat = copy.deepcopy(self.ddlt_accel)
            #add in diagonal (1) term
            for idx in range(len(self.pks)):
                lst = new_mat.setdefault((idx, idx), [])
                lst.append((1,))
            #add in second-order term
            for (idx_out, idx_in), lt1 in self.ddlt_accel.items():
                for idx_in2 in range(len(self.pks)):
                    lt2 = self.ddlt_accel.get((idx_in, idx_in2), None)
                    if lt2 is None:
                        continue
                    lt_mult = []
                    for l1 in lt1:
                        lt_mult.extend([(l1[0] * l2[0] / 2, ) + l1[1:] + l2[1:] for l2 in lt2])
                    #clear out now couplings for speed
                    lt_mult = sort_collect([L for L in lt_mult if abs(L[0]) >= 1e-8])
                    #if lt_mult:
                    #    print("1: ", lt1)
                    #    print("2: ", lt2)
                    #    print("LT_MULT: ", lt_mult)
                    nl = new_mat.setdefault((idx_out, idx_in2), [])
                    nl.extend(lt_mult)

            for (idx_out, idx_in) in list(new_mat.keys()):
                lt1 = new_mat[idx_out, idx_in]
                lt = sort_collect(lt1)
                nl = [L for L in nl if abs(L[0]) >= 1e-5]
                if lt:
                    lt1[:] = lt
                else:
                    del new_mat[idx_out, idx_in]
            self.ddlt_accel_premult2 = new_mat

        #pprint("PKS:")
        #pprint(pks)

    _prev_sol_vector = None

    def edge_func(self, pk_in, pk_out, sol_vector, sB):
        if sol_vector != self._prev_sol_vector:
            self._prev_sol_vector = sol_vector
            if self.order > 0:
                #self.generate_solution(sol_vector)
                self.generate_solution_premult(sol_vector)
            elif self.order == 0:
                self.generate_solution_RK(sol_vector)
        return self.solution.get((pk_in, pk_out), 0)

    def generate_solution_premult(self, sol_vector):
        pks = self.pks
        pkv = np.empty(len(pks), dtype=object)
        for idx, pk in enumerate(pks):
            #print("PK_G: ", (ins_p, pk))
            pkv[idx] = sol_vector.get(self.in_map[pk], 0)
        pkO = pkv.copy()
        #print("PKV: ", pkv)
        #try:
        #    import tabulate
        #    tabular_data = [[str(label)] + [pk] for label, pk in zip(pks, pkv)]
        #    print("PKs:")
        #    print(tabulate.tabulate(tabular_data))
        #except ImportError:
        #    print("XXXX")

        def lt_val(lt):
            if isinstance(lt, list):
                val = 0
                for sublt in lt:
                    val = lt_val(sublt) + val
            elif isinstance(lt, tuple):
                val = lt[0]
                #print("LT0: ", val)
                for pk_idx in lt[1:]:
                    val = val * pkv[pk_idx]  # sol_vector.get(pk, 0)
            else:
                raise RuntimeError("BOO")
            return val

        eye = np.eye(len(pks), dtype = object)
        Mexp_tot = eye
        for idx_N in range(self.N_ode):
            print('pe_A premult')
            Mexp = np.zeros([len(pks), len(pks)], dtype = object)
            for (idx_out, idx_in), lt in self.ddlt_accel_premult2.items():
                    val = lt_val(lt)
                    Mexp[idx_out, idx_in] = val
            #print("Mexp: ", Mexp)
            print('pe_B')
            Mexp_tot = np.dot(Mexp, Mexp_tot)
            print("Sparsity: ", idx_N, np.sum(Mexp_tot.flatten() != 0)/ len(Mexp_tot.flatten()))
            #print(Mexp.shape, pkv.shape)
            pkv = np.dot(Mexp, pkv.reshape(-1, 1)).reshape(-1)
            #try:
            #    import tabulate
            #    tabular_data = [[str(label)] + [str(pk), str(pkk)] for label, pk, pkk in zip(pks, pkv, pkX)]
            #    print("PKs2:")
            #    print(tabulate.tabulate(tabular_data))
            #except ImportError:
            #    print("XXXX")
            #print("pkv2:", type(pkv), pkv.shape)
            #print(pkv)

        #print(m1)
        #print(Mexp)
        try:
            import tabulate
            tabular_data = [[str(label)] + [str(pk)] for label, pk in zip(pks, pkv)]
            print("PKs2:")
            print(tabulate.tabulate(tabular_data))
        except ImportError:
            print("XXXX")
        try:
            import tabulate
            tabular_data = [[str(label)] + list(abs(x) for x in td) for idx, (label, td) in enumerate(zip(pks, Mexp_tot))]
            print("Mexp_tot", idx)
            print(Mexp.dtype)
            print(tabulate.tabulate(tabular_data))
        except ImportError:
            print("XXXX")

        N_sparsity = 0
        solution = dict()
        for idx_in in range(len(pks)):
            for idx_out in range(len(pks)):
                edge = Mexp_tot[idx_out, idx_in]
                if np.any(edge != 0):
                    N_sparsity += 1
                    #pk_in = pks[idx_in]
                    #pk_out = pks[idx_out]
                    ###TODO: add debug config reference for this print
                    #print(pk_in)
                    #print(pk_out)
                    #print(idx_in, idx_out, edge)
                    pkin = self.in_map[pks[idx_in]]
                    pkout = self.out_map[pks[idx_out]]
                    if pkout is None:
                        continue
                    solution[pkin, pkout] = edge
        #print("Sparsity: ", N_sparsity / len(Mexp_tot.flatten()))

        #pprint(pks)

        self.solution = solution

    def generate_solution(self, sol_vector):
        pks = self.pks
        pkv = np.empty(len(pks), dtype=object)
        for idx, pk in enumerate(pks):
            #print("PK_G: ", (ins_p, pk))
            pkv[idx] = sol_vector.get(self.in_map[pk], 0)
        pkO = pkv.copy()
        #print("PKV: ", pkv)
        #try:
        #    import tabulate
        #    tabular_data = [[str(label)] + [pk] for label, pk in zip(pks, pkv)]
        #    print("PKs:")
        #    print(tabulate.tabulate(tabular_data))
        #except ImportError:
        #    print("XXXX")

        def lt_val(lt):
            if isinstance(lt, list):
                val = 0
                for sublt in lt:
                    val = lt_val(sublt) + val
            elif isinstance(lt, tuple):
                val = lt[0]
                #print("LT0: ", val)
                for pk_idx in lt[1:]:
                    val = val * pkv[pk_idx]  # sol_vector.get(pk, 0)
            else:
                raise RuntimeError("BOO")
            return val

        eye = np.eye(len(pks), dtype = object)
        Mexp_tot = eye
        for idx_N in range(self.N_ode):
            print('pe_A')
            m1 = np.zeros([len(pks), len(pks)], dtype = object)
            for (idx_out, idx_in), lt in self.ddlt_accel.items():
                    val = lt_val(lt)
                    m1[idx_out, idx_in] = val
            #print("M1: ", m1)
            print('pe_B')
            Mexp = m1 + eye
            mmem = m1
            #try:
            #    import tabulate
            #    tabular_data = [[str(idx)] + list(td) for idx, (label, td) in enumerate(zip(pks, m1))]
            #    print("M1")
            #    print(m1.dtype)
            #    print(tabulate.tabulate(tabular_data))
            #except ImportError:
            #    print("XXXX")
            #try:
            #    import tabulate
            #    tabular_data = [[str(idx)] + list(str(x) for x in td) for idx, (label, td) in enumerate(zip(pks, Mexp))]
            #    print("Mexp", 1)
            #    print(Mexp.dtype)
            #    print(tabulate.tabulate(tabular_data))
            #except ImportError:
            #    print("XXXX")
            for idx in range(2, self.order+1):
                mmem = (1 / idx) * np.dot(m1, mmem)
                #try:
                #    import tabulate
                #    tabular_data = [[str(idx)] + list(td) for idx, (label, td) in enumerate(zip(pks, mmem))]
                #    print("mmem", idx)
                #    print(mmem.dtype)
                #    print(tabulate.tabulate(tabular_data))
                #except ImportError:
                #    print("XXXX")
                Mexp = Mexp + mmem
            print('pe_C')
            #try:
            #    import tabulate
            #    tabular_data = [[str(idx)] + list(str(x) for x in td) for idx, (label, td) in enumerate(zip(pks, Mexp))]
            #    print("Mexp", idx)
            #    print(Mexp.dtype)
            #    print(tabulate.tabulate(tabular_data))
            #except ImportError:
            #    print("XXXX")
            #import scipy.linalg
            #Mexpe_2 = scipy.linalg.expm(m1.astype(complex))
            #try:
            #    import tabulate
            #    tabular_data = [[str(idx)] + list(str(x) for x in td) for idx, (label, td) in enumerate(zip(pks, Mexpe_2))]
            #    print("Mexpe_2", idx)
            #    print(Mexpe_2.dtype)
            #    print(tabulate.tabulate(tabular_data))
            #except ImportError:
            #    print("XXXX")
            ### IMPROVE POWER CONSERVATION
            #for idx in range(len(pks)):
            #    NORMsq = np.dot(Mexp[idx], Mexp[idx].conjugate())
            #    #print("pwr ", idx, " VAL: ", NORMsq)
            #    Mexp[idx] = Mexp[idx] / (NORMsq.real)**.5
            Mexp_tot = np.dot(Mexp, Mexp_tot)
            print("Sparsity: ", idx_N, np.sum(Mexp_tot.flatten() != 0)/ len(Mexp_tot.flatten()))
            #print(Mexp.shape, pkv.shape)
            pkv = np.dot(Mexp, pkv.reshape(-1, 1)).reshape(-1)
            #try:
            #    import tabulate
            #    tabular_data = [[str(label)] + [str(pk), str(pkk)] for label, pk, pkk in zip(pks, pkv, pkX)]
            #    print("PKs2:")
            #    print(tabulate.tabulate(tabular_data))
            #except ImportError:
            #    print("XXXX")
            #print("pkv2:", type(pkv), pkv.shape)
            #print(pkv)

        #print(m1)
        #print(Mexp)
        #try:
        #    import tabulate
        #    tabular_data = [[str(label)] + [pk] for label, pk in zip(pks, pkv)]
        #    print("PKs2:")
        #    print(tabulate.tabulate(tabular_data))
        #except ImportError:
        #    print("XXXX")
        #try:
        #    import tabulate
        #    tabular_data = [[str(idx)] + list(str(x) for x in td) for idx, (label, td) in enumerate(zip(pks, Mexp_tot))]
        #    print("Mexp_tot", idx)
        #    print(Mexp.dtype)
        #    print(tabulate.tabulate(tabular_data))
        #except ImportError:
        #    print("XXXX")

        N_sparsity = 0
        solution = dict()
        for idx_in in range(len(pks)):
            for idx_out in range(len(pks)):
                edge = Mexp_tot[idx_out, idx_in]
                if np.any(edge != 0):
                    N_sparsity += 1
                    #pk_in = pks[idx_in]
                    #pk_out = pks[idx_out]
                    ###TODO: add debug config reference for this print
                    #print(pk_in)
                    #print(pk_out)
                    #print(idx_in, idx_out, edge)
                    pkin = self.in_map[pks[idx_in]]
                    pkout = self.out_map[pks[idx_out]]
                    if pkout is None:
                        continue
                    solution[pkin, pkout] = edge
        #print("Sparsity: ", N_sparsity / len(Mexp_tot.flatten()))

        #pprint(pks)

        self.solution = solution

    def generate_solution_RK(self, sol_vector):
        pks = self.pks
        pkv = np.empty(len(pks), dtype=object)
        for idx, pk in enumerate(pks):
            #print("PK_G: ", (ins_p, pk))
            pkv[idx] = sol_vector.get(self.in_map[pk], 0)

        def lt_val(lt, pkv):
            if isinstance(lt, list):
                val = 0
                for sublt in lt:
                    val = lt_val(sublt, pkv) + val
            elif isinstance(lt, tuple):
                val = lt[0]
                #print("LT0: ", val)
                for pk_idx in lt[1:]:
                    val = val * pkv[pk_idx]  # sol_vector.get(pk, 0)
            else:
                raise RuntimeError("BOO")
            return val

        eye = np.eye(len(pks), dtype = object)
        Mexp_tot = eye
        for idx_N in range(self.N_ode):
            mk1 = np.zeros([len(pks), len(pks)], dtype = object)
            #the current ddlt_accel already incorporates h, so we must reverse that for the Runge Kutta Solver
            h = 1 / self.N_ode
            for (idx_out, idx_in), lt in self.ddlt_accel.items():
                    val = lt_val(lt, pkv) / h
                    mk1[idx_out, idx_in] = val

            pkv_k1 = np.dot(mk1, pkv.reshape(-1, 1)).reshape(-1)
            mk2 = np.zeros([len(pks), len(pks)], dtype = object)
            for (idx_out, idx_in), lt in self.ddlt_accel.items():
                    val = lt_val(lt, pkv + h/2 * pkv_k1) / h
                    mk2[idx_out, idx_in] = val

            pkv_k2 = np.dot(mk2, pkv.reshape(-1, 1)).reshape(-1)
            mk3 = np.zeros([len(pks), len(pks)], dtype = object)
            for (idx_out, idx_in), lt in self.ddlt_accel.items():
                    val = lt_val(lt, pkv + h/2 * pkv_k2) / h
                    mk3[idx_out, idx_in] = val

            pkv_k3 = np.dot(mk3, pkv.reshape(-1, 1)).reshape(-1)
            mk4 = np.zeros([len(pks), len(pks)], dtype = object)
            for (idx_out, idx_in), lt in self.ddlt_accel.items():
                    val = lt_val(lt, pkv + h * pkv_k3) / h
                    mk4[idx_out, idx_in] = val

            #try:
            #    import tabulate
            #    tabular_data = [[str(idx)] + list(str(t) for t in td) for idx, (label, td) in enumerate(zip(pks, mk1))]
            #    print("MK1")
            #    print(tabulate.tabulate(tabular_data))
            #    tabular_data = [[str(idx)] + list(str(t) for t in td) for idx, (label, td) in enumerate(zip(pks, mk2))]
            #    print("MK2")
            #    print(tabulate.tabulate(tabular_data))
            #    tabular_data = [[str(idx)] + list(str(t) for t in td) for idx, (label, td) in enumerate(zip(pks, mk3))]
            #    print("MK3")
            #    print(tabulate.tabulate(tabular_data))
            #    tabular_data = [[str(idx)] + list(str(t) for t in td) for idx, (label, td) in enumerate(zip(pks, mk4))]
            #    print("MK4")
            #    print(tabulate.tabulate(tabular_data))
            #except ImportError:
            #    print("XXXX")

            Mexp = eye + h/6 * (mk1 + 2 * mk2 + 2 * mk3 + mk4)
            Mexp_tot = np.dot(Mexp, Mexp_tot)
            pkv = np.dot(Mexp, pkv.reshape(-1, 1)).reshape(-1)

        solution = dict()
        for idx_in in range(len(pks)):
            for idx_out in range(len(pks)):
                edge = Mexp_tot[idx_out, idx_in]
                if np.any(edge != 0):
                    #pk_in = pks[idx_in]
                    #pk_out = pks[idx_out]
                    ###TODO: add debug config reference for this print
                    #print(pk_in)
                    #print(pk_out)
                    #print(idx_in, idx_out, edge)
                    pkin = self.in_map[pks[idx_in]]
                    pkout = self.out_map[pks[idx_out]]
                    if pkout is None:
                        continue
                    solution[pkin, pkout] = edge

        #pprint(pks)

        self.solution = solution

"""


#!/usr/bin/env python3

###########################################################################
#  This file is part of rsp2, and is licensed under EITHER the MIT license
#  or the Apache 2.0 license, at your option.
#
#      http://www.apache.org/licenses/LICENSE-2.0
#      http://opensource.org/licenses/MIT
#
#  Be aware that not all of rsp2 is provided under this permissive license,
#  and that the project as a whole is licensed under the GPL 3.0.
###########################################################################

import numpy as np
import json
import sys
import scipy.sparse
import typing as tp

from . import eigsh_custom, get_OPinv
from rsp2.internals import info
from rsp2.io import dynmat, eigensols

# The default tolerance for eigsh is machine precision, which I feel is
# overkill. Hopefully a lighter tolerance will save some time.
#
# TODO maybe: experiment with this
TOL = 1e-10

# absolute cosines greater than this are deemed non-orthogonal
OVERLAP_THRESH = 1e-6

DEFAULT_MAX_SOLUTIONS = 12
DEFAULT_SHIFT_INVERT_ATTEMPTS = 4
DEFAULT_NCV = 0

def main_from_rust():
    """
    Entry point when called from rsp2's rust code.

    Communicates through JSON over the standard IO streams.
    """
    info('trace: sending dynmat from rust to python')
    d = json.load(sys.stdin)
    m = dynmat.from_dict(d.pop('matrix'))
    shift_invert_attempts = d.pop('shift-invert-attempts')
    dense = d.pop('dense')
    max_solutions = d.pop('max-solutions')
    assert not d

    out = run(m,
              dense=dense,
              shift_invert_attempts=shift_invert_attempts,
              plain_ncv=DEFAULT_NCV, # FIXME add to input json from rust
              shift_invert_ncv=DEFAULT_NCV, # FIXME add to input json from rust
              use_fallback=True, # FIXME add to input json from rust
              max_solutions=max_solutions,
              search_solutions=None, # FIXME add to input json from rust
              )

    info('trace: sending eigensolutions from python to rust')
    json.dump(eigensols.to_cereal(out), sys.stdout)
    print(file=sys.stdout) # newline

def main_from_cli():
    """
    Entry point for the standalone CLI wrapper.
    """
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument('DYNMAT', help='dynmat file (.npz, .json, .json.gz, ...)')
    p.add_argument('--output', '-o', type=str, required=True)
    p.add_argument(
        '--dense', action='store_true',
        help="Use a dense eigenvalue solver. Almost all other options will be "
             "ignored in this case."
    )
    p.add_argument('--shift-invert-attempts', type=int, default=DEFAULT_SHIFT_INVERT_ATTEMPTS)
    p.add_argument(
        '--no-fallback', dest='use_fallback', action='store_false',
        help="Disable non-shift-invert-based fallback method."
    )
    p.add_argument(
        '--max-solutions', type=int, default=None,
        help="max number of solutions to seek. Default is 12 unless --dense is given."
    )
    p.add_argument(
        '--shift-invert-ncv', type=int, default=DEFAULT_NCV,
        help="suggested number of Lanczos vectors. This will automatically be "
             "clipped into the range of [min(2*max_solutions + 1, n), n]"
    )
    p.add_argument(
        '--plain-ncv', type=int, default=DEFAULT_NCV,
        help="suggested number of Lanczos vectors. This will automatically be "
             "clipped into the range of [min(2*max_solutions + 1, n), n]"
    )
    p.add_argument(
        '--search-solutions', type=int, default=None,
        help="actually ask the sparse solver for this many solutions instead "
             "of --max-solutions.  The sparse solver can converge much, much "
             "faster when a few hundred solutions are requested rather than "
             "just 12."
    )
    args = p.parse_args()

    if (not args.dense
            and args.search_solutions is not None
            and args.max_solutions is not None
            and args.search_solutions < args.max_solutions
    ):
        p.error("--max-solutions must not exceed --search-solutions")

    out = run(
        m=dynmat.from_path(args.DYNMAT),
        dense=args.dense,
        shift_invert_attempts=args.shift_invert_attempts,
        shift_invert_ncv=args.shift_invert_ncv,
        plain_ncv=args.plain_ncv,
        max_solutions=args.max_solutions,
        use_fallback=args.use_fallback,
        search_solutions=args.search_solutions,
    )
    eigensols.to_path(args.output, out)

def run(m: scipy.sparse.bsr_matrix,
        dense: bool,
        shift_invert_attempts: int,
        shift_invert_ncv: int,
        plain_ncv: int,
        max_solutions: tp.Optional[int],
        use_fallback: bool,
        search_solutions: tp.Optional[int],
        ):
    """
    A suitable entry point from pure python code.
    """
    if max_solutions is None:
        if dense:
            max_solutions = m.shape[0]
        else:
            max_solutions = DEFAULT_MAX_SOLUTIONS

    if search_solutions is None:
        search_solutions = max_solutions

    # Logic for deciding when to use shift invert results
    def inner():
        if dense:
            return try_dense(m, max_solutions)

        if shift_invert_attempts:
            esols = try_shift_invert(m,
                                     max_solutions=search_solutions,
                                     shift_invert_attempts=shift_invert_attempts,
                                     ncv=shift_invert_ncv,
                                     )
            if not all(acousticness(v) > 1. - 1e-3 for v in esols[1]):
                return esols

        if use_fallback:
            return try_regular(m,
                               max_solutions=search_solutions,
                               ncv=plain_ncv,
                               )
        else:
            raise RuntimeError('Failed to diagonalize matrix!')

    # Cutting out other solutions
    evals, evecs = inner()
    if not len(evals):
        raise RuntimeError('No solutions found!')
    return evals[:max_solutions], evecs[:max_solutions]

# As an optimization, begin by using shift-invert mode, which can converge
# in **significantly** fewer iterations than regular mode.
# noinspection PyUnreachableCode
def try_shift_invert(m, *, shift_invert_attempts, max_solutions, ncv):
    info('trace: precomputing OPinv for shift-invert')

    # From what I have seen, shift_invert mode tends to find most of its
    # solutions fairly quickly, but some may be incorrect. Furthermore, it does
    # not always find all valid negative solutions.
    #
    # I fear that when incorrect solutions appear, it may compromise those found
    # afterwards. So we make multiple calls with a small number of iterations.
    MAX_ITER = max(30, int(10 * m.shape[0] ** (1/3)))

    # A heavy computational step at the beginning of shift-invert mode is
    # factorizing the matrix; do that ahead of time.
    OPinv = get_OPinv(m, sigma=0, tol=TOL)

    found_evals = []
    found_evecs = []

    # debug info
    counts = []

    for call_i in range(shift_invert_attempts):
        info('trace: shift-invert call', call_i + 1)
        (evals, evecs) = eigsh_custom(
            m,
            k=max_solutions,
            maxiter=MAX_ITER,
            sigma=0,
            which='SA',
            tol=TOL,
            OPinv=OPinv,
            ncv=ncv,
            allow_fewer_solutions=True,
            auto_adjust_k=True,
            auto_adjust_ncv=True,
        )
        evecs = np.array(list(map(normalize, evecs)))

        # a tree of counts based on direct field assignment so that static
        # linters can catch typos. (CLion handles it very impressively!)
        class Count:
            def total(self): return int(self)
            def __int__(self): return sum(map(int, self.__dict__.values()))

        count = Count() # total solutions found
        count.good = 0 # total solutions kept
        count.bad = Count() # total solutions rejected
        count.bad.repeat = 0 # linearly dependent with prior solutions
        count.bad.wrong = 0  # non-eigenvector solutions
        count.bad.ortho_bad = 0 # tried to orthogonalize, got a non-eigenvector
        count.bad.ortho_fail = 0 # tried to orthogonalize, and failed

        for (eval, ev) in zip(evals, evecs):
            # Is it ACTUALLY an eigenvector?
            if not is_good_esol(m, eval, ev):
                count.bad.wrong += 1
                continue

            # Linearly dependent with existing solutions?
            if sum(np.abs(np.vdot(ev, other))**2 for other in found_evecs) > 0.95:
                count.bad.repeat += 1
                continue

            # Prepare it for possible insertion.
            ortho_ev = mgs_step(ev, found_evecs)

            # We didn't ruin it, did we?
            if not is_good_esol(m, eval, ortho_ev):
                count.bad.ortho_bad += 1
                continue

            if sum(np.abs(np.vdot(ortho_ev, other))**2 for other in found_evecs) > 1e-6:
                count.bad.ortho_fail += 1
                continue

            # ship it
            count.good += 1
            found_evecs.append(ortho_ev)
            found_evals.append(eval)

        counts.append(count)

    info(" Good -- Bad (Old Wrong OrthoFail OrthoBad)")
    for count in counts:
        info(
            " {:^4} -- {:^3} ({:^3} {:^5} {:^9} {:^8})".format(
                count.good,
                count.bad.total(),
                count.bad.repeat,
                count.bad.wrong,
                count.bad.ortho_fail,
                count.bad.ortho_bad,
            )
        )

    perm = np.argsort(found_evals)
    evals = np.array(found_evals)[perm]
    evecs = np.array(found_evecs)[perm]
    for val, v in zip(evals, evecs):
        if not is_good_esol(m, val, v):
            np.save('bad-mat.npy', m)
            np.save('bad-vec.npy', v)
            assert False, "bad evec"
    for i in range(len(evecs)):
        for j in range(i):
            if is_overlapping(evecs[i], evecs[j]):
                np.save('bad-a.npy', evecs[i])
                np.save('bad-b.npy', evecs[j])
                assert False, "overlap"
    return evals, evecs

def mgs_step(a, b_hats):
    """
    This is the function such that

    >>> def mgs(original_vecs):
    ...     out = []
    ...     for vec in original_vecs:
    ...         out.append(mgs_step(vec, out))
    ...     return out

    is a correct implementation of Modified Gram Schmidt method.

    Many sources present the Modified Gram Schmidt (MGS) method in a misleading
    light by presenting it as having a different order of iteration from the
    famously unstable Classical Gram-Schmidt (CGS) method, and suggesting that
    the change in iteration order is the cause for MGS' improved numerical
    stability.  More specifically, MGS is typically presented as taking each
    vector in turn and using it to modify the vectors that come **after** it,
    in contrast to CGS which modifies each vector based on those **prior** to
    it. This could lead a naive reader to believe that the function `mgs_step`
    cannot exist!

    This is a complete ruse, however. The change in iteration order does not
    change the dependency tree of floating point operations. (it merely
    reintroduces opportunities for parallelism that would otherwise be lost
    compared to CGS)

    The ACTUAL difference between CGS and MGS is best understood by contrasting
    their step functions. Here's what `cgs_step` would look like. Notice how
    all dot products involve the original vector:

    >>> def cgs_step(a, b_hats):
    ...     original = a.copy()
    ...     for b_hat in b_hats:
    ...         a = a - par(original, b_hat)
    ...     return normalize(a)

    Contrast that with the following:
    """

    # Yes, for all of that text, *it really is this simple.*
    for b_hat in b_hats:
        a = a - par(a, b_hat)
    return normalize(a)

# The part of `a` that points along `b_hat`.
def par(a, b_hat):
    return np.vdot(b_hat, a) * b_hat

def acousticness(v_hat):
    sum = np.reshape(v_hat, (-1, 3)).sum(axis=0)
    return abs(np.vdot(sum, sum))

def normalize(v):
    return v / np.sqrt(np.vdot(v, v))

def is_overlapping(a_hat, b_hat):
    return abs(np.vdot(a_hat, b_hat)) > OVERLAP_THRESH

def is_good_esol(m, eval, evec):
    assert abs(abs(np.vdot(evec, evec)) - 1) < 1e-12
    return lazy_any([
        lambda: acousticness(evec) > 1. - 1e-3,
        lambda: lazy_all([
            lambda: abs(abs(np.vdot(normalize(m @ evec), evec)) - 1.0) < 1e-2,
            lambda: (np.abs(m @ evec - eval * evec) < TOL * 10).all(),
        ])
    ])

def lazy_any(it): return any(pred() for pred in it)
def lazy_all(it): return all(pred() for pred in it)

# If shift-invert hasn't produced anything satisfactory, try regular mode.
# From what I've seen, this always produces legitimate solutions, but generally
# takes long to converge onto anything.
def try_regular(m, *, max_solutions, ncv):
    info('trace: trying non-shift-invert')

    return eigsh_custom(
        m,
        k=max_solutions,
        which='SA',
        tol=TOL,
        ncv=ncv,
        allow_fewer_solutions=True,
        auto_adjust_k=True,
        auto_adjust_ncv=True,
    )

def try_dense(m, max_solutions: int):
    info('trace: using dense eigensolver')

    # note: order for returned eigenvalues is the same as 'SA'
    # note: raises LinAlgError if the eigenvalue computation does not converge
    evals, evecs = np.linalg.eigh(m.todense())
    return evals[:max_solutions], evecs.T[:max_solutions]

if __name__ == '__main__':
    main_from_rust()