#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from typing import List
from typing import Optional
from typing import Tuple
from typing import Union

import numpy as np
import torch


def chunk_audio(audio: np.ndarray, max_duration_seconds: int = 30, sample_rate: int = 16000) -> List[np.ndarray]:
    """split long audio into chunks

    Args:
        audio:
        max_duration_seconds:
        sample_rate:

    Returns:
        chunks
    """
    max_len = int(max_duration_seconds * sample_rate)

    if len(audio) <= max_len:
        return [audio]

    chunks = []
    for i in range(0, len(audio), max_len):
        chunk = audio[i : i + max_len]
        chunks.append(chunk)

    return chunks


def process_audio_batch(
    audios: Union[np.ndarray, List[np.ndarray], List[List[np.ndarray]]],
    feature_extractor,
    sampling_rate: int = 16000,
    max_duration_seconds: int = 30,
    return_attention_mask: bool = True,
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
    """extract audio mel features

    Args:
        audios:
        feature_extractor: WhisperFeatureExtractor
        sampling_rate:
        max_duration_seconds:
        return_attention_mask:

    Returns:
        (audio_features, audio_feature_lens)
        audio_features: [batch_size, n_mels, max_frames]
        audio_feature_lens:
    """
    if isinstance(audios, np.ndarray):
        audios_list = [[audios]]
    elif len(audios) > 0 and isinstance(audios[0], np.ndarray):
        audios_list = [audios]
    else:
        audios_list = audios

    audio_features_all = []
    audio_feature_lens_list = []

    for batch_audios in audios_list:
        batch_lens = []

        for audio in batch_audios:
            chunks = chunk_audio(audio, max_duration_seconds, sampling_rate)

            for chunk in chunks:
                audio_input = feature_extractor(
                    chunk,
                    sampling_rate=sampling_rate,
                    return_tensors="pt",
                    padding="max_length",
                    return_attention_mask=return_attention_mask,
                )

                audio_feature = audio_input["input_features"]  # [1, 80, frames]

                if return_attention_mask:
                    actual_len = audio_input["attention_mask"].sum(dim=1)  # Tensor([frames])
                    audio_feature = audio_feature[:, :, : actual_len[0]]
                    batch_lens.append(actual_len[0])
                else:
                    batch_lens.append(torch.tensor(audio_feature.shape[2]))

                audio_features_all.append(audio_feature.squeeze(0))  # [80, frames]

        if len(batch_lens) > 0:
            audio_feature_lens_list.append(torch.hstack(batch_lens))
        else:
            audio_feature_lens_list.append(torch.tensor([]))

    # pad to same length
    if audio_features_all:
        audio_features = torch.nn.utils.rnn.pad_sequence(
            [feat.transpose(0, 1) for feat in audio_features_all], batch_first=True, padding_value=0.0
        ).transpose(
            1, 2
        )  # [batch, 80, max_frames]
    else:
        audio_features = torch.tensor([])

    return audio_features, audio_feature_lens_list


def regroup_audio_features(
    audio_features: torch.Tensor, audio_feature_lens: List[torch.Tensor], regroup_seconds: int, fps: int = 100
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
    """regroup audio features to fixed duration

    Args:
        audio_features: [batch, n_mels, frames]
        audio_feature_lens: each batch's actual length
        regroup_seconds: regroup duration (seconds)
        fps: frames per second

    Returns:
        (regrouped_features, regrouped_lens)
    """
    # flatten to continuous frames sequence
    all_lens = []
    for lens in audio_feature_lens:
        if isinstance(lens, torch.Tensor):
            all_lens.extend(lens.tolist())
        elif isinstance(lens, list):
            all_lens.extend([int(x) for x in lens])

    if len(all_lens) == 0:
        return torch.tensor([]), []

    # concatenate all valid features
    flat_slices = [audio_features[i, :, :L] for i, L in enumerate(all_lens)]  # [n_mels, L]

    if len(flat_slices) == 1:
        full_feat = flat_slices[0]
    else:
        full_feat = torch.cat(flat_slices, dim=1)  # [n_mels, total_frames]

    # split to fixed frames
    frames_per_seg = int(regroup_seconds * fps)
    segments = []

    for start in range(0, full_feat.size(1), frames_per_seg):
        seg = full_feat[:, start : start + frames_per_seg]
        if seg.size(1) > 0:
            segments.append(seg)

    if len(segments) == 0:
        return torch.tensor([]), []

    # pad and convert to batch
    seg_lens = [s.size(1) for s in segments]
    segs_transposed = [s.transpose(0, 1) for s in segments]

    padded = torch.nn.utils.rnn.pad_sequence(segs_transposed, batch_first=True, padding_value=0.0)  # [N, max_T, n_mels]

    padded = padded.transpose(1, 2)  # [N, n_mels, max_T]
    lens_tensor = torch.tensor(seg_lens, dtype=torch.int32, device=padded.device)

    return padded, [lens_tensor]


def calculate_mel_frames(audio_samples: int, n_fft: int = 400, hop_length: int = 160) -> int:
    """calculate mel frames

    Args:
        audio_samples:
        n_fft:
        hop_length:

    Returns:
        mel frames
    """
    if audio_samples < n_fft:
        return 0
    return 1 + (audio_samples - n_fft) // hop_length


def samples_to_ms(samples: int, sample_rate: int = 16000) -> float:
    """convert samples to milliseconds

    Args:
        samples:
        sample_rate:

    Returns:
        milliseconds
    """
    return samples / sample_rate * 1000


def ms_to_samples(ms: float, sample_rate: int = 16000) -> int:
    """convert milliseconds to samples

    Args:
        ms: milliseconds
        sample_rate:

    Returns:
        samples
    """
    return int(ms * sample_rate / 1000)


def validate_audio_format(audio: np.ndarray, expected_rate: int = 16000, max_duration: Optional[float] = None) -> bool:
    """validate audio format

    Args:
        audio: audio data
        expected_rate:
        max_duration:

    Returns:
        whether valid audio format or not
    """
    if not isinstance(audio, np.ndarray):
        return False

    if audio.ndim != 1:
        return False

    if audio.dtype not in [np.float32, np.float64, np.int16, np.int32]:
        return False

    if max_duration is not None:
        max_samples = int(max_duration * expected_rate)
        if len(audio) > max_samples:
            return False

    return True