import math
from typing import Optional

import numpy as np

from .decoupled_analysis import Decoupled
from .ground_motion import GroundMotion
from .utilities import sample_ground_motions


class Coupled(Decoupled):
    """
    Coupled analysis for sliding block and ground motion interaction.

    Parameters
    ----------
    ky : float or tuple[list[float], list[float]] or tuple[np.ndarray, np.ndarray] or callable
        Yield acceleration or function defining yield acceleration.
    ground_motion : GroundMotion
        Ground motion object containing acceleration time history and time step.
    height : int or float
        Height of the sliding block.
    vs_slope : int or float
        Shear wave velocity slope.
    vs_base : int or float
        Base shear wave velocity.
    damp_ratio : float
        Damping ratio.
    ref_strain : float, optional
        Reference strain for equivalent linear analysis. Required when soil_model is 'equivalent_linear'.
    scale_factor : float, optional
        Scale factor for input acceleration, by default 1.
    target_pga : float, optional
        Target peak ground acceleration, by default None.
    soil_model : str, optional
        Soil model type, by default "linear_elastic".
    si_units : bool, optional
        Use SI units, by default True.
    lite : bool, optional
        Lite mode, by default False.

    Attributes
    ----------
    HEA : np.ndarray
        Horizontal equivalent acceleration time history from dynamic response.
    gamma : float
        Integration parameter gamma.
    block_acc : np.ndarray
        Block acceleration time history.
    ground_acc : np.ndarray
        Ground acceleration time history.
    max_sliding_disp : float
        Maximum sliding displacement.
    """

    def __init__(
        self,
        ky: float
        or tuple[list[float], list[float]]
        or tuple[np.ndarray, np.ndarray]
        or callable,
        ground_motion: GroundMotion,
        height: int or float,
        vs_slope: int or float,
        vs_base: int or float,
        damp_ratio: float,
        ref_strain: Optional[float] = None,
        scale_factor: float = 1,
        target_pga: float = None,
        soil_model: str = "linear_elastic",
        si_units: bool = True,
        lite: bool = False,
        inverse: bool = False,
    ):
        super().__init__(
            ky,
            ground_motion,
            height,
            vs_slope,
            vs_base,
            damp_ratio,
            ref_strain,
            scale_factor,
            target_pga,
            soil_model,
            si_units,
            lite,
            inverse,
        )

        self.s1 = self.sdot1 = self.sdotdot1 = 0.0
        self.s2 = self.sdot2 = self.sdotdot2 = 0.0
        self.u1 = self.udot1 = self.udotdot1 = 0.0
        self.u2 = self.udot2 = self.udotdot2 = 0.0
        self.baseacc = self.basef = self.acc11 = self.acc22 = 0.0
        self.normalf1 = self.normalf2 = self.gameff1 = 0.0
        self.mx = self.mx1 = self.mmax = 0.0
        self.s = np.zeros(self.npts)
        self.u = np.zeros(self.npts)
        self.udotdot = np.zeros(self.npts)
        self.HEA = np.zeros(self.npts)
        self.sliding_vel = np.zeros(self.npts)
        self.udot = np.zeros(self.npts)
        self.angle = 0
        self.COS = math.cos(self.angle * math.pi / 180.0)
        self.SIN = math.sin(self.angle * math.pi / 180.0)

        self.gCOS = self.g * self.COS
        self.gSIN = self.g * self.SIN

        self.beta = 0.25  # TODO: move to global constants
        self.gamma = 0.5  # TODO: move to global constants

        self._block_acc_ = np.zeros(self.npts)
        # Sign reversal below to match decoupled sliding direction
        self.a_in *= -1
        if type(self) is Coupled:
            self.run_sliding_analysis()

        # Sign reversal corrected for plotting
        self._ground_acc_ = -self.a_in * self.g

    def __str__(self):
        return (
            f"Coupled:\n"
            f"  ky: {self.ky} g,\n"
            f"  {self.ground_motion},\n"
            f"  Scale factor: {self.scale_factor},\n"
            f"  Displacement: {100 * getattr(self, 'max_sliding_disp', 0):.1f} cm"
        )

    def __eq__(self, other):
        if not isinstance(other, Coupled):
            return NotImplemented
        return super().__eq__(other) and (
            self.height == other.height
            and self.vs_slope == other.vs_slope
            and self.vs_base == other.vs_base
            and self.damp_ratio == other.damp_ratio
            and self.ref_strain == other.ref_strain
            and self.soil_model == other.soil_model
            and self.SI_units == other.SI_units
            and self.lite == other.lite
        )

    def run_sliding_analysis(self):  # TODO: add ca to inputs
        if self.soil_model == "equivalent_linear":
            self.equivalent_linear()

        for i in range(1, self.npts + 1):
            self.dynamic_response(i)

        # calculate coupled displacements
        for i in range(1, self.npts + 1):
            self.coupled_sliding(i)

        # return self.max_sliding_disp
        self.block_disp = self.s
        self.max_sliding_disp = self.block_disp[-1]

    def coupled_sliding(self, i):
        self.coupled_setupstate(i)
        # solve for x_resp, v_resp, a_resp at next time step
        self.solvu(i)
        self.udotdot[i - 1] = self.udotdot2

        # update sliding acceleration based on calc'd response
        self.c_slideacc(i)

        # check if sliding has started
        self.c_slidingcheck(i)
        self.HEA[i - 1] = self.basef / self.mass  # Horizontal equivalent acceleration
        self._block_acc_[i - 1] = self.HEA[i - 1] - self.sdotdot1

        self.s[i - 1] = self.s2

    def coupled_setupstate(self, i):
        # set up state from previous time step
        if i == 1:
            self.u1 = self.udot1 = self.udotdot1 = self.s1 = self.sdot1 = (
                self.sdotdot1
            ) = self.normalf1 = 0.0
        else:
            self.u1 = self.u2
            self.udot1 = self.udot2
            self.udotdot1 = self.udotdot2
            self.s1 = self.s2
            self.sdot1 = self.sdot2
            self.sdotdot1 = self.sdotdot2
            self.normalf1 = self.normalf2

        # Set up acceleration loading. Normal force corrected for vertical component of a_in.
        self.normalf2 = (
            self.mass * self.gCOS
            + self.mass * self.a_in[i - 1] * self.gSIN  # * self.scale_factor
        )

        if i == 1:
            self.acc11 = 0.0
            self.acc22 = self.a_in[i - 1] * self.gCOS  # * self.scale_factor
        elif not self._slide:
            self.acc11 = self.a_in[i - 2] * self.gCOS  # * self.scale_factor
            self.acc22 = self.a_in[i - 1] * self.gCOS  # * self.scale_factor
        else:
            self.acc11 = self.gSIN - self.k_y(self.s[i - 2]) * self.normalf1 / self.mass
            self.acc22 = self.gSIN - self.k_y(self.s[i - 2]) * self.normalf2 / self.mass

    def solvu(self, i):
        khat = a = b = deltp = deltu = deltudot = d1 = 0.0

        delt = self.dt

        if self._slide:
            d1 = 1.0 - (self.L1**2) / (self.M1 * self.mass)
        else:
            d1 = 1.0

        khat = (
            (self._omega**2)
            + 2.0 * self._damp_tot * self._omega * self.gamma / (self.beta * delt)
            + d1 / (self.beta * (delt**2))
        )
        a = (
            d1 / (self.beta * delt)
            + 2.0 * self._damp_tot * self._omega * self.gamma / self.beta
        )
        b = d1 / (2.0 * self.beta) + delt * 2.0 * self._damp_tot * self._omega * (
            self.gamma / (2.0 * self.beta) - 1.0
        )

        if i == 1:
            deltp = -self.L1 / self.M1 * (self.acc22 - self.acc11)
            deltu = deltp / khat
            deltudot = self.gamma / (self.beta * delt) * deltu
            self.u2 = deltu
            self.udot2 = deltudot
            self.udotdot2 = (
                -(self.L1 / self.M1) * self.acc22
                - 2.0 * self._damp_tot * self._omega * self.udot2
                - (self._omega**2) * self.u2
            ) / d1
        else:
            deltp = (
                -self.L1 / self.M1 * (self.acc22 - self.acc11)
                + a * self.udot1
                + b * self.udotdot1
            )
            deltu = deltp / khat
            deltudot = (
                self.gamma / (self.beta * delt) * deltu
                - self.gamma / self.beta * self.udot1
                + delt * (1.0 - self.gamma / (2.0 * self.beta)) * self.udotdot1
            )
            self.u2 = self.u1 + deltu
            self.udot2 = self.udot1 + deltudot
            self.udotdot2 = (
                -(self.L1 / self.M1) * self.acc22
                - 2.0 * self._damp_tot * self._omega * self.udot2
                - (self._omega**2) * self.u2
            ) / d1

        self.u[i - 1] = self.u2

    def c_slideacc(self, i):
        # update sliding acceleration based on calc'd response
        if self._slide:
            self.sdotdot2 = (
                -self.a_in[i - 1] * self.gCOS  # * self.scale_factor
                - self.k_y(self.s[i - 2]) * self.normalf2 / self.mass
                - self.L1 * self.udotdot2 / self.mass
                + self.gSIN
            )

        # calc. base force based on a_resp calc
        self.basef = (
            -self.mass * self.a_in[i - 1] * self.gCOS  # * self.scale_factor
            - self.L1 * self.udotdot2
            + self.mass * self.gSIN
        )

        # If sliding is occurring, integrate sdotdot, using trapezoid rule, to get block_vel and block_disp.
        if self._slide:
            self.sdot2 = self.sdot1 + 0.5 * self.dt * (self.sdotdot2 + self.sdotdot1)
            self.s2 = self.s1 + 0.5 * self.dt * (self.sdot2 + self.sdot1)

    def c_slidingcheck(self, i):
        # check if sliding has started
        if not self._slide:
            if self.basef > self.k_y(self.s[i - 2]) * self.normalf2:
                self._slide = True
        elif self._slide:
            if self.sdot2 <= 0.0:
                self.slidestop(i)
                self._slide = False
                self.sdot2 = 0.0
                self.sdotdot2 = 0.0
        self.sliding_vel[i - 1] = self.sdot2

    def slidestop(self, i):
        ddt = acc11 = acc22 = acc1b = delt = dd = 0.0
        khat = deltp = a = b = 0.0

        delt = self.dt

        # Time of end of sliding is taken as where block_vel=0 from previous analysis
        dd = -self.sdot1 / (self.sdot2 - self.sdot1)
        ddt = dd * delt
        acc11 = self.gSIN - self.k_y(self.s[i - 2]) * (
            self.gCOS + self.a_in[i - 1] * self.gSIN  # * self.scale_factor
        )
        acc1b = (
            self.a_in[i - 2] * self.g  # * self.scale_factor
            + dd * (self.a_in[i - 1] - self.a_in[i - 2]) * self.g  # * self.scale_factor
        )
        acc22 = self.gSIN - self.k_y(self.s[i - 2]) * (self.gCOS + acc1b * self.SIN)

        # if dd=0, sliding has already stopped and skip this solution
        if dd == 0:
            return

        self.solvu(i)
        self.u1 = self.u2
        self.udot1 = self.udot2
        self.udotdot1 = self.udotdot2
        self.normalf2 = self.mass * self.gCOS + self.mass * acc1b * self.SIN
        self.sdotdot2 = (
            -acc1b * self.COS
            - self.k_y(self.s[i - 2]) * self.normalf2 / self.mass
            - self.L1 * self.udotdot2 / self.mass
            + self.gSIN
        )
        self.sdot2 = self.sdot1 + 0.5 * ddt * (self.sdotdot2 + self.sdotdot1)
        self.s2 = self.s1 + 0.5 * ddt * (self.sdot1 + self.sdot2)

        # Solve for non sliding response during remaining part of dt
        ddt = (1.0 - dd) * delt
        self._slide = False
        acc11 = acc22
        acc22 = self.a_in[i - 1] * self.gCOS  # * self.scale_factor

        khat = (
            1.0
            + 2.0 * self._damp_tot * self._omega * self.gamma * ddt
            + (self._omega**2) * self.beta * (ddt**2)
        )
        a = (
            (1.0 - (self.L1**2) / (self.mass * self.M1))
            + 2.0 * self._damp_tot * self._omega * ddt * (self.gamma - 1.0)
            + (self._omega**2) * (ddt**2) * (self.beta - 0.5)
        )
        b = (self._omega**2) * ddt
        deltp = (
            -self.L1 / self.M1 * (acc22 - acc11)
            + a * (self.udotdot1)
            - b * (self.udot1)
        )
        self.udotdot2 = deltp / khat

        self.udot2 = (
            self.udot1
            + (1.0 - self.gamma) * ddt * (self.udotdot1)
            + self.gamma * ddt * (self.udotdot2)
        )
        self.u2 = (
            self.u1
            + self.udot1 * ddt
            + (0.5 - self.beta) * (ddt**2) * (self.udotdot1)
            + self.beta * (ddt**2) * (self.udotdot2)
        )


equivalent_linear_testing = False
k_y_testing = False


def some_ky_func(disp):
    initial = 0.15
    minimum = 0.05
    factor = 0.005
    exponent = -1.5
    value = max(factor * (disp + minimum) ** exponent + 0.4 * disp, minimum)
    return min(initial, value)


if __name__ == "__main__":
    # import pyslammer as slam
    histories = sample_ground_motions()
    ky_const = 0.15
    ky_interp = ([0.2, 0.3, 0.4, 0.5], [0.15, 0.14, 0.13, 0.12])
    ky_func = some_ky_func
    motion = histories["Chi-Chi_1999_TCU068-090"]
    # t_step = motion[0][1] - motion[0][0]
    # input_acc = motion[1] / 9.80665

    ca = Coupled(
        ky=ky_const,
        ground_motion=motion,
        height=50.0,
        vs_slope=600.0,
        vs_base=600.0,
        damp_ratio=0.05,
        ref_strain=0.0005,
        scale_factor=1.0,
        soil_model="equivalent_linear",
        si_units=True,
        lite=False,
    )

    ca.run_sliding_analysis()
    print(ca.s[-1] * 100)
