app.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

import gradio as gr

from tools import (
    DEFAULT_INPUT_CSV,
    HF_MODEL,
    OUTPUT_DIR,
    compare_with_taxonomy,
    consolidate_into_themes,
    export_narrative,
    generate_comparison_csv,
    label_topics_with_llm,
    load_scopus_csv,
    run_bertopic_discovery,
    run_full_pipeline,
)


CUSTOM_CSS = """
.status-ok { color: #1b5e20; font-weight: 600; }
.status-note { color: #37474f; }
"""


def _resolve_output_dir(value: str) -> Path:
    return Path(value) if value else OUTPUT_DIR


def ui_load(file_path: str, output_dir: str) -> str:
    stats = load_scopus_csv(file_path, _resolve_output_dir(output_dir))
    return json.dumps(stats, indent=2)


def ui_discover(text_type: str, output_dir: str) -> str:
    payload = run_bertopic_discovery(text_type, _resolve_output_dir(output_dir))
    return json.dumps(payload, indent=2)


def ui_label(text_type: str, output_dir: str) -> str:
    payload = label_topics_with_llm(text_type, _resolve_output_dir(output_dir))
    return json.dumps(payload, indent=2)


def ui_theme(text_type: str, output_dir: str) -> str:
    payload = consolidate_into_themes(text_type, 15, _resolve_output_dir(output_dir))
    return json.dumps(payload, indent=2)


def ui_taxonomy(text_type: str, output_dir: str) -> str:
    payload = compare_with_taxonomy(text_type, _resolve_output_dir(output_dir))
    return json.dumps(payload, indent=2)


def ui_compare(output_dir: str) -> str:
    payload = generate_comparison_csv(_resolve_output_dir(output_dir))
    return json.dumps(payload, indent=2)


def ui_narrative(output_dir: str) -> str:
    payload = export_narrative(_resolve_output_dir(output_dir))
    return json.dumps(payload, indent=2)


def ui_full_pipeline(file_path: str, output_dir: str) -> str:
    payload = run_full_pipeline(file_path=file_path, output_dir=_resolve_output_dir(output_dir))
    return json.dumps(payload, indent=2)


def create_interface() -> gr.Blocks:
    with gr.Blocks(css=CUSTOM_CSS, title="CHB BERTopic V3") as app:
        gr.Markdown(
            f"""
            # CHB BERTopic V3
            ### Hugging Face deployment: SPECTER2 + UMAP + HDBSCAN + BERTopic

            **Default input:** `{DEFAULT_INPUT_CSV}`  
            **Default output:** `{OUTPUT_DIR}`  
            **LLM backend:** Hugging Face Inference `{HF_MODEL}`  
            **Secret required for LLM phases:** `HF_TOKEN`
            """
        )

        with gr.Row():
            file_path = gr.Textbox(label="Input CSV", value=str(DEFAULT_INPUT_CSV), lines=1)
            output_dir = gr.Textbox(label="Output directory", value=str(OUTPUT_DIR), lines=1)

        full_run = gr.Button("Run full pipeline", variant="primary")
        full_output = gr.Textbox(label="Full pipeline result", lines=18)
        full_run.click(ui_full_pipeline, inputs=[file_path, output_dir], outputs=[full_output])

        with gr.Tabs():
            with gr.Tab("Phase 1"):
                load_btn = gr.Button("Load corpus")
                load_out = gr.Textbox(lines=16, label="Load output")
                load_btn.click(ui_load, inputs=[file_path, output_dir], outputs=[load_out])

            with gr.Tab("Phase 2"):
                abs_disc_btn = gr.Button("Discover abstract topics")
                title_disc_btn = gr.Button("Discover title topics")
                abs_disc_out = gr.Textbox(lines=16, label="Abstract discovery")
                title_disc_out = gr.Textbox(lines=16, label="Title discovery")
                abs_disc_btn.click(ui_discover, inputs=[gr.State("abstract"), output_dir], outputs=[abs_disc_out])
                title_disc_btn.click(ui_discover, inputs=[gr.State("title"), output_dir], outputs=[title_disc_out])

            with gr.Tab("Phase 3"):
                abs_label_btn = gr.Button("Label abstract topics")
                title_label_btn = gr.Button("Label title topics")
                abs_label_out = gr.Textbox(lines=16, label="Abstract labels")
                title_label_out = gr.Textbox(lines=16, label="Title labels")
                abs_label_btn.click(ui_label, inputs=[gr.State("abstract"), output_dir], outputs=[abs_label_out])
                title_label_btn.click(ui_label, inputs=[gr.State("title"), output_dir], outputs=[title_label_out])

            with gr.Tab("Phase 4"):
                abs_theme_btn = gr.Button("Consolidate abstract themes")
                title_theme_btn = gr.Button("Consolidate title themes")
                abs_theme_out = gr.Textbox(lines=16, label="Abstract themes")
                title_theme_out = gr.Textbox(lines=16, label="Title themes")
                abs_theme_btn.click(ui_theme, inputs=[gr.State("abstract"), output_dir], outputs=[abs_theme_out])
                title_theme_btn.click(ui_theme, inputs=[gr.State("title"), output_dir], outputs=[title_theme_out])

            with gr.Tab("Phase 5"):
                abs_tax_btn = gr.Button("Map abstract themes to PAJAIS")
                title_tax_btn = gr.Button("Map title themes to PAJAIS")
                abs_tax_out = gr.Textbox(lines=16, label="Abstract taxonomy")
                title_tax_out = gr.Textbox(lines=16, label="Title taxonomy")
                abs_tax_btn.click(ui_taxonomy, inputs=[gr.State("abstract"), output_dir], outputs=[abs_tax_out])
                title_tax_btn.click(ui_taxonomy, inputs=[gr.State("title"), output_dir], outputs=[title_tax_out])

            with gr.Tab("Phase 6"):
                compare_btn = gr.Button("Generate comparison")
                compare_out = gr.Textbox(lines=16, label="Comparison")
                compare_btn.click(ui_compare, inputs=[output_dir], outputs=[compare_out])

            with gr.Tab("Phase 7"):
                narrative_btn = gr.Button("Generate narrative")
                narrative_out = gr.Textbox(lines=16, label="Narrative result")
                narrative_btn.click(ui_narrative, inputs=[output_dir], outputs=[narrative_out])

    return app


def main() -> None:
    parser = argparse.ArgumentParser(description="CHB BERTopic Hugging Face app")
    parser.add_argument("--mode", choices=["ui", "pipeline"], default="ui")
    parser.add_argument("--input", default=str(DEFAULT_INPUT_CSV))
    parser.add_argument("--output-dir", default=str(OUTPUT_DIR))
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7860)
    args = parser.parse_args()

    if args.mode == "pipeline":
        payload = run_full_pipeline(file_path=args.input, output_dir=Path(args.output_dir))
        print(json.dumps(payload, indent=2))
        return

    app = create_interface()
    app.launch(server_name=args.host, server_port=args.port)


if __name__ == "__main__":
    main()