import copy
import os
import random
import sys

import librosa
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio

sys.path.append(os.path.join("..", ".."))
import bat_detect.utils.audio_utils as au


def generate_gt_heatmaps(spec_op_shape, sampling_rate, ann, params):
    # spec may be resized on input into the network
    num_classes = len(params["class_names"])
    op_height = spec_op_shape[0]
    op_width = spec_op_shape[1]
    freq_per_bin = (params["max_freq"] - params["min_freq"]) / op_height

    # start and end times
    x_pos_start = au.time_to_x_coords(
        ann["start_times"],
        sampling_rate,
        params["fft_win_length"],
        params["fft_overlap"],
    )
    x_pos_start = (params["resize_factor"] * x_pos_start).astype(np.int)
    x_pos_end = au.time_to_x_coords(
        ann["end_times"],
        sampling_rate,
        params["fft_win_length"],
        params["fft_overlap"],
    )
    x_pos_end = (params["resize_factor"] * x_pos_end).astype(np.int)

    # location on y axis i.e. frequency
    y_pos_low = (ann["low_freqs"] - params["min_freq"]) / freq_per_bin
    y_pos_low = (op_height - y_pos_low).astype(np.int)
    y_pos_high = (ann["high_freqs"] - params["min_freq"]) / freq_per_bin
    y_pos_high = (op_height - y_pos_high).astype(np.int)
    bb_widths = x_pos_end - x_pos_start
    bb_heights = y_pos_low - y_pos_high

    valid_inds = np.where(
        (x_pos_start >= 0)
        & (x_pos_start < op_width)
        & (y_pos_low >= 0)
        & (y_pos_low < (op_height - 1))
    )[0]

    ann_aug = {}
    ann_aug["x_inds"] = x_pos_start[valid_inds]
    ann_aug["y_inds"] = y_pos_low[valid_inds]
    keys = [
        "start_times",
        "end_times",
        "high_freqs",
        "low_freqs",
        "class_ids",
        "individual_ids",
    ]
    for kk in keys:
        ann_aug[kk] = ann[kk][valid_inds]

    # if the number of calls is only 1, then it is unique
    # TODO would be better if we found these unique calls at the merging stage
    if len(ann_aug["individual_ids"]) == 1:
        ann_aug["individual_ids"][0] = 0

    y_2d_det = np.zeros((1, op_height, op_width), dtype=np.float32)
    y_2d_size = np.zeros((2, op_height, op_width), dtype=np.float32)
    # num classes and "background" class
    y_2d_classes = np.zeros(
        (num_classes + 1, op_height, op_width), dtype=np.float32
    )

    # create 2D ground truth heatmaps
    for ii in valid_inds:
        draw_gaussian(
            y_2d_det[0, :],
            (x_pos_start[ii], y_pos_low[ii]),
            params["target_sigma"],
        )
        # draw_gaussian(y_2d_det[0,:], (x_pos_start[ii], y_pos_low[ii]), params['target_sigma'], params['target_sigma']*2)
        y_2d_size[0, y_pos_low[ii], x_pos_start[ii]] = bb_widths[ii]
        y_2d_size[1, y_pos_low[ii], x_pos_start[ii]] = bb_heights[ii]

        cls_id = ann["class_ids"][ii]
        if cls_id > -1:
            draw_gaussian(
                y_2d_classes[cls_id, :],
                (x_pos_start[ii], y_pos_low[ii]),
                params["target_sigma"],
            )
            # draw_gaussian(y_2d_classes[cls_id, :], (x_pos_start[ii], y_pos_low[ii]), params['target_sigma'], params['target_sigma']*2)

    # be careful as this will have a 1.0 places where we have event but dont know gt class
    # this will be masked in training anyway
    y_2d_classes[num_classes, :] = 1.0 - y_2d_classes.sum(0)
    y_2d_classes = y_2d_classes / y_2d_classes.sum(0)[np.newaxis, ...]
    y_2d_classes[np.isnan(y_2d_classes)] = 0.0

    return y_2d_det, y_2d_size, y_2d_classes, ann_aug


