import io
import functools
import logging

# import soundfile as sf
import numpy as np
import matplotlib
import matplotlib.pylab as plt

# from IPython.display import display, Audio

from nara_wpe.utils import stft, istft

from pb_bss.distribution import CACGMMTrainer
from pb_bss.evaluation import InputMetrics, OutputMetrics
from dataclasses import dataclass

# from beamforming_wrapper import beamform_mvdr_souden_from_masks
from pb_chime5.utils.numpy_utils import segment_axis_v2
from textgrid_processor import read_textgrid_from_file


def get_time_activity(dur_list, wavlen, sr):
    time_activity = [False] * wavlen
    for dur in dur_list:
        xmax = int(dur[1] * sr)
        xmin = int(dur[0] * sr)
        if xmax > wavlen:
            continue
        for i in range(xmin, xmax):
            time_activity[i] = True
    logging.info("Num of actived samples {}".format(time_activity.count(True)))
    return time_activity

def get_frequency_activity(
    time_activity,
    stft_window_length,
    stft_shift,
    stft_fading=True,
    stft_pad=True,
):
    time_activity = np.asarray(time_activity)

    if stft_fading:
        pad_width = np.array([(0, 0)] * time_activity.ndim)
        pad_width[-1, :] = stft_window_length - stft_shift  # Consider fading
        time_activity = np.pad(time_activity, pad_width, mode="constant")

    return segment_axis_v2(
        time_activity,
        length=stft_window_length,
        shift=stft_shift,
        end="pad" if stft_pad else "cut",
    ).any(axis=-1)


@dataclass
class Beamformer:
    type: str
    postfilter: str

    def __call__(self, Obs, target_mask, distortion_mask, debug=False):
        bf = self.type

        if bf == "mvdrSouden_ban":
            from pb_chime5.speech_enhancement.beamforming_wrapper import (
                beamform_mvdr_souden_from_masks,
            )

            X_hat = beamform_mvdr_souden_from_masks(
                Y=Obs,
                X_mask=target_mask,
                N_mask=distortion_mask,
                ban=True,
            )
        elif bf == "ch0":
            X_hat = Obs[0]
        elif bf == "sum":
            X_hat = np.sum(Obs, axis=0)
        else:
            raise NotImplementedError(bf)

        if self.postfilter is None:
            pass
        elif self.postfilter == "mask_mul":
            X_hat = X_hat * target_mask
        else:
            raise NotImplementedError(self.postfilter)

        return X_hat


@dataclass
class GSS:
    iterations: int = 20
    iterations_post: int = 0
    verbose: bool = True
    # use_pinv: bool = False
    # stable: bool = True

    def __call__(self, Obs, acitivity_freq=None, debug=False):
        initialization = np.asarray(acitivity_freq, dtype=np.float64)
        initialization = np.where(initialization == 0, 1e-10, initialization)
        initialization = initialization / np.sum(initialization, keepdims=True, axis=0)
        initialization = np.repeat(initialization[None, ...], 257, axis=0)

        source_active_mask = np.asarray(acitivity_freq, dtype=bool)
        source_active_mask = np.repeat(source_active_mask[None, ...], 257, axis=0)

        cacGMM = CACGMMTrainer()

        if debug:
            learned = []
        all_affiliations = []
        F = Obs.shape[-1]
        T = Obs.T.shape[-2]

        for f in range(F):
            if self.verbose:
                if f % 50 == 0:
                    logging.info(f"{f}/{F}")

            # T: Consider end of signal.
            # This should not be nessesary, but activity is for inear and not for
            # array.
            cur = cacGMM.fit(
                y=Obs.T[f, ...],
                initialization=initialization[f, ..., :T],
                iterations=self.iterations,
                source_activity_mask=source_active_mask[f, ..., :T],
            )
            affiliation = cur.predict(
                Obs.T[f, ...],
                source_activity_mask=source_active_mask[f, ..., :T],
            )

            all_affiliations.append(affiliation)

        posterior = np.array(all_affiliations).transpose(1, 2, 0)

        return posterior