"""
Defines the CloudNoiseModel class and supporting functions
"""
#***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# 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 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************

import collections as _collections
import itertools as _itertools
import warnings as _warnings

import numpy as _np
import scipy.sparse as _sps

from pygsti.baseobjs import statespace as _statespace
from pygsti.models.implicitmodel import ImplicitOpModel as _ImplicitOpModel, _init_spam_layers
from pygsti.models.layerrules import LayerRules as _LayerRules
from pygsti.models.memberdict import OrderedMemberDict as _OrderedMemberDict
from pygsti.evotypes import Evotype as _Evotype
from pygsti.forwardsims.forwardsim import ForwardSimulator as _FSim
from pygsti.forwardsims.mapforwardsim import MapForwardSimulator as _MapFSim
from pygsti.forwardsims.matrixforwardsim import MatrixForwardSimulator as _MatrixFSim
from pygsti.modelmembers import operations as _op
from pygsti.modelmembers import povms as _povm
from pygsti.modelmembers import states as _state
from pygsti.modelmembers.operations import opfactory as _opfactory
from pygsti.modelmembers.modelmembergraph import ModelMemberGraph as _MMGraph
from pygsti.baseobjs.basis import BuiltinBasis as _BuiltinBasis, ExplicitBasis as _ExplicitBasis
from pygsti.baseobjs.label import Label as _Lbl, CircuitLabel as _CircuitLabel
from pygsti.baseobjs.verbosityprinter import VerbosityPrinter as _VerbosityPrinter
from pygsti.baseobjs.qubitgraph import QubitGraph as _QubitGraph
from pygsti.tools import basistools as _bt
from pygsti.tools import internalgates as _itgs
from pygsti.tools import optools as _ot
from pygsti.baseobjs.basisconstructors import sqrt2, id2x2, sigmax, sigmay, sigmaz
from pygsti.processors.processorspec import ProcessorSpec as _ProcessorSpec, QubitProcessorSpec as _QubitProcessorSpec

class CloudNoiseModel(_ImplicitOpModel):
    """
    A n-qudit model using a low-weight and geometrically local error model with a common "global idle" operation.

    Parameters
    ----------
    processor_spec : ProcessorSpec
        The processor specification to create a model for.  This object specifies the
        gate names and unitaries for the processor, and their availability on the
        processor.

    gatedict : dict
        A dictionary (an `OrderedDict` if you care about insertion order) that
        associates with string-type gate names (e.g. `"Gx"`) :class:`LinearOperator`,
        `numpy.ndarray`, or :class:`OpFactory` objects. When the objects may act on
        fewer than the total number of qudits (determined by their dimension/shape) then
        they are repeatedly embedded into operation on the entire state space as specified
        by their availability within `processor_spec`.  These operations represent the ideal
        target operations, and thus, any `LinearOperator` or `OpFactory` objects must be *static*,
        i.e., have zero parameters.

    prep_layers, povm_layers : None or operator or dict or list, optional
        The SPAM operations as n-qudit layer operations.  If `None`, then
        no preps (or POVMs) are created.  If a dict, then the keys are
        labels and the values are layer operators.  If a list, then the
        elements are layer operators and the labels will be assigned as
        "rhoX" and "MX" where X is an integer starting at 0.  If a single
        layer operation is given, then this is used as the sole prep or
        POVM and is assigned the label "rho0" or "Mdefault" respectively.

    build_cloudnoise_fn : function, optional
        A function which takes a single :class:`Label` as an argument and
        returns the cloud-noise operation for that primitive layer
        operation.  Note that if `errcomp_type="gates"` the returned
        operator should be a superoperator whereas if
        `errcomp_type="errorgens"` then the returned operator should be
        an error generator (not yet exponentiated).

    build_cloudkey_fn : function, optional
        An function which takes a single :class:`Label` as an argument and
        returns a "cloud key" for that primitive layer.  The "cloud" is the
        set of qudits that the error (the operator returned from
        `build_cloudnoise_fn`) touches -- and the "key" returned from this
        function is meant to identify that cloud.  This is used to keep track
        of which primitive layer-labels correspond to the same cloud - e.g.
        the cloud-key for ("Gx",2) and ("Gy",2) might be the same and could
        be processed together when selecing sequences that amplify the parameters
        in the cloud-noise operations for these two labels.  The return value
        should be something hashable with the property that two noise
        which act on the same qudits should have the same cloud key.

    simulator : ForwardSimulator or {"auto", "matrix", "map"}
        The circuit simulator used to compute any
        requested probabilities, e.g. from :method:`probs` or
        :method:`bulk_probs`.  The default value of `"auto"` automatically
        selects the simulation type, and is usually what you want. Other
        special allowed values are:

        - "matrix" : op_matrix-op_matrix products are computed and
          cached to get composite gates which can then quickly simulate
          a circuit for any preparation and outcome.  High memory demand;
          best for a small number of (1 or 2) qubits.
        - "map" : op_matrix-state_vector products are repeatedly computed
          to simulate circuits.  Slower for a small number of qubits, but
          faster and more memory efficient for higher numbers of qubits (3+).

    evotype : Evotype or str, optional
        The evolution type of this model, describing how states are
        represented.  The special value `"default"` is equivalent
        to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.

    errcomp_type : {"gates","errorgens"}
        How errors are composed when creating layer operations in the created
        model.  `"gates"` means that the errors on multiple gates in a single
        layer are composed as separate and subsequent processes.  Specifically,
        the layer operation has the form `Composed(target,idleErr,cloudErr)`
        where `target` is a composition of all the ideal gate operations in the
        layer, `idleErr` is idle error (`.operation_blks['layers']['globalIdle']`),
        and `cloudErr` is the composition (ordered as layer-label) of cloud-
        noise contributions, i.e. a map that acts as the product of exponentiated
        error-generator matrices.  `"errorgens"` means that layer operations
        have the form `Composed(target, error)` where `target` is as above and
        `error` results from composing the idle and cloud-noise error
        *generators*, i.e. a map that acts as the exponentiated sum of error
        generators (ordering is irrelevant in this case).

    implicit_idle_mode : {'none', 'add_global', 'pad_1Q'}
        The way idle operations are added implicitly within the created model. `"none"`
        doesn't add any "extra" idle operations when there is a layer that contains some
        gates but not gates on all the qubits.  `"add_global"` adds the global idle operation,
        i.e., the operation for a global idle layer (zero gates - a completely empty layer),
        to every layer that is simulated, using the global idle as a background idle that always
        occurs regardless of the operation.  `"pad_1Q"` applies the 1-qubit idle gate (if one
        exists) to all idling qubits within a circuit layer.

    verbosity : int, optional
        An integer >= 0 dictating how must output to send to stdout.
    """

    def __init__(self, processor_spec, gatedict,
                 prep_layers=None, povm_layers=None,
                 build_cloudnoise_fn=None, build_cloudkey_fn=None,
                 simulator="map", evotype="default", errcomp_type="gates",
                 implicit_idle_mode="none", verbosity=0):

        qudit_labels = processor_spec.qudit_labels
        state_space = _statespace.QubitSpace(qudit_labels) if isinstance(processor_spec, _QubitProcessorSpec) \
            else _statespace.QuditSpace(qudit_labels, processor_spec.qudit_udims)

        simulator = _FSim.cast(simulator,
                               state_space.num_qubits if isinstance(state_space, _statespace.QubitSpace) else None)
        prefer_dense_reps = isinstance(simulator, _MatrixFSim)
        evotype = _Evotype.cast(evotype, default_prefer_dense_reps=prefer_dense_reps)

        # Build gate dictionaries. A value of `gatedict` can be an array, a LinearOperator, or an OpFactory.
        # For later processing, we'll create mm_gatedict to contain each item as a ModelMember.  For cloud-
        # noise models, these gate operations should be *static* (no parameters) as they represent the target
        # operations and all noise (and parameters) are assumed to enter through the cloudnoise members.
        mm_gatedict = _collections.OrderedDict()  # static *target* ops as ModelMembers
        for key, gate in gatedict.items():
            if isinstance(gate, _op.LinearOperator):
                assert(gate.num_params == 0), "Only *static* ideal operators are allowed in `gatedict`!"
                mm_gatedict[key] = gate
            elif isinstance(gate, _opfactory.OpFactory):
                assert(gate.num_params == 0), "Only *static* ideal factories are allowed in `gatedict`!"
                mm_gatedict[key] = gate
            else:  # presumably a numpy array or something like it:
                mm_gatedict[key] = _op.StaticArbitraryOp(gate, evotype, state_space=None)  # use default state space
            assert(mm_gatedict[key]._evotype == evotype), \
                ("Custom gate object supplied in `gatedict` for key %s has evotype %s (!= expected %s)"
                 % (str(key), str(mm_gatedict[key]._evotype), str(evotype)))

        #Set other members
        self.processor_spec = processor_spec
        self.errcomp_type = errcomp_type

        idle_names = self.processor_spec.idle_gate_names
        global_idle_name = self.processor_spec.global_idle_gate_name

        # Set noisy_global_idle_name == global_idle_name if the global idle gate isn't the perfect identity
        #  and if we're generating cloudnoise members (if we're not then layer rules could encouter a key error
        #  if we let noisy_global_idle_name be non-None).
        global_idle_gate = mm_gatedict.get(global_idle_name, None)
        if (global_idle_gate is not None) and (build_cloudnoise_fn is not None) \
           and (build_cloudnoise_fn(self.processor_spec.global_idle_layer_label) is not None):
            noisy_global_idle_name = global_idle_name
        else:
            noisy_global_idle_name = None

        singleq_idle_layer_labels = {}
        for idle_name in idle_names:
            if self.processor_spec.gate_num_qubits(idle_name) == 1:
                for idlelayer_sslbls in self.processor_spec.resolved_availability(idle_name, 'tuple'):
                    if idlelayer_sslbls is None: continue  # case of 1Q model with "global" idle
                    assert(len(idlelayer_sslbls) == 1)  # should be a 1-qubit gate!
                    if idlelayer_sslbls not in singleq_idle_layer_labels:
                        singleq_idle_layer_labels[idlelayer_sslbls] = _Lbl(idle_name, idlelayer_sslbls)
        #assert(set(idle_names).issubset([global_idle_name])), \
        #    "Only global idle operations are allowed in a CloudNoiseModel!"

        layer_rules = CloudNoiseLayerRules(errcomp_type, qudit_labels, implicit_idle_mode, singleq_idle_layer_labels,
                                           noisy_global_idle_name)
        super(CloudNoiseModel, self).__init__(state_space, layer_rules, "pp", simulator=simulator, evotype=evotype)

        flags = {'auto_embed': False, 'match_parent_statespace': False,
                 'match_parent_evotype': True, 'cast_to_type': None}
        self.prep_blks['layers'] = _OrderedMemberDict(self, None, None, flags)
        self.povm_blks['layers'] = _OrderedMemberDict(self, None, None, flags)
        self.operation_blks['gates'] = _OrderedMemberDict(self, None, None, flags)
        self.operation_blks['cloudnoise'] = _OrderedMemberDict(self, None, None, flags)
        self.operation_blks['layers'] = _OrderedMemberDict(self, None, None, flags)
        self.instrument_blks['layers'] = _OrderedMemberDict(self, None, None, flags)
        self.factories['gates'] = _OrderedMemberDict(self, None, None, flags)
        self.factories['cloudnoise'] = _OrderedMemberDict(self, None, None, flags)
        self.factories['layers'] = _OrderedMemberDict(self, None, None, flags)

        printer = _VerbosityPrinter.create_printer(verbosity)
        printer.log("Creating a %d-qudit cloud-noise model" % self.processor_spec.num_qudits)

        # a dictionary of "cloud" objects
        # keys = cloud identifiers, e.g. (target_qudit_indices, cloud_qudit_indices) tuples
        # values = list of gate-labels giving the gates (primitive layers?) associated with that cloud (necessary?)
        self._clouds = _collections.OrderedDict()

        for gn in self.processor_spec.gate_names:
            # process gate names (no sslbls, e.g. "Gx", not "Gx:0") - we'll check for the
            # latter when we process the corresponding gate name's availability

            gate_unitary = self.processor_spec.gate_unitaries[gn]
            resolved_avail = self.processor_spec.resolved_availability(gn)
            gate = mm_gatedict.get(gn, None)  # a static op or factory, no need to consider if "independent" (no params)
            gate_is_factory = callable(gate_unitary) or isinstance(gate, _opfactory.OpFactory)
            #gate_is_noiseless_identity = (gate is None) or \
            #    (isinstance(gate, _op.ComposedOp) and len(gate.factorops) == 0)

            if gate is not None:  # (a gate name may not be in gatedict if it's an identity without any noise)
                if gate_is_factory:
                    self.factories['gates'][_Lbl(gn)] = gate
                else:
                    self.operation_blks['gates'][_Lbl(gn)] = gate

            if callable(resolved_avail) or resolved_avail == '*':

                # Target operation
                if gate is not None:
                    allowed_sslbls_fn = resolved_avail if callable(resolved_avail) else None
                    gate_nQudits = self.processor_spec.gate_num_qudits(gn)
                    printer.log("Creating %dQ %s gate on arbitrary qudits!!" % (gate_nQudits, gn))
                    self.factories['layers'][_Lbl(gn)] = _opfactory.EmbeddingOpFactory(
                        state_space, gate, num_target_labels=gate_nQudits, allowed_sslbls_fn=allowed_sslbls_fn)
                    # add any primitive ops for this embedding factory?

                # Cloudnoise operation
                if build_cloudnoise_fn is not None:
                    cloudnoise = build_cloudnoise_fn(_Lbl(gn))
                    if cloudnoise is not None:  # build function can return None to signify no noise
                        assert (isinstance(cloudnoise, _opfactory.EmbeddingOpFactory)), \
                            ("`build_cloudnoise_fn` must return an EmbeddingOpFactory for gate %s"
                             " with arbitrary availability") % gn
                        self.factories['cloudnoise'][_Lbl(gn)] = cloudnoise

            else:  # resolved_avail is a list/tuple of available sslbls for the current gate/factory
                for inds in resolved_avail:  # inds are target qudit labels

                    #Target operation
                    if gate is not None:
                        printer.log("Creating %dQ %s gate on qudits %s!!"
                                    % ((len(qudit_labels) if inds is None else len(inds)), gn, inds))
                        assert(inds is None or _Lbl(gn, inds) not in gatedict), \
                            ("Cloudnoise models do not accept primitive-op labels, e.g. %s, in `gatedict` as this dict "
                             "specfies the ideal target gates. Perhaps make the cloudnoise depend on the target qudits "
                             "of the %s gate?") % (str(_Lbl(gn, inds)), gn)

                        if gate_is_factory:
                            self.factories['layers'][_Lbl(gn, inds)] = gate if (inds is None) else \
                                _opfactory.EmbeddedOpFactory(state_space, inds, gate)
                            # add any primitive ops for this factory?
                        else:
                            self.operation_blks['layers'][_Lbl(gn, inds)] = gate if (inds is None) else \
                                _op.EmbeddedOp(state_space, inds, gate)

                    #Cloudnoise operation
                    if build_cloudnoise_fn is not None:
                        cloudnoise = build_cloudnoise_fn(_Lbl(gn, inds))
                        if cloudnoise is not None:  # build function can return None to signify no noise
                            if isinstance(cloudnoise, _opfactory.OpFactory):
                                self.factories['cloudnoise'][_Lbl(gn, inds)] = cloudnoise
                            else:
                                self.operation_blks['cloudnoise'][_Lbl(gn, inds)] = cloudnoise

                    if build_cloudkey_fn is not None:
                        # TODO: is there any way to get a default "key", e.g. the
                        # qudits touched by the corresponding cloudnoise op?
                        # need a way to identify a clound (e.g. Gx and Gy gates on some qudit will have *same* cloud)
                        cloud_key = build_cloudkey_fn(_Lbl(gn, inds))
                        if cloud_key not in self.clouds: self.clouds[cloud_key] = []
                        self.clouds[cloud_key].append(_Lbl(gn, inds))
                    #keep track of the primitive-layer labels in each cloud,
                    # used to specify which gate parameters should be amplifiable by germs for a given cloud (?)
                    # TODO CHECK THIS

        _init_spam_layers(self, prep_layers, povm_layers)  # SPAM

        printer.log("DONE! - created Model with nqudits=%d and op-blks=" % self.state_space.num_qudits)
        for op_blk_lbl, op_blk in self.operation_blks.items():
            printer.log("  %s: %s" % (op_blk_lbl, ', '.join(map(str, op_blk.keys()))))
        self._clean_paramvec()

    def create_processor_spec(self):
        import copy as _copy
        return _copy.deepcopy(self.processor_spec)

    @property
    def clouds(self):
        """
        Returns the set of cloud-sets used when creating sequences which amplify the parameters of this model.

        Returns
        -------
        dict
        """
        return self._clouds

    def _to_nice_serialization(self):
        state = super()._to_nice_serialization()
        state.update({'processor_spec': self.processor_spec.to_nice_serialization(),
                      'error_composition_mode': self.errcomp_type,
                      })
        mmgraph = self.create_modelmember_graph()
        state['modelmembers'] = mmgraph.create_serialization_dict()
        return state

    @classmethod
    def _from_nice_serialization(cls, state):
        state_space = _statespace.StateSpace.from_nice_serialization(state['state_space'])
        #basis = _nice_serialization(state['basis'])
        modelmembers = _MMGraph.load_modelmembers_from_serialization_dict(state['modelmembers'])
        simulator = _FSim.from_nice_serialization(state['simulator'])
        layer_rules = _LayerRules.from_nice_serialization(state['layer_rules'])
        processor_spec = _ProcessorSpec.from_nice_serialization(state['processor_spec'])

        # __init__ does too much, so we need to create an alternate __init__ function here:
        mdl = cls.__new__(cls)
        mdl.processor_spec = processor_spec
        mdl.errcomp_type = state['error_composition_mode']
        _ImplicitOpModel.__init__(mdl, state_space, layer_rules, 'pp',
                                  simulator=simulator, evotype=state['evotype'])

        flags = {'auto_embed': False, 'match_parent_statespace': False,
                 'match_parent_evotype': True, 'cast_to_type': None}
        mdl.prep_blks['layers'] = _OrderedMemberDict(mdl, None, None, flags, modelmembers.get('prep_blks|layers', []))
        mdl.povm_blks['layers'] = _OrderedMemberDict(mdl, None, None, flags, modelmembers.get('povm_blks|layers', []))
        mdl.operation_blks['gates'] = _OrderedMemberDict(mdl, None, None, flags,
                                                         modelmembers.get('operation_blks|gates', []))
        mdl.operation_blks['cloudnoise'] = _OrderedMemberDict(mdl, None, None, flags,
                                                              modelmembers.get('operation_blks|cloudnoise', []))
        mdl.operation_blks['layers'] = _OrderedMemberDict(mdl, None, None, flags,
                                                          modelmembers.get('operation_blks|layers', []))
        mdl.instrument_blks['layers'] = _OrderedMemberDict(mdl, None, None, flags,
                                                           modelmembers.get('instrument_blks|layers', []))
        mdl.factories['gates'] = _OrderedMemberDict(mdl, None, None, flags, modelmembers.get('factories|gates', []))
        mdl.factories['cloudnoise'] = _OrderedMemberDict(mdl, None, None, flags,
                                                         modelmembers.get('factories|cloudnoise', []))
        mdl.factories['layers'] = _OrderedMemberDict(mdl, None, None, flags, modelmembers.get('factories|layers', []))

        mdl._clouds = _collections.OrderedDict()
        mdl._clean_paramvec()

        return mdl

