# Ke Chen
# knutchen@ucsd.edu
# HTS-AT: A HIERARCHICAL TOKEN-SEMANTIC AUDIO TRANSFORMER FOR SOUND CLASSIFICATION AND DETECTION
# Some Useful Common Methods

import numpy as np
import torch
import torch.nn as nn
from torch import Tensor
from typing import Optional
import logging
import os
import sys
import h5py
import csv
import time
import json
import museval
import librosa
from datetime import datetime
from tqdm import tqdm
from scipy import stats
import torch.nn as nn
import torch.nn.functional as F


# import from https://github.com/Alibaba-MIIL/ASL/blob/main/src/loss_functions/losses.py
class AsymmetricLoss(nn.Module):
    def __init__(self, gamma_neg=4, gamma_pos=1, clip=0.05, eps=1e-8, disable_torch_grad_focal_loss=True):
        super(AsymmetricLoss, self).__init__()

        self.gamma_neg = gamma_neg
        self.gamma_pos = gamma_pos
        self.clip = clip
        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
        self.eps = eps

    def forward(self, x, y):
        """"
        Parameters
        ----------
        x: input logits
        y: targets (multi-label binarized vector)
        """

        # Calculating Probabilities
        # x_sigmoid = torch.sigmoid(x)
        x_sigmoid = x # without sigmoid since it has been computed
        xs_pos = x_sigmoid
        xs_neg = 1 - x_sigmoid

        # Asymmetric Clipping
        if self.clip is not None and self.clip > 0:
            xs_neg = (xs_neg + self.clip).clamp(max=1)

        # Basic CE calculation
        los_pos = y * torch.log(xs_pos.clamp(min=self.eps))
        los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))
        loss = los_pos + los_neg

        # Asymmetric Focusing
        if self.gamma_neg > 0 or self.gamma_pos > 0:
            if self.disable_torch_grad_focal_loss:
                torch.set_grad_enabled(False)
            pt0 = xs_pos * y
            pt1 = xs_neg * (1 - y)  # pt = p if t > 0 else 1-p
            pt = pt0 + pt1
            one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)
            one_sided_w = torch.pow(1 - pt, one_sided_gamma)
            if self.disable_torch_grad_focal_loss:
                torch.set_grad_enabled(True)
            loss *= one_sided_w

        return -loss.mean()


def get_mix_lambda(mixup_alpha, batch_size):
    mixup_lambdas = [np.random.beta(mixup_alpha, mixup_alpha, 1)[0] for _ in range(batch_size)]
    return np.array(mixup_lambdas).astype(np.float32)

def create_folder(fd):
    if not os.path.exists(fd):
        os.makedirs(fd)

def dump_config(config, filename, include_time = False):
    save_time = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
    config_json = {}
    for key in dir(config):
        if not key.startswith("_"):
            config_json[key] = eval("config." + key)
    if include_time:
        filename = filename + "_" + save_time
    with open(filename + ".json", "w") as f:      
        json.dump(config_json, f ,indent=4)

def int16_to_float32(x):
    return (x / 32767.).astype(np.float32)

def float32_to_int16(x):
    x = np.clip(x, a_min = -1., a_max = 1.)
    return (x * 32767.).astype(np.int16)


# index for each class
def process_idc(index_path, classes_num, filename):
    # load data
    logging.info("Load Data...............")
    idc = [[] for _ in range(classes_num)]
    with h5py.File(index_path, "r") as f:
        for i in tqdm(range(len(f["target"]))):
            t_class = np.where(f["target"][i])[0]
            for t in t_class:
                idc[t].append(i)
    print(idc)
    np.save(filename, idc)
    logging.info("Load Data Succeed...............")

def clip_bce(pred, target):
    """Binary crossentropy loss.
    """
    return F.binary_cross_entropy(pred, target)
    # return F.binary_cross_entropy(pred, target)


def clip_ce(pred, target):
    return F.cross_entropy(pred, target)

def d_prime(auc):
    d_prime = stats.norm().ppf(auc) * np.sqrt(2.0)
    return d_prime


def get_loss_func(loss_type):
    if loss_type == 'clip_bce':
        return clip_bce
    if loss_type == 'clip_ce':
        return clip_ce
    if loss_type == 'asl_loss':
        loss_func = AsymmetricLoss(gamma_neg=4, gamma_pos=0,clip=0.05)
        return loss_func

def do_mixup_label(x):
    out = torch.logical_or(x, torch.flip(x, dims = [0])).float()
    return out

def do_mixup(x, mixup_lambda):
    """
    Args:
      x: (batch_size , ...)
      mixup_lambda: (batch_size,)

    Returns:
      out: (batch_size, ...)
    """
    out = (x.transpose(0,-1) * mixup_lambda + torch.flip(x, dims = [0]).transpose(0,-1) * (1 - mixup_lambda)).transpose(0,-1)
    return out
    
def interpolate(x, ratio):
    """Interpolate data in time domain. This is used to compensate the 
    resolution reduction in downsampling of a CNN.
    
    Args:
      x: (batch_size, time_steps, classes_num)
      ratio: int, ratio to interpolate

    Returns:
      upsampled: (batch_size, time_steps * ratio, classes_num)
    """
    (batch_size, time_steps, classes_num) = x.shape
    upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)
    upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)
    return upsampled


def pad_framewise_output(framewise_output, frames_num):
    """Pad framewise_output to the same length as input frames. The pad value 
    is the same as the value of the last frame.

    Args:
      framewise_output: (batch_size, frames_num, classes_num)
      frames_num: int, number of frames to pad

    Outputs:
      output: (batch_size, frames_num, classes_num)
    """
    pad = framewise_output[:, -1 :, :].repeat(1, frames_num - framewise_output.shape[1], 1)
    """tensor for padding"""

    output = torch.cat((framewise_output, pad), dim=1)
    """(batch_size, frames_num, classes_num)"""

    return output

# set the audio into the format that can be fed into the model
# resample -> convert to mono -> output the audio  
# track [n_sample, n_channel]
def prepprocess_audio(track, ofs, rfs, mono_type = "mix"):
    if track.shape[-1] > 1:
        # stereo
        if mono_type == "mix":
            track = np.transpose(track, (1,0))
            track = librosa.to_mono(track)
        elif mono_type == "left":
            track = track[:, 0]
        elif mono_type == "right":
            track = track[:, 1]
    else:
        track = track[:, 0]
    # track [n_sample]
    if ofs != rfs:
        track = librosa.resample(track, ofs, rfs)
    return track

def init_hier_head(class_map, num_class):
    class_map = np.load(class_map, allow_pickle = True)
    
    head_weight = torch.zeros(num_class,num_class).float()
    head_bias = torch.zeros(num_class).float()

    for i in range(len(class_map)):
        for d in class_map[i][1]:
            head_weight[d][i] = 1.0
        for d in class_map[i][2]:
            head_weight[d][i] = 1.0 / len(class_map[i][2])
        head_weight[i][i] = 1.0
    return head_weight, head_bias