#!/usr/bin/env python3
"""Comprehensive quality scoring across all paper formats + MCP tools."""
import re
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent / "src"))

from paper2md import convert

TEST_PAPERS = Path("/Volumes/EXTERNAL_USB/test_papers")
TEST_PDF = Path("/Volumes/EXTERNAL_USB/2603.04390v1.pdf")

# Paper configs: name -> (path, expected_title_substr, expected_author, min_refs, format_type)
PAPERS = {
    "attention": (TEST_PAPERS / "attention.pdf", "Attention Is All You Need", "Vaswani", 20, "arXiv/NeurIPS"),
    "resnet": (TEST_PAPERS / "resnet.pdf", "Deep Residual Learning", "He", 20, "arXiv/CVPR"),
    "bert": (TEST_PAPERS / "bert.pdf", "BERT", "Devlin", 10, "arXiv/NAACL"),
    "gpt3": (TEST_PAPERS / "gpt3.pdf", "Few-Shot Learners", "Brown", 50, "arXiv/NeurIPS"),
    "ddpm": (TEST_PAPERS / "ddpm.pdf", "Denoising Diffusion", "Ho", 30, "arXiv/NeurIPS"),
    "vit": (TEST_PAPERS / "vit.pdf", "16x16 Words", "Dosovitskiy", 40, "ICLR"),
    "llama": (TEST_PAPERS / "llama.pdf", "LLaMA", "Touvron", 50, "Meta AI"),
    "ldm": (TEST_PAPERS / "ldm.pdf", "Latent Diffusion", "Rombach", 80, "CVPR"),
    "kan_survey": (TEST_PAPERS / "kan_survey.pdf", "Kolmogorov-Arnold Network", "Somvanshi", 50, "ACM"),
    "mamba": (TEST_PAPERS / "mamba.pdf", "Mamba", "Gu", 2, "arXiv"),
    "batch_norm": (TEST_PAPERS / "batch_norm.pdf", "Batch Normalization", "Ioffe", 30, "arXiv/ICML"),
    "dropout": (TEST_PAPERS / "dropout.pdf", "neural networks", "Hinton", 15, "arXiv"),
    "yolo_ieee": (TEST_PAPERS / "yolo_ieee.pdf", "You Only Look Once", "Redmon", 30, "IEEE/CVPR"),
    "webgis": (TEST_PDF, "Dual-Helix", None, 30, "Custom"),
    "lora": (TEST_PAPERS / "lora.pdf", "Low-Rank Adaptation", "Hu", 50, "ICLR"),
    "chinchilla": (TEST_PAPERS / "instructgpt.pdf", "Compute-Optimal", "Hoffmann", 50, "DeepMind"),
    "llava": (TEST_PAPERS / "llava.pdf", "Visual Instruction", "Liu", 30, "NeurIPS"),
    "codex": (TEST_PAPERS / "alphafold2.pdf", "Evaluating Large Language", "Chen", 50, "OpenAI"),
    "sklearn": (TEST_PAPERS / "jmlr_sklearn.pdf", "Scikit-learn", "Pedregosa", 10, "JMLR"),
    "adam": (TEST_PAPERS / "pmlr_adam.pdf", "Adam", "Kingma", 15, "ICLR"),
    "unet": (TEST_PAPERS / "elsevier_unet.pdf", "U-Net", "Ronneberger", 10, "Springer/LNCS"),
    "gan": (TEST_PAPERS / "neurips_gan.pdf", "Generative Adversarial", "Goodfellow", 20, "NeurIPS"),
}

SECTIONS_EXPECTED = {
    "attention": ["introduction", "background", "model architecture"],
    "resnet": ["introduction"],
    "bert": ["introduction"],
    "gpt3": ["introduction"],
    "ddpm": ["introduction", "background"],
    "vit": ["introduction", "related work"],
    "llama": ["introduction", "approach"],
    "ldm": ["introduction", "related work"],
    "kan_survey": ["introduction"],
    "mamba": ["introduction"],
    "batch_norm": ["introduction"],
    "dropout": [],
    "yolo_ieee": ["introduction"],
    "webgis": ["introduction"],
    "lora": ["introduction"],
    "chinchilla": ["introduction"],
    "llava": ["introduction", "related work"],
    "codex": ["evaluation framework"],
    "sklearn": [],
    "adam": ["introduction"],
    "unet": ["introduction"],
    "gan": ["introduction"],
}


