# app_full_reimpl.py
"""
Full app: MDES (BPObox) reimplementation + audio/text emotion helpers + Gradio UI + CLI test.
- Default ASR: Whisper 'tiny' (fast & small). Change WHISPER_MODEL near top to 'base' if you have memory.
- CLI quick test:
    python app_full_reimpl.py --agent agent_mono16.wav --customer customer_mono16.wav --debug
- To run Gradio:
    python app_full_reimpl.py
- Debug output: /tmp/mdes_debug.json
"""

import argparse
import json
import math
import os
import sys
import warnings
from typing import Dict, List

import librosa
import numpy as np

warnings.filterwarnings("ignore")

# ---------------- CONFIG ----------------
WHISPER_MODEL = os.environ.get("WHISPER_MODEL", "tiny")  # default tiny (change to "base" if you want)
EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"  # optional
TEXT_CLASSIFIER_NAME = "bhadresh-savani/distilbert-base-uncased-emotion"  # optional
AUDIO_EMOTION_MODEL_PATH = os.environ.get("AUDIO_EMOTION_MODEL", "CNN_LSTM.h5")  # optional existing model
DEFAULT_EMPATHY_THRESHOLD = 0.40
DEBUG_PERMISSIVE_THRESHOLD = 0.18
DEBUG_DUMP_PATH = "/tmp/mdes_debug.json"
USE_GRADIO = True

# ---------------- lazy imports for optional libs ----------------
try:
    import whisper
except Exception:
    whisper = None

try:
    import torch
    from transformers import AutoModel, AutoTokenizer, pipeline
except Exception:
    torch = None
    AutoModel = None
    AutoTokenizer = None
    pipeline = None

# ---------------- small helpers & salvage of existing runner ----------------
# The code below preserves your audio-emotion runner (prepare_data + runner)
try:
    import tensorflow as tf
    from keras.initializers import Orthogonal
except Exception:
    tf = None

CATEGORIES = ['Neutral', 'Happy', 'Sad', 'Angry', 'Fear', 'Disgust']

def prepare_data_for_audio_emotion(audio_path):
    """Salvaged prepare_data from your prior app (keeps same frame settings)."""
    try:
        raw_audio, sr = librosa.load(audio_path, sr=16000)
    except Exception:
        return None
    raw_audio, _ = librosa.effects.trim(raw_audio, top_db=25, frame_length=256, hop_length=64)
    audio_duration = len(raw_audio) / 16000.0
    if audio_duration > 4:
        raw_audio = raw_audio[:4 * 16000]
    else:
        raw_audio = np.pad(raw_audio, (0, (4 * 16000) - len(raw_audio)), 'constant')

    FRAME_LENGTH = 400
    HOP_LENGTH = 160
    sr = 16000
    y = raw_audio

    # may raise if audio very short; handle in runner
    try:
        zcr = librosa.feature.zero_crossing_rate(y, frame_length=FRAME_LENGTH, hop_length=HOP_LENGTH)
        rms = librosa.feature.rms(y=y, frame_length=FRAME_LENGTH, hop_length=HOP_LENGTH)
        mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=20, hop_length=HOP_LENGTH)
        pda = np.concatenate((zcr, rms, mfccs), axis=0)
        pda = pda.astype('float32')
        pda = np.expand_dims(pda, axis=0)
        return pda
    except Exception:
        return None

# Attempt to load your audio model if present
audio_emotion_model = None
if tf is not None and os.path.exists(AUDIO_EMOTION_MODEL_PATH):
    try:
        audio_emotion_model = tf.keras.models.load_model(AUDIO_EMOTION_MODEL_PATH, custom_objects={'Orthogonal': Orthogonal})
        print("[AUDIO_MODEL] Loaded audio emotion model:", AUDIO_EMOTION_MODEL_PATH)
    except Exception as e:
        print("[AUDIO_MODEL] Could not load model at", AUDIO_EMOTION_MODEL_PATH, ":", e)
        audio_emotion_model = None
