import torch
import torchaudio
from torchaudio.transforms import Resample
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC

# -----------------------------
# CONFIG
# -----------------------------
MODEL_NAME = "vitouphy/wav2vec2-xls-r-300m-english"   # YOUR MODEL
TARGET_SR = 16000                                     # wav2vec2 expected SR
TARGET_RMS = 0.05                                     # loudness normalization

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# -----------------------------
# LOAD PROCESSOR + MODEL ONCE
# -----------------------------
processor = Wav2Vec2Processor.from_pretrained(MODEL_NAME)
model = Wav2Vec2ForCTC.from_pretrained(MODEL_NAME).to(device)
model.eval()

print("✅ Loaded ASR model:"+MODEL_NAME)

def preprocess_audio(path: str,
                     target_sr: int = TARGET_SR,
                     use_vad: bool = True) -> tuple[torch.Tensor, int]:
    """
    Load audio file, convert to mono 16 kHz, optional VAD trimming,
    normalize loudness, and return (waveform, sr).

    waveform: [1, time]
    """
    # 1) load audio
    wav, sr = torchaudio.load(path)

    # 2) mix stereo to mono
    if wav.size(0) > 1:
        wav = wav.mean(dim=0, keepdim=True)

    # 3) resample if needed
    if sr != target_sr:
        resampler = Resample(orig_freq=sr, new_freq=target_sr)
        wav = resampler(wav)
        sr = target_sr

    # 4) apply VAD safely
    if use_vad:
        try:
            trimmed = torchaudio.functional.vad(
                wav.squeeze(0),
                sample_rate=sr,
            )
            if trimmed.numel() > 0:
                wav = trimmed.unsqueeze(0)
        except Exception as e:
            print("⚠️ VAD failed, using untrimmed audio:", e)

    # 5) normalize volume (RMS normalization)
    rms = torch.sqrt(torch.mean(wav ** 2))
    if rms > 0:
        wav = wav * (TARGET_RMS / (rms + 1e-9))

    return wav, sr

@torch.no_grad()
def transcribe_audio(path: str,
                     use_vad: bool = True) -> dict:
    """
    High-level transcription function using your model.
    Returns:
      {
        "transcript": "...",
        "sample_rate": 16000,
        "num_samples": int
      }
    """
    # 1) preprocess
    wav, sr = preprocess_audio(path, use_vad=use_vad)

    # 2) processor expects a 1D numpy array
    inputs = processor(
        wav.squeeze(0).cpu().numpy(),
        sampling_rate=sr,
        return_tensors="pt",
        padding=True,
    ).to(device)

    # 3) forward pass
    logits = model(inputs.input_values).logits
    predicted_ids = torch.argmax(logits, dim=-1)

    # 4) decode
    text = processor.batch_decode(predicted_ids)[0]

    return {
        "transcript": text,
        "sample_rate": sr,
        "num_samples": wav.size(1)
    }

LEXICON = {
    # Stops
    "piano":  ["P", "IY", "AE", "N", "OW"],   # P initial
    "cup":    ["K", "AH", "P"],              # P final
    "ball":   ["B", "AO", "L"],              # B initial
    "crib":   ["K", "R", "IH", "B"],         # B final
    "tiger":  ["T", "AY", "G", "ER"],        # T initial
    "cat":    ["K", "AE", "T"],              # T final, K initial
    "door":   ["D", "AO", "R"],              # D initial
    "slide":  ["S", "L", "AY", "D"],         # D final
    "goat":   ["G", "OW", "T"],              # G initial
    "bag":    ["B", "AE", "G"],              # G final

    # Nasals
    "monkey": ["M", "AH", "NG", "K", "IY"],  # M initial, NG medial
    "ham":    ["HH", "AE", "M"],             # M final
    "nose":   ["N", "OW", "Z"],              # N initial, Z final
    "sunny":  ["S", "AH", "N", "IY"],        # N medial
    "ring":   ["R", "IH", "NG"],             # NG final

    # Fricatives
    "fish":   ["F", "IH", "SH"],             # F initial, SH final
    "leaf":   ["L", "IY", "F"],              # F final
    "violin": ["V", "AY", "AH", "L", "IH", "N"], # V initial
    "five":   ["F", "AY", "V"],              # V final
    "sun":    ["S", "AH", "N"],              # S initial
    "bus":    ["B", "AH", "S"],              # S final
    "zebra":  ["Z", "IY", "B", "R", "AH"],   # Z initial
    "treasure": ["T", "R", "EH", "ZH", "ER"],# ZH medial
    "thumb":  ["TH", "AH", "M"],             # voiceless TH initial
    "tooth":  ["T", "UW", "TH"],             # voiceless TH final
    "this":   ["DH", "IH", "S"],             # voiced TH initial
    "feather":["F", "EH", "DH", "ER"],       # voiced TH medial

    # Affricates
    "chair":  ["CH", "EH", "R"],             # CH initial
    "peach":  ["P", "IY", "CH"], 
    "duck":   ["D", "AH", "K"],                          # CH final
    "jam":    ["JH", "AE", "M"],             # J initial
    "cage":   ["K", "EY", "JH"],             # J final

    # Glides / liquids / H
    "window": ["W", "IH", "N", "D", "OW"],   # W initial
    "cow":    ["K", "AW"],                   # W not needed final; still useful
    "yellow": ["Y", "EH", "L", "OW"],        # Y initial
    "lion":   ["L", "AY", "AH", "N"],        # L initial
    "bell":   ["B", "EH", "L"],              # L final
    "rabbit": ["R", "AE", "B", "IH", "T"],   # R initial
    "car":    ["K", "AA", "R"],              # R final
    "hat":    ["HH", "AE", "T"],             # H initial
}

