import os
# Set expandable segments at the very top, before any torch/CUDA import, to avoid
# the ZeroGPU CUDACachingAllocator NVML assert during transient allocation spikes.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

# The SHINE custom modeling code targets transformers 4.57.x (the version used to
# train the checkpoint). That release pins huggingface-hub<1.0, which conflicts with
# the Space runtime's hub 1.x (required by gradio). We install transformers 4.57.6
# with --no-deps so the platform-managed huggingface_hub 1.x stays intact, then relax
# transformers' hard hub version check (which is only an advisory upper bound; the
# 4.57 API works fine against hub 1.x). Its other deps come from requirements.txt.
import subprocess  # noqa: E402
import sys  # noqa: E402


def _install_transformers():
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "--no-deps", "transformers==4.57.6"],
        check=True,
    )
    # Relax the huggingface-hub pin baked into transformers' version table so the
    # import-time dependency check does not reject the platform hub 1.x. We must edit
    # the files on disk WITHOUT importing transformers (importing it triggers the very
    # check we are trying to disable), so locate the package via site-packages paths.
    import sysconfig
    import glob
    candidates = set()
    for key in ("purelib", "platlib"):
        p = sysconfig.get_paths().get(key)
        if p:
            candidates.add(p)
    for sp in list(candidates):
        for fname in ("dependency_versions_table.py", "dependency_versions_check.py"):
            for path in glob.glob(os.path.join(sp, "transformers", fname)):
                try:
                    with open(path, "r", encoding="utf-8") as fh:
                        src = fh.read()
                except OSError:
                    continue
                patched = src.replace(
                    "huggingface-hub>=0.34.0,<1.0", "huggingface-hub>=0.34.0"
                )
                if patched != src:
                    with open(path, "w", encoding="utf-8") as fh:
                        fh.write(patched)


_install_transformers()

import spaces  # noqa: E402  must come before torch / transformers
import gc  # noqa: E402
import re  # noqa: E402
import time  # noqa: E402
import logging  # noqa: E402
from typing import Tuple  # noqa: E402

import torch  # noqa: E402
import gradio as gr  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402
from transformers import AutoTokenizer, AutoConfig  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402

from LoraQwen import LoraQwen3ForCausalLM  # noqa: E402
from metanetwork_family import Metanetwork  # noqa: E402
from utils.myseed import set_seed  # noqa: E402
from utils.myfreeze import freeze  # noqa: E402
from utils.mysaveload import load_checkpoint  # noqa: E402

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("shine")

torch.backends.cuda.matmul.allow_tf32 = True

# ---------------------------------------------------------------------------
# Configuration (mirrors the authors' inference.ipynb for SHINE-ift_mqa)
# ---------------------------------------------------------------------------
BACKBONE_ID = "Qwen/Qwen3-8B"
HYPERNET_ID = "Yewei-Liu/SHINE-ift_mqa"
CONTEXT_MAX_LENGTH = 1550
CONVERSATION_MAX_LENGTH = 5000

CONF_DICT = {
    "name": "shine",
    "mode": "train",
    "run": {"seed": 42, "device": "cuda"},
    "model": {
        "lora_r": 8,
        "metalora_r": 128,
        "ift_additional_metalora_r": -1,
        "num_mem_token": 4,
        "metamodel_class_path": "LoraQwen.LoraQwen3ForCausalLM",
        "config_class_path": "LoraQwen.Qwen3Config",
    },
    "metanetwork": {
        "type": "transformer",
        "method": "rl",
        "transformer_cfg": {
            "encoder_cfg": {
                "d_model": 4096, "nhead": 32, "dim_feedforward": 8192,
                "dropout": 0, "activation": "gelu", "layer_norm_eps": 1e-05,
                "batch_first": True, "norm_first": False, "bias": True,
            },
            "couple_encoder_cfg": {
                "d_model": 4096, "nhead": 32, "dim_feedforward": 8192,
                "dropout": 0, "activation": "gelu", "layer_norm_eps": 1e-05,
                "batch_first": True, "norm_first": False, "bias": True,
            },
            "layer_transformer_first": True,
            "mean_pool_size": 1,
            "num_layers": 4,
            "couple_num_layers": 0,
            "scale": 0.001,
        },
    },
    "hidden_size": -1,
    "num_layers": -1,
    "num_mem_token": 4,
}

