from typing import Dict, List, Tuple, Union

import numpy as np


def energy_spectrum(hs: float, tp: float, gamma: float, duration: int) -> List[float]:
    """
    Calculate JONSWAP energy spectrum for given wave parameters.

    This function computes the energy spectrum S(f) using the JONSWAP spectral model,
    which is commonly used for wind-generated waves in deep water.

    Parameters
    ----------
    hs : float
        Significant wave height in meters.
    tp : float
        Peak period in seconds.
    gamma : float
        JONSWAP peak enhancement factor (typically 1-7, default 3.3).
    duration : int
        Number of frequency points for the spectrum.

    Returns
    -------
    List[float]
        Energy spectrum values S(f) for each frequency.

    Notes
    -----
    The JONSWAP spectrum combines the Pierson-Moskowitz spectrum with a peak
    enhancement factor. The spectrum is defined as:

    S(f) = Beta * (Hs^2) * (Tp^-4) * (f^-5) * exp(-1.25 * (Tp*f)^-4) * gamma^exp(-((Tp*f-1)^2)/(2*sigma^2))

    where sigma varies based on frequency relative to peak frequency.
    """

    # Defining frequency series - tend length
    freqs = np.linspace(0.02, 1, duration)

    S = []
    fp = 1 / tp

    for f in freqs:
        if f <= fp:
            sigma = 0.07
        if f > fp:
            sigma = 0.09

        Beta = (0.06238 / (0.23 + 0.0336 * gamma - 0.185 * (1.9 + gamma) ** -1)) * (
            1.094 - 0.01915 * np.log(gamma)
        )
        Sf = (
            Beta
            * (hs**2)
            * (tp**-4)
            * (f**-5)
            * np.exp(-1.25 * (tp * f) ** -4)
            * gamma ** (np.exp((-((tp * f - 1) ** 2)) / (2 * sigma**2)))
        )
        S.append(Sf)

    return S


def series_Jonswap(waves: Dict[str, Union[float, int]]) -> np.ndarray:
    """
    Generate surface elevation time series from JONSWAP spectrum.

    This function generates a time series of surface elevation using the JONSWAP
    spectral model. The elevation is computed by summing harmonic components
    with random phases.

    Parameters
    ----------
    waves : Dict[str, Union[float, int]]
        Dictionary containing wave parameters:
        - 'H': Significant wave height (m)
        - 'T': Peak period (s)
        - 'gamma': JONSWAP peak enhancement factor
        - 'warmup': Spin-up time (s)
        - 'deltat': Time step (s)
        - 'tendc': Simulation period (s)

    Returns
    -------
    np.ndarray
        2D array with shape (n_time_steps, 2) containing:
        - Column 0: Time values (s)
        - Column 1: Surface elevation values (m)

    Notes
    -----
    The surface elevation is generated by:
    1. Computing the JONSWAP energy spectrum
    2. Converting spectrum to amplitude components
    3. Summing harmonic components with random phases
    """

    # waves properties
    hs = waves["H"]
    tp = waves["T"]
    gamma = waves["gamma"]
    warmup = waves["warmup"]
    deltat = waves["deltat"]
    tendc = waves["tendc"]

    # series duration
    duration = int(tendc + warmup)
    time = np.arange(0, duration, deltat)

    # series frequency
    # why? wee user manual
    freqs = np.linspace(0.02, 1, duration)
    delta_f = freqs[1] - freqs[0]

    # calculate energy spectrum
    S = energy_spectrum(hs, tp, gamma, duration)

    # series elevation
    teta = np.zeros((len(time)))

    # calculate aij
    for f in range(len(freqs)):
        ai = np.sqrt(S[f] * 2 * delta_f)
        eps = np.random.rand() * (2 * np.pi)

        # calculate elevation
        teta = teta + ai * np.cos(2 * np.pi * freqs[f] * time + eps)

    # generate series dataframe
    series = np.zeros((len(time), 2))
    series[:, 0] = time
    series[:, 1] = teta

    return series