SCREENING_ITEMS = {
    # P
    "p_initial": { "word": "piano", "phonemeKey": "piano", "targetPhoneme": "P",  "targetIndex": 0, "position": "initial", "masteryAge": 2 },
    "p_final":   { "word": "cup",   "phonemeKey": "cup",   "targetPhoneme": "P",  "targetIndex": 2, "position": "final",   "masteryAge": 2 },

    # B
    "b_initial": { "word": "ball",  "phonemeKey": "ball",  "targetPhoneme": "B",  "targetIndex": 0, "position": "initial", "masteryAge": 2 },
    "b_final":   { "word": "crib",  "phonemeKey": "crib",  "targetPhoneme": "B",  "targetIndex": 3, "position": "final",   "masteryAge": 2 },

    # M
    "m_initial": { "word": "monkey","phonemeKey": "monkey","targetPhoneme": "M",  "targetIndex": 0, "position": "initial", "masteryAge": 2 },
    "m_final":   { "word": "ham",   "phonemeKey": "ham",   "targetPhoneme": "M",  "targetIndex": 2, "position": "final",   "masteryAge": 2 },

    # N
    "n_initial": { "word": "nose",  "phonemeKey": "nose",  "targetPhoneme": "N",  "targetIndex": 0, "position": "initial", "masteryAge": 3 },
    "n_medial":  { "word": "sunny", "phonemeKey": "sunny", "targetPhoneme": "N",  "targetIndex": 2, "position": "medial",  "masteryAge": 3 },

    # NG
    "ng_medial": { "word": "monkey","phonemeKey": "monkey","targetPhoneme": "NG", "targetIndex": 2, "position": "medial",  "masteryAge": 3 },
    "ng_final":  { "word": "ring",  "phonemeKey": "ring",  "targetPhoneme": "NG", "targetIndex": 2, "position": "final",   "masteryAge": 3 },

    # F
    "f_initial": { "word": "fish",  "phonemeKey": "fish",  "targetPhoneme": "F",  "targetIndex": 0, "position": "initial", "masteryAge": 4 },
    "f_final":   { "word": "leaf",  "phonemeKey": "leaf",  "targetPhoneme": "F",  "targetIndex": 2, "position": "final",   "masteryAge": 4 },

    # V
    "v_initial": { "word": "violin","phonemeKey": "violin","targetPhoneme": "V",  "targetIndex": 0, "position": "initial", "masteryAge": 4 },
    "v_final":   { "word": "five",  "phonemeKey": "five",  "targetPhoneme": "V",  "targetIndex": 2, "position": "final",   "masteryAge": 4 },

    # S
    "s_initial": { "word": "sun",   "phonemeKey": "sun",   "targetPhoneme": "S",  "targetIndex": 0, "position": "initial", "masteryAge": 4 },
    "s_final":   { "word": "bus",   "phonemeKey": "bus",   "targetPhoneme": "S",  "targetIndex": 2, "position": "final",   "masteryAge": 4 },

    # Z
    "z_initial": { "word": "zebra", "phonemeKey": "zebra", "targetPhoneme": "Z",  "targetIndex": 0, "position": "initial", "masteryAge": 5 },
    "z_final":   { "word": "nose",  "phonemeKey": "nose",  "targetPhoneme": "Z",  "targetIndex": 2, "position": "final",   "masteryAge": 5 },

    # SH
    "sh_final":  { "word": "fish",  "phonemeKey": "fish",  "targetPhoneme": "SH", "targetIndex": 2, "position": "final",   "masteryAge": 4 },

    # ZH
    "zh_medial": { "word": "treasure","phonemeKey": "treasure","targetPhoneme": "ZH","targetIndex": 3,"position": "medial","masteryAge": 6 },

    # TH (voiceless) θ
    "th_initial":{ "word": "thumb", "phonemeKey": "thumb", "targetPhoneme": "TH", "targetIndex": 0, "position": "initial", "masteryAge": 6 },
    "th_final":  { "word": "tooth", "phonemeKey": "tooth", "targetPhoneme": "TH", "targetIndex": 2, "position": "final",   "masteryAge": 6 },

    # DH (voiced) ð
    "dh_initial":{ "word": "this",  "phonemeKey": "this",  "targetPhoneme": "DH", "targetIndex": 0, "position": "initial", "masteryAge": 7 },
    "dh_medial": { "word": "feather","phonemeKey": "feather","targetPhoneme": "DH","targetIndex": 2,"position": "medial", "masteryAge": 7 },

    # CH
    "ch_initial":{ "word": "chair", "phonemeKey": "chair", "targetPhoneme": "CH", "targetIndex": 0, "position": "initial", "masteryAge": 5 },
    "ch_final":  { "word": "peach", "phonemeKey": "peach", "targetPhoneme": "CH", "targetIndex": 2, "position": "final",   "masteryAge": 5 },

    # J
    "j_initial": { "word": "jam",   "phonemeKey": "jam",   "targetPhoneme": "JH", "targetIndex": 0, "position": "initial", "masteryAge": 5 },
    "j_final":   { "word": "cage",  "phonemeKey": "cage",  "targetPhoneme": "JH", "targetIndex": 2, "position": "final",   "masteryAge": 5 },

    # K
    "k_initial": { "word": "cat",   "phonemeKey": "cat",   "targetPhoneme": "K",  "targetIndex": 0, "position": "initial", "masteryAge": 3 },
    "k_final":   { "word": "duck",  "phonemeKey": "duck",  "targetPhoneme": "K",  "targetIndex": 3, "position": "final",   "masteryAge": 3 },

    # G
    "g_initial": { "word": "goat",  "phonemeKey": "goat",  "targetPhoneme": "G",  "targetIndex": 0, "position": "initial", "masteryAge": 3 },
    "g_final":   { "word": "bag",   "phonemeKey": "bag",   "targetPhoneme": "G",  "targetIndex": 2, "position": "final",   "masteryAge": 3 },

    # W
    "w_initial": { "word": "window","phonemeKey": "window","targetPhoneme": "W",  "targetIndex": 0, "position": "initial", "masteryAge": 3 },

    # Y
    "y_initial": { "word": "yellow","phonemeKey": "yellow","targetPhoneme": "Y",  "targetIndex": 0, "position": "initial", "masteryAge": 4 },

    # L
    "l_initial": { "word": "leaf",  "phonemeKey": "leaf",  "targetPhoneme": "L",  "targetIndex": 0, "position": "initial", "masteryAge": 4 },
    "l_final":   { "word": "bell",  "phonemeKey": "bell",  "targetPhoneme": "L",  "targetIndex": 2, "position": "final",   "masteryAge": 4 },

    # R
    "r_initial": { "word": "rabbit","phonemeKey": "rabbit","targetPhoneme": "R",  "targetIndex": 0, "position": "initial", "masteryAge": 6 },
    "r_final":   { "word": "car",   "phonemeKey": "car",   "targetPhoneme": "R",  "targetIndex": 2, "position": "final",   "masteryAge": 6 },

    # H
    "h_initial": { "word": "hat",   "phonemeKey": "hat",   "targetPhoneme": "HH", "targetIndex": 0, "position": "initial", "masteryAge": 3 },
}

