import argparse
import sys
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
import torchaudio
from transformers import WavLMModel

# Ensure UTF-8 output encoding for Windows terminal compatibility
if sys.stdout.encoding and sys.stdout.encoding.lower() != 'utf-8':
    try:
        sys.stdout.reconfigure(encoding='utf-8')
    except Exception:
        pass

# ============================================================
# CONFIGURATION & CONSTANTS
# ============================================================

WAVLM_MODEL_NAME = "microsoft/wavlm-base-plus"
CHECKPOINT_PATH = Path("checkpoints/best_model.pt")
INPUT_DIR = Path("input")
TARGET_SR = 16000

EMOTION_CLASSES = ["Anger", "Disgust", "Fear", "Happy", "Neutral", "Sad"]

# ============================================================
# MODEL ARCHITECTURE DEFINITION (Self-Contained Single File)
# ============================================================

class BiLSTMFeatureExtractor(nn.Module):
    def __init__(self, input_size=768, hidden_size=128, num_layers=1, dropout=0.0):
        super(BiLSTMFeatureExtractor, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bidirectional = True
        
        self.bilstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if num_layers > 1 else 0.0
        )
        
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
        
        if mask is not None:
            lengths = mask.sum(dim=1).cpu()
            packed_x = pack_padded_sequence(
                x,
                lengths,
                batch_first=True,
                enforce_sorted=False
            )
            packed_out, (hn, cn) = self.bilstm(packed_x)
            out, _ = pad_packed_sequence(
                packed_out,
                batch_first=True,
                total_length=seq_len
            )
        else:
            out, (hn, cn) = self.bilstm(x)
            
        return out


class TemporalAttention(nn.Module):
    def __init__(self, input_dim=256):
        super(TemporalAttention, self).__init__()
        self.input_dim = input_dim
        self.w = nn.Linear(input_dim, 1, bias=False)
        
    def forward(self, h, mask=None):
        scores = self.w(torch.tanh(h)).squeeze(-1)  # [B, T]
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
            
        attn_weights = torch.softmax(scores, dim=1)  # [B, T]
        context = torch.bmm(attn_weights.unsqueeze(1), h).squeeze(1)  # [B, 256]
        
        return context, attn_weights


class BiLSTMAttentionClassifier(nn.Module):
    def __init__(self, input_size=768, hidden_size=128, num_classes=6, dropout=0.3):
        super(BiLSTMAttentionClassifier, self).__init__()
        self.bilstm = BiLSTMFeatureExtractor(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
            dropout=dropout
        )
        context_dim = hidden_size * 2  # 256
        self.attention = TemporalAttention(input_dim=context_dim)
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(context_dim, num_classes)
        
    def forward(self, x, mask=None):
        bilstm_out = self.bilstm(x, mask=mask)
        context, attn_weights = self.attention(bilstm_out, mask=mask)
        dropped_context = self.dropout(context)
        logits = self.classifier(dropped_context)
        return logits, attn_weights

# ============================================================
# EXACT AUDIO PREPROCESSING LOGIC (Used During Training)
# ============================================================

def preprocess_audio(audio_path):
    """
    Load raw audio file (.wav, .mp3, .flac, .ogg), convert to 16 kHz mono, and peak-normalize amplitude.
    """
    audio_path = Path(audio_path)
    
    # 1. Load audio with torchaudio, fallback to soundfile for mp3/flac if needed
    try:
        waveform, sample_rate = torchaudio.load(str(audio_path))
    except Exception:
        import soundfile as sf
        data, sample_rate = sf.read(str(audio_path))
        waveform = torch.tensor(data, dtype=torch.float32)
        if waveform.ndim == 1:
            waveform = waveform.unsqueeze(0)
        elif waveform.ndim == 2:
            waveform = waveform.T
            
    # 2. Convert to mono if multi-channel
    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)
        
    # 3. Resample to 16 kHz if necessary
    if sample_rate != TARGET_SR:
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=TARGET_SR)
        waveform = resampler(waveform)
        
    # 4. Squeeze channel dimension [1, num_samples] -> [num_samples]
    waveform = waveform.squeeze(0)
    
    # 5. Amplitude peak normalization
    max_val = waveform.abs().max()
    if max_val > 0:
        waveform = waveform / max_val
        
    return waveform

# ============================================================
# FINAL PREDICTOR PIPELINE CLASS
# ============================================================