def series_TMA(waves: Dict[str, Union[float, int]], depth: float) -> np.ndarray:
    """
    Generate surface elevation time series from TMA spectrum.

    This function generates a time series of surface elevation using the TMA
    spectral model, which extends the JONSWAP spectrum to include depth effects.

    Parameters
    ----------
    waves : Dict[str, Union[float, int]]
        Dictionary containing wave parameters:
        - 'H': Significant wave height (m)
        - 'T': Peak period (s)
        - 'gamma': JONSWAP peak enhancement factor
        - 'warmup': Spin-up time (s)
        - 'deltat': Time step (s)
        - 'comptime': Simulation period (s)
    depth : float
        Water depth in meters.

    Returns
    -------
    np.ndarray
        2D array with shape (n_time_steps, 2) containing:
        - Column 0: Time values (s)
        - Column 1: Surface elevation values (m)

    Notes
    -----
    The TMA spectrum modifies the JONSWAP spectrum to account for depth effects
    using a depth function fi that depends on wave number and depth.
    """

    # waves properties
    hs = waves["H"]
    tp = waves["T"]
    gamma = waves["gamma"]
    warmup = waves["warmup"]
    deltat = waves["deltat"]
    comptime = waves["comptime"]
    _g = 9.801

    # series duration
    duration = int(comptime + warmup)
    time = np.arange(0, duration, deltat)

    # series frequency
    # why? wee user manual
    freqs = np.linspace(0.02, 1, duration)
    delta_f = freqs[1] - freqs[0]

    # calculate JONSWAP energy spectrum
    S = energy_spectrum(hs, tp, gamma, duration)
    m0_Jonswap = np.trapz(S, x=freqs)

    # Transform JONSWAP into TMA by (omega, h)
    S_TMA = []
    for pf, f in enumerate(freqs):
        """
        omega = 2 * np.pi * f
        
        thv = omega * np.sqrt(depth / g)
        if thv < 1: 
            fi = 0.5 * (thv)**2
        elif (thv >= 1) and (thv < 2):
            fi = 1 - 0.5 * (2 - thv)**2
        else:
            fi = 1
        """
        # S_TMA.append(S[pf] * fi)

        # As in Holthuijsen
        L = waves_dispersion(1 / f, depth)[0]
        k = 2 * np.pi / L
        n = 0.5 * (1 + ((2 * k * depth) / (np.sinh(2 * k * depth))))
        fi = (1 / (2 * n)) * np.tanh(k * depth)

        S_TMA.append(S[pf] * fi)

    m0_TMA = np.trapz(S_TMA, x=freqs)

    S_TMA_mod = [i * (m0_Jonswap / m0_TMA) for i in S_TMA]

    # series elevation
    teta = np.zeros((len(time)))

    # calculate aij
    for f in range(len(freqs)):
        ai = np.sqrt(S_TMA_mod[f] * 2 * delta_f)
        eps = np.random.rand() * (2 * np.pi)

        # calculate elevation
        teta = teta + ai * np.cos(2 * np.pi * freqs[f] * time + eps)

    # generate series dataframe
    series = np.zeros((len(time), 2))
    series[:, 0] = time
    series[:, 1] = teta

    return series


