paper-doc/backend/scripts/seed.py

"""Seed the template-field library and a couple of starter templates.

Idempotent: fields are matched by name and templates by name, so running it
twice adds nothing. ``--reset`` empties the three tables first, which is the
way to get back to the seeded baseline after experimenting.

Usage (from ``backend/``)::

    .venv/bin/python scripts/seed.py
    .venv/bin/python scripts/seed.py --reset
"""

from __future__ import annotations

import argparse
import sys
from decimal import Decimal
from pathlib import Path

# Running this as a plain script puts scripts/ on sys.path, not backend/, so
# `import app` would fail. Anchor to the backend directory instead.
BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

from sqlalchemy import delete, select  # noqa: E402
from sqlalchemy.orm import Session  # noqa: E402

from app.db.session import SessionLocal  # noqa: E402
from app.models import Template, TemplateFieldLibrary, TemplateField  # noqa: E402

# --- the field library -------------------------------------------------------
#
# (name, level, font_size, font_color)
#
# The number in the name is written by hand and stored verbatim — nothing in
# the system derives or rewrites it. Level is only a rendering hint: it decides
# indentation in the outline, never parentage.
#
# Colours: headings are near-black, running text is black, and the two
# unnumbered front/back sections (Abstract, References) are grey so they read
# as apparatus rather than as body sections.
BODY = "#000000"
HEADING = "#1F1F1F"
SUB = "#404040"
GREY = "#595959"

FIELDS: list[tuple[str, int, Decimal, str]] = [
    ("0 Abstract", 1, Decimal("10.5"), GREY),
    ("1 Introduction", 1, Decimal("12.0"), HEADING),
    ("1.1 Research Background", 2, Decimal("10.5"), SUB),
    ("1.2 Problem Statement", 2, Decimal("10.5"), SUB),
    ("1.3 Contributions", 2, Decimal("10.5"), SUB),
    ("2 Related Work", 1, Decimal("12.0"), HEADING),
    ("2.1 Prior Approaches", 2, Decimal("10.5"), SUB),
    ("2.2 Limitations of Existing Work", 2, Decimal("10.5"), SUB),
    ("3 Method", 1, Decimal("12.0"), HEADING),
    ("3.1 Problem Formulation", 2, Decimal("10.5"), SUB),
    ("3.2 Framework Overview", 2, Decimal("10.5"), SUB),
    ("3.3 Implementation Details", 2, Decimal("10.5"), SUB),
    ("4 Experiments", 1, Decimal("12.0"), HEADING),
    ("4.1 Datasets", 2, Decimal("10.5"), SUB),
    ("4.2 Experimental Setup", 2, Decimal("10.5"), SUB),
    ("4.3 Main Results", 2, Decimal("10.5"), SUB),
    ("4.4 Ablation Study", 2, Decimal("10.5"), SUB),
    ("5 Discussion", 1, Decimal("12.0"), BODY),
    ("6 Conclusion", 1, Decimal("12.0"), BODY),
    ("7 References", 1, Decimal("10.5"), GREY),
]

# --- starter templates -------------------------------------------------------
#
# Each entry is (template name, abstract, [field names in display order]).
# `sort` is assigned from the list position (1, 2, 3 ...), which is exactly
# what the UI does when the user picks fields and orders them.
TEMPLATES: list[tuple[str, str, list[str]]] = [
    (
        "标准学术论文（通用）",
        "完整的通用学术论文骨架，含摘要、引言、相关工作、方法、实验、讨论与结论。"
        "适合期刊或会议长文；先按此模板把结构填满，再按目标期刊微调字段。",
        [name for name, *_ in FIELDS],
    ),
    (
        "四段式短文",
        "紧凑的四段式结构，只保留摘要、引言、方法、实验与结论。"
        "适合短文、技术报告或初稿阶段的结构搭建。",
        [
            "0 Abstract",
            "1 Introduction",
            "3 Method",
            "4 Experiments",
            "6 Conclusion",
        ],
    ),
    (
        "方法创新型论文",
        "以方法贡献为主线的结构，弱化相关工作、强化方法细节与消融实验。"
        "适合以新框架、新算法为主要贡献的投稿。",
        [
            "0 Abstract",
            "1 Introduction",
            "1.2 Problem Statement",
            "1.3 Contributions",
            "3 Method",
            "3.1 Problem Formulation",
            "3.2 Framework Overview",
            "3.3 Implementation Details",
            "4 Experiments",
            "4.2 Experimental Setup",
            "4.3 Main Results",
            "4.4 Ablation Study",
            "6 Conclusion",
        ],
    ),
]


def reset(db: Session) -> None:
    """Empty the three tables in dependency order."""
    db.execute(delete(TemplateField))
    db.execute(delete(Template))
    db.execute(delete(TemplateFieldLibrary))
    db.commit()


def seed_fields(db: Session) -> dict[str, TemplateFieldLibrary]:
    """Insert any missing library fields and return name -> field."""
    existing = {field.name: field for field in db.scalars(select(TemplateFieldLibrary)).all()}

    created = 0
    for name, level, font_size, font_color in FIELDS:
        if name in existing:
            continue
        field = TemplateFieldLibrary(
            name=name, level=level, font_size=font_size, font_color=font_color
        )
        db.add(field)
        existing[name] = field
        created += 1

    db.commit()
    for field in existing.values():
        db.refresh(field)

    print(
        f"  template_field_library : {created} created, "
        f"{len(existing) - created} already present"
    )
    return existing


def seed_templates(db: Session, fields: dict[str, TemplateFieldLibrary]) -> None:
    """Insert any missing starter templates."""
    existing = {name for name in db.scalars(select(Template.name)).all()}

    created = 0
    for name, abstract, field_names in TEMPLATES:
        if name in existing:
            continue

        missing = [field_name for field_name in field_names if field_name not in fields]
        if missing:
            raise SystemExit(f"template {name!r} references unknown fields: {missing}")

        template = Template(name=name, abstract=abstract)
        # sort is the list position: 1-based, ascending, with gaps allowed
        # because the UI lets the user type any integer.
        template.items = [
            TemplateField(field_id=fields[field_name].id, sort=index)
            for index, field_name in enumerate(field_names, start=1)
        ]
        db.add(template)
        created += 1

    db.commit()
    print(f"  template: {created} created, {len(existing)} already present")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--reset",
        action="store_true",
        help="delete all fields and templates before seeding",
    )
    args = parser.parse_args()

    with SessionLocal() as db:
        if args.reset:
            reset(db)
            print("  reset         : all rows deleted")

        fields = seed_fields(db)
        seed_templates(db, fields)

    print("seed complete")


if __name__ == "__main__":
    main()