# -*- coding: utf-8 -*-
# pylint: disable=wrong-import-position, too-many-locals, consider-using-enumerate

"""
Function to generate the funciton for finding expected number of photons to
survive from a 5D CLSim table.
"""

from __future__ import absolute_import, division, print_function

__all__ = [
    'MACHINE_EPS',
    'MAX_RAD_SQ',
    'USE_JITTER',
    'generate_pexp_function',
]

__author__ = 'P. Eller, J.L. Lanfranchi'
__license__ = '''Copyright 2017 Philipp Eller and Justin L. Lanfranchi

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.'''

from collections import OrderedDict
import math
from os.path import abspath, dirname
import sys

import numpy as np
from scipy import stats

if __name__ == '__main__' and __package__ is None:
    RETRO_DIR = dirname(dirname(dirname(abspath(__file__))))
    if RETRO_DIR not in sys.path:
        sys.path.append(RETRO_DIR)
from retro import DFLT_NUMBA_JIT_KWARGS, numba_jit
from retro.const import SPEED_OF_LIGHT_M_PER_NS, SRC_OMNI, SRC_CKV_BETA1
from retro.utils.geom import generate_digitizer

MACHINE_EPS = 1e-10

# TODO: not currently using MAX_RAD_SQ...
MAX_RAD_SQ = 500**2
"""Maximum radius to consider, squared (units of m^2)"""

# TODO: a "proper" jitter (and transit time spread) implementation should treat each DOM
# independently and pick the time offset for each DOM that maximizes LLH (_not_ expected
# photon detections)
USE_JITTER = True
"""Whether to use a crude jitter implementation"""

