import os
import gradio as gr
from transformers import pipeline
import torch
from openai import OpenAI
from pypdf import PdfReader


device = "cuda" if torch.cuda.is_available() else "cpu"

vision_pipe = pipeline(
    "image-to-text",
    model="nlpconnect/vit-gpt2-image-captioning",
    device=0 if device == "cuda" else -1
)

api_key = os.environ.get("YUNWU_API_KEY")
if not api_key:
    raise RuntimeError("YUNWU_API_KEY not set in Space secrets.")

client = OpenAI(
    api_key=api_key,
    base_url="https://yunwu.ai/v1"
)


def call_llm(prompt, model="deepseek-chat", temperature=0.2, max_tokens=512):
    resp = client.chat.completions.create(
        model=model,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are an academic figure interpretation assistant. "
                    "Write accurate, natural-sounding text for scientific use."
                ),
            },
            {"role": "user", "content": prompt},
        ],
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return resp.choices[0].message.content.strip()


def extract_pdf_snippet(pdf_path, max_chars=2500):
    if not pdf_path:
        return ""
    try:
        reader = PdfReader(pdf_path)
        texts = []
        for page in reader.pages:
            txt = page.extract_text() or ""
            texts.append(txt)
            if sum(len(t) for t in texts) > max_chars * 1.5:
                break
        full = " ".join(texts)
        return full[:max_chars]
    except Exception:
        return ""


def analyze_figure(image, style, pdf_path):
    if image is None:
        return None, "Please upload a figure first.", "", ""

    vision_raw = vision_pipe(image)[0]["generated_text"]
    pdf_context = extract_pdf_snippet(pdf_path)

    step1_prompt = f"""
You are looking at a scientific figure from a paper.

A vision model produced this rough description:
\"\"\"{vision_raw}\"\"\"

Paper context (may be noisy or incomplete):
\"\"\"{pdf_context}\"\"\"

Task:
Write a clear, paper-style explanation (4–6 sentences) of what the figure shows.
- Describe what is compared on the x-axis/panels and what the y-axis measures.
- Summarize the main pattern/trend across conditions.
- Do NOT invent exact numbers, statistics, or p-values.
- Do NOT restate the full experimental design.

Write in formal academic English, but keep it readable.
"""
    step1_text = call_llm(step1_prompt, model="deepseek-chat", max_tokens=420)

    step2_prompt = f"""
You are helping a student annotate this scientific figure for a presentation.

Rough visual description:
\"\"\"{vision_raw}\"\"\"

Paper context:
\"\"\"{pdf_context}\"\"\"

Give practical suggestions for how to annotate the figure directly on the image.
Constraints:
- Output 4–6 bullet points.
- Use plain hyphen bullets only (no numbering, no bold, no asterisks, no markdown headings).
- Sound like a helpful human TA, not an AI.
- Focus on labels, arrows, callouts, grouping, legend clarity, and highlighting key contrasts.
"""
    step2_text = call_llm(step2_prompt, model="deepseek-chat", temperature=0.4, max_tokens=260)
    step2_text = step2_text.replace("**", "").replace("*", "").strip()

    style_map = {
        "formal": "formal but still plain language, suitable for a report",
        "fluency": "smooth, narrative, easy to speak aloud in a presentation",
        "simple": "very simple words for quick student notes"
    }
    style_instruction = style_map.get(style, style_map["fluency"])

    step3_prompt = f"""
You are writing a plain-language explanation of the figure for a student.

Paper-style meaning:
\"\"\"{step1_text}\"\"\"

Paper context:
\"\"\"{pdf_context}\"\"\"

Now paraphrase/explain the figure in {style_instruction}.
- 3–5 sentences.
- Keep it accurate to the paper-style meaning above.
- No numbers or p-values unless they are explicitly visible in the figure.
- Make it easy to reuse in slides or homework.
"""
    step3_text = call_llm(step3_prompt, model="deepseek-chat", temperature=0.5, max_tokens=240)

    return image, step1_text, step2_text, step3_text


with gr.Blocks() as demo:
    gr.Markdown("## ChartSmith – AI Figure Explainer (multi-model workflow)")

    with gr.Row():
        with gr.Column():
            img_in = gr.Image(
                type="pil",
                label="Upload your scientific figure (screenshot is fine)"
            )

            style = gr.Radio(
                ["formal", "fluency", "simple"],
                value="fluency",
                label="Explanation style for Step 3"
            )

            pdf_in = gr.File(
                label="Upload the paper PDF (optional, for context)",
                type="filepath"
            )

            run_btn = gr.Button("Run workflow", variant="primary")

        with gr.Column():
            preview_img = gr.Image(label="Figure preview")

            step1_box = gr.Textbox(
                label="Step 1: Explanation of what the figure shows (paper-style)",
                lines=8
            )

            step2_box = gr.Textbox(
                label="Step 2: Suggestions for annotating the figure",
                lines=7
            )

            step3_box = gr.Textbox(
                label="Step 3: Plain-language explanation (style-adapted)",
                lines=6
            )

    run_btn.click(
        analyze_figure,
        inputs=[img_in, style, pdf_in],
        outputs=[preview_img, step1_box, step2_box, step3_box],
    )

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