def draw_gaussian(heatmap, center, sigmax, sigmay=None):
    # center is (x, y)
    # this edits the heatmap inplace

    if sigmay is None:
        sigmay = sigmax
    tmp_size = np.maximum(sigmax, sigmay) * 3
    mu_x = int(center[0] + 0.5)
    mu_y = int(center[1] + 0.5)
    w, h = heatmap.shape[0], heatmap.shape[1]
    ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
    br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]

    if ul[0] >= h or ul[1] >= w or br[0] < 0 or br[1] < 0:
        return False

    size = 2 * tmp_size + 1
    x = np.arange(0, size, 1, np.float32)
    y = x[:, np.newaxis]
    x0 = y0 = size // 2
    # g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
    g = np.exp(
        -((x - x0) ** 2) / (2 * sigmax**2)
        - ((y - y0) ** 2) / (2 * sigmay**2)
    )
    g_x = max(0, -ul[0]), min(br[0], h) - ul[0]
    g_y = max(0, -ul[1]), min(br[1], w) - ul[1]
    img_x = max(0, ul[0]), min(br[0], h)
    img_y = max(0, ul[1]), min(br[1], w)
    heatmap[img_y[0] : img_y[1], img_x[0] : img_x[1]] = np.maximum(
        heatmap[img_y[0] : img_y[1], img_x[0] : img_x[1]],
        g[g_y[0] : g_y[1], g_x[0] : g_x[1]],
    )
    return True


def pad_aray(ip_array, pad_size):
    return np.hstack((ip_array, np.ones(pad_size, dtype=np.int) * -1))


def warp_spec_aug(spec, ann, return_spec_for_viz, params):
    # This is messy
    # Augment spectrogram by randomly stretch and squeezing
    # NOTE this also changes the start and stop time in place

    # not taking care of spec for viz
    if return_spec_for_viz:
        assert False

    delta = params["stretch_squeeze_delta"]
    op_size = (spec.shape[1], spec.shape[2])
    resize_fract_r = np.random.rand() * delta * 2 - delta + 1.0
    resize_amt = int(spec.shape[2] * resize_fract_r)
    if resize_amt >= spec.shape[2]:
        spec_r = torch.cat(
            (
                spec,
                torch.zeros(
                    (1, spec.shape[1], resize_amt - spec.shape[2]),
                    dtype=spec.dtype,
                ),
            ),
            2,
        )
    else:
        spec_r = spec[:, :, :resize_amt]
    spec = F.interpolate(
        spec_r.unsqueeze(0), size=op_size, mode="bilinear", align_corners=False
    ).squeeze(0)
    ann["start_times"] *= 1.0 / resize_fract_r
    ann["end_times"] *= 1.0 / resize_fract_r
    return spec


def mask_time_aug(spec, params):
    # Mask out a random block of time - repeat up to 3 times
    # SpecAugment: A Simple Data Augmentation Methodfor Automatic Speech Recognition
    fm = torchaudio.transforms.TimeMasking(
        int(spec.shape[1] * params["mask_max_time_perc"])
    )
    for ii in range(np.random.randint(1, 4)):
        spec = fm(spec)
    return spec


def mask_freq_aug(spec, params):
    # Mask out a random frequncy range - repeat up to 3 times
    # SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition
    fm = torchaudio.transforms.FrequencyMasking(
        int(spec.shape[1] * params["mask_max_freq_perc"])
    )
    for ii in range(np.random.randint(1, 4)):
        spec = fm(spec)
    return spec


def scale_vol_aug(spec, params):
    return spec * np.random.random() * params["spec_amp_scaling"]


def echo_aug(audio, sampling_rate, params):
    sample_offset = (
        int(params["echo_max_delay"] * np.random.random() * sampling_rate) + 1
    )
    audio[:-sample_offset] += np.random.random() * audio[sample_offset:]
    return audio


def resample_aug(audio, sampling_rate, params):
    sampling_rate_old = sampling_rate
    sampling_rate = np.random.choice(params["aug_sampling_rates"])
    audio = librosa.resample(
        audio, sampling_rate_old, sampling_rate, res_type="polyphase"
    )

    audio = au.pad_audio(
        audio,
        sampling_rate,
        params["fft_win_length"],
        params["fft_overlap"],
        params["resize_factor"],
        params["spec_divide_factor"],
        params["spec_train_width"],
    )
    duration = audio.shape[0] / float(sampling_rate)
    return audio, sampling_rate, duration


def resample_audio(num_samples, sampling_rate, audio2, sampling_rate2):
    if sampling_rate != sampling_rate2:
        audio2 = librosa.resample(
            audio2, sampling_rate2, sampling_rate, res_type="polyphase"
        )
        sampling_rate2 = sampling_rate
    if audio2.shape[0] < num_samples:
        audio2 = np.hstack(
            (
                audio2,
                np.zeros((num_samples - audio2.shape[0]), dtype=audio2.dtype),
            )
        )
    elif audio2.shape[0] > num_samples:
        audio2 = audio2[:num_samples]
    return audio2, sampling_rate2