def series_TMA_bimodal(waves: Dict[str, Union[float, int]], depth: float) -> np.ndarray:
    """
    Generate surface elevation time series from bimodal TMA spectrum.

    This function generates a time series of surface elevation using a bimodal
    TMA spectral model, combining two wave systems with different characteristics.

    Parameters
    ----------
    waves : Dict[str, Union[float, int]]
        Dictionary containing wave parameters for two wave systems:
        - 'Hs1': Significant wave height of system 1 (m)
        - 'Tp1': Peak period of system 1 (s)
        - 'gamma1': JONSWAP peak enhancement factor for system 1
        - 'Hs2': Significant wave height of system 2 (m)
        - 'Tp2': Peak period of system 2 (s)
        - 'gamma2': JONSWAP peak enhancement factor for system 2
        - 'warmup': Spin-up time (s)
        - 'deltat': Time step (s)
        - 'tendc': Simulation period (s)
    depth : float
        Water depth in meters.

    Returns
    -------
    np.ndarray
        2D array with shape (n_time_steps, 2) containing:
        - Column 0: Time values (s)
        - Column 1: Surface elevation values (m)

    Notes
    -----
    The bimodal spectrum is created by summing two TMA spectra, each representing
    a different wave system (e.g., wind sea and swell).
    """

    # waves properties
    hs1 = waves["Hs1"]
    tp1 = waves["Tp1"]
    gamma1 = waves["gamma1"]
    hs2 = waves["Hs2"]
    tp2 = waves["Tp2"]
    gamma2 = waves["gamma2"]
    warmup = waves["warmup"]
    deltat = waves["deltat"]
    tendc = waves["tendc"]

    # series duration
    # TODO: puede que haya algun problema en esta suma
    duration = int(tendc + warmup)
    time = np.arange(0, duration, deltat)

    # series frequency
    freqs = np.linspace(0.02, 1, duration)
    delta_f = freqs[1] - freqs[0]

    # calculate energy spectrum for each wave system
    S1 = energy_spectrum(hs1, tp1, gamma1, duration)
    S2 = energy_spectrum(hs2, tp2, gamma2, duration)

    m01_Jonswap = np.trapz(S1, x=freqs)
    m02_Jonswap = np.trapz(S2, x=freqs)

    # Transform JONSWAP into TMA by (omega, h)
    S_TMA_1, S_TMA_2 = [], []
    for pf, f in enumerate(freqs):
        # As in Holthuijsen
        L = waves_dispersion(1 / f, depth)[0]
        k = 2 * np.pi / L
        n = 0.5 * (1 + ((2 * k * depth) / (np.sinh(2 * k * depth))))
        fi = (1 / (2 * n)) * np.tanh(k * depth)

        S_TMA_1.append(S1[pf] * fi)
        S_TMA_2.append(S2[pf] * fi)

    m0_TMA_1 = np.trapz(S_TMA_1, x=freqs)
    m0_TMA_2 = np.trapz(S_TMA_2, x=freqs)

    S_TMA_mod_1 = [i * (m01_Jonswap / m0_TMA_1) for i in S_TMA_1]
    S_TMA_mod_2 = [i * (m02_Jonswap / m0_TMA_2) for i in S_TMA_2]

    # bimodal wave spectra
    S = np.sum([S_TMA_mod_1, S_TMA_mod_2], axis=0)

    # series elevation
    teta = np.zeros((len(time)))

    # calculate aij
    for f in range(len(freqs)):
        ai = np.sqrt(S[f] * 2 * delta_f)
        eps = np.random.rand() * (2 * np.pi)

        # calculate elevation
        teta = teta + ai * np.cos(2 * np.pi * freqs[f] * time + eps)

    # generate series dataframe
    series = np.zeros((len(time), 2))
    series[:, 0] = time
    series[:, 1] = teta

    return series


def series_Jonswap_bimodal(waves: Dict[str, Union[float, int]]) -> np.ndarray:
    """
    Generate surface elevation time series from bimodal JONSWAP spectrum.

    This function generates a time series of surface elevation using a bimodal
    JONSWAP spectral model, combining two wave systems with different characteristics.

    Parameters
    ----------
    waves : Dict[str, Union[float, int]]
        Dictionary containing wave parameters for two wave systems:
        - 'Hs1': Significant wave height of system 1 (m)
        - 'Tp1': Peak period of system 1 (s)
        - 'gamma1': JONSWAP peak enhancement factor for system 1
        - 'Hs2': Significant wave height of system 2 (m)
        - 'Tp2': Peak period of system 2 (s)
        - 'gamma2': JONSWAP peak enhancement factor for system 2
        - 'warmup': Spin-up time (s)
        - 'deltat': Time step (s)
        - 'tendc': Simulation period (s)

    Returns
    -------
    np.ndarray
        2D array with shape (n_time_steps, 2) containing:
        - Column 0: Time values (s)
        - Column 1: Surface elevation values (m)

    Notes
    -----
    The bimodal spectrum is created by summing two JONSWAP spectra, each representing
    a different wave system (e.g., wind sea and swell).
    """

    # waves properties
    hs1 = waves["Hs1"]
    tp1 = waves["Tp1"]
    gamma1 = waves["gamma1"]
    hs2 = waves["Hs2"]
    tp2 = waves["Tp2"]
    gamma2 = waves["gamma2"]
    warmup = waves["warmup"]
    deltat = waves["deltat"]
    tendc = waves["tendc"]

    # series duration
    # TODO: puede que haya algun problema en esta suma
    duration = int(tendc + warmup)
    time = np.arange(0, duration, deltat)

    # series frequency
    freqs = np.linspace(0.02, 1, duration)
    delta_f = freqs[1] - freqs[0]

    # calculate energy spectrum for each wave system
    S1 = energy_spectrum(hs1, tp1, gamma1, duration)
    S2 = energy_spectrum(hs2, tp2, gamma2, duration)
    # bimodal wave spectra
    S = np.sum([S1, S2], axis=0)

    # series elevation
    teta = np.zeros((len(time)))

    # calculate aij
    for f in range(len(freqs)):
        ai = np.sqrt(S[f] * 2 * delta_f)
        eps = np.random.rand() * (2 * np.pi)

        # calculate elevation
        teta = teta + ai * np.cos(2 * np.pi * freqs[f] * time + eps)

    # generate series dataframe
    series = np.zeros((len(time), 2))
    series[:, 0] = time
    series[:, 1] = teta

    return series