def generate_pexp_function(
    dom_tables,
    tdi_tables=None,
    tdi_metas=None,
):
    """Generate a numba-compiled function for computing expected photon counts
    at a DOM, where the table's binning info is used to pre-compute various
    constants for the compiled function to use.

    Parameters
    ----------
    dom_tables : Retro5DTables
        Fully-loaded set of single-DOM tables (time-dependent and, if no `tdi_tables`,
        time-independent)

    tdi_tables : sequence of 1 or 2 arrays, optional
        Time- and DOM-independent tables.

    tdi_metas : sequence of 1 or 2 mappings, optional
        If provided, sequence must contain two mappings where the first
        corresponds to the finely-binned TDI table and the second corresponds
        to the coarsely-binned table (the first table takes precedence over the
        second for looking up sources). Each of the mappings must contain keys
        "bin_edges", itself a mapping containing "x", "y", "z", "costhetadir",
        and "phidir"; values of these are arrays of the bin edges in each of
        these dimensions. "costhetadir" must span [-1, 1] (inclusive) and
        "phidir" must span [-pi, pi] inclusive). All edges must be strictly
        monotonic and increasing.

    Returns
    -------
    pexp : callable
        Function to find detected-photon expectations given a hypothesis; "raw" function
        that requires passing tables & norms

    pexp_wrapper : callable
        Function to find detected-photon expectations given a hypothesis; "wrapped"
        function that bakes in tables & norms, exposing a more simple interface

    pexp_meta : OrderedDict
        Parameters, including the binning, that uniquely identify what the
        capabilities of the returned `pexp`. (Use this to eliminate
        redundant pexp functions.)

    """
    if tdi_tables is None:
        tdi_tables = ()
    if tdi_metas is None:
        tdi_metas = ()

    tbl_is_ckv = dom_tables.table_kind in ['ckv_uncompr', 'ckv_templ_compr']
    tbl_is_templ_compr = dom_tables.table_kind in ['raw_templ_compr', 'ckv_templ_compr']
    if not tbl_is_ckv:
        raise NotImplementedError('Only Ckv tables are implemented.')

    # TODO: sanity checks that all TDI metadata is compatible with DOM tables
    for tdi_meta in tdi_metas:
        assert tdi_meta['bin_edges']['phidir'][0] == -np.pi
        assert tdi_meta['bin_edges']['phidir'][-1] == np.pi

    pexp_meta = OrderedDict()
    pexp_meta['table_kind'] = dom_tables.table_kind
    pexp_meta['table_binning'] = OrderedDict()
    for key in (
        'r_bin_edges', 'costhetadir_bin_edges', 't_bin_edges', 'costhetadir_bin_edges',
        'deltaphidir_bin_edges'
    ):
        pexp_meta['table_binning'][key] = dom_tables.table_meta[key]

    pexp_meta['tdi'] = tdi_metas
    if len(tdi_tables) == 1:
        tdi_tables = (tdi_tables[0], tdi_tables[0])

    # NOTE: For now, we only support absolute value of deltaphidir (which
    # assumes azimuthal symmetry). In future, this could be revisited (and then
    # the abs(...) applied before binning in the pexp code will have to be
    # removed or replaced with behavior that depend on the range of the
    # deltaphidir_bin_edges).
    assert dom_tables.table_meta['deltaphidir_bin_edges'][0] == 0, 'only abs(deltaphidir) supported'
    assert dom_tables.table_meta['deltaphidir_bin_edges'][-1] == np.pi

    # -- Define things used by `pexp*` closures defined below -- #

    # Constants
    rsquared_max = np.max(dom_tables.table_meta['r_bin_edges'])**2
    t_max = np.max(dom_tables.table_meta['t_bin_edges'])
    recip_max_group_vel = dom_tables.table_meta['group_refractive_index'] / SPEED_OF_LIGHT_M_PER_NS

    # Digitization functions for each binning dimension
    digitize_r = generate_digitizer(
        dom_tables.table_meta['r_bin_edges'],
        clip=True
    )
    digitize_costheta = generate_digitizer(
        dom_tables.table_meta['costheta_bin_edges'],
        clip=True
    )
    digitize_t = generate_digitizer(
        dom_tables.table_meta['t_bin_edges'],
        clip=True
    )
    digitize_costhetadir = generate_digitizer(
        dom_tables.table_meta['costhetadir_bin_edges'],
        clip=True
    )
    digitize_deltaphidir = generate_digitizer(
        dom_tables.table_meta['deltaphidir_bin_edges'],
        clip=True
    )

    num_tdi_tables = len(tdi_metas)
    if num_tdi_tables == 0:
        # Numba needs an object that it can determine type of
        tdi_tables = 0
    else:
        x_edges = tdi_metas[0]['bin_edges']['x']
        y_edges = tdi_metas[0]['bin_edges']['y']
        z_edges = tdi_metas[0]['bin_edges']['z']
        tdi0_xmin, tdi0_xmax = x_edges[[0, -1]]
        tdi0_ymin, tdi0_ymax = y_edges[[0, -1]]
        tdi0_zmin, tdi0_zmax = z_edges[[0, -1]]
        digitize_tdi0_x = generate_digitizer(x_edges, clip=True)
        digitize_tdi0_y = generate_digitizer(y_edges, clip=True)
        digitize_tdi0_z = generate_digitizer(z_edges, clip=True)
        digitize_tdi0_costhetadir = generate_digitizer(
            tdi_metas[0]['bin_edges']['costhetadir'], clip=True
        )
        digitize_tdi0_phidir = generate_digitizer(
            tdi_metas[0]['bin_edges']['phidir'], clip=True
        )

        if num_tdi_tables == 1:
            idx = 0
        elif num_tdi_tables == 2:
            idx = 1
        else:
            raise ValueError(
                'Can only handle 0, 1, or 2 TDI tables; got {}'
                .format(num_tdi_tables)
            )

        x_edges = tdi_metas[idx]['bin_edges']['x']
        y_edges = tdi_metas[idx]['bin_edges']['y']
        z_edges = tdi_metas[idx]['bin_edges']['z']
        tdi1_xmin, tdi1_xmax = x_edges[[0, -1]]
        tdi1_ymin, tdi1_ymax = y_edges[[0, -1]]
        tdi1_zmin, tdi1_zmax = z_edges[[0, -1]]
        digitize_tdi1_x = generate_digitizer(x_edges, clip=True)
        digitize_tdi1_y = generate_digitizer(y_edges, clip=True)
        digitize_tdi1_z = generate_digitizer(z_edges, clip=True)
        digitize_tdi1_costhetadir = generate_digitizer(
            tdi_metas[idx]['bin_edges']['costhetadir'], clip=True
        )
        digitize_tdi1_phidir = generate_digitizer(
            tdi_metas[idx]['bin_edges']['phidir'], clip=True
        )

    dom_tables_ = dom_tables

    dom_tables = dom_tables_.tables
    dom_table_norms = dom_tables_.table_norms
    dom_tables_template_library = dom_tables_.template_library
    t_indep_dom_tables = dom_tables_.t_indep_tables
    t_indep_dom_table_norms = dom_tables_.t_indep_table_norms
    t_is_residual_time = dom_tables_.t_is_residual_time

    if not isinstance(dom_tables, np.ndarray):
        dom_tables = np.stack(dom_tables, axis=0)
        print('dom_tables.shape:', dom_tables.shape)
    if not isinstance(dom_table_norms, np.ndarray):
        dom_table_norms = np.stack(dom_table_norms, axis=0)
        print('dom_table_norms.shape:', dom_table_norms.shape)
    if not isinstance(t_indep_dom_tables, np.ndarray):
        t_indep_dom_tables = np.stack(t_indep_dom_tables, axis=0)
        print('t_indep_dom_tables.shape:', t_indep_dom_tables.shape)
    if not isinstance(t_indep_dom_table_norms, np.ndarray):
        t_indep_dom_table_norms = np.stack(t_indep_dom_table_norms, axis=0)
        print('t_indep_dom_table_norms.shape:', t_indep_dom_table_norms.shape)

    dom_tables.flags.writeable = False
    dom_table_norms.flags.writeable = False
    dom_tables_template_library.flags.writeable = False
    t_indep_dom_tables.flags.writeable = False
    t_indep_dom_table_norms.flags.writeable = False

    if USE_JITTER:
        # Time offsets to sample for DOM jitter
        jitter_dt = np.arange(-10, 11, 2)

        # Weight at each time offset
        jitter_weights = stats.norm.pdf(jitter_dt, 0, 5)
        jitter_weights /= np.sum(jitter_weights)
    else:
        jitter_dt = np.array([0.])
        jitter_weights = np.array([1.])
    num_jitter_time_offsets = len(jitter_dt)

    # Indexing functions for table types omni / directional lookups
    if tbl_is_templ_compr:
        @numba_jit(**DFLT_NUMBA_JIT_KWARGS)
        def table_lookup_mean(
            tables, table_idx, r_bin_idx, costheta_bin_idx, t_bin_idx
        ): # pylint: disable=missing-docstring
            templ = tables[table_idx][r_bin_idx, costheta_bin_idx, t_bin_idx]
            return templ['weight'] / dom_tables_template_library[templ['index']].size

        @numba_jit(**DFLT_NUMBA_JIT_KWARGS)
        def table_lookup(
            tables, table_idx, r_bin_idx, costheta_bin_idx, t_bin_idx,
            costhetadir_bin_idx, deltaphidir_bin_idx
        ): # pylint: disable=missing-docstring
            templ = tables[table_idx][r_bin_idx, costheta_bin_idx, t_bin_idx]
            return (
                templ['weight'] * dom_tables_template_library[
                    templ['index'],
                    costhetadir_bin_idx,
                    deltaphidir_bin_idx,
                ]
            )

    else: # table is not template-compressed
        @numba_jit(**DFLT_NUMBA_JIT_KWARGS)
        def table_lookup_mean(
            tables, table_idx, r_bin_idx, costheta_bin_idx, t_bin_idx
        ): # pylint: disable=missing-docstring
            return np.mean(
                tables[table_idx][r_bin_idx, costheta_bin_idx, t_bin_idx]
            )

        @numba_jit(**DFLT_NUMBA_JIT_KWARGS)
        def table_lookup(
            tables, table_idx, r_bin_idx, costheta_bin_idx, t_bin_idx,
            costhetadir_bin_idx, deltaphidir_bin_idx
        ): # pylint: disable=missing-docstring
            return tables[table_idx][
                r_bin_idx,
                costheta_bin_idx,
                t_bin_idx,
                costhetadir_bin_idx,
                deltaphidir_bin_idx,
            ]

    table_lookup_mean.__doc__ = (
        """Helper function for directionality-averaged table lookup"""
    )
    table_lookup.__doc__ = """Helper function for directional table lookup"""

    # -- Define `pexp*` closures (functions that use things defined above) -- #

    # Note in the following that we need to invert the costheta
    # direction of the sources to match the directions that
    # Retro simulation comes up with. Thus the angles
    # associated with this that we want to work with are
    #   cos(pi - thetadir) = -cos(thetadir),
    #   sin(pi - thetadir) = sin(thetadir),
    #   cos(phidir) = cos(-phidir),
    #   sin(phidir) = -sin(-phidir)
    #
    # This should be seen as a bug, but not sure how to address
    # it without modifying existing tables, so sticking with it
    # for now.
    #
    # We bin cos(pi - thetadir), so need to simply bin the
    # quantity `-src_dir_costheta`.
    #
    # We want to bin abs(deltaphidir), which is described now:
    # Just look at vectors in the xy-plane, since we want
    # difference of angle in this plane. Use dot product:
    #   dot(dir_vec_xy, pos_vec_xy) = |dir_vec_xy| |pos_vec_xy| cos(deltaphidir)
    # where the length of the directionality vector in the xy-plane is
    #   |dir_vec_xy| = rhodir = rdir * sin(pi - thetadir)
    # and since rdir = 1 and the inversion of the angle above
    #   |dir_vec_xy| = rhodir = sin(thetadir).
    # The length of the position vector in the xy-plane is
    #   |pos_vec_xy| = rho = sqrt(dx^2 + dy^2)
    # where dx and dy are src_x - dom_x and src_y - dom_y.
    # Solving for cos(deltaphidir):
    #   cos(deltaphidir) = dot(dir_vec_xy, pos_vec_xy) / (rhodir * rho)
    # we just need to write out the components of the dot
    # product in terms of quantites we have:
    #   dir_vec_x = rhodir * cos(phidir)
    #   dir_vec_y = rhodir * sin(phidir)
    #   pos_vec_x = dx
    #   pos_vec_y = dy
    # giving
    #   cos(deltaphidir) = -(rhodir*cos(phidir)*dx + rhodir*sin(phidir)*dy)/(rhodir*rho)
    # (where we use the negative to account for the inverted
    # costhetadir in the tables); cancel rhodir out
    #   cos(deltaphidir) = (cos(phidirpi)*dx + sin(phidirpi)*dy) / rho
    # and substitute the identities above
    #   cos(deltaphidir) = (cos(phidir)*dx + sin(phidir)*dy) / rho
    # Finally, solve for deltaphidir
    #   deltaphidir = acos((cos(phidir)*dx + sin(phidir)*dy) / rho)

    # A photon that starts immediately in the past (before the
    # DOM was hit) will show up in the Retro DOM tables in bin
    # 0; the further in the past the photon started, the
    # higher the time bin index. Therefore, subract source
    # time from hit time.

    # TODO: integrate tdi into same pexp function?

    pexp_docstr = (
        r"""For a set of generated photons `sources`, compute the expected
        photons in a particular DOM at `hit_time` and the total expected
        photons, independent of time.

        This function utilizes the relative space-time coordinates _and_
        directionality of the generated photons (via "raw" 5D CLSim tables) to
        determine how many photons are expected to arrive at the DOM.

        Retro DOM tables applied to the generated photon info `sources`,
        and the total expected photon count (time integrated) -- the
        normalization of the pdf.

        Parameters
        ----------
        sources : shape (num_sources,) array of dtype SRC_T
            A discrete sequence of points describing expected sources of
            photons that result from a hypothesized event.

        sources_start, sources_stop : int
            Starting and stopping indices for the part of the array on which to
            work. Note that the latter is exclusive, i.e., following Python
            range / slice syntax. Hence, the following section of `sources` will
            be operated upon: .. ::

                sources[sources_start:sources_stop]

        event_dom_info : shape (n_operational_doms,) array of dtype EVT_DOM_INFO_T

        event_hit_info : shape (n_hits,) array of dtype EVT_HIT_INFO_T

        hit_exp : shape (n_hits,) array of floats
            Time-dependent hit expectation at each (actual) hit time;
            initialize outside of this function, as values are incremented
            within this function. Values in `hit_exp` correspond to the values
            in `event_hit_info`.

        dom_tables : array
            DOM time-dependent photon survival probability tables. If using an
            uncompressed table, these will have shape
                (n_r, n_costheta, n_t, n_costhetadir, n_deltaphidir)
            while if you use a template-compressed table, this will have shape
                (n_templates, n_costhetadir, n_deltaphidir)

        dom_table_norms : shape (n_tables, n_r, n_t) array
            Normalization to apply to `table`, which is assumed to depend on
            both r- and t-dimensions.

        t_indep_dom_tables : array
            Time-independent photon survival probability table. If using an
            uncompressed table, this will have shape
                (n_r, n_costheta, n_costhetadir, n_deltaphidir)
            while if using a

        t_indep_dom_dom_table_norms : shape (n_tables, n_r) array
            r-dependent normalization (any t-dep normalization is assumed to
            already have been applied to generate the t_indep_table).

        tdi_tables : {type}
            {text}

        Returns
        -------
        t_indep_exp : float
            Expectation of total hits for all operational DOMs

        Out
        ---
        hit_exp
            `hit_exp` is modified by the function; see Parameters section for
            detailed explanation of parameter `hit_exp`

        """
    )

    if num_tdi_tables == 0:

        @numba_jit(**DFLT_NUMBA_JIT_KWARGS)
        def pexp(
            sources,
            sources_start,
            sources_stop,
            event_dom_info,
            event_hit_info,
            hit_exp,
            dom_tables,
            dom_table_norms,
            t_indep_dom_tables,
            t_indep_dom_table_norms,
            tdi_tables, # pylint: disable=unused-argument
        ): # pylint: disable=missing-docstring, too-many-arguments
            num_operational_doms = len(event_dom_info)
            t_indep_exp = 0.
            for source_idx in range(sources_start, sources_stop):
                src = sources[source_idx]

                for op_dom_idx in range(num_operational_doms):
                    dom_info = event_dom_info[op_dom_idx]
                    dom_tbl_idx = dom_info['table_idx']
                    dom_qe = dom_info['quantum_efficiency']
                    dom_hits_start_idx = dom_info['hits_start_idx']
                    dom_hits_stop_idx = dom_info['hits_stop_idx']

                    dx = src['x'] - dom_info['x']
                    dy = src['y'] - dom_info['y']
                    dz = src['z'] - dom_info['z']

                    rhosquared = max(MACHINE_EPS, dx**2 + dy**2)
                    rsquared = rhosquared + dz**2

                    if rsquared > rsquared_max:
                        continue

                    r = max(MACHINE_EPS, math.sqrt(rsquared))
                    r_bin_idx = digitize_r(r)

                    costheta_bin_idx = digitize_costheta(dz/r)

                    if src['kind'] == SRC_OMNI:
                        t_indep_surv_prob = np.mean(
                            t_indep_dom_tables[dom_tbl_idx][r_bin_idx, costheta_bin_idx, :, :]
                        )

                    else: # SRC_CKV_BETA1:
                        rho = math.sqrt(rhosquared)

                        if rho <= MACHINE_EPS:
                            absdeltaphidir = 0.
                        else:
                            absdeltaphidir = abs(math.acos(
                                max(-1., min(1., -(src['dir_cosphi']*dx + src['dir_sinphi']*dy) / rho))
                            ))

                        costhetadir_bin_idx = digitize_costhetadir(src['dir_costheta'])
                        deltaphidir_bin_idx = digitize_deltaphidir(absdeltaphidir)

                        t_indep_surv_prob = t_indep_dom_tables[dom_tbl_idx][
                            r_bin_idx,
                            costheta_bin_idx,
                            costhetadir_bin_idx,
                            deltaphidir_bin_idx
                        ]

                    ti_norm = t_indep_dom_table_norms[dom_tbl_idx][r_bin_idx]
                    t_indep_exp += src['photons'] * ti_norm * t_indep_surv_prob * dom_qe

                    for hit_idx in range(dom_hits_start_idx, dom_hits_stop_idx):
                        hit_info = event_hit_info[hit_idx]
                        if t_is_residual_time:
                            nominal_dt = hit_info['time'] - src['time'] - r * recip_max_group_vel
                        else:
                            nominal_dt = hit_info['time'] - src['time']

                        for jitter_idx in range(num_jitter_time_offsets):
                            dt = nominal_dt + jitter_dt[jitter_idx]

                            # Note the comparison is written such that it will evaluate
                            # to True if `dt` is NaN or less than zero.
                            if (not dt >= 0) or dt > t_max:
                                continue

                            t_bin_idx = digitize_t(dt)

                            if src['kind'] == SRC_OMNI:
                                surv_prob_at_hit_t = table_lookup_mean(
                                    tables=dom_tables,
                                    table_idx=dom_tbl_idx,
                                    r_bin_idx=r_bin_idx,
                                    costheta_bin_idx=costheta_bin_idx,
                                    t_bin_idx=t_bin_idx,
                                )

                            else: # SRC_CKV_BETA1
                                surv_prob_at_hit_t = table_lookup(
                                    tables=dom_tables,
                                    table_idx=dom_tbl_idx,
                                    r_bin_idx=r_bin_idx,
                                    costheta_bin_idx=costheta_bin_idx,
                                    t_bin_idx=t_bin_idx,
                                    costhetadir_bin_idx=costhetadir_bin_idx,
                                    deltaphidir_bin_idx=deltaphidir_bin_idx,
                                )

                            r_t_bin_norm = dom_table_norms[dom_tbl_idx][r_bin_idx, t_bin_idx]
                            hit_exp[hit_idx] += jitter_weights[jitter_idx] * (
                                src['photons'] * r_t_bin_norm * surv_prob_at_hit_t * dom_qe
                            )

            return t_indep_exp

        pexp.__doc__ = pexp_docstr.format(
            type='int',
            text="""Dummy argument for this version of `pexp` since it doesn't use TDI
            tables (but this argument needs to be present to maintain same
            interface)"""
        )

    else: # pexp function given we are using TDI tables

        @numba_jit(**DFLT_NUMBA_JIT_KWARGS)
        def pexp(
            sources,
            sources_start,
            sources_stop,
            event_dom_info,
            event_hit_info,
            hit_exp,
            dom_tables,
            dom_table_norms,
            t_indep_dom_tables, # pylint: disable=unused-argument
            t_indep_dom_table_norms, # pylint: disable=unused-argument
            tdi_tables,
        ): # pylint: disable=missing-docstring, too-many-arguments
            # -- Time- and DOM-independent photon-detection expectation -- #

            t_indep_exp = 0.
            for source_idx in range(sources_start, sources_stop):
                src = sources[source_idx]
                src_opposite_dir_costheta = -src['dir_costheta']
                src_opposite_dir_phi = ((src['dir_phi'] + 2*np.pi) % (2*np.pi)) - np.pi

                if (
                    tdi0_xmin <= src['x'] <= tdi0_xmax
                    and tdi0_ymin <= src['y'] <= tdi0_ymax
                    and tdi0_zmin <= src['z'] <= tdi0_zmax
                ):
                    t_indep_exp += 0.45 * src['photons'] * tdi_tables[0][
                        digitize_tdi0_x(src['x']),
                        digitize_tdi0_y(src['y']),
                        digitize_tdi0_z(src['z']),
                        digitize_tdi0_costhetadir(src_opposite_dir_costheta),
                        digitize_tdi0_phidir(src_opposite_dir_phi),
                    ]
                elif num_tdi_tables >= 2 and (
                    tdi1_xmin <= src['x'] <= tdi1_xmax
                    and tdi1_ymin <= src['y'] <= tdi1_ymax
                    and tdi1_zmin <= src['z'] <= tdi1_zmax
                ):
                    t_indep_exp += 0.45 * src['photons'] * tdi_tables[1][
                        digitize_tdi1_x(src['x']),
                        digitize_tdi1_y(src['y']),
                        digitize_tdi1_z(src['z']),
                        digitize_tdi1_costhetadir(src_opposite_dir_costheta),
                        digitize_tdi1_phidir(src_opposite_dir_phi),
                    ]
                else:
                    continue

            # -- Time-dependent photon-det expectation for each hit DOM -- #

            for hit_idx, hit_info in enumerate(event_hit_info):
                dom_info = event_dom_info[hit_info['event_dom_idx']]
                dom_tbl_idx = dom_info['table_idx']
                dom_qe = dom_info['quantum_efficiency']

                for source_idx in range(sources_start, sources_stop):
                    src = sources[source_idx]

                    dx = src['x'] - dom_info['x']
                    dy = src['y'] - dom_info['y']
                    dz = src['z'] - dom_info['z']

                    rhosquared = max(MACHINE_EPS, dx**2 + dy**2)
                    rsquared = rhosquared + dz**2

                    if rsquared > rsquared_max:
                        continue

                    r = max(MACHINE_EPS, math.sqrt(rsquared))
                    r_bin_idx = digitize_r(r)

                    costheta_bin_idx = digitize_costheta(dz/r)

                    if src['kind'] == SRC_CKV_BETA1:
                        rho = math.sqrt(rhosquared)

                        if rho <= MACHINE_EPS:
                            absdeltaphidir = 0.
                        else:
                            absdeltaphidir = abs(math.acos(
                                max(-1., min(1., -(src['dir_cosphi']*dx + src['dir_sinphi']*dy) / rho))
                            ))

                        costhetadir_bin_idx = digitize_costhetadir(src['dir_costheta'])
                        deltaphidir_bin_idx = digitize_deltaphidir(absdeltaphidir)

                    if t_is_residual_time:
                        nominal_dt = hit_info['time'] - src['time'] - r * recip_max_group_vel
                    else:
                        nominal_dt = hit_info['time'] - src['time']

                    # Note: caching last `t_bin_idx`, `r_t_bin_norm`, and
                    # `surv_prob_at_hit_t` and checking for identical `t_bin_idx` seems
                    # to take about the same time as not caching these values, so
                    # choosing the simpler way

                    for jitter_idx in range(num_jitter_time_offsets):
                        dt = nominal_dt + jitter_dt[jitter_idx]

                        # Note the comparison is written such that it will evaluate to
                        # True if `dt` is NaN or less than zero.
                        if (not dt >= 0) or dt > t_max:
                            continue

                        t_bin_idx = digitize_t(dt)

                        if src['kind'] == SRC_OMNI:
                            surv_prob_at_hit_t = table_lookup_mean(
                                tables=dom_tables,
                                table_idx=dom_tbl_idx,
                                r_bin_idx=r_bin_idx,
                                costheta_bin_idx=costheta_bin_idx,
                                t_bin_idx=t_bin_idx,
                            )

                        else: # SRC_CKV_BETA1
                            surv_prob_at_hit_t = table_lookup(
                                tables=dom_tables,
                                table_idx=dom_tbl_idx,
                                r_bin_idx=r_bin_idx,
                                costheta_bin_idx=costheta_bin_idx,
                                t_bin_idx=t_bin_idx,
                                costhetadir_bin_idx=costhetadir_bin_idx,
                                deltaphidir_bin_idx=deltaphidir_bin_idx,
                            )

                        r_t_bin_norm = dom_table_norms[dom_tbl_idx][r_bin_idx, t_bin_idx]
                        hit_exp[hit_idx] += jitter_weights[jitter_idx] * (
                            src['photons'] * r_t_bin_norm * surv_prob_at_hit_t * dom_qe
                        )

            return t_indep_exp

        pexp.__doc__ = pexp_docstr.format(
            type='tuple of 1 or 2 arrays',
            text="""TDI tables"""
        )

    # -- Define pexp closure to bake-in the tables -- #

    # Note: faster to _not_ jit-compile this function (why, though?)
    def pexp_wrapper(
        sources,
        sources_start,
        sources_stop,
        event_dom_info,
        event_hit_info,
        hit_exp,
    ):
        return pexp(
            sources=sources,
            sources_start=sources_start,
            sources_stop=sources_stop,
            event_dom_info=event_dom_info,
            event_hit_info=event_hit_info,
            hit_exp=hit_exp,
            dom_tables=dom_tables,
            dom_table_norms=dom_table_norms,
            t_indep_dom_tables=t_indep_dom_tables,
            t_indep_dom_table_norms=t_indep_dom_table_norms,
            tdi_tables=tdi_tables,
        )

    return pexp, pexp_wrapper, pexp_meta