cfg = OmegaConf.create(CONF_DICT)

# ---------------------------------------------------------------------------
# Load the backbone + hypernetwork checkpoint at module scope
# ---------------------------------------------------------------------------
set_seed(int(cfg.run.seed))

logger.info("Downloading backbone %s ...", BACKBONE_ID)
backbone_dir = snapshot_download(BACKBONE_ID)
logger.info("Downloading hypernetwork checkpoint %s ...", HYPERNET_ID)
ckpt_dir = snapshot_download(HYPERNET_ID)

logger.info("Loading config ...")
config = AutoConfig.from_pretrained(backbone_dir)
config.num_mem_token = -1
cfg.hidden_size = config.hidden_size
cfg.num_layers = config.num_hidden_layers

# Compute num_mem_token from the LoRA param count / (hidden * layers)
with torch.device("meta"):
    _tmp = LoraQwen3ForCausalLM(config)
lora_params = _tmp.lora_params_numel(cfg.model.lora_r)
base_params = cfg.hidden_size * cfg.num_layers
assert lora_params % base_params == 0
config.num_mem_token = lora_params // base_params
cfg.num_mem_token = config.num_mem_token
del _tmp
gc.collect()
logger.info("num_mem_token = %d", config.num_mem_token)

logger.info("Loading tokenizer ...")
tokenizer = AutoTokenizer.from_pretrained(backbone_dir, padding_side="left", use_fast=True)
tokenizer.add_tokens(["<RECON>", "<COMP>", "<NOTHING>"])
# Chat template used by the authors (Qwen3 im_start/im_end format)
tokenizer.chat_template = (
    "{%- if messages[0].role == 'system' %}"
    "{{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}"
    "{%- endif %}"
    "{%- for message in messages %}"
    "{%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}"
    "{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}"
    "{%- elif message.role == \"assistant\" %}"
    "{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}"
    "{%- endif %}"
    "{%- endfor %}"
    "{%- if add_generation_prompt %}"
    "{{- '<|im_start|>assistant\n' }}"
    "{%- if enable_thinking is not defined or enable_thinking != false %}"
    "{{- '<think>\n\n</think>\n\n' }}"
    "{%- endif %}"
    "{%- endif %}"
)

DTYPE = torch.bfloat16


def _cast_tree(obj, dtype):
    """Recursively cast every floating-point tensor in a nested dict/list to dtype,
    preserving the leaf/requires_grad status set up by load_checkpoint."""
    if torch.is_tensor(obj):
        if obj.is_floating_point():
            return obj.to(dtype)
        return obj
    if isinstance(obj, dict):
        return {k: _cast_tree(v, dtype) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_cast_tree(v, dtype) for v in obj]
    return obj


logger.info("Loading backbone weights ...")
# The reference inference.ipynb runs in float32, but the full-precision Qwen3-8B (~36 GB)
# plus generation activations exceed the ZeroGPU allocator's transient budget (surfaces as
# the NVML CUDACachingAllocator assert). We run the whole stack in bfloat16 instead: the
# backbone, the hypernetwork, and the metalora tensors are all cast to bf16, and the memory
# tokens produced inside LoraQwen inherit the hidden-state dtype, so there is no dtype
# mismatch anywhere in the SHINE / in-context / question-only paths.
metamodel = LoraQwen3ForCausalLM.from_pretrained(backbone_dir, config=config, torch_dtype=DTYPE)
metamodel.reset_mem_tokens()
metamodel.resize_token_embeddings(len(tokenizer))