def score_paper(name, path, title_substr, author_name, min_refs, fmt):
    """Score a single paper. Returns (scores_dict, details_dict)."""
    scores = {}
    details = {}

    try:
        t0 = time.time()
        r = convert(path)
        elapsed = time.time() - t0
        details["time"] = f"{elapsed:.1f}s"
    except Exception as e:
        return {"error": 100}, {"error": str(e)}

    # 1. Title (0 or 100)
    if title_substr.lower() in r.title.lower():
        scores["title"] = 100
    else:
        scores["title"] = 0
        details["title"] = f"Expected '{title_substr}', got '{r.title[:80]}'"

    # 2. Abstract (0-100)
    if r.abstract and len(r.abstract) > 50:
        scores["abstract"] = 100
    elif r.abstract and len(r.abstract) > 10:
        scores["abstract"] = 60
        details["abstract"] = f"Short abstract ({len(r.abstract)} chars)"
    else:
        scores["abstract"] = 0
        details["abstract"] = "No abstract detected"

    # 3. Authors (0-100)
    if author_name is None:
        if len(r.authors) >= 2:
            scores["authors"] = 100
        elif len(r.authors) >= 1:
            scores["authors"] = 50
        else:
            scores["authors"] = 0
    else:
        has_author = any(author_name in a for a in r.authors)
        if has_author and len(r.authors) >= 2:
            scores["authors"] = 100
        elif has_author:
            scores["authors"] = 70
        elif len(r.authors) >= 2:
            scores["authors"] = 50
            details["authors"] = f"Missing '{author_name}', got: {r.authors[:5]}"
        else:
            scores["authors"] = 0
            details["authors"] = f"No authors found, expected '{author_name}'"

    # 4. Sections (0-100)
    h2s = " ".join(l for l in r.markdown.split("\n") if l.startswith("## ")).lower()
    expected = SECTIONS_EXPECTED.get(name, [])
    if expected:
        found = sum(1 for s in expected if s in h2s)
        scores["sections"] = int(100 * found / len(expected))
        if found < len(expected):
            missing = [s for s in expected if s not in h2s]
            details["sections"] = f"Missing: {missing}"
    else:
        # No specific sections expected, check we have at least some
        h2_count = len([l for l in r.markdown.split("\n") if l.startswith("## ")])
        scores["sections"] = 100 if h2_count >= 3 else (70 if h2_count >= 1 else 0)

    # 5. References (0-100)
    ref_count = len(r.references)
    if ref_count >= min_refs:
        scores["references"] = 100
    elif ref_count >= min_refs * 0.5:
        scores["references"] = 70
        details["references"] = f"{ref_count}/{min_refs} refs"
    else:
        scores["references"] = 0
        details["references"] = f"Only {ref_count}/{min_refs} refs"

    # 6. No garbled chars (0-100)
    garbled = sum(1 for ch in r.markdown if ord(ch) < 0x20 and ch not in '\n\t\r')
    if garbled == 0:
        scores["no_garbled"] = 100
    elif garbled < 5:
        scores["no_garbled"] = 80
        details["no_garbled"] = f"{garbled} garbled chars"
    else:
        scores["no_garbled"] = 0
        details["no_garbled"] = f"{garbled} garbled chars"

    # 6b. No footnote leaks (footnotes should not appear in markdown)
    footnote_leaks = len(re.findall(r'^\[\^\]:', r.markdown, re.MULTILINE))
    if footnote_leaks == 0:
        scores["no_footnote_leaks"] = 100
    elif footnote_leaks < 5:
        scores["no_footnote_leaks"] = 70
        details["no_footnote_leaks"] = f"{footnote_leaks} footnote leaks"
    else:
        scores["no_footnote_leaks"] = 0
        details["no_footnote_leaks"] = f"{footnote_leaks} footnote leaks"

    # 7. Math quality (where applicable)
    if name in ("ddpm", "batch_norm", "attention"):
        if "$" in r.markdown:
            all_math = re.findall(r"\$([^$]+)\$", r.markdown)
            substantive = [m for m in all_math if len(m.strip()) > 3]
            latex_cmds = sum(1 for m in all_math if "\\" in m)

            if name == "ddpm":
                # DDPM should have many math expressions
                if len(substantive) >= 20 and latex_cmds >= 10:
                    scores["math"] = 100
                elif len(substantive) >= 10:
                    scores["math"] = 70
                    details["math"] = f"{len(substantive)} substantive, {latex_cmds} with LaTeX"
                else:
                    scores["math"] = 40
                    details["math"] = f"Only {len(substantive)} substantive math"
            elif name == "batch_norm":
                # Check ^{(k)} pattern
                bad = len(re.findall(r'\^\{\(\}_\{', r.markdown))
                good = len(re.findall(r'\^\{\([a-zA-Z0-9]+\)\}', r.markdown))
                if bad == 0 and good >= 5:
                    scores["math"] = 100
                elif bad == 0:
                    scores["math"] = 80
                    details["math"] = f"No bad patterns, {good} good patterns"
                else:
                    scores["math"] = 40
                    details["math"] = f"{bad} broken superscripts"
            elif name == "attention":
                # Check d_{model} and Attention equation
                has_dmodel = "d_{\\text{model}}" in r.markdown
                display_eqs = re.findall(r'\$\$(.+?)\$\$', r.markdown)
                attn_eqs = [eq for eq in display_eqs if 'Attention' in eq and 'softmax' in eq]
                if has_dmodel and attn_eqs:
                    scores["math"] = 100
                elif has_dmodel or attn_eqs:
                    scores["math"] = 70
                    details["math"] = f"d_model: {has_dmodel}, attn_eq: {bool(attn_eqs)}"
                else:
                    scores["math"] = 40
        else:
            scores["math"] = 0
            details["math"] = "No math found"

    # 8. Tables (where applicable)
    if name in ("attention", "resnet", "gpt3", "llama", "ldm", "vit", "kan_survey"):
        if len(r.tables) >= 3:
            scores["tables"] = 100
        elif len(r.tables) >= 1:
            scores["tables"] = 60
            details["tables"] = f"Only {len(r.tables)} tables"
        else:
            scores["tables"] = 0

    # 9. Figures (where applicable)
    if name in ("resnet", "ldm", "webgis"):
        if len(r.figures) >= 3:
            scores["figures"] = 100
        elif len(r.figures) >= 1:
            scores["figures"] = 60
        else:
            scores["figures"] = 0

    # 10. Section ordering (for papers with clear structure)
    if name in ("attention", "bert", "ddpm"):
        lines = r.markdown.split("\n")
        positions = {}
        for i, line in enumerate(lines):
            if line.startswith("## "):
                lower = line.lower()
                if "introduction" in lower and "introduction" not in positions:
                    positions["introduction"] = i
                elif "conclusion" in lower and "conclusion" not in positions:
                    positions["conclusion"] = i
        if "introduction" in positions:
            if "conclusion" in positions:
                if positions["introduction"] < positions["conclusion"]:
                    scores["ordering"] = 100
                else:
                    scores["ordering"] = 0
                    details["ordering"] = "Introduction after Conclusion"
            else:
                scores["ordering"] = 80  # No conclusion found
        else:
            scores["ordering"] = 0
            details["ordering"] = "No Introduction found"

    # 11. Dehyphenation (for multi-column papers)
    if name in ("attention", "vit", "batch_norm"):
        broken = re.findall(r"[a-z]- [a-z]", r.markdown)
        if len(broken) < 3:
            scores["dehyphenation"] = 100
        elif len(broken) < 10:
            scores["dehyphenation"] = 70
        else:
            scores["dehyphenation"] = 40
            details["dehyphenation"] = f"{len(broken)} broken hyphens"

    return scores, details