from collections import deque
import os
import numpy as np
import torch
import torch.nn.functional as F
from torch.distributions import Normal

from .base import BaseAgent
from core.network import Network
from core.optimizer import Optimizer
from core.buffer import ReplayBuffer

class MPO(BaseAgent):
    """Maximum A Posteriori Policy Optimization (MPO) agent.

    Args:
        state_size (int): dimension of state.
        action_size (int): dimension of action.
        hidden_size (int): dimension of hidden unit.
        optim_config (dict): dictionary of the optimizer info.
            (key: 'name', value: name of optimizer)
        actor (str): key of actor network class in _network_dict.txt.
        critic (str): key of critic network class in _network_dict.txt.
        head (str): key of head in _head_dict.txt.
        buffer_size (int): the size of the memory buffer.
        batch_size (int): the number of samples in the one batch.
        start_train_step (int): steps to start learning.
        n_epoch (int): Number of epoch when optimizing the surrogate.
        n_step (int): The number of steps to run for each environment per update.
        clip_grad_norm (float): gradient clipping threshold.
        gamma (float): discount factor.
        device (str): device to use.
            (e.g. 'cpu' or 'gpu'. None can also be used, and in this case, the cpu is used.)
        num_workers: the number of agents in distributed learning.
        critic_loss_type (str): type of critic loss. One of ['1step_TD', 'retrace'].
        num_sample (int): the number of samples.
        min_eta (float): minimum value of eta.
        min_alpha_mu (float): minimum value of alpha_mu.
        min_alpha_sigma (float): minimum value of alpha_sigma.
        eps_eta (float): threshold of temperature loss term.
        eps_alpha_mu (float): threshold of mean part of Gaussian-KL constraint term.
        eps_alpha_sigma (float): threshold of variance part of Gaussian-KL constraint term.
        eta (float): Lagrange multipliers of temperature loss term.
        alpha_mu (float): Lagrange multipliers of mean part of Gaussian-KL constraint term (trust-region loss).
        alpha_sigma (float): Lagrange multipliers of variance part of Gaussian-KL constraint term.
    """

    def __init__(
        self,
        state_size,
        action_size,
        hidden_size=512,
        optim_config={"name": "adam"},
        actor="discrete_policy",
        critic="dqn",
        head="mlp",
        buffer_size=50000,
        batch_size=64,
        start_train_step=2000,
        n_epoch=64,
        n_step=8,
        clip_grad_norm=1.0,
        gamma=0.99,
        device=None,
        num_workers=1,
        # parameters unique to MPO
        critic_loss_type="retrace",  # one of ['1step_TD', 'retrace']
        num_sample=30,
        min_eta=1e-8,
        min_alpha_mu=1e-8,
        min_alpha_sigma=1e-8,
        eps_eta=0.01,
        eps_alpha_mu=0.01,
        eps_alpha_sigma=5 * 1e-5,
        eta=1.0,
        alpha_mu=1.0,
        alpha_sigma=1.0,
        **kwargs,
    ):
        self.device = (
            torch.device(device)
            if device
            else torch.device("cuda" if torch.cuda.is_available() else "cpu")
        )
        self.head = head
        self.action_type = actor.split("_")[0]
        assert self.action_type in ["continuous", "discrete"]
        self.action_size = action_size

        self.actor = Network(
            actor, state_size, action_size, D_hidden=hidden_size, head=head
        ).to(self.device)
        self.target_actor = Network(
            actor, state_size, action_size, D_hidden=hidden_size, head=head
        ).to(self.device)
        self.target_actor.load_state_dict(self.actor.state_dict())

        assert critic_loss_type in ["1step_TD", "retrace"]
        self.critic_loss_type = critic_loss_type
        self.critic = Network(
            critic, state_size, action_size, D_hidden=hidden_size, head=head
        ).to(self.device)
        self.target_critic = Network(
            critic, state_size, action_size, D_hidden=hidden_size, head=head
        ).to(self.device)
        self.target_critic.load_state_dict(self.critic.state_dict())

        self.batch_size = batch_size
        self.n_step = n_step if critic_loss_type == "retrace" else 1
        self.clip_grad_norm = clip_grad_norm

        self.num_learn = 0
        self.time_t = 0
        self.start_train_step = start_train_step
        self.n_epoch = n_epoch

        self.num_sample = num_sample

        self.min_eta = torch.tensor(min_eta, device=self.device)
        self.min_alpha_mu = torch.tensor(min_alpha_mu, device=self.device)
        self.min_alpha_sigma = torch.tensor(min_alpha_sigma, device=self.device)

        self.eps_eta = eps_eta
        self.eps_alpha_mu = eps_alpha_mu
        self.eps_alpha_sigma = eps_alpha_sigma

        self.eta = torch.nn.Parameter(
            torch.tensor(eta, requires_grad=True).to(self.device)
        )
        self.alpha_mu = torch.nn.Parameter(
            torch.tensor(alpha_mu, requires_grad=True).to(self.device)
        )
        self.alpha_sigma = torch.nn.Parameter(
            torch.tensor(alpha_sigma, requires_grad=True).to(self.device)
        )

        self.reset_lgr_muls()

        self.actor_optimizer = Optimizer(
            **optim_config,
            params=list(self.actor.parameters())
            + [self.eta, self.alpha_mu, self.alpha_sigma],
        )
        self.critic_optimizer = Optimizer(
            **optim_config, params=list(self.critic.parameters())
        )

        self.gamma = gamma
        self.tmp_buffer = deque(maxlen=n_step)
        self.memory = ReplayBuffer(buffer_size)

    @torch.no_grad()
    def act(self, state, training=True):
        self.actor.train(training)
        if self.action_type == "continuous":
            mu, std = self.actor(self.as_tensor(state))
            m = Normal(mu, std)
            z = m.sample() if training else mu
            action = torch.tanh(z)
            action = action.data.cpu().numpy()
            prob = m.log_prob(z).sum(axis=-1, keepdims=True)
            prob = prob.exp().cpu().numpy()

        else:
            pi = self.actor(self.as_tensor(state))
            action = (
                torch.multinomial(pi, 1)
                if training
                else torch.argmax(pi, dim=-1, keepdim=True)
            )
            action = action.cpu().numpy()
            prob = np.take(pi.cpu().numpy(), action)
        return {
            "action": action,
            "prob": prob,
        }

    def learn(self):
        transitions = self.memory.sample(self.batch_size)
        for key in transitions.keys():
            # reshape: (batch_size, len_tr, item_dim)
            #        -> (batch_size * len_tr, item_dim)
            transitions[key] = self.as_tensor(transitions[key]).view(
                -1, *transitions[key].shape[2:]
            )

        state = transitions["state"]
        action = transitions["action"]
        reward = transitions["reward"]
        next_state = transitions["next_state"]
        done = transitions["done"]
        prob_b = transitions["prob"]

        if self.action_type == "continuous":
            mu, std = self.actor(state)
            Q = self.critic(state, action)
            m = Normal(mu, std)
            z = torch.atanh(torch.clamp(action, -1 + 1e-7, 1 - 1e-7))
            log_pi = m.log_prob(z)
            log_prob = log_pi.sum(axis=-1, keepdims=True)
            prob = torch.exp(log_prob)

            with torch.no_grad():
                mut, stdt = self.target_actor(state)
                mt = Normal(mut, stdt)
                zt = torch.atanh(torch.clamp(action, -1 + 1e-7, 1 - 1e-7))
                log_pit = mt.log_prob(zt)
                log_probt = log_pit.sum(axis=-1, keepdims=True)

                mu_old = mut
                std_old = stdt
                prob_t = torch.exp(log_probt)

                Qt_a = self.target_critic(state, action)

                next_mu, next_std = self.actor(next_state)
                mn = Normal(next_mu, next_std)
                zn = mn.sample(
                    (self.num_sample,)
                )  # (num_sample, batch_size * len_tr, dim_action)
                next_action = torch.tanh(zn)

                Qt_next = self.target_critic(
                    next_state.unsqueeze(0).repeat_interleave(self.num_sample, dim=0),
                    next_action,
                )  # (num_sample, batch_size * len_tr, 1)

                c = torch.clip(prob / (prob_b + 1e-6), max=1.0)

                if self.critic_loss_type == "1step_TD":
                    Qret = reward + self.gamma * (1 - done) * Qt_next.mean(axis=0)
                elif self.critic_loss_type == "retrace":
                    Qret = reward + self.gamma * Qt_next.mean(axis=0) * (1 - done)

                    # temporarily reshaping values
                    # (batch_size * len_tr, item_dim) -> (batch_size, len_tr, item_dim)
                    Qret = Qret.view(self.batch_size, -1, *Qret.shape[1:])
                    Qt_a = Qt_a.view(self.batch_size, -1, *Qt_a.shape[1:])
                    c = c.view(self.batch_size, -1, *c.shape[1:])
                    done = done.view(self.batch_size, -1, *done.shape[1:])
                    for i in reversed(range(Qret.shape[1] - 1)):
                        Qret[:, i] += (
                            self.gamma
                            * c[:, i + 1]
                            * (1 - done[:, i])
                            * (Qret[:, i + 1] - Qt_a[:, i + 1])
                        )
                    Qret = Qret.view(-1, *Qret.shape[2:])

            zt_add = mt.sample(
                (self.num_sample,)
            )  # (num_sample, batch_size * len_tr, dim_action)
            action_add = torch.tanh(zt_add)
            log_pi_add = m.log_prob(zt_add)
            log_prob_add = log_pi_add.sum(axis=-1, keepdims=True)
            Qt_add = self.target_critic(
                state.unsqueeze(0).repeat_interleave(self.num_sample, dim=0), action_add
            )

            critic_loss = F.mse_loss(Q, Qret).mean()

            # Calculate Vt_add, At_add using Qt_add
            Vt_add = torch.mean(Qt_add, axis=0, keepdims=True)
            At_add = Qt_add - Vt_add
            At = At_add

            """ variational distribution q uses exp(At / eta) instead of exp(Qt / eta), for stable learning"""
            q = torch.softmax(At_add / self.eta, axis=0)
            actor_loss = -torch.mean(torch.sum(q.detach() * log_prob_add, axis=0))

            eta_loss = self.eta * self.eps_eta + self.eta * torch.mean(
                torch.log(torch.exp((At_add) / self.eta).mean(axis=0))
            )

            ss = 1.0 / (std ** 2)  # (batch_size * len_tr, action_dim)
            ss_old = 1.0 / (std_old ** 2)

            """
            KL-Divergence losses(related to alpha) implemented using methods introduced from V-MPO paper
            https://arxiv.org/abs/1909.12238
            """

            # mu
            d_mu = mu - mu_old.detach()  # (batch_size * len_tr, action_dim)
            KLD_mu = 0.5 * torch.sum(d_mu * 1.0 / ss_old.detach() * d_mu, axis=-1)
            mu_loss = torch.mean(
                self.alpha_mu * (self.eps_alpha_mu - KLD_mu.detach())
                + self.alpha_mu.detach() * KLD_mu
            )

            # sigma
            KLD_sigma = 0.5 * (
                torch.sum(1.0 / ss * ss_old.detach(), axis=-1)
                - ss.shape[-1]
                + torch.log(
                    torch.prod(ss, axis=-1) / torch.prod(ss_old.detach(), axis=-1)
                )
            )
            sigma_loss = torch.mean(
                self.alpha_sigma * (self.eps_alpha_sigma - KLD_sigma.detach())
                + self.alpha_sigma.detach() * KLD_sigma
            )

            alpha_loss = mu_loss + sigma_loss

        else:
            pi = self.actor(state)  # pi,Q: (batch_size, len_tr, dim_action)
            pi_next = self.actor(next_state)
            Q = self.critic(state)
            Q_a = Q.gather(1, action.long())

            with torch.no_grad():
                # calculate Q_ret using Retrace
                Qt = self.target_critic(state)  # Q_target
                Qt_next = self.target_critic(next_state)
                pit = self.target_actor(state)

                Qt_a = Qt.gather(1, action.long())
                prob_t = pi.gather(
                    1, action.long()
                )  # (batch_size * len_tr, 1), target policy probability

                c = torch.clip(
                    prob_t / (prob_b + 1e-6), max=1.0
                )  # (batch_size * len_tr, 1), prod of importance ratio and gamma

                if self.critic_loss_type == "1step_TD":
                    Qret = reward + self.gamma * (1 - done) * torch.sum(
                        pi_next * Qt_next, axis=-1, keepdim=True
                    )
                elif self.critic_loss_type == "retrace":
                    Qret = reward + self.gamma * torch.sum(
                        pi_next * Qt_next, axis=-1, keepdim=True
                    ) * (1 - done)

                    # temporarily reshaping values
                    # (batch_size * len_tr, item_dim) -> (batch_size, len_tr, item_dim)
                    Qret = Qret.view(self.batch_size, -1, *Qret.shape[1:])
                    Qt_a = Qt_a.view(self.batch_size, -1, *Qt_a.shape[1:])
                    c = c.view(self.batch_size, -1, *c.shape[1:])
                    done = done.view(self.batch_size, -1, *done.shape[1:])
                    for i in reversed(
                        range(Qret.shape[1] - 1)
                    ):  # along the trajectory length
                        Qret[:, i] += (
                            self.gamma
                            * c[:, i + 1]
                            * (Qret[:, i + 1] - Qt_a[:, i + 1])
                            * (1 - done[:, i])
                        )
                    Qret = Qret.view(-1, *Qret.shape[2:])

                pi_old = pit

            critic_loss = F.mse_loss(Q_a, Qret).mean()

            # calculate V, Advantage of Qt
            Vt = torch.sum(pi_old * Qt, axis=-1, keepdims=True)
            At = Qt - Vt

            """ variational distribution q uses exp(At / eta) instead of exp(Qt / eta), for stable learning"""
            q = torch.softmax(At / self.eta, axis=-1)
            actor_loss = -torch.mean(torch.sum(q.detach() * torch.log(pi), axis=-1))

            eta_loss = self.eta * self.eps_eta + self.eta * torch.mean(
                torch.log(torch.sum(pi_old * torch.exp(At / self.eta), axis=-1))
            )

            """
            KL-Divergence losses(related to alpha) implemented using methods introduced from V-MPO paper
            https://arxiv.org/abs/1909.12238
            """

            KLD_pi = pi_old.detach() * (torch.log(pi_old.detach()) - torch.log(pi))
            KLD_pi = torch.sum(KLD_pi, axis=len(pi_old.shape) - 1)
            alpha_loss = torch.mean(
                self.alpha_mu * (self.eps_alpha_mu - KLD_pi.detach())
                + self.alpha_mu.detach() * KLD_pi
            )

        loss = critic_loss + actor_loss + eta_loss + alpha_loss

        self.actor_optimizer.zero_grad()
        self.critic_optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.clip_grad_norm)
        torch.nn.utils.clip_grad_norm_(self.critic.parameters(), self.clip_grad_norm)
        self.actor_optimizer.step()
        self.critic_optimizer.step()
        self.reset_lgr_muls()

        self.num_learn += 1

        result = {
            "actor_loss": actor_loss.item(),
            "critic_loss": critic_loss.item(),
            "eta_loss": eta_loss.item(),
            "alpha_loss": alpha_loss.item(),
            "eta": self.eta.item(),
            "alpha_mu": self.alpha_mu.item(),
            "alpha_sigma": self.alpha_sigma.item(),
            "min_Q": Q.detach().cpu().numpy().min(),
            "max_Q": Q.detach().cpu().numpy().max(),
            "min_At": At.detach().cpu().numpy().min(),
            "max_At": At.detach().cpu().numpy().max(),
        }

        return result

    # reset Lagrange multipliers: eta, alpha_{mu, sigma}
    def reset_lgr_muls(self):
        self.eta.data = torch.max(self.eta, self.min_eta)
        self.alpha_mu.data = torch.max(self.alpha_mu, self.min_alpha_mu)
        self.alpha_sigma.data = torch.max(self.alpha_sigma, self.min_alpha_sigma)

    def update_target(self):
        self.target_actor.load_state_dict(self.actor.state_dict())
        self.target_critic.load_state_dict(self.critic.state_dict())

    def save(self, path):
        print(f"...Save model to {path}...")
        torch.save(
            {
                "actor": self.actor.state_dict(),
                "critic": self.critic.state_dict(),
                "actor_optimizer": self.actor_optimizer.state_dict(),
                "critic_optimizer": self.critic_optimizer.state_dict(),
            },
            os.path.join(path, "ckpt"),
        )

    def load(self, path):
        print(f"...Load model from {path}...")
        checkpoint = torch.load(os.path.join(path, "ckpt"), map_location=self.device)
        self.actor.load_state_dict(checkpoint["actor"])
        self.target_actor.load_state_dict(self.actor.state_dict())
        self.critic.load_state_dict(checkpoint["critic"])
        self.target_critic.load_state_dict(self.critic.state_dict())
        self.actor_optimizer.load_state_dict(checkpoint["actor_optimizer"])
        self.critic_optimizer.load_state_dict(checkpoint["critic_optimizer"])

    def process(self, transitions, step):
        result = {}

        # Process per step
        self.memory.store(transitions)
        delta_t = step - self.time_t
        self.time_t = step

        if self.memory.size >= self.batch_size and self.time_t >= self.start_train_step:
            for i in range(self.n_epoch):
                result = self.learn()
            self.update_target()

        return result

    def sync_in(self, weights):
        self.actor.load_state_dict(weights)

    def sync_out(self, device="cpu"):
        weights = self.actor.state_dict()
        for k, v in weights.items():
            weights[k] = v.to(device)
        sync_item = {
            "weights": weights,
        }
        return sync_item

    def interact_callback(self, transition):
        _transition = {}
        self.tmp_buffer.append(transition)
        if len(self.tmp_buffer) == self.n_step:
            for key in self.tmp_buffer[0].keys():
                _transition[key] = np.stack([t[key] for t in self.tmp_buffer], axis=1)

        return _transition


