# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2012-2018 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:`ChiouYoungs2014`.
"""
import numpy as np
import math

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

class ChiouYoungs2014(GMPE):
    """
    Implements GMPE developed by Brian S.-J. Chiou and Robert R. Youngs
    and published as "Updated of the Chiou and Youngs NGA Model for the
    Average Horizontal Component of Peak Ground Motion and Response Spectra"
    (2014, Earthquake Spectra).
    """
    #: Supported tectonic region type is active shallow crust
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.ACTIVE_SHALLOW_CRUST

    #: Supported intensity measure types are spectral acceleration,
    #: peak ground velocity and peak ground acceleration
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([
        PGA,
        PGV,
        SA
    ])

    #: Supported intensity measure component is orientation-independent
    #: measure :attr:`~openquake.hazardlib.const.IMC.RotD50`,
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.RotD50

    #: Supported standard deviation types are inter-event, intra-event
    #: and total, see chapter "Variance model".
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL,
        const.StdDev.INTER_EVENT,
        const.StdDev.INTRA_EVENT
    ])

    #: Required site parameters are Vs30, Vs30 measured flag
    #: and Z1.0.
    REQUIRES_SITES_PARAMETERS = set(('vs30', 'vs30measured', 'z1pt0'))

    #: Required rupture parameters are magnitude, rake,
    #: dip and ztor.
    REQUIRES_RUPTURE_PARAMETERS = set(('dip', 'rake', 'mag', 'ztor'))

    #: Required distance measures are RRup, Rjb and Rx.
    REQUIRES_DISTANCES = set(('rrup', 'rjb', 'rx'))

    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[imt]
        # intensity on a reference soil is used for both mean
        # and stddev calculations.
        ln_y_ref = self._get_ln_y_ref(rup, dists, C)
        # exp1 and exp2 are parts of eq. 12 and eq. 13,
        # calculate it once for both.
        exp1 = np.exp(C['phi3'] * (sites.vs30.clip(-np.inf, 1130) - 360))
        exp2 = np.exp(C['phi3'] * (1130 - 360))
        mean = self._get_mean(sites, C, ln_y_ref, exp1, exp2)
        stddevs = self._get_stddevs(sites, rup, C, stddev_types,
                                    ln_y_ref, exp1, exp2)

        return mean, stddevs

    def _get_mean(self, sites, C, ln_y_ref, exp1, exp2):
        """
        Add site effects to an intensity.

        Implements eq. 13b.
        """
        # we do not support estimating of basin depth and instead
        # rely on it being available (since we require it).
        # centered_z1pt0
        centered_z1pt0 = self._get_centered_z1pt0(sites)
        # we consider random variables being zero since we want
        # to find the exact mean value.
        eta = epsilon = 0.

        ln_y = (
            # first line of eq. 12
            ln_y_ref + eta
            # second line
            + C['phi1'] * np.log(sites.vs30 / 1130).clip(-np.inf, 0)
            # third line
            + C['phi2'] * (exp1 - exp2)
            * np.log((np.exp(ln_y_ref) * np.exp(eta) + C['phi4']) / C['phi4'])
            # fourth line
            + C['phi5']
            * (1.0 - np.exp(-1. * centered_z1pt0 / C['phi6']))
            # fifth line
            + epsilon
        )

        return ln_y

    def _get_stddevs(self, sites, rup, C, stddev_types, ln_y_ref, exp1, exp2):
        """
        Get standard deviation for a given intensity on reference soil.

        Implements equations 13 for inter-event, intra-event
        and total standard deviations.
        """
        Fmeasured = sites.vs30measured
        Finferred = 1 - sites.vs30measured

        # eq. 13 to calculate inter-event standard error
        mag_test = min(max(rup.mag, 5.0), 6.5) - 5.0
        tau = C['tau1'] + (C['tau2'] - C['tau1']) / 1.5 * mag_test

        # b and c coeffs from eq. 10
        b = C['phi2'] * (exp1 - exp2)
        c = C['phi4']

        y_ref = np.exp(ln_y_ref)
        # eq. 13
        NL = b * y_ref / (y_ref + c)
        sigma = ((C['sig1'] + (C['sig2'] - C['sig1']) * mag_test / 1.5)
                 * np.sqrt((C['sig3'] * Finferred + 0.7 * Fmeasured) +
                           (1. + NL) ** 2.))

        ret = []
        for stddev_type in stddev_types:
            assert stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
            if stddev_type == const.StdDev.TOTAL:
                # eq. 13
                ret += [np.sqrt(((1 + NL) ** 2) * (tau ** 2) + (sigma ** 2))]
            elif stddev_type == const.StdDev.INTRA_EVENT:
                ret.append(sigma)
            elif stddev_type == const.StdDev.INTER_EVENT:
                # this is implied in eq. 21
                ret.append(np.abs((1 + NL) * tau))
        return ret

    def _get_ln_y_ref(self, rup, dists, C):
        """
        Get an intensity on a reference soil.

        Implements eq. 13a.
        """
        # reverse faulting flag
        Frv = 1. if 30 <= rup.rake <= 150 else 0.
        # normal faulting flag
        Fnm = 1. if -120 <= rup.rake <= -60 else 0.
        # hanging wall flag

        Fhw = np.zeros_like(dists.rx)
        idx = np.nonzero(dists.rx >= 0.)
        Fhw[idx] = 1.

        # a part in eq. 11
        mag_test1 = np.cosh(2. * max(rup.mag - 4.5, 0))

        # centered DPP
        centered_dpp = self._get_centered_cdpp(dists)
        # centered_ztor
        centered_ztor = self._get_centered_ztor(rup, Frv)
        #
        dist_taper = np.fmax(1 - (np.fmax(dists.rrup - 40,
                                  np.zeros_like(dists)) / 30.),
                             np.zeros_like(dists))
        dist_taper = dist_taper.astype(np.float64)
        ln_y_ref = (
            # first part of eq. 11
            C['c1']
            + (C['c1a'] + C['c1c'] / mag_test1) * Frv
            + (C['c1b'] + C['c1d'] / mag_test1) * Fnm
            + (C['c7'] + C['c7b'] / mag_test1) * centered_ztor
            + (C['c11'] + C['c11b'] / mag_test1) *
            np.cos(math.radians(rup.dip)) ** 2
            # second part
            + C['c2'] * (rup.mag - 6)
            + ((C['c2'] - C['c3']) / C['cn'])
            * np.log(1 + np.exp(C['cn'] * (C['cm'] - rup.mag)))
            # third part
            + C['c4']
            * np.log(dists.rrup + C['c5']
                     * np.cosh(C['c6'] * max(rup.mag - C['chm'], 0)))
            + (C['c4a'] - C['c4'])
            * np.log(np.sqrt(dists.rrup ** 2 + C['crb'] ** 2))
            # forth part
            + (C['cg1'] + C['cg2'] / (np.cosh(max(rup.mag - C['cg3'], 0))))
            * dists.rrup
            # fifth part
            + C['c8'] * dist_taper
            * min(max(rup.mag - 5.5, 0) / 0.8, 1.0)
            * np.exp(-1 * C['c8a'] * (rup.mag - C['c8b']) ** 2) * centered_dpp
            # sixth part
            + C['c9'] * Fhw * np.cos(math.radians(rup.dip)) *
            (C['c9a'] + (1 - C['c9a']) * np.tanh(dists.rx / C['c9b']))
            * (1 - np.sqrt(dists.rjb ** 2 + rup.ztor ** 2)
               / (dists.rrup + 1.0))
        )

        return ln_y_ref

    def _get_centered_z1pt0(self, sites):
        """
        Get z1pt0 centered on the Vs30- dependent avarage z1pt0(m)
        California and non-Japan regions

        """
        #: California and non-Japan regions

        mean_z1pt0 = (-7.15 / 4.) * np.log(((sites.vs30) ** 4. + 570.94 ** 4.)
                                           / (1360 ** 4. + 570.94 ** 4.))
        centered_z1pt0 = sites.z1pt0 - np.exp(mean_z1pt0)

        return centered_z1pt0

    def _get_centered_ztor(self, rup, Frv):
        """
        Get ztor centered on the M- dependent avarage ztor(km)
        by different fault types.
        """
        if Frv == 1:

            mean_ztor = max(2.704 - 1.226 * max(rup.mag - 5.849, 0.0), 0.) ** 2
            centered_ztor = rup.ztor - mean_ztor
        else:

            mean_ztor = max(2.673 - 1.136 * max(rup.mag - 4.970, 0.0), 0.) ** 2
            centered_ztor = rup.ztor - mean_ztor

        return centered_ztor

    def _get_centered_cdpp(self, dists):
        """
        Get directivity prediction parameter centered on the avgerage
        directivity prediction parameter. Here we set the centered_dpp
        equals to zero, since the near fault directivity effect prediction is
        off in our calculation.

        """
        centered_dpp = 0.

        return centered_dpp

    #: Coefficient tables are constructed from values in tables 1 - 5

    COEFFS = CoeffsTable(sa_damping=5, table="""\
