"""
Morse Code Demo — Gradio app
  Tab 1: Text → Morse audio
  Tab 2: Morse audio → transcription + English translation
"""

import math
import tempfile

import gradio as gr
import librosa
import numpy as np
import soundfile as sf
import torch
from transformers import WhisperForConditionalGeneration, WhisperProcessor

# ── Morse table ────────────────────────────────────────────────────────────────
MORSE_TABLE = {
    "A": ".-",
    "B": "-...",
    "C": "-.-.",
    "D": "-..",
    "E": ".",
    "F": "..-.",
    "G": "--.",
    "H": "....",
    "I": "..",
    "J": ".---",
    "K": "-.-",
    "L": ".-..",
    "M": "--",
    "N": "-.",
    "O": "---",
    "P": ".--.",
    "Q": "--.-",
    "R": ".-.",
    "S": "...",
    "T": "-",
    "U": "..-",
    "V": "...-",
    "W": ".--",
    "X": "-..-",
    "Y": "-.--",
    "Z": "--..",
    "0": "-----",
    "1": ".----",
    "2": "..---",
    "3": "...--",
    "4": "....-",
    "5": ".....",
    "6": "-....",
    "7": "--...",
    "8": "---..",
    "9": "----.",
    ".": ".-.-.-",
    ",": "--..--",
    "?": "..--..",
    "'": ".----.",
    "!": "-.-.--",
    "/": "-..-.",
    "(": "-.--.",
    ")": "-.--.-",
}

SAMPLE_RATE = 16_000
RISE_TIME_MS = 8
PAD_MS = 300
WPM_DEFAULT = 20
WPM_MAX = 60
MAX_DURATION = 30.0
SNR_DB = 15.0


def _ms(ms: float) -> int:
    return int(round(ms * SAMPLE_RATE / 1000))


