import math
import numpy as np
from pojagi_dsp.channel.ecg import Segments
from pojagi_dsp.channel.ecg.generator.wavetable import ECGWaveTable

# NOTE: larger table required for avoiding aliasing at different srates than 125Hz
sinus_data = np.array(
    [
        # R-S: 0
        2000,
        1822,
        374,
        # S-Jp: 3
        -474,
        -271,
        -28,
        18,
        66,
        # Jp-T: 9
        63,
        73,
        91,
        101,
        101,
        101,
        116,
        124,
        124,
        # T: 17
        141,
        171,
        186,
        196,
        229,
        265,
        297,
        327,
        363,
        406,
        446,
        475,
        493,
        508,
        526,
        533,
        518,
        475,
        403,
        327,
        272,
        222,
        174,
        138,
        109,
        88,
        73,
        66,
        69,
        69,
        66,
        73,
        81,
        76,
        73,
        76,
        76,
        66,
        58,
        58,
        63,
        63,
        41,
        26,
        26,
        18,
        8,
        8,
        8,
        # U: 66 -- not found
        # T-P: 66
        2,
        3,
        2,
        2,
        2,
        -1,
        2,
        2,
        2,
        -1,
        0,
        -1,
        -1,
        3,
        2,
        1,
        3,
        2,
        1,
        0,
        1,
        # P: 87
        0,
        3,
        11,
        11,
        0,
        8,
        18,
        18,
        18,
        15,
        8,
        18,
        26,
        26,
        26,
        8,
        32,
        61,
        116,
        164,
        182,
        159,
        131,
        116,
        116,
        109,
        91,
        73,
        58,
        55,
        58,
        63,
        69,
        # P-R: 120
        48,
        -14,
        # Q-R: 122
        -40,
        131,
        931,
    ]
)  # len == 125


def parabolic_curve(t_idx, pr_idx):
    """Calculates a smooth, physiologically mimetic curve for the
    T-P-R segment of the ECG waveform.
    """
    # Compute the difference in indices to determine the segment length
    t_pr_idx_diff = pr_idx - t_idx

    # Generate a symmetric range of values centered around zero for the segment
    t_pr = np.arange(-math.floor(t_pr_idx_diff / 2), math.ceil(t_pr_idx_diff / 2))

    # Apply a parabolic transformation to create a smooth transition
    t_pr_curve: np.ndarray = t_pr**2 * -0.25

    # Normalize the curve so it starts at the T wave amplitude (141)
    return t_pr_curve - t_pr_curve[0]


tachycardia = np.array(
    [  # 58-107 flat
        # R-S: 0
        2000,
        1822,
        374,
        # S-Jp: 3
        -474,
        -271,
        -28,
        18,
        66,
        # Jp-T: 8
        63,
        73,
        91,
        101,
        101,
        101,
        116,
        124,
        124,
        # T: 17
        *parabolic_curve(17, 119) + 141,
        # P-R: 120
        124,
        48,
        -14,
        # Q-R: 122
        -40,
        731,
        1231,
    ]
)  # len == 125


def SinusWaveTable():
    segments = Segments(
        S=3,
        S_T=9,
        T=17,
        T_P=66,
        P=87,
        P_R=120,
        Q=122,
    )

    return ECGWaveTable(
        data=sinus_data,
        segments=segments,
    )


def TachycardiaWaveTable():
    segments = Segments(
        S=3,
        S_T=8,
        T=17,
        T_P=66,
        P=87,
        P_R=119,
        Q=122,
    )

    return ECGWaveTable(
        data=tachycardia,
        segments=segments,
        # Tachy is weaker than sinus, so we inflate the range here by 3/2,
        # which effectively attenuates the signal by 1/3 (i.e., it is 2/3 of
        # the amplitude of the data definition).
        top=2000 * (3 / 2),
        bottom=0,
    )


def FastTachycardiaWaveTable():
    segments = Segments(
        S=3,
        S_T=8,
        T=17,
        T_P=66,
        P=87,
        P_R=119,
        Q=122,
    )

    return ECGWaveTable(
        data=np.arange(-50, 51) ** 11 / 1e19,
        segments=segments,
        tachy_compress=("R", "R"),
        top=2,
        bottom=0,
    )


def ImpulseWaveTable():
    return ECGWaveTable(
        data=np.array([1, *([0] * 124)]),
        segments=Segments(
            S=3,
            S_T=9,
            T=17,
            T_P=66,
            P=87,
            P_R=119,
            Q=122,
        ),
        top=1,
        bottom=0,
    )


if __name__ == "__main__":
    from matplotlib import pyplot as plt

    samples = np.tile(FastTachycardiaWaveTable().data, 3)

    plt.plot(range(len(samples)), samples)
    plt.show()