else:
    if tf is None:
        print("[AUDIO_MODEL] TensorFlow not available; audio model disabled.")
    else:
        print("[AUDIO_MODEL] No audio model found at", AUDIO_EMOTION_MODEL_PATH)

def runner_audio_emotion(audio_path):
    """Return audio-only emotion probs or zeros if model missing."""
    if audio_emotion_model is None:
        return {c: 0.0 for c in CATEGORIES}
    features = prepare_data_for_audio_emotion(audio_path)
    if features is None:
        return {c: 0.0 for c in CATEGORIES}
    try:
        pr = audio_emotion_model.predict(features)
        out = {CATEGORIES[i]: float(np.round(pr[0, i], 3)) for i in range(len(CATEGORIES))}
        # normalize if needed
        s = sum(out.values())
        if s > 0:
            for k in out:
                out[k] = round(out[k] / s, 3)
        return out
    except Exception as e:
        print("[AUDIO_MODEL] inference error:", e)
        return {c: 0.0 for c in CATEGORIES}

# ---------------- ASR loader (Whisper) with fallback ----------------
def load_whisper_model(name=WHISPER_MODEL):
    global whisper
    if whisper is None:
        try:
            import whisper as _w
            whisper = _w
        except Exception as e:
            print("[ASR] whisper import failed:", e)
            return None
    try:
        print(f"[ASR] Loading Whisper model '{name}' ...")
        model = whisper.load_model(name)
        print(f"[ASR] Loaded Whisper model '{name}'")
        return model
    except Exception as e:
        print(f"[ASR] Failed to load Whisper '{name}':", e)
        if name != "tiny":
            try:
                print("[ASR] Attempting fallback to 'tiny' ...")
                model = whisper.load_model("tiny")
                print("[ASR] Loaded Whisper model 'tiny' fallback")
                return model
            except Exception as e2:
                print("[ASR] Failed to load fallback 'tiny':", e2)
                return None
        return None

# ---------------- embedding and text classifier loaders ----------------
def try_load_embedding(name=EMBEDDING_MODEL_NAME):
    if AutoTokenizer is None or AutoModel is None:
        print("[EMB] transformers not available; embedding disabled.")
        return None, None
    try:
        print("[EMB] Loading embedding model (may take time)...")
        tok = AutoTokenizer.from_pretrained(name)
        mdl = AutoModel.from_pretrained(name)
        device = "cuda" if torch and torch.cuda.is_available() else "cpu"
        mdl.to(device)
        print("[EMB] Embedding model loaded.")
        return tok, mdl
    except Exception as e:
        print("[EMB] Failed to load embedding model:", e)
        return None, None

def try_load_text_classifier(name=TEXT_CLASSIFIER_NAME):
    if pipeline is None:
        print("[TXT] transformers pipeline not available; text classifier disabled.")
        return None
    try:
        print("[TXT] Loading text classifier (may take time)...")
        clf = pipeline("text-classification", model=name, return_all_scores=True,
                       device=0 if torch and torch.cuda.is_available() else -1)
        print("[TXT] Text classifier loaded.")
        return clf
    except Exception as e:
        print("[TXT] Failed to load text classifier:", e)
        return None

# ---------------- semantic helpers ----------------
def compute_embedding_sentence(text: str, tokenizer, model):
    if not text or tokenizer is None or model is None:
        return None
    try:
        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
        if next(model.parameters()).is_cuda:
            inputs = {k: v.cuda() for k, v in inputs.items()}
        with torch.no_grad():
            out = model(**inputs, return_dict=True)
        token_embeddings = out.last_hidden_state
        attention_mask = inputs['attention_mask'].unsqueeze(-1)
        masked = token_embeddings * attention_mask
        summed = masked.sum(1)
        denom = attention_mask.sum(1).clamp(min=1e-9)
        emb = (summed / denom).squeeze().cpu().numpy()
        return emb
    except Exception as e:
        print("[EMB] embedding error:", e)
        return None