def score_mcp():
    """Score MCP tool output consistency."""
    scores = {}
    details = {}

    try:
        from paper2md.server import create_server
    except ImportError:
        return {"mcp_available": 0}, {"mcp_available": "fastmcp not installed"}

    try:
        server = create_server()
    except ImportError as e:
        return {"mcp_available": 0}, {"mcp_available": str(e)}

    scores["mcp_available"] = 100

    # Test convert_pdf tool
    test_paper = TEST_PAPERS / "attention.pdf"
    if test_paper.exists():
        from paper2md import convert as direct_convert
        direct_result = direct_convert(test_paper)

        # Get convert_pdf result
        from paper2md.core import convert
        from paper2md.types import PipelineConfig
        result = convert(test_paper, PipelineConfig())

        # Check structured result
        structured = {
            "title": result.title,
            "authors": result.authors,
            "abstract": result.abstract,
            "page_count": result.page_count,
            "figure_count": len(result.figures),
            "table_count": len(result.tables),
            "reference_count": len(result.references),
        }

        # Verify title matches
        if structured["title"] == direct_result.title:
            scores["mcp_title_match"] = 100
        else:
            scores["mcp_title_match"] = 0
            details["mcp_title_match"] = f"Mismatch: {structured['title']}"

        # Verify authors match
        if structured["authors"] == direct_result.authors:
            scores["mcp_authors_match"] = 100
        else:
            scores["mcp_authors_match"] = 0
            details["mcp_authors_match"] = f"Mismatch"

        # Verify reference count
        if structured["reference_count"] == len(direct_result.references):
            scores["mcp_refs_match"] = 100
        else:
            scores["mcp_refs_match"] = 0

        # Check extract_metadata sections
        sections = []
        for line in result.markdown.split("\n"):
            stripped = line.strip()
            if stripped.startswith("## "):
                sections.append({"level": 2, "title": stripped[3:].strip()})
            elif stripped.startswith("### "):
                sections.append({"level": 3, "title": stripped[4:].strip()})

        if len(sections) >= 3:
            scores["mcp_sections"] = 100
        else:
            scores["mcp_sections"] = 0
            details["mcp_sections"] = f"Only {len(sections)} sections"

        # Check figure captions
        captions = [f.caption for f in result.figures if f.caption]
        if len(captions) >= 1:
            scores["mcp_figures"] = 100
        else:
            scores["mcp_figures"] = 60

    return scores, details


