import os
import sys
import json
from typing import Any, Dict, List, Optional, Tuple

import gradio as gr
import torch
import numpy as np


from Eyettention.raw_text_inference import EyettentionRawTextInference

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

BSC_CHECKPOINT = os.environ.get(
    "EYETTENTION_BSC_CHECKPOINT",
    os.path.join(os.path.dirname(__file__), "checkpoints", "bsc.pt"),
)
CELER_CHECKPOINT = os.environ.get(
    "EYETTENTION_CELER_CHECKPOINT",
    os.path.join(os.path.dirname(__file__), "checkpoints", "celer.pt"),
)

_MODELS: Dict[str, EyettentionRawTextInference] = {}


def get_model(dataset: str) -> EyettentionRawTextInference:
    """Load and cache an Eyettention model for the given dataset."""
    if dataset not in _MODELS:
        if dataset == "BSC":
            ckpt = BSC_CHECKPOINT
        elif dataset == "celer":
            ckpt = CELER_CHECKPOINT
        else:
            raise ValueError(f"Unsupported dataset: {dataset}")

        if not os.path.exists(ckpt):
            raise FileNotFoundError(
                f"Checkpoint not found at {ckpt!r}. "
                f"Set EYETTENTION_{dataset.upper()}_CHECKPOINT to the correct path."
            )

        device = "cuda" if torch.cuda.is_available() else "cpu"
        _MODELS[dataset] = EyettentionRawTextInference(
            checkpoint_path=ckpt,
            dataset=dataset,
            device=device,
        )
    return _MODELS[dataset]


def words_for_text(text: str, dataset: str) -> List[str]:
    """
    Return the list of words that the model's fixation indices refer to.

    - For BSC (Chinese): the model tokenizes per character, so we return the
      list of characters in the input (excluding whitespace).
    - For CELER (English): the model tokenizes with a BERT wordpiece tokenizer
      on whitespace-split words, so we return the whitespace-split words.
    """
    if dataset == "BSC":

        return [ch for ch in text if not ch.isspace()]
    else:
        return text.split()


# ---------------------------------------------------------------------------
# Decoding helpers
# ---------------------------------------------------------------------------
def decode_scanpath(
    scanpath_indices: torch.Tensor,
    words: List[str],
    dataset: str,
) -> List[Dict[str, Any]]:
    """
    Convert a single predicted scanpath (indices into the word/char sequence)
    into a list of dicts describing each fixation.

    Convention in Eyettention's `scanpath_generation`:
      * index 0 corresponds to the CLS token (sentence start),
      * indices 1..len(words) correspond to words/characters,
      * index len(words)+1 corresponds to SEP (sentence end).

    We drop the leading CLS and any trailing SEP here.
    """
    if isinstance(scanpath_indices, torch.Tensor):
        scanpath_indices = scanpath_indices.detach().cpu().tolist()

    fixations: List[Dict[str, Any]] = []
    n_words = len(words)

    for step, idx in enumerate(scanpath_indices):
        idx = int(idx)

        if step == 0 and idx == 0:
            continue

        if idx == 0:
            word = "<CLS>"
        elif idx == n_words + 1:
            word = "<SEP>"
        elif 1 <= idx <= n_words:
            word = words[idx - 1]
        else:
            word = f"<OOR:{idx}>"

        fixations.append(
            {
                "step": len(fixations) + 1,
                "word": word,
                "word_index": idx,
            }
        )

    return fixations


def fixations_to_markdown(
    fixations: List[Dict[str, Any]],
    original_text: str,
    dataset: str,
) -> str:
    """Render a scanpath as a Markdown table."""
    lines = [
        f"**Input ({dataset}):** {original_text}",
        "",
        "| Step | Fixated unit | Index |",
        "|------|--------------|-------|",
    ]
    for f in fixations:
        lines.append(f"| {f['step']} | {f['word']} | {f['word_index']} |")
    return "\n".join(lines)


