backend/main.py

import os
import re
import json
import tempfile
from pathlib import Path
from typing import List

import fitz
import pytesseract
from pdf2image import convert_from_path
from groq import Groq
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse
from pydantic import BaseModel

app = FastAPI(title="ContraGenAI - A Protocol Contradiction Detector")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)
# ── Static files + SPA ───────────────────────────────────────────────────────
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import pathlib

_static_dir = pathlib.Path(__file__).parent / "static"

if _static_dir.exists():
    app.mount("/assets", StaticFiles(directory=str(_static_dir / "assets")), name="assets")

@app.get("/")
async def serve_index():
    return FileResponse(str(_static_dir / "index.html"))
# ── Config ────────────────────────────────────────────────────────────────────
GROQ_MODEL   = "llama-3.3-70b-versatile"
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")

# ── Prompts ───────────────────────────────────────────────────────────────────
EXTRACTION_SYSTEM = """Extract protocol parameters from this research paper.
IMPORTANT:
Do not extract numbers randomly.
Only extract values that are clearly associated with an experimental step or procedure. Always preserve context.
Return ONLY valid JSON:
{
  "title": "paper title",
  "authors": "first author et al.",
  "year": "year",
  "protocol_type": "e.g. Western blot / ELISA / PCR",
  "parameters": {
    "reagents": [{"name": "...", "concentration": "...", "vendor": "..."}],
    "cell_lines": [{"name": "...", "culture_conditions": "...", "serum": "..."}],
    "temperatures": [{"step": "...", "value": "...", "unit": "C"}],
    "timings": [{"step": "...", "duration": "..."}],
    "equipment": [{"name": "...", "settings": "..."}],
    "buffers": [{"name": "...", "composition": "...", "pH": "..."}],
    "antibodies": [{"target": "...", "dilution": "...", "vendor": "..."}],
    "other": [{"parameter": "...", "value": "..."}]
  }
}"""

COMPARISON_SYSTEM = """You are an expert in bioengineering reproducibility and experimental methodology.
Compare protocols from multiple research papers and identify ALL methodological contradictions. Always preserve context.
SEVERITY ASSIGNMENT MUST FOLLOW THE RULES BELOW.

Do NOT decide severity based on subjective judgment.

Always determine severity using the predefined parameter hierarchy.

If a parameter appears in the hierarchy below, assign the corresponding severity level.

If multiple rules apply, assign the highest severity.

If a parameter is not listed, assign LOW severity by default.

PARAMETER SEVERITY HIERARCHY

HIGH SEVERITY — Directly affects experimental outcome or biological behavior

Assign HIGH severity if contradiction involves:

- cell line identity or organism
- incubation temperature
- antibiotic or drug concentration
- enzyme concentration
- transfection conditions
- MOI (multiplicity of infection)
- centrifugation speed
- incubation duration
- buffer composition
- pH
- reagent concentration
- sample storage conditions
- reaction volume
- PCR cycle conditions

These parameters directly impact reproducibility or experimental success.


MEDIUM SEVERITY — Affects efficiency, yield, or data quality

Assign MEDIUM severity if contradiction involves:

- incubation time variations within acceptable ranges
- reagent vendor differences
- equipment model differences
- washing conditions
- dilution ratios
- culture media composition
- cell density
- mixing speed
- incubation environment (shaking/static)


LOW SEVERITY — Minor procedural or documentation differences

Assign LOW severity if contradiction involves:

- equipment brand differences
- container type
- labeling differences
- formatting differences
- reporting style differences
- minor descriptive variations

Return ONLY valid JSON (no markdown fences, no extra text) using this schema:
{
  "protocol_type": "detected protocol type",
  "contradictions": [
    {
      "parameter": "parameter name",
      "category": "reagents|cell_lines|temperatures|timings|equipment|buffers|antibodies|other",
      "severity": "high|medium|low",
      "values": {"paper_0": "value from paper 1", "paper_1": "value from paper 2"},
      "explanation": "why this is a contradiction and its likely impact on reproducibility"
    }
  ],
  "ranked_issues": [
    {"rank": 1, "parameter": "...", "severity": "high|medium|low", "brief": "one-sentence impact"}
  ],
  "optimal_protocol": {
    "rationale": "overall recommendation rationale",
    "parameters": [
      {"label": "parameter name", "value": "recommended value", "reason": "why this is optimal"}
    ]
  },
  "summary": {"total_contradictions": 0, "high": 0, "medium": 0, "low": 0}
}"""