logger.info("Initializing hypernetwork ...")
metanetwork = Metanetwork(metamodel, cfg, metamodel.lora_params_numel(cfg.model.lora_r))
freeze(metamodel)

logger.info("Loading hypernetwork checkpoint from %s ...", ckpt_dir)
metanetwork, metalora, _ = load_checkpoint(metanetwork, ckpt_dir, "cuda")

# Unify precision across the whole stack (backbone is already bf16 from_pretrained).
metanetwork = metanetwork.to("cuda").to(DTYPE)
metamodel = metamodel.to("cuda")
metalora = _cast_tree(metalora, DTYPE)
metanetwork.eval()
logger.info("Model ready.")

SYSTEM_PROMPT_CONTEXT = (
    "You are a helpful assistant, answer the questions based on the given context. "
    "Each answer must be directly extractable from the context (i.e., an exact span or "
    "minor paraphrase for fluency). Do not invent information. Answer the question directly "
    "and output nothing else. Never enter think mode.\n\nContext: "
)
SYSTEM_PROMPT_NOCTX = (
    "You are a helpful assistant. Answer the question concisely with short words or phrases. "
    "Answer the question directly and output nothing else. Never say you don't know the answer. "
    "Never enter think mode."
)


def _extract_answer(text: str) -> str:
    answer = text
    if "<think>" in text:
        rest = text.split("<think>", 1)[1]
        answer = rest.split("</think>", 1)[1] if "</think>" in rest else ""
    answer = answer.strip()
    answer = re.sub(r"^(final answer|answer)\s*:\s*", "", answer, flags=re.IGNORECASE).strip()
    return answer


def _encode_evidence(context: str):
    enc = tokenizer(
        [context],
        max_length=CONTEXT_MAX_LENGTH,
        truncation=True,
        return_tensors="pt",
        padding="max_length",
    )
    return enc["input_ids"].to("cuda"), enc["attention_mask"].to("cuda")


@torch.no_grad()
def _generate(messages, lora_dict, max_new_tokens):
    enc = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_tensors="pt",
        max_length=CONVERSATION_MAX_LENGTH,
        truncation=True,
        return_dict=True,
    )
    input_ids = enc["input_ids"].to("cuda")
    attention_mask = enc["attention_mask"].to("cuda")
    out = metamodel.generate(
        input_ids=input_ids,
        attention_mask=attention_mask,
        max_new_tokens=max_new_tokens,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
        do_sample=False,
        ignore_mem_token=True,
        loradict=lora_dict,
    )
    new_tokens = out[0, input_ids.shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True)


def _estimate_duration(context, question, history, mode, max_new_tokens=128):
    # bf16 Qwen3-8B: a fixed ~1.5k-token evidence encode (SHINE) plus autoregressive
    # decoding. Keep a comfortable cap so a run is never killed mid-generation.
    return min(180, 45 + int(max_new_tokens * 0.6))