def predict(
    text: str,
    dataset: str,
    max_pred_len: int,
    use_previous_scanpath: bool,
    previous_scanpath: str,
    progress=gr.Progress(track_tqdm=True),
) -> Tuple[str, str, str]:
    """
    Run Eyettention on the input text and return:
      - a Markdown table of the predicted scanpath,
      - a space-separated string of the predicted word/char sequence,
      - the raw JSON of the scanpath indices and fixated units.
    """
    if text is None or text.strip() == "":
        raise gr.Error("Please provide some input text.")

    if dataset == "celer" and not any(c.isalpha() for c in text):
        raise gr.Error("CELER is the English model — please provide English text.")
    if dataset == "BSC" and not any("\u4e00" <= c <= "\u9fff" for c in text):

        pass

    if max_pred_len <= 0:
        raise gr.Error("max_pred_len must be a positive integer.")

    progress(0.05, desc="Loading Eyettention model...")
    model = get_model(dataset)

    prev: Optional[List[int]] = None
    if use_previous_scanpath and previous_scanpath.strip():
        try:
            prev = [int(x) for x in previous_scanpath.replace(",", " ").split()]
        except ValueError:
            raise gr.Error(
                "previous_scanpath must be a whitespace/comma-separated list of integers."
            )

    progress(0.25, desc="Running autoregressive scanpath generation...")
    with torch.no_grad():
        if dataset == "BSC":
            scanpath, _density = model.generate_from_chinese_text(
                text=text,
                max_pred_len=max_pred_len,
                previous_scanpath=prev,
            )
        else:
            scanpath, _density = model.generate_from_english_text(
                text=text,
                max_pred_len=max_pred_len,
                previous_scanpath=prev,
            )

    scanpath = scanpath[0]

    progress(0.9, desc="Formatting output...")
    words = words_for_text(text, dataset)
    fixations = decode_scanpath(scanpath, words, dataset)

    markdown = fixations_to_markdown(fixations, text, dataset)
    word_seq = " ".join(f["word"] for f in fixations)
    json_out = json.dumps(
        {
            "dataset": dataset,
            "input": text,
            "scanpath_indices": [int(i) for i in scanpath.detach().cpu().tolist()],
            "fixated_units": [f["word"] for f in fixations],
        },
        ensure_ascii=False,
        indent=2,
    )

    return markdown, word_seq, json_out


DESCRIPTION = """
# Eyettention — Scanpath Prediction

**Eyettention** predicts human-like **eye-movement scanpaths** from raw text.
Given a sentence, it autoregressively generates a sequence of fixation
locations (word or character indices) that approximate where a reader would
look, in order.

Two checkpoints are supported:

| Dataset | Language | Tokenization |
|---------|----------|--------------|
| **BSC** | Chinese | Character-level |
| **CELER** | English | Word-level (BERT wordpieces pooled) |

### How to use
1. Choose the **dataset / language**.
2. Paste your text.
3. Adjust `max_pred_len` if you want shorter or longer scanpaths.
4. Click **Run**.

### Optional: replay an observed prefix
Enable **Use previous scanpath** and provide a whitespace- or comma-separated
list of integer fixation indices. Those fixations will be replayed before the
model starts sampling new ones (useful for conditional generation / prefix
completion).
"""

EXAMPLES = [
    [
        "He said BankEast's offer appears to be \"attractive to the bank's shareholders.\"",
        "celer",
        20,
        False,
        "",
    ],
    ["中国选手在男子滑雪比赛中有望蝉联冠军", "BSC", 20, False, ""],
]


def build_demo() -> gr.Blocks:
    with gr.Blocks(
        title="Eyettention — Scanpath Prediction",
        theme=gr.themes.Soft(),
    ) as demo:
        gr.Markdown(DESCRIPTION)

        with gr.Row():
            with gr.Column(scale=3):
                text_in = gr.Textbox(
                    label="Input text",
                    placeholder="Enter a sentence...",
                    lines=5,
                )
                dataset_in = gr.Radio(
                    choices=["celer", "BSC"],
                    value="celer",
                    label="Dataset / language",
                    info="celer = English, BSC = Chinese",
                )
                max_len_in = gr.Slider(
                    minimum=1,
                    maximum=120,
                    value=30,
                    step=1,
                    label="max_pred_len",
                    info="Maximum number of fixations to generate.",
                )
                with gr.Accordion("Advanced: replay an observed prefix", open=False):
                    use_prev_in = gr.Checkbox(
                        value=False,
                        label="Use previous scanpath",
                    )
                    prev_in = gr.Textbox(
                        label="Previous scanpath (integer indices)",
                        placeholder="e.g. 0 1 3 5 4",
                        lines=2,
                    )
                with gr.Row():
                    run_btn = gr.Button("Run", variant="primary")
                    clear_btn = gr.Button("Clear")

            with gr.Column(scale=4):
                table_out = gr.Markdown(
                    label="Predicted scanpath",
                    value="_Results will appear here._",
                )
                seq_out = gr.Textbox(
                    label="Fixated units (sequence)",
                    lines=3,
                    interactive=False,
                )
                json_out = gr.Code(
                    label="Raw JSON output",
                    language="json",
                    value="{}",
                )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[text_in, dataset_in, max_len_in, use_prev_in, prev_in],
        )

        def _run(text, dataset, max_pred_len, use_prev, prev):
            md, seq, js = predict(
                text=text,
                dataset=dataset,
                max_pred_len=int(max_pred_len),
                use_previous_scanpath=use_prev,
                previous_scanpath=prev,
            )
            return md, seq, js

        run_btn.click(
            fn=_run,
            inputs=[text_in, dataset_in, max_len_in, use_prev_in, prev_in],
            outputs=[table_out, seq_out, json_out],
        )
        clear_btn.click(
            fn=lambda: ("", "celer", 30, False, "", "_Results will appear here._", "", "{}"),
            inputs=None,
            outputs=[
                text_in,
                dataset_in,
                max_len_in,
                use_prev_in,
                prev_in,
                table_out,
                seq_out,
                json_out,
            ],
        )

    return demo


if __name__ == "__main__":
    demo = build_demo()
    demo.queue(max_size=16).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        show_error=True,
    )