import torch
import librosa
import numpy as np
import sys
from transformers import AutoFeatureExtractor, ASTForAudioClassification

# Modelo AST (Audio Spectrogram Transformer)
# Usamos o modelo base do MIT como referência para análise espectral
MODEL_NAME = "MIT/ast-finetuned-audioset-10-10-0.4593"

# Singleton para carregar o modelo apenas uma vez
_extractor = None
_model = None

def get_ast_resources():
    global _extractor, _model
    if _extractor is None or _model is None:
        print(f"Carregando motor AST: {MODEL_NAME}...")
        _extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
        model = ASTForAudioClassification.from_pretrained(MODEL_NAME)
        
        # --- OTIMIZAÇÃO: Quantização Dinâmica (CPU) ---
        if not torch.cuda.is_available():
            print("Aplicando Quantização Dinâmica no AST...", file=sys.stderr)
            model = torch.quantization.quantize_dynamic(
                model, {torch.nn.Linear}, dtype=torch.qint8
            )
        
        _model = model
        _model.eval()
    return _extractor, _model

def run_ast_inference(file_path):
    """
    Executa a análise via Audio Spectrogram Transformer.
    Focado em detectar artefatos de compressão e texturas sintéticas.
    """
    try:
        extractor, model = get_ast_resources()

        # Carrega áudio (resample para 16kHz conforme exigido pelo AST)
        audio, _ = librosa.load(file_path, sr=16000)
        
        # Padronização para 10 segundos
        if len(audio) > 160000:
            audio = audio[:160000]
        else:
            audio = np.pad(audio, (0, 160000 - len(audio)), mode='constant')

        # Extração de Features
        inputs = extractor(audio, sampling_rate=16000, return_tensors="pt")

        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            
        # Refinamento do Score Espectral:
        # Em vez de apenas variância, olhamos para a dispersão das top-k classes.
        # Vozes sintéticas tendem a gerar incerteza em modelos genéricos.
        probs = torch.nn.functional.softmax(logits, dim=-1)
        
        # Cálculo de Entropia como medida de 'estranheza' do sinal
        entropy = -torch.sum(probs * torch.log(probs + 1e-9)).item()
        
        # Normalização (Heurística baseada em observação do AudioSet)
        # Sinais naturais tendem a ter entropia mais baixa (classes mais claras)
        # Sinais sintéticos/ruidosos espalham a probabilidade.
        risk_score = min(max((entropy - 3.0) / 4.0, 0.0), 1.0) 

        return {
            "risk_score": risk_score,
            "engine": "AST-Transformer",
            "status": "success"
        }

    except Exception as e:
        print(f"Erro no motor AST: {e}")
        return {"error": str(e), "risk_score": 0.0}

if __name__ == "__main__":
    # Teste simples
    import sys
    if len(sys.argv) > 1:
        print(run_ast_inference(sys.argv[1]))