def cosine_sim(a, b):
    if a is None or b is None:
        return None
    num = float(np.dot(a, b))
    da = float(np.linalg.norm(a))
    db = float(np.linalg.norm(b))
    if da == 0 or db == 0:
        return None
    return num / (da * db)

def fallback_token_overlap(a_text, b_text):
    if not a_text or not b_text:
        return 0.0
    a_set = set(w.strip(".,!?\"'").lower() for w in a_text.split() if w.strip())
    b_set = set(w.strip(".,!?\"'").lower() for w in b_text.split() if w.strip())
    if not a_set or not b_set:
        return 0.0
    inter = len(a_set & b_set)
    avg_len = (len(a_set) + len(b_set)) / 2.0
    return inter / avg_len if avg_len > 0 else 0.0

def semantic_similarity(a_text, b_text, tokenizer=None, emb_model=None):
    emb_a = compute_embedding_sentence(a_text, tokenizer, emb_model) if tokenizer and emb_model else None
    emb_b = compute_embedding_sentence(b_text, tokenizer, emb_model) if tokenizer and emb_model else None
    sim = cosine_sim(emb_a, emb_b)
    if sim is not None:
        return float(max(0.0, min(1.0, (sim + 1.0) / 2.0)))
    return float(max(0.0, min(1.0, fallback_token_overlap(a_text, b_text))))

# ---------------- text emotion helpers (HF fallback) ----------------
LABEL_MAP = {
    "neutral": "Neutral", "joy": "Happy", "happy": "Happy", "happiness": "Happy",
    "sad": "Sad", "sadness": "Sad", "anger": "Angry", "angry": "Angry",
    "fear": "Fear", "disgust": "Disgust", "surprise": "Neutral", "love": "Happy"
}

def hf_text_to_probs(text: str, clf_pipeline):
    zeros = {c: 0.0 for c in CATEGORIES}
    if not text or clf_pipeline is None:
        return zeros
    try:
        out = clf_pipeline(text)
    except Exception as e:
        print("[TXT] classifier error:", e)
        return zeros
    scores = out[0] if isinstance(out, list) and isinstance(out[0], list) else out
    mapping = {}
    total = 0.0
    for it in scores:
        lbl = (it.get("label") or "").lower()
        sc = float(it.get("score") or 0.0)
        tgt = LABEL_MAP.get(lbl)
        if tgt:
            mapping[tgt] = mapping.get(tgt, 0.0) + sc
            total += sc
    if total <= 0:
        return zeros
    for k in mapping:
        mapping[k] = mapping[k] / total
    for c in CATEGORIES:
        if c not in mapping:
            mapping[c] = 0.0
    return mapping

def fallback_text_probs(text: str):
    out = {c: 0.0 for c in CATEGORIES}
    if not text:
        out["Neutral"] = 1.0
        return out
    t = text.lower()
    if any(w in t for w in ["sorry", "apolog", "understand", "i'm sorry", "i am sorry"]):
        out["Sad"] = 0.6; out["Neutral"] = 0.4
    elif any(w in t for w in ["happy", "great", "good", "awesome"]):
        out["Happy"] = 1.0
    elif any(w in t for w in ["angry", "mad", "furious", "upset"]):
        out["Angry"] = 1.0
    else:
        out["Neutral"] = 1.0
    return out

def text_valence_from_probs_map(prob_map: Dict[str, float]):
    v = 0.0
    v += prob_map.get("Happy", 0.0) * 1.0
    v += prob_map.get("Neutral", 0.0) * 0.0
    neg = prob_map.get("Sad", 0.0) + prob_map.get("Angry", 0.0) + prob_map.get("Fear", 0.0) + prob_map.get("Disgust", 0.0)
    v += neg * -1.0
    return float(v)

