"""This module contains the tabular Critic implementation."""
from __future__ import annotations

import math
from typing import Hashable, List, Tuple, Dict, Callable

import attr
import numpy as np
from scipy.special import logsumexp, softmax
from scipy.optimize import brentq

from improvisers.game_graph import GameGraph, Node
from improvisers.critic import Critic, Distribution, DistLike
from improvisers.explicit import ExplicitDist as Dist

oo = float('inf')

CacheKey = Tuple[Node, Hashable, float]

@attr.s(frozen=True, auto_attribs=True)
class Cache:
    data: Dict[Tuple[Node, Hashable], Tuple[float, float]] = attr.ib(
        factory=dict
    )

    def __contains__(self, key: CacheKey) -> bool:
        node, stat_key, rationality = key
        if (node, stat_key) not in self.data:
            return False
        return self.data[node, stat_key][1] == rationality

    def __getitem__(self, key: CacheKey) -> float:
        node, stat_key, _ = key
        if key not in self:
            raise ValueError(f"key: {key} not in cache.")
        return self.data[node, stat_key][0]

    def __setitem__(self, key: CacheKey, val: float) -> None:
        node, stat_key, rationality = key
        self.data[node, stat_key] = (val, rationality)

def cached_stat(func: NodeStatFunc) -> NodeStatFunc:
    """Cache node level statistics. Only once per rationality."""
    def wrap(critic: TabularCritic,
             node_dist: DistLike,
             rationality: float) -> float:
        if isinstance(node_dist, Dist):  # Don't cache distributions.
            return func(critic, node_dist, rationality)
        node = node_dist
        if (node, func, rationality) in critic.cache:
            return critic.cache[node, func, rationality]
        val = func(critic, node, rationality)
        critic.cache[node, func, rationality] = val
        return val
    return wrap

@attr.s(auto_attribs=True, frozen=True)
class TabularCritic:
    game: GameGraph
    cache: Cache = attr.ib(factory=Cache)
    _min_ent_moves: Dict[Node, List[Node]] = attr.ib(factory=dict)

    def min_ent_moves(self, node: Node) -> List[Node]:
        """Return moves which minimizes the *achievable* entropy."""
        if node in self._min_ent_moves:
            return self._min_ent_moves[node]

        moves, worst = [], oo
        for node2 in self.game.moves(node):
            entropy = self.entropy(node2, 0)
            if entropy < worst:
                moves, worst = [node2], entropy
            elif entropy == worst:
                moves.append(node2)
        self._min_ent_moves[node] = moves
        return moves

    def min_ent_move(self, node: Node, rationality: float) -> Node:
        moves = self.min_ent_moves(node)

        # Optimization. If all values are the same, the resulting
        # policy will assign same probability to transitioning to this
        # node. Commonly happens when two subtrees are equivalent.
        val0 = self.value(moves[0], 0)

        other_vals = (self.value(move, 0) for move in moves[1:])
        if all(val == val0 for val in other_vals):
            return moves[0]

        # Break ties with psat.
        # Note 1: Triggering this is fairly difficult to arrange in
        #   practice, since entropy and values both sensitive to exact
        #   model.
        # Note 2: Unlike in general min psat move case, rationality
        #   need note be updated since entropy is already matched.
        # Note 3: This step cannot be cached since psat will, in general,
        #   depend on the rationality.
        return min(moves, key=lambda n: self.psat(n, rationality))

    def min_psat_move(
            self, node: Node, rationality: float) -> Tuple[Node, float]:
        assert self.game.label(node) == 'p2'

        # Compute entropy of planned move.
        planned_move = self.min_ent_move(node, rationality)
        entropy = self.entropy(planned_move, rationality)

        # p1 will increase rationality until target entropy matched.
        def replanned_psat(move: Node) -> float:
            replanned_rationality = rationality
            if rationality < oo:  # Note: can't increase rationality past oo.
                replanned_rationality = self.match_entropy(move, entropy)
            return self.psat(move, max(replanned_rationality, 0))

        # p2 will take the minimum psat of the replanned moves.
        moves = self.game.moves(node)
        p2_move = min(moves, key=replanned_psat)

        if rationality < oo:
            rationality = self.match_entropy(p2_move, entropy)

        return p2_move, rationality

    @cached_stat
    def value(self, node: Node, rationality: float) -> float:
        label = self.game.label(node)

        if isinstance(label, bool):              # Terminal node.
            return rationality * label if rationality < oo else float(label)

        moves = list(self.game.moves(node))  # Fix order of moves.

        if label == 'p2':                        # Player 2 case.
            p2_move = self.min_ent_move(node, rationality)
            return self.value(p2_move, rationality)

        values = [self.value(move, rationality) for move in moves]

        if label == 'p1':                        # Player 1 case.
            return logsumexp(values) if rationality < oo else max(values)

        dist = label                             # Environment case.
        probs = [dist.prob(move) for move in moves]
        return np.average(values, weights=probs)

    @cached_stat
    def lsat(self, node_dist: DistLike, rationality: float) -> float:
        if isinstance(node_dist, Dist):  # Reduce dist to calls over support.
            dist = node_dist
            probs = [dist.prob(n) for n in dist.support()]
            lsats = [self.lsat(n, rationality) for n in dist.support()]
            return logsumexp(lsats, b=probs)
        node = node_dist

        label = self.game.label(node)
        if isinstance(label, bool):
            return 0 if label else -oo
        elif label == 'p2':
            # Plan against optimal deterministic p2 policy.
            p2_move, rationality = self.min_psat_move(node, rationality)

            return self.lsat(p2_move, rationality)

        node_dist2 = self.move_dist(node, rationality)
        return self.lsat(node_dist2, rationality)

    def psat(self, node: Node, rationality: float) -> float:
        sat_prob = math.exp(self.lsat(node, rationality))
        assert sat_prob < 1.2
        return min(sat_prob, 1)  # Clip at 1 due to numerics.

    def _rationality(self, node: Node, target: float,
                     match_entropy: bool = False,
                     num_iter: int = 100) -> float:
        """Bracketed search for rationality to match either psat or entropy."""
        assert target >= 0, "Entropy or probabilities must be positive."
        if not match_entropy:  # Matching psat.
            assert target <= 1, "Probabilities are less than 1!"

        stat = self.entropy if match_entropy else self.psat

        def f(coeff: float) -> float:
            return stat(node, coeff) - target

        # TODO: properly support negative rationality.
        if f(-100) > 0:
            return -100   # TODO: support -oo.
        elif f(oo) < 0:
            return oo

        top = 1
        for _ in range(num_iter):
            try:
                return brentq(f, -top, top)
            except ValueError:
                top *= 2

        return oo  # Effectively infinite.

    @cached_stat
    def match_entropy(self, node: Node, target: float) -> float:
        return self._rationality(node, target, match_entropy=True)

    @cached_stat
    def match_psat(self, node: Node, target: float) -> float:
        return self._rationality(node, target, match_entropy=False)

    @cached_stat
    def entropy(self, node_dist: DistLike, rationality: float) -> float:
        if isinstance(node_dist, Dist):  # Reduce dist to calls over support.
            dist = node_dist
            entropy = dist.entropy
            # Contribution from children. H(A[t+1:T] || S[t+1: T], S[:t]).
            for node in dist.support():
                entropy += dist.prob(node) * self.entropy(node, rationality)
            return entropy

        node = node_dist
        label = self.game.label(node)
        if isinstance(label, bool):
            return 0.0  # Terminal node has no entropy.

        node_dist2 = self.move_dist(node, rationality)
        return self.entropy(node_dist2, rationality)

    def move_dist(self, state: Node, rationality: float) -> Distribution:
        label = self.game.label(state)
        if isinstance(label, bool):
            return Dist({})
        elif label == 'p2':
            p2_move = self.min_ent_move(state, rationality)
            return Dist({p2_move: 1})  # Assume worst case.

        moves = self.game.moves(state)

        if label == 'p1':
            vals = [self.value(move, rationality) for move in moves]

            if rationality < oo:
                probs = softmax(vals)
                return Dist({move: p for move, p in zip(moves, probs)})

            # If rationality = oo, then we pick uniformly from the best move.
            optimal = max(vals)
            support = [a for a, v in zip(moves, vals) if v == optimal]
            return Dist({node: 1 / len(support) for node in support})

        return label  # Environment Case. label *is* the distribution.

    def state_dist(self, move: Node, rationality: float) -> Distribution:
        stack = [(0.0, move, rationality)]
        node2prob = {}
        while stack:
            lprob, node, rationality = stack.pop()
            label = self.game.label(node)

            if isinstance(label, bool) or label == 'p1':
                node2prob[node] = lprob
                continue
            elif label == 'p2':  # Plan against deterministic adversary.
                p2_move = self.min_ent_move(node, rationality)
                stack.append((lprob, p2_move, rationality))
                continue
            else:
                dist = label
                for node2 in dist.support():
                    lprob2 = lprob + math.log(dist.prob(node2))
                    stack.append((lprob2, node2, rationality))
        node2prob = {k: math.exp(v) for k, v in node2prob.items()}
        return Dist(node2prob)

    @staticmethod
    def from_game_graph(game_graph: GameGraph) -> Critic:
        return TabularCritic(game_graph)

NodeStatFunc = Callable[[TabularCritic, Node, float], float]

