import contextlib
import io
import logging
from functools import lru_cache

import numpy as np
import torch

try:
    with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
        from pyctcdecode import build_ctcdecoder
except ImportError:  # pragma: no cover
    build_ctcdecoder = None

logging.getLogger("pyctcdecode").setLevel(logging.ERROR)


_FOLD_PREDICTED = {
    "ɹ̩": "ɝ",
    "l̩": "l",
    "m̩": "m",
    "n̩": "n",
    "ɾ": "t",
    "ʔ": "t",
}

_VOWELS = {
    "i", "u", "ɪ", "ʊ", "eɪ", "oʊ", "ɔɪ", "aɪ", "aʊ", "æ", "ɑ", "ɔ", "ə", "ɛ", "ɝ",
}

_PROXY_BASE = 0xE000


@lru_cache(maxsize=16)
def _build_proxy_decoder(vocab_size: int, blank_id: int):
    if build_ctcdecoder is None:
        raise RuntimeError("Beam decoder requested but pyctcdecode is not installed.")
    labels = []
    for token_id in range(vocab_size):
        if token_id == blank_id:
            labels.append("")
        else:
            labels.append(chr(_PROXY_BASE + token_id))
    return build_ctcdecoder(labels)


def ctc_prefix_beam_decode(frame_probs: torch.Tensor, blank_id: int, beam_size: int = 8) -> list[int]:
    decoder = _build_proxy_decoder(frame_probs.size(-1), blank_id)
    text = decoder.decode(frame_probs.detach().cpu().numpy().astype(np.float32), beam_width=beam_size)
    return [ord(ch) - _PROXY_BASE for ch in text]


def fold_phone_for_display(phone: str) -> str:
    return _FOLD_PREDICTED.get(phone, phone)


def normalize_decoded_display(phones: list[str]) -> list[str]:
    out: list[str] = []
    i = 0
    while i < len(phones):
        ph = phones[i]
        if ph == "n̩":
            if out and out[-1] in _VOWELS:
                out.append("n")
            else:
                out.extend(("ə", "n"))
            i += 1
            continue
        if ph == "m̩":
            out.append("m")
            i += 1
            continue
        if ph == "l̩":
            if out and out[-1] in _VOWELS:
                out.append("l")
            else:
                out.extend(("ə", "l"))
            i += 1
            continue
        if i + 1 < len(phones) and phones[i + 1] == "ɹ̩":
            if ph in ("ɑ", "ə"):
                out.extend((ph, "r"))
                i += 2
                continue
            if ph == "r":
                out.append("r")
                i += 2
                continue
        out.append(ph)
        i += 1

    deduped: list[str] = []
    for ph in out:
        if not deduped or deduped[-1] != ph:
            deduped.append(ph)
    return deduped


def _transform_phone_confs(
    phones: list[str],
    confs: list[float],
    apply_display_norm: bool,
) -> tuple[list[str], list[float]]:
    if not apply_display_norm:
        return list(phones), list(confs)

    out_phones: list[str] = []
    out_confs: list[float] = []
    i = 0
    while i < len(phones):
        ph = phones[i]
        conf = confs[i]
        if ph == "n̩":
            if out_phones and out_phones[-1] in _VOWELS:
                out_phones.append("n")
                out_confs.append(conf)
            else:
                out_phones.extend(("ə", "n"))
                out_confs.extend((conf, conf))
            i += 1
            continue
        if ph == "m̩":
            out_phones.append("m")
            out_confs.append(conf)
            i += 1
            continue
        if ph == "l̩":
            if out_phones and out_phones[-1] in _VOWELS:
                out_phones.append("l")
                out_confs.append(conf)
            else:
                out_phones.extend(("ə", "l"))
                out_confs.extend((conf, conf))
            i += 1
            continue
        if i + 1 < len(phones) and phones[i + 1] == "ɹ̩":
            next_conf = confs[i + 1]
            if ph in ("ɑ", "ə"):
                out_phones.extend((ph, "r"))
                out_confs.extend((conf, next_conf))
                i += 2
                continue
            if ph == "r":
                out_phones.append("r")
                out_confs.append((conf + next_conf) / 2.0)
                i += 2
                continue
        out_phones.append(ph)
        out_confs.append(conf)
        i += 1
    return out_phones, out_confs