# ---------------- lexical empathy & action heuristics (salvaged) ----------------
EMPATHY_PHRASES = ["i understand", "i completely understand", "i know how", "sorry to hear", "i'm sorry", "i'm so sorry", "i apologise", "i apologize", "we'll prioritize", "i'll escalate", "i'll create a ticket"]
ACTION_WORDS = ["ticket", "escalat", "callback", "call back", "schedule", "refund", "transfer", "open a ticket", "follow up"]

def detect_empathy_phrase(text: str):
    if not text:
        return False
    t = text.lower()
    return any(p in t for p in EMPATHY_PHRASES)

def detect_action_presence(text: str):
    if not text:
        return False
    t = text.lower()
    return any(w in t for w in ACTION_WORDS)

# ---------------- MDES subcomponents (BPObox mapping) ----------------
def compute_PU(agent_text, customer_context, tokenizer=None, emb_model=None):
    if not agent_text or not customer_context:
        return 0.0
    return semantic_similarity(agent_text, customer_context, tokenizer, emb_model)

def compute_CC(agent_text, customer_context, tokenizer=None, emb_model=None):
    if not agent_text or not customer_context:
        return 0.0
    sem = semantic_similarity(agent_text, customer_context, tokenizer, emb_model)
    if sem >= 0.80:
        return 0.5
    if sem >= 0.60:
        return 0.3
    a_set = set(w.strip(".,!?\"'").lower() for w in agent_text.split() if w.strip())
    c_set = set(w.strip(".,!?\"'").lower() for w in customer_context.split() if w.strip())
    inter = len(a_set & c_set)
    if inter >= 2:
        return 0.2
    if inter == 1:
        return 0.1
    return 0.0

def compute_EM(agent_text, customer_text, text_classify_fn):
    if not agent_text or not customer_text:
        return 0.0
    cust_probs = text_classify_fn(customer_text)
    ag_probs = text_classify_fn(agent_text)
    top_c = max(cust_probs.items(), key=lambda x: x[1])[0]
    top_a = max(ag_probs.items(), key=lambda x: x[1])[0]
    if top_c == top_a:
        return 0.95
    val_c = text_valence_from_probs_map(cust_probs)
    val_a = text_valence_from_probs_map(ag_probs)
    if val_c == 0 and val_a == 0:
        return 0.7
    if val_c * val_a > 0:
        return 0.7
    return 0.25

def compute_IR(agent_prosody, customer_prosody):
    if not agent_prosody or not customer_prosody:
        return 0.8
    ae = float(agent_prosody.get("energy", 0.0))
    ce = float(customer_prosody.get("energy", 0.0))
    if ce <= 0:
        return 0.9
    ratio = ae / max(ce, 1e-9)
    if 0.75 <= ratio <= 1.33:
        return 1.0
    return float(max(0.2, 1.0 - abs(math.log2(ratio)) / 4.0))

def compute_AA(agent_text, action_logs=None):
    if not agent_text:
        return 0.0
    t = agent_text.lower()
    if action_logs and isinstance(action_logs, dict) and (action_logs.get("ticket_created") or action_logs.get("escalated")):
        return 1.0
    if any(w in t for w in ACTION_WORDS):
        if any(tok.isdigit() for tok in t.split()) or "at " in t or "tomorrow" in t or "today" in t:
            return 1.0
        return 0.7
    return 0.0

def compute_FR(agent_text):
    if not agent_text:
        return 0.0
    t = agent_text.lower()
    if "call back" in t or "callback" in t or "at " in t or "tomorrow" in t:
        return 1.0
    if "i'll" in t or "i will" in t or "we will" in t:
        return 0.6
    return 0.0

def compute_ST(agent_text):
    if not agent_text:
        return 0.0
    candidates = ["reset","send","escalate","transfer","create","schedule","verify","confirm","refund"]
    steps = sum(1 for w in candidates if w in agent_text.lower())
    return float(min(1.0, steps / 3.0))