#test these khaula
def get_screening_item(item_id: str) -> dict:
    """Fetch a screening item or raise a clear error."""
    try:
        return SCREENING_ITEMS[item_id]
    except KeyError:
        raise ValueError(f"Unknown screening item: {item_id}")


def get_canonical_phonemes(item: dict) -> list[str]:
    """
    Get the canonical phoneme sequence for this item,
    using the lexicon and phonemeKey.
    """
    key = item.get("phonemeKey", item["word"]).lower()
    if key not in LEXICON:
        raise ValueError(f"Lexicon has no entry for '{key}'")
    return LEXICON[key]


import re

# ----------------------------------------
# Simple rule-based G2P (Option C)
# ----------------------------------------

# Consonant and digraph mappings (ARPABET-ish)
_CONSONANT_MAP = {
    "b": ["B"],
    "c": ["K"],       # 'c' -> K (cat) – good enough for MVP
    "d": ["D"],
    "f": ["F"],
    "g": ["G"],
    "h": ["HH"],
    "j": ["JH"],
    "k": ["K"],
    "l": ["L"],
    "m": ["M"],
    "n": ["N"],
    "p": ["P"],
    "q": ["K"],       # 'qu' handled separately as K+W
    "r": ["R"],
    "s": ["S"],
    "t": ["T"],
    "v": ["V"],
    "w": ["W"],
    "x": ["K", "S"],  # 'x' ~ /ks/
    "y": ["Y"],
    "z": ["Z"],
}