#!/usr/bin/env python
# ~*~ coding: utf8 ~*~
"""Find the domain average flux uncertainties."""
from __future__ import division, print_function
import sys
import glob
import os.path
import datetime

import numpy as np
import dateutil.tz
import netCDF4
import cf_units
import xarray

try:
    THIS_DIR = os.path.dirname(__file__)
except NameError:
    THIS_DIR = os.getcwd()

sys.path.insert(0, os.path.join(
    THIS_DIR, "..", "src"))
sys.path.append(THIS_DIR)

import atmos_flux_inversion.correlations
import atmos_flux_inversion.covariances
from atmos_flux_inversion.util import kronecker_product
from atmos_flux_inversion.linalg import kron

############################################################
# A few utility functions
def flush_output_streams():
    """Flush stdout and stderr."""
    sys.stdout.flush()
    sys.stderr.flush()

def write_progress_message(msg):
    """Write the message to stdout with time.

    Parameters
    ----------
    msg: str
    """
    flush_output_streams()
    print(datetime.datetime.now(UTC).strftime("%c"), msg)
    flush_output_streams()

UTC = dateutil.tz.tzutc()
SECONDS_PER_HOUR = 3600
HOURS_PER_DAY = 24
DAYS_PER_WEEK = 7
DAYS_PER_MONTH = 30