def compute_OT(customer_segment, customer_prosody, customer_valence):
    if customer_valence < -0.5:
        return 3.0
    if customer_prosody and customer_prosody.get("energy", 0.0) > 1e-4:
        return 3.0
    return 6.0

def compute_EF_for_segment(agent_seg_start, customer_segments, text_classify_fn, action_present=False):
    before_start = agent_seg_start - 8.0
    before_end = agent_seg_start
    after_start = agent_seg_start
    after_end = agent_seg_start + 12.0
    before_texts = [s['text'] for s in customer_segments if s['start'] >= before_start and s['end'] <= before_end]
    after_texts = [s['text'] for s in customer_segments if s['start'] > after_start and s['end'] <= after_end]
    before_text = " ".join(before_texts).strip()
    after_text = " ".join(after_texts).strip()
    def val_from_text(t):
        if not t:
            return 0.0
        pm = text_classify_fn(t)
        return text_valence_from_probs_map(pm)
    b = val_from_text(before_text)
    a = val_from_text(after_text)
    delta = a - b
    raw = 1.0 + (delta / 2.0) * 0.2
    raw = float(max(0.5, min(1.5, raw)))
    if action_present:
        ef = 0.6 * raw + 0.4 * 1.0
    else:
        ef = raw
    ef = float(max(0.8, min(1.2, ef)))
    return ef