def main():
    print("=" * 90)
    print("PAPER2MD COMPREHENSIVE QUALITY REPORT")
    print("=" * 90)

    all_scores = {}
    all_details = {}
    dimension_scores = {}  # dimension -> list of scores

    # Score each paper
    for name, (path, title, author, min_refs, fmt) in PAPERS.items():
        if not path.exists():
            print(f"\n  SKIP {name} ({fmt}) - file not found")
            continue

        print(f"\n  Scoring {name} ({fmt})...", end="", flush=True)
        scores, details = score_paper(name, path, title, author, min_refs, fmt)
        all_scores[name] = scores
        all_details[name] = details

        # Aggregate by dimension
        for dim, score in scores.items():
            if dim not in dimension_scores:
                dimension_scores[dim] = []
            dimension_scores[dim].append((name, score))

        avg = sum(scores.values()) / len(scores) if scores else 0
        print(f" {avg:.0f}/100 avg", end="")
        if details.get("time"):
            print(f" ({details['time']})", end="")

    # Score MCP
    print(f"\n\n  Scoring MCP tools...", end="", flush=True)
    mcp_scores, mcp_details = score_mcp()
    all_scores["MCP"] = mcp_scores
    all_details["MCP"] = mcp_details
    mcp_avg = sum(mcp_scores.values()) / len(mcp_scores) if mcp_scores else 0
    print(f" {mcp_avg:.0f}/100 avg")

    # ============================
    # Per-dimension report
    # ============================
    print("\n" + "=" * 90)
    print("SCORES BY DIMENSION")
    print("=" * 90)

    core_dims = ["title", "abstract", "authors", "sections", "references", "no_garbled", "no_footnote_leaks"]
    optional_dims = ["math", "tables", "figures", "ordering", "dehyphenation"]

    for dim in core_dims + optional_dims:
        if dim not in dimension_scores:
            continue
        entries = dimension_scores[dim]
        avg = sum(s for _, s in entries) / len(entries)
        passed = sum(1 for _, s in entries if s >= 90)
        total = len(entries)
        failures = [(n, s) for n, s in entries if s < 90]

        status = "PASS" if avg >= 90 else "NEEDS WORK"
        label = f"  {dim:<16s}"
        print(f"{label} {avg:5.1f}/100  ({passed}/{total} papers ≥90)  [{status}]")
        for n, s in failures:
            detail = all_details.get(n, {}).get(dim, "")
            print(f"    {'↳':>4s} {n}: {s}/100  {detail}")

    # ============================
    # Per-paper report
    # ============================
    print("\n" + "=" * 90)
    print("SCORES BY PAPER")
    print("=" * 90)

    for name in sorted(all_scores.keys()):
        if name == "MCP":
            continue
        scores = all_scores[name]
        details = all_details[name]
        fmt = PAPERS[name][4] if name in PAPERS else ""
        avg = sum(scores.values()) / len(scores) if scores else 0
        status = "PASS" if avg >= 90 else "NEEDS WORK"
        print(f"\n  {name} ({fmt})  — {avg:.0f}/100 [{status}]")
        for dim, score in sorted(scores.items()):
            marker = "✓" if score >= 90 else "✗"
            detail = details.get(dim, "")
            print(f"    {marker} {dim:<16s} {score:3d}/100  {detail}")

    # ============================
    # MCP report
    # ============================
    print("\n" + "=" * 90)
    print("MCP TOOL QUALITY")
    print("=" * 90)
    for dim, score in sorted(mcp_scores.items()):
        marker = "✓" if score >= 90 else "✗"
        detail = mcp_details.get(dim, "")
        print(f"  {marker} {dim:<20s} {score:3d}/100  {detail}")

    # ============================
    # Overall summary
    # ============================
    all_dim_avgs = []
    for dim in core_dims:
        if dim in dimension_scores:
            entries = dimension_scores[dim]
            avg = sum(s for _, s in entries) / len(entries)
            all_dim_avgs.append((dim, avg))

    overall = sum(a for _, a in all_dim_avgs) / len(all_dim_avgs) if all_dim_avgs else 0

    print("\n" + "=" * 90)
    print("OVERALL SUMMARY")
    print("=" * 90)
    for dim, avg in all_dim_avgs:
        bar = "█" * int(avg / 5) + "░" * (20 - int(avg / 5))
        print(f"  {dim:<16s} {bar} {avg:.1f}")

    mcp_overall = sum(mcp_scores.values()) / len(mcp_scores) if mcp_scores else 0
    bar = "█" * int(mcp_overall / 5) + "░" * (20 - int(mcp_overall / 5))
    print(f"  {'MCP':<16s} {bar} {mcp_overall:.1f}")

    print(f"\n  OVERALL CORE SCORE: {overall:.1f}/100")
    print(f"  MCP SCORE: {mcp_overall:.1f}/100")

    below_90 = [(dim, avg) for dim, avg in all_dim_avgs if avg < 90]
    if below_90:
        print(f"\n  DIMENSIONS BELOW 90:")
        for dim, avg in below_90:
            print(f"    - {dim}: {avg:.1f}")
    else:
        print(f"\n  ALL CORE DIMENSIONS ≥ 90 ✓")

    return 0 if overall >= 90 and mcp_overall >= 90 else 1


if __name__ == "__main__":
    sys.exit(main())