__all__ = ['TabularCritic']


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2012-2016 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Module exports
:class:`AtkinsonBoore2003SInter`,
:class:`AtkinsonBoore2003SSlab`,
:class:`AtkinsonBoore2003SInterNSHMP2008`,
:class:`AtkinsonBoore2003SSlabNSHMP2008`,
:class:`AtkinsonBoore2003SSlabCascadia`,
:class:`AtkinsonBoore2003SSlabCascadiaNSHMP2008`,
:class:`AtkinsonBoore2003SSlabJapan`
:class:`AtkinsonBoore2003SSlabJapanNSHMP2008`
"""
from __future__ import division

import numpy as np
# standard acceleration of gravity in m/s**2
from scipy.constants import g
import copy

from openquake.hazardlib.gsim.base import GMPE, CoeffsTable
from openquake.hazardlib import const
from openquake.hazardlib.imt import PGA, SA

class AtkinsonBoore2003SInter(GMPE):
    """
    Implements GMPE developed by G. M  Atkinson and D. Boore and published as
    "Empirical Ground-Motion Relations for Subduction-Zone Earthquakes and
    Their Application to Cascadia and Other Regions" (Bulletin of the
    Seismological Society of America, Volume 93, Number 4, pages 1703-1929,
    2003) and includes correction for subduction interface equations as
    described in "Erratum to 'Empirical Ground Motion Relations for
    Subduction-Zone Earthquakes and their application to Cascadia and other
    regions'", Gail M. Atkinson and David M. Boore, Volume 98, Number 5,
    pp.2567-2569, 2008. The class implements the global model but not the
    corrections for Japan/Cascadia. SA values at 4 s (not supported by the
    original equations) are obtained from mean value at 3 s divided by a
    factor equal to 0.550 (scaling factor computed in the context of the SHARE
    project and obtained as average ratio between median values at 4 and 3
    seconds as predicted by SHARE subduction GMPEs). The class implements the
    equations for 'Subduction Interface' (that's why the class name ends with
    'SInter').
    """
    #: Supported tectonic region type is subduction interface
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.SUBDUCTION_INTERFACE

    #: Supported intensity measure types are spectral acceleration,
    #: and peak ground acceleration, see table 1, page 1715
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([
        PGA,
        SA
    ])

    #: Supported intensity measure component is the random horizontal
    #component :
    #attr:`~openquake.hazardlib.const.IMC.RANDOM_HORIZONTAL`, see
    #paragraph 'Functional : Form', page 1706
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.RANDOM_HORIZONTAL

    #: Supported standard deviation types are inter-event, intra-event
    #: and total, see table 1, page 1715
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL,
        const.StdDev.INTER_EVENT,
        const.StdDev.INTRA_EVENT
    ])

    #: Required site parameters is Vs30, used to distinguish between NEHRP
    #: soil classes, see paragraph 'Functional Form', page 1706
    REQUIRES_SITES_PARAMETERS = set(('vs30', ))

    #: Required rupture parameters are magnitude and focal depth, see equation
    #: 1, page 1706
    REQUIRES_RUPTURE_PARAMETERS = set(('mag', 'hypo_depth'))

    #: Required distance measure is closest distance to rupture, see equation
    #: 1, page 1706
    REQUIRES_DISTANCES = set(('rrup', ))

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        # extracting dictionary of coefficients specific to required
        # intensity measure type.
        C = self.COEFFS_SINTER[imt]

        # cap magnitude values at 8.5, see page 1709
        mag = rup.mag
        if mag > 8.5:
            mag = 8.5

        # compute PGA on rock (needed for site amplification calculation)
        G = 10 ** (1.2 - 0.18 * mag)
        pga_rock = self._compute_mean(self.COEFFS_SINTER[PGA()], G, mag,
                                      rup.hypo_depth, dists.rrup, sites.vs30,
                                      # by passing pga_rock > 500 the soil
                                      # amplification is 0
                                      np.zeros_like(sites.vs30) + 600,
                                      PGA())
        pga_rock = 10 ** (pga_rock)

        # periods 0.4 s (2.5 Hz) and 0.2 s (5 Hz) need a special case because
        # of the erratum. SA for 0.4s and 0.2s is computed and a weighted sum
        # is returned
        if isinstance(imt, SA) and imt.period in (0.2, 0.4):
            C04 = self.COEFFS_SINTER[SA(period=0.4, damping=5.0)]
            C02 = self.COEFFS_SINTER[SA(period=0.2, damping=5.0)]
            mean04 = self._compute_mean(C04, G, mag, rup.hypo_depth,
                                        dists.rrup, sites.vs30, pga_rock, imt)
            mean02 = self._compute_mean(C02, G, mag, rup.hypo_depth,
                                        dists.rrup, sites.vs30, pga_rock, imt)

            if imt.period == 0.2:
                mean = 0.333 * mean02 + 0.667 * mean04
            else:
                mean = 0.333 * mean04 + 0.667 * mean02
        else:
            mean = self._compute_mean(C, G, mag, rup.hypo_depth, dists.rrup,
                                      sites.vs30, pga_rock, imt)

        # convert from log10 to ln and units from cm/s**2 to g
        mean = np.log((10 ** mean) * 1e-2 / g)

        if isinstance(imt, SA) and imt.period == 4.0:
            mean /= 0.550

        stddevs = self._get_stddevs(C, stddev_types, sites.vs30.shape[0])

        return mean, stddevs

    def _compute_mean(self, C, g, mag, hypo_depth, rrup, vs30, pga_rock, imt):
        """
        Compute mean according to equation 1, page 1706.
        """
        if hypo_depth > 100:
            hypo_depth = 100
        delta = 0.00724 * 10 ** (0.507 * mag)
        R = np.sqrt(rrup ** 2 + delta ** 2)

        s_amp = self._compute_soil_amplification(C, vs30, pga_rock, imt)

        mean = (
            # 1st term
            C['c1'] + C['c2'] * mag +
            # 2nd term
            C['c3'] * hypo_depth +
            # 3rd term
            C['c4'] * R -
            # 4th term
            g * np.log10(R) +
            # 5th, 6th and 7th terms
            s_amp
        )

        return mean

    @classmethod
    def _compute_soil_amplification(cls, C, vs30, pga_rock, imt):
        """
        Compute soil amplification (5th, 6th, and 7th terms in equation 1,
        page 1706).
        """
        Sc, Sd, Se = cls._compute_site_class_dummy_variables(vs30)
        sl = cls._compute_soil_linear_factor(pga_rock, imt)

        return C['c5'] * sl * Sc + C['c6'] * sl * Sd + C['c7'] * sl * Se

    @classmethod
    def _compute_site_class_dummy_variables(cls, vs30):
        """
        Compute site class dummy variables as explained in paragraph
        'Functional Form', page 1706.
        """
        Sc = np.zeros_like(vs30)
        Sd = np.zeros_like(vs30)
        Se = np.zeros_like(vs30)

        Sc[(vs30 > 360) & (vs30 <= 760)] = 1
        Sd[(vs30 >= 180) & (vs30 <= 360)] = 1
        Se[vs30 < 180] = 1

        return Sc, Sd, Se

    @classmethod
    def _compute_soil_linear_factor(cls, pga_rock, imt):
        """
        Compute soil linear factor as explained in paragraph 'Functional
        Form', page 1706.
        """
        if isinstance(imt, SA) and imt.period >= 1:
            return np.ones_like(pga_rock)
        else:
            sl = np.zeros_like(pga_rock)

            pga_between_100_500 = (pga_rock > 100) & (pga_rock < 500)
            pga_greater_equal_500 = pga_rock >= 500

            is_SA_between_05_1 = isinstance(imt, SA) and 0.5 < imt.period < 1

            is_SA_less_equal_05 = isinstance(imt, SA) and (imt.period <= 0.5)

            if is_SA_between_05_1:
                sl[pga_between_100_500] = (1 - (1. / imt.period - 1) *
                                           (pga_rock[pga_between_100_500] -
                                           100) / 400)
                sl[pga_greater_equal_500] = 1 - (1. / imt.period - 1)

            if is_SA_less_equal_05 or isinstance(imt, PGA):
                sl[pga_between_100_500] = (1 - (pga_rock[pga_between_100_500] -
                                           100) / 400)

            sl[pga_rock <= 100] = 1

            return sl

    def _get_stddevs(self, C, stddev_types, num_sites):
        """
        Return standard deviations as defined in table 1, pag 1715.
        """
        stddevs = []
        for stddev_type in stddev_types:
            assert stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
            if stddev_type == const.StdDev.TOTAL:
                stddevs.append(np.log(10 ** C['sigma']) + np.zeros(num_sites))
            elif stddev_type == const.StdDev.INTRA_EVENT:
                stddevs.append(np.log(10 ** C['s1']) + np.zeros(num_sites))
            elif stddev_type == const.StdDev.INTER_EVENT:
                stddevs.append(np.log(10 ** C['s2']) + np.zeros(num_sites))
        return stddevs

    COEFFS_SINTER = CoeffsTable(sa_damping=5, table="""\
    IMT      c1          c2          c3           c4          c5          c6          c7          sigma       s1          s2
    pga      2.991000    0.035250    0.007590    -0.002060   0.190000    0.240000    0.290000    0.230000    0.200000    0.110000
    0.0400   2.875300    0.070520    0.010040    -0.002780   0.150000    0.200000    0.200000    0.260000    0.220000    0.140000
    0.1000   2.778900    0.098410    0.009740    -0.002870   0.150000    0.230000    0.200000    0.270000    0.250000    0.100000
    0.2000   2.663800    0.123860    0.008840    -0.002800   0.150000    0.270000    0.250000    0.280000    0.250000    0.130000
    0.4000   2.524900    0.147700    0.007280    -0.002350   0.130000    0.370000    0.380000    0.290000    0.250000    0.150000
    1.0000   2.144200    0.134500    0.005210    -0.001100   0.100000    0.300000    0.550000    0.340000    0.280000    0.190000
    2.0000   2.190700    0.071480    0.002240     0.000000   0.100000    0.250000    0.400000    0.340000    0.290000    0.180000
    3.0000   2.301000    0.022370    0.000120     0.000000   0.100000    0.250000    0.360000    0.360000    0.310000    0.180000
    4.0000   2.301000    0.022370    0.000120     0.000000   0.100000    0.250000    0.360000    0.360000    0.310000    0.180000
    """)

class AtkinsonBoore2003SSlab(AtkinsonBoore2003SInter):
    """
    Implements GMPE developed by G. M  Atkinson and D. Boore and published as
    "Empirical Ground-Motion Relations for Subduction-Zone Earthquakes and
    Their Application to Cascadia and Other Regions" (Bulletin of the
    Seismological Society of America, Volume 93, Number 4, pages 1703-1929,
    2003). The class implements the global model but not the corrections for
    Japan/Cascadia. SA values at 4 s (not supported by the original equations)
    are obtained from mean value at 3 s divided by a factor equal to 0.550
    (scaling factor computed in the context of the SHARE project and obtained
    as average ratio between median values at 4 and 3 seconds as predicted by
    SHARE subduction GMPEs). The class implements the equations for 'Subduction
    IntraSlab' (that's why the class name ends with 'SSlab').
    """
    #: Supported tectonic region type is subduction interface
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.SUBDUCTION_INTRASLAB

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        # extracting dictionary of coefficients specific to required
        # intensity measure type.
        C = self.COEFFS_SSLAB[imt]

        # cap magnitude values at 8.0, see page 1709
        mag = rup.mag
        if mag >= 8.0:
            mag = 8.0

        # compute PGA on rock (needed for site amplification calculation)
        G = 10 ** (0.301 - 0.01 * mag)
        pga_rock = self._compute_mean(self.COEFFS_SSLAB[PGA()], G, mag,
                                      rup.hypo_depth, dists.rrup, sites.vs30,
                                      # by passing pga_rock > 500 the soil
                                      # amplification is 0
                                      np.zeros_like(sites.vs30) + 600,
                                      PGA())
        pga_rock = 10 ** (pga_rock)

        # compute actual mean and convert from log10 to ln and units from
        # cm/s**2 to g
        mean = self._compute_mean(C, G, mag, rup.hypo_depth, dists.rrup,
                                  sites.vs30, pga_rock, imt)
        mean = np.log((10 ** mean) * 1e-2 / g)

        if isinstance(imt, SA) and imt.period == 4.0:
            mean /= 0.550

        stddevs = self._get_stddevs(C, stddev_types, sites.vs30.shape[0])

        return mean, stddevs

    COEFFS_SSLAB = CoeffsTable(sa_damping=5, table="""\
    IMT      c1         c2         c3         c4         c5          c6         c7         sigma      s1        s2
    pga     -0.04713    0.69090    0.01130    -0.00202    0.19000    0.24000    0.29000    0.27000    0.23000    0.14000
    0.0400   0.50697    0.63273    0.01275    -0.00234    0.15000    0.20000    0.20000    0.25000    0.24000    0.07000
    0.1000   0.43928    0.66675    0.01080    -0.00219    0.15000    0.23000    0.20000    0.28000    0.27000    0.07000
    0.2000   0.51589    0.69186    0.00572    -0.00192    0.15000    0.27000    0.25000    0.28000    0.26000    0.10000
    0.4000   0.00545    0.77270    0.00173    -0.00178    0.13000    0.37000    0.38000    0.28000    0.26000    0.10000
    1.0000  -1.02133    0.87890    0.00130    -0.00173    0.10000    0.30000    0.55000    0.29000    0.27000    0.11000
    2.0000  -2.39234    0.99640    0.00364    -0.00118    0.10000    0.25000    0.40000    0.30000    0.28000    0.11000
    3.0000  -3.70012    1.11690    0.00615    -0.00045    0.10000    0.25000    0.36000    0.30000    0.29000    0.08000
    4.0000  -3.70012    1.11690    0.00615    -0.00045    0.10000    0.25000    0.36000    0.30000    0.29000    0.08000
    """)

class AtkinsonBoore2003SInterNSHMP2008(AtkinsonBoore2003SInter):
    """
    Extend :class:`AtkinsonBoore2003SInter` and introduces site amplification
    for B/C site condition and fixed rupture hypocentral depth (20 km) as
    defined by the National Seismic Hazard Mapping Project (NSHMP) for the
    2008 US hazard model

    Site amplification for B/C is triggered when vs30 > 760 and it is
    computed as site amplification for C soil scaled by a factor equal to 0.5

    The class implements the equation as coded in ``subroutine getABsub``
    in ``hazSUBXnga.f`` Fortran code available at:
    http://earthquake.usgs.gov/hazards/products/conterminous/2008/software/
    """

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.

        Call super class method with hypocentral depth fixed at 20 km
        """
        # fix hypocentral depth to 20 km. Create new rupture context to avoid
        # changing the original one
        new_rup = copy.deepcopy(rup)
        new_rup.hypo_depth = 20.

        mean, stddevs = super(AtkinsonBoore2003SInterNSHMP2008, self). \
            get_mean_and_stddevs(sites, new_rup, dists, imt, stddev_types)

        return mean, stddevs

    @classmethod
    def _compute_soil_amplification(cls, C, vs30, pga_rock, imt):
        """
        Compute soil amplification (5th, 6th, and 7th terms in equation 1,
        page 1706) and add the B/C site condition as implemented by NSHMP.
        """
        Sbc, Sc, Sd, Se = cls._compute_site_class_dummy_variables(vs30)
        sl = cls._compute_soil_linear_factor(pga_rock, imt)

        return (
            C['c5'] * sl * Sbc * 0.5 +
            C['c5'] * sl * Sc +
            C['c6'] * sl * Sd +
            C['c7'] * sl * Se
        )

    @classmethod
    def _compute_site_class_dummy_variables(cls, vs30):
        """
        Extend
        :meth:`AtkinsonBoore2003SInter._compute_site_class_dummy_variables`
        and includes dummy variable for B/C site conditions (vs30 > 760.)
        """
        Sbc = np.zeros_like(vs30)
        Sc = np.zeros_like(vs30)
        Sd = np.zeros_like(vs30)
        Se = np.zeros_like(vs30)

        Sbc[vs30 > 760.] = 1
        Sc[(vs30 > 360) & (vs30 <= 760)] = 1
        Sd[(vs30 >= 180) & (vs30 <= 360)] = 1
        Se[vs30 < 180] = 1

        return Sbc, Sc, Sd, Se

