import dataclasses
import logging
import math
from numbers import Number
from typing import Dict, List, Optional, Tuple

import numpy as np
from scipy.interpolate import CubicSpline

from pojagi_dsp.channel.ecg import Segments
from pojagi_dsp.channel.ecg.generator import AECGSynthesizer

logger = logging.getLogger(__name__)


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,
        bottom: Optional[Number] = None,
        top: Optional[Number] = None,
        table_length: int = 1 << 11,  # 2048
    ):
        """
        The table size is increased to `table_length` upon initialization
        using linear interpolation (via `numpy.interp`).
        """

        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)

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

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

    def linear_interpolation(
        self,
        index: float,
        floor: Optional[int] = None,
        ceiling: Optional[int] = None,
    ) -> float:
        """
        Handles the situation where the floor would produce duplicate values,
        which makes the waveform chunky with aliasing; instead, we obtain a
        value weighted between the floor/ceiling, trading time (CPU) for
        memory (table size).
        """
        dl = len(self.data)
        floor = floor if floor is not None else math.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[floor] * floor_weight + self[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),
            top=1,
            bottom=0,
        )


class ECGWaveTableSynthesizer(AECGSynthesizer):
    def __init__(
        self,
        /,
        tables: Dict[Tuple[float, float], ECGWaveTable],
        heart_rate: int,
        srate: Optional[float] = None,
    ):
        super().__init__(heart_rate, srate)
        self.inc: float = 0.0
        self.tables = tables
        self._segments: Segments = None
        self.phase = 0.0
        self.q_lock = False

    def samples(self):
        # phase: float = 0.0
        inc: float = None
        idx: int = 0
        heart_rate = self.heart_rate

        self._calibrate()

        while idx < self.wavelength:
            phase = self.phase
            floor = math.floor(phase)

            yield self.table.linear_interpolation(phase, floor=floor)

            if (
                heart_rate < 60
                and self.table.segments.T_P <= floor < self.table.segments.P
            ):
                inc = self.brady_inc
                logger.info(["brady", inc, self.inc])
                self.q_lock = True
            elif (
                heart_rate > 60
                and self.table.segments.S_T <= floor < self.table.segments.Q
            ):
                # FIXME: this is probably only good below a certain
                # `heart_rate` threshold, because at some high frequency, even
                # the QRS complex will not have enough room to complete.
                inc = self.tachy_inc
                self.q_lock = True
            else:
                inc = None
                if self.q_lock:
                    phase = self.table.segments.Q
                    self.q_lock = False
                logger.info(f"\n{[
                    self.table.linear_interpolation(phase, floor=floor),
                    inc,
                    self.inc,
                    self.table.segments.Q,
                    phase,
                    self.table.segments.S_T
                ]}")

            phase += inc if inc is not None else self.inc
            phase %= len(self.table)
            self.phase = phase
            idx += 1

    @AECGSynthesizer.heart_rate.setter
    def heart_rate(self, val):
        AECGSynthesizer.heart_rate.fset(self, val)

    def _calibrate(self):
        heart_rate = self.heart_rate

        table_matches = {
            k: v for k, v in self.tables.items() if k[0] <= heart_rate < k[1]
        }

        if not table_matches:
            raise ValueError(
                f"No table found corresponding to heart rate: {heart_rate}."
            )

        # Since we may have more than two tables that match, we loop
        # through all the matches, applying them in key order.
        keys = iter(sorted(table_matches))
        key = next(keys)
        table = table_matches[key]

        for next_key in keys:
            next_table = table_matches[next_key]

            if next_key[1] < key[1]:
                # `next_key` is fully contained within `key`
                floor, ceiling = next_key
                next_weight = (heart_rate - floor) / (ceiling - floor)
                weight = 1 - next_weight

                if (heart_rate - floor) > ((ceiling - floor) / 2):
                    # Weights form an "X" shape; i.e., crossfade to 50%
                    # and back.
                    weight, next_weight = next_weight, weight
            else:
                floor = next_key[0]  # i.e., the bottom of the top
                ceiling = key[1]  # i.e., the top of the bottom
                next_weight = (heart_rate - floor) / (ceiling - floor)

            table = table.merge(next_table, next_weight)
            key = next_key

        self.table = table

        # ECG Tables are designed for 1Hz, and as a default, we don't want to
        # stretch anything; hence, no reference to `self.heart_rate` here,
        # instead constant 60:
        self.inc = len(self.table) / (self.srate * (60 / 60))

        # Stretch only the T_P segment to compensate, rather than
        # stretching the whole wave.
        table_segment_length = self.table.segments.P - self.table.segments.T_P
        self.brady_inc = self.stretch_inc(table_segment_length)

        # Preserve QRS-J-point; compress Jp-Q to compensate.
        table_segment_length = self.table.segments.Q - self.table.segments.S_T
        self.tachy_inc = self.stretch_inc(table_segment_length)

    def stretch_inc(self, table_segment_length):
        # Get the missing samples by subtracting the number of samples
        # contributed by the 1Hz table default, minus the segment we
        # want to stretch.
        tmp_wavelength = (
            self.wavelength - (len(self.table) - table_segment_length) / self.inc
        )

        return table_segment_length / tmp_wavelength

    @property
    def segments(self):
        if self._segments:
            return self._segments

        table_length = len(self.table)
        table_segments = self.table.segments
        # FIXME: this is a lie since we stretch T_P, etc.
        self._segments = Segments(
            **{
                k: math.floor(v * self.srate / table_length) if v else v
                for k, v in [
                    (f.name, getattr(table_segments, f.name))
                    for f in dataclasses.fields(table_segments)
                ]
            }
        )
        return self._segments