# Multi-letter consonant clusters first (so they don't get split)
_DIGRAPHS = {
    "ch": ["CH"],
    "sh": ["SH"],
    "th": ["TH"],   # voiceless; DH for voiced you can add later if you want
    "ph": ["F"],
    "ng": ["NG"],
    "wh": ["W"],
}

# Very rough vowel mapping – we mainly need them as placeholders
# so that consonant indices line up with your LEXICON.
def _map_vowel(ch: str) -> list[str]:
    if ch in "ae":
        return ["AE"]
    if ch == "i":
        return ["IH"]
    if ch == "o":
        return ["AO"]
    if ch == "u":
        return ["UH"]
    if ch == "y":
        return ["IH"]   # sometimes vowel
    return ["AH"]       # fallback


def simple_g2p_word(word: str) -> list[str]:
    """
    Very simple grapheme→phoneme for English-like words.
    - Handles digraphs (ch, sh, th, ph, ng, wh)
    - Maps consonant letters → ARPABET-like phonemes
    - Gives approximate vowels (we only care about consonant slots)
    """
    w = re.sub(r"[^a-zA-Z]", "", word.lower())
    phones: list[str] = []
    i = 0
    while i < len(w):
        # 1) digraphs first
        if i + 1 < len(w):
            pair = w[i:i+2]
            if pair in _DIGRAPHS:
                phones.extend(_DIGRAPHS[pair])
                i += 2
                continue

        ch = w[i]

        # 2) consonant letters
        if ch in _CONSONANT_MAP:
            phones.extend(_CONSONANT_MAP[ch])

        # 3) vowels
        elif ch in "aeiouy":
            phones.extend(_map_vowel(ch))

        # 4) ignore anything else
        else:
            pass

        i += 1

    return phones


def run_g2p(word: str) -> list[str]:
    """
    Option C: rule-based G2P.
    Used when the word is NOT found in LEXICON.

    Later, you can replace the internals with a real NeMo G2P call.
    """
    phones = simple_g2p_word(word)
    # For debugging, you can uncomment:
    # print(f"[G2P] {word} -> {phones}")
    return phones


def get_or_g2p(word: str) -> list[str]:
    """
    Unified access:
      - if word exists in LEXICON, use canonical entry
      - else, fall back to rule-based G2P
    """
    w = word.lower()
    if w in LEXICON:
        return LEXICON[w]
    return run_g2p(w)

import re

def edit_distance(a: str, b: str) -> int:
    """
    Classic Levenshtein edit distance between two strings.
    Used to judge how close transcript tokens are to target_word.
    """
    a = a.lower()
    b = b.lower()
    dp = [[0] * (len(b) + 1) for _ in range(len(a) + 1)]
    for i in range(len(a) + 1):
        dp[i][0] = i
    for j in range(len(b) + 1):
        dp[0][j] = j
    for i in range(1, len(a) + 1):
        for j in range(1, len(b) + 1):
            cost = 0 if a[i - 1] == b[j - 1] else 1
            dp[i][j] = min(
                dp[i - 1][j] + 1,      # deletion
                dp[i][j - 1] + 1,      # insertion
                dp[i - 1][j - 1] + cost  # substitution
            )
    return dp[-1][-1]