def combine_audio_aug(audio, sampling_rate, ann, audio2, sampling_rate2, ann2):

    # resample so they are the same
    audio2, sampling_rate2 = resample_audio(
        audio.shape[0], sampling_rate, audio2, sampling_rate2
    )

    # # set mean and std to be the same
    # audio2 = (audio2 - audio2.mean())
    # audio2 = (audio2/audio2.std())*audio.std()
    # audio2 = audio2 + audio.mean()

    if (
        ann["annotated"]
        and (ann2["annotated"])
        and (sampling_rate2 == sampling_rate)
        and (audio.shape[0] == audio2.shape[0])
    ):
        comb_weight = 0.3 + np.random.random() * 0.4
        audio = comb_weight * audio + (1 - comb_weight) * audio2
        inds = np.argsort(np.hstack((ann["start_times"], ann2["start_times"])))
        for kk in ann.keys():

            # when combining calls from different files, assume they come from different individuals
            if kk == "individual_ids":
                if (ann[kk] > -1).sum() > 0:
                    ann2[kk][ann2[kk] > -1] += (
                        np.max(ann[kk][ann[kk] > -1]) + 1
                    )

            if (kk != "class_id_file") and (kk != "annotated"):
                ann[kk] = np.hstack((ann[kk], ann2[kk]))[inds]

    return audio, ann