class AtkinsonBoore2003SSlabNSHMP2008(AtkinsonBoore2003SSlab):
    """
    Extend :class:`AtkinsonBoore2003SSlab` and introduces site amplification
    for B/C site condition as defined by the National Seismic Hazard Mapping
    Project (NSHMP) for the 2008 US hazard model.

    Site amplification for B/C is triggered when vs30 > 760 and it is
    computed as site amplification for C soil scaled by a factor equal to 0.5

    The class replicates the equation as coded in ``subroutine getABsub``
    in ``hazgridXnga2.f`` Fortran code available at:
    http://earthquake.usgs.gov/hazards/products/conterminous/2008/software/
    """
    @classmethod
    def _compute_soil_amplification(cls, C, vs30, pga_rock, imt):
        """
        Compute soil amplification (5th, 6th, and 7th terms in equation 1,
        page 1706) and add the B/C site condition as implemented by NSHMP.

        Call
        :meth:`AtkinsonBoore2003SInterNSHMP2008._compute_soil_amplification`
        """
        return AtkinsonBoore2003SInterNSHMP2008._compute_soil_amplification(
            C, vs30, pga_rock, imt)

    @classmethod
    def _compute_site_class_dummy_variables(cls, vs30):
        """
        Extend
        :meth:`AtkinsonBoore2003SInter._compute_site_class_dummy_variables`
        and includes dummy variable for B/C site conditions (vs30 > 760.)

        Call
        meth:`AtkinsonBoore2003SInter._compute_site_class_dummy_variables`
        """
        return AtkinsonBoore2003SInterNSHMP2008. \
            _compute_site_class_dummy_variables(vs30)

class AtkinsonBoore2003SSlabCascadia(AtkinsonBoore2003SSlab):
    """
    Extends :class:`AtkinsonBoore2003SSlab` but uses coefficients for
    Cascadia region

    The class replicates the equation as coded in ``subroutine getABsub``
    in ``hazgridXnga2.f`` Fortran code available at:
    http://earthquake.usgs.gov/hazards/products/conterminous/2008/software/
    """
    COEFFS_SSLAB = CoeffsTable(sa_damping=5, table="""\
    IMT      c1      c2         c3         c4         c5          c6         c7         sigma      s1        s2
    pga     -0.25    0.69090    0.01130    -0.00202    0.19000    0.24000    0.29000    0.27000    0.23000    0.14000
    0.0400   0.23    0.63273    0.01275    -0.00234    0.15000    0.20000    0.20000    0.25000    0.24000    0.07000
    0.1000   0.16    0.66675    0.01080    -0.00219    0.15000    0.23000    0.20000    0.28000    0.27000    0.07000
    0.2000   0.40    0.69186    0.00572    -0.00192    0.15000    0.27000    0.25000    0.28000    0.26000    0.10000
    0.4000  -0.01    0.77270    0.00173    -0.00178    0.13000    0.37000    0.38000    0.28000    0.26000    0.10000
    1.0000  -0.98    0.87890    0.00130    -0.00173    0.10000    0.30000    0.55000    0.29000    0.27000    0.11000
    2.0000  -2.25    0.99640    0.00364    -0.00118    0.10000    0.25000    0.40000    0.30000    0.28000    0.11000
    3.0000  -3.64    1.11690    0.00615    -0.00045    0.10000    0.25000    0.36000    0.30000    0.29000    0.08000
    """)

class AtkinsonBoore2003SSlabCascadiaNSHMP2008(AtkinsonBoore2003SSlabCascadia,
                                              AtkinsonBoore2003SSlabNSHMP2008):
    """
    Combines :class:`AtkinsonBoore2003SSlabNSHMP2008` for NSHMP site
    amplification with :class:`AtkinsonBoore2003SSlabCascadia` for Cascadia.
    """
    pass

class AtkinsonBoore2003SSlabJapan(AtkinsonBoore2003SSlab):
    """
    Extends :class:`AtkinsonBoore2003SSlab` but substitutes values for c1 from
    Table 3 which incorporate correction factors for Japan.
    """

    COEFFS_SSLAB = CoeffsTable(sa_damping=5, table="""\
    IMT      c1      c2         c3         c4         c5          c6         c7         sigma      s1        s2
    pga      0.10    0.69090    0.01130    -0.00202    0.19000    0.24000    0.29000    0.27000    0.23000    0.14000
    0.0400   0.68    0.63273    0.01275    -0.00234    0.15000    0.20000    0.20000    0.25000    0.24000    0.07000
    0.1000   0.61    0.66675    0.01080    -0.00219    0.15000    0.23000    0.20000    0.28000    0.27000    0.07000
    0.2000   0.70    0.69186    0.00572    -0.00192    0.15000    0.27000    0.25000    0.28000    0.26000    0.10000
    0.4000   0.07    0.77270    0.00173    -0.00178    0.13000    0.37000    0.38000    0.28000    0.26000    0.10000
    1.0000  -0.98    0.87890    0.00130    -0.00173    0.10000    0.30000    0.55000    0.29000    0.27000    0.11000
    2.0000  -2.44    0.99640    0.00364    -0.00118    0.10000    0.25000    0.40000    0.30000    0.28000    0.11000
    3.0000  -3.73    1.11690    0.00615    -0.00045    0.10000    0.25000    0.36000    0.30000    0.29000    0.08000
    """)

class AtkinsonBoore2003SSlabJapanNSHMP2008(AtkinsonBoore2003SSlabJapan,
                                           AtkinsonBoore2003SSlabNSHMP2008):
    """
    Combines :class:`AtkinsonBoore2003SSlabNSHMP2008` for NSHMP site
    amplification with :class:`AtkinsonBoore2003SSlabJapan` for Japan.

    Validation test vector was generated by applying increments in columns 1
    and 2 of Table 3 to test vector for
    AtkinsonBoore2003SSlabCascadiaNSHMP2008.
    """
    pass


""" graph functions that depend on resonance structure

BEFORE ADDING ANYTHING, SEE IMPORT HIERARCHY IN __init__.py!!!!
"""
import itertools
import functools
import numpy
from automol.util import dict_
from automol.graph.base._algo import atom_groups
from automol.graph.base._algo import isomorphism
from automol.graph.base._algo import connected_components
from automol.graph.base._core import atoms
from automol.graph.base._core import atom_keys
from automol.graph.base._core import bond_keys
from automol.graph.base._core import bond_orders
from automol.graph.base._core import set_bond_orders
from automol.graph.base._core import atom_hybridizations
from automol.graph.base._core import atom_unsaturated_valences
from automol.graph.base._core import remove_atoms
from automol.graph.base._core import remove_bonds
from automol.graph.base._core import without_dummy_atoms
from automol.graph.base._core import without_bond_orders
from automol.graph.base._core import without_dummy_bonds
from automol.graph.base._core import without_fractional_bonds
from automol.graph.base._core import maximum_spin_multiplicity
from automol.graph.base._core import implicit
from automol.graph.base._core import subgraph
from automol.graph.base._core import atoms_bond_keys
from automol.graph.base._core import atoms_neighbor_atom_keys
from automol.graph.base._core import dummy_atoms_neighbor_atom_key