class CloudNoiseLayerRules(_LayerRules):

    def __init__(self, errcomp_type, qubit_labels, implicit_idle_mode, singleq_idle_layer_labels,
                 implied_global_idle_label):
        self.qubit_labels = qubit_labels
        self.errcomp_type = errcomp_type
        self.implied_global_idle_label = implied_global_idle_label
        self.single_qubit_idle_layer_labels = singleq_idle_layer_labels
        self.implicit_idle_mode = implicit_idle_mode  # how to handle implied idles ("blanks") in circuits
        self._add_global_idle_to_all_layers = False
        self._add_padded_idle = False

        if implicit_idle_mode is None or implicit_idle_mode == "none":  # no noise on idles
            pass  # just use defaults above
        elif implicit_idle_mode == "add_global" and self.implied_global_idle_label is not None:
            self._add_global_idle_to_all_layers = True    # add global idle to all layers
        elif implicit_idle_mode == "pad_1Q" and self.single_qubit_idle_layer_labels is not None:
            self._add_padded_idle = True
        else:
            raise ValueError("Invalid `implicit_idle_mode`: '%s'" % str(implicit_idle_mode))

    def _to_nice_serialization(self):
        state = super()._to_nice_serialization()
        assert(all([len(k) == 1 for k in self.single_qubit_idle_layer_labels.keys()])), \
            "All keys of single_qubit_idle_layer_labels should be 1-tuples of a *single* sslbl!"
        state.update({'error_composition_mode': self.errcomp_type,
                      'implied_global_idle_label': (str(self.implied_global_idle_label)
                                                    if (self.implied_global_idle_label is not None) else None),
                      'single_qubit_idle_layer_labels': ({str(sslbls[0]): str(idle_lbl) for sslbls, idle_lbl
                                                          in self.single_qubit_idle_layer_labels.items()}
                                                         if self.single_qubit_idle_layer_labels is not None else None),
                      'implicit_idle_mode': self.implicit_idle_mode,
                      'qubit_labels': list(self.qubit_labels),
                      })
        return state

    @classmethod
    def _from_nice_serialization(cls, state):
        from pygsti.circuits.circuitparser import parse_label as _parse_label

        def _to_int(x):  # (same as in slowcircuitparser.py)
            return int(x) if x.isdigit() else x

        gi_label = _parse_label(state['implied_global_idle_label']) \
            if (state['implied_global_idle_label'] is not None) else None

        if state.get('single_qubit_idle_layer_labels', None) is not None:
            singleQ_idle_lbls = {(_to_int(k),): _parse_label(v)
                                 for k, v in state['single_qubit_idle_layer_labels'].items()}
        else:
            singleQ_idle_lbls = None

        qubit_labels = tuple(state['qubit_labels']) if ('qubit_labels' in state) else None

        return cls(state['error_composition_mode'], qubit_labels, state['implicit_idle_mode'],
                   singleQ_idle_lbls, gi_label)

    def prep_layer_operator(self, model, layerlbl, caches):
        """
        Create the operator corresponding to `layerlbl`.

        Parameters
        ----------
        layerlbl : Label
            A circuit layer label.

        Returns
        -------
        State
        """
        #No cache for preps
        return model.prep_blks['layers'][layerlbl]  # prep_blks['layer'] are full prep ops

    def povm_layer_operator(self, model, layerlbl, caches):
        """
        Create the operator corresponding to `layerlbl`.

        Parameters
        ----------
        layerlbl : Label
            A circuit layer label.

        Returns
        -------
        POVM or POVMEffect
        """
        # caches['povm-layers'] *are* just complete layers
        if layerlbl in caches['povm-layers']: return caches['povm-layers'][layerlbl]
        if layerlbl in model.povm_blks['layers']:
            return model.povm_blks['layers'][layerlbl]
        else:
            # See if this effect label could correspond to a *marginalized* POVM, and
            # if so, create the marginalized POVM and add its effects to model.effect_blks['layers']
            #assert(isinstance(layerlbl, _Lbl))  # Sanity check
            povmName = _ot.effect_label_to_povm(layerlbl)
            if povmName in model.povm_blks['layers']:
                # implicit creation of marginalized POVMs whereby an existing POVM name is used with sslbls that
                # are not present in the stored POVM's label.
                mpovm = _povm.MarginalizedPOVM(model.povm_blks['layers'][povmName],
                                               model.state_space, layerlbl.sslbls)  # cache in FUTURE
                mpovm_lbl = _Lbl(povmName, layerlbl.sslbls)
                caches['povm-layers'].update(mpovm.simplify_effects(mpovm_lbl))
                assert(layerlbl in caches['povm-layers']), "Failed to create marginalized effect!"
                return caches['povm-layers'][layerlbl]
            else:
                raise KeyError("Could not build povm/effect for %s!" % str(layerlbl))

    def operation_layer_operator(self, model, layerlbl, caches):
        """
        Create the operator corresponding to `layerlbl`.

        Parameters
        ----------
        layerlbl : Label
            A circuit layer label.

        Returns
        -------
        LinearOperator
        """
        #Note: cache uses 'op-layers' for *simple target* layers, not complete ones
        if layerlbl in caches['complete-layers']: return caches['complete-layers'][layerlbl]

        if isinstance(layerlbl, _CircuitLabel):
            op = self._create_op_for_circuitlabel(model, layerlbl)
            caches['complete-layers'][layerlbl] = op
            return op

        Composed = _op.ComposedOp
        ExpErrorgen = _op.ExpErrorgenOp
        Sum = _op.ComposedErrorgen
        add_global_idle = self._add_global_idle_to_all_layers
        add_padded_idle = self._add_padded_idle

        #print("DB: CloudNoiseLayerLizard building gate %s for %s w/comp-type %s" %
        #      (('matrix' if dense else 'map'), str(oplabel), self.errcomp_type) )

        components = layerlbl.components
        if len(components) == 0:
            if add_global_idle:
                if self.errcomp_type == "gates":
                    return model.operation_blks['cloudnoise'][self.implied_global_idle_label]  # idle!
                elif self.errcomp_type == "errorgens":
                    return ExpErrorgen(model.operation_blks['cloudnoise'][self.implied_global_idle_label])
                else:
                    raise ValueError("Invalid errcomp_type in CloudNoiseLayerRules: %s" % str(self.errcomp_type))
            elif add_padded_idle:
                idle_factors = [model.operation_blks['cloudnoise'][self.single_qubit_idle_layer_labels[(sslbl,)]]
                                for sslbl in self.qubit_labels]
                if self.errcomp_type == "gates":
                    ret = Composed(idle_factors, evotype=model.evotype, state_space=model.state_space)
                elif self.errcomp_type == "errorgens":
                    ret = ExpErrorgen(Sum(idle_factors, state_space=model.state_space, evotype=model.evotype))
                else:
                    raise ValueError("Invalid errcomp_type in CloudNoiseLayerRules: %s" % str(self.errcomp_type))
                model._init_virtual_obj(ret)  # so ret's gpindices get set
                return ret
            else:
                #Perfect no-noise idle
                return Composed([], evotype=model.evotype, state_space=model.state_space)  # no need to init_virtual

        #Compose target operation from layer's component labels, which correspond
        # to the perfect (embedded) target ops in op_blks
        if len(components) > 1:
            #Note: _layer_component_targetop can return `None` for a (static) identity op
            to_compose = [self._layer_component_targetop(model, l, caches['op-layers']) for l in components]
            targetOp = Composed([op for op in to_compose if op is not None],
                                evotype=model.evotype, state_space=model.state_space)
        else:
            targetOp = self._layer_component_targetop(model, components[0], caches['op-layers'])

        ops_to_compose = [targetOp] if (targetOp is not None) else []

        if self.errcomp_type == "gates":
            if add_global_idle:
                ops_to_compose.append(model.operation_blks['cloudnoise'][self.implied_global_idle_label])
            # Note: add_padded_idle handled within _layer_component_cloudnoises
            component_cloudnoise_ops = self._layer_component_cloudnoises(model, components, caches['op-cloudnoise'])
            if len(component_cloudnoise_ops) > 0:
                if len(component_cloudnoise_ops) > 1:
                    localErr = Composed(component_cloudnoise_ops,
                                        evotype=model.evotype, state_space=model.state_space)
                else:
                    localErr = component_cloudnoise_ops[0]
                ops_to_compose.append(localErr)

        elif self.errcomp_type == "errorgens":
            #We compose the target operations to create a
            # final target op, and compose this with a *single* ExpErrorgen operation which has as
            # its error generator the composition (sum) of all the factors' error gens.
            # Note: add_padded_idle handled within _layer_component_cloudnoises
            errorGens = [model.operation_blks['cloudnoise'][self.implied_global_idle_label]] if add_global_idle else []
            errorGens.extend(self._layer_component_cloudnoises(model, components, caches['op-cloudnoise']))
            if len(errorGens) > 0:
                if len(errorGens) > 1:
                    error = ExpErrorgen(Sum(errorGens, state_space=model.state_space, evotype=model.evotype))
                else:
                    error = ExpErrorgen(errorGens[0])
                ops_to_compose.append(error)
        else:
            raise ValueError("Invalid errcomp_type in CloudNoiseLayerRules: %s" % str(self.errcomp_type))

        ret = Composed(ops_to_compose, evotype=model.evotype, state_space=model.state_space) \
            if len(ops_to_compose) > 1 else ops_to_compose[0]
        model._init_virtual_obj(ret)  # so ret's gpindices get set
        caches['complete-layers'][layerlbl] = ret  # cache the final label value
        return ret

    def _layer_component_targetop(self, model, complbl, cache):
        """
        Retrieves the target- or ideal-operation portion of one component of a layer operation.

        Parameters
        ----------
        complbl : Label
            A component label of a larger layer label.

        Returns
        -------
        LinearOperator
        """
        if complbl in cache:
            return cache[complbl]  # caches['op-layers'] would hold "simplified" instrument members

        if complbl == self.implied_global_idle_label:
            # special case of the implied global idle, which give `None` instead of the
            # identity as its target operation since we don't want to include an unnecesseary idle op.
            return None

        if isinstance(complbl, _CircuitLabel):
            raise NotImplementedError("Cloud noise models cannot simulate circuits with partial-layer subcircuits.")
            # In the FUTURE, could easily implement this for errcomp_type == "gates", but it's unclear what to
            #  do for the "errorgens" case - how do we gate an error generator of an entire (mulit-layer) sub-circuit?
            # Maybe we just need to expand the label and create a composition of those layers?
        elif complbl in model.operation_blks['layers']:
            return model.operation_blks['layers'][complbl]
        else:
            return _opfactory.op_from_factories(model.factories['layers'], complbl)

    def _layer_component_cloudnoises(self, model, complbl_list, cache):
        """
        Retrieves cloud-noise portion of the components of a layer operation.

        Get any present cloudnoise ops from a list of components.  This function processes
        a list rather than an item because it's OK if some components don't have
        corresponding cloudnoise ops - we just leave those off.

        Parameters
        ----------
        complbl_list : list
            A list of circuit-layer component labels.

        Returns
        -------
        list
        """
        ret = []
        if self._add_padded_idle:
            component_sslbls = [c.sslbls for c in complbl_list]
            if None not in component_sslbls:  # sslbls == None => label covers *all* labels, no padding needed
                present_sslbl_components = set(_itertools.chain(*[sslbls for sslbls in component_sslbls]))
                absent_sslbls = [sslbl for sslbl in self.qubit_labels if (sslbl not in present_sslbl_components)]
                factors = {sslbl: model.operation_blks['cloudnoise'][self.single_qubit_idle_layer_labels[(sslbl,)]]
                           for sslbl in absent_sslbls}  # key = *lowest* (and only) sslbl
            else:
                factors = {}

            for complbl in complbl_list:
                complbl_lowest_sslbl = sorted(complbl.sslbls)[0] if (complbl.sslbls is not None) else 0
                if complbl in cache:
                    factors[complbl_lowest_sslbl] = cache[complbl]
                elif complbl in model.operation_blks['cloudnoise']:
                    factors[complbl_lowest_sslbl] = model.operation_blks['cloudnoise'][complbl]
                else:
                    try:
                        factors[complbl_lowest_sslbl] = _opfactory.op_from_factories(
                            model.factories['cloudnoise'], complbl)
                    except KeyError: pass  # OK if cloudnoise doesn't exist (means no noise)

            ret = [factors[k] for k in sorted(factors.keys())]

        else:
            for complbl in complbl_list:
                if complbl in cache:
                    ret.append(cache[complbl])  # caches['cloudnoise-layers'] would hold "simplified" instrument members
                elif complbl in model.operation_blks['cloudnoise']:
                    ret.append(model.operation_blks['cloudnoise'][complbl])
                else:
                    try:
                        ret.append(_opfactory.op_from_factories(model.factories['cloudnoise'], complbl))
                    except KeyError: pass  # OK if cloudnoise doesn't exist (means no noise)

        return ret


# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# 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.
#
# Copyright 2018-2019, Mingkun Huang
#
# 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.

import math

import torch
from numba import cuda

from nemo.collections.asr.parts.numba.rnnt_loss.utils import rnnt_helper

GPU_RNNT_THREAD_SIZE = 128

@cuda.jit(device=True, inline=True)
def logp(
    denom: torch.Tensor, acts: torch.Tensor, maxT: int, maxU: int, alphabet_size: int, mb: int, t: int, u: int, v: int
):
    """
    Compute the sum of log probability from the activation tensor and its denominator.

    Args:
        denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
            across entire vocabulary.
        acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
        maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
        maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
        alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
        mb: Batch indexer.
        t: Acoustic sequence timestep indexer.
        u: Target sequence timestep indexer.
        v: Vocabulary token indexer.

    Returns:
        The sum of logprobs[mb, t, u, v] + denom[mb, t, u]
    """
    col = (mb * maxT + t) * maxU + u
    return denom[col] + acts[col * alphabet_size + v]

@cuda.jit()
def compute_alphas_kernel(
    acts: torch.Tensor,
    denom: torch.Tensor,
    alphas: torch.Tensor,
    llForward: torch.Tensor,
    xlen: torch.Tensor,
    ylen: torch.Tensor,
    mlabels: torch.Tensor,  # [B]
    minibatch: int,
    maxT: int,
    maxU: int,
    alphabet_size: int,
    blank_: int,
):
    """
    Compute alpha (forward variable) probabilities over the transduction step.

    Args:
        acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
        denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
            across entire vocabulary.
        alphas: Zero tensor of shape [B, T, U]. Will be updated inside the kernel with the forward variable
            probabilities.
        llForward: Zero tensor of shape [B]. Represents the log-likelihood of the forward pass.
            Returned as the forward pass loss that is reduced by the optimizer.
        xlen: Vector of length B which contains the actual acoustic sequence lengths in the padded
            activation tensor.
        ylen: Vector of length B which contains the actual target sequence lengths in the padded
            activation tensor.
        mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token - usually the RNNT blank).
            The matrix contains the padded target transcription that must be predicted.
        minibatch: Int representing the batch size.
        maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
        maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
        alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
        blank_: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.

    Updates:
        Kernel inplace updates the following inputs:
        -   alphas: forward variable scores.
        -   llForward: log-likelihood of forward variable.
    """
    # // launch B blocks, each block has U threads
    b = cuda.blockIdx.x  # // batch id
    u = cuda.threadIdx.x  # label id, u
    T = xlen[b]  # select AM length of current sample
    U = ylen[b] + 1  # select target length of current sample, +1 for the blank token

    labels: torch.Tensor = mlabels[b]  # mb label start point, equivalent to mlabels + b * (maxU - 1)
    offset = b * maxT * maxU  # pointer indexing offset

    # alphas += offset # pointer offset, ignored since we explicitly add offset

    # Initilize alpha[b, t=0, u=0] for all b in B
    if u == 0:
        alphas[offset] = 0

    # sync until all alphas are initialized
    cuda.syncthreads()

    # Ordinary alpha calculations, broadcast across B=b and U=u
    # Look up forward variable calculation from rnnt_numpy.forward_pass()
    for n in range(1, T + U - 1):
        t = n - u

        if u == 0:
            # for t in range(1, T) step to initialize alphas[b, t, 0]
            if t > 0 and t < T:
                alphas[offset + t * maxU + u] = alphas[offset + (t - 1) * maxU + u] + logp(
                    denom, acts, maxT, maxU, alphabet_size, b, t - 1, 0, blank_
                )
        elif u < U:
            # for u in range(1, U) step to initialize alphas[b, 0, u]
            if t == 0:
                alphas[offset + u] = alphas[offset + u - 1] + logp(
                    denom, acts, maxT, maxU, alphabet_size, b, 0, u - 1, labels[u - 1]
                )

            # for t in range(1, T) for u in range(1, U) step to compute alphas[b, t, u]
            elif t > 0 and t < T:
                no_emit = alphas[offset + (t - 1) * maxU + u] + logp(
                    denom, acts, maxT, maxU, alphabet_size, b, t - 1, u, blank_
                )
                emit = alphas[offset + t * maxU + u - 1] + logp(
                    denom, acts, maxT, maxU, alphabet_size, b, t, u - 1, labels[u - 1]
                )

                alphas[offset + t * maxU + u] = rnnt_helper.log_sum_exp(emit, no_emit)

        # sync across all B=b and U=u
        cuda.syncthreads()

    # After final sync, alphas[b, T-1, U - 1] + logprobs[b, T-1, U-1, blank] + denom[b, T-1, U-1] gives
    # log-likelihood of forward pass.
    if u == 0:
        loglike = alphas[offset + (T - 1) * maxU + U - 1] + logp(
            denom, acts, maxT, maxU, alphabet_size, b, T - 1, U - 1, blank_
        )
        llForward[b] = loglike