def _tone(duration_ms: float, freq: float) -> np.ndarray:
    n = _ms(duration_ms)
    t = np.arange(n) / SAMPLE_RATE
    sig = np.sin(2 * np.pi * freq * t).astype(np.float32)
    rise = min(_ms(RISE_TIME_MS), n // 2)
    if rise > 0:
        w = np.ones(n, dtype=np.float32)
        h = np.hanning(2 * rise)[:rise]
        w[:rise], w[-rise:] = h, h[::-1]
        sig *= w
    return sig


def _silence(ms: float) -> np.ndarray:
    return np.zeros(_ms(ms), dtype=np.float32)


def _count_units(text: str) -> int:
    units = 0
    for w_idx, word in enumerate(text.upper().split()):
        if w_idx > 0:
            units += 7
        chars = [c for c in word if c in MORSE_TABLE]
        for c_idx, ch in enumerate(chars):
            if c_idx > 0:
                units += 3
            for e_idx, e in enumerate(MORSE_TABLE[ch]):
                if e_idx > 0:
                    units += 1
                units += 1 if e == "." else 3
    return units


def _choose_wpm(text: str) -> int:
    units = _count_units(text)
    if units == 0:
        return WPM_DEFAULT
    max_dit = (MAX_DURATION - 2 * PAD_MS / 1000) * 1000 / units
    return max(
        WPM_DEFAULT, min(math.ceil(1200 / max_dit) if max_dit > 0 else WPM_MAX, WPM_MAX)
    )


def text_to_audio(text: str, freq: float = 650.0, wpm: int | None = None) -> np.ndarray:
    if wpm is None:
        wpm = _choose_wpm(text)
    dit_ms = 1200 / wpm
    dash_ms = 3 * dit_ms
    dit = _tone(dit_ms, freq)
    dash = _tone(dash_ms, freq)

    chunks = [_silence(PAD_MS)]
    for w_idx, word in enumerate(text.upper().split()):
        if w_idx > 0:
            chunks.append(_silence(7 * dit_ms))
        chars = [c for c in word if c in MORSE_TABLE]
        for c_idx, ch in enumerate(chars):
            if c_idx > 0:
                chunks.append(_silence(3 * dit_ms))
            for e_idx, e in enumerate(MORSE_TABLE[ch]):
                if e_idx > 0:
                    chunks.append(_silence(dit_ms))
                chunks.append(dit if e == "." else dash)
    chunks.append(_silence(PAD_MS))

    audio = np.concatenate(chunks)
    peak = np.max(np.abs(audio))
    if peak > 0:
        audio /= peak

    # add mild noise
    p = np.mean(audio**2)
    if p > 0:
        std = math.sqrt(p / (10 ** (SNR_DB / 10)))
        audio = np.clip(
            audio + np.random.normal(0, std, audio.shape).astype(np.float32), -1, 1
        )

    return audio


# ── Whisper model (loaded once) ────────────────────────────────────────────────
# Path to fine-tuned HF checkpoint directory; set to a base model id (e.g.
# "openai/whisper-small") to fall back to the unmodified base model.
FINETUNED_CHECKPOINT: str = "michaellin/whisper-small-morse"

_processor: WhisperProcessor | None = None
_model: WhisperForConditionalGeneration | None = None
_device: torch.device | None = None


def _get_whisper() -> (
    tuple[WhisperProcessor, WhisperForConditionalGeneration, torch.device]
):
    global _processor, _model, _device
    if _model is None:
        _device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        _processor = WhisperProcessor.from_pretrained(FINETUNED_CHECKPOINT)
        _model = WhisperForConditionalGeneration.from_pretrained(FINETUNED_CHECKPOINT)
        _model.to(_device)
        _model.eval()
    return _processor, _model, _device


def _forced_decoder_ids(
    processor: WhisperProcessor, task: str
) -> list[tuple[int, int]]:
    """
    Build forced_decoder_ids for the given task ("transcribe" or "translate").
    Uses <|startoflm|> as the language token when available (fine-tuned model);
    otherwise falls back to English.
    """
    tokenizer = processor.tokenizer
    morse_id = tokenizer.convert_tokens_to_ids("<|startoflm|>")
    unk_id = (
        tokenizer.convert_tokens_to_ids(tokenizer.unk_token)
        if tokenizer.unk_token
        else None
    )
    if morse_id is None or morse_id == unk_id:
        # Stock base model — no morse token, fall back to English
        return processor.get_decoder_prompt_ids(language="en", task=task)
    task_id = tokenizer.convert_tokens_to_ids(f"<|{task}|>")
    notimestamps_id = tokenizer.convert_tokens_to_ids("<|notimestamps|>")
    # HF Whisper prefix: position 1 = language, 2 = task, 3 = (no)timestamps
    return [(1, morse_id), (2, task_id), (3, notimestamps_id)]


def _run_whisper(audio_path: str, task: str) -> str:
    processor, model, device = _get_whisper()

    audio, _ = librosa.load(audio_path, sr=SAMPLE_RATE, mono=True)
    inputs = processor(audio, sampling_rate=SAMPLE_RATE, return_tensors="pt")
    input_features = inputs.input_features.to(device)

    forced = _forced_decoder_ids(processor, task)
    with torch.no_grad():
        pred_ids = model.generate(input_features, forced_decoder_ids=forced)

    return processor.batch_decode(pred_ids, skip_special_tokens=True)[0].strip()


# ── Gradio handlers ────────────────────────────────────────────────────────────
def generate_morse(text: str, freq: float, wpm: int):
    if not text.strip():
        return None, "Please enter some text."
    audio = text_to_audio(text.strip(), freq=freq, wpm=wpm if wpm > 0 else None)
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        sf.write(f.name, audio, SAMPLE_RATE, subtype="PCM_16")
        path = f.name
    morse_str = " / ".join(
        " ".join(MORSE_TABLE.get(c, "") for c in w.upper() if c in MORSE_TABLE)
        for w in text.split()
    )
    return path, f"Morse: {morse_str}"


def decode_morse(audio_path: str):
    if audio_path is None:
        return "Upload a WAV file first.", ""
    transcription = _run_whisper(audio_path, task="transcribe")
    translation = _run_whisper(audio_path, task="translate")
    return transcription, translation


# ── UI ─────────────────────────────────────────────────────────────────────────
WEBSDR_MD = """
## Live Radio Demo (WebSDR)
1. Open **[websdr.ewi.utwente.nl:8901](http://websdr.ewi.utwente.nl:8901/m.html)** in your browser
2. Tune to **~7000 kHz**, switch mode to **CW**
3. Drag the trapezoid shape onto a signal line — you'll hear Morse code
4. Use your browser's audio capture or any screen-recorder to save a WAV clip
5. Upload the clip to the **Decode** tab above

> Once the model is fine-tuned with a `<|startoflm|>` language token,
> decoding will use that token instead of falling back to English.
"""

with gr.Blocks(title="Morse Code Demo") as demo:
    gr.Markdown("# Morse Code Demo")

    with gr.Tab("Generate: Text → Morse Audio"):
        txt_in = gr.Textbox(label="Input text", placeholder="HELLO WORLD")
        freq_sl = gr.Slider(400, 1000, value=650, step=50, label="Tone frequency (Hz)")
        wpm_sl = gr.Slider(0, 60, value=20, step=5, label="Speed (WPM)  —  0 = auto")
        gen_btn = gr.Button("Generate")
        audio_out = gr.Audio(label="Morse audio", type="filepath")
        morse_out = gr.Textbox(label="Morse code")

    with gr.Tab("Decode: Morse Audio → Text"):
        audio_in = gr.Audio(label="Morse audio (WAV)", type="filepath")
        dec_btn = gr.Button("Decode")
        transcription_out = gr.Textbox(label="Transcription")
        translation_out = gr.Textbox(label="Translation")
        dec_btn.click(decode_morse, [audio_in], [transcription_out, translation_out])

    # TODO: This should eventually run live as a browser extension. Disabling for pilot study demo.
    #with gr.Tab("Live Radio (WebSDR)"):
    #    gr.Markdown(WEBSDR_MD)

    # Wire up generate → decode tab. The Generate button populates audio_out,
    # which then propagates to audio_in via .change(). The user can still
    # manually upload or clear audio_in independently.
    gen_btn.click(generate_morse, [txt_in, freq_sl, wpm_sl], [audio_out, morse_out])
    audio_out.change(lambda x: x, [audio_out], [audio_in])

if __name__ == "__main__":
    demo.launch()