# # core functions
def dominant_resonance(rgr):
    """ *a* dominant (minimum spin/maximum pi) resonance graph
    """
    return next(iter(dominant_resonances(rgr)))

def dominant_resonances(rgr):
    """ all dominant (minimum spin/maximum pi) resonance graphs
    """
    rgr = without_fractional_bonds(rgr)
    rgrs = resonances(rgr)
    mult_min = min(map(maximum_spin_multiplicity, rgrs))
    dom_rgrs = tuple(
        rgr for rgr in rgrs if maximum_spin_multiplicity(rgr) == mult_min)
    return dom_rgrs

def resonances(rgr):
    """ all resonance graphs with this connectivity
    """
    return subresonances(without_bond_orders(rgr))

def subresonances(rgr):
    """ this connected graph and its lower-spin (more pi-bonded) resonances
    """
    rgr = without_fractional_bonds(rgr)
    # get the bond capacities (room for increasing bond order), filtering out
    # the negative ones to avoid complications with hypervalent atoms in TSs
    bnd_cap_dct = dict_.by_value(_bond_capacities(rgr), lambda x: x > 0)

    ret_rgrs = []
    if bnd_cap_dct:
        bnd_keys, bnd_caps = zip(*bnd_cap_dct.items())
        atm_keys = list(functools.reduce(frozenset.union, bnd_keys))

        # Loop over all possible combinations of bond order increments (amounts
        # by which to increase the bond order), filtering out combinations that
        # exceed the valences of the atoms involved.
        # (Note that we are only testing the bonds with available pi electrons,
        # so this is compatible with having hypervalent atoms elsewhere in the
        # molecule)
        bnd_ord_inc_ranges = [range(bnd_cap+1) for bnd_cap in bnd_caps]
        for bnd_ord_incs in itertools.product(*bnd_ord_inc_ranges):
            bnd_ord_inc_dct = dict(zip(bnd_keys, bnd_ord_incs))
            ret_rgr = _add_pi_bonds(rgr, bnd_ord_inc_dct)

            max_bnd_ord = max(bond_orders(ret_rgr).values())

            atm_unsat_vlcs = dict_.values_by_key(
                atom_unsaturated_valences(ret_rgr), atm_keys)

            if not any(atm_unsat_vlc < 0 for atm_unsat_vlc in atm_unsat_vlcs):
                if max_bnd_ord < 4:
                    ret_rgrs.append(ret_rgr)

    if not ret_rgrs:
        ret_rgrs = (rgr,)
    else:
        ret_rgrs = tuple(ret_rgrs)

    return ret_rgrs

def resonance_dominant_bond_orders(rgr):
    """ resonance-dominant bond orders, by bond
    """
    rgr = without_fractional_bonds(rgr)
    bnd_keys = list(bond_keys(rgr))
    bnd_ords_by_res = [
        dict_.values_by_key(bond_orders(dom_rgr), bnd_keys)
        for dom_rgr in dominant_resonances(rgr)]
    bnd_ords_lst = list(map(frozenset, zip(*bnd_ords_by_res)))
    bnd_dom_res_ords_dct = dict(zip(bnd_keys, bnd_ords_lst))
    return bnd_dom_res_ords_dct

def one_resonance_dominant_bond_orders(rgr):
    """ resonance-dominant bond orders, by bond
    """
    rgr = without_fractional_bonds(rgr)
    bnd_keys = list(bond_keys(rgr))
    bnd_ords_by_res = [
        dict_.values_by_key(bond_orders(dom_rgr), bnd_keys)
        for dom_rgr in dominant_resonances(rgr)]
    first_bnd_ords = [bnd_ords_by_res[0]]
    bnd_ords_lst = list(map(frozenset, zip(*first_bnd_ords)))
    bnd_dom_res_ords_dct = dict(zip(bnd_keys, bnd_ords_lst))
    return bnd_dom_res_ords_dct

def resonance_avg_bond_orders(rgr):
    """ resonance-dominant bond orders, by bond
    """
    rgr = without_fractional_bonds(rgr)
    bnd_keys = list(bond_keys(rgr))
    bnd_ords_by_res = [
        dict_.values_by_key(bond_orders(dom_rgr), bnd_keys)
        for dom_rgr in dominant_resonances(rgr)]
    nres = len(bnd_ords_by_res)
    bnd_ords_lst = zip(*bnd_ords_by_res)
    avg_bnd_ord_lst = [sum(bnd_ords)/nres for bnd_ords in bnd_ords_lst]
    avg_bnd_ord_dct = dict(zip(bnd_keys, avg_bnd_ord_lst))
    return avg_bnd_ord_dct

# # derived properties
def linear_atom_keys(rgr, dummy=True):
    """ atoms forming linear bonds, based on their hybridization

    :param rgr: the graph
    :param dummy: whether or not to consider atoms connected to dummy atoms as
        linear, if different from what would be predicted based on their
        hybridization
    :returns: the linear atom keys
    :rtype: tuple[int]
    """
    rgr = without_fractional_bonds(rgr)
    atm_hyb_dct = resonance_dominant_atom_hybridizations(implicit(rgr))
    lin_atm_keys = set(dict_.keys_by_value(atm_hyb_dct, lambda x: x == 1))

    if dummy:
        dum_ngb_key_dct = dummy_atoms_neighbor_atom_key(rgr)
        lin_atm_keys |= set(dum_ngb_key_dct.values())

    lin_atm_keys = tuple(sorted(lin_atm_keys))
    return lin_atm_keys

def linear_segments_atom_keys(gra, lin_keys=None):
    """ atom keys for linear segments in the graph
    """
    ngb_keys_dct = atoms_neighbor_atom_keys(without_dummy_atoms(gra))

    lin_keys = (linear_atom_keys(gra, dummy=True)
                if lin_keys is None else lin_keys)

    lin_keys = [k for k in lin_keys if len(ngb_keys_dct[k]) <= 2]

    lin_segs = connected_components(subgraph(gra, lin_keys))

    lin_keys_lst = []
    for lin_seg in lin_segs:
        lin_seg_keys = atom_keys(lin_seg)
        if len(lin_seg_keys) == 1:
            key, = lin_seg_keys
            lin_keys_lst.append([key])
        else:
            end_key1, end_key2 = sorted([
                key for key, ngb_keys in
                atoms_neighbor_atom_keys(lin_seg).items()
                if len(ngb_keys) == 1])
            ngb_keys_dct = atoms_neighbor_atom_keys(lin_seg)

            key = None
            keys = [end_key1]
            while key != end_key2:
                key, = ngb_keys_dct[keys[-1]] - set(keys)
                keys.append(key)
            lin_keys_lst.append(keys)

    lin_keys_lst = tuple(map(tuple, lin_keys_lst))
    return lin_keys_lst

def radical_atom_keys(gra, single_res=False, min_valence=1.):
    """ Radical atom keys for this molecular graph

    Radical atoms are based on the lowest-spin resonance structures for this
    graph. If the `single_res` flag is set, a single low-spin resonance
    structure will be chosen when there are multiple such structures.

    This function should eventually replace both
    `resonance_dominant_radical_atom_keys` and
    `sing_res_dom_radical_atom_keys` for a more user-friendly interface.

    Note that this function ignores the bond orders in `gra`. If you wish to
    identify radical atom keys based on the bond orders in `gra`, this can be
    done by using the `atom_unsaturated_valences` function.

    :param gra: the molecular graph
    :param single_res: only include radical keys for a single (arbitrary)
        resonance structure, or include all atoms that are radicals in any of
        the low-spin resonance structures?
    :type single_res: bool
    :param min_valence: optionally, specify that only sites with at least a
        certain number of radical electrons be included
    :type min_valence: int
    :returns: the radical atom keys
    :rtype: frozenset[int]

    """
    gra = without_fractional_bonds(gra)
    atm_keys = list(atom_keys(gra))

    if single_res:
        atm_rad_vlcs = dict_.values_by_key(
            atom_unsaturated_valences(dominant_resonance(gra)), atm_keys)
    else:
        atm_rad_vlcs_by_res = [
            dict_.values_by_key(atom_unsaturated_valences(dom_gra), atm_keys)
            for dom_gra in dominant_resonances(gra)]
        atm_rad_vlcs = [
            max(rad_vlcs) for rad_vlcs in zip(*atm_rad_vlcs_by_res)]

    atm_rad_keys = frozenset(atm_key for atm_key, atm_rad_vlc
                             in zip(atm_keys, atm_rad_vlcs)
                             if atm_rad_vlc >= min_valence)
    return atm_rad_keys

def has_separated_radical_sites(gra):
    """ does this radical have two or more separated radical sites?

    The identification is performed based on one of its lowest-spin resonance
    structures. It shouldn't matter which of the low-spin resonance structures
    is used -- if one of them has separated radical sites, they all should.

    This identifies polyradical molecules, but excludes things like carbenes.

    :param gra: the graph
    :returns: True if it has, False if not
    :rtype: bool
    """
    rad_atm_keys = radical_atom_keys(gra, single_res=True)
    return len(rad_atm_keys) > 1

def nonresonant_radical_atom_keys(rgr):
    """ keys for radical atoms that are not in resonance
    """
    rgr = without_fractional_bonds(rgr)
    atm_keys = list(atom_keys(rgr))
    atm_rad_vlcs_by_res = [
        dict_.values_by_key(atom_unsaturated_valences(dom_rgr), atm_keys)
        for dom_rgr in dominant_resonances(rgr)]
    atm_rad_vlcs = [min(rad_vlcs) for rad_vlcs in zip(*atm_rad_vlcs_by_res)]
    atm_rad_keys = frozenset(atm_key for atm_key, atm_rad_vlc
                             in zip(atm_keys, atm_rad_vlcs) if atm_rad_vlc)
    return atm_rad_keys

def sigma_radical_atom_keys(rgr):
    """ keys for sigma radical atoms
    """
    rgr = without_fractional_bonds(rgr)
    atm_rad_keys = nonresonant_radical_atom_keys(rgr)
    bnd_ords_dct = resonance_dominant_bond_orders(rgr)
    atm_bnd_keys_dct = atoms_bond_keys(rgr)
    atm_sig_keys = []
    for atm_key in atm_rad_keys:
        for bnd_key in atm_bnd_keys_dct[atm_key]:
            if 3 in bnd_ords_dct[bnd_key]:
                atm_sig_keys.append(atm_key)
                break
    atm_sig_keys = frozenset(atm_sig_keys)
    return atm_sig_keys

