app.py

"""
SlideScholar — app.py
HuggingFace Spaces deployment (CPU basic — free tier).
"""

import os, re, json
from pathlib import Path
from typing  import List, Dict

import numpy as np
import faiss
import torch

import gradio_client.utils as _gcu
_orig = _gcu.json_schema_to_python_type
def _safe(schema, defs=None):
    if not isinstance(schema, dict): return "Any"
    try: return _orig(schema)
    except Exception: return "Any"
_gcu.json_schema_to_python_type = _safe

import gradio as gr
from sentence_transformers import SentenceTransformer

HF_TOKEN    = os.environ.get("HF_TOKEN", "")
CHUNKS_PATH = Path("chunks.json")
FAISS_PATH  = Path("slidescholar.faiss")
MISTRAL_ID  = "mistralai/Mistral-7B-Instruct-v0.2"
_db         = None
_PLACEHOLDER = "SSANSWER"


# ══════════════════════════════════════════════════════════════════════════════
# VDB
# ══════════════════════════════════════════════════════════════════════════════

class vdb:
    def __init__(self, model_name="intfloat/e5-large-v2"):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model  = SentenceTransformer(model_name, device=self.device)
        self.dim    = 1024
        self.index  = None
        self.chunks: List[Dict] = []

    def _embed(self, texts, prefix):
        return self.model.encode(
            [f"{prefix}{t}" for t in texts],
            batch_size=32, normalize_embeddings=True,
            convert_to_numpy=True, show_progress_bar=False,
        ).astype(np.float32)

    def load(self, faiss_path, chunks):
        self.index  = faiss.read_index(faiss_path)
        self.chunks = chunks

    def search(self, query, top_k=8):
        if self.index is None: raise RuntimeError("Index not loaded.")
        q = self._embed([query], "query: ")
        distances, indices = self.index.search(q, top_k)
        return [
            {"score": float(d), "chunk": self.chunks[i]}
            for d, i in zip(distances[0], indices[0])
            if i != -1 and i < len(self.chunks)
        ]

    def add_texts(self, texts, metadatas):
        self.chunks.extend([{"text": t, "metadata": m} for t, m in zip(texts, metadatas)])
        emb = self._embed(texts, "passage: ")
        if self.index is None: self.index = faiss.IndexFlatIP(self.dim)
        self.index.add(emb)


def _load_index():
    global _db
    if _db is not None: return _db
    if not CHUNKS_PATH.exists(): raise FileNotFoundError("chunks.json not found.")
    if not FAISS_PATH.exists():  raise FileNotFoundError("slidescholar.faiss not found.")
    with open(CHUNKS_PATH) as f: chunks = json.load(f)
    _db = vdb()
    _db.load(str(FAISS_PATH), chunks)
    print(f"Index loaded — {_db.index.ntotal} vectors")
    return _db


# ══════════════════════════════════════════════════════════════════════════════
# GENERATION
# ══════════════════════════════════════════════════════════════════════════════

def _generate(prompt, max_new_tokens=700, temperature=0.1):
    if not HF_TOKEN:
        raise RuntimeError("HF_TOKEN secret not set.")
    from huggingface_hub import InferenceClient
    client   = InferenceClient(model=MISTRAL_ID, token=HF_TOKEN)
    response = client.chat_completion(
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_new_tokens,
        temperature=temperature,
    )
    return response.choices[0].message.content.strip()


# ══════════════════════════════════════════════════════════════════════════════
# RAG HELPERS
# ══════════════════════════════════════════════════════════════════════════════

def _fmt_context(results, max_chars=900):
    parts = []
    for i, r in enumerate(results):
        m = r["chunk"]["metadata"]
        parts.append(
            f"[Source {i+1}: {m.get('name','?')} — Slide {m.get('slide','?')}]\n"
            f"{r['chunk']['text'][:max_chars]}"
        )
    return "\n\n" + ("\n\n" + "─"*40 + "\n\n").join(parts)

def _fmt_sources(results):
    lines = ["---", "**📚 Sources used:**"]
    for r in results:
        m = r["chunk"]["metadata"]
        lines.append(f"- **{m.get('name','?')}** — Slide {m.get('slide','?')} *(score: {r['score']:.2f})*")
    return "\n".join(lines)

def _strip_bracket_hints(text):
    """Remove lines that are purely [placeholder hints] the model echoed back."""
    lines   = text.split("\n")
    cleaned = [l for l in lines if not re.match(r"^\s*\[.*\]\s*$", l)]
    return re.sub(r"\n{3,}", "\n\n", "\n".join(cleaned)).strip()