# ---------------- MDES aggregator (per-utterance) ----------------
def compute_mdes_doc(agent_segments: List[Dict], customer_segments: List[Dict],
                     agent_prosody_map: Dict, customer_prosody_map: Dict,
                     tokenizer=None, emb_model=None, text_clf=None,
                     action_logs=None, empathy_threshold=DEFAULT_EMPATHY_THRESHOLD,
                     debug_permissive=False):
    def text_classify_fn(t):
        if text_clf:
            return hf_text_to_probs(t, text_clf)
        return fallback_text_probs(t)

    empathic_indices = []
    per_utt = []

    for i, a in enumerate(agent_segments):
        a_text = a.get("text", "") or ""
        a_start = a.get("start", 0.0)
        ctx = [s for s in customer_segments if s['end'] <= a_start and (a_start - s['end']) <= 12.0]
        cust_context = " ".join([s['text'] for s in ctx]) if ctx else (customer_segments[-1]['text'] if customer_segments else "")
        lex = 1.0 if detect_empathy_phrase(a_text) else 0.0
        sem = semantic_similarity(a_text, cust_context, tokenizer, emb_model) if cust_context else 0.0
        empathy_score = 0.4 * lex + 0.6 * sem
        threshold = DEBUG_PERMISSIVE_THRESHOLD if debug_permissive else empathy_threshold
        if empathy_score >= threshold or (debug_permissive and i == 0):
            empathic_indices.append(i)
        per_utt.append({"agent_index": i, "agent_start": a_start, "agent_text": a_text, "cust_context": cust_context, "empathy_score": empathy_score})

    if len(empathic_indices) == 0:
        return {"CE": 0.0, "AE": 0.0, "BE": 0.0, "TF": 1.0, "EF": 1.0, "MDES": 0.0, "n": 0, "per_utterance": []}

    CE_list = []
    AE_list = []
    BE_list = []
    TF_scores = []
    EF_scores = []

    for idx in empathic_indices:
        a = agent_segments[idx]
        a_text = a.get("text", "")
        a_start = a.get("start", 0.0)
        prior_customers = [s for s in customer_segments if s['end'] <= a_start]
        cust_seg = max(prior_customers, key=lambda x: x['end']) if prior_customers else None
        cust_text = cust_seg['text'] if cust_seg else ""
        # CE
        PU_i = compute_PU(a_text, cust_text, tokenizer, emb_model)
        CC_i = compute_CC(a_text, cust_text, tokenizer, emb_model)
        CE_i = PU_i * (1.0 + CC_i)
        CE_list.append(CE_i)
        # AE
        cust_probs = text_classify_fn(cust_text) if cust_text else {c:0.0 for c in CATEGORIES}
        ag_probs = text_classify_fn(a_text) if a_text else {c:0.0 for c in CATEGORIES}
        CS = text_valence_from_probs_map(cust_probs)
        AS = text_valence_from_probs_map(ag_probs)
        EM_i = compute_EM(a_text, cust_text, text_classify_fn)
        ag_key = round(a_start, 3)
        cust_key = round(cust_seg.get("start", 0.0), 3) if cust_seg else None
        ag_pros = agent_prosody_map.get(ag_key, {"energy":0.0})
        cust_pros = customer_prosody_map.get(cust_key, {"energy":0.0})
        IR_i = compute_IR(ag_pros, cust_pros)
        alignment_term = max(0.0, 1.0 - (abs(CS - AS) / 2.0))
        AE_i = EM_i * IR_i * alignment_term
        AE_list.append(AE_i)
        # BE
        AA_i = compute_AA(a_text, action_logs)
        FR_i = compute_FR(a_text)
        ST_i = compute_ST(a_text)
        BE_i = AA_i * FR_i * ST_i
        BE_list.append(BE_i)
        # TF
        if cust_seg:
            RT_i = a_start - cust_seg.get("end", a_start)
            cust_val = text_valence_from_probs_map(cust_probs)
            OT_i = compute_OT(cust_seg, cust_pros, cust_val)
            TF_score = max(0.0, 1.0 - abs(OT_i - RT_i) / OT_i) if OT_i > 0 else 0.0
        else:
            TF_score = 0.5
        TF_scores.append(TF_score)
        # EF
        action_present = detect_action_presence(a_text) or (action_logs and isinstance(action_logs, dict) and (action_logs.get("ticket_created") or action_logs.get("escalated")))
        EF_i = compute_EF_for_segment(a_start, customer_segments, text_classify_fn, action_present=action_present)
        EF_scores.append(EF_i)

    CE = float(np.mean(CE_list)) * 100.0
    AE = float(np.mean(AE_list)) * 100.0
    BE = float(np.mean(BE_list)) * 100.0
    TF = 1.0 + 0.2 * float(np.mean(TF_scores))
    EF = float(np.mean(EF_scores))
    inner = 0.4 * CE + 0.3 * AE + 0.3 * BE
    MDES = float(inner * TF * EF)

    per_utterance = []
    for j, idx in enumerate(empathic_indices):
        per_utterance.append({
            "agent_index": idx,
            "CE_i": float(CE_list[j]),
            "AE_i": float(AE_list[j]),
            "BE_i": float(BE_list[j]),
            "TF_i": float(TF_scores[j]),
            "EF_i": float(EF_scores[j])
        })

    return {"CE": CE, "AE": AE, "BE": BE, "TF": TF, "EF": EF, "MDES": MDES, "n": len(CE_list), "per_utterance": per_utterance}

# ---------------- ASR / transcription wrapper ----------------
def transcribe_with_whisper(model, path):
    if model is None:
        return {"text": "", "segments": []}
    try:
        res = model.transcribe(path)
        text = res.get("text", "") or ""
        segments = res.get("segments", []) or []
        segs = [{"start": float(s.get("start",0.0)), "end": float(s.get("end",0.0)), "text": s.get("text","").strip()} for s in segments]
        return {"text": text.strip(), "segments": segs}
    except Exception as e:
        print("[ASR] transcribe error:", e)
        return {"text":"", "segments": []}