def resonance_dominant_radical_atom_keys(rgr):
    """ resonance-dominant radical atom keys

    TODO: DEPRECATE

    (keys of resonance-dominant radical sites)
    """
    rgr = without_fractional_bonds(rgr)
    atm_keys = list(atom_keys(rgr))
    atm_rad_vlcs_by_res = [
        dict_.values_by_key(atom_unsaturated_valences(dom_rgr), atm_keys)
        for dom_rgr in dominant_resonances(rgr)]
    atm_rad_vlcs = [max(rad_vlcs) for rad_vlcs in zip(*atm_rad_vlcs_by_res)]
    atm_rad_keys = frozenset(atm_key for atm_key, atm_rad_vlc
                             in zip(atm_keys, atm_rad_vlcs) if atm_rad_vlc)
    return atm_rad_keys

def sing_res_dom_radical_atom_keys(rgr):
    """ resonance-dominant radical atom keys,for one resonance

    TODO: DEPRECATE
    """
    rgr = without_fractional_bonds(rgr)
    atm_keys = list(atom_keys(rgr))
    atm_rad_vlcs_by_res = [
        dict_.values_by_key(atom_unsaturated_valences(dom_rgr), atm_keys)
        for dom_rgr in dominant_resonances(rgr)]
    first_atm_rad_val = [atm_rad_vlcs_by_res[0]]
    atm_rad_vlcs = [max(rad_vlcs) for rad_vlcs in zip(*first_atm_rad_val)]
    atm_rad_keys = frozenset(atm_key for atm_key, atm_rad_vlc
                             in zip(atm_keys, atm_rad_vlcs) if atm_rad_vlc)
    return atm_rad_keys

def radical_groups(gra):
    """ returns a list of lists of groups attached each radical
    """
    gra = without_fractional_bonds(gra)

    groups = []
    rads = sing_res_dom_radical_atom_keys(gra)
    for rad in rads:
        groups.append(atom_groups(gra, rad))
    return groups

def radical_group_dct(gra):
    """ return a dictionary of lists of groups attached each radical
    """
    gra = without_fractional_bonds(gra)

    groups = {}
    rads = list(sing_res_dom_radical_atom_keys(gra))
    atms = atoms(gra)
    for rad in rads:
        key = atms[rad][0]
        if key in groups:
            groups[atms[rad][0]] += atom_groups(gra, rad)
        else:
            groups[atms[rad][0]] = atom_groups(gra, rad)

    return groups

def radical_dissociation_prods(gra, pgra1):
    """ given a dissociation product, determine the other product
    """
    gra = without_fractional_bonds(gra)

    pgra2 = None
    rads = sing_res_dom_radical_atom_keys(gra)
    adj_atms = atoms_neighbor_atom_keys(gra)
    # adj_idxs = tuple(adj_atms[rad] for rad in rads)
    for rad in rads:
        for adj in adj_atms[rad]:
            for group in atom_groups(gra, adj, stereo=False):
                if isomorphism(group, pgra1, backbone_only=True):
                    pgra2 = remove_atoms(gra, atom_keys(group))
                    # pgra2 = remove_bonds(pgra2, bond_keys(group))
                    if bond_keys(group) in pgra2:
                        pgra2 = remove_bonds(pgra2, bond_keys(group))
    return (pgra1, pgra2)

def sp2_bond_keys(gra):
    """ determine the sp2 bonds in this graph
    """
    gra = without_bond_orders(gra)
    bnd_keys = dict_.keys_by_value(
        resonance_dominant_bond_orders(gra), lambda x: 2 in x)

    # make sure both ends are sp^2 (excludes cumulenes)
    atm_hyb_dct = resonance_dominant_atom_hybridizations(gra)
    sp2_atm_keys = dict_.keys_by_value(atm_hyb_dct, lambda x: x == 2)
    bnd_keys = frozenset({bnd_key for bnd_key in bnd_keys
                          if bnd_key <= sp2_atm_keys})
    return bnd_keys

def resonance_dominant_atom_hybridizations(rgr):
    """ resonance-dominant atom hybridizations, by atom
    """
    rgr = without_fractional_bonds(rgr)
    atm_keys = list(atom_keys(rgr))
    atm_hybs_by_res = [
        dict_.values_by_key(atom_hybridizations(dom_rgr), atm_keys)
        for dom_rgr in dominant_resonances(rgr)]
    atm_hybs = [min(hybs) for hybs in zip(*atm_hybs_by_res)]
    atm_hyb_dct = dict(zip(atm_keys, atm_hybs))
    return atm_hyb_dct

