import io
import av
from torch.utils.data import Dataset as TorchDataset, ConcatDataset, WeightedRandomSampler
import torch
import numpy as np
import h5py
import os

from datasets.helpers.audiodatasets import PreprocessDataset, get_roll_func

# specify AudioSet location in 'dataset_dir'
# 3 files have to be located there:
# - balanced_train_segments_mp3.hdf
# - unbalanced_train_segments_mp3.hdf
# - eval_segments_mp3.hdf
# follow the instructions here to get these 3 files:
# https://github.com/kkoutini/PaSST/tree/main/audioset

dataset_dir = './datasets'
assert dataset_dir is not None, "Specify AudioSet location in variable 'dataset_dir'. " \
                                "Check out the Readme file for further instructions. " \
                                "https://github.com/fschmid56/EfficientAT/blob/main/README.md"

dataset_config = {
    'balanced_train_hdf5': os.path.join(dataset_dir, "train_mp3.hdf5"),
    #'unbalanced_train_hdf5': os.path.join(dataset_dir, "unbalanced_train_segments_mp3.hdf"),
    'eval_hdf5': os.path.join(dataset_dir, "test_mp3.hdf5"),
    'num_of_classes': 7
}


def decode_mp3(mp3_arr):
    """
    decodes an array if uint8 representing an mp3 file
    :rtype: np.array
    """
    container = av.open(io.BytesIO(mp3_arr.tobytes()))
    stream = next(s for s in container.streams if s.type == 'audio')
    # print(stream)
    a = []
    for i, packet in enumerate(container.demux(stream)):
        for frame in packet.decode():
            a.append(frame.to_ndarray().reshape(-1))
    waveform = np.concatenate(a)
    if waveform.dtype != 'float32':
        raise RuntimeError("Unexpected wave type")
    return waveform


# def pad_or_truncate(x, audio_length):
#     """Pad all audio to specific length."""
#     if len(x) <= audio_length:
#         return np.concatenate((x, np.zeros(audio_length - len(x), dtype=np.float32)), axis=0)
#     else:
#         return x[0: audio_length]