class FinalPredictor:
    def __init__(self, checkpoint_path=CHECKPOINT_PATH):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("=" * 60)
        print("INITIALIZING F1 DRIVER TONE DETECTOR PIPELINE")
        print("=" * 60)
        print(f"Compute Device : {self.device}")
        if self.device.type == "cuda":
            print(f"GPU            : {torch.cuda.get_device_name(0)}")
            
        print(f"\n[1] Loading WavLM Encoder ({WAVLM_MODEL_NAME})...")
        self.wavlm = WavLMModel.from_pretrained(WAVLM_MODEL_NAME).to(self.device)
        self.wavlm.eval()
        
        print(f"[2] Loading Trained Downstream Model ({checkpoint_path})...")
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Checkpoint file not found at {checkpoint_path}. Please run training first!")
            
        self.classifier = BiLSTMAttentionClassifier(
            input_size=768,
            hidden_size=128,
            num_classes=6
        ).to(self.device)
        
        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        self.classifier.load_state_dict(checkpoint["model_state_dict"])
        self.classifier.eval()
        print("Pipeline initialized and ready for inference!")
        print("=" * 60)
        
    def predict_single(self, audio_path, chunk_duration=2.5, hop_duration=1.5):
        audio_path = Path(audio_path)
        if not audio_path.exists():
            raise FileNotFoundError(f"Audio file not found: {audio_path}")
            
        # 1. Preprocess raw audio waveform [num_samples]
        waveform = preprocess_audio(audio_path)
        num_samples = waveform.size(0)
        duration_sec = num_samples / float(TARGET_SR)
        
        chunk_samples = int(chunk_duration * TARGET_SR)
        hop_samples = int(hop_duration * TARGET_SR)
        
        chunk_probs = []
        last_attn_weights = None
        
        with torch.no_grad():
            # If audio is long (> 4.0s), use 2.5s sliding window chunking
            if duration_sec > 4.0 and num_samples > chunk_samples:
                starts = list(range(0, num_samples - chunk_samples + 1, hop_samples))
                # Ensure the end of audio is covered
                if starts[-1] + chunk_samples < num_samples:
                    starts.append(num_samples - chunk_samples)
                    
                for start in starts:
                    end = start + chunk_samples
                    chunk_wave = waveform[start:end].unsqueeze(0).to(self.device)
                    
                    outputs = self.wavlm(input_values=chunk_wave)
                    embedding = outputs.last_hidden_state
                    mask = torch.ones((1, embedding.size(1)), dtype=torch.int64, device=self.device)
                    
                    logits, attn_weights = self.classifier(embedding, mask=mask)
                    probs = F.softmax(logits, dim=-1).squeeze(0)
                    chunk_probs.append(probs)
                    last_attn_weights = attn_weights.squeeze(0).cpu().numpy()
                    
                probabilities = torch.stack(chunk_probs).mean(dim=0)
            else:
                input_values = waveform.unsqueeze(0).to(self.device)
                outputs = self.wavlm(input_values=input_values)
                embedding = outputs.last_hidden_state
                mask = torch.ones((1, embedding.size(1)), dtype=torch.int64, device=self.device)
                logits, attn_weights = self.classifier(embedding, mask=mask)
                probabilities = F.softmax(logits, dim=-1).squeeze(0)
                last_attn_weights = attn_weights.squeeze(0).cpu().numpy()
            
        pred_id = torch.argmax(probabilities).item()
        pred_emotion = EMOTION_CLASSES[pred_id]
        confidence = probabilities[pred_id].item() * 100.0
        
        probs_dict = {
            EMOTION_CLASSES[i]: probabilities[i].item() * 100.0
            for i in range(len(EMOTION_CLASSES))
        }
        
        return {
            "audio_file": audio_path.name,
            "duration_sec": round(duration_sec, 2),
            "predicted_emotion": pred_emotion,
            "confidence": confidence,
            "probabilities": probs_dict,
            "attention_weights": last_attn_weights,
            "chunks_processed": len(chunk_probs) if chunk_probs else 1
        }

# ============================================================
# MAIN CLI DRIVER
# ============================================================

def print_report(result):
    print("\n" + "=" * 60)
    print(f"PREDICTION REPORT: {result['audio_file']}")
    print("=" * 60)
    print(f"Predicted Emotion : {result['predicted_emotion']} ({result['confidence']:.2f}% confidence)")
    print("-" * 60)
    print("EMOTION PROBABILITY BREAKDOWN:")
    for emotion, prob in result["probabilities"].items():
        bar = "#" * int(prob / 5)
        print(f"  {emotion:10s} : {prob:6.2f}% | {bar}")
    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(description="F1 Driver Tone Predictor (Single File Pipeline)")
    parser.add_argument("--audio_path", type=str, default=None,
                        help="Path to specific raw driver audio file (.wav, .mp3, .flac, .ogg)")
    args = parser.parse_args()

    predictor = FinalPredictor()

    if args.audio_path:
        target_path = Path(args.audio_path)
        result = predictor.predict_single(target_path)
        print_report(result)
    else:
        # Scan input directory
        INPUT_DIR.mkdir(parents=True, exist_ok=True)
        valid_extensions = {".wav", ".mp3", ".flac", ".ogg", ".m4a"}
        audio_files = [
            f for f in INPUT_DIR.iterdir()
            if f.is_file() and f.suffix.lower() in valid_extensions
        ]

        if not audio_files:
            print(f"\n[INFO] No audio files found in input/ directory ({INPUT_DIR.resolve()}).")
            print("Usage Options:")
            print("  1. Place .wav or .mp3 files into the 'input/' folder and re-run python finalpredictor.py")
            print("  2. Run with explicit file path: python finalpredictor.py --audio_path <path_to_audio_file>")
        else:
            print(f"\nFound {len(audio_files)} audio file(s) in {INPUT_DIR}/ directory. Processing...")
            for audio_file in audio_files:
                result = predictor.predict_single(audio_file)
                print_report(result)


if __name__ == "__main__":
    main()