DAILY_FLUX_TIMESCALE = 7
DAILY_FLUX_FUN = "exp"
HOURLY_FLUX_TIMESCALE = 3
CORRELATION_LENGTH = 1000
GRID_RESOLUTION = 27
# I would like to add a fixed minimum at some point.
# full stds would then be sqrt(fixed^2 + varying^2)
# average seasonal variation (or some fraction thereof) might work.
# x2 since MsTMIP spread does not represent full uncertainty
# x5 since MsTMIP spread only represents monthly values and this uses sub-daily
# x10 matches model-model for Raczka for 200km/21d (0.68)
# x4 matches model-model for 1000km/7d (0.65)
FLUX_VARIANCE_VARYING_FRACTION = 4.
NC_ENGINE = "netcdf4"
FLUX_UNITS = "umol/m2/s"
FLUX_WINDOW = HOURS_PER_DAY * DAYS_PER_WEEK * 2

OBS_HOURS = (datetime.time(12), datetime.time(16))
"""Which observation times will be used in the inversion.

Assumed to be local solar. Should give afternoon hours for the domain.

I really hope I can assume this doesn't depend on latitude. That would
make this much more complicated.
"""
OBS_TIMES_PER_DAY = OBS_HOURS[1].hour - OBS_HOURS[0].hour
"""Observations used per site per day."""
CO2_MOLAR_MASS = 16 * 2 + 12.01
"""Molar mass of CO2 (g/mol).

Used to convert WRF fluxes to units expected by observation operator.
"""
DAYS_DROPPED_FROM_END = 1
"""Currently 1 to avoid problems with lack of fluxes in August."""
OBS_DAYS = 30
N_REALIZATIONS = 80
#  1    9m55 (80 realizations)    9m46
# 30 2h48m40 (80 realizations) 3h56m41
"""Number of days of obs to use."""
OBS_WINDOW = OBS_DAYS * OBS_TIMES_PER_DAY
"""Number of observation times."""

############################################################
# Get grid parameters

FLUX_INTERVAL = 6
"""The interval at which fluxes become available in hours.

Fluxes are usually integrated forward from hourly input, but I can't
solve for that in a reasonable timeframe.

This determines both the input and the output time resolution as well
as what the inversion solves for.

Note
----
Must divide twenty-four.
"""
INTERVALS_PER_DAY = HOURS_PER_DAY // FLUX_INTERVAL
FLUX_RESOLUTION = 27
"""Resolution of fluxes and influence functions in kilometers."""
UNCERTAINTY_RESOLUTION_REDUCTION_FACTOR = 6
"""How much coarser uncertainty is than mean estimate in the x direction.

If we compute the uncertainty at full resolution, the resulting file
is huge, starting in the hundreds of terabytes for a month.
Coarsening the resolution for the uncertainties allows us to still
report uncertainties within current computing constraints.
"""
# 4: 2h49 wall 132GiB mem 3h57 cpu
# 3: 5h11 wall 140GiB mem 6h27 cpu
UNCERTAINTY_FLUX_RESOLUTION = (UNCERTAINTY_RESOLUTION_REDUCTION_FACTOR *
                               FLUX_RESOLUTION * 1e3)
"""Resolution of posterior uncertainties in meters."""
UNCERTAINTY_TEMPORAL_RESOLUTION = "2D"
"""The resolution at which the uncertainty is calculated and saved.

Higher resolution means the uncertainties will be more accurate.
"""

INFLUENCE_PATHS = ["/mc1s2/s4/dfw5129/data/LPDM_2010_fpbounds/"
                   "ACT-America_trial5/2010/01/GROUP1",
                   "/mc1s2/s4/dfw5129/data/LPDM_2010_fpbounds/"
                   "candidacy_more_towers/2010/01/GROUP1"]
INFLUENCE_FILES = [
    name
    for path in INFLUENCE_PATHS
    for name in glob.glob(os.path.join(
        path,
        "LPDM_2010_01_{flux_interval:02d}hrly_{res:03d}km_molar_footprints.nc4"
        .format(flux_interval=FLUX_INTERVAL, res=FLUX_RESOLUTION)))]

TEST_DS = netCDF4.Dataset(INFLUENCE_FILES[0])

NX = len(TEST_DS.dimensions["dim_x"])
NY = len(TEST_DS.dimensions["dim_y"])
N_TIMES_BACK = len(TEST_DS.dimensions["time_before_observation"])

N_SITES = len(TEST_DS.dimensions["site"])
N_OBS_TIMES = len(TEST_DS.dimensions["observation_time"])

TEST_DS.close()
del TEST_DS

N_TIMES = INTERVALS_PER_DAY * DAYS_PER_MONTH + FLUX_WINDOW // FLUX_INTERVAL

if N_TIMES_BACK < FLUX_WINDOW / FLUX_INTERVAL:
    raise ValueError("FLUX_WINDOW too long for file")

N_GRID_POINTS = NY * NX
STATE_SIZE = N_GRID_POINTS * FLUX_WINDOW

OBS_VEC_SIZE = N_SITES * OBS_WINDOW
OBS_VEC_TOTAL_SIZE = N_SITES * N_OBS_TIMES

write_progress_message("Getting covariances")
spatial_correlations = (
    atmos_flux_inversion.correlations.HomogeneousIsotropicCorrelation.
    from_function(
        atmos_flux_inversion.correlations.ExponentialCorrelation(
            CORRELATION_LENGTH / GRID_RESOLUTION),
        (NY, NX),
        is_cyclic=False))
write_progress_message("Have spatial correlations")

hour_correlations = (
    atmos_flux_inversion.correlations.HomogeneousIsotropicCorrelation.
    from_function(
        atmos_flux_inversion.correlations.ExponentialCorrelation(
            HOURLY_FLUX_TIMESCALE / FLUX_INTERVAL),
        (INTERVALS_PER_DAY,),
        is_cyclic=True))
hour_correlations_matrix = hour_correlations.dot(np.eye(
    hour_correlations.shape[0]))