def pad_or_truncate(
    x,
    segment_length,
    method="topk_energy",       # "center" | "energy" | "topk_energy" | "sliding"
    top_k=5,                    # dùng khi method="topk_energy"
    stride=None,                # bước nhảy cho sliding window
    random_pick=False,           # nếu True: chọn ngẫu nhiên 1 window trong top-k energy
    random_pad_short=False,     # nếu True: short audio (<segment_length) sẽ random padding
    return_all=False,           # nếu True: trả về tất cả window thay vì 1 window duy nhất
    pad_value=0.0,              # giá trị dùng để padding
    debug=False                 # nếu True: in ra thông tin debug
):
    """
    Hàm tối ưu cho việc pad hoặc truncate audio cho task như emotion detection.
    
    1. Audio ngắn (< segment_length): pad thêm
       - random_pad_short=True: audio nằm ngẫu nhiên trong vùng padding
       - random_pad_short=False: audio được căn giữa (center-pad)
    
    2. Audio dài (>= segment_length): cắt window
       - center: lấy đoạn giữa
       - energy: lấy window có tổng năng lượng lớn nhất
       - topk_energy: chọn 1 window trong top-K energy (có thể random hoặc max)
       - sliding: trả về tất cả window theo stride
    
    Tham số:
    x              : np.array, audio 1D (mono)
    segment_length : int, số mẫu audio mong muốn
    method         : str, phương pháp cắt/pad:
                     "center"       -> lấy chính giữa
                     "energy"       -> lấy cửa sổ có năng lượng cao nhất
                     "topk_energy"  -> chọn top-K cửa sổ nhiều năng lượng, random hoặc max
                     "sliding"      -> tạo tất cả các cửa sổ trượt
    top_k          : int, số cửa sổ top-K năng lượng
    stride         : int, bước trượt cửa sổ, mặc định = segment_length // 2
    random_pick    : bool, nếu True chọn ngẫu nhiên 1 cửa sổ trong top-K
    return_all     : bool, nếu True trả về tất cả các cửa sổ, else chỉ 1 cửa sổ
    pad_value      : float, giá trị dùng để padding audio ngắn
    debug          : bool, in thông tin debug
    
    Trả về:
    out            : np.array (1D hoặc 2D nếu return_all=True)
    """
    x = np.asarray(x)
    if x.ndim != 1:
        raise ValueError("Expect mono audio")
    n = len(x)
    L = segment_length
    if L <= 0:
        raise ValueError("segment_length must >0")
    
    # ======= SHORT AUDIO =======
    if n < L:
        total_pad = L - n  # tổng số lượng pad cần thêm
        if random_pad_short:
            # Chọn left pad ngẫu nhiên trong khoảng [0, total_pad]
            left = np.random.randint(0, total_pad + 1)
        else:
            # Căn giữa mặc định
            left = total_pad // 2
        right = total_pad - left
        # ghép left padding + audio + right padding
        out = np.concatenate((
            np.full(left, pad_value, dtype=np.float32),
            x.astype(np.float32),
            np.full(right, pad_value, dtype=np.float32)
        ))
        return out if not return_all else out[np.newaxis,:]  # nếu return_all=True, shape=(1,L)

    # Long audio: sliding windows
    if stride is None:
        stride = L // 2
    starts = np.arange(0, n - L + 1, stride)
    if starts[-1] != n - L:
        starts = np.concatenate([starts, np.array([n-L])])  # ensure last window

    xsq = x.astype(np.float64)**2
    csum = np.cumsum(xsq)
    window_sums = csum[starts+L-1] - np.concatenate(([0.0], csum[starts[:-1]]))

    if debug:
        print(f"[Optimized] n={n}, L={L}, n_windows={len(starts)}")

    if method == "center":
        start = (n - L) // 2
        out = x[start:start+L].astype(np.float32)
        return out if not return_all else out[np.newaxis,:]

    if method == "energy":
        best_idx = np.argmax(window_sums)
        out = x[starts[best_idx]:starts[best_idx]+L].astype(np.float32)
        return out if not return_all else x[starts].astype(np.float32)

    if method == "topk_energy":
        k = min(top_k, len(starts))
        topk_idx = np.argsort(window_sums)[-k:]
        if return_all:
            ordered = topk_idx[np.argsort(window_sums[topk_idx])[::-1]]
            windows = np.stack([x[starts[i]:starts[i]+L].astype(np.float32) for i in ordered], axis=0)
            return windows
        else:
            chosen = np.random.choice(topk_idx) if random_pick else topk_idx[np.argmax(window_sums[topk_idx])]
            out = x[starts[chosen]:starts[chosen]+L].astype(np.float32)
            return out

    if method == "sliding":
        windows = np.stack([x[s:s+L].astype(np.float32) for s in starts], axis=0)
        return windows if return_all else windows[len(starts)//2]

    raise ValueError(f"Unknown method: {method}")


def pydub_augment(waveform, gain_augment=0):
    if gain_augment:
        gain = torch.randint(gain_augment * 2, (1,)).item() - gain_augment
        amp = 10 ** (gain / 20)
        waveform = waveform * amp
    return waveform


class MixupDataset(TorchDataset):
    """ Mixing Up wave forms
    """

    def __init__(self, dataset, beta=2, rate=0.5):
        self.beta = beta
        self.rate = rate
        self.dataset = dataset
        print(f"Mixing up waveforms from dataset of len {len(dataset)}")

    def __getitem__(self, index):
        if torch.rand(1) < self.rate:
            x1, f1, y1 = self.dataset[index]
            idx2 = torch.randint(len(self.dataset), (1,)).item()
            x2, f2, y2 = self.dataset[idx2]
            l = np.random.beta(self.beta, self.beta)
            l = max(l, 1. - l)
            x1 = x1 - x1.mean()
            x2 = x2 - x2.mean()
            x = (x1 * l + x2 * (1. - l))
            x = x - x.mean()
            return x, f1, (y1 * l + y2 * (1. - l))
        return self.dataset[index]

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


class AddIndexDataset(TorchDataset):
    def __init__(self, ds):
        self.ds = ds

    def __getitem__(self, index):
        x, f, y = self.ds[index]
        return x, f, y, index

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


class EmotionDataset(TorchDataset):
    def __init__(self, hdf5_file, sample_rate=32000, resample_rate=32000, classes_num=527,
                 clip_length=4, in_mem=False, gain_augment=0, is_train = True):
        """
        Reads the mp3 bytes from HDF file decodes using av and returns a fixed length audio wav
        """
        self.sample_rate = sample_rate
        self.resample_rate = resample_rate
        self.hdf5_file = hdf5_file
        if in_mem:
            print("\nPreloading in memory\n")
            with open(hdf5_file, 'rb') as f:
                self.hdf5_file = io.BytesIO(f.read())
        with h5py.File(hdf5_file, 'r') as f:
            self.length = len(f['audio_name'])
            print(f"Dataset from {hdf5_file} with length {self.length}.")
        self.dataset_file = None  # lazy init
        self.clip_length = clip_length * sample_rate
        self.classes_num = classes_num
        self.gain_augment = gain_augment
        self.is_train = is_train

    def open_hdf5(self):
        self.dataset_file = h5py.File(self.hdf5_file, 'r')
        #print(self.dataset_file)

    def __len__(self):
        return self.length

    def __del__(self):
        if self.dataset_file is not None:
            self.dataset_file.close()
            self.dataset_file = None

    def __getitem__(self, index):
        """Load waveform and target of an audio clip.
        Args:
          'index': int
        Returns:
          data_dict: {
            'audio_name': str,
            'waveform': (clip_samples,),
            'target': (classes_num,)}
        """
        if self.dataset_file is None:
            self.open_hdf5()
        try:
            audio_name = self.dataset_file['audio_name'][index].decode()
            # convert our modified filenames to official file names
            audio_name = audio_name.replace(".mp3", "")
            waveform = decode_mp3(self.dataset_file['mp3'][index])
            waveform = pydub_augment(waveform, self.gain_augment)
            if self.is_train:
                waveform = pad_or_truncate(waveform, self.clip_length, method="energy", random_pad_short = True, top_k=5, random_pick=True)
            else:
                waveform = pad_or_truncate(waveform, self.clip_length, method="energy", random_pad_short = True, return_all=False)
            waveform = self.resample(waveform)
            target = self.dataset_file['target'][index]
            target = np.unpackbits(target, axis=-1,
                                   count=self.classes_num).astype(np.float32)
            return waveform.reshape(1, -1), audio_name, target
        except Exception as e:
            print(f"Error processing index {index}: {e}")
            return self.__getitem__(index % len(self))
        

    def resample(self, waveform):
        """Resample.
        Args:
          waveform: (clip_samples,)
        Returns:
          (resampled_clip_samples,)
        """
        if self.resample_rate == 32000:
            return waveform
        elif self.resample_rate == 16000:
            return waveform[0:: 2]
        elif self.resample_rate == 8000:
            return waveform[0:: 4]
        else:
            raise Exception('Incorrect sample rate!')


# def get_ft_weighted_sampler(epoch_len=100000, sampler_replace=False):
#     samples_weights = get_ft_cls_balanced_sample_weights()
    
#     # Giới hạn epoch_len nếu replacement=False
#     if not sampler_replace:
#         epoch_len = min(epoch_len, len(samples_weights))
    
#     return WeightedRandomSampler(samples_weights, num_samples=epoch_len, replacement=sampler_replace)


# def get_ft_cls_balanced_sample_weights(sample_weight_offset=100, sample_weight_sum=True):
#     """
#     :return: float tensor of shape len(full_training_set) representing the weights of each sample.
#     """
#     # the order of balanced_train_hdf5, unbalanced_train_hdf5 is important.
#     # should match get_full_training_set
#     unbalanced_train_hdf5 = dataset_config['unbalanced_train_hdf5']
#     balanced_train_hdf5 = dataset_config['balanced_train_hdf5']
#     num_of_classes = dataset_config['num_of_classes']

#     all_y = []
#     for hdf5_file in [balanced_train_hdf5, unbalanced_train_hdf5]:
#         with h5py.File(hdf5_file, 'r') as dataset_file:
#             target = dataset_file['target']
#             target = np.unpackbits(target, axis=-1, count=num_of_classes)
#             all_y.append(target)

#     all_y = np.concatenate(all_y, axis=0)
#     all_y = torch.as_tensor(all_y)
    
#     per_class = all_y.long().sum(0).float().reshape(1, -1)  # frequencies per class

#     # Thêm offset cho các lớp có tần suất thấp
#     per_class = sample_weight_offset + per_class
#     if sample_weight_offset > 0:
#         print(f"Warning: sample_weight_offset={sample_weight_offset} minnow={per_class.min()}")
    
#     # Tính trọng số cho mỗi lớp, tránh chia cho 0
#     per_class_weights = torch.where(per_class > 0, torch.log(1000.0 / per_class), torch.tensor(float('inf')))
    
#     # Tính trọng số cho từng mẫu
#     all_weight = (all_y * per_class_weights)
    
#     # Đảm bảo rằng trọng số không có giá trị âm hoặc bằng 0
#     all_weight = torch.clamp(all_weight, min=1e-6)  # Đảm bảo không có trọng số bằng 0 hoặc âm
    
#     if sample_weight_sum:
#         all_weight = all_weight.sum(dim=1)
#     else:
#         all_weight, _ = all_weight.max(dim=1)
    
#     return all_weight



def get_ft_weighted_sampler(
    epoch_len=100000,
    sampler_replace=False,
    num_classes=dataset_config['num_of_classes'],
    balanced_train_hdf5=dataset_config['balanced_train_hdf5']
):
    """
    Create a WeightedRandomSampler for single-class classification
    using per-sample weights from packed one-hot targets.
    """
    # lấy sample_weights từ hàm single-class đã sửa
    sample_weights = get_ft_cls_balanced_sample_weights_single_class(
        num_of_classes=num_classes,
        balanced_train_hdf5=balanced_train_hdf5,
        normalize_weights=True,   # có thể bật/tắt theo nhu cầu
        debug=False
    )

    num_samples_available = len(sample_weights)

    # Giới hạn epoch_len
    if epoch_len is None or epoch_len > num_samples_available:
        epoch_len = num_samples_available

    return WeightedRandomSampler(
        sample_weights,
        num_samples=epoch_len,
        replacement=sampler_replace
    )


def get_ft_cls_balanced_sample_weights_single_class(
    sample_weight_offset=1e-6,
    num_of_classes=dataset_config['num_of_classes'],
    balanced_train_hdf5=dataset_config['balanced_train_hdf5'],
    normalize_weights=True,
    debug=False
):
    if num_of_classes is None:
        raise ValueError("num_of_classes must be provided")
    if not os.path.exists(balanced_train_hdf5):
        raise FileNotFoundError(f"HDF5 file not found: {balanced_train_hdf5}")

    with h5py.File(balanced_train_hdf5, "r") as f:
        target_ds = f["target"]
        packed = target_ds[:]   # (N, packed_len), dtype=uint8

    # --- unpack bits back to one-hot ---
    unpacked = np.unpackbits(packed, axis=-1, count=num_of_classes)
    if unpacked.shape[1] != num_of_classes:
        unpacked = unpacked[:, :num_of_classes]  # trim if over

    # --- convert one-hot to class indices ---
    all_y = np.argmax(unpacked, axis=1).astype(np.int64)

    if debug:
        print(f"Decoded targets: {all_y.shape}, unique={np.unique(all_y)}")

    # --- class frequency ---
    class_counts = np.bincount(all_y, minlength=num_of_classes).astype(np.float64)
    class_counts += sample_weight_offset

    n_samples = len(all_y)
    class_weights = n_samples / (num_of_classes * class_counts)

    if normalize_weights:
        class_weights /= class_weights.mean()

    sample_weights = class_weights[all_y]
    sample_weights = np.clip(sample_weights, 1e-8, 1e6).astype(np.float32)
    return torch.from_numpy(sample_weights)



def get_base_full_training_set(resample_rate=32000, gain_augment=0, num_classes = 527):
    sets = [get_base_training_set(resample_rate=resample_rate, gain_augment=gain_augment, num_classes = num_classes),
            get_unbalanced_training_set(resample_rate=resample_rate, gain_augment=gain_augment, num_classes = num_classes)]
    ds = ConcatDataset(sets)
    return ds


def get_base_training_set(resample_rate=32000, gain_augment=0, num_classes = 527):
    balanced_train_hdf5 = dataset_config['balanced_train_hdf5']
    ds = EmotionDataset(balanced_train_hdf5, resample_rate=resample_rate, gain_augment=gain_augment, classes_num = num_classes, is_train = True)
    return ds


def get_unbalanced_training_set(resample_rate=32000, gain_augment=0, num_classes = 527):
    unbalanced_train_hdf5 = dataset_config['unbalanced_train_hdf5']
    ds = EmotionDataset(unbalanced_train_hdf5, resample_rate=resample_rate, gain_augment=gain_augment, classes_num = num_classes, is_train = True)
    return ds


def get_base_test_set(resample_rate=32000, num_classes = 527):
    eval_hdf5 = dataset_config['eval_hdf5']
    ds = EmotionDataset(eval_hdf5, resample_rate=resample_rate, classes_num = num_classes, is_train = False)
    return ds


def get_training_set(add_index=True, roll=False, wavmix=False, gain_augment=0, resample_rate=32000, num_classes = 527):
    ds = get_base_training_set(resample_rate=resample_rate, gain_augment=gain_augment, num_classes = num_classes)
    if roll:
        ds = PreprocessDataset(ds, get_roll_func())
    if wavmix:
        ds = MixupDataset(ds)
    if add_index:
        ds = AddIndexDataset(ds)
    return ds


def get_full_training_set(add_index=True, roll=False, wavmix=False, gain_augment=0, resample_rate=32000, num_classes = 527):
    # global dataset_config
    # dataset_config['num_of_classes'] = num_classes
    ds = get_base_full_training_set(resample_rate=resample_rate, gain_augment=gain_augment, num_classes = num_classes)
    if roll:
        ds = PreprocessDataset(ds, get_roll_func())
    if wavmix:
        ds = MixupDataset(ds)
    if add_index:
        ds = AddIndexDataset(ds)
    return ds


def get_test_set(resample_rate=32000, num_classes = 527):
    ds = get_base_test_set(resample_rate=resample_rate, num_classes = num_classes)
    return ds
