import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# --- 模型加载配置 ---
ADAPTER_REPO_ID = "jinv2/gpt2-lora-trajectory-prediction"
BASE_MODEL_NAME = "gpt2"

# --- 加载模型和分词器 ---
# 这是一个耗时操作，Gradio应用启动时会执行一次
print(f"开始加载模型: {BASE_MODEL_NAME} 和适配器: {ADAPTER_REPO_ID}")
try:
    base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_NAME)
    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO_ID) # 适配器仓库通常包含分词器配置

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        print("tokenizer.pad_token 设置为 tokenizer.eos_token")

    model = PeftModel.from_pretrained(base_model, ADAPTER_REPO_ID)
    model.eval() # 设置为评估模式

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    print(f"模型和分词器加载完成，运行在: {device}")
except Exception as e:
    print(f"模型加载失败: {e}")
    # 如果模型加载失败，Gradio界面可能无法正常工作，这里可以抛出错误或设置一个标志
    model = None
    tokenizer = None
    raise RuntimeError(f"无法加载模型: {e}")


# --- 推理函数 ---
def predict_trajectory(history_text_input):
    if model is None or tokenizer is None:
        return "错误: 模型未能成功加载，请检查Space的日志。"

    if not history_text_input or not history_text_input.strip():
        return "错误: 请输入有效的历史轨迹。"

    # 格式化为模型期望的输入
    # 假设用户只输入历史点，例如 "1.00,1.00,0.50,0.00; 1.05,1.00,0.50,0.00"
    prompt = f"历史: {history_text_input.strip()}; 预测:"
    print(f"收到的提示: {prompt}")

    try:
        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=60,  # 调整以适应预期的输出长度 (例如2-3个点)
                num_return_sequences=1,
                pad_token_id=tokenizer.pad_token_id, # 使用pad_token_id
                eos_token_id=tokenizer.eos_token_id,
                # temperature=0.7, # 如果想要一些随机性
                # do_sample=True,  # 如果想要一些随机性
                do_sample=False, # 为了演示的确定性
                num_beams=1      # 使用贪婪解码
            )

        generated_text_full = tokenizer.decode(outputs[0], skip_special_tokens=True)
        print(f"完整生成文本: {generated_text_full}")

        predicted_part = ""
        # 从完整输出中提取预测部分
        if "预测:" in generated_text_full:
            split_output = generated_text_full.split("预测:", 1)
            if len(split_output) > 1:
                predicted_part = split_output[1].strip()
                # 清理可能的末尾分号或eos token的文本残留
                if predicted_part.endswith(tokenizer.eos_token):
                    predicted_part = predicted_part[:-len(tokenizer.eos_token)].strip()
                if predicted_part.endswith(';'):
                    predicted_part = predicted_part[:-1].strip()
        else:
            # 如果模型输出不包含 "预测:"，尝试从提示后截取
            # 这部分逻辑可能需要根据模型的实际输出行为调整
            if prompt in generated_text_full:
                predicted_part = generated_text_full[len(prompt):].strip()
            else: # 假设模型只输出了预测部分 (可能需要更鲁棒的逻辑)
                predicted_part = generated_text_full.strip() # 基本回退

            if predicted_part.endswith(tokenizer.eos_token):
                predicted_part = predicted_part[:-len(tokenizer.eos_token)].strip()
            if predicted_part.endswith(';'):
                predicted_part = predicted_part[:-1].strip()


        print(f"提取的预测部分: {predicted_part}")
        return predicted_part

    except Exception as e:
        print(f"推理时发生错误: {e}")
        import traceback
        traceback.print_exc()
        return f"推理错误: {str(e)}"

# --- 创建 Gradio 界面 ---
# 使用 gr.Markdown 来显示更丰富的文本和说明
readme_url = f"https://huggingface.co/{ADAPTER_REPO_ID}"
description = f"""
# GPT-2 LoRA 轨迹预测 Demo
这是一个使用微调后的 `gpt2` 模型进行轨迹预测的简单演示。
模型仓库: [{ADAPTER_REPO_ID}]({readme_url})

**如何使用:**
1. 在下面的 "历史轨迹输入" 框中输入历史轨迹点。
2. 格式应为: `x1,y1,vx1,vy1; x2,y2,vx2,vy2` (例如，两个历史点，用分号隔开)。
3. 每个点包含四个逗号分隔的数值: x坐标, y坐标, x方向速度, y方向速度。
4. 点击 "预测轨迹" 按钮查看模型生成的未来轨迹点。
"""

# 示例输入
example_history = "0.00,0.00,1.00,0.00; 0.10,0.00,1.00,0.00"


# 定义界面组件
iface = gr.Interface(
    fn=predict_trajectory,
    inputs=gr.Textbox(
        lines=3,
        placeholder="例如: 0.00,0.00,1.00,0.00; 0.10,0.00,1.00,0.00",
        label="历史轨迹输入 (格式: x1,y1,vx1,vy1; x2,y2,vx2,vy2; ...)",
        value=example_history # 设置一个默认示例值
    ),
    outputs=gr.Textbox(
        lines=3,
        label="模型预测的未来轨迹 (文本格式)"
    ),
    title="基于LLM的轨迹预测",
    description=description,
    examples=[
        ["1.00,1.00,0.50,0.00; 1.05,1.00,0.50,0.00"],
        ["-2.0,0.5,0.0,1.0; -2.0,0.6,0.0,1.0; -2.0,0.7,0.0,1.0"], # 三个历史点
        ["0.0,0.0,0.2,0.2; 0.02,0.02,0.2,0.2"]
    ],
    allow_flagging='never' # 通常用于演示，不需要用户标记
)

# 启动 Gradio 应用 (在 Hugging Face Spaces 上会自动处理)
if __name__ == "__main__":
    if model is not None and tokenizer is not None: # 仅当模型加载成功时启动
        print("正在本地启动Gradio应用...")
        iface.launch()
    else:
        print("模型未能加载，Gradio应用无法启动。请检查日志。")