import torch
from torch.utils import data
import json
import os
import numpy as np
import soundfile as sf
from .wsj0_mix import wsj0_license
from .wham_dataset import wham_noise_license

DATASET = "WHAMR"

# WHAMR tasks
# Many tasks can be considered with this dataset, we only consider the 4 core
# separation tasks presented in the paper for now.
sep_clean = {
    "mixture": "mix_clean_anechoic",
    "sources": ["s1_anechoic", "s2_anechoic"],
    "infos": [],
    "default_nsrc": 2,
}
sep_noisy = {
    "mixture": "mix_both_anechoic",
    "sources": ["s1_anechoic", "s2_anechoic"],
    "infos": ["noise"],
    "default_nsrc": 2,
}
sep_reverb = {
    "mixture": "mix_clean_reverb",
    "sources": ["s1_anechoic", "s2_anechoic"],
    "infos": [],
    "default_nsrc": 2,
}
sep_reverb_noisy = {
    "mixture": "mix_both_reverb",
    "sources": ["s1_anechoic", "s2_anechoic"],
    "infos": ["noise"],
    "default_nsrc": 2,
}

WHAMR_TASKS = {
    "sep_clean": sep_clean,
    "sep_noisy": sep_noisy,
    "sep_reverb": sep_reverb,
    "sep_reverb_noisy": sep_reverb_noisy,
}
# Support both order, confusion is easy
WHAMR_TASKS["sep_noisy_reverb"] = WHAMR_TASKS["sep_reverb_noisy"]


class WhamRDataset(data.Dataset):
    """ Dataset class for WHAMR source separation and speech enhancement tasks.

    Args:
        json_dir (str): The path to the directory containing the json files.
        task (str): One of ``'sep_clean'``, ``'sep_noisy'``, ``'sep_reverb'``
            or ``'sep_reverb_noisy'``.

            * ``'sep_clean'`` for two-speaker clean (anechoic) source
                separation.
            * ``'sep_noisy'`` for two-speaker noisy (anechoic) source
                separation.
            * ``'sep_reverb'`` for two-speaker clean reverberant
                source separation.
            * ``'sep_reverb_noisy'`` for two-speaker noisy reverberant source
                separation.

        sample_rate (int, optional): The sampling rate of the wav files.
        segment (float, optional): Length of the segments used for training,
            in seconds. If None, use full utterances (e.g. for test).
        nondefault_nsrc (int, optional): Number of sources in the training
            targets.
            If None, defaults to one for enhancement tasks and two for
            separation tasks.

    References:
        "WHAMR!: Noisy and Reverberant Single-Channel Speech Separation",
        Maciejewski et al. 2020
    """

    dataset_name = "WHAMR"

    def __init__(self, json_dir, task, sample_rate=8000, segment=4.0, nondefault_nsrc=None):
        super(WhamRDataset, self).__init__()
        if task not in WHAMR_TASKS.keys():
            raise ValueError(
                "Unexpected task {}, expected one of " "{}".format(task, WHAMR_TASKS.keys())
            )
        # Task setting
        self.json_dir = json_dir
        self.task = task
        self.task_dict = WHAMR_TASKS[task]
        self.sample_rate = sample_rate
        self.seg_len = None if segment is None else int(segment * sample_rate)
        if not nondefault_nsrc:
            self.n_src = self.task_dict["default_nsrc"]
        else:
            assert nondefault_nsrc >= self.task_dict["default_nsrc"]
            self.n_src = nondefault_nsrc
        self.like_test = self.seg_len is None
        # Load json files
        mix_json = os.path.join(json_dir, self.task_dict["mixture"] + ".json")
        sources_json = [
            os.path.join(json_dir, source + ".json") for source in self.task_dict["sources"]
        ]
        with open(mix_json, "r") as f:
            mix_infos = json.load(f)
        sources_infos = []
        for src_json in sources_json:
            with open(src_json, "r") as f:
                sources_infos.append(json.load(f))
        # Filter out short utterances only when segment is specified
        orig_len = len(mix_infos)
        drop_utt, drop_len = 0, 0
        if not self.like_test:
            for i in range(len(mix_infos) - 1, -1, -1):  # Go backward
                if mix_infos[i][1] < self.seg_len:
                    drop_utt += 1
                    drop_len += mix_infos[i][1]
                    del mix_infos[i]
                    for src_inf in sources_infos:
                        del src_inf[i]

        print(
            "Drop {} utts({:.2f} h) from {} (shorter than {} samples)".format(
                drop_utt, drop_len / sample_rate / 36000, orig_len, self.seg_len
            )
        )
        self.mix = mix_infos
        # Handle the case n_src > default_nsrc
        while len(sources_infos) < self.n_src:
            sources_infos.append([None for _ in range(len(self.mix))])
        self.sources = sources_infos

    def __add__(self, wham):
        if self.n_src != wham.n_src:
            raise ValueError(
                "Only datasets having the same number of sources"
                "can be added together. Received "
                "{} and {}".format(self.n_src, wham.n_src)
            )
        if self.seg_len != wham.seg_len:
            self.seg_len = min(self.seg_len, wham.seg_len)
            print(
                "Segment length mismatched between the two Dataset"
                "passed one the smallest to the sum."
            )
        self.mix = self.mix + wham.mix
        self.sources = [a + b for a, b in zip(self.sources, wham.sources)]

    def __len__(self):
        return len(self.mix)

    def __getitem__(self, idx):
        """ Gets a mixture/sources pair.
        Returns:
            mixture, vstack([source_arrays])
        """
        # Random start
        if self.mix[idx][1] == self.seg_len or self.like_test:
            rand_start = 0
        else:
            rand_start = np.random.randint(0, self.mix[idx][1] - self.seg_len)
        if self.like_test:
            stop = None
        else:
            stop = rand_start + self.seg_len
        # Load mixture
        x, _ = sf.read(self.mix[idx][0], start=rand_start, stop=stop, dtype="float32")
        seg_len = torch.as_tensor([len(x)])
        # Load sources
        source_arrays = []
        for src in self.sources:
            if src[idx] is None:
                # Target is filled with zeros if n_src > default_nsrc
                s = np.zeros((seg_len,))
            else:
                s, _ = sf.read(src[idx][0], start=rand_start, stop=stop, dtype="float32")
            source_arrays.append(s)
        sources = torch.from_numpy(np.vstack(source_arrays))
        return torch.from_numpy(x), sources

    def get_infos(self):
        """ Get dataset infos (for publishing models).

        Returns:
            dict, dataset infos with keys `dataset`, `task` and `licences`.
        """
        infos = dict()
        infos["dataset"] = self.dataset_name
        infos["task"] = self.task
        if self.task == "sep_clean":
            data_license = [wsj0_license]
        else:
            data_license = [wsj0_license, wham_noise_license]
        infos["licenses"] = data_license
        return infos
