import argparse
import os
from unittest import result
import librosa
from pyparsing import results
from speechmos import dnsmos

def main(audio_files, output_file):
    # Configuración de sample rate original, ajústalo si es necesario
    config = {'sample_rate': 24000}  # Cambia 24000 por el sample rate real de tus audios si es distinto
    results_txt = []
    results = []
    file_names = []
    ovrl_cmos, sig_cmos, bak_mos, p808_mos = None, None, None, None
    for audio_path in audio_files:
        head, tail = os.path.split(audio_path)
        file_names.append(tail)
        if not os.path.isfile(audio_path):
            print(f"Archivo no encontrado: {audio_path}")
            results.append({'ovrl_mos': None, 'sig_mos': None, 'bak_mos': None, 'p808_mos': None})
            continue
        try:
            wav, sr = librosa.load(audio_path, sr=None)
            wav_resampled = librosa.resample(wav, orig_sr=sr,
                                            target_sr=16000, res_type='kaiser_best', fix=True)
            mos_dict = dnsmos.run(wav_resampled, sr=16000)
            results.append(mos_dict)
        except Exception as e:
            print(f"{tail}: ERROR - {str(e)}")
            results.append({'ovrl_mos': None, 'sig_mos': None, 'bak_mos': None, 'p808_mos': None})

    # Imprimir tabla de resultados
    header = f"{'Audio':<40} {'OVRL':>8} {'SIG':>8} {'BAK':>8} {'P808':>8}"
    print(header)
    print('-' * len(header))
    for i, res in enumerate(results):
        if res['ovrl_mos'] is not None:
            print(f"{file_names[i]:<40} {res['ovrl_mos']:8.3f} {res['sig_mos']:8.3f} {res['bak_mos']:8.3f} {res['p808_mos']:8.3f}")
        else:
            print(f"{file_names[i]:<40} {'ERROR':>8} {'ERROR':>8} {'ERROR':>8} {'ERROR':>8}")

    # Calcular CMOS si hay dos audios válidos
    if len(results) == 2 and all(r['ovrl_mos'] is not None for r in results):
        ovrl_cmos = results[0]['ovrl_mos'] - results[1]['ovrl_mos']
        sig_cmos = results[0]['sig_mos'] - results[1]['sig_mos']
        bak_mos = results[0]['bak_mos'] - results[1]['bak_mos']
        p808_mos = results[0]['p808_mos'] - results[1]['p808_mos']
        print("\nCMOS (Audio1 - Audio2):")
        print(f"{'CMOS OVRL':<12}: {ovrl_cmos:.3f}")
        print(f"{'CMOS SIG':<12}: {sig_cmos:.3f}")
        print(f"{'CMOS BAK':<12}: {bak_mos:.3f}")
        print(f"{'CMOS P808':<12}: {p808_mos:.3f}")
    else:
        print("\nNo es posible calcular CMOS")

    # Guardar tabla en archivo
    with open(output_file, "w") as f:
        f.write(header + "\n")
        f.write('-' * len(header) + "\n")
        for i, res in enumerate(results):
            if res['ovrl_mos'] is not None:
                f.write(f"{file_names[i]:<40} {res['ovrl_mos']:8.3f} {res['sig_mos']:8.3f} {res['bak_mos']:8.3f} {res['p808_mos']:8.3f}\n")
            else:
                f.write(f"{file_names[i]:<40} {'ERROR':>8} {'ERROR':>8} {'ERROR':>8} {'ERROR':>8}\n")
        if len(results) == 2 and all(r['ovrl_mos'] is not None for r in results):
            f.write("\nCMOS (Audio1 - Audio2):\n")
            f.write(f"{'CMOS OVRL':<12}: {ovrl_cmos:.3f}\n")
            f.write(f"{'CMOS SIG':<12}: {sig_cmos:.3f}\n")
            f.write(f"{'CMOS BAK':<12}: {bak_mos:.3f}\n")
            f.write(f"{'CMOS P808':<12}: {p808_mos:.3f}\n")
        else:
            f.write("\nNo es posible calcular CMOS\n")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evalúa audios con DNSMOS")
    parser.add_argument("--audios", nargs="+", help="Lista de archivos de audio a evaluar")
    parser.add_argument("--output", default="dnsmos_results.txt", help="Archivo de salida para los resultados")
    args = parser.parse_args()
    main(args.audios, args.output)