def pick_candidate_word(transcript: str, target_word: str) -> str | None:
    """
    From the ASR transcript, pick the token that is most likely
    the child's attempt at `target_word`.

    Returns:
      - best token string, or
      - None if nothing is close enough (wrong word / no attempt).
    """
    # tokenize transcript into plain alphabetic tokens
    tokens = re.findall(r"[a-zA-Z]+", transcript.lower())
    if not tokens:
        return None

    target = target_word.lower()
    best = None
    best_dist = 999

    for tok in tokens:
        d = edit_distance(target, tok)
        if d < best_dist:
            best_dist = d
            best = tok

    # Threshold: if it's too far, treat as "no valid attempt"
    # Example: if distance > half the target length, it's probably not the same word.
    max_allowed = max(2, len(target) // 2)
    if best is None or best_dist > max_allowed:
        return None

    return best

def compare_target_phoneme(item: dict,
                           canonical_phonemes: list[str],
                           predicted_phonemes: list[str]) -> dict:
    """
    Compare the target phoneme at targetIndex between:
      - canonical_phonemes (from lexicon for target word)
      - predicted_phonemes (from candidate word via get_or_g2p)

    Returns a dict:
      {
        "status": "correct" | "substitution" | "omission" | "wrong_word" | "unknown",
        "target": str,
        "observed": str | None
      }
    """
    target_idx = item["targetIndex"]
    target_ph = item["targetPhoneme"]

    # If we don't even have a predicted sequence, we can't assess
    if not predicted_phonemes:
        return {
            "status": "unknown",
            "target": target_ph,
            "observed": None
        }

    # If predicted sequence is too short to have this position → omission
    if target_idx >= len(predicted_phonemes):
        return {
            "status": "omission",
            "target": target_ph,
            "observed": None
        }

    observed = predicted_phonemes[target_idx]

    if observed == target_ph:
        status = "correct"
    else:
        status = "substitution"

    return {
        "status": status,
        "target": target_ph,
        "observed": observed
    }

def score_item_from_result(result: dict) -> int:
    """
    Convert a phoneme comparison result into a numeric score.

    Simple MVP rule:
      correct       → 100
      substitution  → 50
      omission      → 0
      wrong_word    → 0
      unknown       → 0
    """
    status = result.get("status")

    if status == "correct":
        return 100
    if status == "substitution":
        return 50
    if status in ("omission", "wrong_word", "unknown"):
        return 0

    # fallback
    return 0

def assess_screening_item(audio_path: str, item_id: str) -> dict:
    """
    Core function for Module 2 (per item).

    Inputs:
      - audio_path: path to child's recording
      - item_id: key from SCREENING_ITEMS

    Output dict (ready to return via API later):
      {
        "itemId": ...,
        "targetWord": ...,
        "transcript": ...,
        "candidateWord": ...,
        "canonicalPhonemes": [...],
        "predictedPhonemes": [...],
        "phonemeResult": {...},
        "itemScore": 0/50/100
      }
    """
    # 1) get item + canonical phonemes
    item = get_screening_item(item_id)
    target_word = item["word"]
    canonical_phonemes = get_canonical_phonemes(item)

    # 2) ASR
    asr_out = transcribe_audio(audio_path)  # using JSON version
    transcript = asr_out["transcript"]

    # 3) pick candidate word from transcript
    candidate = pick_candidate_word(transcript, target_word)

    # If no good candidate, treat as wrong word
    if candidate is None:
        phoneme_result = {
            "status": "wrong_word",
            "target": item["targetPhoneme"],
            "observed": None
        }
        predicted_phonemes = []
        score = score_item_from_result(phoneme_result)
    else:
        # 4) get predicted phonemes from candidate (lexicon or G2P)
        predicted_phonemes = get_or_g2p(candidate)

        # 5) compare target phoneme
        phoneme_result = compare_target_phoneme(
            item,
            canonical_phonemes,
            predicted_phonemes
        )

        # 6) score
        score = score_item_from_result(phoneme_result)

    return {
        "itemId": item_id,
        "targetWord": target_word,
        "transcript": transcript,
        "candidateWord": candidate,
        "canonicalPhonemes": canonical_phonemes,
        "predictedPhonemes": predicted_phonemes,
        "phonemeResult": phoneme_result,
        "itemScore": score
    }