write_progress_message("Have hourly correlations")
day_correlations = (
    atmos_flux_inversion.correlations.make_matrix(
        atmos_flux_inversion.correlations.ExponentialCorrelation(
            DAILY_FLUX_TIMESCALE
        ),
        (N_TIMES // INTERVALS_PER_DAY,)))
write_progress_message("Have daily correlations")
temporal_correlations = kron(day_correlations,
                             hour_correlations_matrix)
write_progress_message("Have temporal correlations")

full_correlations = kronecker_product(
    temporal_correlations,
    spatial_correlations)
write_progress_message("Have combined correlations")
flux_std_pattern = xarray.open_dataset(
    "../data_files/2010_MsTMIP_flux_std.nc4",
    engine=NC_ENGINE
).get(
    ["E_TRA{:d}".format(i + 1) for i in range(1)]
).sel(
    Time=slice("2010-07-01", "2010-07-30")
).mean(
    dim="Time",
    keep_attrs=True,
)

# Ensure units work out
for flux_part in flux_std_pattern.data_vars.values():
    unit = (cf_units.Unit(flux_part.attrs["units"]))
    if unit is not FLUX_UNITS:
        flux_part *= unit.convert(1, FLUX_UNITS)
        flux_part.attrs["units"] = str(FLUX_UNITS)

reduced_flux_stds = (
    FLUX_VARIANCE_VARYING_FRACTION *
    flux_std_pattern["E_TRA1"].data)
write_progress_message("Have standard deviations")

spatial_covariance = (
    atmos_flux_inversion.covariances.CorrelationStandardDeviation(
        spatial_correlations, reduced_flux_stds
    )
)
write_progress_message("Have full spatial covariance")

prior_covariance = kronecker_product(
    temporal_correlations,
    spatial_covariance)

averager = np.full(
    prior_covariance.shape[0], 1. / prior_covariance.shape[0], dtype=np.float32
)

write_progress_message("Getting mean covariance for all")
cov_of_avg = averager.dot(prior_covariance.dot(averager))
write_progress_message("Got mean covariance for all")

write_progress_message("Covariance is {0:5.3f}\nStandard Deviation: {1:5.3f}"
                       .format(cov_of_avg, np.sqrt(cov_of_avg)))


# _*_ coding: utf-8 _*_

import tensorflow as tf
from tensorflow.python.ops.rnn_cell_impl import BasicLSTMCell, RNNCell
from tensorflow.python.ops.rnn_cell_impl import LSTMStateTuple
import numpy as np

'''
本欲想调整LatticeLSTM, 使其不依赖于固定的batch_size，但由于Lattice内部Cell的参数需要指定Shape，
而与batch_size对应的Shape在后续融合相关信息时，具有明显的实际意义，每个词汇对字符的贡献统计，
这也造成了本模型适用于处理批处理任务，对单个任务需进行调整后再处理较合适。

对于引入外部知识，以便对LSTM进行变形来说，应该是一种折中平衡吧。
'''

class CharLSTM(BasicLSTMCell):
    def __init__(self, lexicon_num_units, dtype, batch_size,
                 reuse=None, name=None, **kwargs):
        super(CharLSTM, self).__init__(reuse=reuse, name=name, **kwargs)
        self._lexicon_num_units = lexicon_num_units
        self._dtype = dtype
        self._char_state_tensor = tf.Variable(tf.zeros(shape=[batch_size, self._num_units]),
                                              dtype=self._dtype,
                                              trainable=False)

    def build(self, inputs_shape):
        # inputs_shape should be in the shape of [batch_size, char_embedding_size]
        if inputs_shape[-1].value is None:
            raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape)

        input_depth = inputs_shape[-1].value
        h_depth = self._num_units
        lexicon_state_depth = self._lexicon_num_units

        self._kernel = self.add_variable(name='multi_input_kernel',
                                         shape=[input_depth + h_depth, 4 * self._num_units])
        self._bias = self.add_variable(name='multi_input_bias',
                                       shape=[4 * self._num_units],
                                       initializer=tf.zeros_initializer(dtype=self._dtype))

        self._linking_kernel = self.add_variable(name='linking_kernel',
                                                 shape=[input_depth + lexicon_state_depth,
                                                        self._num_units])
        self._linking_bias = self.add_variable(name='linking_bias',
                                               shape=[self._num_units],
                                               initializer=tf.zeros_initializer(dtype=self._dtype))
        self.built = True

    def call(self, inputs, state):
        char_inputs = inputs[0]   # shape = [batch_size, input_dimension]
        state_inputs = inputs[1]  # shape = [batch_size, max_num_of_lexicon words, lexicon_state_dimension]

        # check whether the last dimension of state_inputs are all zero.
        # check_state_0 should be in the shape of [batch_size, max_num_of_lexicon words]
        check_state_0 = tf.reduce_sum(state_inputs, axis=-1)
        # check_state_1 should be in the shape of [batch_size]
        check_state_1 = tf.reduce_sum(check_state_0, axis=-1)

        # 查找匹配含有词汇的索引，只处理该部分信息，避免较多无词库匹配的信息参与计算消耗资源
        # state_inputs_indices_for_lexicon should be in the shape of [batch_size, 2]
        state_inputs_indices_for_lexicon = tf.where(tf.not_equal(check_state_0, 0))

        # 查找不含有词汇的索引，避免较多无词库匹配的信息参与计算消耗资源
        # tf.where(tf.equal(check_state_1, 0)) should be in the shape of [batch_size, 1]
        # state_inputs_indices_for_not_lexicon should be in the shape of [batch_size]
        state_inputs_indices_for_not_lexicon = tf.squeeze(tf.where(tf.equal(check_state_1, 0)))

        # 对不含词汇的细胞状态进行选择，主要是针对标量数据，因其秩为0，需进行维度扩展
        # in case `[i]` is squeezed to scalar `i`, change it back to 1-dimension tensor `[i]` by `tf.expand_dims()`
        # otherwise, `[]` and `[i, j]` will remain as-is after tf.squeeze() and further conversion on it
        state_inputs_indices_for_not_lexicon = tf.cond(pred=tf.equal(tf.rank(state_inputs_indices_for_not_lexicon), 0),
                                                       true_fn=lambda: tf.expand_dims(
                                                           state_inputs_indices_for_not_lexicon, axis=0),
                                                       false_fn=lambda: state_inputs_indices_for_not_lexicon)

        # 含有词汇匹配的字符索引
        # char_inputs_indices_for_lexicon should be in the shape of [batch_size, 1]
        char_inputs_indices_for_lexicon = tf.where(tf.not_equal(check_state_1, 0))

        # 不含有词汇匹配的字符索引
        # char_inputs_indices_for_not_lexicon should be in the shape of [batch_size, 1]
        char_inputs_indices_for_not_lexicon = tf.where(tf.equal(check_state_1, 0))

        if self._state_is_tuple:
            c, h = state
        else:
            c, h = tf.split(value=state, num_or_size_splits=2, axis=1)

        # tf.concat([char_inputs, h], 1) should be in the shape of
        # [batch_size, char_embedding_size + state_dimension]
        # h should be in the shape of [batch_size, state_dimension]
        # self._kernel should be in the shape of [char_embedding_size + state_dimension, X]
        # gate_inputs should be in the shape of [batch_size, 4 * state_dimension]
        gate_inputs = tf.matmul(tf.concat([char_inputs, h], 1), self._kernel)
        gate_inputs = tf.nn.bias_add(gate_inputs, self._bias)

        i, j, f, o = tf.split(value=gate_inputs, num_or_size_splits=4, axis=1)

        new_c_without_lexicon = self._new_c_without_lexicon(i=i, f=f, j=j, c=c,
                                                            indices_tensor=state_inputs_indices_for_not_lexicon)
        new_c = tf.scatter_nd_update(self._char_state_tensor,
                                     indices=char_inputs_indices_for_not_lexicon,
                                     updates=new_c_without_lexicon)

        new_c = tf.cond(tf.not_equal(tf.shape(state_inputs_indices_for_not_lexicon)[-1],
                                     tf.shape(state_inputs)[0]),
                        true_fn=lambda: self._if_not_empty_lexicon_state(i, j, char_inputs, state_inputs,
                                                                         char_inputs_indices_for_lexicon,
                                                                         state_inputs_indices_for_lexicon,
                                                                         new_c),
                        false_fn=lambda: new_c)

        # 计算输出隐状态
        new_h = tf.multiply(self._activation(new_c), tf.nn.sigmoid(o))

        if self._state_is_tuple:
            new_state = LSTMStateTuple(new_c, new_h)
        else:
            new_state = tf.concat([new_c, new_h], 1)

        return new_h, new_state

    def _new_c_without_lexicon(self, i, f, j, c, indices_tensor):
        # indices_tensor should be in the shape of [batch_size]
        f_without_lexicon_state_input = tf.gather(f, indices=indices_tensor)
        i_without_lexicon_state_input = tf.gather(i, indices=indices_tensor)
        j_without_lexicon_state_input = tf.gather(j, indices=indices_tensor)
        # j_without_lexicon_state_input should be in the shape of [batch_size]

        # 运行常规LSTM描述逻辑
        forget_bias_tensor = tf.constant(self._forget_bias, dtype=f.dtype)
        new_c_without_lexicon_state = tf.add(
            tf.multiply(c, tf.nn.sigmoid(tf.add(f_without_lexicon_state_input,
                                                forget_bias_tensor))),
            tf.multiply(tf.nn.sigmoid(i_without_lexicon_state_input),
                        self._activation(j_without_lexicon_state_input)))

        return new_c_without_lexicon_state

    def _new_c_with_lexicon(self, i, j, char_inputs, state_inputs, indices_tensor):
        # char_inputs should be in the shape of [batch_size, char_embedding]
        # state_inputs should be in the shape of
        # [batch_size, max_num_of_lexicon words, lexicon_state_dimension]
        # indices_tensor is state_inputs_indices_for_lexicon, should be in the shape of [batch_size]
        char_inputs_with_lexicon_state = tf.gather_nd(char_inputs, indices=[indices_tensor])

        # 提取指定索引下的词汇状态信息
        # lexicon_state_inputs should be in the shape of [max_num_of_lexicon words, lexicon_state_dimension]
        lexicon_state_inputs = tf.gather_nd(state_inputs, indices=indices_tensor)

        i_with_lexicon_state_input = tf.gather_nd(i, indices=[indices_tensor])
        j_with_lexicon_state_input = tf.gather_nd(j, indices=[indices_tensor])

        # 常规输入门操作
        state_input_gate = tf.matmul(tf.concat([char_inputs_with_lexicon_state,
                                                lexicon_state_inputs], axis=-1),
                                     self._linking_kernel)
        state_input_gate = tf.nn.sigmoid(tf.nn.bias_add(state_input_gate, self._linking_bias))

        # 为了后面评估子词库对字符的贡献概率，引入了额外的门控单元 state_char_input_gate
        state_char_input_gate = tf.concat([state_input_gate,
                                           tf.nn.sigmoid(i_with_lexicon_state_input)], axis=1)

        # softmax 评估子词库对字符的贡献概率
        state_gate_weights, char_gate_weight = tf.split(
            tf.nn.softmax(state_char_input_gate, axis=0),
            num_or_size_splits=[tf.shape(lexicon_state_inputs)[0], 1],
            axis=1)

        # 常规LSTM操作
        new_c_with_lexicon_state = tf.add(
            tf.reduce_sum(tf.multiply(state_gate_weights, lexicon_state_inputs), axis=0),
            tf.multiply(char_gate_weight, j_with_lexicon_state_input))

        return new_c_with_lexicon_state

    def _if_not_empty_lexicon_state(self, i, j,
                                    char_inputs, state_inputs,
                                    char_inputs_indices_for_lexicon,
                                    state_inputs_indices_for_lexicon, new_c_in):
        new_c_with_lexicon = self._new_c_with_lexicon(i=i, j=j,
                                                      char_inputs=char_inputs,
                                                      state_inputs=state_inputs,
                                                      indices_tensor=state_inputs_indices_for_lexicon)
        # 根据新生成的含词库信息的字符状态更新字符的细胞状态
        new_c_out = tf.scatter_nd_update(new_c_in,
                                         indices=char_inputs_indices_for_lexicon,
                                         updates=new_c_with_lexicon)

        return new_c_out

class LexiconLSTM(BasicLSTMCell):
    def __init__(self, dtype, reuse=None, name=None, **kwargs):
        super(LexiconLSTM, self).__init__(reuse=reuse, name=name, **kwargs)
        self._dtype = dtype

    def build(self, inputs_shape):
        if inputs_shape[-1].value is None:
            raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape)

        input_depth = inputs_shape[-1].value
        h_depth = self._num_units

        # 需要指出此处的超参数3与LatticeLSTM构成某种对应关系
        self._kernel = self.add_variable(name='lexicon_kernel',
                                         shape=[input_depth + h_depth, 3 * self._num_units])
        self._bias = self.add_variable(name='lexicon_bias',
                                       shape=[3 * self._num_units],
                                       initializer=tf.zeros_initializer(dtype=self._dtype))
        self.built = True

    def call(self, inputs, state):
        sigmoid = tf.nn.sigmoid
        add = tf.add
        multiply = tf.multiply

        if self._state_is_tuple:
            c, h = state
        else:
            c, h = tf.split(value=state, num_or_size_splits=2, axis=1)

        gate_inputs = tf.matmul(tf.concat([inputs, h], 1), self._kernel)
        gate_inputs = tf.nn.bias_add(gate_inputs, self._bias)

        i, j, f = tf.split(value=gate_inputs, num_or_size_splits=3, axis=1)

        forget_bias_tensor = tf.constant(self._forget_bias, dtype=f.dtype)

        new_c = add(multiply(c, sigmoid(add(f, forget_bias_tensor))),
                    multiply(sigmoid(i), self._activation(j)))
        return new_c

class LatticeLSTMCell(RNNCell):
    ''' inherit from: tf.nn.rnn_cell.RNNCell
        Lattice Long short-term memory unit recurrent network cell. the implementation is based on
        https://arxiv.org/pdf/1805.02023.pdf
        Please be noted that LatticeLSTMCell should be called within tf.nn.dynamic_rnn
    '''
    def __init__(self, char_num_units, lexicon_num_units, batch_size, max_lexicon_words_num,
                 word_length_tensor, seq_len, dtype, **kwargs):
        super(LatticeLSTMCell, self).__init__(**kwargs)
        '''
        Parameters
        ----------
        char_num_units: int
            the num_units of char_lstm cell units.
            this is expected to be the same as lexicon_num_units by the paper.

        lexicon_num_units: int
            the num_units of lexicon_lstm cell units.
            this is expected to be the same as char_num_units by the paper.

        max_lexicon_words_num: int
            the upper bound of the lexicon words per characters.

        batch_size: int
            batch_size of the input data

        seq_len: int
            sequence_length of the input data

        dtype:
            data type defined for LatticeLSTMCell variable

        word_length_tensor: tensor
            tensor, it contains a batch of lexicon word length.
            this should be padded with 0 if there is matched
            lexicon word for the respective character.
            example:  char： 南 京 市 長 江 大 橋
                      word： 南京 市長
                      max_lexicon_words_num = 5
            then the word length tensor should be like tf.constant([[[2,0,0,0,0],
                                                                     [0,0,0,0,0],
                                                                     [2,0,0,0,0],
                                                                     [0,0,0,0,0],
                                                                     [0,0,0,0,0],
                                                                     [0,0,0,0,0],
                                                                     [0,0,0,0,0]]],
                                                                   dtype=tf.float32)
        '''
        # 简要描述逻辑：先针对字符，以该字符结束匹配词库中词汇信息，最大词汇个数为 max_lexicon_words_num
        # 然后计算匹配的词汇对该字符的贡献，并融入在字符的细胞状态中；
        # 其中为了加快计算逻辑，对含有词汇的字符与不含有词汇的字符进行了分离，
        # 此外为了对LSTM进行变种计算，进行了数据的堆叠计算工作.
        self._char_lstm = CharLSTM(dtype=dtype,
                                   num_units=char_num_units,
                                   batch_size=batch_size,
                                   lexicon_num_units=lexicon_num_units,
                                   name='character_lstm')

        self._lexicon_lstm = LexiconLSTM(dtype=dtype,
                                         num_units=lexicon_num_units,
                                         name='lexicon_word_lstm')

        # word_length_tensor should be in the shape of [batched_size, seq_len, max_lexicon_words_num]
        self.word_length_tensor = word_length_tensor
        self.max_lexicon_words_num = max_lexicon_words_num
        self.seq_len = seq_len
        self.time_step = 0
        self._dtype = dtype

        lexicon_state_init_value = tf.zeros(shape=[batch_size, self.seq_len,
                                                   self.max_lexicon_words_num,
                                                   lexicon_num_units])

        # lexicon_state_tensor should be in the shape of
        # [batch_size, seq_len, max_lexicon_words_num, state_dimension]
        self.lexicon_state_tensor = tf.Variable(initial_value=lexicon_state_init_value,
                                                trainable=False,
                                                dtype=self._dtype)

    def build(self, inputs_shape):
        # inputs shape should be in the shape
        # [[batch_size, char_embedding_size],
        #  [batch_size, max_lexicon_words_num, lexicon_word_embedding_size]]

        self._char_lstm.build(inputs_shape[0])
        self._lexicon_lstm.build(inputs_shape[1])
        self.lexicon_shape = inputs_shape[1]

        if self.lexicon_shape[1] != self.max_lexicon_words_num:
            raise ValueError('max_lexicon_words_num should be equal to lexicon input')

        self.built = True

    @property
    def state_size(self):
        return self._char_lstm.state_size

    @property
    def output_size(self):
        return self._char_lstm.output_size

    def zero_state(self, batch_size, dtype):
        return self._char_lstm.zero_state(batch_size, dtype)

    def call(self, inputs, state):
        '''
        Parameters
        ----------
        inputs: list of tensors
            inputs here should be a tensors of character_inputs(character embedding inputs) and
            lexicon_inputs(lexicon word embedding inputs). char_embedding_inputs has shape [batch_size, char_embedding_size]
            lexicon_embedding_inputs has shape [batch_size, max_lexicon_words_num, lexicon_word_embedding_size]
            please be noted that number of lexicon words is expect to be upper bounded. so if the number of words
            for a character is less than max_lexicon_words_num, the word embedding inputs should be padded with zero.
            example: inputs = [char_embedding_inputs, word_embedding_inputs]

        state: tensors
            Either a single 2-D tensor, or a tuple of tensors matching the arity and shapes of state.

        Returns
        -------
        output: tensor
            tensor of hidden output of character units with shape [batch_size, self.output_size]

        new_state: tensor or a tuple of tensors
            Either a single 2-D tensor, or a tuple of tensors matching the arity and shapes of state.
        '''
        char_input = inputs[0]      # shape = [batch_size, char_embedding_size]
        lexicon_inputs = inputs[1]  # shape = [batch_size, max_lexicon_words_num, lexicon_word_embedding_size]

        # 根据时间步获取对应词库状态, 其中 self.lexicon_state_tensor 会迭代更新
        # lexicon_state_tensor should be in the shape of [batch_size, max_lexicon_words_num, state_dimension]
        lexicon_state_tensor = tf.gather(self.lexicon_state_tensor, axis=1, indices=self.time_step)
        char_hidden_output, char_state = self._char_lstm.call([char_input, lexicon_state_tensor], state)

        # max_lexicon_words_num可视作局部字符的time_step
        # 经过融合处理后，作为字符融合了词库的过程
        for word_index in range(self.max_lexicon_words_num):
            self.lexicon_state_tensor = self._update_lexicon_state_per_word(lexicon_inputs=lexicon_inputs,
                                                                            word_index=word_index,
                                                                            char_state=char_state)

        self.time_step = self.time_step + 1  # time_step 向后推进

        # reset the lexicon_state_tensor after finish a loop, 一个循环后进行重置
        self.lexicon_state_tensor = tf.cond(tf.equal(tf.mod(self.time_step, self.seq_len - 1), 0),
                                            true_fn=lambda: tf.assign(ref=self.lexicon_state_tensor,
                                                                      value=tf.zeros_like(
                                                                          self.lexicon_state_tensor)),
                                            false_fn=lambda: self.lexicon_state_tensor)

        self.time_step = np.remainder(self.time_step, self.seq_len - 1)  # 求模以便重置

        return char_hidden_output, char_state

    def _update_lexicon_state_per_word(self, lexicon_inputs, word_index, char_state):
        # 根据词索引获取对应对应词库值
        # lexicon_inputs should be in the shape of
        # [batch_size, max_lexicon_words_num, lexicon_word_embedding_size]

        # lexicon_input_per_word should be in the shape of [batch_size, state_dimension]
        lexicon_input_per_word = tf.gather(lexicon_inputs, axis=1, indices=word_index)

        # 根据时间步获取对应词长度
        # self.word_length_tensor should be in the shape of [batch_size, seq_len, max_lexicon_words_num]
        # word_length_per_time_step should be in the shape of [batch_size, max_lexicon_words_num]
        word_length_per_time_step = tf.gather(self.word_length_tensor, axis=1, indices=self.time_step)

        # 根据词索引获取对应时间步 词长度中的对应值, 即一个整数描述词汇的长度
        # word_length should be in the shape of [batch_size]
        word_length = tf.gather(word_length_per_time_step, axis=1, indices=word_index)

        # 引入char_state 为了与该字符进行关联，并根据依据索引变化的 lexicon_input_per_word 词库调整词库状态
        # lexicon_state should be in the shape of [batch_size, state_dimension]
        lexicon_state = self._lexicon_lstm.call(lexicon_input_per_word, char_state)

        # temp_lexicon_state_to_char_index should be an integer, 词库中词汇长度字符进行映射，经分词后词汇长度会不一致
        temp_lexicon_state_to_char_index = self.time_step + word_length - 1

        # 通过判断词汇长度不为0，获取对应的状态索引值，并组织为 batch_size 个数据
        # not equal result should be an bool, assert the match state
        # lexicon_state_index should be in the shape of [batch_size, 1]
        lexicon_state_index = tf.where(tf.not_equal(temp_lexicon_state_to_char_index,
                                                    self.time_step - 1))

        # 根据词库匹配命中的索引值，提取相应的索引值，每个对应一个索引，并组织为 batch_size 个数据
        # lexicon_state_to_char_index should be in the shape of [batch_size]
        lexicon_state_to_char_index = tf.gather_nd(temp_lexicon_state_to_char_index,
                                                   indices=lexicon_state_index)

        # 根据词汇长度的索引，提取对应的词库状态值，只针对词库长度不为0进行处理
        # lexicon_state_update should be in the shape of [batch_size, state_dimension]
        lexicon_state_update = tf.gather_nd(lexicon_state, indices=lexicon_state_index)

        # 根据词汇匹配索引值，在词汇大小方向上扩展并设置为1
        # word_index_for_stack should be in the shape of [batch_size]
        word_index_for_stack = tf.ones_like(lexicon_state_to_char_index) * word_index

        # word_index_for_stack should be in the shape of [batch_size]
        # tf.squeeze(lexicon_state_index)将剔除列方向上的维度
        lexicon_state_index_for_stack = tf.cast(tf.squeeze(lexicon_state_index), dtype=self._dtype)

        # 堆叠的逻辑，实际描述同一个逻辑，可理解为递进描述
        # indices should be in the shape of [batch_size, 3]
        # 此处采用堆叠3个向量与LexiconLSTM 中kernel与bias中超参数3保持一致，
        # 并与tf.split中的超参数3保持一致，即后续向量能被整除
        indices = tf.stack([lexicon_state_index_for_stack,   # 长度不为0的词汇对应的索引
                            lexicon_state_to_char_index,     # 提取长度不为0的词汇对应的索引
                            word_index_for_stack], axis=-1)  # 词汇长度对应的索引扩展后设置为1

        # updated_lexicon_state_tensor should be in the shape of
        # [batch_size, seq_len, max_lexicon_words_num, state_dimension]
        updated_lexicon_state_tensor = tf.scatter_nd_update(ref=self.lexicon_state_tensor,
                                                            indices=tf.cast(indices, dtype=tf.int32),
                                                            updates=lexicon_state_update)
        return updated_lexicon_state_tensor


import os
import subprocess
import numpy as np
import pandas as pd
import tempfile

BACKGROUND_FREQS = np.array([0.25, 0.25, 0.25, 0.25])
DATABASE_PATH = "/users/amtseng/tfmodisco/data/processed/motif_databases/HOCOMOCO_JASPAR_motifs.txt"

def import_database_pfms(database_path):
    """
    Imports the database of PFMs by reading through the entire database and
    constructing a dictionary mapping motif IDs to NumPy arrays of PFMs.
    """
    motif_dict = {}
    with open(database_path, "r") as f:
        try:
            while True:
                line = next(f)
                if line.startswith("MOTIF"):
                    key = line.strip().split()[1]
                    header = next(f)
                    motif_width = int(header.split()[5])
                    motif = np.empty((motif_width, 4))
                    for i in range(motif_width):
                        motif[i] = np.array([
                            float(x) for x in next(f).strip().split()
                        ])
                    motif_dict[key] = motif
        except StopIteration:
            pass
    return motif_dict

def export_pfms_to_meme_format(
    pfms, outfile, background_freqs=None, names=None
):
    """
    Exports a set of PFMs to MEME motif format. Includes the background
    frequencies `BACKGROUND_FREQS`.
    Arguments:
        `pfms`: a list of L x 4 PFMs (where L can be different for each PFM)
        `outfile`: path to file to output the MEME-format PFMs
        `background_freqs`: background frequencies of A, C, G, T as a length-4
            NumPy array; defaults to `BACKGROUND_FREQS`
        `names`: if specified, a list of unique names to give to each PFM, must
            be parallel to `pfms`
    """
    if names is None:
        names = [str(i) for i in range(len(pfms))]
    else:
        assert len(names) == pfms
        assert len(names) == len(np.unique(names))
    if background_freqs is None:
        background_freqs = BACKGROUND_FREQS

    os.makedirs(os.path.dirname(outfile), exist_ok=True)
    with open(outfile, "w") as f:
        f.write("MEME version 5\n\n")
        f.write("ALPHABET= ACGT\n\n")
        f.write("Background letter frequencies\n")
        f.write("A %f C %f G %f T %f\n\n" % tuple(background_freqs))
        for i in range(len(pfms)):
            pfm, name = pfms[i], names[i]
            f.write("MOTIF %s\n" % name)
            f.write("letter-probability matrix:\n")
            for row in pfm:
                f.write(" ".join([str(freq) for freq in row]) + "\n")
            f.write("\n")

def run_tomtom(
    query_motif_file, target_motif_file, outdir, show_output=True
):
    """
    Runs TOMTOM given the target and query motif files. The default threshold
    of q < 0.5 is used to filter for matches.
    Arguments:
        `query_motif_file`: file containing motifs in MEME format, which will
            be the query motifs for which matches are found
        `target_motif_file`: file containing motifs in MEME format, which will
            be used to search for matches
        `outdir`: path to directory to store results
        `show_output`: whether or not to show TOMTOM output
    """
    comm = ["tomtom"]
    comm += [query_motif_file, target_motif_file]
    comm += ["-oc", outdir]
    comm += ["-no-ssc"]
    comm += ["-dist", "pearson"]
    comm += ["-min-overlap", "5"]
    proc = subprocess.run(comm, capture_output=(not show_output))

def import_tomtom_results(tomtom_dir):
    """
    Imports the TOMTOM output directory as a Pandas DataFrame.
    Arguments:
        `tomtom_dir`: TOMTOM output directory, which contains the output file
            "tomtom.tsv"
    Returns a Pandas DataFrame.
    """
    return pd.read_csv(
        os.path.join(tomtom_dir, "tomtom.tsv"), sep="\t", header=0,
        index_col=False, comment="#"
    )

def match_motifs_to_targets(
    query_pfms, target_pfms, temp_dir=None, show_tomtom_output=False
):
    """
    For each motif in the query PFMs, finds the best match to the target PFMs,
    based on TOMTOM q-value.
    Arguments:
        `query_pfms`: list of L x 4 PFMs to look for matches for
        `target_pfms`: list of L x 4 PFMs to match to
        `temp_dir`: a temporary directory to store intermediates; defaults to
            a randomly created directory
        `show_tomtom_output`: whether to show TOMTOM output when running
    Returns an array of indices parallel to `query_pfms`, where each index is
    denotes the best PFM within `target_pfms` that matches the query PFM. If
    a good match is not found (i.e. based on TOMTOM's threshold), the index will
    be -1.
    """
    if temp_dir is None:
        temp_dir_obj = tempfile.TemporaryDirectory()
        temp_dir = temp_dir_obj.name
    else:
        temp_dir_obj = None

    # Convert motifs to MEME format
    query_motif_file = os.path.join(temp_dir, "query_motifs.txt")
    target_motif_file = os.path.join(temp_dir, "target_motifs.txt")
    export_pfms_to_meme_format(query_pfms, query_motif_file)
    export_pfms_to_meme_format(target_pfms, target_motif_file)

    # Run TOMTOM
    tomtom_dir = os.path.join(temp_dir, "tomtom")
    run_tomtom(
        query_motif_file, target_motif_file, tomtom_dir,
        show_output=show_tomtom_output
    )

    # Find results, mapping each query motif to target index
    # The query/target IDs are the indices
    tomtom_table = import_tomtom_results(tomtom_dir)
    match_inds = []
    for i in range(len(query_pfms)):
        rows = tomtom_table[tomtom_table["Query_ID"] == i]
        if rows.empty:
            match_inds.append(-1)
            continue
        target_id = rows.loc[rows["q-value"].idxmin()]["Target_ID"]
        match_inds.append(target_id)

    if temp_dir_obj is not None:
        temp_dir_obj.cleanup()

    return np.array(match_inds)

def match_motifs_to_database(
    query_pfms, top_k=5, temp_dir=None, database_path=DATABASE_PATH,
    show_tomtom_output=False
):
    """
    For each motif in the query PFMs, finds the best matches to the TOMTOM
    database, ranked by TOMTOM q-value.
    Arguments:
        `query_pfms`: list of L x 4 PFMs to look for matches for
        `top_k`: the number of motifs to return based on q-value
        `temp_dir`: a temporary directory to store intermediates; defaults to
            a randomly created directory
        `database_path`: the path to a TOMTOM motif database; defaults to
            DATABASE_PATH
        `show_tomtom_output`: whether to show TOMTOM output when running
    Returns a list of lists of (motif name, motif PFM, q-value) tuples
    parallel to `query_pfms`, where each sublist of tuples is the set of motif
    names, motif PFMs (as NumPy arrays), and q-values for the corresponding
    query motif. Each sublit is sorted in ascending order by q-value. If fewer
    than `top_k` matches are found (based on TOMTOM's threshold), the returned
    sublist will be shorter (and may even be empty).
    """
    # First, import the database PFMs
    database_pfms = import_database_pfms(database_path)

    if temp_dir is None:
        temp_dir_obj = tempfile.TemporaryDirectory()
        temp_dir = temp_dir_obj.name
    else:
        temp_dir_obj = None

    # Convert motifs to MEME format
    query_motif_file = os.path.join(temp_dir, "query_motifs.txt")
    export_pfms_to_meme_format(query_pfms, query_motif_file)

    # Run TOMTOM
    tomtom_dir = os.path.join(temp_dir, "tomtom")
    run_tomtom(
        query_motif_file, database_path, tomtom_dir,
        show_output=show_tomtom_output
    )

    # Find results, mapping each query motif to target index
    # The query/target IDs are the indices
    tomtom_table = import_tomtom_results(tomtom_dir)
    matches = []
    for i in range(len(query_pfms)):
        rows = tomtom_table[tomtom_table["Query_ID"] == i]
        if rows.empty:
            matches.append([])
            continue
        rows = rows.sort_values("q-value").head(top_k)
        tups = list(zip(rows["Target_ID"], rows["q-value"]))
        tups = [
            (tup[0], database_pfms[tup[0]], tup[1]) for tup in tups
        ]
        matches.append(tups)

    if temp_dir_obj is not None:
        temp_dir_obj.cleanup()

    return matches