@cuda.jit()
def compute_betas_kernel(
    acts: torch.Tensor,
    denom: torch.Tensor,
    betas: torch.Tensor,
    llBackward: torch.Tensor,
    xlen: torch.Tensor,
    ylen: torch.Tensor,
    mlabels: torch.Tensor,  # [B, U]
    minibatch: int,
    maxT: int,
    maxU: int,
    alphabet_size: int,
    blank_: int,
):
    """
    Compute beta (backward variable) probabilities over the transduction step.

    Args:
        acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
        denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
            across entire vocabulary.
        betas: Zero tensor of shape [B, T, U]. Will be updated inside the kernel with the backward variable
            probabilities.
        llBackward: Zero tensor of shape [B]. Represents the log-likelihood of the backward pass.
            Returned as the backward pass loss that is reduced by the optimizer.
        xlen: Vector of length B which contains the actual acoustic sequence lengths in the padded
            activation tensor.
        ylen: Vector of length B which contains the actual target sequence lengths in the padded
            activation tensor.
        mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token - usually the RNNT blank).
            The matrix contains the padded target transcription that must be predicted.
        minibatch: Int representing the batch size.
        maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
        maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
        alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
        blank_: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.

    Updates:
        Kernel inplace updates the following inputs:
        -   betas: backward variable scores.
        -   llBackward: log-likelihood of backward variable.
    """
    # // launch B blocks, each block has U threads
    b = cuda.blockIdx.x  # // batch id
    u = cuda.threadIdx.x  # label id, u
    T = xlen[b]  # select AM length of current sample
    U = ylen[b] + 1  # select target length of current sample, +1 for the blank token

    labels: torch.Tensor = mlabels[b]  # mb label start point, equivalent to mlabels + b * (maxU - 1)
    offset = b * maxT * maxU  # pointer indexing offset

    # betas += offset # pointer offset, ignored since we explicitly add offset

    # Initilize beta[b, t=T-1, u=U-1] for all b in B with log_probs[b, t=T-1, u=U-1, blank]
    if u == 0:
        betas[offset + (T - 1) * maxU + U - 1] = logp(denom, acts, maxT, maxU, alphabet_size, b, T - 1, U - 1, blank_)

    # sync until all betas are initialized
    cuda.syncthreads()

    # Ordinary beta calculations, broadcast across B=b and U=u
    # Look up backward variable calculation from rnnt_numpy.backward_pass()
    for n in range(T + U - 2, -1, -1):
        t = n - u

        if u == (U - 1):
            # for t in reversed(range(T - 1)) step to initialize betas[b, t, U-1]
            if t >= 0 and t < (T - 1):
                betas[offset + t * maxU + U - 1] = betas[offset + (t + 1) * maxU + U - 1] + logp(
                    denom, acts, maxT, maxU, alphabet_size, b, t, U - 1, blank_
                )
        elif u < U:
            if t == T - 1:
                # for u in reversed(range(U - 1)) step to initialize betas[b, T-1, u]
                betas[offset + (T - 1) * maxU + u] = betas[offset + (T - 1) * maxU + u + 1] + logp(
                    denom, acts, maxT, maxU, alphabet_size, b, T - 1, u, labels[u]
                )
            elif (t >= 0) and (t < T - 1):
                # for t in reversed(range(T - 1)) for u in reversed(range(U - 1)) step to compute betas[b, t, u]
                no_emit = betas[offset + (t + 1) * maxU + u] + logp(
                    denom, acts, maxT, maxU, alphabet_size, b, t, u, blank_
                )
                emit = betas[offset + t * maxU + u + 1] + logp(
                    denom, acts, maxT, maxU, alphabet_size, b, t, u, labels[u]
                )
                betas[offset + t * maxU + u] = rnnt_helper.log_sum_exp(emit, no_emit)

        # sync across all B=b and U=u
        cuda.syncthreads()

    # After final sync, betas[b, 0, 0] gives
    # log-likelihood of backward pass.
    if u == 0:
        llBackward[b] = betas[offset]

@cuda.jit()
def compute_grad_kernel(
    grads: torch.Tensor,
    acts: torch.Tensor,
    denom: torch.Tensor,
    alphas: torch.Tensor,
    betas: torch.Tensor,
    logll: torch.Tensor,
    xlen: torch.Tensor,
    ylen: torch.Tensor,
    mlabels: torch.Tensor,  # [B, U]
    minibatch: int,
    maxT: int,
    maxU: int,
    alphabet_size: int,
    blank_: int,
    fastemit_lambda: float,
):
    """
    Compute gradients over the transduction step.

    Args:
        grads: Zero Tensor of shape [B, T, U, V+1]. Is updated by this kernel to contain the gradients
            of this batch of samples.
        acts: Tensor of shape [B, T, U, V+1] flattened. Represents the logprobs activation tensor.
        denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor
            across entire vocabulary.
        alphas: Alpha variable, contains forward probabilities. A tensor of shape [B, T, U].
        betas: Beta varoable, contains backward probabilities. A tensor of shape [B, T, U].
        logll: Log-likelihood of the forward variable, represented as a vector of shape [B].
            Represents the log-likelihood of the forward pass.
        xlen: Vector of length B which contains the actual acoustic sequence lengths in the padded
            activation tensor.
        ylen: Vector of length B which contains the actual target sequence lengths in the padded
            activation tensor.
        mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token - usually the RNNT blank).
            The matrix contains the padded target transcription that must be predicted.
        minibatch: Int representing the batch size.
        maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
        maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
        alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
        blank_: Index of the RNNT blank token in the vocabulary. Generally the first or last token in the vocab.
        fastemit_lambda: Float scaling factor for FastEmit regularization. Refer to
            FastEmit: Low-latency Streaming ASR with Sequence-level Emission Regularization.

    Updates:
        Kernel inplace updates the following inputs:
        -   grads: Gradients with respect to the log likelihood (logll).
    """
    # Kernel call:
    # blocks_per_grid = minibatch (b) * maxT (t) * maxU (u)
    # threads_per_block = constant buffer size of parallel threads (v :: Constant)
    tid = cuda.threadIdx.x  # represents v, taking steps of some constant size
    idx = tid  # index of v < V+1; in steps of constant buffer size
    col = cuda.blockIdx.x  # represents a fused index of b * t * u

    # Decompose original indices from fused `col`
    u = col % maxU  # (b * t * u) % u = u
    bt = (col - u) // maxU  # (b * t * u - u) // U = b * t
    t = bt % maxT  # (b * t) % t = t
    mb = (bt - t) // maxT  # (b * t - t) // T = b

    # constants
    T = xlen[mb]  # select AM length of current sample
    U = ylen[mb] + 1  # select target length of current sample, +1 for the blank token
    labels: torch.Tensor = mlabels[mb]  # labels = mlabels + mb * (maxU - 1);

    # Buffered gradient calculations, broadcast across B=b, T=t and U=u, looped over V with some constant stride.
    # Look up gradient calculation from rnnt_numpy.compute_gradient()
    if t < T and u < U:
        # For cuda kernels, maximum number of threads per block is limited to some value.
        # However, it may be the case that vocabulary size is larger than this limit
        # To work around this, an arbitrary thread buffer size is chosen such that,
        # 1) each element within the thread pool operates independently of the other
        # 2) An inner while loop moves the index of each buffer element by the size of the buffer itself,
        #    such that all elements of the vocabulary size are covered in (V + 1 // thread_buffer) number of steps.
        # As such, each thread will perform the while loop at least (V + 1 // thread_buffer) number of times
        while idx < alphabet_size:
            # remember, `col` represents the tri-index [b, t, u]
            # therefore; logpk = denom[b, t, u] + acts[b, t, u, v]
            logpk = denom[col] + acts[col * alphabet_size + idx]
            # initialize the grad of the sample acts[b, t, u, v]
            grad = math.exp(alphas[col] + betas[col] + logpk - logll[mb])

            # If FastEmit regularization is enabled, calculate the gradeint of probability of predicting the next label
            # at the current timestep.
            # The formula for this is Equation 9 in https://arxiv.org/abs/2010.11148, multiplied by the log probability
            # of the current step (t, u), normalized by the total log likelihood.
            # Once the gradient has been calculated, scale it by `fastemit_lambda`, as in Equation 10.
            if fastemit_lambda > 0.0 and u < U - 1:
                fastemit_grad = fastemit_lambda * math.exp(
                    alphas[col]  # alphas(t, u)
                    + (denom[col] + acts[col * alphabet_size + labels[u]])  # y_hat(t, u)
                    + betas[col + 1]  # betas(t, u+1)
                    + logpk  # log Pr(k|t, u)
                    - logll[mb]  # total log likelihood for normalization
                )
            else:
                fastemit_grad = 0.0

            # Update the gradient of act[b, t, u, v] with the gradient from FastEmit regularization
            grad = grad + fastemit_grad

            # // grad to last blank transition
            # grad[b, T-1, U-1, v=blank] -= exp(alphas[b, t, u) + logpk - logll[b])
            if (idx == blank_) and (t == T - 1) and (u == U - 1):
                grad -= math.exp(alphas[col] + logpk - logll[mb])

            # grad of blank across t < T;
            # grad[b, t<T-1, u, v=blank] -= exp(alphas[b, t, u] + logpk - logll[b] betas[b, t + 1, u])
            if (idx == blank_) and (t < T - 1):
                grad -= math.exp(alphas[col] + logpk - logll[mb] + betas[col + maxU])

            # grad of correct token across u < U;
            # grad[b, t, u<U-1, v=label[u]] -= exp(alphas[b, t, u] + logpk - logll[b] + betas[b, t, u+1])
            # Scale the gradient by (1.0 + FastEmit_lambda) in log space, then exponentiate
            if (u < U - 1) and (idx == labels[u]):
                # exp(log(1 + fastemit_lambda) + ...) is numerically more stable than
                # multiplying (1.0 + fastemit_lambda) with result.
                grad -= math.exp(math.log1p(fastemit_lambda) + alphas[col] + logpk - logll[mb] + betas[col + 1])

            # update grads[b, t, u, v] = grad
            grads[col * alphabet_size + idx] = grad

            # update internal index through the thread_buffer;
            # until idx < V + 1, such that entire vocabulary has been updated.
            idx += GPU_RNNT_THREAD_SIZE