@spaces.GPU(duration=_estimate_duration)
def respond(context: str, question: str, history, mode: str, max_new_tokens: int = 128) -> str:
    """Answer a question about a context using SHINE.

    Args:
        context: the background document the model should ground its answer in.
        question: the user question.
        history: prior (question, answer) turns for the multi-turn conversation.
        mode: 'SHINE' (context encoded into a generated LoRA, no context in the prompt),
            'In-Context' (context placed in the system prompt), or
            'Only Question' (no context at all).
        max_new_tokens: maximum number of tokens to generate.
    """
    context = (context or "").strip()
    question = (question or "").strip()
    if not question:
        return "Please enter a question."

    max_new_tokens = int(max_new_tokens)
    history = history or []

    lora_dict = None
    messages = []

    if mode == "SHINE":
        if not context:
            return "SHINE mode needs a context to encode into a LoRA."
        evidence_ids, evidence_mask = _encode_evidence(context)
        lora_dict = metanetwork.generate_lora_dict(evidence_ids, evidence_mask, metalora)
        # No system prompt / no context in the conversation — knowledge is in the LoRA.
    elif mode == "In-Context":
        if context:
            messages.append({"role": "system", "content": SYSTEM_PROMPT_CONTEXT + context})
    else:  # Only Question
        messages.append({"role": "system", "content": SYSTEM_PROMPT_NOCTX})

    for user_q, assistant_a in history:
        messages.append({"role": "user", "content": user_q})
        messages.append({"role": "assistant", "content": assistant_a})
    messages.append({"role": "user", "content": question})

    raw = _generate(messages, lora_dict, max_new_tokens)
    answer = _extract_answer(raw)
    return answer if answer else "[no answer produced]"


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1000px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["Apple is green.", "What color is an apple?", "SHINE"],
    ["Chinese food is the best food on earth.", "Which food is the best?", "SHINE"],
    [
        "If the light is on, somebody must be at home. If the light is off, often nobody is at home. "
        "But this holds true only during the day. In the night people are all sleeping so there will "
        "always be no lights.",
        "What does it mean if the light is on?",
        "SHINE",
    ],
    [
        "Whoever organizes cheating in a national examination prescribed by law shall be sentenced to "
        "fixed-term imprisonment of not more than three years or criminal detention and shall also be "
        "fined, or shall be fined only; if the circumstances are serious, he shall be sentenced to "
        "fixed-term imprisonment of not less than three years but not more than seven years and shall "
        "also be fined.",
        "What will happen if one organizes cheating?",
        "SHINE",
    ],
]


def _run_and_show(context, question, mode, max_new_tokens=128):
    answer = respond(context, question, [], mode, max_new_tokens)
    return answer


def _run_example(context, question, mode):
    return respond(context, question, [], mode, 128)


with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🔆 SHINE: Context → LoRA in a Single Pass\n"
            "SHINE is an in-context hypernetwork that reads a **context document** and generates a "
            "**LoRA adapter** for the Qwen3-8B backbone in a single forward pass. In **SHINE** mode "
            "the context is encoded entirely into the LoRA weights — it is *not* placed in the prompt — "
            "and the model answers questions using only that adapter. Compare against **In-Context** "
            "(context in the prompt) and **Only Question** (no context).\n\n"
            "Paper: [arXiv 2602.06358](https://huggingface.co/papers/2602.06358) · "
            "Code: [MuLabPKU/SHINE](https://github.com/MuLabPKU/SHINE) · "
            "Checkpoint: [Yewei-Liu/SHINE-ift_mqa](https://huggingface.co/Yewei-Liu/SHINE-ift_mqa)"
        )
        context = gr.Textbox(
            label="Context document",
            lines=5,
            placeholder="Paste the background knowledge SHINE should encode into a LoRA...",
        )
        with gr.Row():
            question = gr.Textbox(
                label="Question", scale=4,
                placeholder="Ask something answerable from the context...",
            )
            run = gr.Button("Answer", variant="primary", scale=1)
        mode = gr.Radio(
            choices=["SHINE", "In-Context", "Only Question"],
            value="SHINE",
            label="Mode",
        )
        answer = gr.Textbox(label="Answer", lines=3)
        with gr.Accordion("Advanced settings", open=False):
            max_new_tokens = gr.Slider(
                minimum=16, maximum=512, value=128, step=16, label="Max new tokens"
            )

        run.click(
            _run_and_show,
            inputs=[context, question, mode, max_new_tokens],
            outputs=answer,
            api_name="respond",
        )
        question.submit(
            _run_and_show,
            inputs=[context, question, mode, max_new_tokens],
            outputs=answer,
        )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[context, question, mode],
            outputs=answer,
            fn=_run_example,
            cache_examples=True,
            cache_mode="lazy",
        )

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