# ══════════════════════════════════════════════════════════════════════════════
# EXAM FORMATTER
# ══════════════════════════════════════════════════════════════════════════════

def _fix_exam_format(text):
    text = text.replace("**", "")
    text = re.sub(r"\[([ABCD])\]\s*",    r"\1) ", text)
    text = re.sub(r"\(([ABCD])\)\s*",    r"\1) ", text)
    text = re.sub(r"([ABCD])\)([^\s)])", r"\1) \2", text)
    text = re.sub(
        r"[✅✓]?\s*\[?Answer:?\s*([ABCD])[\])]?\s*",
        lambda m: f" {_PLACEHOLDER}_{m.group(1)} ", text
    )
    text = re.sub(r"([^\n]) *([ABCD]\) )", r"\1\n\2", text)
    text = re.sub(r"^([ABCD])\) ", r"**\1)** ", text, flags=re.MULTILINE)
    text = re.sub(rf"\s*{_PLACEHOLDER}_([ABCD])\s*", r"\n\n✅ **Answer: \1)** ", text)
    text = re.sub(r" {2,}",  " ",    text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


# ══════════════════════════════════════════════════════════════════════════════
# PROMPTS — context FIRST, instructions LAST
# ══════════════════════════════════════════════════════════════════════════════

def _study_prompt(context, query):
    return (
        "You are an AI tutor helping students study for exams.\n"
        "STRICT RULES:\n"
        "1. Use ONLY the lecture slide content below. Do not add outside knowledge.\n"
        "2. Copy key terms, formulas, and definitions verbatim from the slides.\n"
        "3. Cite every point with [Source N].\n"
        "4. If something is not covered, say 'Not covered in provided slides.'\n\n"
        f"Context:\n{context}\n\n"
        f"Question: {query}\n\n"
        "Output (use these exact headers):\n"
        "## Key Concepts\n"
        "## Definitions\n"
        "## Important Formulas\n"
        "## Common Exam Topics\n"
        "## Summary"
    )


def _flashcard_prompt(context, query, batch=1):
    extra = "" if batch == 1 else (
        f"Generate 5 MORE flashcards covering DIFFERENT aspects of {query} not yet covered.\n"
    )
    return (
        f"Use ONLY these lecture slides to make 5 flashcard Q&A pairs about \"{query}\".\n"
        f"{extra}\n"
        f"LECTURE SLIDES:\n{context}\n\n"
        "Write exactly 5 cards in this format. No other text.\n\n"
        "Q: first question\n"
        "A: first answer (under 20 words)\n"
        "Source: slide reference\n\n"
        "Q: second question\n"
        "A: second answer (under 20 words)\n"
        "Source: slide reference\n\n"
        "Q: third question\n"
        "A: third answer (under 20 words)\n"
        "Source: slide reference\n\n"
        "Q: fourth question\n"
        "A: fourth answer (under 20 words)\n"
        "Source: slide reference\n\n"
        "Q: fifth question\n"
        "A: fifth answer (under 20 words)\n"
        "Source: slide reference\n\n"
        f"Now write 5 real flashcards from the slides about \"{query}\":\n\n"
        "Q:"
    )


def _exam_prompt(context, query):
    return (
        "You are a professor creating a practice exam from lecture slides only.\n"
        "Do not use outside knowledge.\n\n"
        f"LECTURE SLIDES:\n{context}\n\n"
        f"TOPIC: {query}\n\n"
        "Create a practice exam with EXACTLY this structure:\n\n"
        f"## Practice Exam: {query}\n\n"
        "### Multiple Choice (5 questions)\n\n"
        "For each MCQ:\n"
        "**Question N.** [Difficulty] Question text?\n"
        "**A)** Option\n**B)** Option\n**C)** Option\n**D)** Option\n"
        "✅ **Answer: X)** One-sentence explanation. [Source N]\n\n"
        "---\n\n"
        "### Short Answer (3 questions)\n\n"
        "**Question N.** [Difficulty] Question text?\n"
        "**Model Answer:** Full answer in 2-3 sentences. [Source N]\n\n"
        "---\n\n"
        "### Answer Key\n"
        "1-?  2-?  3-?  4-?  5-?\n\n"
        "Now write the full exam. Start with Question 1:"
    )


def _eli5_prompt(context, query):
    return (
        "You are a friendly tutor. Use ONLY the lecture slides to explain a concept simply.\n\n"
        f"LECTURE SLIDES:\n{context}\n\n"
        f"CONCEPT: {query}\n\n"
        "Write your explanation with these four sections:\n\n"
        "## Simple Explanation\n\n"
        "**The core idea in one sentence:**\n\n"
        "**Real-world analogy:**\n\n"
        "**How it works, step by step:**\n"
        "1.\n2.\n3.\n\n"
        "**Why it matters:**\n\n"
        "Begin now:"
    )


def _gap_prompt(context, question, student_ans, correct_ans):
    return (
        "A student got an exam question wrong. Use ONLY these lecture slides to explain why.\n\n"
        f"LECTURE SLIDES:\n{context}\n\n"
        f"EXAM QUESTION: {question}\n"
        f"STUDENT ANSWERED: {student_ans}\n"
        f"CORRECT ANSWER: {correct_ans}\n\n"
        "Write a re-explanation with these four sections:\n\n"
        "## Why the correct answer is right\n"
        "Cite the relevant slide using [Source N].\n\n"
        "## Why the student answer was wrong\n"
        "Be specific and constructive.\n\n"
        "## Key concept to remember\n"
        "State it in one sentence.\n\n"
        "## Memory aid\n"
        "Give a simple analogy or mnemonic.\n\n"
        "Begin now:"
    )


# ══════════════════════════════════════════════════════════════════════════════
# FLASHCARD PARSER — Q:/A:/Source: line format
# ══════════════════════════════════════════════════════════════════════════════

def _parse_flashcards(raw):
    # Model starts mid-card-1 since prompt ends with "Q:"
    if not raw.strip().upper().startswith("Q:"):
        raw = "Q:" + raw

    pat_q = re.compile(r"^Q:\s*(.+)", re.MULTILINE)
    pat_a = re.compile(r"^A:\s*(.+)", re.MULTILINE)
    pat_s = re.compile(r"^Source:\s*(.+)", re.MULTILINE | re.IGNORECASE)

    SKIP_Q = {"first question", "second question", "third question", "fourth question",
              "fifth question", "write the question here", "[question]", "question here"}
    SKIP_A = {"first answer (under 20 words)", "second answer (under 20 words)",
              "third answer (under 20 words)", "fourth answer (under 20 words)",
              "fifth answer (under 20 words)", "write a short answer here", "[answer]"}

    cards = []
    blocks = re.split(r"(?=^Q:)", raw, flags=re.MULTILINE)
    for block in blocks:
        block = block.strip()
        if not block:
            continue
        q_m = pat_q.match(block)
        a_m = pat_a.search(block)
        s_m = pat_s.search(block)
        if not (q_m and a_m):
            continue
        q = q_m.group(1).strip().rstrip("?") + ("?" if not q_m.group(1).strip().endswith("?") else "")
        a = a_m.group(1).strip()
        if q.lower().rstrip("?") in SKIP_Q or a.lower() in SKIP_A:
            continue
        cards.append({
            "q":      q,
            "a":      a,
            "source": s_m.group(1).strip() if s_m else "",
        })

    if not cards:
        raise ValueError("No flashcards parsed from output")
    return cards


def _render_flashcards(cards, query):
    lines = [f"## 🃏 Flashcards: *{query}*", f"*{len(cards)} cards generated*", ""]
    for i, card in enumerate(cards, 1):
        q = card.get("q", "").strip()
        a = card.get("a", "").strip()
        s = card.get("source", "").strip()
        lines += ["---", f"**Q{i}.** {q}", "", f"> {a}"]
        if s: lines += ["> ", f"> *📖 {s}*"]
        lines.append("")
    return "\n".join(lines)


# ══════════════════════════════════════════════════════════════════════════════
# TAB HANDLERS
# ══════════════════════════════════════════════════════════════════════════════

def handle_study_guide(query, k):
    if not query.strip(): return "Please enter a topic or question.", ""
    try:
        db      = _load_index()
        results = db.search(query, top_k=int(k))
        output  = _generate(_study_prompt(_fmt_context(results, max_chars=800), query), max_new_tokens=900)
        return output, _fmt_sources(results)
    except Exception as e:
        return f"❌ {e}", ""


def handle_flashcards(query, k):
    if not query.strip(): return "Please enter a topic.", ""
    try:
        db      = _load_index()
        top_k   = min(int(k), 4)
        results = db.search(query, top_k=top_k)
        ctx     = _fmt_context(results, max_chars=500)

        # Two batches of 5 — more reliable than asking for 10 at once
        all_cards = []
        for batch in [1, 2]:
            try:
                raw   = _generate(_flashcard_prompt(ctx, query, batch=batch),
                                   max_new_tokens=550, temperature=0.2)
                cards = _parse_flashcards(raw)
                all_cards.extend(cards)
            except Exception:
                pass  # show whatever we have if a batch fails

        if not all_cards:
            return "❌ Could not generate flashcards — try a more specific topic.", _fmt_sources(results)

        display = _render_flashcards(all_cards[:10], query)
        if len(all_cards) < 10:
            display += f"\n\n*Note: {len(all_cards)} cards generated*"
        return display, _fmt_sources(results)
    except Exception as e:
        return f"❌ {e}", ""


def handle_exam(query, k):
    if not query.strip(): return "Please enter a topic.", ""
    try:
        db      = _load_index()
        top_k   = min(int(k), 5)
        results = db.search(query, top_k=top_k)
        raw     = _generate(_exam_prompt(_fmt_context(results, max_chars=550), query),
                            max_new_tokens=900, temperature=0.2)
        return _fix_exam_format(raw), _fmt_sources(results)
    except Exception as e:
        return f"❌ {e}", ""


def handle_eli5(query, k):
    if not query.strip(): return "Please enter a concept.", ""
    try:
        db      = _load_index()
        top_k   = min(int(k), 5)
        results = db.search(query, top_k=top_k)
        output  = _strip_bracket_hints(
            _generate(_eli5_prompt(_fmt_context(results, max_chars=600), query),
                      max_new_tokens=600, temperature=0.1)
        )
        return output, _fmt_sources(results)
    except Exception as e:
        return f"❌ {e}", ""


def handle_gap(question, student_ans, correct_ans, k):
    if not all([question.strip(), student_ans.strip(), correct_ans.strip()]):
        return "Please fill in all three fields.", ""
    try:
        db      = _load_index()
        top_k   = min(int(k), 5)
        results = db.search(f"{question} {correct_ans}", top_k=top_k)
        output  = _strip_bracket_hints(
            _generate(_gap_prompt(_fmt_context(results, max_chars=550), question, student_ans, correct_ans),
                      max_new_tokens=600, temperature=0.1)
        )
        return output, _fmt_sources(results)
    except Exception as e:
        return f"❌ {e}", ""


def handle_upload(files):
    if not files: return "No files selected."
    try:
        import pdfplumber
    except ImportError:
        return "pdfplumber not installed."
    try:
        db = _load_index()
        total, names = 0, []
        for file in files:
            path = Path(file if isinstance(file, str) else file.name)
            texts, metas = [], []
            with pdfplumber.open(str(path)) as pdf:
                for i, page in enumerate(pdf.pages):
                    raw = (page.extract_text() or "").strip()
                    if raw:
                        texts.append(raw)
                        metas.append({
                            "name": path.stem, "slide": i+1, "lecture_num": None,
                            "source": str(path), "filename": path.name,
                            "filetype": "pdf", "is_scanned": False,
                            "char_count": len(raw), "chunk_id": f"upload_{path.stem}_{i+1}",
                        })
            db.add_texts(texts, metas)
            total += len(texts)
            names.append(f"{path.name} ({len(texts)} slides)")
        return (
            f"✅ Added {total} slides from {len(files)} file(s):\n"
            + "\n".join(f"  • {n}" for n in names)
            + f"\n\nTotal index size: {len(db.chunks)} chunks"
        )
    except Exception as e:
        return f"❌ Upload failed: {e}"


# ══════════════════════════════════════════════════════════════════════════════
# GRADIO UI
# ══════════════════════════════════════════════════════════════════════════════

CSS = """
.tab-nav button { font-size: 15px !important; padding: 10px 18px !important; }
.sources-box    { border-left: 4px solid #0D9488; padding: 12px 16px;
                  border-radius: 6px; font-size: 13px; margin-top: 8px; }
.output-box     { min-height: 280px; }
footer          { display: none !important; }
"""

HEADER = """
# 📚 SlideScholar
### AI Study Assistant — STATGR5293 · GenAI Using LLMs · Spring 2026

Powered by **Mistral-7B-Instruct** + **FAISS** retrieval over your actual lecture slides.
All outputs are grounded in course content — not generic AI responses.

> ⏱️ First generation may take 30–60s while the model warms up on HuggingFace servers.
"""


def build_app():
    with gr.Blocks(css=CSS, title="SlideScholar") as app:
        gr.Markdown(HEADER)
        with gr.Row():
            k_slider = gr.Slider(minimum=3, maximum=15, value=8, step=1,
                                 label="Slides to retrieve (k)",
                                 info="More slides = richer context but slower.")

        with gr.Tab("📝 Study Guide"):
            gr.Markdown("Get structured notes with citations grounded in your lecture slides.")
            sg_query   = gr.Textbox(label="Topic or question",
                                    placeholder="e.g.  attention mechanism and transformers", lines=2)
            sg_btn     = gr.Button("Generate Study Guide", variant="primary")
            sg_output  = gr.Markdown(elem_classes=["output-box"])
            sg_sources = gr.Markdown(elem_classes=["sources-box"])
            sg_btn.click(handle_study_guide, [sg_query, k_slider], [sg_output, sg_sources])

        with gr.Tab("🃏 Flashcards"):
            gr.Markdown("Generate Q&A flashcard pairs from your slides.")
            fc_query   = gr.Textbox(label="Topic",
                                    placeholder="e.g.  gradient descent and optimization", lines=2)
            fc_btn     = gr.Button("Generate Flashcards", variant="primary")
            fc_output  = gr.Markdown(elem_classes=["output-box"])
            fc_sources = gr.Markdown(elem_classes=["sources-box"])
            fc_btn.click(handle_flashcards, [fc_query, k_slider], [fc_output, fc_sources])

        with gr.Tab("📋 Practice Exam"):
            gr.Markdown("Generate a practice exam: 5 MCQs + 3 short-answer questions with answer key.")
            pe_query   = gr.Textbox(label="Topic",
                                    placeholder="e.g.  transformer architecture and self-attention", lines=2)
            pe_btn     = gr.Button("Generate Exam", variant="primary")
            pe_output  = gr.Markdown(elem_classes=["output-box"])
            pe_sources = gr.Markdown(elem_classes=["sources-box"])
            pe_btn.click(handle_exam, [pe_query, k_slider], [pe_output, pe_sources])

        with gr.Tab("💡 ELI5"):
            gr.Markdown("Explain a complex concept in simple terms using your lecture slides.")
            e5_query   = gr.Textbox(label="Concept",
                                    placeholder="e.g.  what is the attention mechanism?", lines=2)
            e5_btn     = gr.Button("Explain Simply", variant="primary")
            e5_output  = gr.Markdown(elem_classes=["output-box"])
            e5_sources = gr.Markdown(elem_classes=["sources-box"])
            e5_btn.click(handle_eli5, [e5_query, k_slider], [e5_output, e5_sources])

        with gr.Tab("🎯 Gap Analysis"):
            gr.Markdown("Got a question wrong? Enter the question, your answer, and the correct answer.")
            gap_q   = gr.Textbox(label="Exam question", lines=2)
            gap_s   = gr.Textbox(label="Your answer", lines=2)
            gap_c   = gr.Textbox(label="Correct answer", lines=2)
            gap_btn = gr.Button("Explain My Mistake", variant="primary")
            gap_out = gr.Markdown(elem_classes=["output-box"])
            gap_src = gr.Markdown(elem_classes=["sources-box"])
            gap_btn.click(handle_gap, [gap_q, gap_s, gap_c, k_slider], [gap_out, gap_src])

        with gr.Tab("📂 Upload Slides"):
            gr.Markdown(
                "Upload additional PDF lecture slides to extend the knowledge base.\n\n"
                "> **Note:** Text is extracted directly — no vision model on Spaces."
            )
            upload_files  = gr.File(label="Upload PDF files", file_count="multiple",
                                    file_types=[".pdf"], type="filepath")
            upload_btn    = gr.Button("Add to Index", variant="primary")
            upload_status = gr.Textbox(label="Status", interactive=False, lines=5)
            upload_btn.click(handle_upload, [upload_files], [upload_status])

        gr.Markdown("---\n*SlideScholar · STATGR5293 · GenAI Using LLMs · Spring 2026 · Columbia University*")
    return app


if __name__ == "__main__":
    print("Pre-loading index...")
    try:
        _load_index()
        print(f"Index ready — {_db.index.ntotal} vectors")
    except Exception as e:
        print(f"Warning: {e}")
    build_app().launch(server_name="0.0.0.0", server_port=7860, show_api=False)