import dataclasses
from numbers import Number

import numpy as np
from scipy.interpolate import CubicSpline

from pojagi_dsp.channel.ecg import Segments


class ECGWaveTable:
    """
    This type of wavetable is designed around the P and R. By
    convention, R will always be equal to 1, and the baseline (P) will always
    be 0. (That doesn't mean, however, that the other values can't cross these
    boundaries. E.g., Q and S are often negative.)
    """

    def __init__(
        self,
        /,
        data: list[Number],
        segments: Segments,
        brady_stretch: tuple[str, str] | None = None,
        tachy_compress: tuple[str, str] | None = None,
        bottom: Number | None = None,
        top: Number | None = None,
        table_length: int = (1 << 10) * 2,
    ):
        """
        Initialize an ECGWaveTable.

        Args:
            data (List[Number]): The raw waveform data points for one cardiac cycle.
            segments (Segments): Segment indices (e.g., P, Q, R, S, T) marking key features in the waveform.
            brady_stretch (tuple[str, str] | None): Segment interval to stretch for bradycardia (slow heart rate).
            tachy_compress (tuple[str, str] | None): Segment interval to compress for tachycardia (fast heart rate).
            bottom (Optional[Number]): Value to use as the baseline (P segment) for normalization. If None, uses data[segments.P].
            top (Optional[Number]): Value to use as the peak (R segment) for normalization. If None, uses data[segments.R].
            table_length (int): Number of samples in the expanded wavetable (default: 2048).
        """

        if len(data) == table_length:
            self.data = np.array(data)
        else:
            # We generate a larger table for use with linear interpolation,
            # trading time (CPU) for memory (table size).
            #
            # Here we use cubic spline interpolation instead of linear for
            # wavetable construction, since it usually only happens once at
            # startup, and should provide a much better quality table from
            # limited data, making it possible to work with small, manually
            # composed tables that we JIT convert to the larger table.
            ### FIXME: use the sinc function instead:
            ### f(x) = sin(x)/x where x =/= 0 and f(x) = 1 if x = 0
            ### you have to apply this scaled to each sample in the table and
            ### then add all of the resulting signals together.
            ### I think this is the same as summing the dft of each impulse
            ### as if the impulse is a member of a larger table.
            cs = CubicSpline(
                range(len(data)),
                data,
                bc_type="natural",
            )

            self.data = np.array(
                [
                    cs(x)
                    for x in np.linspace(
                        0, len(data), table_length
                    )
                ]
            )

        # Scale the declared segments to the table_length
        self.segments = Segments(
            **{
                k: (
                    int(v / (len(data) / table_length))
                    if v
                    else v
                )
                for k, v in (
                    (f.name, getattr(segments, f.name))
                    for f in dataclasses.fields(segments)
                )
            },
        )

        # NOTE: these are not the data min/max, but the normal min/max, either
        # provided as kwargs, or derived from P and R segment starts, by
        # convention.
        bottom = (
            bottom if bottom is not None else data[segments.P]
        )
        top = top if top is not None else data[segments.R]

        if not (0 == bottom and 1 == top):
            # Normalize between 0 and 1:
            self.data = (self.data - bottom) / (top - bottom)

        self.brady_stretch = (
            brady_stretch
            if brady_stretch is not None
            else ("S_T", "P")
        )

        self.tachy_compress = (
            tachy_compress
            if tachy_compress is not None
            else ("S_T", "Q")
        )

    def __getitem__(self, k):
        return self.data[k]

    def __len__(self):
        return len(self.data)  # O(1)

    def linear_interpolation(
        self,
        index: float,
        floor: Number | None = None,
        ceiling: Number | None = None,
    ) -> float:
        """
        Returns a smoothly interpolated value from the wavetable at a fractional index.

        Instead of returning discrete values (which can cause aliasing and a "chunky" waveform),
        this method linearly interpolates between the nearest lower (floor) and upper (ceiling)
        indices, weighted by their distance from the requested index. This improves waveform
        smoothness and reduces artifacts when sampling at arbitrary positions.

        Args:
            index (float): The fractional index to sample.
            floor (Optional[Number]): Override for the lower index (default: floor of index).
            ceiling (Optional[Number]): Override for the upper index (default: floor + 1).

        Returns:
            float: The interpolated value at the given index.
        """
        dl = len(self.data)
        floor = (
            floor if floor is not None else np.floor(index) % dl
        )
        ceiling = (
            ceiling if ceiling is not None else (floor + 1) % dl
        )

        # e.g., a. 124.75 - 124 == 0.75
        #       b. 123 - 123 == 0 (no weight goes to ceiling)
        ceiling_weight = index - floor
        # e.g., a. 1 - 0.75 == 0.25
        #       b. 1 - 0 == 1 (all weight goes to floor)
        floor_weight = 1 - ceiling_weight

        return (
            self[int(floor)] * floor_weight
            + self[int(ceiling)] * ceiling_weight
        )

    def merge(
        self,
        other: "ECGWaveTable",
        weight: float,
    ):
        self_weight = 1 - weight
        return ECGWaveTable(
            data=(self.data * self_weight + other.data * weight),
            segments=self.segments.merge(other.segments, weight),
            brady_stretch=(
                other.brady_stretch
                if self_weight < 0.5
                else self.brady_stretch
            ),
            tachy_compress=(
                other.tachy_compress
                if self_weight < 0.5
                else self.tachy_compress
            ),
            top=1,
            bottom=0,
        )
