import numpy as np
import torch
from typing import Dict, List, Tuple, Optional

def chunk_audio_with_stride(
    audio: np.ndarray,
    sampling_rate: int,
    chunk_length_s: float = 30.0,
    stride_length_s: Tuple[float, float] = (5.0, 5.0)
) -> List[Dict]:
    """
    Split audio into overlapping chunks with stride.
    
    Args:
        audio: Audio array (mono, 1D)
        sampling_rate: Sampling rate of the audio
        chunk_length_s: Length of each chunk in seconds
        stride_length_s: Tuple of (left_stride, right_stride) in seconds
        
    Returns:
        List of dicts with keys:
            - 'audio': audio chunk array
            - 'stride': tuple of (left_stride_samples, right_stride_samples)
            - 'is_first': whether this is the first chunk
            - 'is_last': whether this is the last chunk
    """
    chunk_length = int(chunk_length_s * sampling_rate)
    if stride_length_s is None:
        stride_length_s = chunk_length_s / 6
    if isinstance(stride_length_s, (int, float)):
        stride_length_s = [stride_length_s, stride_length_s]
    stride_left = int(stride_length_s[0] * sampling_rate)
    stride_right = int(stride_length_s[1] * sampling_rate)

    audio_length = len(audio)
    
    # If audio is shorter than chunk, return as single chunk
    if audio_length <= chunk_length:
        return [{
            'audio': audio,
            'stride': (0, 0),
            'is_first': True,
            'is_last': True,
            'start_sample': 0,
            'end_sample': audio_length
        }]
    
    chunks = []
    start = 0
    
    while start < audio_length:

        end = min(start + chunk_length, audio_length)

        actual_stride_left = stride_left if start > 0 else 0
        actual_stride_right = stride_right if end < audio_length else 0

        chunk_start = max(0, start - actual_stride_left)
        chunk_end = min(audio_length, end + actual_stride_right)
        
        chunk_audio = audio[chunk_start:chunk_end]
        
        chunks.append({
            'audio': chunk_audio,
            'stride': (actual_stride_left, actual_stride_right),
            'is_first': (start == 0),
            'is_last': (end >= audio_length),
            'start_sample': start,
            'end_sample': end,
            'chunk_start': chunk_start,
            'chunk_end': chunk_end
        })

        start = end
    
    return chunks

def process_logits_with_stride(
    all_logits: List[torch.Tensor],
    chunks_info: List[Dict],
    model_ratio: float
) -> torch.Tensor:
    """
    Process logits from all chunks, removing stride parts and concatenating.
    
    Args:
        all_logits: List of logit tensors from each chunk [batch, time, vocab]
        chunks_info: List of chunk info dicts from chunk_audio_with_stride
        model_ratio: The downsampling ratio of the model (e.g. 320 for wav2vec2)
                    This is the ratio between input samples and output frames
        
    Returns:
        Concatenated logits tensor
    """
    processed_logits = []
    
    for i, (logits, chunk_info) in enumerate(zip(all_logits, chunks_info)):
        stride_left_frames = int(chunk_info['stride'][0] / model_ratio)
        stride_right_frames = int(chunk_info['stride'][1] / model_ratio)

        if len(logits.shape) == 3:
            logits = logits[0]  # [time, vocab]

        start_frame = stride_left_frames
        end_frame = logits.shape[0] - stride_right_frames
        
        trimmed_logits = logits[start_frame:end_frame]
        processed_logits.append(trimmed_logits)

    concatenated_logits = torch.cat(processed_logits, dim=0)
    
    return concatenated_logits.unsqueeze(0)