# ---------------- top-level process function ----------------
def process_pair(agent_wav: str, customer_wav: str, asr_model,
                 tokenizer=None, emb_model=None, text_clf=None,
                 debug_permissive=False, action_logs=None, write_debug=True):
    # transcribe
    agent_asr = transcribe_with_whisper(asr_model, agent_wav)
    customer_asr = transcribe_with_whisper(asr_model, customer_wav)
    agent_segments = agent_asr.get("segments", [])
    customer_segments = customer_asr.get("segments", [])

    # compute per-segment prosody by slicing audio arrays
    agent_y = None
    cust_y = None
    try:
        agent_y, _ = librosa.load(agent_wav, sr=16000)
    except Exception:
        agent_y = None
    try:
        cust_y, _ = librosa.load(customer_wav, sr=16000)
    except Exception:
        cust_y = None

    agent_map = {}
    cust_map = {}
    sr = 16000
    for s in agent_segments:
        st = s.get("start", 0.0); en = s.get("end", st + 0.5)
        if agent_y is not None:
            b = int(max(0, st*sr)); e = int(min(len(agent_y), en*sr))
            seg = agent_y[b:e] if e>b else agent_y[b:b+1]
            pros = {"energy": float(np.mean(seg**2)) if seg.size else 0.0}
            agent_map[round(st,3)] = pros
        else:
            agent_map[round(st,3)] = {"energy": 0.0}
    for s in customer_segments:
        st = s.get("start", 0.0); en = s.get("end", st + 0.5)
        if cust_y is not None:
            b = int(max(0, st*sr)); e = int(min(len(cust_y), en*sr))
            seg = cust_y[b:e] if e>b else cust_y[b:b+1]
            pros = {"energy": float(np.mean(seg**2)) if seg.size else 0.0}
            cust_map[round(st,3)] = pros
        else:
            cust_map[round(st,3)] = {"energy": 0.0}

    # debug dump
    if write_debug:
        dbg = {
            "agent_asr": agent_asr,
            "customer_asr": customer_asr,
            "agent_segments": agent_segments,
            "customer_segments": customer_segments,
            "env": {
                "whisper_model": WHISPER_MODEL,
                "audio_model_loaded": audio_emotion_model is not None,
                "python_version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
            }
        }
        try:
            with open(DEBUG_DUMP_PATH, "w", encoding="utf-8") as fh:
                json.dump(dbg, fh, indent=2, ensure_ascii=False)
            print("[DEBUG] wrote", DEBUG_DUMP_PATH)
        except Exception as e:
            print("[DEBUG] failed to write debug dump:", e)

    # fused emotions for UI
    agent_audio_probs = runner_audio_emotion(agent_wav)
    customer_audio_probs = runner_audio_emotion(customer_wav)
    agent_text_probs = hf_text_to_probs(agent_asr.get("text",""), text_clf) if text_clf else fallback_text_probs(agent_asr.get("text",""))
    customer_text_probs = hf_text_to_probs(customer_asr.get("text",""), text_clf) if text_clf else fallback_text_probs(customer_asr.get("text",""))
    def fuse(audio_p, text_p, a_w=0.6, t_w=0.4):
        fused = {}
        for c in CATEGORIES:
            fused[c] = a_w * audio_p.get(c,0.0) + t_w * text_p.get(c,0.0)
        s = sum(fused.values())
        if s > 0:
            for k in fused:
                fused[k] = round(float(fused[k] / s), 3)
        return fused
    agent_fused = fuse(agent_audio_probs, agent_text_probs)
    customer_fused = fuse(customer_audio_probs, customer_text_probs)

    # compute MDES
    mdes = compute_mdes_doc(agent_segments, customer_segments, agent_map, cust_map,
                             tokenizer=tokenizer, emb_model=emb_model, text_clf=text_clf,
                             action_logs=action_logs, empathy_threshold=DEFAULT_EMPATHY_THRESHOLD,
                             debug_permissive=debug_permissive)

    out = {
        "mdes_summary": {
            "MDES": round(mdes["MDES"], 3),
            "CE": round(mdes["CE"], 3),
            "AE": round(mdes["AE"], 3),
            "BE": round(mdes["BE"], 3),
            "TF": round(mdes["TF"], 3),
            "EF": round(mdes["EF"], 3),
            "n_empathic_segments": mdes["n"]
        },
        "per_segment": mdes["per_utterance"],
        "agent_transcript": agent_asr.get("text",""),
        "customer_transcript": customer_asr.get("text",""),
        "agent_fused_emotion": agent_fused,
        "customer_fused_emotion": customer_fused
    }
    return out