#! /usr/bin/env python

import numpy as np
import inspect
from pathlib import Path
# import functools

from astropy.wcs import WCS
from astropy.coordinates import SkyCoord, ICRS, AltAz, Angle
from astropy.coordinates.baseframe import frame_transform_graph
from kidsproc.kidsmodel import _Model as Model  # this allows complex inputs
from astropy.modeling import Parameter
import astropy.units as u
from astropy.modeling import models
# from astropy.table import Table
from astropy.io import fits
from astropy.table import Table, QTable

from gwcs import coordinate_frames as cf
from scipy import interpolate

from tollan.utils.log import get_logger, timeit
from tollan.utils import getobj
from tollan.utils.fmt import pformat_yaml
from tollan.utils.namespace import NamespaceMixin
from tollan.utils.registry import Registry

def _get_skyoffset_frame(c):
    """This function creates a skyoffset_frame and ensures
    the cached origin frame attribute is the correct instance.
    """
    frame = c.skyoffset_frame()
    frame_transform_graph._cached_frame_attributes['origin'] = \
        frame.frame_attributes['origin']
    return frame

class _Model(Model, NamespaceMixin):

    _namespace_type_key = 'model'

    @classmethod
    def _namespace_from_dict_op(cls, d):
        # we resolve the model here so that we can allow
        # one use only the model class name to specify a model class.
        if cls._namespace_type_key not in d:
            raise ValueError(
                    f'unable to load model: '
                    f'missing required key "{cls._namespace_type_key}"')
        model_cls = cls._resolve_model_cls(d[cls._namespace_type_key])
        return dict(d, **{cls._namespace_type_key: model_cls})

    @staticmethod
    def _resolve_model_cls(arg):
        """Return a template class specified by `arg`.

        If `arg` is string, it is resolved using `tollan.utils.getobj`.
        """
        logger = get_logger()

        _arg = arg  # for logging
        if isinstance(arg, str):
            arg = getobj(arg)
        # check if _resolve_template_cls attribute is present
        if inspect.ismodule(arg):
            raise ValueError(f"cannot resolve model class from {arg}")
        if not isinstance(arg, Model):
            raise ValueError(f"cannot resolve model class from {arg}")
        model_cls = arg
        logger.debug(
                f"resolved model {_arg} as {model_cls}")
        return model_cls

# class SkyOffsetModel(_Model):
#     """This computes the relative offsets between two sets of coordinates.

#     """
#     n_inputs = 2
#     n_outputs = 2

#     def __init__(self, ref_frame=None, *args, **kwargs):
#         super().__init__(*args, **kwargs)
#         self._t0 = t0
#         self._target = target
#         self._ref_frame = ref_frame

#     def evaluate(self, x, y):
#         return NotImplemented

#     def evaluate_at(self, ref_coord, *args):
#         """Returns the mapping pattern as evaluated at given coordinates.
#         """
#         frame = _get_skyoffset_frame(ref_coord)
#         return coord.SkyCoord(*self(*args), frame=frame).transform_to(
#                 ref_coord.frame)

#     # TODO these three method seems to be better live in a wrapper
#     # model rather than this model
#     @property
#     def t0(self):
#         return self._t0

#     @property
#     def target(self):
#         return self._target

#     input_frame = cf.CelestialFrame(
#             name='icrs',
#             reference_frame=coord.ICRS(),
#             unit=(u.deg, u.deg)
#             )

#     def __init__(self, *args, **kwargs):
#         inputs = kwargs.pop('inputs', self.input_frame.axes_names)
#         super().__init__(*args, **kwargs)
#         # self.inputs =

class SourceModel(_Model):
    """Base class for models that compute the optical signal.
    """

    _subclasses = Registry.create()

    def __init_subclass__(cls, *args, **kwargs):
        super().__init_subclass__(*args, **kwargs)
        cls._subclasses.register(cls, cls)

    def __init__(self, *args, **kwargs):
        inputs = kwargs.pop('inputs', self.input_frame.axes_names)
        outputs = ('S', )
        super().__init__(*args, **kwargs)
        self.inputs = inputs
        self.outputs = outputs

    @property
    def data(self):
        return self._data