def infer_fairseq_with_chunking(
    audio: np.ndarray,
    sampling_rate: int,
    model: torch.nn.Module,
    dictionary: Dict,
    device: torch.device,
    normalize_audio: bool = True,
    chunk_length_s: float = 30.0,
    stride_length_s: Tuple[float, float] = (5.0, 5.0),
    model_downsample_ratio: float = 320.0,
    decode_fn: callable = None
) -> str:
    """
    Run fairseq inference with chunking and stride.
    
    Args:
        audio: Audio array (mono, 1D)
        sampling_rate: Sampling rate
        model: Fairseq checkpoint
        dictionary: Token dictionary for decoding
        device: torch device
        chunk_length_s: Length of each chunk in seconds
        stride_length_s: Tuple of (left_stride, right_stride) in seconds
        model_downsample_ratio: Downsampling ratio of the model
        decode_fn: Optional custom decoding function. If None, will use simple argmax
        
    Returns:
        ASR transcript
    """

    chunks = chunk_audio_with_stride(
        audio, 
        sampling_rate, 
        chunk_length_s, 
        stride_length_s
    )
    
    all_logits = []
    model.eval()
    
    with torch.no_grad():
        for chunk_info in chunks:
            chunk_audio = chunk_info['audio']

            # Normalize if needed
            if normalize_audio:
                chunk_audio = (chunk_audio - np.mean(chunk_audio)) / (np.std(chunk_audio) + 1e-5)

            try:
                chunk_tensor = torch.from_numpy(chunk_audio).float().unsqueeze(0).to(device)
            except TypeError:
                chunk_tensor = torch.FloatTensor(chunk_audio).unsqueeze(0).to(device)

            output = model(source=chunk_tensor, padding_mask=None)

            if isinstance(output, dict):
                logits = output['encoder_out']  # [time, batch, vocab] for fairseq
                logits = logits.transpose(0, 1)  # [batch, time, vocab]
            else:
                logits = output
            
            all_logits.append(logits.cpu())

    concatenated_logits = process_logits_with_stride(
        all_logits,
        chunks,
        model_downsample_ratio
    )

    if decode_fn is not None:
        transcription = decode_fn(concatenated_logits, dictionary)
    else:
        transcription = simple_ctc_decode(concatenated_logits, dictionary)
    
    return transcription

def simple_ctc_decode(emissions: torch.Tensor, dictionary: Dict) -> str:
    """
    Simple greedy CTC decoding.
    
    Args:
        emissions: Logits tensor [batch, time, vocab]
        dictionary: Token dictionary {id: token}
        
    Returns:
        Decoded string
    """
    # Get the most likely token at each timestep
    predictions = emissions.argmax(dim=-1)
    
    vocab_size = emissions.shape[-1]
    blank_idx = vocab_size - 1
    
    # Convert to tokens
    tokens = []
    prev_token_id = None
    
    for pred in predictions[0]:
        token_id = pred.item()
        
        # Skip blank tokens
        if token_id == blank_idx:
            prev_token_id = None
            continue
        
        # Skip special tokens and repeated tokens
        if token_id == prev_token_id or token_id < 4:  # Skip <s>, <pad>, </s>, <unk>
            if token_id != prev_token_id and token_id < 4:
                prev_token_id = None
            continue
        
        # Get token from dictionary
        token = dictionary.get(token_id, f'<unk_{token_id}>')
        tokens.append(token)
        prev_token_id = token_id
    
    # Join character tokens and replace | with space
    transcription = ''.join(tokens).replace('|', ' ').strip()
    return transcription

def map_to_pred_fairseq_chunked(
    batch: Dict,
    model: torch.nn.Module,
    dictionary: Dict,
    device: torch.device,
    normalize_audio: bool = True,
    chunk_length_s: float = 30.0,
    stride_length_s: Tuple[float, float] = (5.0, 5.0),
    model_downsample_ratio: float = 320.0,
    decode_fn: callable = None
) -> Dict:
    """
    Map function compatible with HuggingFace datasets that uses chunking.
    
    Args:
        batch: Batch dict with 'audio' and 'sampling_rate' keys
        model: Fairseq checkpoint
        dictionary: Token dictionary
        device: torch device
        chunk_length_s: Chunk length in seconds
        stride_length_s: Stride lengths (left, right) in seconds
        model_downsample_ratio: Model downsampling ratio
        decode_fn: Optional custom decoding function
        
    Returns:
        Dict with 'prediction' key added
    """
    transcriptions = []
    
    for i in range(len(batch['audio'])):
        audio = batch['audio'][i]
        sr = batch['sampling_rate'][i]

        transcription = infer_fairseq_with_chunking(
            audio=audio,
            sampling_rate=sr,
            model=model,
            dictionary=dictionary,
            device=device,
            normalize_audio=normalize_audio,
            chunk_length_s=chunk_length_s,
            stride_length_s=stride_length_s,
            model_downsample_ratio=model_downsample_ratio,
            decode_fn=decode_fn
        )
        
        transcriptions.append(transcription)
    
    batch['prediction'] = transcriptions
    return batch