def _collapse_phone_confs(phones: list[str], confs: list[float]) -> tuple[list[str], list[float]]:
    collapsed: list[str] = []
    collapsed_confs: list[float] = []
    prev_phone = None
    run_conf: list[float] = []

    for phone, conf in zip(phones, confs):
        if phone != prev_phone:
            if prev_phone is not None:
                collapsed.append(prev_phone)
                collapsed_confs.append(sum(run_conf) / len(run_conf))
            prev_phone = phone
            run_conf = [conf]
        else:
            run_conf.append(conf)

    if prev_phone is not None:
        collapsed.append(prev_phone)
        collapsed_confs.append(sum(run_conf) / len(run_conf))
    return collapsed, collapsed_confs


def _align_token_confidences(
    src_phones: list[str],
    src_confs: list[float],
    dst_phones: list[str],
    default_conf: float,
) -> list[float]:
    n = len(src_phones)
    m = len(dst_phones)
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    back = [[None] * (m + 1) for _ in range(n + 1)]

    for i in range(1, n + 1):
        dp[i][0] = i
        back[i][0] = "del"
    for j in range(1, m + 1):
        dp[0][j] = j
        back[0][j] = "ins"

    for i in range(1, n + 1):
        for j in range(1, m + 1):
            sub_cost = 0 if src_phones[i - 1] == dst_phones[j - 1] else 1
            candidates = [
                (dp[i - 1][j - 1] + sub_cost, "sub"),
                (dp[i - 1][j] + 1, "del"),
                (dp[i][j - 1] + 1, "ins"),
            ]
            dp[i][j], back[i][j] = min(candidates, key=lambda x: x[0])

    aligned_confs = [default_conf] * m
    i, j = n, m
    while i > 0 or j > 0:
        op = back[i][j]
        if op == "sub":
            aligned_confs[j - 1] = src_confs[i - 1]
            i -= 1
            j -= 1
        elif op == "del":
            i -= 1
        elif op == "ins":
            j -= 1
        else:
            break
    return aligned_confs


def decode_ctc_predictions(
    frame_probs: torch.Tensor,
    blank_id: int,
    id2phone: dict,
    apply_display_norm: bool = False,
    frame_confidence: torch.Tensor | None = None,
) -> tuple[list[str], list[float]]:
    pred_ids = frame_probs.argmax(dim=-1).tolist()
    raw_collapsed: list[str] = []
    raw_confs: list[float] = []

    start = 0
    while start < len(pred_ids):
        pid = pred_ids[start]
        end = start + 1
        while end < len(pred_ids) and pred_ids[end] == pid:
            end += 1
        if pid != blank_id:
            phone = fold_phone_for_display(id2phone[str(pid)])
            if frame_confidence is None:
                run_conf = float(frame_probs[start:end, pid].mean().item())
            else:
                run_conf = float(frame_confidence[start:end].mean().item())
            raw_collapsed.append(phone)
            raw_confs.append(run_conf)
        start = end

    disp_phones, disp_confs = _transform_phone_confs(raw_collapsed, raw_confs, apply_display_norm=apply_display_norm)
    disp_phones = [fold_phone_for_display(ph) for ph in disp_phones]
    return _collapse_phone_confs(disp_phones, disp_confs)


def decode_ctc_predictions_beam(
    frame_probs: torch.Tensor,
    blank_id: int,
    id2phone: dict,
    beam_size: int = 8,
    apply_cambridge_norm: bool = False,
    frame_confidence: torch.Tensor | None = None,
) -> tuple[list[str], list[float]]:
    pred_ids = ctc_prefix_beam_decode(frame_probs, blank_id, beam_size=beam_size)
    phones = [id2phone[str(pid)] for pid in pred_ids]
    if apply_cambridge_norm:
        phones = normalize_decoded_display(phones)
    phones = [fold_phone_for_display(ph) for ph in phones]

    collapsed: list[str] = []
    for ph in phones:
        if not collapsed or collapsed[-1] != ph:
            collapsed.append(ph)

    greedy_phones, greedy_confs = decode_ctc_predictions(
        frame_probs,
        blank_id,
        id2phone,
        apply_display_norm=apply_cambridge_norm,
        frame_confidence=frame_confidence,
    )
    if collapsed:
        if frame_confidence is None:
            default_conf = float(frame_probs.max(dim=-1).values.mean().item())
        else:
            default_conf = float(frame_confidence.mean().item())
    else:
        default_conf = 0.0
    confs = _align_token_confidences(greedy_phones, greedy_confs, collapsed, default_conf)
    return collapsed, confs