# ---------------- CLI + Gradio entrypoints ----------------
def main_cli(args):
    asr = load_whisper_model()
    tok, emb = try_load_embedding()
    txt_clf = try_load_text_classifier()
    if not asr:
        print("[ASR] No ASR model loaded. Exiting.")
        sys.exit(1)
    res = process_pair(args.agent, args.customer, asr, tokenizer=tok, emb_model=emb, text_clf=txt_clf,
                       debug_permissive=args.debug, action_logs=None, write_debug=True)
    print(json.dumps(res, indent=2))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Full MDES app (clean reimplementation).")
    parser.add_argument("--agent", type=str, help="Agent WAV path (mono 16k recommended)")
    parser.add_argument("--customer", type=str, help="Customer WAV path (mono 16k recommended)")
    parser.add_argument("--debug", action="store_true", help="Permissive empathy threshold for debugging")
    parser.add_argument("--no-gradio", action="store_true", help="Do not launch Gradio UI")
    args = parser.parse_args()

    if args.agent and args.customer:
        main_cli(args)
        sys.exit(0)

    # else try to launch Gradio UI (if available and not disabled)
    if not args.no_gradio and USE_GRADIO:
        try:
            import gradio as gr
            # load models
            asr = load_whisper_model()
            tok, emb = try_load_embedding()
            txt_clf = try_load_text_classifier()
            def gr_runner(agent_audio, customer_audio, debug_mode=False):
                a_path = agent_audio[0] if isinstance(agent_audio, tuple) else agent_audio
                c_path = customer_audio[0] if isinstance(customer_audio, tuple) else customer_audio
                res = process_pair(a_path, c_path, asr, tokenizer=tok, emb_model=emb, text_clf=txt_clf,
                                   debug_permissive=debug_mode, action_logs=None, write_debug=True)
                # Gradio outputs: agent fused, agent audio-only, customer fused, customer audio-only, agent transcript, customer transcript, JSON
                return (res["agent_fused_emotion"], res["agent_fused_emotion"], res["customer_fused_emotion"], res["customer_fused_emotion"], res["agent_transcript"], res["customer_transcript"], json.dumps(res["mdes_summary"], indent=2))
            with gr.Blocks() as demo:
                gr.Markdown("## MDES (BPObox) — Clean Reimplementation")
                a_in = gr.Audio(label="Agent audio (wav)", type="filepath")
                c_in = gr.Audio(label="Customer audio (wav)", type="filepath")
                debug_checkbox = gr.Checkbox(label="Debug (permissive empathy)", value=False)
                out1 = gr.Label(label="Agent Combined Emotion")
                out2 = gr.Label(label="Agent Audio-only Emotion")
                out3 = gr.Label(label="Customer Combined Emotion")
                out4 = gr.Label(label="Customer Audio-only Emotion")
                t1 = gr.Textbox(label="Agent Transcript")
                t2 = gr.Textbox(label="Customer Transcript")
                jbox = gr.Textbox(label="MDES Summary JSON")
                run_btn = gr.Button("Run")
                run_btn.click(fn=gr_runner, inputs=[a_in, c_in, debug_checkbox], outputs=[out1, out2, out3, out4, t1, t2, jbox])
                demo.launch()
        except Exception as e:
            print("[GRADIO] Failed to launch Gradio UI:", e)
            print("You can still use CLI mode: python app_full_reimpl.py --agent agent.wav --customer customer.wav")
    else:
        print("Provide --agent and --customer to run in CLI mode, or remove --no-gradio to launch the UI.")