# ── Core helpers ──────────────────────────────────────────────────────────────

def extract_text_from_pdf(pdf_path: str) -> str:
    text = ""

    # ---------- Try PyMuPDF extraction ----------
    try:
        doc = fitz.open(pdf_path)
        for page in doc:
            text += page.get_text()
    except Exception:
        pass

    # ---------- OCR fallback if no text ----------
    if not text.strip():
        try:
            images = convert_from_path(pdf_path)
            for img in images:
                text += pytesseract.image_to_string(img)
        except Exception:
            raise ValueError("PDF text extraction failed.")

    if not text.strip():
        raise ValueError("PDF text extraction failed.")

    if len(text.strip()) < 500:
        text = text[:2000]

    # ---------- Try to isolate Methods section ----------
    methods_match = re.search(
        r'(?:^\s*(?:\d+\.?\d*\.?\s+)?'
        r'(?:materials?\s+and\s+methods?|experimental\s+procedures?'
        r'|methods?\s+and\s+materials?|methods?)\s*$)'
        r'(.*?)'
        r'(?=^\s*(?:\d+\.?\d*\.?\s+)?'
        r'(?:results?|discussion|conclusion|references|acknowledgements?))',
        text, re.IGNORECASE | re.DOTALL | re.MULTILINE
    )

    if methods_match and len(methods_match.group(1).strip()) > 200:
        text = methods_match.group(1)
    else:
        chars = len(text)
        text = text[chars // 10: chars // 10 + 15000]

    words = text.split()
    if len(words) > 1500:
        text = " ".join(words[:1500])

    text = text.replace("\n", " ")

    return text


def parse_json_response(raw: str) -> dict:
    cleaned = re.sub(r"```(?:json)?", "", raw).strip()
    cleaned = re.sub(r"```", "", cleaned).strip()
    brace_start = cleaned.find("{")
    if brace_start > 0:
        cleaned = cleaned[brace_start:]

    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        pass

    try:
        open_braces   = cleaned.count("{") - cleaned.count("}")
        open_brackets = cleaned.count("[") - cleaned.count("]")
        repaired      = cleaned.rstrip(",\n ")
        repaired     += "]" * max(open_brackets, 0)
        repaired     += "}" * max(open_braces, 0)
        return json.loads(repaired)
    except json.JSONDecodeError:
        pass

    match = re.search(r"\{.*\}", cleaned, re.DOTALL)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass

    raise ValueError("No valid JSON found in response.\n" + raw[:500])


def call_groq_with_backoff(client: Groq, messages: list) -> str:
    word_limits = [2500, 2000, 1500, 1000]
    attempt     = 0
    while attempt < len(word_limits):
        try:
            response = client.chat.completions.create(
                model=GROQ_MODEL,
                messages=messages,
                temperature=0.1,
                max_tokens=2000,
            )
            return response.choices[0].message.content
        except Exception as e:
            err = str(e)
            if "413" in err or "too large" in err.lower() or "rate_limit_exceeded" in err.lower():
                attempt += 1
                if attempt >= len(word_limits):
                    if "429" in err:

                        retry_after = None
        
                        try:
                            # Try to read headers if available
                            if hasattr(e, "response") and e.response:
        
                                headers = getattr(e.response, "headers", {})
        
                                retry_after = headers.get("retry-after")
        
                        except Exception:
                            pass
        
                        if retry_after:
                            raise ValueError(
                                f"Rate limit exceeded. Retry after {retry_after} seconds."
                            )
        
                        else:
                            raise ValueError(
                                "Rate limit exceeded. Retry later."
                            )
                    raise ValueError("Request too large even after maximum reductions.")
                new_limit = word_limits[attempt]
                for msg in messages:
                    if msg["role"] == "user":
                        words          = msg["content"].split()
                        msg["content"] = " ".join(words[:new_limit])
            else:
                raise


def extract_protocol(client: Groq, pdf_path: str, filename: str) -> dict:
    paper_text = extract_text_from_pdf(pdf_path)
    messages = [
        {"role": "system", "content": EXTRACTION_SYSTEM},
        {"role": "user", "content":
            "Extract the complete experimental protocol from this research paper. "
            "Be thorough - capture all concentrations, timings, temperatures, cell lines, "
            "reagents, antibodies, and equipment settings.\n\nPAPER TEXT:\n" + paper_text
        }
    ]
    raw = call_groq_with_backoff(client, messages)

    try:
        result = parse_json_response(raw)
    except Exception as e:
        raise ValueError(f"Invalid JSON returned by LLM: {str(e)}")

    result["_filename"] = filename
    return result


def compare_protocols(client: Groq, extracted: list, paper_names: list) -> dict:
    sections = []
    for i, (e, name) in enumerate(zip(extracted, paper_names)):
        sections.append(
            f"=== Paper {i+1}: {e.get('title', name)} ===\n"
            f"{json.dumps(e.get('parameters', {}), indent=2)}"
        )
    prompt   = (
        f"Compare these {len(extracted)} research paper protocols "
        f"and identify ALL methodological contradictions:\n\n"
        + "\n\n".join(sections)
    )
    messages = [
        {"role": "system", "content": COMPARISON_SYSTEM},
        {"role": "user",   "content": prompt},
    ]
    raw    = call_groq_with_backoff(client, messages)
    return parse_json_response(raw)


def generate_html_report(extracted: list, comparison: dict, paper_names: list) -> str:
    summary  = comparison.get("summary", {})
    contras  = comparison.get("contradictions", [])
    ranked   = comparison.get("ranked_issues", [])
    optimal  = comparison.get("optimal_protocol", {})

    sev_color = {"high": "#ff4d6d", "medium": "#ff9f1c", "low": "#2ec4b6"}
    sev_bg    = {"high": "#fff0f3", "medium": "#fff8ee", "low": "#f0fafa"}

    def badge(sev):
        col = sev_color.get(sev, "#888")
        return (f'<span style="background:{col}20;color:{col};border-radius:4px;'
                f'padding:2px 8px;font-size:11px;font-weight:600;text-transform:uppercase;">{sev}</span>')

    extraction_rows = ""
    for i, e in enumerate(extracted):
        cats  = e.get("parameters", {})
        total = sum(len(v) for v in cats.values() if isinstance(v, list))
        extraction_rows += f"""
        <tr>
          <td><b>Paper {i+1}</b><br><span style="color:#888;font-size:12px;">{e.get('_filename','')}</span></td>
          <td>{e.get('title','—')}</td>
          <td>{e.get('authors','—')}</td>
          <td>{e.get('year','—')}</td>
          <td>{e.get('protocol_type','—')}</td>
          <td style="color:#7b2d8b;font-weight:600;">{total}</td>
        </tr>"""

    contra_rows = ""
    for c in contras:
        sev  = c.get("severity", "low")
        col  = sev_color.get(sev, "#888")
        vals = c.get("values", {})
        paper_cols = "".join(
            f'<td style="font-family:monospace;font-size:12px;">{vals.get(f"paper_{i}","-")}</td>'
            for i in range(len(paper_names))
        )
        contra_rows += f"""
        <tr>
          <td><b>{c.get('parameter','')}</b></td>
          <td style="color:#888;font-size:12px;">{c.get('category','')}</td>
          <td>{badge(sev)}</td>
          {paper_cols}
          <td style="font-size:12px;color:#444;line-height:1.5;">{c.get('explanation','')}</td>
        </tr>"""

    ranked_html = ""
    for item in ranked:
        sev = item.get("severity", "low")
        col = sev_color.get(sev, "#888")
        bg  = sev_bg.get(sev, "#fafafa")
        ranked_html += f"""
        <div style="display:flex;gap:14px;background:{bg};border:1px solid {col}30;
                    border-radius:8px;padding:14px 16px;margin-bottom:8px;">
          <div style="font-size:24px;font-weight:800;color:{col};min-width:34px;">#{item.get('rank','?')}</div>
          <div>
            <div style="font-size:14px;font-weight:600;">{item.get('parameter','')} {badge(sev)}</div>
            <div style="font-size:12px;color:#555;margin-top:4px;line-height:1.5;">{item.get('brief','')}</div>
          </div>
        </div>"""

    optimal_html = ""
    for p in optimal.get("parameters", []):
        optimal_html += f"""
        <div style="border-left:3px solid #2ec4b6;padding-left:12px;margin-bottom:14px;">
          <div style="font-size:11px;text-transform:uppercase;color:#aaa;">{p.get('label','')}</div>
          <div style="font-size:14px;font-weight:600;color:#1a1a2e;margin:3px 0;">{p.get('value','')}</div>
          <div style="font-size:12px;color:#666;line-height:1.5;">{p.get('reason','')}</div>
        </div>"""

    paper_th = "".join(
        f'<th>Paper {i+1}<br><span style="font-weight:400;font-size:11px;">{n}</span></th>'
        for i, n in enumerate(paper_names)
    )

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width,initial-scale=1.0">
  <title>ContraGenAI - Protocol Contradiction Report</title>
  <style>
    *{{box-sizing:border-box;margin:0;padding:0}}
    body{{font-family:'Segoe UI',sans-serif;background:#f7f7fa;color:#1a1a2e;padding:40px 24px}}
    .container{{max-width:1100px;margin:0 auto}}
    .header{{background:#1a1a2e;color:white;border-radius:12px;padding:32px 36px;margin-bottom:28px}}
    .header h1{{font-size:28px;font-weight:800;margin-bottom:6px}}
    .header p{{color:#aab;font-size:13px}}
    .section{{background:white;border-radius:12px;padding:24px 28px;margin-bottom:20px;box-shadow:0 1px 4px rgba(0,0,0,.07)}}
    .section-title{{font-size:16px;font-weight:700;border-left:4px solid #7b2d8b;padding-left:12px;margin-bottom:18px}}
    .stats{{display:flex;gap:14px;flex-wrap:wrap}}
    .stat{{border-radius:10px;padding:16px 22px;min-width:120px}}
    table{{width:100%;border-collapse:collapse;font-size:13px}}
    th{{background:#f5f0ff;font-size:11px;text-transform:uppercase;letter-spacing:.06em;padding:10px 14px;text-align:left}}
    td{{padding:10px 14px;border-bottom:1px solid #f0f0f5;vertical-align:top}}
    tr:hover td{{background:#fafafa}}
    .optimal-grid{{display:grid;grid-template-columns:repeat(auto-fill,minmax(240px,1fr));gap:14px;margin-top:14px}}
    .footer{{text-align:center;font-size:12px;color:#aaa;margin-top:32px}}
  </style>
</head>
<body><div class="container">
  <div class="header">
    <h1>Protocol Contradiction Report</h1>
    <p>{len(paper_names)} papers compared &nbsp;|&nbsp; Protocol: {comparison.get('protocol_type','Unknown')}</p>
  </div>
  <div class="section">
    <div class="section-title">Summary</div>
    <div class="stats">
      <div class="stat" style="background:#f5f0ff;border-left:4px solid #7b2d8b">
        <div style="font-size:36px;font-weight:800;color:#7b2d8b">{summary.get('total_contradictions',0)}</div>
        <div style="font-size:11px;color:#888;text-transform:uppercase">Total</div>
      </div>
      <div class="stat" style="background:#fff0f3;border-left:4px solid #ff4d6d">
        <div style="font-size:36px;font-weight:800;color:#ff4d6d">{summary.get('high',0)}</div>
        <div style="font-size:11px;color:#888;text-transform:uppercase">High</div>
      </div>
      <div class="stat" style="background:#fff8ee;border-left:4px solid #ff9f1c">
        <div style="font-size:36px;font-weight:800;color:#ff9f1c">{summary.get('medium',0)}</div>
        <div style="font-size:11px;color:#888;text-transform:uppercase">Medium</div>
      </div>
      <div class="stat" style="background:#f0fafa;border-left:4px solid #2ec4b6">
        <div style="font-size:36px;font-weight:800;color:#2ec4b6">{summary.get('low',0)}</div>
        <div style="font-size:11px;color:#888;text-transform:uppercase">Low</div>
      </div>
    </div>
  </div>
  <div class="section">
    <div class="section-title">Extracted Parameters</div>
    <table><thead><tr><th>Paper</th><th>Title</th><th>Authors</th><th>Year</th><th>Protocol Type</th><th>Parameters</th></tr></thead>
    <tbody>{extraction_rows}</tbody></table>
  </div>
  <div class="section">
    <div class="section-title">Ranked by Severity</div>
    {ranked_html}
  </div>
  <div class="section">
    <div class="section-title">Side-by-side Comparison</div>
    <div style="overflow-x:auto">
      <table><thead><tr><th>Parameter</th><th>Category</th><th>Severity</th>{paper_th}<th>Impact on Reproducibility</th></tr></thead>
      <tbody>{contra_rows}</tbody></table>
    </div>
  </div>
  <div class="section">
    <div class="section-title">Recommended Optimal Protocol</div>
    <p style="font-size:13px;color:#444;line-height:1.6">{optimal.get('rationale','')}</p>
    <div class="optimal-grid">{optimal_html}</div>
  </div>
  <div class="footer">ContraGenAI &nbsp;·&nbsp; Protocol Contradiction Detector &nbsp;·&nbsp; Powered by Groq + LLaMA 3.3 70B</div>
</div></body></html>"""


# ── API Routes ────────────────────────────────────────────────────────────────

@app.get("/health")
def health():
    return {"status": "ok", "model": GROQ_MODEL}


@app.post("/analyze")
async def analyze(files: List[UploadFile] = File(...)):
    if len(files) < 2:
        raise HTTPException(status_code=400, detail="Please upload at least 2 PDF files.")

    api_key = GROQ_API_KEY
    if not api_key:
        raise HTTPException(status_code=500, detail="GROQ_API_KEY not set in environment.")

    client = Groq(api_key=api_key)

    extracted   = []
    paper_names = []
    tmp_paths   = []

    try:
        # Save uploads to temp files
        for f in files:
            if not f.filename.lower().endswith(".pdf"):
                raise HTTPException(status_code=400, detail=f"{f.filename} is not a PDF.")
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
            tmp.write(await f.read())
            tmp.close()
            tmp_paths.append(tmp.name)
            paper_names.append(f.filename)

        # Extract each paper
        for path, name in zip(tmp_paths, paper_names):
         try:
            result = extract_protocol(client, path, name)
            extracted.append(result)
         except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Failed to process {name}: {str(e)}"
            )

        # Compare
        comparison = compare_protocols(client, extracted, paper_names)

        # Build HTML report
        html_report = generate_html_report(extracted, comparison, paper_names)

        return JSONResponse({
            "extracted":   extracted,
            "comparison":  comparison,
            "html_report": html_report,
            "paper_names": paper_names,
        })

    finally:
        for p in tmp_paths:
            try:
                os.unlink(p)
            except Exception:
                pass