class SourceImageModel(SourceModel):
    """
    A model given by 2-d images.
    """

    logger = get_logger()

    n_inputs = 2
    n_outputs = 1
    input_frame = cf.CelestialFrame(
            name='icrs',
            reference_frame=ICRS(),
            unit=(u.deg, u.deg)
            )

    def __init__(self, data=None, grouping=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._data = data
        self._grouping = grouping

    @classmethod
    def from_fits(cls, filepath, extname_map=None, **kwargs):
        """
        Parameters
        ----------
        filepath : str, pathlib.Path
            The path to the FITS file.

        extname_map : dict, optional
            Specify the extensions to included in the returned data.
            If None, an educated guess will be made.
        """
        filepath = Path(filepath)
        hdulist = fits.open(filepath)

        extname_hdu = dict()
        for i, hdu in enumerate(hdulist):
            cls.logger.debug('HDU {}: {}'.format(
                i, hdu.header.tostring(sep='\n')
                ))
            label = hdu.header.get('EXTNAME', i)
            extname_hdu[label] = hdu

        if extname_map is None:
            extname_map = {
                    k: k for k in extname_hdu.keys()
                    }
        cls.logger.debug(f"use extname_map: {pformat_yaml(extname_map)}")

        data = dict()
        for k, extname in extname_map.items():
            data[k] = extname_hdu[extname]
        cls.logger.debug(f'data keys: {list(data.keys())}')
        return cls(data=data, name=filepath.as_posix(), **kwargs)

    def evaluate_tod(self, tbl, lon, lat):
        """Extract flux for given array property table.

        Parameters
        ==========
        tbl : astropy.table.Table
            The array property table for mapping the data keys.

        """
        data = self._data
        grouping = self._grouping
        if grouping not in tbl.colnames:
            raise ValueError(
                    "unable to map data keys to array property table.")
        # make masks
        data_groups = []
        for g in np.unique(tbl[grouping]):
            if g not in data:
                self.logger.debug(f"group {g} not found in data")
                continue
            d = self._data[g]
            m = tbl[grouping] == g
            data_groups.append([d, m])
            self.logger.debug(f"group {g}: {m.sum()}/{len(m)}")
        self.logger.debug(f"evaluate {len(data_groups)} groups")

        s_out = np.zeros(lon.shape) << u.MJy / u.sr
        lon = lon.to_value(u.deg)
        lat = lat.to_value(u.deg)
        for d, m in data_groups:
            wcsobj = WCS(d.header)
            s_out_unit = u.Unit(d.header.get('SIGUNIT', 'adu'))
            # check lon lat range
            # because here we check longitude ranges
            # we need to take into account wrapping issue
            lon_m = lon[m, :]
            lat_m = lat[m, :]
            # s_out_m = s_out[m, :]
            # w, e = np.min(lon_m), np.max(lon_m)
            # s, n = np.min(lat_m), np.max(lat_m)
            # check pixel range
            ny, nx = d.data.shape
            # lon lat range of pixel edges
            lon_e, lat_e = wcsobj.wcs_pix2world(
                    np.array([0, 0, nx - 1, nx - 1]),
                    np.array([0, ny - 1, 0, ny - 1]),
                    0)
            xx, yy = wcsobj.wcs_world2pix(lon_e, lat_e, 0)
            # fix potential wrapping issue by check at 360 and 180 wrapping
            lon_e = Angle(lon_e << u.deg).wrap_at(360. << u.deg).degree
            lon_e_180 = Angle(lon_e << u.deg).wrap_at(180. << u.deg).degree

            w_e, e_e = np.min(lon_e), np.max(lon_e)
            w_e_180, e_e_180 = np.min(lon_e_180), np.max(lon_e_180)
            s_e, n_e = np.min(lat_e), np.max(lat_e)
            # take the one with smaller size as the coordinte
            if (e_e_180 - w_e_180) < (e_e - w_e):
                # use wrapping at 180.d
                w_e = w_e_180
                e_e = e_e_180
                lon_m = Angle(lon_m << u.deg).wrap_at(180. << u.deg).degree
                self.logger.debug("re-wrapping coordinates at 180d")
            self.logger.debug(f"data bbox: w={w_e} e={e_e} s={s_e} n={n_e}")
            self.logger.debug(f'data shape: {d.data.shape}')

            # mask to include in range lon lat
            g = (
                    (lon_m > w_e) & (lon_m < e_e)
                    & (lat_m > s_e) & (lat_m < n_e)
                    )
            self.logger.debug(f"data mask {g.sum()}/{lon_m.size}")
            if g.sum() == 0:
                continue
            # convert all lon lat to x y
            x_g, y_g = wcsobj.wcs_world2pix(lon_m[g], lat_m[g], 0)
            ii = np.rint(y_g).astype(int)
            jj = np.rint(x_g).astype(int)
            self.logger.debug(
                    f"pixel range: [{ii.min()}, {ii.max()}] "
                    f"[{jj.min()}, {jj.max()}]")
            # check ii and jj for valid pixel range
            gp = (ii >= 0) & (ii < ny) & (jj >= 0) & (jj < nx)
            # update g to include only valid pixels
            g[g] = gp
            # convert all lon lat to x y
            x_g, y_g = wcsobj.wcs_world2pix(lon_m[g], lat_m[g], 0)
            ii = np.rint(y_g).astype(int)
            jj = np.rint(x_g).astype(int)
            self.logger.debug(
                    f"pixel range updated: [{ii.min()}, {ii.max()}] "
                    f"[{jj.min()}, {jj.max()}]")
            # take values in data
            # ii, jj = np.meshgrid(
            #         np.rint(y_g).astype(int),
            #         np.rint(x_g).astype(int), indexing='ij')

            # x, y = w.wcs_world2pix(
            #         lon[m, :].ravel(), lat[m, :].ravel(), 0)
            # g = (x >= 0) & (y < imshape[0]) & (jj >=0) & (jj < imshape[1])
            # import matplotlib.pyplot as plt
            # fig, ax = plt.subplots(1, 1, figsize=(10, 10))
            # ax.scatter(ii, jj, c=d.data[ii, jj])
            # ax.set_aspect('equal')
            # plt.show()
            ig, jg = np.where(g)
            s_out[np.flatnonzero(m)[ig], jg] = d.data[ii, jj] << s_out_unit
        self.logger.debug(
                f'signal range: [{s_out.min()}, {s_out.max()}]')
        return s_out

    def evaluate(self, lon, lat):
        pass

class SourceCatalogModel(SourceModel):
    """
    A model with point sources.
    """
    logger = get_logger()

    n_inputs = 2
    n_outputs = 1
    input_frame = cf.CelestialFrame(
            name='icrs',
            reference_frame=ICRS(),
            unit=(u.deg, u.deg)
            )

    def __init__(self, pos, data, grouping=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._pos = pos
        self._data = data
        self._grouping = grouping

    def make_image_model(self, beam_models, pixscale):

        pixscale = u.pixel_scale(pixscale)
        delta_pix = (1. << u.pix).to(u.arcsec, equivalencies=pixscale)

        # convert beam_models to pixel unit
        if next(iter(beam_models.values())).x_mean.unit.is_equivalent(u.pix):
            m_beams = beam_models
        else:
            m_beams = dict()
            for k, m in beam_models.items():
                m_beams[k] = m.__class__(
                        amplitude=m.amplitude.quantity,
                        x_mean=m.x_mean.quantity.to_value(
                            u.pix, equivalencies=pixscale),
                        y_mean=m.y_mean.quantity.to_value(
                            u.pix, equivalencies=pixscale),
                        x_stddev=m.x_stddev.quantity.to_value(
                            u.pix, equivalencies=pixscale),
                        y_stddev=m.y_stddev.quantity.to_value(
                            u.pix, equivalencies=pixscale),
                        )
        # use the first pos as the reference
        ref_coord = self.pos[0]

        wcsobj = WCS(naxis=2)
        wcsobj.wcs.crpix = [1.5, 1.5]
        wcsobj.wcs.cdelt = np.array([
            -delta_pix.to_value(u.deg),
            delta_pix.to_value(u.deg),
            ])
        wcsobj.wcs.ctype = ["RA---TAN", "DEC--TAN"]
        wcsobj.wcs.crval = [ref_coord.ra.degree, ref_coord.dec.degree]

        # compute the pixel range
        x, y = wcsobj.wcs_world2pix(self.pos.ra, self.pos.dec, 0)
        l, r = np.min(x), np.max(x)
        b, t = np.min(y), np.max(y)
        w, h = r - l, t - b
        # size of the square bbox, with added padding on the edge
        s = int(np.ceil(np.max([w, h]) + 10 * np.max(
                [m.x_fwhm for m in m_beams.values()])))
        self.logger.debug(f'source image size: {s}')
        # figure out center coord
        c_ra, c_dec = wcsobj.wcs_pix2world((l + r) / 2, (b + t) / 2, 0)
        # re-center the wcs to pixel center
        wcsobj.wcs.crpix = [s / 2 + 1, s / 2 + 1]
        wcsobj.wcs.crval = c_ra, c_dec
        header = wcsobj.to_header()
        # compute the pixel positions
        x, y = wcsobj.wcs_world2pix(self.pos.ra, self.pos.dec, 0)
        assert ((x < 0) | (x > s)).sum() == 0
        assert ((y < 0) | (y > s)).sum() == 0
        # get the pixel
        # render the image
        hdus = dict()
        for k, m in m_beams.items():
            amp = (self.data[k] * m.amplitude).to(u.MJy / u.sr)
            img = np.zeros((s, s), dtype=float) << u.MJy / u.sr
            m = m.copy()
            for xx, yy, aa in zip(x, y, amp):
                m.amplitude = aa
                m.x_mean = xx
                m.y_mean = yy
                m.render(img)
            hdu = fits.ImageHDU(img.to_value(u.MJy / u.sr), header=header)
            hdu.header['SIGUNIT'] = 'MJy / sr'
            hdus[k] = hdu
            self.logger.debug('HDU {}: {}'.format(
                k, hdu.header.tostring(sep='\n')
                ))
        return SourceImageModel(
                data=hdus, grouping=self._grouping,
                name=self.name,
                )

    @property
    def pos(self):
        return self._pos

    @classmethod
    def from_file(cls, filepath, **kwargs):
        """Create instance from file path.

        Parameters
        ----------
        filepath : str, `pathlib.Path`
            The path to the catalog file.

        **kwargs
            Arguments passed to `from_table`.
        """
        tbl = Table.read(filepath, format='ascii')
        # use this to keep track of the original table filepath
        tbl.meta['_source'] = Path(filepath)
        return cls.from_table(tbl, **kwargs)

    @classmethod
    def from_table(cls, tbl, colname_map=None, **kwargs):
        """
        Parameters
        ----------
        tbl : `astropy.table.Table`
            The table containing the source catalog.

        colname_map : dict, optional
            Specify the column names to included in the returned data.
            If None, an educated guess will be made.
        """
        # TODO: add code to guess colname map
        if colname_map is None:
            colname_map = dict()
        colname_map = dict(**colname_map)
        cls.logger.debug(f"use colname_map: {pformat_yaml(colname_map)}")

        def getcol_quantity(tbl, colname, unit):
            col = tbl[colname]
            if col.unit is None and unit is not None:
                cls.logger.debug(f"assume unit {unit} for column {colname}")
                col.unit = unit
            if col.unit is None:
                return col
            return col.quantity

        # we use a qtable to hold the data internally
        data = QTable()
        for k, c in colname_map.items():
            if c.startswith('flux'):
                unit = u.mJy
            elif c in ['ra', 'dec']:
                unit = u.deg
            else:
                unit = None
            data[k] = getcol_quantity(tbl, c, unit)
        cls.logger.debug(f'data keys: {data.colnames}')

        # figure out the position
        # TODO: add handling of different coordinate system in input
        pos = SkyCoord(
                data['ra'], data['dec'],
                frame=ICRS()).transform_to(cls.input_frame.reference_frame)

        filepath = tbl.meta.get('_source', None)
        if filepath is not None:
            name = filepath.as_posix()
        else:
            name = None
        return cls(pos=pos, data=data, name=name, **kwargs)

    def evaluate(self, lon, lat):
        coo = self.input_frame.coordinates(lon, lat)
        sep = self._source_pos.separation(coo)
        return self._psfmodel(sep) * self._source_flux

class ProjModel(_Model):
    """Base class for models that transform the detector locations.
    """

    def __init__(self, t0=None, target=None, *args, **kwargs):
        inputs = kwargs.pop('inputs', self.input_frame.axes_names)
        outputs = kwargs.pop('outputs', self.output_frame.axes_names)
        kwargs.setdefault('name', self._name)
        super().__init__(*args, **kwargs)
        self.inputs = inputs
        self.outputs = outputs
        self._t0 = t0
        self._target = target

    def mpl_axes_params(self):
        return dict(aspect='equal')

class SkyMapModel(_Model):
    """A model that describes mapping patterns on the sky.

    It computes the sky coordinates as a function of the time.
    """

    n_inputs = 1
    n_outputs = 2

    def __init__(self, t0=None, target=None, ref_frame=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._t0 = t0
        self._target = target
        self._ref_frame = ref_frame

    def evaluate(self, x, y):
        return NotImplemented

    @timeit
    def evaluate_at(self, ref_coord, *args):
        """Returns the mapping pattern as evaluated at given coordinates.
        """
        frame = _get_skyoffset_frame(ref_coord)
        return SkyCoord(*self(*args), frame=frame).transform_to(
                ref_coord.frame)

    # TODO these three method seems to be better live in a wrapper
    # model rather than this model
    @property
    def t0(self):
        return self._t0

    @property
    def target(self):
        return self._target

    @property
    def ref_frame(self):
        return self._ref_frame

class RasterScanModelMeta(SkyMapModel.__class__):
    """A meta class that defines a raster scan pattern.

    This is implemented as a meta class so that we can reuse it
    in any map model of any coordinate frame.
    """

    @staticmethod
    def _evaluate(
            inst, t, length, space, n_scans, rot, speed, t_turnaround,
            return_holdflag_only=False):
        """This computes a raster patten around the origin.

        This assumes a circular turn over trajectory where the
        speed of the turn over is implicitly controlled by `t_turnaround`.
        """
        t_per_scan = length / speed

        holdflag = np.zeros(t.shape, dtype=bool)
        if n_scans == 1:
            # TODO this is ugly, will revisit later
            if return_holdflag_only:
                return holdflag
            # have to make a special case here
            x = (t / t_per_scan - 0.5) * length
            y = np.zeros(t.shape) << length.unit
        else:
            n_spaces = n_scans - 1

            # bbox_width = length
            # bbox_height = space * n_spaces
            # # (x0, y0, w, h)
            # bbox = (
            #         -bbox_width / 2., -bbox_height / 2.,
            #         bbox_width, bbox_height)
            t_per_scan = length / speed
            ratio_scan_to_si = (
                    t_per_scan / (t_turnaround + t_per_scan))
            ratio_scan_to_turnover = (t_per_scan / t_turnaround)

            # scan index
            _si = (t / (t_turnaround + t_per_scan))
            si = _si.astype(int)
            si_frac = _si - si

            # get scan and turnover part
            scan_frac = np.empty_like(si_frac)
            turnover_frac = np.empty_like(si_frac)

            turnover = si_frac > ratio_scan_to_si
            if return_holdflag_only:
                holdflag[turnover] = True
                return holdflag
            scan_frac[turnover] = 1.
            scan_frac[~turnover] = si_frac[~turnover] / ratio_scan_to_si
            turnover_frac[turnover] = si_frac[turnover] - (
                    1. - si_frac[turnover]) * ratio_scan_to_turnover
            turnover_frac[~turnover] = 0.

            x = (scan_frac - 0.5) * length
            y = (si / n_spaces - 0.5) * n_spaces * space

            # turnover part
            radius_t = space / 2
            theta_t = turnover_frac[turnover] * np.pi * u.rad
            dy = radius_t * (1 - np.cos(theta_t))
            dx = radius_t * np.sin(theta_t)
            x[turnover] = x[turnover] + dx
            y[turnover] = y[turnover] + dy
            # make continuous
            x = x * (-1) ** si

        # apply rotation
        m_rot = models.AffineTransformation2D(
            models.Rotation2D._compute_matrix(
                angle=rot.to_value('rad')) * inst.frame_unit,
            translation=(0., 0.) * inst.frame_unit)
        xx, yy = m_rot(x, y)
        return xx, yy

    def __new__(meta, name, bases, attrs):
        frame = attrs['frame']
        frame_unit = frame.unit[0]

        attrs.update(dict(
            frame_unit=frame_unit,
            length=Parameter(default=10., unit=frame_unit),
            space=Parameter(default=1., unit=frame_unit),
            n_scans=Parameter(default=10., unit=u.dimensionless_unscaled),
            rot=Parameter(default=0., unit=u.deg),
            speed=Parameter(default=1., unit=frame_unit / u.s),
            # accel=Parameter(default=1., unit=cls.frame_unit / u.s ** 2),
            t_turnaround=Parameter(default=1., unit=u.s),
            pattern='raster',
                ))

        def get_total_time(self):
            return (self.length / self.speed * self.n_scans
                    + self.t_turnaround * (self.n_scans - 1.)).to(u.s)

        attrs['get_total_time'] = get_total_time

        @timeit(name)
        def evaluate(self, t, *args, **kwargs):
            t = np.asarray(t) * t.unit
            return meta._evaluate(self, t, *args, **kwargs)

        attrs['evaluate'] = evaluate
        # TODO refactor this part
        attrs['evaluate_holdflag'] = lambda self, t: evaluate(
                self, t, self.length, self.space, self.n_scans, self.rot,
                self.speed, self.t_turnaround, return_holdflag_only=True
                )
        return super().__new__(meta, name, bases, attrs)

    def __call__(cls, *args, **kwargs):
        inst = super().__call__(*args, **kwargs)
        inst.inputs = ('t', )
        inst.outputs = cls.frame.axes_names
        return inst

class LissajousModelMeta(SkyMapModel.__class__):
    """A meta class that defines a Lissajous scan pattern.

    This is implemented as a meta class so that we can reuse it
    in any map model of any coordinate frame.
    """

    def __new__(meta, name, bases, attrs):
        frame = attrs['frame']
        frame_unit = frame.unit[0]

        attrs.update(dict(
            frame_unit=frame_unit,
            x_length=Parameter(default=10., unit=frame_unit),
            y_length=Parameter(default=10., unit=frame_unit),
            x_omega=Parameter(default=1. * u.rad / u.s),
            y_omega=Parameter(default=1. * u.rad / u.s),
            delta=Parameter(default=0., unit=u.rad),
            rot=Parameter(default=0., unit=u.deg),
            pattern='lissajous',
                ))

        def get_total_time(self):
            t_x = 2 * np.pi * u.rad / self.x_omega
            t_y = 2 * np.pi * u.rad / self.y_omega
            r = (t_y / t_x).to_value(u.dimensionless_unscaled)
            s = 100
            r = np.lcm(int(r * s), s) / s
            return (t_x * r).to(u.s)

        attrs['get_total_time'] = get_total_time

        @timeit(name)
        def evaluate(
                self, t, x_length, y_length, x_omega, y_omega, delta, rot):
            """This computes a lissajous pattern around the origin.

            """
            t = np.asarray(t) * t.unit

            x = x_length * 0.5 * np.sin(x_omega * t + delta)
            y = y_length * 0.5 * np.sin(y_omega * t)

            m_rot = models.AffineTransformation2D(
                models.Rotation2D._compute_matrix(
                    angle=rot.to_value('rad')) * self.frame_unit,
                translation=(0., 0.) * self.frame_unit)
            xx, yy = m_rot(x, y)
            return xx, yy

        attrs['evaluate'] = evaluate
        attrs['evaluate_holdflag'] = \
            lambda self, t: np.zeros(t.shape, dtype=bool)
        return super().__new__(meta, name, bases, attrs)

    def __call__(cls, *args, **kwargs):
        inst = super().__call__(*args, **kwargs)
        inst.inputs = ('t', )
        inst.outputs = cls.frame.axes_names
        return inst

class DoubleLissajousModelMeta(SkyMapModel.__class__):
    """A meta class that defines a Double Lissajous scan pattern.

    """

    @staticmethod
    def _evaluate(
            inst, t,
            x_length_0, y_length_0, x_omega_0, y_omega_0, delta_0,
            x_length_1, y_length_1, x_omega_1, y_omega_1, delta_1,
            delta, rot):
        """This computes a double lissajous pattern around the origin.

        """
        x_0 = x_length_0 * 0.5 * np.sin(x_omega_0 * t + delta + delta_0)
        y_0 = y_length_0 * 0.5 * np.sin(y_omega_0 * t + delta)
        x_1 = x_length_1 * 0.5 * np.sin(x_omega_1 * t + delta_1)
        y_1 = y_length_1 * 0.5 * np.sin(y_omega_1 * t)

        x = x_0 + x_1
        y = y_0 + y_1

        m_rot = models.AffineTransformation2D(
            models.Rotation2D._compute_matrix(
                angle=rot.to_value('rad')) * inst.frame_unit,
            translation=(0., 0.) * inst.frame_unit)
        xx, yy = m_rot(x, y)
        return xx, yy

    def __new__(meta, name, bases, attrs):
        frame = attrs['frame']
        frame_unit = frame.unit[0]

        attrs.update(dict(
            frame_unit=frame_unit,
            x_length_0=Parameter(default=10., unit=frame_unit),
            y_length_0=Parameter(default=10., unit=frame_unit),
            x_omega_0=Parameter(default=1. * u.rad / u.s),
            y_omega_0=Parameter(default=1. * u.rad / u.s),
            delta_0=Parameter(default=0., unit=u.rad),
            x_length_1=Parameter(default=5., unit=frame_unit),
            y_length_1=Parameter(default=5., unit=frame_unit),
            x_omega_1=Parameter(default=1. * u.rad / u.s),
            y_omega_1=Parameter(default=1. * u.rad / u.s),
            delta_1=Parameter(default=0., unit=u.rad),
            delta=Parameter(default=0., unit=u.rad),
            rot=Parameter(default=0., unit=u.deg),
            pattern='double_lissajous',
                ))

        def get_total_time(self):
            # make the total time the longer one among the two
            def _get_total_time(x_omega, y_omega):
                t_x = 2 * np.pi * u.rad / x_omega
                t_y = 2 * np.pi * u.rad / y_omega
                r = (t_y / t_x).to_value(u.dimensionless_unscaled)
                s = 100
                r = np.lcm(int(r * s), s) / s
                return (t_x * r).to(u.s)
            t0 = _get_total_time(self.x_omega_0, self.y_omega_0)
            t1 = _get_total_time(self.x_omega_1, self.y_omega_1)
            return t0 if t0 > t1 else t1

        attrs['get_total_time'] = get_total_time

        @timeit(name)
        def evaluate(self, t, *args, **kwargs):
            t = np.asarray(t) * t.unit
            return meta._evaluate(self, t, *args, **kwargs)

        attrs['evaluate'] = evaluate
        attrs['evaluate_holdflag'] = \
            lambda self, t: np.zeros(t.shape, dtype=bool)
        return super().__new__(meta, name, bases, attrs)

    def __call__(cls, *args, **kwargs):
        inst = super().__call__(*args, **kwargs)
        inst.inputs = ('t', )
        inst.outputs = cls.frame.axes_names
        return inst

class RastajousModelMeta(SkyMapModel.__class__):
    """A meta class that defines a Rastajous scan pattern.

    """

    def __new__(meta, name, bases, attrs):
        frame = attrs['frame']
        frame_unit = frame.unit[0]

        attrs.update(dict(
            frame_unit=frame_unit,
            length=Parameter(default=10., unit=frame_unit),
            space=Parameter(default=1., unit=frame_unit),
            n_scans=Parameter(default=10., unit=u.dimensionless_unscaled),
            rot=Parameter(default=0., unit=u.deg),
            speed=Parameter(default=1., unit=frame_unit / u.s),
            t_turnaround=Parameter(default=1., unit=u.s),
            x_length_0=Parameter(default=10., unit=frame_unit),
            y_length_0=Parameter(default=10., unit=frame_unit),
            x_omega_0=Parameter(default=1. * u.rad / u.s),
            y_omega_0=Parameter(default=1. * u.rad / u.s),
            delta_0=Parameter(default=0., unit=u.rad),
            x_length_1=Parameter(default=5., unit=frame_unit),
            y_length_1=Parameter(default=5., unit=frame_unit),
            x_omega_1=Parameter(default=1. * u.rad / u.s),
            y_omega_1=Parameter(default=1. * u.rad / u.s),
            delta_1=Parameter(default=0., unit=u.rad),
            delta=Parameter(default=0., unit=u.rad),
            pattern='rastajous',
                ))

        def get_total_time(self):
            # make the total time based on the raster
            return (self.length / self.speed * self.n_scans
                    + self.t_turnaround * (self.n_scans - 1.)).to(u.s)

        attrs['get_total_time'] = get_total_time

        @timeit(name)
        def evaluate(
                self, t,
                length, space, n_scans, rot, speed, t_turnaround,
                x_length_0, y_length_0, x_omega_0, y_omega_0, delta_0,
                x_length_1, y_length_1, x_omega_1, y_omega_1, delta_1,
                delta):
            """This computes a rastajous pattern around the origin.

            """
            t = np.asarray(t) * t.unit

            x_r, y_r = RasterScanModelMeta._evaluate(
                    inst=self, t=t, length=length, space=space,
                    n_scans=n_scans, rot=0. << u.deg,
                    speed=speed, t_turnaround=t_turnaround,
                    return_holdflag_only=False)

            x_l, y_l = DoubleLissajousModelMeta._evaluate(
                    inst=self, t=t,
                    x_length_0=x_length_0, y_length_0=y_length_0,
                    x_omega_0=x_omega_0, y_omega_0=y_omega_0, delta_0=delta_0,
                    x_length_1=x_length_1, y_length_1=y_length_1,
                    x_omega_1=x_omega_1, y_omega_1=y_omega_1, delta_1=delta_1,
                    delta=delta, rot=0. << u.deg
                    )
            x = x_r + x_l
            y = y_r + y_l
            m_rot = models.AffineTransformation2D(
                models.Rotation2D._compute_matrix(
                    angle=rot.to_value('rad')) * self.frame_unit,
                translation=(0., 0.) * self.frame_unit)
            xx, yy = m_rot(x, y)
            return xx, yy

        attrs['evaluate'] = evaluate
        # TODO use the raster hold flag
        attrs['evaluate_holdflag'] = \
            lambda self, t: RasterScanModelMeta._evaluate(
                self, t, self.length, self.space, self.n_scans, self.rot,
                self.speed, self.t_turnaround, return_holdflag_only=True
                )
        return super().__new__(meta, name, bases, attrs)

    def __call__(cls, *args, **kwargs):
        inst = super().__call__(*args, **kwargs)
        inst.inputs = ('t', )
        inst.outputs = cls.frame.axes_names
        return inst

class TrajectoryModelMeta(SkyMapModel.__class__):
    """A meta class that defines a trajectory.

    This is implemented as a meta class so that we can reuse it
    in any map model of any coordinate frame.
    """

    def __new__(meta, name, bases, attrs):
        frame = attrs['frame']
        frame_unit = frame.unit[0]

        attrs.update(dict(
            frame_unit=frame_unit,
                ))

        def get_total_time(self):
            return self._time[-1]

        attrs['get_total_time'] = get_total_time

        lon_attr, lat_attr = {
                'icrs': ('_ra', '_dec'),
                'altaz': ('_az', '_alt'),
                }[frame.name]

        @property
        def _lon(self):
            return getattr(self, lon_attr)

        attrs['_lon_attr'] = lon_attr
        attrs['_lon'] = _lon

        @property
        def _lat(self):
            return getattr(self, lat_attr)

        attrs['_lat_attr'] = lat_attr
        attrs['_lat'] = _lat

        @timeit(name)
        def evaluate(
                self, t):
            """This computes the position based on interpolation.

            """
            t = t.to_value(u.s)
            return self._lon_interp(t) << u.deg, self._lat_interp(t) << u.deg

        attrs['evaluate'] = evaluate
        attrs['evaluate_holdflag'] = \
            lambda self, t: self._holdflag_interp(t).astype(int)

        return super().__new__(meta, name, bases, attrs)

    def __call__(
            cls, *args, **kwargs):
        data = dict()
        for attr in ('time', 'ra', 'dec', 'az', 'alt', 'holdflag'):
            data[f'_{attr}'] = kwargs.pop(attr, None)
        inst = super().__call__(*args, **kwargs)
        inst.__dict__.update(data)
        inst.inputs = ('t', )
        inst.outputs = cls.frame.axes_names
        inst._lon_interp = interpolate.interp1d(
                inst._time.to_value(u.s), inst._lon.to_value(u.deg))
        inst._lat_interp = interpolate.interp1d(
                inst._time.to_value(u.s), inst._lat.to_value(u.deg))
        inst._holdflag_interp = interpolate.interp1d(
                inst._time, inst._holdflag, kind='previous')
        return inst

class SkyRasterScanModel(SkyMapModel, metaclass=RasterScanModelMeta):
    frame = cf.Frame2D(
            name='skyoffset', axes_names=('lon', 'lat'),
            unit=(u.deg, u.deg))

class SkyLissajousModel(SkyMapModel, metaclass=LissajousModelMeta):
    frame = cf.Frame2D(
            name='skyoffset', axes_names=('lon', 'lat'),
            unit=(u.deg, u.deg))

class SkyDoubleLissajousModel(SkyMapModel, metaclass=DoubleLissajousModelMeta):
    frame = cf.Frame2D(
            name='skyoffset', axes_names=('lon', 'lat'),
            unit=(u.deg, u.deg))

class SkyRastajousModel(SkyMapModel, metaclass=RastajousModelMeta):
    frame = cf.Frame2D(
            name='skyoffset', axes_names=('lon', 'lat'),
            unit=(u.deg, u.deg))

class SkyMapTrajModel(SkyMapModel):

    @timeit
    def evaluate_at(self, ref_coord, *args):
        """Returns the mapping pattern as evaluated at given coordinates.
        """
        # ref_coord is ignored in this case
        return SkyCoord(*self(*args), frame=self.ref_frame)

class SkyICRSTrajModel(SkyMapTrajModel, metaclass=TrajectoryModelMeta):
    frame = cf.CelestialFrame(
            name='icrs',
            reference_frame=ICRS(),
            unit=(u.deg, u.deg)
            )

class SkyAltAzTrajModel(SkyMapTrajModel, metaclass=TrajectoryModelMeta):
    frame = cf.CelestialFrame(
            name='altaz',
            reference_frame=AltAz(),
            unit=(u.deg, u.deg)
            )

def resolve_sky_map_ref_frame(ref_frame, observer=None, time_obs=None):
    """
    Return a frame with respect to which sky map offset model can be
    rendered.
    """
    if isinstance(ref_frame, str):
        # this is not public API so be careful for future changes.
        from astropy.coordinates.sky_coordinate_parsers import (
                _get_frame_class)
        ref_frame = _get_frame_class(ref_frame)
    if ref_frame is AltAz:
        return observer.altaz(time=time_obs)
    return ref_frame

import numpy,from scipy,import astropy,from astropy


#!/usr/bin/env python3
"""
Plots 3 components for seismograms.
"""
from argparse import ArgumentParser
from functools import partial
from glob import glob
from multiprocessing import Pool
import os

import matplotlib as mpl

mpl.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from qcore.timeseries import BBSeis, LFSeis, HFSeis, read_ascii
from visualization.util import intersection

# files that contain the 3 components (text based)
extensions = [".090", ".000", ".ver"]
BINARY_FORMATS = {"BB": BBSeis, "LF": LFSeis, "HF": HFSeis}
colours = ["black", "red", "blue", "magenta", "darkgreen", "orange"]

def load_args():
    """
    Process command line arguments.
    """
    # read
    parser = ArgumentParser(description="Plots components for seismograms. ")

    parser.add_argument(
        "--waveforms",
        help="directory to text data or binary file followed by label",
        nargs=2,
        action="append",
        required=True,
    )
    parser.add_argument(
        "--out",
        help="output folder to place plots",
        type=os.path.abspath,
        default="waveforms",
    )
    parser.add_argument(
        "--n-stations",
        help="Limit number of stations to plot (selected randomly).",
        type=int,
    )
    parser.add_argument("--stations", help="Specific stations to plot.", nargs="+")
    parser.add_argument("-v", help="verbose messages", action="store_true")
    parser.add_argument(
        "-n", "--nproc", help="number of processes to use", type=int, default=1
    )
    parser.add_argument(
        "-t", "--tmax", type=float, help="maximum duration of waveform simulation"
    )
    args = parser.parse_args()

    # validate
    for source in args.waveforms:
        if not os.path.exists(source[0]):
            parser.error(f"Cannot find waveform source: {source[0]}")

    if args.tmax is not None and args.tmax <= 0:
        parser.error("Duration -t / --tmax must be greater than 0")

    os.makedirs(args.out, exist_ok=True)

    return args

def load_location(path, verbose=False):
    """
    Return opened binary file or text directory (automatically detected).
    """
    if os.path.isfile(path):
        try:
            binary = HFSeis(path)
            if verbose:
                print(f"HF: {path}")
        except ValueError:
            # file is not an HF seis file
            binary = BBSeis(path)
            if verbose:
                print(f"BB: {path}")
    else:
        try:
            binary = LFSeis(path)
            if verbose:
                print(f"LF: {path}")
        except ValueError:
            # cannot find e3d.par... if text data
            if verbose:
                print(f"TEXT: {path}")
            return path
    return binary

def load_stations(source):
    """
    Retrieve stations for waveforms.
    """
    if type(source).__name__ != "str":
        # opened binary object
        return list(source.stations.name)

    # path to directory containing text data
    files = glob(os.path.join(source, f"*{extensions[0]}"))
    stations = list(map(lambda f: os.path.basename(f)[:-4], files))
    return stations

def plot_station(
    output,
    sources,
    labels,
    tmax,
    verbose,
    station,
):
    """Creates a waveform plot for a specific station."""

    if verbose:
        print("Plotting station: {}...".format(station))

    timeseries = []
    for source in sources:
        if type(source).__name__ != "str":
            # opened binary object
            timeline = (
                np.arange(source.nt, dtype=np.float32) * source.dt + source.start_sec
            )
            timeseries.append(np.vstack((source.vel(station).T, timeline)))
        else:
            # text directory
            meta = read_ascii(
                os.path.join(source, f"{station}{extensions[0]}"), meta=True
            )[1]
            vals = np.array(
                [
                    read_ascii(os.path.join(source, f"{station}{ext}"))
                    for ext in extensions
                ]
            )
            timeline = (
                np.arange(meta["nt"], dtype=np.float32) * meta["dt"] + meta["sec"]
            )
            timeseries.append(np.vstack((vals, timeline)))
    x_max = max([ts[-1, -1] for ts in timeseries])
    if tmax is not None:
        x_max = min(tmax, x_max)

    all_y = np.concatenate([ts[:-1] for ts in timeseries], axis=1)
    # get axis min/max
    y_min, y_max = np.min(all_y), np.max(all_y)
    y_diff = y_max - y_min
    pgvs = np.max(np.abs(all_y), axis=1)
    ppgvs = np.max(all_y, axis=1)
    npgvs = np.min(all_y, axis=1)

    scale_length = max(int(round(x_max / 25.0)) * 5, 5)

    # start plot
    f, axis = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(20, 4), dpi=96)
    f.subplots_adjust(
        left=0.08, bottom=0.12, right=0.96, top=None, wspace=0.08, hspace=0
    )
    plt.suptitle(
        station,
        fontsize=20,
        x=0.02,
        y=0.5,
        horizontalalignment="left",
        verticalalignment="center",
    )
    plt.xlim([0, x_max])

    # subplots
    for i, s in enumerate(timeseries):
        for j in range(len(extensions)):
            ax = axis[j]
            ax.set_axis_off()
            ax.set_ylim([y_min - y_diff * 0.15, y_max])

            (line,) = ax.plot(
                s[len(extensions)],
                s[j] * min(y_max / ppgvs[j], y_min / npgvs[j]),
                color=colours[i % len(colours)],
                linewidth=1,
            )
            if j == 2:
                line.set_label(labels[i])
                ax.legend()

            if i == 1 and j == 0:
                # Add scale
                ax.plot(
                    [0, scale_length],
                    [y_min - y_diff * 0.1] * 2,
                    color="black",
                    linewidth=1,
                )
                ax.text(
                    0,
                    y_min - y_diff * 0.15,
                    "0",
                    size=12,
                    verticalalignment="top",
                    horizontalalignment="center",
                )
                ax.text(
                    scale_length,
                    y_min - y_diff * 0.15,
                    str(scale_length),
                    size=12,
                    verticalalignment="top",
                    horizontalalignment="center",
                )
                ax.text(
                    scale_length / 2.0,
                    y_min - y_diff * 0.225,
                    "sec",
                    size=12,
                    verticalalignment="top",
                    horizontalalignment="center",
                )

            if i == 0:
                # Add component label
                ax.set_title(extensions[j][1:], fontsize=18)
                ax.text(x_max, y_max, "{:.1f}".format(pgvs[j]), fontsize=14)

    plt.savefig(os.path.join(output, f"{station}.png"))
    plt.close()

if __name__ == "__main__":
    args = load_args()

    # binary class object or text folder location
    sources = [load_location(source[0], args.v) for source in args.waveforms]
    # station list
    stations = intersection([load_stations(source) for source in sources])
    if args.n_stations is not None and args.n_stations < len(stations):
        # random station selection
        stations = np.random.choice(stations, args.n_stations, replace=False)
    elif args.stations is not None:
        # specific station selection
        stations = np.intersect1d(stations, args.stations)
    assert len(stations) > 0

    p = Pool(args.nproc)
    single_station = partial(
        plot_station,
        args.out,
        sources,
        [source[1] for source in args.waveforms],
        args.tmax,
        args.v,
    )
    p.map(single_station, stations)