def resonance_dominant_atom_centered_cumulene_keys(rgr):
    """ resonance dominant keys for atom-centered cumulenes

    the bond-centered cumulenes are described by
        (frozenset({end_atm_key1, end_atm_key2}), cent_atm_key)
    where the first pair contains the sp2 atoms at the cumulene ends and
    `cent_atm_key` is the key of the central atom
    """
    rgr = without_fractional_bonds(rgr)
    cum_chains = _cumulene_chains(rgr)
    cum_keys = set()
    for cum_chain in cum_chains:
        size = len(cum_chain)
        if size % 2 == 1:
            cum_keys.add(
                (frozenset({cum_chain[0], cum_chain[-1]}),
                 cum_chain[size // 2])
            )
    cum_keys = frozenset(cum_keys)
    return cum_keys

def resonance_dominant_bond_centered_cumulene_keys(rgr):
    """ resonance dominant keys for bond-centered cumulenes

    the bond-centered cumulenes are described by
        (frozenset({end_atm_key1, end_atm_key2}),
         frozenset({cent_atm_key1, cent_atm_key2}))
    where the first pair contains the sp2 atoms at the cumulene ends and the
    second pair is the bond key for the central bond
    """
    rgr = without_fractional_bonds(rgr)
    cum_chains = _cumulene_chains(rgr)
    cum_keys = set()
    for cum_chain in cum_chains:
        size = len(cum_chain)
        if size % 2 == 0:
            cum_keys.add(
                (frozenset({cum_chain[0], cum_chain[-1]}),
                 frozenset({cum_chain[size // 2 - 1], cum_chain[size // 2]}))
            )
    cum_keys = frozenset(cum_keys)
    return cum_keys

# # helpers
def _bond_capacities(rgr):
    """ the number of electron pairs available for further pi-bonding, by bond
    """
    rgr = without_dummy_bonds(rgr)
    atm_unsat_vlc_dct = atom_unsaturated_valences(rgr)

    def _pi_capacities(bnd_key):
        return min(map(atm_unsat_vlc_dct.__getitem__, bnd_key))

    bnd_keys = list(bond_keys(rgr))
    bnd_caps = tuple(map(_pi_capacities, bnd_keys))
    bnd_cap_dct = dict(zip(bnd_keys, bnd_caps))
    return bnd_cap_dct

def _add_pi_bonds(rgr, bnd_ord_inc_dct):
    """ add pi bonds to this graph
    """
    bnd_keys = bond_keys(rgr)
    assert set(bnd_ord_inc_dct.keys()) <= bnd_keys

    bnd_keys = list(bnd_keys)
    bnd_ords = dict_.values_by_key(bond_orders(rgr), bnd_keys)
    bnd_ord_incs = dict_.values_by_key(bnd_ord_inc_dct, bnd_keys, fill_val=0)
    new_bnd_ords = numpy.add(bnd_ords, bnd_ord_incs)
    bnd_ord_dct = dict(zip(bnd_keys, new_bnd_ords))
    rgr = set_bond_orders(rgr, bnd_ord_dct)
    return rgr

def _cumulene_chains(rgr):
    atm_hyb_dct = resonance_dominant_atom_hybridizations(rgr)
    sp1_atm_keys = dict_.keys_by_value(atm_hyb_dct, lambda x: x == 1)
    sp2_atm_keys = dict_.keys_by_value(atm_hyb_dct, lambda x: x == 2)

    atm_ngb_keys_dct = atoms_neighbor_atom_keys(rgr)

    def _cumulene_chain(chain):
        ret = None
        atm_key = chain[-1]
        next_atm_keys = atm_ngb_keys_dct[atm_key] - {chain[-2]}
        if next_atm_keys:
            assert len(next_atm_keys) == 1
            next_atm_key, = next_atm_keys
            if next_atm_key in sp1_atm_keys:
                chain.append(next_atm_key)
                ret = _cumulene_chain(chain)
            elif next_atm_key in sp2_atm_keys:
                chain.append(next_atm_key)
                ret = chain
        return ret

    cum_chains = []
    for atm_key in sp2_atm_keys:
        sp1_atm_ngb_keys = atm_ngb_keys_dct[atm_key] & sp1_atm_keys
        chains = [[atm_key, atm_ngb_key] for atm_ngb_key in sp1_atm_ngb_keys]
        for chain in chains:
            cum_chain = _cumulene_chain(chain)
            if cum_chain is not None:
                cum_chains.append(cum_chain)

    cum_chains = tuple(map(tuple, cum_chains))
    return cum_chains


import numpy as np
import pandas as pd
import os

class HemibrainALSim:
    def __init__(self, file_folder = 'data/', dfpath = 'adult_hallem06.csv', threshold=5., return_flycircuit=True):
        """Creates the AL with a simple model for Hemibrain and FlyCircuit.

        # Arguments:
            file_folder (str): The folder with the matrices to load for presynaptically acting and postsynaptically acting LNs.
            dfpath (str): The path to the csv file to be used for loading the affinity values.
        """
        df = pd.read_csv(os.path.join(file_folder, dfpath) )
        all_odorant_names = list(df.iloc[:,0])

        self.preLN_field = np.load(os.path.join(file_folder, 'preLN_field.npy'))
        self.postLN_field = np.load(os.path.join(file_folder, 'postLN_field.npy'))
        self.preLN_field_a = np.load(os.path.join(file_folder, 'preLN_field_a.npy')) # OSN-to-LN
        self.preLN_field_b = np.load(os.path.join(file_folder, 'preLN_field_b.npy')) # LN-to-PN

        glom_names = ['VP3',
                      'VC3l',
                      'VP5',
                      'DL2v',
                      'V',
                      'VL2a',
                      'VC5',
                      'DM4',
                      'DM3',
                      'DA4m',
                      'VP2',
                      'VP1l',
                      'DL2d',
                      'VP1m',
                      'DM5',
                      'DC4',
                      'DA1',
                      'VA3',
                      'VM2',
                      'D',
                      'VL2p',
                      'VM5d',
                      'VA1v',
                      'DL3',
                      'VA7m',
                      'DA2',
                      'VM7d',
                      'VC3m',
                      'VM1',
                      'VM4',
                      'VA4',
                      'DL1',
                      'DC1',
                      'DA4l',
                      'DP1m',
                      'VA2',
                      'VA1d',
                      'DM2',
                      'DP1l',
                      'VC4',
                      'VM7v',
                      'VA5',
                      'VA6',
                      'DC2',
                      'DM1',
                      'DL4',
                      'VA7l',
                      'DM6',
                      'VM3',
                      'VM5v',
                      'VC2',
                      'DA3',
                      'DC3',
                      'DL5',
                      'VC1',
                      'VL1']

        GL_to_OR = {
                'D': {'receptors': ['OR69a','OR69b'], 'name': 'ab9', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DA1': {'receptors': ['OR67d'], 'name': 'at1A', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'DA2': {'receptors': ['OR33a','OR56a'], 'name': 'ab4B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DA3': {'receptors': ['OR23a'], 'name': 'at2B', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'DA4l': {'receptors': ['OR43a'], 'name': 'at3', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'DA4m': {'receptors': ['OR2a'], 'name': 'at3', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'DC1': {'receptors': ['OR19a','OR19b'], 'name': 'at3A', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'DC2': {'receptors': ['OR13a'], 'name': 'ab6A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DC3': {'receptors': ['OR83c'], 'name': 'at2A', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'DC4': {'receptors': ['IR64a'], 'name': 'Sac III', 'co-receptors': ['IR8a'], 'sensillum': 'sacculus', 'sensillum location': 'antenna'},
                'DL1': {'receptors': ['OR10a','GR10a'], 'name': 'ab1D', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DL2d': {'receptors': ['IR75b'], 'name': 'ac3A', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'DL2v': {'receptors': ['IR75c'], 'name': 'ac3A', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'DL3': {'receptors': ['OR65a','OR65b','OR65c'], 'name': 'at4B', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'DL4': {'receptors': ['OR49a','OR85f'], 'name': 'ab10B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DL5': {'receptors': ['OR7a'], 'name': 'ab4A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DM1': {'receptors': ['OR42b'], 'name': 'ab1A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DM2': {'receptors': ['OR22a','OR22b'], 'name': 'ab3A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DM3': {'receptors': ['OR47a','OR33b'], 'name': 'ab5B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DM4': {'receptors': ['OR59b'], 'name': 'ab2A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DM5': {'receptors': ['OR33b','OR85a'], 'name': 'ab2B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DM6': {'receptors': ['OR67a'], 'name': 'ab10B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'DP1l': {'receptors': ['IR75a'], 'name': 'ac2', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'DP1m': {'receptors': ['IR64a'], 'name': 'Sac III', 'co-receptors': ['IR8a'], 'sensillum': 'sacculus', 'sensillum location': 'antenna'},
                'V': {'receptors': ['GR21a','GR63a'], 'name': 'ab1C', 'co-receptors': [], 'sensillum': 'sacculus', 'sensillum location': 'antenna'},
                'DA4l': {'receptors': ['OR43a'], 'name': 'at3', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'VA1d': {'receptors': ['OR88a'], 'name': 'at4C', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'VA1v': {'receptors': ['OR47b'], 'name': 'at4A', 'co-receptors': ['Orco'], 'sensillum': 'trichodea', 'sensillum location': 'antenna'},
                'VA2': {'receptors': ['OR92a'], 'name': 'ab1B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'VA3': {'receptors': ['OR67b'], 'name': 'ab9', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'VA4': {'receptors': ['OR85d'], 'name': 'pb3B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'maxillary pulp'},
                'VA5': {'receptors': ['OR49b'], 'name': 'ab6B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna pulp'},
                'VA6': {'receptors': ['OR82a'], 'name': 'ab5A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna pulp'},
                'VA7l': {'receptors': ['OR46a'], 'name': 'pb2B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'maxillary pulp'},
                'VA7m': {'receptors': [], 'name': '', 'co-receptors': [], 'sensillum': 'unknown', 'sensillum location': 'unknown'},
                'VC1': {'receptors': ['OR33c,OR85e'], 'name': 'pb2A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'maxillary pulp'},
                'VC2': {'receptors': ['OR71a'], 'name': 'pb1B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'maxillary pulp'},
                'VC3l': {'receptors': ['OR35a'], 'name': 'ac1', 'co-receptors': ['Orco'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VC3m': {'receptors': [], 'name': 'unknown', 'co-receptors': ['Orco'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VC4': {'receptors': ['OR67c'], 'name': 'ab7B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'VC5': {'receptors': ['IR41a'], 'name': 'IR25a,IR76b', 'co-receptors': ['Orco'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VL1': {'receptors': ['IR75d'], 'name': 'ac1', 'co-receptors': ['IR25a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VM7d': {'receptors': ['OR42a'], 'name': 'pb1A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'maxillary pulp'},
                'VM7v': {'receptors': ['OR59c'], 'name': 'pb3A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'maxillary pulp'},
                'VM5v': {'receptors': ['OR98a'], 'name': 'ab7A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'VM5d': {'receptors': ['OR85b','OR98b'], 'name': 'ab3B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'VM4': {'receptors': ['IR76a'], 'name': 'ac4', 'co-receptors': ['IR25a','IR76b'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VM3': {'receptors': ['OR9a'], 'name': 'ab8B', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'VM2': {'receptors': ['OR43b'], 'name': 'ab8A', 'co-receptors': ['Orco'], 'sensillum': 'basiconica', 'sensillum location': 'antenna'},
                'VM1': {'receptors': ['OR92a'], 'name': 'ac1', 'co-receptors': ['IR25a','IR76b'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VL2p': {'receptors': ['IR31a'], 'name': 'ac1', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VL2a': {'receptors': ['IR84a'], 'name': 'ac4', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VP3': {'receptors': [], 'name': 'ac4', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VP5': {'receptors': [], 'name': 'ac4', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VP2': {'receptors': [], 'name': 'ac4', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VP1l': {'receptors': [], 'name': 'ac4', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
                'VP1m': {'receptors': [], 'name': 'ac4', 'co-receptors': ['IR8a'], 'sensillum': 'coeloconica', 'sensillum location': 'antenna'},
            }

        idx_a = all_odorant_names.index('putrescine')

        b_array = np.array(list(df.iloc[idx_a,1:]))

        or_names = [i.lower() for i in list(df.columns[1:])]
        bj = []
        or_names_found = []
        or_inds_found = []
        for i, val in enumerate(glom_names):
            receptors = GL_to_OR[val]['receptors']
            found = False
            for j in receptors:
                jreal = j
                j = j.lower()
                if j in or_names:
                    bj.append(b_array[or_names.index(j)])
                    found = True
                    or_names_found.append(jreal)
                    or_inds_found.append(len(bj)-1)
            if found == False:
                bj.append(0)
        self.ngloms = 56
        self.or_inds_found = or_inds_found
        self.or_names_found = or_names_found
        T = 50000
        dt = 1e-4
        bj = np.array(bj)

        I = np.zeros((self.ngloms, T))
        b = bj.reshape((-1,1))
        b = np.repeat(b, T, axis=1)
        I[:,10000:30000] = 5 + 0. * np.repeat(5.*np.random.random((1,20000)), self.ngloms, axis=0) # 1.
        Ib = I*b

        self.OSN = np.diff(np.floor(np.cumsum(Ib,axis=1)/threshold), axis=1)
        self.Ib = Ib
        self.threshold = threshold

        self.base_hpf_gain = 10.
        self.hemi_hpf_gain = 180.
        self.hemi_preLN_gain = 5e3
        self.preLN_linear_inhibition_gain = 1e-2
        self.hemi_amp = 4.
        self.osnf_filter = 20.
        self.alpha = 1.
        self.eps = 1e-1
        self.dt = dt
        self.T = T

    def sim(self, return_flycircuit=True):
        """Simulates the AL with a simple model for Hemibrain and FlyCircuit.

        # Arguments:
            return_flycircuit (bool): Whether to return the FlyCircuit output or not.
        """
        dt = self.dt
        T = self.T
        OSN = np.diff(np.floor(np.cumsum(self.Ib,axis=1)/self.threshold), axis=1)
        W_OSNLN = self.preLN_field_a
        W_PNLN = self.preLN_field_b.T
        preLN = W_PNLN.dot(W_OSNLN.dot(self.OSN))
        preLNbaseline = self.OSN
        OSNf = self.OSN * 0.
        alpha = 1.
        for j in range(1,OSNf.shape[1]):
            OSNf[:,j] = OSNf[:,j-1] + dt * (alpha * self.OSN[:,j]-self.osnf_filter*OSNf[:,j-1])
        XONE = OSN * 0.
        alpha = 1.
        for j in range(1,XONE.shape[1]):
            XONE[:,j] = XONE[:,j-1] + dt * (alpha * (1-XONE[:,j-1])*OSNf[:,j]-10.*XONE[:,j-1])
        XONEf = XONE * 0.
        alpha = 10.
        for j in range(1,OSNf.shape[1]):
            XONEf[:,j] = XONEf[:,j-1] + dt * (alpha * XONE[:,j]-10.0*XONEf[:,j-1])
        OSNh = XONE - XONEf
        OSNh = np.maximum(0., OSNh)
        def lpass(X, alpha=10.):
            Xf = X * 0.
            for j in range(1,X.shape[1]):
                Xf[:,j] = Xf[:,j-1] + dt * (alpha * X[:,j]-alpha*Xf[:,j-1])
            return Xf

        preLN = W_PNLN.dot(W_OSNLN.dot(OSNf))
        preLNbaseline = OSNf
        OSNhnew = W_PNLN.dot(W_OSNLN).dot(OSNf)
        OSNhnew2 = self.postLN_field.dot(OSNh)

        hemimodel = np.maximum(0.,self.hemi_amp * (self.hemi_preLN_gain * OSNf/(self.eps+1.*W_PNLN.dot(W_OSNLN).dot(OSNf))-self.preLN_linear_inhibition_gain*OSNhnew+self.hemi_hpf_gain*OSNhnew2))

        basemodel = preLNbaseline.copy()
        for j in range(1,basemodel.shape[1]):
            if np.sqrt(np.sum((basemodel[:,j])**2))>1e-0:
                basemodel[:,j] = basemodel[:,j] / np.sqrt(np.sum((basemodel[:,j])**2))
        basemodel = lpass(basemodel) + self.base_hpf_gain * OSNh
        if return_flycircuit:
            return hemimodel, basemodel
        else:
            return hemimodel


#!/usr/bin/env python

import io2, re, os, sys
import numpy as np
from io2.gaussian_reader import GaussianReader as GR0
import aqml.cheminfo.molecule.molecule as cmm
from aqml.cheminfo.core import *
import aqml.cheminfo.rdkit.core as crk
from aqml.cheminfo.rw.ctab import *
import scipy.spatial.distance as ssd

h2kc = io2.Units().h2kc
T, F = True, False
np.set_printoptions(formatter={'float': '{: 0.4f}'.format})

class _atoms(object):
    """ `atoms object from file formats other than xyz"""
    def __init__(self, f):
        import ase.io as aio
        m = aio.read(f)
        self.zs = m.numbers
        self.coords = m.positions
        self.na = len(self.zs)

uc = io2.Units() # unit converter

def get_val(dic, key):
    assert key in list(dic.keys()), '#ERROR: prop not found!'
    if key in ['HF','MP2','MP3','MP4','CCSD','CCSD(T)']:
        return '%.4f'%( dic[key]*uc.h2kc )
    elif key in ['dipole']:
        return '%.4f'%dic[key]
    else:
        raise '#ERROR: prop not found!'

class CM(object):
    """
    coulomb matrix object
    """
    def __init__(self, atoms, param={'M':'cml1','rp':1.,'wz':T,'sort':T}):
        self.param = param
        self.atoms = atoms
        self.cml1 = T if param['M'] in ['cml1'] else F

    def generate_coulomb_matrix(self):
        """ Coulomb matrix

        sorted CM has serious limitation when used to identify unique conformers.
        E.g., for CH3-CH3 molecule, as the L1 norm of all H-containing columns are
        the same, so shuffling such columns leads to different CM, though the molecule
        remains unchanged

        The limitation can be alleviated through the use of L1 norm of each column!!
        """
        atoms = self.atoms
        na = atoms.na
        mat = np.zeros((na,na))
        _ds = ssd.squareform( ssd.pdist(atoms.coords) )
        dsp = _ds**self.param['rp']
        np.fill_diagonal(dsp, 1.0)
        zs = atoms.zs
        _X, _Y = np.meshgrid(zs,zs)
        if self.param['wz']:
            mat = _X*_Y/dsp
            diag = -np.array(zs)**2.4
        else:
            mat = 1/dsp
            diag = np.zeros(na)
        np.fill_diagonal(mat,  diag)
        if self.param['sort']:
            L1s = np.array([ np.sum(np.abs(mat[i])) for i in range(na) ])
            ias = np.argsort(L1s)
            if self.cml1:
                x = L1s[ias]
            else:
                x = np.ravel(mat[ias,:][:,ias])
        else:
            x = np.ravel(mat)
        #print 'x = ', x
        return x

def cdist(objs, param={}):
    _param = {'M':'cml1','rp':1.0,'sort':T,'wz':F}
    for key in list(param.keys()):
        if key in list(_param.keys()):
            if param[key] != _param[key]:
                _param[key] = param[key]
    _xs = []
    nc = len(objs)
    for obj in objs:
        if _param['M'] in ['cm','cml1']:
            _xobj = CM(obj,_param)
            xi = _xobj.generate_coulomb_matrix()#; print '              xi = ', xi
            _xs.append( xi )
        else:
            raise '#ERROR: unknown `M'
    xs = np.array(_xs)
    return xs, ssd.squareform( ssd.pdist(xs,'cityblock') )

def get_alternative(s):
    """ c1cccc[n]1 --> c1ccccn1 """
    patt = '\[n\]'
    s = re.sub(patt,'n',s)
    return s

class OptedMols(object):

    """
    postprocess optimized goemetries (by G09)
    so as to retrieve only the unqiue conformers
    and then convert to sdf format with properties
    embedded at the end of the file
    """

    def __init__(self, fs, rsmi, props=['HF'], istart=0):
        self.nc0 = len(fs)
        fsc = []
        cso = [] # mol objects
        ms = [] # ase mols
        ys = []
        #assert '_c' in fs[0]
        #self.filename = '_'.join( fs[0].split('/')[-1].split('_')[:-1] )
        self.fs_diss = []
        self.fs_redundant = []
        cids = []
        for i,f in enumerate(fs):
            fmt = f[-3:]
            if fmt in ['log','out']: #G09 output file
                dic = GR0(f, istart=istart)[-1]
                zs = np.array(dic['Atomic_numbers'],np.int)
                coords = np.array( dic['Positions'] )
                m = atoms(zs, coords)
                try:
                    co = cmm.Mol(zs, coords, ican=True)
                    can2 = get_alternative(co.can)
                    if rsmi not in [co.can,can2]:
                        print("#ERROR: %s has a SMILES %s, differs from %s"%(f,co.can,rsmi))
                        self.fs_diss.append(f)
                        continue
                    else:
                        _ys = {}
                        for key in props:
                            _ys[key] = get_val(dic,key)
                        ys.append(_ys)
                        ms.append(m); cso.append(co); fsc.append(f)
                except:
                    print("#ERROR: this is a radical!")
                    self.fs_diss.append(f)
                    continue
            elif fmt in ['mol','sdf']:
                oo = crk.RDMol(f)
                m = atoms(oo.zs, oo.coords)
                ms.append(m)
                cso.append( oo.prop['smiles_indigo'] )
                fsc.append(f)
                ys.append( [ oo.prop[k] for k in props ] )
        self.cso = cso
        self.ms = ms
        self.fsc = fsc
        self.nc = len(cso)
        self.ys = ys

    def prune_conformers(self, param={'M':'cml1','wz':F,'thresh':0.01}, KeepPat=None):
        """ get unique conformers """
        ccidsr = [0,] # always keep the first conformer!!
        if self.nc > 1:
            xs, ds = cdist(self.ms, param=param)
            #self.ds = ds
            #seq = np.argsort(self.ys[:,0])
            for i in range(1,self.nc):
                #ic = seq[i]
                if (not np.all(ds[i,ccidsr]>param['thresh'])):
                    self.fs_redundant.append( self.fsc[i] )
                    continue
                ccidsr.append(i)
        self.ccidsr = ccidsr
        nc2 = len(ccidsr)
        if nc2 < self.nc:
            print('   %d out of %d conformers survived'%(nc2, self.nc))

    def write_conformers(self):
        """ write conformers to sdf files """
        #print self.ccidsr
        #print self.fsc
        for ic in self.ccidsr:
            fo = self.fsc[ic][:-4] + '.sdf'
            ci = self.cso[ic]
            #zs = [ chemical_symbols[zi] for zi in mi.zs ]
            #si = '%.4f #HF '%(self.ys[cid]*h2kc)
            #write_xyz(fo, (zs, mi.coords), comments=si)
            prop = self.ys[ic]
            prop['smiles_indigo'] = ci.can
            zs, coords, chgs, bom = ci.blk
            write_ctab(zs, chgs, bom, coords, sdf=fo, prop=prop)

if __name__ == "__main__":
    """
    generate conformers for a input molecule

    Attention: most frequently, the input are sdf files of AMONs !!!!!!!!!!
    """
    import stropr as so

    _args = sys.argv[1:]
    if ('-h' in _args) or (len(_args) < 3):
        print("Usage: ")
        print("   geomprune [-r amons.can] [-thresh 0.01] [-M cml1] [folder]")
        sys.exit()

    print(' \n Now executing ')
    print('         geomprune ' + ' '.join(sys.argv[1:]) + '\n')

    idx = 0
    keys = ['-r','-ref']; hask,fsmi,idx = so.parser(_args,keys,'',idx,F)
    assert hask, '#ERROR: a reference smi/can file must be provided'
    assert fsmi[-3:] in ['smi','can'], '#ERROR: ref smiles file format not allowed'

    keys = ['-fmt','-format']; ifmt,ffmt,idx = so.parser(_args,keys,'sdf',idx,F)
    assert ifmt, '#ERROR: plz specify [-ffmt sdf/out]'

    keys = ['-d','-digits']; has_sl,sl,idx = so.parser(_args,keys,'6',idx,F) # frag_000001.sdf ...

    keys = ['-w','-write']; write,idx = so.haskey(_args,keys,idx) # rename & write unique conformers

    thresh = 0.01
    rep = 'cml1'

    fd = _args[idx]
    refs = [ si.strip() for si in file(fsmi).readlines() ]
    ng = len(refs) # number of (unique molecular) graphs

    if has_sl:
        sfmt = '%%0%dd'%( int(sl) )
    else:
        sfmt = '%%0%dd'%( len(str(ng)) )

    for _ in ['diss','redundant']:
        if not os.path.exists(fd+'/'+_):
            os.system('mkdir -p %s/%s'%(fd,_))

    for mid in range(1,ng+1):
        lb = sfmt%mid
        fs = io2.cmdout('ls %s/frag_'%fd + lb + '_*%s'%ffmt)
        print(' ** now processing mid %s'%lb)
        obj = OptedMols(fs,refs[mid-1])
        obj.prune_conformers(param={'M':rep,'wz':False,'thresh':thresh})
        if write:
            obj.write_conformers()
        for f in obj.fs_diss:
            os.system('mv %s.* %s/diss'%(f[:-4],fd))
        for f in obj.fs_redundant:
            cmd = 'mv %s.* %s/redundant'%(f[:-4],fd)
            os.system(cmd)


"""
#;+
#; NAME:
#; utils
#;    Version 1.0
#;
#; PURPOSE:
#;    Module for spectral utilities
#;       Primarily overloads of Spectrum1D
#;   07-Sep-2014 by JXP
#;-
#;------------------------------------------------------------------------------
"""
from __future__ import print_function, absolute_import, division, unicode_literals

import numpy as np
import os
import astropy as apy

from astropy import units as u
from astropy import constants as const
from astropy.io import fits

from specutils import Spectrum1D
from specutils.wcs import BaseSpectrum1DWCS, Spectrum1DLookupWCS
from specutils.wcs.specwcs import Spectrum1DPolynomialWCS

from xastropy.xutils import xdebug as xdb

# Child Class of specutils/Spectrum1D
#    Generated by JXP to add functionality before it gets ingested in the specutils distribution
class XSpectrum1D(Spectrum1D):

    #### ###############################
    #  Instantiate from Spectrum1D [best to avoid!]
    @classmethod
    def from_spec1d(cls, spec1d):

        # Giddy up
        return cls(flux=spec1d.flux, wcs=spec1d.wcs, unit=spec1d.unit,
                   uncertainty=spec1d.uncertainty, mask=spec1d.mask, meta=spec1d.meta)

    #### ###############################
    #  Normalize
    def normalize(self, conti, verbose=False, no_check=False):
        """
        Normalize the spectrum with an input continuum

        Parameters
        ----------
        conti: numpy array
          Continuum
        verbose: bool (False)
        no_check: bool (False)
          Check size of array?
        """
        # Sanity check
        if (len(conti) != len(self.flux)):
            if no_check:
                print('WARNING: Continuum length differs from flux')
                if len(conti) > len(self.flux):
                    self.flux = self.flux / conti[0:len(self.flux)]
                    return
                else:
                    raise ValueError('normalize: Continuum needs to be longer!')
            else:
                raise ValueError('normalize: Continuum needs to be same length as flux array')

        # Adjust the flux
        self.flux = self.flux / conti
        if verbose:
            print('spec.utils: Normalizing the spectrum')

    #### ###############################
    #  Grabs spectrum pixels in a velocity window
    def pix_minmax(self, *args):
        """Pixels in velocity range

        Parameters
        ----------
        Option 1: wvmnx
          wvmnx: Tuple of 2 floats
            wvmin, wvmax in spectral units

        Option 2: zabs, wrest, vmnx  [not as a tuple or list!]
          zabs: Absorption redshift
          wrest: Rest wavelength  (with Units!)
          vmnx: Tuple of 2 floats
            vmin, vmax in km/s

        Returns:
        pix: array
          Integer list of pixels
        """
        if len(args) == 1: # Option 1
            wvmnx = args[0]
        elif len(args) == 3: # Option 2
            from astropy import constants as const
            # args = zabs, wrest, vmnx
            wvmnx = (args[0]+1) * (args[1] + (args[1] * args[2] / const.c.to('km/s')) )
            wvmnx.to(u.AA)

        # Locate the values
        pixmin = np.argmin( np.fabs( self.dispersion-wvmnx[0] ) )
        pixmax = np.argmin( np.fabs( self.dispersion-wvmnx[1] ) )

        gdpix = np.arange(pixmin,pixmax+1)

        # Fill + Return
        self.sub_pix = gdpix
        return gdpix, wvmnx, (pixmin, pixmax)

    #### ###############################
    #  Box car smooth
    def box_smooth(self, nbox, preserve=False):
        """ Box car smooth spectrum and return a new one
        Is a simple wrapper to the rebin routine

        Parameters
        ----------
        nbox: integer
          Number of pixels to smooth over
        preserve: bool (False)
          Keep the new spectrum at the same number of pixels as original
        Returns:
          XSpectrum1D of the smoothed spectrum
        """
        from xastropy.xutils import arrays as xxa
        if preserve:
            from astropy.convolution import convolve, Box1DKernel
            new_fx = convolve(self.flux, Box1DKernel(nbox))
            new_sig = convolve(self.sig, Box1DKernel(nbox))
            new_wv = self.dispersion
        else:
            # Truncate arrays as need be
            npix = len(self.flux)
            try:
                new_npix = npix // nbox # New division
            except ZeroDivisionError:
                xdb.set_trace()
            orig_pix = np.arange( new_npix * nbox )

            # Rebin (mean)
            new_wv = xxa.scipy_rebin( self.dispersion[orig_pix], new_npix )
            new_fx = xxa.scipy_rebin( self.flux[orig_pix], new_npix )
            new_sig = xxa.scipy_rebin( self.sig[orig_pix], new_npix ) / np.sqrt(nbox)

        # Return
        return XSpectrum1D.from_array(new_wv, new_fx,
                                      uncertainty=apy.nddata.StdDevUncertainty(new_sig))

    #### ###############################
    #  Rebin
    def rebin(self, new_wv):
        """ Rebin the existing spectrum to a new wavelength array
        Uses simple linear interpolation.  The default (and only) option
        conserves counts (and flambda).

        WARNING: Do not trust either edge pixel of the new array

        Parameters
        ----------
        new_wv: Quantity array
          New wavelength array

        Returns:
        ----------
          XSpectrum1D of the rebinned spectrum
        """
        from scipy.interpolate import interp1d

        # Endpoints of original pixels
        npix = len(self.dispersion)
        wvh = (self.dispersion + np.roll(self.dispersion, -1))/2.
        wvh[npix-1] = self.dispersion[npix-1] + (self.dispersion[npix-1] - self.dispersion[npix-2])/2.
        dwv = wvh - np.roll(wvh,1)
        dwv[0] = 2*(wvh[0]-self.dispersion[0])

        # Cumulative Sum
        cumsum = np.cumsum(self.flux * dwv)

        # Interpolate
        fcum = interp1d(wvh, cumsum, fill_value=0., bounds_error=False)

        # Endpoints of new pixels
        nnew = len(new_wv)
        nwvh = (new_wv + np.roll(new_wv, -1))/2.
        nwvh[nnew-1] = new_wv[nnew-1] + (new_wv[nnew-1] - new_wv[nnew-2])/2.
        # Pad starting point
        bwv = np.zeros(nnew+1) * new_wv.unit
        #xdb.set_trace()
        bwv[0] = new_wv[0] - (new_wv[1] - new_wv[0])/2.
        bwv[1:] = nwvh

        # Evaluate
        newcum = fcum(bwv)
        # Endpoint
        if (bwv[-1] > wvh[-1]):
            newcum[-1] = cumsum[-1]

        # Rebinned flux
        new_fx = (np.roll(newcum,-1)-newcum)[:-1]

        # Normalize (preserve counts and flambda)
        new_dwv = bwv - np.roll(bwv,1)
        new_fx = new_fx / new_dwv[1:]

        # Return new spectrum
        return XSpectrum1D.from_array(new_wv, new_fx)

    # Quick plot
    def plot(self):
        ''' Plot the spectrum
        Parameters
        ----------
        '''
        if self.sig is not None:
            xdb.xplot(self.dispersion, self.flux, self.sig)
        else:
            xdb.xplot(self.dispersion, self.flux)

    # Velo array
    def relative_vel(self, wv_obs):
        ''' Return a velocity array relative to an input wavelength
        Should consider adding a velocity array to this Class,
        i.e. self.velo

        Parameters
        ----------
        wv_obs : float
          Wavelength to set the zero of the velocity array.
          Often (1+z)*wrest

        Returns:
        ---------
        velo: Quantity array (km/s)
        '''
        return  (self.dispersion-wv_obs) * const.c.to('km/s')/wv_obs

    # Write to fits
    def write_to_fits(self, outfil, clobber=True, add_wave=False):
        ''' Write to a FITS file
        Should generate a separate code to make a Binary FITS table format

        Parameters
        ----------
        outfil: String
          Name of the FITS file
        clobber: bool (True)
          Clobber existing file?
        add_wave: bool (False)
          Force writing of wavelength array
        '''
        # TODO
        #  1. Add unit support for wavelength arrays

        from specutils.io import write_fits as sui_wf
        prihdu = sui_wf._make_hdu(self.data)  # Not for binary table format
        prihdu.name = 'FLUX'
        multi = 0 #  Multi-extension?

        # Type
        if type(self.wcs) is Spectrum1DPolynomialWCS:  # CRVAL1, etc. WCS
            # WCS
            wcs = self.wcs
            wcs.write_fits_header(prihdu.header)
            # Error array?
            if self.sig is not None:
                sighdu = fits.ImageHDU(self.sig)
                sighdu.name='ERROR'
                #
                if add_wave:
                    wvhdu = fits.ImageHDU(self.dispersion.value)
                    wvhdu.name = 'WAVELENGTH'
                    hdu = fits.HDUList([prihdu, sighdu, wvhdu])
                else:
                    hdu = fits.HDUList([prihdu, sighdu])
                multi=1
            else:
                hdu = prihdu

        elif type(self.wcs) is Spectrum1DLookupWCS: # Wavelengths as an array (without units for now)
            # Add sig, wavelength to HDU
            sighdu = fits.ImageHDU(self.sig)
            sighdu.name='ERROR'
            wvhdu = fits.ImageHDU(self.dispersion.value)
            wvhdu.name = 'WAVELENGTH'
            hdu = fits.HDUList([prihdu, sighdu, wvhdu])
            multi=1
        else:
            raise ValueError('write_to_fits: Not ready for this type of spectrum wavelengths')

        # Deal with header
        if hasattr(self,'head'):
            hdukeys = prihdu.header.keys()
            # Append ones to avoid
            hdukeys = hdukeys +ZZ ['BUNIT','COMMENT','', 'NAXIS2', 'HISTORY']
            for key in self.head.keys():
                # Use new ones
                if key in hdukeys:
                    continue
                # Update unused ones
                try:
                    prihdu.header[key] = self.head[key]
                except ValueError:
                    xdb.set_trace()
            # History
            if 'HISTORY' in self.head.keys():
                prihdu.header.add_history(str(self.head['HISTORY']))

        # Write
        hdu.writeto(outfil, clobber=clobber)
        print('Wrote spectrum to {:s}'.format(outfil))

# Quick plot
def bspline_stack(spectra):
    ''' "Stack" a set of spectra with a bspline algorithm
    Might be useful for coadding

    Parameters:
    -----------
    spectra: List of Spectrum1D

    Returns:
    -------
    bspline
    '''

# ################
if __name__ == "__main__":

    flg_test = 0
    #flg_test += 2**0  # Test write (simple)
    #flg_test += 2**1  # Test write with 3 arrays
    #flg_test += 2**2  # Test boxcar
    flg_test += 2**3  # Test rebin

    from xastropy.spec import readwrite as xsr

    if (flg_test % 2**1) >= 2**0:
        # Standard log-linear read + write (MagE)
        fil = '~/PROGETTI/LLSZ3/data/normalize/UM669_nF.fits'
        myspec = xsr.readspec(fil)
        # Write
        myspec.write_to_fits('tmp.fits')

    if (flg_test % 2**2) >= 2**1:
        # Now 2D
        fil = '/Users/xavier/Dropbox/QSOPairs/data/LRIS_redux/SDSSJ231254.65-025403.1_b400_F.fits.gz'
        myspec = xsr.readspec(fil)
        myspec.write_to_fits('tmp.fits')

    if (flg_test % 2**3) >= 2**2: # Boxcar
        fil = '~/PROGETTI/LLSZ3/data/normalize/UM669_nF.fits'
        myspec = xsr.readspec(fil)
        newspec = myspec.box_smooth(3)
        #
        newspec2 = myspec.box_smooth(3, preserve=True)
        xdb.xplot(myspec.dispersion, myspec.flux, newspec2.flux)

    if (flg_test % 2**4) >= 2**3: # Rebin array
        fil = '~/PROGETTI/LLSZ3/data/normalize/UM669_nF.fits'
        myspec = xsr.readspec(fil)

        new_wv = np.arange(3000., 9000., 5) * u.AA
        newspec = myspec.rebin(new_wv)
        #xdb.xplot(myspec.dispersion, myspec.flux,
        #    xtwo=new_wv, ytwo=newspec.flux)
        # Test EW
        wvmnx = np.array((4859., 4961.))*u.AA
        gd1 = np.where( (myspec.dispersion > wvmnx[0]) &
            (myspec.dispersion < wvmnx[1]))[0]
        dwv1 = myspec.dispersion - np.roll(myspec.dispersion,1)
        EW1 = np.sum(dwv1[gd1]*(1.-myspec.flux[gd1].value))
        gd2 = np.where( (newspec.dispersion > wvmnx[0]) &
            (newspec.dispersion < wvmnx[1]))[0]
        dwv2 = newspec.dispersion - np.roll(newspec.dispersion,1)
        EW2 = np.sum(dwv2[gd2]*(1.-newspec.flux[gd2].value))
        print('EW1={:g} and EW2={:g} for wvmnx={:g},{:g}'.format(
            EW1,EW2,wvmnx[0],wvmnx[1]))
        print('Percent diff = {:0.2f}%'.format(100*(EW2-EW1)/EW1))

import numpy,from scipy,import astropy,from astropy