IMT     c1      c1a     c1b     c1c     c1d     cn      cm    c2      c3    c4     c4a  crb   c5      chm     c6      c7      c7b     c8     c8a    c8b       c9     c9a    c9b     c11      c11b        cg1        cg2       cg3     phi1       phi2      phi3     phi4     phi5   phi6  gjpit  gwn      phi1jp  phi5jp   phi6jp     tau1    tau2    sig1    sig2    sig3    sig2jp
pga   -1.5065  0.165  -0.255  -0.165  0.255  16.0875  4.9993  1.06  1.9636  -2.1  -0.5  50  6.4551  3.0956  0.4908  0.0352   0.0462  0.     0.2695  0.4833  0.9228  0.1202  6.8607  0.      -0.4536    -0.007146  -0.006758  4.2542  -0.521   -0.1417   -0.00701   0.102151  0.     300  1.5817  0.7594  -0.6846  0.459    800.        0.4     0.26    0.4912  0.3762  0.8     0.4528
pgv    2.3549  0.165  -0.0626 -0.165  0.0626  3.3024  5.423   1.06  2.3152  -2.1  -0.5  50  5.8096  3.0514  0.4407  0.0324   0.0097  0.2154 0.2695  5.      0.3079  0.1     6.5     0       -0.3834    -0.001852  -0.007403  4.3439  -0.7936  -0.0699   -0.008444  5.41      0.0202 300. 2.2306  0.335   -0.7966  0.9488   800.        0.3894  0.2578  0.4785  0.3629  0.7504  0.3918
0.01  -1.5065  0.165  -0.255  -0.165  0.255  16.0875  4.9993  1.06  1.9636  -2.1  -0.5  50  6.4551  3.0956  0.4908  0.0352   0.0462  0.     0.2695  0.4833  0.9228  0.1202  6.8607  0.      -0.4536    -0.007146  -0.006758  4.2542  -0.521   -0.1417   -0.00701   0.102151  0.     300  1.5817  0.7594  -0.6846  0.459    800.        0.4     0.26    0.4912  0.3762  0.8     0.4528
0.02  -1.4798  0.165  -0.255  -0.165  0.255  15.7118  4.9993  1.06  1.9636  -2.1  -0.5  50  6.4551  3.0963  0.4925  0.0352   0.0472  0.     0.2695  1.2144  0.9296  0.1217  6.8697  0.      -0.4536    -0.007249  -0.006758  4.2386  -0.5055  -0.1364   -0.007279  0.10836   0.     300  1.574   0.7606  -0.6681  0.458    800.        0.4026  0.2637  0.4904  0.3762  0.8     0.4551
0.03  -1.2972  0.165  -0.255  -0.165  0.255  15.8819  4.9993  1.06  1.9636  -2.1  -0.5  50  6.4551  3.0974  0.4992  0.0352   0.0533  0.     0.2695  1.6421  0.9396  0.1194  6.9113  0.      -0.4536    -0.007869  -0.006758  4.2519  -0.4368  -0.1403   -0.007354  0.119888  0.     300  1.5544  0.7642  -0.6314  0.462    800.        0.4063  0.2689  0.4988  0.3849  0.8     0.4571
0.04  -1.1007  0.165  -0.255  -0.165  0.255  16.4556  4.9993  1.06  1.9636  -2.1  -0.5  50  6.4551  3.0988  0.5037  0.0352   0.0596  0.     0.2695  1.9456  0.9661  0.1166  7.0271  0.      -0.4536    -0.008316  -0.006758  4.296   -0.3752  -0.1591   -0.006977  0.133641  0.     300  1.5502  0.7676  -0.5855  0.453    800.        0.4095  0.2736  0.5049  0.391   0.8     0.4642
0.05  -0.9292  0.165  -0.255  -0.165  0.255  17.6453  4.9993  1.06  1.9636  -2.1  -0.5  50  6.4551  3.1011  0.5048  0.0352   0.0639  0.     0.2695  2.181   0.9794  0.1176  7.0959  0.      -0.4536    -0.008743  -0.006758  4.3578  -0.3469  -0.1862   -0.006467  0.148927  0.     300  1.5391  0.7739  -0.5457  0.436    800.        0.4124  0.2777  0.5096  0.3957  0.8     0.4716
0.075 -0.658   0.165  -0.254  -0.165  0.254  20.1772  5.0031  1.06  1.9636  -2.1  -0.5  50  6.4551  3.1094  0.5048  0.0352   0.063   0.     0.2695  2.6087  1.026   0.1171  7.3298  0.      -0.4536    -0.009537  -0.00619   4.5455  -0.3747  -0.2538   -0.005734  0.190596  0.     300  1.4804  0.7956  -0.4685  0.383    800.        0.4179  0.2855  0.5179  0.4043  0.8     0.5022
0.1   -0.5613  0.165  -0.253  -0.165  0.253  19.9992  5.0172  1.06  1.9636  -2.1  -0.5  50  6.8305  3.2381  0.5048  0.0352   0.0532  0.     0.2695  2.9122  1.0177  0.1146  7.2588  0.      -0.4536    -0.00983   -0.005332  4.7603  -0.444   -0.2943   -0.005604  0.230662  0.     300  1.4094  0.7932  -0.4985  0.375    800.        0.4219  0.2913  0.5236  0.4104  0.8     0.523
0.12  -0.5342  0.165  -0.252  -0.165  0.252  18.7106  5.0315  1.06  1.9795  -2.1  -0.5  50  7.1333  3.3407  0.5048  0.0352   0.0452  0.     0.2695  3.1045  1.0008  0.1128  7.2372  0.      -0.4536    -0.009913  -0.004732  4.8963  -0.4895  -0.3077   -0.005696  0.253169  0.     300  1.3682  0.7768  -0.5603  0.377    800.        0.4244  0.2949  0.527   0.4143  0.8     0.5278
0.15  -0.5462  0.165  -0.25   -0.165  0.25   16.6246  5.0547  1.06  2.0362  -2.1  -0.5  50  7.3621  3.43    0.5045  0.0352   0.0345  0.     0.2695  3.3399  0.9801  0.1106  7.2109  0.      -0.4536    -0.009896  -0.003806  5.0644  -0.5477  -0.3113   -0.005845  0.266468  0.     300  1.3241  0.7437  -0.6451  0.379    800.        0.4275  0.2993  0.5308  0.4191  0.8     0.5304
0.17  -0.5858  0.165  -0.248  -0.165  0.248  15.3709  5.0704  1.06  2.0823  -2.1  -0.5  50  7.4365  3.4688  0.5036  0.0352   0.0283  0.     0.2695  3.4719  0.9652  0.115   7.2491  0.      -0.4536    -0.009787  -0.00328   5.1371  -0.5922  -0.3062   -0.005959  0.26506   0.     300  1.3071  0.7219  -0.6981  0.38     800.        0.4292  0.3017  0.5328  0.4217  0.8     0.531
0.2   -0.6798  0.165  -0.2449 -0.165  0.2449 13.7012  5.0939  1.06  2.1521  -2.1  -0.5  50  7.4972  3.5146  0.5016  0.0352   0.0202  0.     0.2695  3.6434  0.9459  0.1208  7.2988  0.      -0.444     -0.009505  -0.00269   5.188   -0.6693  -0.2927   -0.006141  0.255253  0.     300  1.2931  0.6922  -0.7653  0.384    800.        0.4313  0.3047  0.5351  0.4252  0.8     0.5312
0.25  -0.8663  0.165  -0.2382 -0.165  0.2382 11.2667  5.1315  1.06  2.2574  -2.1  -0.5  50  7.5416  3.5746  0.4971  0.0352   0.009   0.     0.2695  3.8787  0.9196  0.1208  7.3691  0.      -0.3539    -0.008918  -0.002128  5.2164  -0.7766  -0.2662   -0.006439  0.231541  0.     300  1.315   0.6579  -0.8469  0.393    800.        0.4341  0.3087  0.5377  0.4299  0.7999  0.5309
0.3   -1.0514  0.165  -0.2313 -0.165  0.2313  9.1908  5.167   1.06  2.344   -2.1  -0.5  50  7.56    3.6232  0.4919  0.0352  -0.0004  0.     0.2695  4.0711  0.8829  0.1175  6.8789  0.      -0.2688    -0.008251  -0.001812  5.1954  -0.8501  -0.2405   -0.006704  0.207277  0.001  300  1.3514  0.6362  -0.8999  0.408    800.        0.4363  0.3119  0.5395  0.4338  0.7997  0.5307
0.4   -1.3794  0.165  -0.2146 -0.165  0.2146  6.5459  5.2317  1.06  2.4709  -2.1  -0.5  50  7.5735  3.6945  0.4807  0.0352  -0.0155  0.     0.2695  4.3745  0.8302  0.106   6.5334  0.      -0.1793    -0.007267  -0.001274  5.0899  -0.9431  -0.1975   -0.007125  0.165464  0.004  300  1.4051  0.6049  -0.9618  0.462    800.        0.4396  0.3165  0.5422  0.4399  0.7988  0.531
0.5   -1.6508  0.165  -0.1972 -0.165  0.1972  5.2305  5.2893  1.06  2.5567  -2.1  -0.5  50  7.5778  3.7401  0.4707  0.0352  -0.0278  0.0991 0.2695  4.6099  0.7884  0.1061  6.526   0.      -0.1428    -0.006492  -0.001074  4.7854  -1.0044  -0.1633   -0.007435  0.133828  0.01   300  1.4402  0.5507  -0.9945  0.524    800.        0.4419  0.3199  0.5433  0.4446  0.7966  0.5313
0.75  -2.1511  0.165  -0.162  -0.165  0.162   3.7896  5.4109  1.06  2.6812  -2.1  -0.5  50  7.5808  3.7941  0.4575  0.0352  -0.0477  0.1982 0.2695  5.0376  0.6754  0.1     6.5     0.      -0.1138    -0.005147  -0.001115  4.3304  -1.0602  -0.1028   -0.00812   0.085153  0.034  300  1.528   0.3582  -1.0225  0.658    800.        0.4459  0.3255  0.5294  0.4533  0.7792  0.5309
1     -2.5365  0.165  -0.14   -0.165  0.14    3.3024  5.5106  1.06  2.7474  -2.1  -0.5  50  7.5814  3.8144  0.4522  0.0352  -0.0559  0.2154 0.2695  5.3411  0.6196  0.1     6.5     0.      -0.1062    -0.004277  -0.001197  4.1667  -1.0941  -0.0699   -0.008444  0.058595  0.067  300  1.6523  0.2003  -1.0002  0.78     800.        0.4484  0.3291  0.5105  0.4594  0.7504  0.5302
1.5   -3.0686  0.165  -0.1184 -0.165  0.1184  2.8498  5.6705  1.06  2.8161  -2.1  -0.5  50  7.5817  3.8284  0.4501  0.0352  -0.063   0.2154 0.2695  5.7688  0.5101  0.1     6.5     0.      -0.102     -0.002979  -0.001675  4.0029  -1.1142  -0.0425   -0.007707  0.031787  0.143  300  1.8872  0.0356  -0.9245  0.96     800.        0.4515  0.3335  0.4783  0.468   0.7136  0.5276
2     -3.4148  0.1645 -0.11   -0.1645 0.11    2.5417  5.7981  1.06  2.8514  -2.1  -0.5  50  7.5818  3.833   0.45    0.0352  -0.0665  0.2154 0.2695  6.0723  0.3917  0.1     6.5     0.      -0.1009    -0.002301  -0.002349  3.8949  -1.1154  -0.0302   -0.004792  0.019716  0.203  300  2.1348  0.      -0.8626  1.11     800.        0.4534  0.3363  0.4681  0.4681  0.7035  0.5167
3     -3.9013  0.1168 -0.104  -0.1168 0.104   2.1488  5.9983  1.06  2.8875  -2.1  -0.5  50  7.5818  3.8361  0.45    0.016   -0.0516  0.2154 0.2695  6.5     0.1244  0.1     6.5     0.      -0.1003    -0.001344  -0.003306  3.7928  -1.1081  -0.0129   -0.001828  0.009643  0.277  300  3.5752  0.      -0.7882  1.291    800.        0.4558  0.3398  0.4617  0.4617  0.7006  0.4917
4     -4.2466  0.0732 -0.102  -0.0732 0.102  1.8957   6.1552  1.06  2.9058  -2.1  -0.5  50  7.5818  3.8369  0.45    0.0062  -0.0448  0.2154 0.2695  6.8035  0.0086  0.1     6.5     0.      -0.1001    -0.001084  -0.003566  3.7443  -1.0603  -0.0016   -0.001523  0.005379  0.309  300  3.8646  0.      -0.7195  1.387    800.        0.4574  0.3419  0.4571  0.4571  0.7001  0.4682
5     -4.5143  0.0484 -0.101  -0.0484 0.101  1.7228   6.2856  1.06  2.9169  -2.1  -0.5  50  7.5818  3.8376  0.45    0.0029  -0.0424  0.2154 0.2695  7.0389  0.      0.1     6.5     0.      -0.1001    -0.00101   -0.00364   3.709   -0.9872   0.       -0.00144   0.003223  0.321  300  3.7292  0.      -0.656   1.433    800.        0.4584  0.3435  0.4535  0.4535  0.7     0.4517
7.5   -5.0009  0.022  -0.101  -0.022  0.101  1.5737   6.5428  1.06  2.932   -2.1  -0.5  50  7.5818  3.838   0.45    0.0007  -0.0348  0.2154 0.2695  7.4666  0.      0.1     6.5     0.      -0.1       -0.000964  -0.003686  3.6632  -0.8274   0.       -0.001369  0.001134  0.329  300  2.3763  0.      -0.5202  1.46     800.        0.4601  0.3459  0.4471  0.4471  0.7     0.4167
10    -5.3461  0.0124 -0.1    -0.0124 0.1    1.5265   6.7415  1.06  2.9396  -2.1  -0.5  50  7.5818  3.838   0.45    0.0003  -0.0253  0.2154 0.2695  7.77    0.      0.1     6.5     0.      -0.1       -0.00095   -0.0037    3.623   -0.7053   0.       -0.001361  0.000515  0.33   300  1.7679  0.      -0.4068  1.464    800.        0.4612  0.3474  0.4426  0.4426  0.7     0.3755
""")

class ChiouYoungs2014PEER(ChiouYoungs2014):
    """
    This implements the Chiou & Youngs (2014) GMPE for use with the PEER
    tests. In this version the total standard deviation is fixed at 0.65
    """
    #: Only the total standars deviation is defined
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL,
    ])
    #: The PEER tests requires only PGA
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([
        PGA,
    ])

    def _get_stddevs(self, sites, rup, C, stddev_types, ln_y_ref, exp1, exp2):
        """
        Returns the standard deviation, which is fixed at 0.65 for every site
        """
        ret = []
        for stddev_type in stddev_types:
            assert stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
            if stddev_type == const.StdDev.TOTAL:
                # eq. 13
                ret.append(0.65 * np.ones_like(sites.vs30))
        return ret

class ChiouYoungs2014NearFaultEffect(ChiouYoungs2014):
    """
    This implements the Chiou & Youngs (2014) GMPE include the near fault
    effect prediction. In this version, we add the distance measure, rcdpp
    for directivity prediction.

    """
    #: Required distance measures are RRup, Rjb, Rx, and Rcdpp
    REQUIRES_DISTANCES = set(('rrup', 'rjb', 'rx', 'rcdpp'))

    def _get_centered_cdpp(self, dists):
        """
        Get directivity prediction parameter centered on the avgerage
        directivity prediction parameter.

        """
        centered_dpp = dists.rcdpp

        return centered_dpp


# -----------------------------------------------------------------------------
#   @author:
#       Tingwu Wang
#   @brief:
# -----------------------------------------------------------------------------
import numpy as np

from .base_reward import base_reward_network
from mbbl.config import init_path
from mbbl.util.common import tf_networks
from mbbl.util.common import tf_utils
from mbbl.util.common import logger
from mbbl.util.il import expert_data_util
import tensorflow as tf

class reward_network(base_reward_network):
    '''
        @brief:
    '''

    def __init__(self, args, session, name_scope,
                 observation_size, action_size):
        '''
            @input:
                @ob_placeholder:
                    if this placeholder is not given, we will make one in this
                    class.

                @trainable:
                    If it is set to true, then the policy weights will be
                    trained. It is useful when the class is a subnet which
                    is not trainable
        '''
        super(reward_network, self).__init__(
            args, session, name_scope, observation_size, action_size
        )
        self._base_dir = init_path.get_abs_base_dir()
        # load the expert data
        self._expert_trajectory_obs = expert_data_util.load_expert_trajectory(
            self.args.expert_data_name, self.args.traj_episode_num
        )

    def build_network(self):
        self._build_ph()
        self._tensor = {}

        self._tensor['normalized_start_state'] = (
            self._input_ph['start_state'] -
            self._whitening_operator['state_mean']
        ) / self._whitening_operator['state_std']
        self._tensor['net_input'] = self._tensor['normalized_start_state']

        # the mlp for policy
        network_shape = [self._observation_size] + \
            self.args.reward_network_shape + [1]
        num_layer = len(network_shape) - 1
        act_type = \
            [self.args.reward_activation_type] * (num_layer - 1) + [None]
        norm_type = \
            [self.args.reward_normalizer_type] * (num_layer - 1) + [None]
        init_data = []
        for _ in range(num_layer):
            init_data.append(
                {'w_init_method': 'normc', 'w_init_para': {'stddev': 1.0},
                 'b_init_method': 'constant', 'b_init_para': {'val': 0.0}}
            )
        # init_data[-1]['w_init_para']['stddev'] = 0.01  # the output layer std
        self._MLP = tf_networks.MLP(
            dims=network_shape, scope='discriminator_mlp', train=True,
            activation_type=act_type, normalizer_type=norm_type,
            init_data=init_data
        )

        self._tensor['logits'] = self._MLP(self._tensor['net_input'])
        self._tensor['discriminator_output'] = \
            tf.nn.sigmoid(self._tensor['logits'])
        # the self.discriminator_output is the sigmoid(logit)
        self._tensor['logOfD'] = \
            tf.log(self._tensor['discriminator_output'] + 1e-8)
        self._tensor['logOf1minusD'] = \
            tf.log(1 - self._tensor['discriminator_output'] + 1e-8)

        self._tensor['reward_output'] = tf.minimum(
            -self._tensor['logOf1minusD'], self.args.GAN_reward_clip_value
        )

    def build_loss(self):
        self._update_operator = {}
        self._input_ph['if_expert_data'] = tf.placeholder(tf.float32, [None, 1],
                                                          name='observation_gt')
        self._tensor['if_fake_data'] = 1 - self._input_ph['if_expert_data']

        # calculate the entropy
        self._update_operator['entropy'] = tf.reduce_mean(
            tf_utils.logit_bernoulli_entropy(self._tensor['logits'])
        )
        self._update_operator['entropy_loss'] = \
            -self.args.GAN_ent_coeff * self._update_operator['entropy']

        self._update_operator['loss'] = \
            -tf.reduce_mean(self._tensor['if_fake_data'] *
                            self._tensor['logOf1minusD']) + \
            -tf.reduce_mean(self._input_ph['if_expert_data'] *
                            self._tensor['logOfD']) + \
            self._update_operator['entropy_loss']

        # logging stats, real traj should be 1
        self._tensor['expert_traj_accuracy'] = tf.reduce_sum(
            tf.to_float(self._tensor['discriminator_output'] > 0.5)
            * self._input_ph['if_expert_data']
        ) / tf.reduce_sum(self._input_ph['if_expert_data'])

        self._tensor['expert_average_reward'] = tf.reduce_sum(
            self._tensor['reward_output'] * self._input_ph['if_expert_data']
        ) / tf.reduce_sum(self._input_ph['if_expert_data'])

        self._tensor['agent_traj_accuracy'] = tf.reduce_sum(
            tf.to_float(self._tensor['discriminator_output'] < 0.5) *
            self._tensor['if_fake_data']
        ) / tf.reduce_sum(self._tensor['if_fake_data'])  # fake traj should be 0

        self._tensor['agent_average_reward'] = tf.reduce_sum(
            self._tensor['reward_output'] * self._tensor['if_fake_data']
        ) / tf.reduce_sum(self._tensor['if_fake_data'])

        self._update_operator['update_gan_op'] = tf.train.AdamOptimizer(
            learning_rate=self.args.reward_lr
        ).minimize(self._update_operator['loss'])

    def train(self, data_dict, replay_buffer, training_info={}):
        """ @brief:
            The GAN training of the reward function
        """

        agent_obs = data_dict['start_state']
        training_stats = []

        for i_epoch in range(self.args.reward_epochs):
            # we start from 1 since the first state will be fixed to be 0 reward
            agent_sample_id = self._npr.randint(
                1, agent_obs.shape[0], self.args.gan_timesteps_per_epoch
            )
            # sample and process the positive data-samples
            expert_sample_id = self._npr.randint(
                0, self._expert_trajectory_obs.shape[0],
                self.args.positive_negative_ratio *
                self.args.gan_timesteps_per_epoch
            )

            feed_dict = {
                self._input_ph['if_expert_data']: np.reshape(
                    np.concatenate(
                        [np.zeros(len(agent_sample_id)),
                         np.ones(len(expert_sample_id))],
                    ),
                    [-1, 1]
                ),
                self._input_ph['start_state']: np.concatenate(
                    [agent_obs[agent_sample_id],
                     self._expert_trajectory_obs[expert_sample_id]]
                )
            }

            # train the network
            fetch_dict = {
                'update_op': self._update_operator['update_gan_op'],
                'reward_gan_loss': self._update_operator['loss'],
                'expert_traj_accuracy': self._tensor['expert_traj_accuracy'],
                'agent_traj_accuracy': self._tensor['agent_traj_accuracy'],
                'expert_average_reward': self._tensor['expert_average_reward'],
                'agent_average_reward': self._tensor['agent_average_reward'],
                'entropy': self._update_operator['entropy']
            }
            i_training_stats = self._session.run(fetch_dict, feed_dict=feed_dict)
            training_stats.append(i_training_stats)

        training_stats = {
            key: np.mean([training_stats[i_epoch][key]
                          for i_epoch in range(len(training_stats))])
            for key in training_stats[-1] if key != 'update_op'
        }
        self._set_whitening_var(data_dict['whitening_stats'])
        return training_stats

    def eval(self, data_dict):
        pass

    def pred(self, data_dict):
        logger.info('This function should not be used!')
        reward = []
        for i_data in range(len(data_dict['action'])):
            i_reward = self._env.reward(
                {key: data_dict[key][i_data]
                 for key in ['start_state', 'action']}
            )
            reward.append(i_reward)
        return np.stack(reward), -1, -1

    def use_groundtruth_network(self):
        return False

    def generate_rewards(self, rollout_data):
        """@brief:
            This function should be called before _preprocess_data
        """
        for path in rollout_data:
            # the predicted value function (baseline function)
            path["raw_rewards"] = path['rewards']  # preserve the raw reward
            # from mbbl.util.common.fpdb import fpdb; fpdb().set_trace()

            # generate the observation pairs
            ob_pairs = path['obs'][1: len(path['obs'])]

            path["rewards"] = self._session.run(
                self._tensor['reward_output'],
                feed_dict={self._input_ph['start_state']: ob_pairs}
            ).flatten()
            assert len(path["rewards"]) == len(path['raw_rewards'])
        return rollout_data


""" Module for mean-flux regulation
Port of K-G Lee IDL code
"""
import pdb

import numpy as np

from astropy.table import Table
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.modeling import FittableModel, Parameter

from pyigm.fN import tau_eff as pyteff

try:
    basestring
except NameError:  # For Python 3
    basestring = str

def fit_forest(wave, flux, sigma, zqso, wavemin=3200., normfac=None,
               user_mask=None, mask_dlas=None, coord=None,
               for_lohi=(1041.,1185)):
    """  Perform mean-flux regulation on an input spectrum

    Parameters
    ----------
    wave : ndarray
      Observed wavelength values in Ang
    flux : ndarray
      flux array
    sigma : ndarray
      sigma array.  0 values are masked
    zqso : float
      Quasar emission redshift
    wavemin : float, optional
      Minimum wavelength to include in analysis
    normfac : float, optional
      optional scaling factor
    user_mask : list (of lists)
      User can mask out regions in the anaylsis `by-hand'
    mask_dlas : Table or str, optional
      Table of DLAs to mask based on coordinates and redshift
        Required columns are RA, DEC, z
      Can be the filename of the table to input
    coord : SkyCoord, optional
      Coordinate of the sightline.  Used to match against DLAs
    for_lohi : tuple
      wavemin, wavemax that define the forest for analysis

    Returns
    -------
    new_ff : ndarray
      Mean-flux regulated flux
    parm : FittableModel
      astropy model of the continuum
    """
    from astropy.modeling import fitting

    # invvar
    ivar = (sigma != 0.) / (sigma**2 + (sigma == 0))

    # Scale?
    if normfac is not None:
        ivar = ivar * normfac**2
        flux = flux / normfac

    # User Mask?
    if user_mask is not None:
        for imask in user_mask:
            msk = (wave >= imask[0]) & (wave <= imask[1])
            ivar[msk] = 0.

    # DLA mask?
    if mask_dlas is not None:
        if coord is None:
            raise IOError("Must input coord to mask DLAs")
        # Load DLAs, as needed
        if isinstance(mask_dlas, basestring):
            dlas = Table.read(mask_dlas)
        elif isinstance(mask_dlas, Table):
            dlas = mask_dlas
            # Check for coord column
            if 'coord' not in dlas.keys():
                dla_coord = SkyCoord(ra=dlas['RA'], dec=dlas['DEC'], unit='deg')
            else:
                dla_coord = dlas['coord']
        else:
            raise IOError("Not ready for this type of DLA input")

        # Search
        sep = coord.separation(dla_coord)
        match = np.where(sep < 2*u.arcsec)[0]
        for imatch in match:
            zdla = dlas['z'][imatch]
            skip_fg = np.abs(wave-1215.6701*(1+zdla)) < 100.
            ivar[skip_fg] = 0.
            #print, 'Removed a DLA'

    # Isolate forest for analysis
    lambda_r = wave / (1+zqso)
    forestrange = (lambda_r >= for_lohi[0]) & (lambda_r <= for_lohi[1]) & (
                    wave > wavemin)
    lamb_forest = lambda_r[forestrange]
    z_for = (lamb_forest/1215.67)*(1.+zqso) -1.  # Redshift at each pixel
    fforest = flux[forestrange]
    ivarforest = ivar[forestrange]

    # Estimate weights for each pixel
    var_F = forestvar(z_for) * (np.exp(-pyteff.lyman_alpha_obs(z_for)))**2
    #var_F = forestvar(z_for) * (np.exp(-old_taueff_evo(z_for)))**2
    var_noise = (ivarforest != 0) / (ivarforest + (ivarforest == 0))
    var_total = var_F + var_noise
    weights_forest = (var_total != 0) / (var_total + (var_total == 0))

    # But need to make sure that masked pixels remain masked...
    maskedpix = ivarforest == 0
    weights_forest[maskedpix] = 0

    # Astropy modeling
    model = mflux_tauevo(p0=0., p1=0., zqso=zqso, lamb_piv=1113.)
    fitter = fitting.LevMarLSQFitter()
    parm = fitter(model, lamb_forest, fforest, weights=weights_forest)

    # Apply
    full_forest = lambda_r < 1220.
    new_ff = flux
    new_ff[full_forest] = flux[full_forest] / mfluxcorr(lambda_r[full_forest],
                                                        parm.p0.value, parm.p1.value,
                                                        lamb_piv=parm.lamb_piv.value)
    # Return
    return new_ff, parm

def forestvar(z_in):
    """  Return intrinsic variance of LyaF variance for weighting. This
    estimate is roughly from McDonald et al 2006

    Parameters
    ----------
    z_in : float or ndarray

    Returns
    -------
    fvar : float or ndarray
      Variance
    """
    fvar = 0.065 * ((1.+z_in)/(1.+2.25))**3.8
    # Return
    return fvar

def mfluxcorr(lambda_r, p0, p1, lamb_piv=1113.):
    """ Correction factor to power-law fit
    Parameters
    ----------
    lambda_r
    p
    lamb_piv

    Returns
    -------

    """
    lamb_piv = 1113. # This is the pivot point in the restframe spectrum
    return p0 + p1*(lambda_r/lamb_piv - 1.)

class mflux_tauevo(FittableModel):
    """ Mean flux evolution * exp(delta*(lambda/1280-1)),
    Meant for use with astropy.modeling to correct the fitted Lya forest continuum
    Abscissa parameter x is the restframe wavelength, and free parameters p0, p1
    set the power law.
    This function NEEDS the quasar redshift zqso to be set

    input: lambda_r :: Rest wavelength Assumed in Angstroms
    output: absorbed, normalized flux
    Parameters: logN,b,z,wrest,f,gamma,fwhm
    """
    inputs = ('lambda_r',)
    outputs = ('flux',)

    # Free parameters (generally)
    p0 = Parameter()
    p1 = Parameter()

    # Fixed parameters
    zqso = Parameter(fixed=True)
    lamb_piv = Parameter(fixed=True)

    @staticmethod
    def evaluate(lambda_r, p0, p1, zqso, lamb_piv): #logN,b,z,wrest,f,gamma,fwhm):
        zfor = (lambda_r/1216.) * (1. + zqso) - 1.
        tau = pyteff.lyman_alpha_obs(zfor)
        #tau = old_taueff_evo(zfor)
        fmean = np.exp(-1*tau)

        mfluxtauevo = fmean * mfluxcorr(lambda_r, p0, p1, lamb_piv=lamb_piv)
        return mfluxtauevo

def old_taueff_evo(z):
    """ F-G taueff.
    Mainly for testing

    Parameters
    ----------
    z

    Returns
    -------

    """
    tauevo  = 0.001845 * (1+z)**3.924
    return tauevo


import numpy as np
from scipy import interpolate
import matplotlib
from datetime import datetime
from scipy.stats.mstats import mquantiles
import os
from astropy.io import fits
import glob
from parameters import *
import sys
import argparse
from astropy.table import Table
##sys.path.insert(0, 'python')
# matplotlib.rc('text', usetex=True)
# matplotlib.rcParams['text.latex.preamble'] = [r"\usepackage{amsmath}",
#                                    r"\usepackage{color}"]
# matplotlib.use('Agg')
from ptemcee import Sampler as PTSampler
import matplotlib.pyplot as plt
##import corner
startTime = datetime.now()

parser = argparse.ArgumentParser(description="Calculate radial velocities and stellar parameters of WEAVE target spectra.")

parser.add_argument("--infile", type=str, required=True, help="input file", nargs=1)
parser.add_argument("--outdir", type=str, required=True, help="output directory", nargs=1)
parser.add_argument("--targlist", type=str, required=True, help="path to list of FIBREIDs or TARGIDs to be analysed. To analyse all BA stars, enter 'all'.", nargs=1)
parser.add_argument("--params", type=str, required=False, default='parameters.py', help="path to parameter file", nargs=1)

#parser.add_argument("--setups", type=str, default=None, required=True, help="input setups", nargs='*')

args = parser.parse_args()
write_directory = args.outdir[0]
target_list = args.targlist[0]
data_file = args.infile[0]

# processing (cropping, smoothing, rebinning, rotational broadening) templates if required

if process_templates==True:
	print('Beginning processing of templates')
	from PyAstronomy import pyasl
	import scipy.stats
	templatelist = np.genfromtxt(template_list_file, dtype=None, encoding=None)

	# Ensuring the number of templates to be processed matches the number of points on the grid.
	notemps = len(templatelist)
	d_size = len(Teff) * len(logg) * len(vsini)
	if notemps < d_size:
		print('ERROR: The amount of templates to be processed is less than the amount of points in the grid. Add more templates to the list or reduce the grid as necessary (see PARAMETER BOUNDARIES in parameters file).')
		sys.exit()
	if notemps > d_size:
		print('ERROR: The amount of templates to be processed is greater than the amount of points in the grid. Remove templates from the list or expand the grid as necessary (see PARAMETER BOUNDARIES in parameters file).')
		sys.exit()

	def restrict_range(w, f):		# to crop the templates
		our_range = (w > template_crop_min_wav) & (w < template_crop_max_wav)
		w, f = w[our_range], f[our_range]
		return w, f

	def smooth(w, f, sig):		# to smooth/broaden templates to match resolution of observed spectrum
		f = pyasl.broadGaussFast(w, f, sig)
		return f

	def rotbroad(w, f, vsini):		# to rotationally broaden the templates
		f = pyasl.rotBroad(w, f, 0.6, float(vsini))
		return f

	def rebin(w, f, samp):		# to rebin the templates to match the sampling of the observed spectrum
		f, bin_edges, binnumber = scipy.stats.binned_statistic(w, f, statistic = 'mean', bins = (w[-1] - w[0]) / samp)
		bin_width = bin_edges[1] - bin_edges[0]
		w = bin_edges[1:] - bin_width/2
		return w, f

	def process_template(w, f, sig=sigma, samp=sampling):
		if sigma != None:
			f = smooth(w, f, sig)
		if sampling != None:
			w, f = rebin(w, f, samp)
		file = open(template_write_directory + '/'+ os.path.splitext(os.path.basename(i))[0] + '_processed', "w")
		for index in range(len(w)):
			file.write(str(w[index]) + " " + str(f[index]) + "\n")
		file.close()
		return w, f

	def process_template_vsini(w, f, vsini, sig=sigma, samp=sampling):
		f = rotbroad(w, f, vsini)
		if sigma != None:
			f = smooth(w, f, sig)
		if sampling != None:
			w, f = rebin(w, f, samp)
		file = open(template_write_directory + '/' + os.path.splitext(os.path.basename(i))[0] + '_processed_vsini' + str(int(vsini)), "w")
		for index in range(len(w)):
			file.write(str(w[index]) + " " + str(f[index]) + "\n")
		file.close()
		return w, f

	def write_to_grid(templatename, f, t_ind=temp_ind, l_ind=logg_ind, v_ind=vsini_ind, vsini0=None):
		# finding the teff and logg information from the template name, to write to corresponding grid point
		t1 = float(templatename[temp_ind[0]:temp_ind[1]])
		l1 = float(templatename[logg_ind[0]:logg_ind[1]])
		if vsini0==None:
			v1 = float(templatename[vsini_ind[0]:vsini_ind[1]])
		else:
			v1 = float(vsini0)

		print('Teff: ' + str(t1) + ', logg: ' + str(l1) + ', vsini: ' + str(v1))

		d[np.where(Teff == t1)[0][0], np.where(logg == l1)[0][0], np.where(vsini == v1)[0][0]] = f

	t_wavelength_0, t_flux_0 = np.genfromtxt(templatelist[0], unpack=True, usecols=(wav_ind, flux_ind), dtype=None, encoding=None)
	t_wavelength_0, t_flux_0 = restrict_range(t_wavelength_0, t_flux_0)
	if sampling != None:
		t_wavelength_0, t_flux_0 = rebin(t_wavelength_0, t_flux_0, sampling)

	d = np.zeros((len(Teff), len(logg), len(vsini), len(t_wavelength_0)))
	d_filled = 0

	for i in templatelist:
		print('Processing: ' + i)
		t_wavelength, t_flux = np.genfromtxt(i, unpack=True, usecols=(wav_ind, flux_ind), dtype=None, encoding=None)
		t_wavelength, t_flux = restrict_range(t_wavelength, t_flux)
		if rotbroads == None:
			_, t_flux = process_template(t_wavelength, t_flux, sigma, sampling)
			print('template processed, writing to grid at:')
			write_to_grid(i, t_flux)
			d_filled += 1
			print('grid is ' + str("{0:.1f}".format(float(d_filled)/float(d_size) * 100.)) + ' percent full')

		else:
			print('processing for vsini: 0')
			_, t_flux_v0 = process_template(t_wavelength, t_flux, sig=sigma, samp=sampling)
			print('template processed, writing to grid at:')
			write_to_grid(i, t_flux_v0, vsini0=0.)
			d_filled += 1
			print('grid is ' + str("{0:.1f}".format(float(d_filled)/float(d_size) * 100.)) + ' percent full')

			for j in rotbroads:
				if j == 0:
					continue
				print('processing for vsini: ' + str(j))
				_, t_flux_vj = process_template_vsini(t_wavelength, t_flux, j, sig=sigma, samp=sampling)
				print('template processed, writing to grid at:')
				write_to_grid(i, t_flux_vj, vsini0=j)
				d_filled += 1
				print('grid is ' + str("{0:.1f}".format(float(d_filled)/float(d_size) * 100.)) + ' percent full')

	np.save(template_write_directory + 'template_grid.npy', d)

# creating /plots and /results folders in write_directory if they dont already exist
if not os.path.exists(write_directory + '/plots'):
    os.mkdir(write_directory + '/plots')
    print(write_directory + 'plots folder created')
if not os.path.exists(write_directory + '/results'):
    os.mkdir(write_directory + '/results')
    print(write_directory + 'results folder created')

if process_templates==False:
	# location of folder containing this script
	templatedirectory = os.path.dirname(os.path.abspath(__file__)) + '/templates/'
	# loading the template flux grid
	flux_data_all = np.load(templatedirectory + 'template_grid.npy')
	#loading the wavelength data
	templatewavelength = np.loadtxt(templatedirectory + 'wavelength_data.dat')
else:
	# the first template for the template wavelength data
	flux_data_all = np.load(template_write_directory + 'template_grid.npy')
	templatewavelength = t_wavelength_0

template_min_wav = min_wav*(1.0 + (min(drv)/299792.458)) - 10.
template_max_wav = max_wav*(1.0 + (max(drv)/299792.458)) + 10.

templatemask = (templatewavelength > template_min_wav) & (templatewavelength < template_max_wav)
flux_data = np.zeros((len(Teff), len(logg), len(vsini), len(templatewavelength[templatemask])))
templatewavelength = templatewavelength[templatemask]

for ii in range(len(Teff)):
	for jj in range(len(logg)):
		for kk in range(len(vsini)):
			flux_data[ii,jj,kk] = flux_data_all[ii,jj,kk][templatemask]

ipo = interpolate.RegularGridInterpolator((Teff, logg, vsini), flux_data, method='linear')

# grid of parameters for walker initial positions
ini_grid = [Teff, logg, vsini, drv, slopes, intercepts]
ndim=6

# define edges of parameter space
teffmin, teffmax, loggmin, loggmax, vsinimin, vsinimax, rvmin, rvmax, slopemin, slopemax, interceptmin, interceptmax = min(Teff), max(Teff), min(logg), max(logg), min(vsini), max(vsini), min(drv), max(drv), min(slopes), max(slopes), min(intercepts), max(intercepts)

def model(X, wavelength):
	i, j, k, l, m, n = X
	# interpolating template grid with teff (i), logg (j), vsini (k) trial parameter
	templateflux = ipo([i, j, k])[0]
	# interpolating on wavelength axis for trial RV (l)
	fi = interpolate.interp1d(templatewavelength*(1.0 + l/299792.458), templateflux)
	return fi(wavelength)

def lnprior(X):
	i, j, k, l, m, n = X
	# flat prior, edges should corespond to template grid
	if (teffmin <= i <= teffmax) & (loggmin <= j <= loggmax) & (vsinimin <= k <= vsinimax) & (rvmin <= l <= rvmax) & (slopemin <= m <= slopemax) & (interceptmin <= n <= interceptmax):
		return 0.0
	else:
		return -np.inf

def lnlike(X, wavelength, flux, noisespec, mask):
	i, j, k, l, m, n = X
	z = m, n
	f = np.poly1d(z)
	if exclude_region == False:
		return -(np.sum((flux - (model(X, wavelength)*f(wavelength)))**2/(2*PPRE*noisespec**2)))
	else:
		return -(np.sum((flux[mask] - (model(X, wavelength)*f(wavelength))[mask])**2/(2*PPRE*noisespec[mask]**2)))

def mcmc_one(t):

	print("Processing: " + t)
	targ_start = datetime.now()

	# checking if spectrum file exists, and if result already written
	if os.path.exists(write_directory + '/results/' + os.path.splitext(os.path.basename(t))[0] + '_results'):
		print('WARNING: result for '+t+' already exists, skipping')
		return

	# specify unnormalised calibrated flux
	with fits.open(data_file) as ALLDATA:
		final_spectra = ALLDATA[1].data[info['TARGID'] == t][0]
		spectra_before_sky_subtraction = ALLDATA[3].data[info['TARGID'] == t][0]
		Calibration_function = ALLDATA[5].data[info['TARGID'] == t][0]
		try:
			idx=list(ALLDATA[6].data['FIBREID']).index(t)
		except:
			idx=list(ALLDATA[6].data['TARGID']).index(t)
		NSPEC = ALLDATA[6].data['Nspec'][idx]
		FIBREID = ALLDATA[6].data['FIBREID'][idx]
		CNAME = ALLDATA[6].data['CNAME'][idx]
	flux = final_spectra * Calibration_function * 1.0e18

	# restrict to desired wavelength range
	flux = flux[targetmask]

	# create corresponsing noise spectrum
	noisespec = np.sqrt((2.*spectra_before_sky_subtraction - final_spectra)*Calibration_function*1.0e18)
	noisespec = noisespec[targetmask]

	# mask for excluding a wavelength region
	if exclude_region == True:
		mask = np.zeros(len(wavelength), dtype=bool)
		for line in lines:
			mask |= (wavelength >= line[0]) & (wavelength <= line[1])
	else:
		mask = np.ones(len(wavelength), dtype=bool)

	# choose initial walker positions
	pos = [[[np.random.choice(i) for i in ini_grid] for i in range(nwalkers)] for i in range(ntemps)]

	# initialise MCMC sampler
	sampler = PTSampler(ntemps=ntemps, nwalkers=nwalkers, dim=ndim, logl=lnlike, logp=lnprior, Tmax=np.inf, loglargs=[wavelength, flux, noisespec, mask])

	# run MCMC sampler for burn period
	if progress_bar==True:
		print("running burn")
	pos, prob, state = sampler.run_mcmc(pos, burn, adapt=True, progress_bar=progress_bar)
	# reset sampler, run MCMC sampler for run period with walkers starting at their positions at the end of burn
	sampler.reset()
	if progress_bar==True:
		print("running runs")
	sampler.run_mcmc(pos, runs, adapt=True, progress_bar=progress_bar)
	samples=sampler.chain[0, :, :, :].reshape((-1, ndim))

	# plot walker paths
	ylabels = ['Teff', 'logg', 'vsini', 'RV', 'slope', 'intercept']
	for m in range(ndim):
		plt.subplot(ndim,1,m+1)
		plt.plot(sampler.chain[0,:,:,m].transpose(), alpha=0.2)
		plt.ylabel(ylabels[m])
	plt.xlabel('Step')
	plt.savefig(write_directory + '/plots/' + t + '_walkers.png', bbox_inches='tight')
	plt.close()

	# calculate 16th, 50th, 84th quantiles of the parameter samples
	quantiles = mquantiles(samples, prob=[0.16, 0.50, 0.84], axis=0)

	targetname = os.path.splitext(os.path.basename(t))[0]
	acceptance_r = np.mean(sampler.acceptance_fraction)

	# print acceptance fraction, should be between 0.2-0.5 for efficient sampling
	print("Mean acceptance fraction: {0:.3f}"
                .format(acceptance_r))

	# The parameter results
	Teff_r = quantiles[1][0]
	Teffminus_r = quantiles[1][0] - quantiles[0][0]
	Teffplus_r = quantiles[2][0] - quantiles[1][0]
	logg_r = quantiles[1][1]
	loggminus_r = quantiles[1][1] - quantiles[0][1]
	loggplus_r = quantiles[2][1] - quantiles[1][1]
	vsini_r = quantiles[1][2]
	vsiniminus_r = quantiles[1][2] - quantiles[0][2]
	vsiniplus_r = quantiles[2][2] - quantiles[1][2]
	RV_r = quantiles[1][3]
	RVminus_r = quantiles[1][3] - quantiles[0][3]
	RVplus_r = quantiles[2][3] - quantiles[1][3]
	slope_r = quantiles[1][4]
	slopeminus_r = quantiles[1][4] - quantiles[0][4]
	slopeplus_r = quantiles[2][4] - quantiles[1][4]
	intercept_r = quantiles[1][5]
	interceptminus_r = quantiles[1][5] - quantiles[0][5]
	interceptplus_r = quantiles[2][5] - quantiles[1][5]

	#fig = corner.corner(samples, quantiles=[0.16, 0.50, 0.84], labels=['Teff', 'log(g)', 'vsini', 'RV', 'slope', 'intercept'], show_titles=True, title_kwargs={"fontsize": 10}, plot_datapoints=True, plot_contours=True, auto_bars=True, data_kwargs={"alpha": 0.005})
	# fig.savefig(write_directory + t + '_cornerplot.png', bbox_inches='tight')
	# plt.close()

	# parameters of best fit (for plotting)
	Xp = Teff_r, logg_r, vsini_r, RV_r, slope_r, intercept_r
	fp = np.poly1d(Xp[-2:])
	fitp = fp(wavelength)

	# plot spectrum and best-fit
	plt.plot(wavelength, flux, wavelength, model(Xp, wavelength) * fitp)
	if exclude_region == True:
		for n,line in enumerate(lines):
			if n == 0:
				plt.vlines([line[1]], flux.min()*0.8, flux.max()*1.2, colors='r', linestyles='dashed')
			elif n == (len(lines)-1):
				plt.vlines([line[0]], flux.min()*0.8, flux.max()*1.2, colors='r', linestyles='dashed')
			else:
				plt.vlines([line[0], line[1]], flux.min()*0.8, flux.max()*1.2, colors='r', linestyles='dashed')

	plt.xlim(min_wav, max_wav)
	plt.ylim(flux.min()*0.8, flux.max()*1.2)
	plt.xlabel(r'Wavelength ($\AA$)')
	plt.ylabel('Calibrated counts')
	plt.savefig(write_directory + '/plots/' + t + '_spectrum.png', bbox_inches='tight')
	plt.close()

	# plot mapping function
	plt.plot(wavelength, fitp, wavelength, flux / model(Xp, wavelength))
	plt.xlim(min_wav, max_wav)
	plt.xlabel(r'Wavelength ($\AA$)')
	plt.ylabel('Calibrated counts')
	plt.savefig(write_directory + '/plots/' + t + '_mapping_function.png', bbox_inches='tight')
	plt.close()

	# write results
	tab = open(write_directory + '/results/' + t +  '_results', "w")
	tab.write(np.str(NSPEC) + " " + np.str(FIBREID) + " " + np.str(CNAME) + " " + t + " " + np.str(acceptance_r) + " " + np.str(Teff_r) + " " + np.str(Teffminus_r) + " " + np.str(Teffplus_r) + " " + np.str(logg_r) + " " + np.str(loggminus_r) + " " + np.str(loggplus_r) + " " + np.str(vsini_r) + " " + np.str(vsiniminus_r) + " " + np.str(vsiniplus_r) + " " + np.str(RV_r) + " " + np.str(RVminus_r) + " " + np.str(RVplus_r) + " " + np.str(slope_r) + " " + np.str(slopeminus_r) + " " + np.str(slopeplus_r) + " " + np.str(intercept_r) + " " + np.str(interceptminus_r) + " " + np.str(interceptplus_r) + "\n")
	tab.close()

	targ_end = datetime.now() - targ_start
	print(targ_end)

def make_output_fits():
	output_files = glob.glob(write_directory + 'results/*_results')
	all_res = []
	for i in output_files:
		with open(i) as outf:
			all_res.append(outf.read().split())
	t = Table(rows=all_res, names=('NSPEC', 'FIBREID', 'CNAME', 'TARGID', 'Acceptance', 'Teff', 'Teff_minus', 'Teff_plus', 'logg', 'logg_minus', 'logg_plus', 'vsini', 'vsini_minus', 'vsini_plus', 'RV', 'RV_minus', 'RV_plus', 'slope', 'slope_minus', 'slope_plus', 'intercept', 'intercept_minus', 'intercept_plus'))
	t.write(write_directory+'results/'+os.path.splitext(os.path.basename(args.infile[0]))[0]+'_ptmcmc.fits', format='fits', overwrite=False)

if __name__ ==  '__main__':
	# checking if results table already exists
	if os.path.exists(write_directory+'results/'+os.path.splitext(os.path.basename(args.infile[0]))[0]+'_ptmcmc.fits'):
		print('WARNING: result table already exists, ending process.')
		sys.exit()
	if target_list == 'all':
		print("Processing all BA stars in fits file")
		BA = []
		with fits.open(data_file) as ALLDATA:
			for n,i in enumerate(ALLDATA[6].data['TARGID']):
				if 'LR-BA' in i:
					BA.append(i)
	else:
		print("Processing BA stars specified in target list")
		targlist = np.genfromtxt(target_list, dtype=None, encoding=None)
		BA = []
		with fits.open(data_file) as ALLDATA:
			for targ in targlist:
				try:
					idx=list(ALLDATA[6].data['FIBREID']).index(targ)
					BA.append(ALLDATA[6].data['TARGID'][idx])
				except:
					try:
						idx=list(ALLDATA[6].data['TARGID']).index(targ)
						BA.append(ALLDATA[6].data['TARGID'][idx])
					except:
						print(str(targ)+": Cant find either FIBREID or TARGID in input table.")

ALLDATA = None
info = None
wavelength = None
targetmask = None

def set_globals():
    global ALLDATA
    global info
    global wavelength
    global targetmask

    with fits.open(data_file) as ALLDATA:
	    head0 = ALLDATA[0].header
	    info = ALLDATA[6].data
	    data1 = ALLDATA[1].data
	    head1 = ALLDATA[1].header
	    wave0 = head1['CRVAL1']
	    increm = head1['CD1_1']
	    wavelength = np.array([wave0+increm*a for a in range(len(data1[1]))])
	    targetmask = (wavelength > min_wav) & (wavelength < max_wav)
	    wavelength = wavelength[targetmask]

if multiprocess==True:
	from multiprocessing import Pool
	if __name__ ==  '__main__':
		pool = Pool(processes=process_no, initializer=set_globals)
		it = pool.imap_unordered(mcmc_one, BA)
		for nn,i in enumerate(range(len(BA))):
			it.next()
		make_output_fits()
		print(datetime.now() - startTime)

else:
	set_globals()
	for nn,i in enumerate(BA):
		mcmc_one(i)
	make_output_fits()
	print(datetime.now() - startTime)

import numpy,import scipy,from scipy,from astropy


from __future__ import division
import numpy as np
from functools import partial
from builtins import zip

from pybasicbayes.distributions import DiagonalRegression, Gaussian, Regression

import pyhsmm
from pyhsmm.util.general import list_split

from pyslds.states import HMMSLDSStatesPython, HMMSLDSStatesEigen, HSMMSLDSStatesPython, HSMMSLDSStatesEigen
from pyslds.states import HMMCountSLDSStatesPython, HMMCountSLDSStatesEigen, HSMMCountSLDSStatesPython, \
    HSMMCountSLDSStatesEigen
from pyslds.util import gaussian_map_estimation, regression_map_estimation, gaussian_logprior, regression_logprior

class _SLDSMixin(object):
    def __init__(self,dynamics_distns,emission_distns,init_dynamics_distns,**kwargs):
        self.init_dynamics_distns = init_dynamics_distns
        self.dynamics_distns = dynamics_distns

        # Allow for a single, shared emission distribution
        if not isinstance(emission_distns, list):
            self._single_emission = True
            self._emission_distn = emission_distns
            self.emission_distns = [emission_distns] * len(self.dynamics_distns)
        else:
            assert isinstance(emission_distns, list) and \
                   len(emission_distns) == len(dynamics_distns)
            self._single_emission = False
            self.emission_distns = emission_distns

        super(_SLDSMixin,self).__init__(
            obs_distns=self.dynamics_distns,**kwargs)

    def generate(self, T=100, keep=True, with_noise=True, initial_condition=None, stateseq=None, **kwargs):
        s = self._states_class(model=self, T=T, initialize_from_prior=True, **kwargs)
        s.generate_states(with_noise=with_noise, initial_condition=initial_condition, stateseq=stateseq)
        data = self._generate_obs(s)
        if keep:
            self.states_list.append(s)
        return data + (s.stateseq,)

    def _generate_obs(self,s):
        if s.data is None:
            s.data = s.generate_obs()
        else:
            # TODO: Handle missing data
            raise NotImplementedError

        return s.data, s.gaussian_states

    def smooth(self, data, inputs=None, mask=None):
        self.add_data(data, inputs=inputs, mask=mask)
        s = self.states_list.pop()
        return s.smooth()

    @property
    def diagonal_noise(self):
        return all([isinstance(ed, DiagonalRegression) for ed in self.emission_distns])

    @property
    def has_missing_data(self):
        return any([s.mask is not None for s in self.states_list])

    def heldout_log_likelihood(self, test_masks=None):
        test_masks = [None] * len(self.states_list) if test_masks is None else test_masks
        assert len(test_masks) == len(self.states_list)

        hll = 0
        for mask, states in zip(test_masks, self.states_list):
            hll += states.heldout_log_likelihood(test_mask=mask)
        return hll

class _SLDSGibbsMixin(_SLDSMixin):
    def resample_parameters(self):
        self.resample_lds_parameters()
        self.resample_hmm_parameters()

    def resample_lds_parameters(self):
        self.resample_init_dynamics_distns()
        self.resample_dynamics_distns()
        self.resample_emission_distns()

    def resample_hmm_parameters(self):
        super(_SLDSGibbsMixin,self).resample_parameters()

    def resample_init_dynamics_distns(self):
        for state, d in enumerate(self.init_dynamics_distns):
            d.resample(
                [s.gaussian_states[0] for s in self.states_list
                    if s.stateseq[0] == state])
        self._clear_caches()

    def resample_dynamics_distns(self):
        zs = [s.stateseq[:-1] for s in self.states_list]
        xs = [np.hstack((s.gaussian_states[:-1], s.inputs[:-1]))
              for s in self.states_list]
        ys = [s.gaussian_states[1:] for s in self.states_list]

        for state, d in enumerate(self.dynamics_distns):
            d.resample(
                [(x[z == state], y[z == state])
                 for x, y, z in zip(xs, ys, zs)])
        self._clear_caches()

    def resample_emission_distns(self):
        if self._single_emission:
            data = [(np.hstack((s.gaussian_states, s.inputs)), s.data)
                    for s in self.states_list]
            mask = [s.mask for s in self.states_list] if self.has_missing_data else None

            if self.has_missing_data:
                self._emission_distn.resample(data=data, mask=mask)
            else:
                self._emission_distn.resample(data=data)
        else:
            for state, d in enumerate(self.emission_distns):
                data = [(np.hstack((s.gaussian_states[s.stateseq == state],
                                    s.inputs[s.stateseq == state])),
                         s.data[s.stateseq == state])
                        for s in self.states_list]

                mask = [s.mask[s.stateseq == state] for s in self.states_list] \
                    if self.has_missing_data else None

                if self.has_missing_data:
                    d.resample(data=data, mask=mask)
                else:
                    d.resample(data=data)

        self._clear_caches()

    def resample_obs_distns(self):
        pass  # handled in resample_parameters

    ### joblib parallel

    def _joblib_resample_states(self,states_list,num_procs):
        from joblib import Parallel, delayed
        import pyslds.parallel as parallel

        if len(states_list) > 0:
            joblib_args = list(map(self._get_joblib_pair, states_list))

            parallel.model = self
            parallel.args = list_split(joblib_args, num_procs)

            idxs = range(len(parallel.args))
            raw_stateseqs = Parallel(n_jobs=num_procs,backend='multiprocessing')\
                    (list(map(delayed(parallel._get_sampled_stateseq), idxs)))

            flatten = lambda lst: [x for y in lst for x in y]
            raw_stateseqs = flatten(raw_stateseqs)

            # since list_split might reorder things, do the same to states_list
            states_list = flatten(list_split(states_list, num_procs))

            for s, tup in zip(states_list, raw_stateseqs):
                s.stateseq, s.gaussian_states, s._normalizer = tup

class _SLDSVBEMMixin(_SLDSMixin):

    def _vb_E_step(self):
        # update the variational approximation for the states
        for state in self.states_list:
            state.vb_E_step()

    def _vb_M_step(self):
        # Update the HMM parameters
        self._M_step_init_state_distn()
        self._M_step_trans_distn()

        # Update the LDS parameters
        self._M_step_init_dynamics_distn()
        self._M_step_dynamics_distn()
        self._M_step_emission_distn()

    def _M_step_init_dynamics_distn(self):
        sum_tuples = lambda lst: list(map(sum, zip(*lst)))
        E_init_stats = lambda i, s: \
            tuple(s.expected_states[0, i] * stat for stat in s.E_init_stats)

        for state, d in enumerate(self.init_dynamics_distns):
            gaussian_map_estimation(sum_tuples(E_init_stats(state, s) for s in self.states_list), d)

    def _M_step_dynamics_distn(self):
        contract = partial(np.tensordot, axes=1)
        sum_tuples = lambda lst: list(map(sum, zip(*lst)))

        E_dyn_stats = lambda i, s: \
            tuple(contract(s.expected_states[:-1, i], stat) for stat in s.E_dynamics_stats)

        for state, d in enumerate(self.dynamics_distns):
            regression_map_estimation(sum_tuples(E_dyn_stats(state, s) for s in self.states_list), d)

    def _M_step_emission_distn(self):
        contract = partial(np.tensordot, axes=1)
        sum_tuples = lambda lst: list(map(sum, zip(*lst)))

        if self._single_emission:
            E_emi_stats = lambda s: \
                tuple(np.sum(stat, axis=0) for stat in s.E_emission_stats)
            stats = sum_tuples(E_emi_stats(s) for s in self.states_list)
            regression_map_estimation(stats, self._emission_distn)
        else:
            E_emi_stats = lambda i, s: \
                tuple(contract(s.expected_states[:, i], stat) for stat in s.E_emission_stats)
            for state, d in enumerate(self.emission_distns):
                regression_map_estimation(sum_tuples(E_emi_stats(state, s) for s in self.states_list), d)

    def VBEM_step(self, n_iter=1):
        for _ in range(n_iter):
            self._vb_E_step()
        self._vb_M_step()

    def VBEM_ELBO(self):
        # log p(theta)
        # todo: include transition distribution and init state distribution!
        elbo = np.sum([gaussian_logprior(id) for id in self.init_dynamics_distns])
        elbo += np.sum([regression_logprior(dd) for dd in self.dynamics_distns])

        if self._single_emission:
            elbo += regression_logprior(self.emission_distns[0])
        else:
            elbo += np.sum([regression_logprior(ed) for ed in self.emission_distns])

        # E_q [log p(z, x, y, theta)]
        elbo += sum(s.vb_elbo() for s in self.states_list)
        return elbo

class _SLDSMeanFieldMixin(_SLDSMixin):
    def meanfield_update_parameters(self):
        self.meanfield_update_init_dynamics_distns()
        self.meanfield_update_dynamics_distns()
        self.meanfield_update_emission_distns()
        super(_SLDSMeanFieldMixin, self).meanfield_update_parameters()

    def meanfield_update_init_dynamics_distns(self):
        sum_tuples = lambda lst: list(map(sum, zip(*lst)))
        E_stats = lambda i, s: \
            tuple(s.expected_states[0,i] * stat for stat in s.E_init_stats)

        for state, d in enumerate(self.init_dynamics_distns):
            d.meanfieldupdate(
                stats=sum_tuples(E_stats(state, s) for s in self.states_list))

    def meanfield_update_dynamics_distns(self):
        contract = partial(np.tensordot, axes=1)
        sum_tuples = lambda lst: list(map(sum, zip(*lst)))
        E_stats = lambda i, s: \
            tuple(contract(s.expected_states[1:,i], stat) for stat in s.E_dynamics_stats)

        for state, d in enumerate(self.dynamics_distns):
            d.meanfieldupdate(
                stats=sum_tuples(E_stats(state, s) for s in self.states_list))

    def meanfield_update_emission_distns(self):
        sum_tuples = lambda lst: list(map(sum, zip(*lst)))

        if self._single_emission:
            E_stats = lambda s: \
                tuple(np.sum(stat, axis=0) for stat in s.E_emission_stats)

            self._emission_distn.meanfieldupdate(
                stats=sum_tuples(E_stats(s) for s in self.states_list))
        else:
            contract = partial(np.tensordot, axes=1)
            E_stats = lambda i, s: \
                tuple(contract(s.expected_states[:, i], stat) for stat in s.E_emission_stats)

            for state, d in enumerate(self.emission_distns):
                d.meanfieldupdate(
                    stats=sum_tuples(E_stats(state, s) for s in self.states_list))

    def meanfield_update_obs_distns(self):
        pass  # handled in meanfield_update_parameters

    ### init
    def _init_mf_from_gibbs(self):
        # Now also update the emission and dynamics params
        for ed in self.emission_distns:
            if hasattr(ed, "_initialize_mean_field"):
                ed._initialize_mean_field()
        for dd in self.dynamics_distns:
            if hasattr(dd, "_initialize_mean_field"):
                dd._initialize_mean_field()

        for s in self.states_list:
            s._init_mf_from_gibbs()

    ### vlb

    def vlb(self, states_last_updated=False):
        vlb = 0.
        vlb += sum(s.get_vlb(states_last_updated) for s in self.states_list)
        vlb += self.trans_distn.get_vlb()
        vlb += self.init_state_distn.get_vlb()
        vlb += sum(d.get_vlb() for d in self.init_dynamics_distns)
        vlb += sum(d.get_vlb() for d in self.dynamics_distns)
        if self._single_emission:
            vlb += self._emission_distn.get_vlb()
        else:
            vlb += sum(d.get_vlb() for d in self.emission_distns)
        return vlb

class HMMSLDSPython(_SLDSGibbsMixin, _SLDSVBEMMixin, _SLDSMeanFieldMixin, pyhsmm.models.HMMPython):
    _states_class = HMMSLDSStatesPython

class HMMSLDS(_SLDSGibbsMixin, _SLDSVBEMMixin, _SLDSMeanFieldMixin, pyhsmm.models.HMM):
    _states_class = HMMSLDSStatesEigen

class HSMMSLDSPython(_SLDSGibbsMixin, _SLDSVBEMMixin, _SLDSMeanFieldMixin, pyhsmm.models.HSMMPython):
    _states_class = HSMMSLDSStatesPython

class HSMMSLDS(_SLDSGibbsMixin, _SLDSVBEMMixin, _SLDSMeanFieldMixin, pyhsmm.models.HSMM):
    _states_class = HSMMSLDSStatesEigen

class WeakLimitHDPHMMSLDS(_SLDSGibbsMixin, _SLDSVBEMMixin, _SLDSMeanFieldMixin,
                          pyhsmm.models.WeakLimitHDPHMM):
    _states_class = HMMSLDSStatesEigen

class WeakLimitStickyHDPHMMSLDS(
        _SLDSGibbsMixin, _SLDSVBEMMixin, _SLDSMeanFieldMixin,
        pyhsmm.models.WeakLimitStickyHDPHMM):
    _states_class = HMMSLDSStatesEigen

class WeakLimitHDPHSMMSLDS(_SLDSGibbsMixin, _SLDSVBEMMixin, _SLDSMeanFieldMixin,
                           pyhsmm.models.WeakLimitHDPHSMM):
    _states_class = HSMMSLDSStatesEigen

## Default constructors

def _default_model(model_class, K, D_obs, D_latent, D_input=0,
                   mu_inits=None, sigma_inits=None,
                   As=None, Bs=None, sigma_statess=None,
                   Cs=None, Ds=None, sigma_obss=None,
                   alpha=3.0, init_state_distn='uniform',
                   **kwargs):

    # Initialize init_dynamics_distns
    init_dynamics_distns = \
        [Gaussian(nu_0=D_latent+3,
                  sigma_0=3.*np.eye(D_latent),
                  mu_0=np.zeros(D_latent),
                  kappa_0=0.01)
         for _ in range(K)]

    if mu_inits is not None:
        assert isinstance(mu_inits, list) and len(mu_inits) == K
        for id, mu in zip(init_dynamics_distns, mu_inits):
            id.mu = mu

    if sigma_inits is not None:
        assert isinstance(sigma_inits, list) and len(sigma_inits) == K
        for id, sigma in zip(init_dynamics_distns, sigma_inits):
            id.sigma = sigma

    # Initialize dynamics distributions
    dynamics_distns = [Regression(
        nu_0=D_latent + 1,
        S_0=D_latent * np.eye(D_latent),
        M_0=np.hstack((.99 * np.eye(D_latent), np.zeros((D_latent, D_input)))),
        K_0=D_latent * np.eye(D_latent + D_input))
        for _ in range(K)]
    if As is not None:
        assert isinstance(As, list) and len(As) == K
        if D_input > 0:
            assert isinstance(Bs, list) and len(Bs) == K
            As = [np.hstack((A, B)) for A,B in zip(As, Bs)]
    else:
        # As = [random_rotation(D_latent) for _ in range(K)]
        As = [np.eye(D_latent) for _ in range(K)]
        if D_input > 0:
            As = [np.hstack((A, np.zeros((D_latent, D_input))))
                  for A in As]
    for dd, A in zip(dynamics_distns, As):
        dd.A = A

    if sigma_statess is not None:
        assert isinstance(sigma_statess, list) and len(sigma_statess) == K
    else:
        sigma_statess = [np.eye(D_latent) for _ in range(K)]

    for dd, sigma in zip(dynamics_distns, sigma_statess):
        dd.sigma = sigma

    # Initialize emission distributions
    _single_emission = (Cs is not None) and (not isinstance(Cs, list))

    if _single_emission:
        if D_input > 0:
            assert Ds is not None and not isinstance(Ds, list)
            Cs = np.hstack((Cs, Ds))

        if sigma_obss is None:
            sigma_obss = np.eye(D_obs)

        emission_distns = Regression(
            nu_0=D_obs + 3,
            S_0=D_obs * np.eye(D_obs),
            M_0=np.zeros((D_obs, D_latent + D_input)),
            K_0=D_obs * np.eye(D_latent + D_input),
            A=Cs, sigma=sigma_obss)

    else:
        emission_distns = [Regression(
            nu_0=D_obs + 1,
            S_0=D_obs * np.eye(D_obs),
            M_0=np.zeros((D_obs, D_latent + D_input)),
            K_0=D_obs * np.eye(D_latent + D_input))
            for _ in range(K)]

        if Cs is not None and sigma_obss is not None:
            assert isinstance(Cs, list) and len(Cs) == K
            assert isinstance(sigma_obss, list) and len(sigma_obss) == K
            if D_input > 0:
                assert isinstance(Ds, list) and len(Ds) == K
                Cs = [np.hstack((C, D)) for C,D in zip(Cs, Ds)]
        else:
            Cs = [np.zeros((D_obs, D_latent + D_input)) for _ in range(K)]
            sigma_obss = [0.05 * np.eye(D_obs) for _ in range(K)]

        for ed, C, sigma in zip(emission_distns, Cs, sigma_obss):
            ed.A = C
            ed.sigma = sigma

    model = model_class(
        init_dynamics_distns=init_dynamics_distns,
        dynamics_distns=dynamics_distns,
        emission_distns=emission_distns,
        init_state_distn=init_state_distn,
        alpha=alpha,
        **kwargs)

    return model

def DefaultSLDS(K, D_obs, D_latent, D_input=0,
                mu_inits=None, sigma_inits=None,
                As=None, Bs=None, sigma_statess=None,
                Cs=None, Ds=None, sigma_obss=None,
                alpha=3.,
                **kwargs):
    return _default_model(HMMSLDS, K, D_obs, D_latent, D_input=D_input,
                          mu_inits=mu_inits, sigma_inits=sigma_inits,
                          As=As, Bs=Bs, sigma_statess=sigma_statess,
                          Cs=Cs, Ds=Ds, sigma_obss=sigma_obss,
                          alpha=alpha,
                          **kwargs)

def DefaultWeakLimitHDPSLDS(K, D_obs, D_latent, D_input=0,
                mu_inits=None, sigma_inits=None,
                As=None, Bs=None, sigma_statess=None,
                Cs=None, Ds=None, sigma_obss=None,
                alpha=3., gamma=3.,
                **kwargs):
    return _default_model(WeakLimitHDPHMMSLDS, K, D_obs, D_latent, D_input=D_input,
                          mu_inits=mu_inits, sigma_inits=sigma_inits,
                          As=As, Bs=Bs, sigma_statess=sigma_statess,
                          Cs=Cs, Ds=Ds, sigma_obss=sigma_obss,
                          alpha=alpha, gamma=gamma,
                          **kwargs)

def DefaultWeakLimitStickyHDPSLDS(K, D_obs, D_latent, D_input=0,
                mu_inits=None, sigma_inits=None,
                As=None, Bs=None, sigma_statess=None,
                Cs=None, Ds=None, sigma_obss=None,
                alpha=3., gamma=3., kappa=10.,
                **kwargs):
    return _default_model(WeakLimitStickyHDPHMMSLDS, K, D_obs, D_latent, D_input=D_input,
                          mu_inits=mu_inits, sigma_inits=sigma_inits,
                          As=As, Bs=Bs, sigma_statess=sigma_statess,
                          Cs=Cs, Ds=Ds, sigma_obss=sigma_obss,
                          kappa=kappa, alpha=alpha, gamma=gamma,
                          **kwargs)

class _CountSLDSMixin(_SLDSGibbsMixin):

    def resample_emission_distns(self):
        if self._single_emission:
            data = [(np.hstack((s.gaussian_states, s.inputs)), s.data)
                    for s in self.states_list]
            mask = [s.mask for s in self.states_list] if self.has_missing_data else None
            omega = [s.omega for s in self.states_list]
            self._emission_distn.resample(data=data, mask=mask, omega=omega)

        else:
            for state, d in enumerate(self.emission_distns):
                data = [(np.hstack((s.gaussian_states[s.stateseq == state],
                                    s.inputs[s.stateseq == state])),
                         s.data[s.stateseq == state])
                        for s in self.states_list]
                mask = [s.mask[s.stateseq == state] for s in self.states_list] \
                    if self.has_missing_data else None
                omega = [s.omega[s.stateseq == state] for s in self.states_list]
                d.resample(data=data, mask=mask, omega=omega)

        self._clear_caches()

class HMMCountSLDSPython(_CountSLDSMixin, pyhsmm.models.HMMPython):
    _states_class = HMMCountSLDSStatesPython

class HMMCountSLDS(_CountSLDSMixin, pyhsmm.models.HMM):
    _states_class = HMMCountSLDSStatesEigen

class HSMMCountSLDSPython(_CountSLDSMixin, pyhsmm.models.HSMMPython):
    _states_class = HSMMCountSLDSStatesPython

class HSMMCountSLDS(_CountSLDSMixin, pyhsmm.models.HSMM):
    _states_class = HSMMCountSLDSStatesEigen

class WeakLimitHDPHMMCountSLDS(_CountSLDSMixin, pyhsmm.models.WeakLimitHDPHMM):
    _states_class = HMMCountSLDSStatesEigen

class WeakLimitStickyHDPHMMCountSLDS(
    _CountSLDSMixin, pyhsmm.models.WeakLimitStickyHDPHMM):
    _states_class = HMMCountSLDSStatesEigen

class WeakLimitHDPHSMMCountSLDS(
    _CountSLDSMixin, pyhsmm.models.WeakLimitHDPHSMM):
    _states_class = HSMMCountSLDSStatesEigen


from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass
from functools import cache
from itertools import permutations
from itertools import product
from typing import Literal

import numpy as np

from chemex.model import model

_BASES = {
    "ixy": ["ix", "iy"],
    "iz": ["iz"],
    "izsz": ["iz", "2izsz"],
    "iz_eq": ["ie", "iz"],
    "ixyz": ["ix", "iy", "iz"],
    "ixyz_eq": ["ie", "ix", "iy", "iz"],
    "ixysxy": ["2ixsx", "2ixsy", "2iysx", "2iysy"],
    "ixy_ixysxy": ["ix", "iy", "2ixsx", "2ixsy", "2iysx", "2iysy"],
    "ixyzsz": ["ix", "iy", "iz", "2ixsz", "2iysz", "2izsz"],
    "ixyzsz_diff": ["ix", "iy", "iz", "2ixsz", "2iysz", "2izsz"],
    "ixyzsz_eq": ["ie", "ix", "iy", "iz", "2ixsz", "2iysz", "2izsz"],
    "ixyzsxyz": [
        "ix",
        "iy",
        "iz",
        "sx",
        "sy",
        "sz",
        "2ixsz",
        "2iysz",
        "2izsx",
        "2izsy",
        "2ixsx",
        "2ixsy",
        "2iysx",
        "2iysy",
        "2izsz",
    ],
    "ixyzsxyz_eq": [
        "ie",
        "se",
        "ix",
        "iy",
        "iz",
        "sx",
        "sy",
        "sz",
        "2ixsz",
        "2iysz",
        "2izsx",
        "2izsy",
        "2ixsx",
        "2ixsy",
        "2iysx",
        "2iysy",
        "2izsz",
    ],
}
_TRANSITIONS: dict[str, tuple[tuple[str, str, float], ...]] = {
    "r2_i_{state}": (("ix", "ix", -1.0), ("iy", "iy", -1.0)),
    "r2_s_{state}": (("sx", "sx", -1.0), ("sy", "sy", -1.0)),
    "r1_i_{state}": (("iz", "iz", -1.0), ("iz", "ie", +1.0)),
    "r1_s_{state}": (("sz", "sz", -1.0), ("sz", "se", +1.0)),
    "r2a_i_{state}": (("2ixsz", "2ixsz", -1.0), ("2iysz", "2iysz", -1.0)),
    "r2a_s_{state}": (("2izsx", "2izsx", -1.0), ("2izsy", "2izsy", -1.0)),
    "r2mq_is_{state}": (
        ("2ixsx", "2ixsx", -1.0),
        ("2ixsy", "2ixsy", -1.0),
        ("2iysx", "2iysx", -1.0),
        ("2iysy", "2iysy", -1.0),
    ),
    "r1a_is_{state}": (("2izsz", "2izsz", -1.0),),
    "etaxy_i_{state}": (
        ("ix", "2ixsz", -1.0),
        ("iy", "2iysz", -1.0),
        ("2ixsz", "ix", -1.0),
        ("2iysz", "iy", -1.0),
    ),
    "etaxy_s_{state}": (
        ("sx", "2izsx", -1.0),
        ("sy", "2izsy", -1.0),
        ("2izsx", "sx", -1.0),
        ("2izsy", "sy", -1.0),
    ),
    "etaz_i_{state}": (
        ("iz", "2izsz", -1.0),
        ("2izsz", "iz", -1.0),
        ("2izsz", "ie", +1.0),
    ),
    "etaz_s_{state}": (
        ("sz", "2izsz", -1.0),
        ("2izsz", "sz", -1.0),
        ("2izsz", "se", +1.0),
    ),
    "sigma_is_{state}": (
        ("iz", "sz", -1.0),
        ("sz", "iz", -1.0),
        ("sz", "ie", +1.0),
        ("iz", "se", +1.0),
    ),
    "mu_is_{state}": (
        ("2ixsx", "2iysy", +1.0),
        ("2ixsy", "2iysx", -1.0),
        ("2iysx", "2ixsy", -1.0),
        ("2iysy", "2ixsx", +1.0),
    ),
    "rotz_i": (
        ("ix", "iy", -1.0),
        ("iy", "ix", +1.0),
        ("2ixsx", "2iysx", -1.0),
        ("2iysx", "2ixsx", +1.0),
        ("2ixsy", "2iysy", -1.0),
        ("2iysy", "2ixsy", +1.0),
        ("2ixsz", "2iysz", -1.0),
        ("2iysz", "2ixsz", +1.0),
    ),
    "rotz_s": (
        ("sx", "sy", -1.0),
        ("sy", "sx", +1.0),
        ("2ixsx", "2ixsy", -1.0),
        ("2ixsy", "2ixsx", +1.0),
        ("2iysx", "2iysy", -1.0),
        ("2iysy", "2iysx", +1.0),
        ("2izsx", "2izsy", -1.0),
        ("2izsy", "2izsx", +1.0),
    ),
    "cs_i_{state}": (
        ("ix", "iy", -1.0),
        ("iy", "ix", +1.0),
        ("2ixsx", "2iysx", -1.0),
        ("2iysx", "2ixsx", +1.0),
        ("2ixsy", "2iysy", -1.0),
        ("2iysy", "2ixsy", +1.0),
        ("2ixsz", "2iysz", -1.0),
        ("2iysz", "2ixsz", +1.0),
    ),
    "cs_s_{state}": (
        ("sx", "sy", -1.0),
        ("sy", "sx", +1.0),
        ("2ixsx", "2ixsy", -1.0),
        ("2ixsy", "2ixsx", +1.0),
        ("2iysx", "2iysy", -1.0),
        ("2iysy", "2iysx", +1.0),
        ("2izsx", "2izsy", -1.0),
        ("2izsy", "2izsx", +1.0),
    ),
    "carrier_i": (
        ("ix", "iy", +1.0),
        ("iy", "ix", -1.0),
        ("2ixsx", "2iysx", +1.0),
        ("2iysx", "2ixsx", -1.0),
        ("2ixsy", "2iysy", +1.0),
        ("2iysy", "2ixsy", -1.0),
        ("2ixsz", "2iysz", +1.0),
        ("2iysz", "2ixsz", -1.0),
    ),
    "carrier_s": (
        ("sx", "sy", +1.0),
        ("sy", "sx", -1.0),
        ("2ixsx", "2ixsy", +1.0),
        ("2ixsy", "2ixsx", -1.0),
        ("2iysx", "2iysy", +1.0),
        ("2iysy", "2iysx", -1.0),
        ("2izsx", "2izsy", +1.0),
        ("2izsy", "2izsx", -1.0),
    ),
    "offset_i": (
        ("ix", "iy", +2.0 * np.pi),
        ("iy", "ix", -2.0 * np.pi),
        ("2ixsx", "2iysx", +2.0 * np.pi),
        ("2iysx", "2ixsx", -2.0 * np.pi),
        ("2ixsy", "2iysy", +2.0 * np.pi),
        ("2iysy", "2ixsy", -2.0 * np.pi),
        ("2ixsz", "2iysz", +2.0 * np.pi),
        ("2iysz", "2ixsz", -2.0 * np.pi),
    ),
    "offset_s": (
        ("sx", "sy", +2.0 * np.pi),
        ("sy", "sx", -2.0 * np.pi),
        ("2ixsx", "2ixsy", +2.0 * np.pi),
        ("2ixsy", "2ixsx", -2.0 * np.pi),
        ("2iysx", "2iysy", +2.0 * np.pi),
        ("2iysy", "2iysx", -2.0 * np.pi),
        ("2izsx", "2izsy", +2.0 * np.pi),
        ("2izsy", "2izsx", -2.0 * np.pi),
    ),
    "jeff_i": (
        ("ix", "iy", -2.0 * np.pi),
        ("iy", "ix", +2.0 * np.pi),
        ("2ixsx", "2iysx", -2.0 * np.pi),
        ("2iysx", "2ixsx", +2.0 * np.pi),
        ("2ixsy", "2iysy", -2.0 * np.pi),
        ("2iysy", "2ixsy", +2.0 * np.pi),
        ("2ixsz", "2iysz", -2.0 * np.pi),
        ("2iysz", "2ixsz", +2.0 * np.pi),
    ),
    "j_is_{state}": (
        ("ix", "2iysz", -np.pi),
        ("2iysz", "ix", +np.pi),
        ("2ixsz", "iy", -np.pi),
        ("iy", "2ixsz", +np.pi),
        ("sx", "2izsy", -np.pi),
        ("2izsy", "sx", +np.pi),
        ("2izsx", "sy", -np.pi),
        ("sy", "2izsx", +np.pi),
    ),
    "d_{state}": (
        ("ix", "ix", -1.0),
        ("iy", "iy", -1.0),
        ("iz", "iz", -1.0),
        ("sx", "sx", -1.0),
        ("sy", "sy", -1.0),
        ("sz", "sz", -1.0),
        ("2ixsz", "2ixsz", -1.0),
        ("2iysz", "2iysz", -1.0),
        ("2izsx", "2izsx", -1.0),
        ("2izsy", "2izsy", -1.0),
        ("2ixsx", "2ixsx", -1.0),
        ("2ixsy", "2ixsy", -1.0),
        ("2iysx", "2iysx", -1.0),
        ("2iysy", "2iysy", -1.0),
        ("2izsz", "2izsz", -1.0),
    ),
    "b1x_i": (
        ("iy", "iz", -2.0 * np.pi),
        ("iz", "iy", +2.0 * np.pi),
        ("2iysx", "2izsx", -2.0 * np.pi),
        ("2izsx", "2iysx", +2.0 * np.pi),
        ("2iysy", "2izsy", -2.0 * np.pi),
        ("2izsy", "2iysy", +2.0 * np.pi),
        ("2iysz", "2izsz", -2.0 * np.pi),
        ("2izsz", "2iysz", +2.0 * np.pi),
    ),
    "b1y_i": (
        ("iz", "ix", -2.0 * np.pi),
        ("ix", "iz", +2.0 * np.pi),
        ("2izsx", "2ixsx", -2.0 * np.pi),
        ("2ixsx", "2izsx", +2.0 * np.pi),
        ("2izsy", "2ixsy", -2.0 * np.pi),
        ("2ixsy", "2izsy", +2.0 * np.pi),
        ("2izsz", "2ixsz", -2.0 * np.pi),
        ("2ixsz", "2izsz", +2.0 * np.pi),
    ),
    "b1x_s": (
        ("sy", "sz", -2.0 * np.pi),
        ("sz", "sy", +2.0 * np.pi),
        ("2ixsy", "2ixsz", -2.0 * np.pi),
        ("2ixsz", "2ixsy", +2.0 * np.pi),
        ("2iysy", "2iysz", -2.0 * np.pi),
        ("2iysz", "2iysy", +2.0 * np.pi),
        ("2izsy", "2izsz", -2.0 * np.pi),
        ("2izsz", "2izsy", +2.0 * np.pi),
    ),
    "b1y_s": (
        ("sz", "sx", -2.0 * np.pi),
        ("sx", "sz", +2.0 * np.pi),
        ("2ixsz", "2ixsx", -2.0 * np.pi),
        ("2ixsx", "2ixsz", +2.0 * np.pi),
        ("2iysz", "2iysx", -2.0 * np.pi),
        ("2iysx", "2iysz", +2.0 * np.pi),
        ("2izsz", "2izsx", -2.0 * np.pi),
        ("2izsx", "2izsz", +2.0 * np.pi),
    ),
}
_ATOMS = {
    atoms: {"i": atoms[0], "s": atoms[1]} for atoms in ("hn", "hc", "nh", "ch", "cn")
}

@cache
def _build_vectors(basis: Basis) -> dict[str, np.ndarray]:
    size = len(basis) * len(model.states)
    vectors: defaultdict[str, np.ndarray] = defaultdict(lambda: np.zeros((size, 1)))
    for index, (state, name) in enumerate(product(model.states, basis.components)):
        vectors[f"{name}_{state}"][index] = 1.0
        vectors[name][index] = 1.0
    return dict(vectors)

def _get_indices(
    basis: Basis, transition_name: str, state: str
) -> tuple[tuple[list[int], list[int]], list[float]]:
    rows: list[int] = []
    cols: list[int] = []
    vals: list[float] = []
    offset = model.states.index(state) * len(basis)
    for start, end, value in _TRANSITIONS[transition_name]:
        if {start, end}.issubset(basis.components):
            rows.append(basis.components.index(start) + offset)
            cols.append(basis.components.index(end) + offset)
            vals.append(value)
    return (rows, cols), vals

def _build_spin_matrices(basis: Basis) -> dict[str, np.ndarray]:
    size = len(basis) * len(model.states)
    matrices: dict[str, np.ndarray] = defaultdict(lambda: np.zeros((size, size)))
    for transition_name, state in product(_TRANSITIONS, model.states):
        if not basis.type.endswith("_diff") and transition_name.startswith("d_"):
            continue
        name = transition_name.format(state=state)
        indices, values = _get_indices(basis, transition_name, state)
        if values:
            matrices[name][indices] = values
    return matrices

def _build_exchange_matrices(basis: Basis) -> dict[str, np.ndarray]:
    matrices: dict[str, np.ndarray] = {}
    for (i1, s1), (i2, s2) in permutations(enumerate(model.states), r=2):
        name = f"k{s1}{s2}"
        matrix = np.zeros((len(model.states), len(model.states)))
        matrix[((i1, i2), i1)] = -1.0, 1.0
        matrices[name] = np.kron(matrix, np.eye(len(basis)))
    return matrices

@cache
def _build_matrices(basis: Basis) -> dict[str, np.ndarray]:
    return _build_exchange_matrices(basis) | _build_spin_matrices(basis)

@dataclass(frozen=True)
class Basis:
    type: str
    extension: Literal["", "dq", "tq"] = ""
    spin_system: str = ""

    @property
    def name(self):
        return ".".join([self.type, self.extension, self.spin_system])

    @property
    def components(self):
        return _BASES[self.type]

    @property
    def atoms(self):
        return {
            letter: atom
            for letter, atom in _ATOMS.get(self.spin_system, {}).items()
            if letter in self.type
        }

    @property
    def vectors(self) -> dict[str, np.ndarray]:
        return dict(_build_vectors(self))

    @property
    def matrices(self) -> dict[str, np.ndarray]:
        return _build_matrices(self)

    @property
    def required_names(self) -> set[str]:
        required_names = set(self.matrices)
        required_names |= {f"p{state}" for state in model.states}
        return required_names

    def __len__(self):
        return len(self.components)