def series_regular_monochromatic(waves: Dict[str, Union[float, int]]) -> np.ndarray:
    """
    Generate monochromatic regular wave time series.

    This function generates a time series of surface elevation for monochromatic
    regular waves with constant amplitude and frequency.

    Parameters
    ----------
    waves : Dict[str, Union[float, int]]
        Dictionary containing wave parameters:
        - 'T': Wave period (s)
        - 'H': Wave height (m)
        - 'warmup': Spin-up time (s)
        - 'deltat': Time step (s)
        - 'comptime': Simulation period (s)

    Returns
    -------
    np.ndarray
        2D array with shape (n_time_steps, 2) containing:
        - Column 0: Time values (s)
        - Column 1: Surface elevation values (m)

    Notes
    -----
    The surface elevation follows a simple cosine function:
    η(t) = (H/2) * cos(2πt/T)
    """

    # waves properties
    T = waves["T"]
    H = waves["H"]
    # WL = waves["WL"]
    warmup = waves["warmup"]
    deltat = waves["deltat"]
    tendc = waves["comptime"]

    # series duration
    duration = tendc + int(warmup)
    time = np.arange(0, duration, deltat)

    # series elevation
    teta = (H / 2) * np.cos((2 * np.pi / T) * time)

    # generate series dataframe
    series = np.zeros((len(time), 2))
    series[:, 0] = time
    series[:, 1] = teta

    return series


def series_regular_bichromatic(waves: Dict[str, Union[float, int]]) -> np.ndarray:
    """
    Generate bichromatic regular wave time series.

    This function generates a time series of surface elevation for bichromatic
    regular waves, combining two monochromatic components.

    Parameters
    ----------
    waves : Dict[str, Union[float, int]]
        Dictionary containing wave parameters:
        - 'T1': Period of component 1 (s)
        - 'T2': Period of component 2 (s)
        - 'H': Wave height (m)
        - 'WL': Water level (m)
        - 'warmup': Spin-up time (s)
        - 'deltat': Time step (s)
        - 'tendc': Simulation period (s)

    Returns
    -------
    np.ndarray
        2D array with shape (n_time_steps, 2) containing:
        - Column 0: Time values (s)
        - Column 1: Surface elevation values (m)

    Notes
    -----
    The surface elevation is the sum of two cosine components:
    η(t) = (H/2) * cos(2πt/T1) + (H/2) * cos(2πt/T2)
    """

    # waves properties
    T1 = waves["T1"]
    T2 = waves["T2"]
    H = waves["H"]
    _WL = waves["WL"]
    warmup = waves["warmup"]
    deltat = waves["deltat"]
    tendc = waves["tendc"]

    # series duration
    duration = tendc + int(warmup)
    time = np.arange(0, duration, deltat)

    # series elevation
    teta = (H / 2) * np.cos((2 * np.pi / T1) * time) + (H / 2) * np.cos(
        (2 * np.pi / T2) * time
    )

    # generate series dataframe
    series = np.zeros((len(time), 2))
    series[:, 0] = time
    series[:, 1] = teta

    return series


def waves_dispersion(T: float, h: float) -> Tuple[float, float, float]:
    """
    Solve the wave dispersion relation to calculate wavelength, wave number, and phase speed.

    This function solves the linear wave dispersion relation to determine wave
    properties in water of finite depth.

    Parameters
    ----------
    T : float
        Wave period in seconds.
    h : float
        Water depth in meters.

    Returns
    -------
    Tuple[float, float, float]
        A tuple containing:
        - L: Wavelength in meters
        - k: Wave number in radians per meter
        - c: Phase speed in meters per second

    Notes
    -----
    The dispersion relation is solved iteratively using:
    L = (g*T²)/(2π) * tanh(2πh/L)

    where g is gravitational acceleration (9.81 m/s²).

    The wave number k = 2π/L and phase speed c = √(g*tanh(kh)/k).
    """

    L1 = 1
    L2 = ((9.81 * T**2) / 2 * np.pi) * np.tanh(h * 2 * np.pi / L1)
    umbral = 1

    while umbral > 0.1:
        L2 = ((9.81 * T**2) / (2 * np.pi)) * np.tanh((h * 2 * np.pi) / L1)
        umbral = np.abs(L2 - L1)
        L1 = L2

    L = L2
    k = (2 * np.pi) / L
    c = np.sqrt(9.8 * np.tanh(k * h) / k)

    return (L, k, c)