class AudioLoader(torch.utils.data.Dataset):
    def __init__(
        self, data_anns_ip, params, dataset_name=None, is_train=False
    ):

        self.data_anns = []
        self.is_train = is_train
        self.params = params
        self.return_spec_for_viz = False

        for ii in range(len(data_anns_ip)):
            dd = copy.deepcopy(data_anns_ip[ii])

            # filter out unused annotation here
            filtered_annotations = []
            for ii, aa in enumerate(dd["annotation"]):

                if "individual" in aa.keys():
                    aa["individual"] = int(aa["individual"])

                    # if only one call labeled it has to be from the same individual
                    if len(dd["annotation"]) == 1:
                        aa["individual"] = 0

                # convert class name into class label
                if aa["class"] in self.params["class_names"]:
                    aa["class_id"] = self.params["class_names"].index(
                        aa["class"]
                    )
                else:
                    aa["class_id"] = -1

                if aa["class"] not in self.params["classes_to_ignore"]:
                    filtered_annotations.append(aa)

            dd["annotation"] = filtered_annotations
            dd["start_times"] = np.array(
                [aa["start_time"] for aa in dd["annotation"]]
            )
            dd["end_times"] = np.array(
                [aa["end_time"] for aa in dd["annotation"]]
            )
            dd["high_freqs"] = np.array(
                [float(aa["high_freq"]) for aa in dd["annotation"]]
            )
            dd["low_freqs"] = np.array(
                [float(aa["low_freq"]) for aa in dd["annotation"]]
            )
            dd["class_ids"] = np.array(
                [aa["class_id"] for aa in dd["annotation"]]
            ).astype(np.int)
            dd["individual_ids"] = np.array(
                [aa["individual"] for aa in dd["annotation"]]
            ).astype(np.int)

            # file level class name
            dd["class_id_file"] = -1
            if "class_name" in dd.keys():
                if dd["class_name"] in self.params["class_names"]:
                    dd["class_id_file"] = self.params["class_names"].index(
                        dd["class_name"]
                    )

            self.data_anns.append(dd)

        ann_cnt = [len(aa["annotation"]) for aa in self.data_anns]
        self.max_num_anns = 2 * np.max(
            ann_cnt
        )  # x2 because we may be combining files during training

        print("\n")
        if dataset_name is not None:
            print("Dataset     : " + dataset_name)
        if self.is_train:
            print("Split type  : train")
        else:
            print("Split type  : test")
        print("Num files   : " + str(len(self.data_anns)))
        print("Num calls   : " + str(np.sum(ann_cnt)))

    def get_file_and_anns(self, index=None):

        # if no file specified, choose random one
        if index == None:
            index = np.random.randint(0, len(self.data_anns))

        audio_file = self.data_anns[index]["file_path"]
        sampling_rate, audio_raw = au.load_audio_file(
            audio_file,
            self.data_anns[index]["time_exp"],
            self.params["target_samp_rate"],
            self.params["scale_raw_audio"],
        )

        # copy annotation
        ann = {}
        ann["annotated"] = self.data_anns[index]["annotated"]
        ann["class_id_file"] = self.data_anns[index]["class_id_file"]
        keys = [
            "start_times",
            "end_times",
            "high_freqs",
            "low_freqs",
            "class_ids",
            "individual_ids",
        ]
        for kk in keys:
            ann[kk] = self.data_anns[index][kk].copy()

        # if train then grab a random crop
        if self.is_train:
            nfft = int(self.params["fft_win_length"] * sampling_rate)
            noverlap = int(self.params["fft_overlap"] * nfft)
            length_samples = (
                self.params["spec_train_width"] * (nfft - noverlap) + noverlap
            )

            if audio_raw.shape[0] - length_samples > 0:
                sample_crop = np.random.randint(
                    audio_raw.shape[0] - length_samples
                )
            else:
                sample_crop = 0
            audio_raw = audio_raw[sample_crop : sample_crop + length_samples]
            ann["start_times"] = ann["start_times"] - sample_crop / float(
                sampling_rate
            )
            ann["end_times"] = ann["end_times"] - sample_crop / float(
                sampling_rate
            )

        # pad audio
        if self.is_train:
            op_spec_target_size = self.params["spec_train_width"]
        else:
            op_spec_target_size = None
        audio_raw = au.pad_audio(
            audio_raw,
            sampling_rate,
            self.params["fft_win_length"],
            self.params["fft_overlap"],
            self.params["resize_factor"],
            self.params["spec_divide_factor"],
            op_spec_target_size,
        )
        duration = audio_raw.shape[0] / float(sampling_rate)

        # sort based on time
        inds = np.argsort(ann["start_times"])
        for kk in ann.keys():
            if (kk != "class_id_file") and (kk != "annotated"):
                ann[kk] = ann[kk][inds]

        return audio_raw, sampling_rate, duration, ann

    def __getitem__(self, index):

        # load audio file
        audio, sampling_rate, duration, ann = self.get_file_and_anns(index)

        # augment on raw audio
        if self.is_train and self.params["augment_at_train"]:
            # augment - combine with random audio file
            if (
                self.params["augment_at_train_combine"]
                and np.random.random() < self.params["aug_prob"]
            ):
                (
                    audio2,
                    sampling_rate2,
                    duration2,
                    ann2,
                ) = self.get_file_and_anns()
                audio, ann = combine_audio_aug(
                    audio, sampling_rate, ann, audio2, sampling_rate2, ann2
                )

            # simulate echo by adding delayed copy of the file
            if np.random.random() < self.params["aug_prob"]:
                audio = echo_aug(audio, sampling_rate, self.params)

            # resample the audio
            # if np.random.random() < self.params['aug_prob']:
            #   audio, sampling_rate, duration = resample_aug(audio, sampling_rate, self.params)

        # create spectrogram
        spec, spec_for_viz = au.generate_spectrogram(
            audio, sampling_rate, self.params, self.return_spec_for_viz
        )
        rsf = self.params["resize_factor"]
        spec_op_shape = (
            int(self.params["spec_height"] * rsf),
            int(spec.shape[1] * rsf),
        )

        # resize the spec
        spec = torch.from_numpy(spec).unsqueeze(0).unsqueeze(0)
        spec = F.interpolate(
            spec, size=spec_op_shape, mode="bilinear", align_corners=False
        ).squeeze(0)

        # augment spectrogram
        if self.is_train and self.params["augment_at_train"]:

            if np.random.random() < self.params["aug_prob"]:
                spec = scale_vol_aug(spec, self.params)

            if np.random.random() < self.params["aug_prob"]:
                spec = warp_spec_aug(
                    spec, ann, self.return_spec_for_viz, self.params
                )

            if np.random.random() < self.params["aug_prob"]:
                spec = mask_time_aug(spec, self.params)

            if np.random.random() < self.params["aug_prob"]:
                spec = mask_freq_aug(spec, self.params)

        outputs = {}
        outputs["spec"] = spec
        if self.return_spec_for_viz:
            outputs["spec_for_viz"] = torch.from_numpy(spec_for_viz).unsqueeze(
                0
            )

        # create ground truth heatmaps
        (
            outputs["y_2d_det"],
            outputs["y_2d_size"],
            outputs["y_2d_classes"],
            ann_aug,
        ) = generate_gt_heatmaps(
            spec_op_shape, sampling_rate, ann, self.params
        )

        # hack to get around requirement that all vectors are the same length in
        # the output batch
        pad_size = self.max_num_anns - len(ann_aug["individual_ids"])
        outputs["is_valid"] = pad_aray(
            np.ones(len(ann_aug["individual_ids"])), pad_size
        )
        keys = [
            "class_ids",
            "individual_ids",
            "x_inds",
            "y_inds",
            "start_times",
            "end_times",
            "low_freqs",
            "high_freqs",
        ]
        for kk in keys:
            outputs[kk] = pad_aray(ann_aug[kk], pad_size)

        # convert to pytorch
        for kk in outputs.keys():
            if type(outputs[kk]) != torch.Tensor:
                outputs[kk] = torch.from_numpy(outputs[kk])

        # scalars
        outputs["class_id_file"] = ann["class_id_file"]
        outputs["annotated"] = ann["annotated"]
        outputs["duration"] = duration
        outputs["sampling_rate"] = sampling_rate
        outputs["file_id"] = index

        return outputs

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