import torch
import random
import numpy as np
import copy
import librosa
import torch.nn.functional as F
import torchaudio
import os

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