#!/usr/bin/env python
# encoding: utf-8

"""Abrahamson, Silva, and Kamai (2014) ground motion model."""

from __future__ import division

import numpy as np
from scipy.interpolate import interp1d

from . import model

__author__ = 'Albert Kottke'


class AbrahamsonSilvaKamai2014(model.Model):
    """The Abrahamson, Silva, and Kamai (2014) :cite:`abrahamson14` ground
    motion model for active tectonic regions.
    """

    NAME = 'Abrahamson, Silva, & Kamai (2014)'
    ABBREV = 'ASK14'

    # Reference velocity (m/sec)
    V_REF = 1180.

    # Load the coefficients for the model
    COEFF = model.load_data_file('abrahamson_silva_kamai_2014.csv', 2)

    PERIODS = COEFF['period']

    INDICES_PSA = np.arange(22)
    INDEX_PGA = -2
    INDEX_PGV = -1

    PARAMS = [
        model.NumericParameter('dist_rup', True, None, 300),
        model.NumericParameter('dist_jb', True),
        model.NumericParameter('mag', True, 3, 8.5),
        model.NumericParameter('v_s30', True, 180, 1000),

        model.NumericParameter('depth_1_0', False),
        model.NumericParameter('depth_tor', False),
        model.NumericParameter('dip', True),
        model.NumericParameter('dist_crjb', False, default=15),
        model.NumericParameter('dist_x', False),
        model.NumericParameter('dist_y0', False),
        model.NumericParameter('width', False),

        model.CategoricalParameter('mechanism', True, ['SS', 'NS', 'RS']),
        model.CategoricalParameter(
            'region', False,
            ['global', 'california', 'china', 'italy', 'japan', 'taiwan'],
            'global'
        ),
        model.CategoricalParameter(
            'vs_source', False, ['measured', 'inferred'], 'measured'),
        model.CategoricalParameter(
            'is_aftershock', False, [True, False], False),
        model.CategoricalParameter('on_hanging_wall', False,
                                   [True, False], False),
    ]

    def _check_inputs(self, **kwds):
        super(AbrahamsonSilvaKamai2014, self)._check_inputs(**kwds)

        p = self.params

        if p['width'] is None:
            p['width'] = self.calc_width(p['mag'], p['dip'])

        if p['depth_tor'] is None:
            p['depth_tor'] = self.calc_depth_tor(p['mag'])

    def __init__(self, **kwds):
        """Initialize the model.

        Keyword Args:
            depth_1_0 (Optional[float]): depth to the 1.0 km∕s shear-wave
                velocity horizon beneath the site, :math:`Z_{1.0}` in (km).
                Used to estimate `depth_2_5`.

            depth_2_5 (Optional[float]): depth to the 2.5 km∕s shear-wave
                velocity horizon beneath the site, :math:`Z_{2.5}` in (km).
                If *None*, then it is computed from `depth_1_0` or `v_s30`
                and the `region` parameter.

            depth_tor (Optional[float]): depth to the top of the rupture
                plane (:math:`Z_{tor}`, km). If *None*, then  the average
                model is used.

            depth_bor (Optional[float]): depth to the bottom of the rupture
                plane (:math:`Z_{bor}`, km). If *None*, then  the average
                model is used.

            dip (float): fault dip angle (:math:`\phi`, deg).

            dist_jb (float): Joyner-Boore distance to the rupture plane
                (:math:`R_\\text{JB}`, km)

            dist_rup (float): closest distance to the rupture plane
                (:math:`R_\\text{rup}`, km)

            dist_x (float): site coordinate measured perpendicular to the
                fault strike from the fault line with the down-dip direction
                being positive (:math:`R_x`, km).

            dist_y0 (Optional[float]): the horizontal distance off the end of
                the rupture measured parallel to strike (:math:`R_{y0}`, km).

            mag (float): moment magnitude of the event (:math:`M_w`)

            mechanism (str): fault mechanism. Valid options: "SS", "NS", "RS".

            on_hanging_wall (Optional[bool]): If the site is located on the
                hanging wall of the fault. If *None*, then *False* is assumed.

            region (Optional[str]): region. Valid options: "global",
                "california", "china", "italy", "japan", "taiwan". If *None*,
                then "global" is used as a default value.

            v_s30 (float): time-averaged shear-wave velocity over the top 30 m
                of the site (:math:`V_{s30}`, m/s).

            vs_source (Optional[str]): source of the `v_s30` value.  Valid
                options: "measured", "inferred"

            width (Optional[float]): Down-dip width of the fault. If *None*,
                then  the model average is used.
        """
        super(AbrahamsonSilvaKamai2014, self).__init__(**kwds)

        # Compute the response at the reference velocity
        resp_ref = np.exp(self._calc_ln_resp(self.V_REF, np.nan))

        self._ln_resp = self._calc_ln_resp(self.params['v_s30'], resp_ref)
        self._ln_std = self._calc_ln_std(resp_ref)

    def _calc_ln_resp(self, v_s30, resp_ref):
        """Calculate the natural logarithm of the response.

        Args:
            v_s30 (float): site condition. Set `v_s30` to the reference
                velocity (e.g., 1180 m/s) for the reference response.

            resp_ref (Optional[:class:`np.array`]): response at the reference
                condition. Required if `v_s30` is not equal to reference
                velocity.

        Returns:
            :class:`np.array`: Natural log of the response.
        """
        c = self.COEFF
        p = self.params

        # Magnitude scaling
        f1 = self._calc_f1()

        if p['on_hanging_wall']:
            # Hanging-wall term
            f4 = self._calc_f4()
        else:
            f4 = 0

        # Depth to top of rupture term
        f6 = c.a15 * np.clip(p['depth_tor'] / 20, 0, 1)
        # Style of faulting
        if p['mechanism'] == 'RS':
            f7 = c.a11 * np.clip(p['mag'] - 4, 0, 1)
            f8 = 0
        elif p['mechanism'] == 'NS':
            f7 = 0
            f8 = c.a12 * np.clip(p['mag'] - 4, 0, 1)
        else:
            f7, f8 = 0, 0

        # Site term
        ###########
        v_1 = np.exp(-0.35 * np.log(np.clip(c.period, 0.5, 3) / 0.5) +
                     np.log(1500))

        vs_ratio = np.minimum(v_s30, v_1) / c.v_lin
        # Linear site model
        f5 = (c.a10 + c.b * c.n) * np.log(vs_ratio)
        # Nonlinear model
        mask = vs_ratio < 1
        f5[mask] = (
            c.a10 * np.log(vs_ratio) -
            c.b * np.log(resp_ref + c.c) +
            c.b * np.log(resp_ref + c.c * vs_ratio ** c.n)
        )[mask]

        # Basin term
        if v_s30 == self.V_REF or p['depth_1_0'] is None:
            # No basin response
            f10 = 0
        else:
            # Ratio between site depth_1_0 and model center
            ln_depth_ratio = np.log(
                (p['depth_1_0'] + 0.01) /
                (self.calc_depth_1_0(v_s30, p['region']) + 0.01)
            )
            slope = interp1d(
                [150, 250, 400, 700],
                np.c_[c.a43, c.a44, c.a45, c.a46],
                copy=False,
                bounds_error=False,
                fill_value=(c.a43, c.a46),
            )(v_s30)
            f10 = slope * ln_depth_ratio

        # Aftershock term
        if p['is_aftershock']:
            f11 = c.a14 * np.clip(1 - (p['dist_crjb'] - 5) / 10, 0, 1)
        else:
            f11 = 0

        if p['region'] == 'taiwan':
            freg = c.a31 * np.log(vs_ratio) + c.a25 * p['dist_rup']
        elif p['region'] == 'china':
            freg = c.a28 * p['dist_rup']
        elif p['region'] == 'japan':
            f13 = interp1d(
                [150, 250, 350, 450, 600, 850, 1150],
                np.c_[c.a36, c.a37, c.a38, c.a39, c.a40, c.a41, c.a42],
                copy=False,
                bounds_error=False,
                fill_value=(c.a36, c.a42),
            )(v_s30)
            freg = f13 + c.a29 * p['dist_rup']
        else:
            freg = 0

        return f1 + f4 + f5 + f6 + f7 + f8 + f10 + f11 + freg

    def _calc_ln_std(self, psa_ref):
        """Calculate the logarithmic standard deviation.

        Returns:
            :class:`np.array`: Logarithmic standard deviation.
        """
        p = self.params
        c = self.COEFF

        if p['region'] == 'japan':
            phi_al = c.s5 + (c.s6 - c.s5) * np.clip((p['dist_rup'] - 30) / 50,
                                                    0, 1)
        else:
            transition = np.clip((p['mag'] - 4) / 2, 0, 1)
            if p['vs_source'] == 'measured':
                phi_al = c.s1m + (c.s2m - c.s1m) * transition
            else:
                phi_al = c.s1e + (c.s2e - c.s1e) * transition

        tau_al = c.s3 + (c.s4 - c.s3) * np.clip((p['mag'] - 5) / 2, 0, 1)
        tau_b = tau_al

        # Remove period independent site amplification uncertainty of 0.4
        phi_amp = 0.4
        phi_b = np.sqrt(phi_al ** 2 - phi_amp ** 2)

        # The partial derivative of the amplification with respect to
        # the reference intensity
        deriv = ((-c.b * psa_ref) / (psa_ref + c.c) +
                 (c.b * psa_ref) /
                 (psa_ref + c.c * (p['v_s30'] / c.v_lin) ** c.n))
        deriv[p['v_s30'] >= c.v_lin] = 0

        tau = tau_b * (1 + deriv)
        phi = np.sqrt(phi_b ** 2 * (1 + deriv) ** 2 + phi_amp ** 2)

        ln_std = np.sqrt(phi ** 2 + tau ** 2)

        return ln_std

    @staticmethod
    def calc_width(mag, dip):
        """Compute the fault width based on equation in NGW2 spreadsheet.

        This equation is not provided in the paper.

        Args:
            mag (float): moment magnitude of the event (:math:`M_w`)
            dip (float): Fault dip angle (:math:`\phi`, deg)

        Returns:
            float: estimated fault width (:math:`W`, km)
        """
        return min(
            18 / np.sin(np.radians(dip)),
            10 ** (-1.75 + 0.45 * mag)
        )

    @staticmethod
    def calc_depth_tor(mag):
        """Calculate the depth to top of rupture (km).

        Args:
            mag (float): moment magnitude of the event (:math:`M_w`)

        Returns:
            float: estimated depth (km)
        """
        return np.interp(mag, [5., 7.2], [7.8, 0])

    @staticmethod
    def calc_depth_1_0(v_s30, region='california'):
        """Estimate the depth to 1 km/sec horizon (:math:`Z_{1.0}`) based on
        :math:`V_{s30}` and region.

        This is based on equations 18 and 19 in the :cite:`abrahamson14`
        and differs from the equations in the :cite:`chiou14`.

        Args:
            v_s30 (float): time-averaged shear-wave velocity over the top 30 m
                of the site (:math:`V_{s30}`, m/s).

        Keyword Args:
            region (Optional[str]): region of basin model. Valid options:
                "california", "japan". If *None*, then "california" is used as
                the default value.

        Returns:
            float: depth to a shear-wave velocity of 1,000 m/sec
                (:math:`Z_{1.0}`, km).

        """
        if region in ['japan']:
            # Japan
            power = 2
            v_ref = 412
            slope = -5.23 / power
        else:
            # Global
            power = 4
            v_ref = 610
            slope = -7.67 / power

        return np.exp(slope * np.log((v_s30 ** power + v_ref ** power) /
                                     (1360. ** power + v_ref ** power))) / 1000

    def _calc_f1(self):
        """Calculate the magnitude scaling parameter f1.

        Returns:
            :class:`np.array`: Model parameter f1.
        """
        c = self.COEFF
        p = self.params

        # Magnitude dependent taper
        dist = np.sqrt(
            p['dist_rup'] ** 2 +
            (c.c4 - (c.c4 - 1) * np.clip(5 - p['mag'], 0, 1)) ** 2
        )

        # Magnitude scaling

        # Need to copy c.a1 to that it isn't modified during the following
        # operations.
        f1 = np.array(c.a1)
        ma1 = (p['mag'] <= c.m2)
        f1[ma1] += (
            c.a4 * (c.m2 - c.m1) + c.a8 * (8.5 - c.m2) ** 2 +
            c.a6 * (p['mag'] - c.m2) +
            c.a7 * (p['mag'] - c.m2) +
            (c.a2 + c.a3 * (c.m2 - c.m1)) * np.log(dist) +
            c.a17 * p['dist_rup']
        )[ma1]

        f1[~ma1] += (
            c.a8 * (8.5 - p['mag']) ** 2 +
            (c.a2 + c.a3 * (p['mag'] - c.m1)) * np.log(dist) +
            c.a17 * p['dist_rup']
        )[~ma1]

        ma2 = np.logical_and(~ma1, p['mag'] <= c.m1)
        f1[ma2] += (c.a4 * (p['mag'] - c.m1))[ma2]

        ma3 = np.logical_and(~ma1, p['mag'] > c.m1)
        f1[ma3] += (c.a5 * (p['mag'] - c.m1))[ma3]

        return f1

    def _calc_f4(self):
        """Calculate the hanging-wall parameter f4.

        Returns:
            :class:`np.array`: Model parameter f4.
        """
        # Move this into a decorator?
        c = self.COEFF
        p = self.params

        t1 = min(90 - p['dip'], 60) / 45
        # Constant from page 1041
        a2hw = 0.2
        if p['mag'] <= 5.5:
            t2 = 0
        elif p['mag'] < 6.5:
            t2 = (1 + a2hw * (p['mag'] - 6.5) -
                  (1 - a2hw) * (p['mag'] - 6.5) ** 2)
        else:
            t2 = 1 + a2hw * (p['mag'] - 6.5)

        # Constants defined on page 1040
        r1 = p['width'] * np.cos(np.radians(p['dip']))
        r2 = 3 * r1
        h1 = 0.25
        h2 = 1.5
        h3 = -0.75
        if p['dist_x'] < r1:
            t3 = h1 + h2 * (p['dist_x'] / r1) + h3 * (p['dist_x'] / r1) ** 2
        elif p['dist_x'] < r2:
            t3 = 1 - ((p['dist_x'] - r1) / (r2 - r1))
        else:
            t3 = 0

        t4 = np.clip(1 - p['depth_tor'] ** 2 / 100, 0, 1)

        if p['dist_y0'] is None:
            t5 = np.clip(1 - p['dist_jb'] / 30, 0, 1)
        else:
            dist_y1 = p['dist_x'] * np.tan(np.radians(20))
            t5 = np.clip(1 - (p['dist_y0'] - dist_y1) / 5, 0, 1)

        f4 = c.a13 * t1 * t2 * t3 * t4 * t5

        return f4
