import numpy as np
from pojagi_dsp.channel.ecg.generator import AECGSynthesizer
from pojagi_dsp.channel.ecg.generator.wavetable.wavetable import (
    ECGWaveTable,
)


class ECGWaveTableSynthesizer(AECGSynthesizer):
    def __init__(
        self,
        /,
        tables: dict[tuple[float, float], ECGWaveTable],
        heart_rate: int,
        srate: float | None = None,
    ):
        super().__init__(heart_rate, srate)
        self.inc: float = 0.0
        self.tables = tables
        self.table: ECGWaveTable | None = None
        self.phase: float = 0.0

        self.brady_start: int = 0
        self.brady_end: int = 0
        self.brady_inc: float = 0.0

        self.tachy_start: int = 0
        self.tachy_end: int = 0
        self.tachy_inc: float = 0.0

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

        self._calibrate()

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

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

            if heart_rate < 60 and (
                self.brady_start <= phase < self.brady_end
            ):
                inc = self.brady_inc
                if phase + inc > self.brady_end:
                    inc = self.brady_end - phase
            elif heart_rate > 60 and (
                self.tachy_start <= phase < self.tachy_end
            ):
                # FIXME: this might 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
                if phase + inc > self.tachy_end:
                    inc = self.tachy_end - phase
            else:
                inc = None

            phase += inc if inc is not None else self.inc
            if phase > len(self.table):
                phase = 0.0

            self.phase = phase
            idx += 1

    @AECGSynthesizer.heart_rate.setter
    def heart_rate(self, val):
        AECGSynthesizer.heart_rate.fset(self, val)
        print("setter", hex(id(self)))

    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))

        self.brady_start, self.brady_end = (
            getattr(self.table.segments, x)
            for x in self.table.brady_stretch
        )
        if self.table.brady_stretch == "R":
            self.brady_end == len(self.table) - 1

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

        self.tachy_start, self.tachy_end = (
            getattr(self.table.segments, x)
            for x in self.table.tachy_compress
        )
        if self.table.tachy_compress[1] == "R":
            self.tachy_end = len(self.table) - 1

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

    def stretch_inc(self, table_segment_length: int) -> float:
        # 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
