from dataclasses import dataclass
from math import inf
from pathlib import Path
from typing import Final

import librosa
import numpy as np
import numpy.random as rand
import numpy.typing as npt

from .color_noise import pink as pk_noise
from .color_noise import white as wh_noise


@dataclass
class _TimeAlignment:
    def __init__(self, start, end, phoneme) -> None:
        self.start: Final[float] = start
        self.end: Final[float] = end
        self.phoneme: Final[str] = phoneme


class WavLabHandler:
    __x: Final[npt.NDArray[np.floating]]
    __sr: Final[float]
    __monophone_label: Final[list[_TimeAlignment]]

    def __init__(self, wav_path: Path, lab_path: Path) -> None:
        # load audio
        self.__x, self.__sr = librosa.load(wav_path, sr=None)
        self.__x.flags.writeable = False  # x is immutable

        # load label
        with lab_path.open() as f:
            self.__monophone_label = []
            for line in f:
                sp: list[str] = line.split()
                align = _TimeAlignment(
                    start=int(sp[0]) / 1e7,
                    end=int(sp[1]) / 1e7,
                    phoneme=sp[2],
                )
                self.__monophone_label.append(align)
        if len(self.__monophone_label) < 3:
            raise ValueError("Invalid label format")

    def get_answer(self) -> tuple[float, float]:
        return (
            self.__monophone_label[1].start,
            self.__monophone_label[-1].start,
        )

    def get_signal(self) -> tuple[npt.NDArray[np.floating], float]:
        return self.__x, self.__sr

    def get_noise_signal(
        self, snr: float, is_white: bool, with_pulse: bool, noise_seed: int
    ) -> tuple[npt.NDArray, int]:
        x: npt.NDArray[np.floating] = self.__x.copy()
        sr: int = int(self.__sr)
        ans_s_sec, ans_e_sec = self.get_answer()
        speech_start_idx: int = int(ans_s_sec * sr)
        speech_end_idx: int = int(ans_e_sec * sr)

        # generate noise (white or pink)
        noise: npt.NDArray[np.floating] = (
            (
                wh_noise(len(x), noise_seed)
                if is_white
                else pk_noise(len(x), sr, noise_seed)
            )
            .cpu()
            .numpy()
        )

        # mix stationary noise and signal (in specified snr)
        if snr == inf:
            noised_x = x
        elif snr == -inf:
            noised_x = noise
        else:
            noise_scale = _determine_noise_scale(
                x[speech_start_idx:speech_end_idx], noise, int(snr)
            )
            noised_x = x + noise * noise_scale

        # add pulse noise
        if with_pulse:
            rand.seed(noise_seed)
            pulse: npt.NDArray[np.floating] = rand.random(2) - 0.5 * 2
            # determine index adding pulse noise
            start_pulse_index: int = np.random.randint(0, speech_start_idx)
            end_pulse_index: int = np.random.randint(speech_end_idx, len(x) - 1)
            # add pulse noise
            noised_x[start_pulse_index] = pulse[0]
            noised_x[end_pulse_index] = pulse[1]
        return noised_x, sr

    def get_noise_signal2(
        self, snr: float, noise_type: str, noise_seed: int
    ) -> tuple[npt.NDArray, float]:
        x: npt.NDArray = self.__x
        ans_s_sec, ans_e_sec = self.get_answer()
        speech_start_idx: Final[int] = int(ans_s_sec * self.__sr)
        speech_end_idx: Final[int] = int(ans_e_sec * self.__sr)

        color_flag = False

        for nt in noise_type.split():
            if not color_flag and (nt == "white" or nt == "pink"):
                # generate color noise
                noise: npt.NDArray[np.floating] = (
                    (
                        wh_noise(len(x), noise_seed)
                        if nt == "white"
                        else pk_noise(len(x), self.__sr, noise_seed)
                    )
                    .cpu()
                    .numpy()
                )

                # mix stationary noise and signal (in specified snr)
                if snr == inf:
                    noised_x = x.copy()
                elif snr == -inf:
                    noised_x = noise
                else:
                    noise_scale = _determine_noise_scale(
                        x[speech_start_idx:speech_end_idx], noise, int(snr)
                    )
                    noised_x = x + noise * noise_scale

                color_flag = True

            # add pulse noise
            if nt == "pulse":
                rand.seed(noise_seed)
                pulse: npt.NDArray[np.floating] = rand.random(2) - 0.5 * 2
                # determine index adding pulse noise
                start_pulse_index: int = np.random.randint(0, speech_start_idx)
                end_pulse_index: int = np.random.randint(speech_end_idx, len(x) - 1)
                # add pulse noise
                noised_x[start_pulse_index] = pulse[0]
                noised_x[end_pulse_index] = pulse[1]

        return noised_x, self.__sr


def _determine_noise_scale(
    signal: npt.NDArray[np.floating],
    noise: npt.NDArray[np.floating],
    desired_snr_db: int,
) -> float:
    desired_snr_linear: float = 10 ** (desired_snr_db / 10)
    signal_power: float = np.mean(signal**2)
    noise_power: float = np.mean(noise**2)
    scaling_factor: float = np.sqrt(signal_power / (desired_snr_linear * noise_power))
    return scaling_factor
