tools.py

"""
tools.py
--------
Tool definitions for the Journal Topic Modelling Agent.
Each tool is a plain Python function; the agent (agent.py) calls these directly.
All functions return a dict with a "status" key ("ok" or "error").
"""

import json
import re
import csv
import io
from typing import Any, List, Dict


# ── Phase 1: Ingestion ────────────────────────────────────────────────────────

def parse_csv(csv_text: str) -> dict:
    """
    Parse raw CSV text and return structured records.
    Returns: {"status", "records", "columns", "count"}
    """
    try:
        reader = csv.DictReader(io.StringIO(csv_text.strip()))
        records = [dict(row) for row in reader]
        columns = list(reader.fieldnames or [])
        return {
            "status": "ok",
            "records": records,
            "columns": columns,
            "count": len(records),
        }
    except Exception as e:
        return {"status": "error", "message": str(e), "records": [], "columns": [], "count": 0}


def extract_text_corpus(records: list, text_fields: list = None) -> dict:
    """
    Extract and combine text from specified fields across all records.
    Defaults to ['abstract', 'title'].
    Returns: {"status", "corpus", "total_words", "doc_count"}
    """
    if text_fields is None:
        text_fields = ["abstract", "title"]

    corpus = []
    total_words = 0
    for i, rec in enumerate(records):
        parts = []
        for field in text_fields:
            for k, v in rec.items():
                if k.strip().lower() == field.lower() and v and str(v).strip():
                    parts.append(str(v).strip())
        text = " ".join(parts)
        if text:
            corpus.append({"id": i, "text": text, "record": rec})
            total_words += len(text.split())

    return {
        "status": "ok",
        "corpus": corpus,
        "total_words": total_words,
        "doc_count": len(corpus),
    }


# ── Phase 2: Topic Extraction & Labelling ─────────────────────────────────────

def chunk_corpus_for_topic_extraction(corpus: list, chunk_size: int = 20) -> dict:
    """
    Split corpus into chunks for batched LLM topic extraction.
    Returns: {"status", "chunks", "chunk_count"}
    """
    chunks = []
    for i in range(0, len(corpus), chunk_size):
        chunks.append(corpus[i : i + chunk_size])
    return {"status": "ok", "chunks": chunks, "chunk_count": len(chunks)}


def merge_and_deduplicate_topics(topic_lists: list) -> dict:
    """
    Merge multiple raw topic lists, normalise, and deduplicate.
    Returns: {"status", "unique_topics", "raw_count", "unique_count"}
    """
    all_topics = []
    for lst in topic_lists:
        if isinstance(lst, list):
            all_topics.extend(lst)
        elif isinstance(lst, str) and lst.strip():
            all_topics.append(lst)

    seen = set()
    unique = []
    for t in all_topics:
        if not isinstance(t, str):
            continue
        norm = re.sub(r"\s+", " ", t.strip().lower())
        if norm and norm not in seen:
            seen.add(norm)
            unique.append(t.strip())

    return {
        "status": "ok",
        "unique_topics": unique,
        "raw_count": len(all_topics),
        "unique_count": len(unique),
    }


def assign_topic_labels(topics: list, labels_map: dict) -> dict:
    """
    Assign human-readable labels to a list of topic strings.
    labels_map: {topic_string: label_string}
    Returns: {"status", "labelled_topics", "count"}
    """
    labelled = []
    for t in topics:
        label = labels_map.get(t) or labels_map.get(t.lower()) or t
        labelled.append({"topic": t, "label": label})
    return {"status": "ok", "labelled_topics": labelled, "count": len(labelled)}


def build_topic_frequency_table(corpus: list, labelled_topics: list) -> dict:
    """
    Count how many documents mention each topic (simple keyword match).
    Returns: {"status", "frequency_table", "doc_count"}
    """
    doc_count = len(corpus)
    freq_table = []
    for item in labelled_topics:
        topic_words = item["topic"].lower().split()
        count = sum(
            1 for doc in corpus if all(w in doc["text"].lower() for w in topic_words)
        )
        pct = round(100.0 * count / doc_count, 2) if doc_count else 0.0
        freq_table.append(
            {
                "topic": item["topic"],
                "label": item["label"],
                "frequency": count,
                "percentage": pct,
            }
        )
    freq_table.sort(key=lambda x: x["frequency"], reverse=True)
    return {"status": "ok", "frequency_table": freq_table, "doc_count": doc_count}


# ── Phase 3: Title vs Abstract Comparison ────────────────────────────────────

def split_corpus_by_field(records: list) -> dict:
    """
    Separate title and abstract corpora from the records list.
    Returns: {"status", "title_corpus", "abstract_corpus", "title_count", "abstract_count"}
    """
    title_corpus = []
    abstract_corpus = []
    for i, rec in enumerate(records):
        title_text = ""
        abstract_text = ""
        for k, v in rec.items():
            kl = k.strip().lower()
            if kl == "title" and v:
                title_text = str(v).strip()
            elif kl == "abstract" and v:
                abstract_text = str(v).strip()
        if title_text:
            title_corpus.append({"id": i, "text": title_text})
        if abstract_text:
            abstract_corpus.append({"id": i, "text": abstract_text})
    return {
        "status": "ok",
        "title_corpus": title_corpus,
        "abstract_corpus": abstract_corpus,
        "title_count": len(title_corpus),
        "abstract_count": len(abstract_corpus),
    }


def compare_topic_distributions(title_topics: list, abstract_topics: list) -> dict:
    """
    Compare topic frequency distributions between title and abstract corpora.
    Each input list: [{"topic", "frequency", "percentage"}, ...]
    Returns: {"status", "comparison", "title_only", "abstract_only", "shared"}
    """
    title_map = {t["topic"].lower(): t for t in title_topics}
    abstract_map = {t["topic"].lower(): t for t in abstract_topics}
    all_keys = set(title_map.keys()) | set(abstract_map.keys())

    comparison = []
    for key in all_keys:
        t_item = title_map.get(key, {})
        a_item = abstract_map.get(key, {})
        topic_name = (t_item.get("topic") or a_item.get("topic") or key)
        comparison.append(
            {
                "topic": topic_name,
                "title_freq": t_item.get("frequency", 0),
                "abstract_freq": a_item.get("frequency", 0),
                "title_pct": t_item.get("percentage", 0.0),
                "abstract_pct": a_item.get("percentage", 0.0),
                "delta_pct": round(
                    a_item.get("percentage", 0.0) - t_item.get("percentage", 0.0), 2
                ),
            }
        )
    comparison.sort(key=lambda x: abs(x["delta_pct"]), reverse=True)

    title_only = [c["topic"] for c in comparison if c["title_freq"] > 0 and c["abstract_freq"] == 0]
    abstract_only = [c["topic"] for c in comparison if c["abstract_freq"] > 0 and c["title_freq"] == 0]
    shared = [c["topic"] for c in comparison if c["title_freq"] > 0 and c["abstract_freq"] > 0]

    return {
        "status": "ok",
        "comparison": comparison,
        "title_only": title_only,
        "abstract_only": abstract_only,
        "shared": shared,
    }


def save_comparison_csv(comparison: list) -> dict:
    """
    Serialise the comparison list to CSV format string.
    Returns: {"status", "csv_text", "row_count"}
    """
    if not comparison:
        return {"status": "error", "message": "No comparison data provided", "csv_text": "", "row_count": 0}

    output = io.StringIO()
    fieldnames = ["topic", "title_freq", "abstract_freq", "title_pct", "abstract_pct", "delta_pct"]
    writer = csv.DictWriter(output, fieldnames=fieldnames, extrasaction="ignore")
    writer.writeheader()
    writer.writerows(comparison)
    return {"status": "ok", "csv_text": output.getvalue(), "row_count": len(comparison)}


# ── Phase 4: PAJAIS Taxonomy Mapping ─────────────────────────────────────────

PAJAIS_THEMES = [
    "Human-Computer Interaction",
    "Decision Support Systems",
    "Knowledge Management",
    "Information Retrieval",
    "Machine Learning & AI",
    "Natural Language Processing",
    "Big Data & Analytics",
    "Privacy & Security",
    "Social Media & Web 2.0",
    "Healthcare Informatics",
    "Education & e-Learning",
    "Business Intelligence",
    "Recommender Systems",
    "Cloud Computing",
    "Internet of Things",
    "Ethical AI & Fairness",
    "Digital Transformation",
    "Ontologies & Semantic Web",
    "Supply Chain & Operations",
    "Sentiment Analysis",
]


def map_topics_to_pajais(labelled_topics: list, mapping: dict) -> dict:
    """
    Map discovered topics to PAJAIS taxonomy themes.
    mapping: {topic_string: pajais_theme_or_"NOVEL"}
    Returns: {"status", "taxonomy_map", "mapped", "novel",
              "coverage_pct", "pajais_themes_used", "novel_count", "mapped_count"}
    """
    taxonomy_map = {theme: [] for theme in PAJAIS_THEMES}
    taxonomy_map["NOVEL"] = []

    mapped = []
    novel = []

    for item in labelled_topics:
        topic = item["topic"]
        theme = mapping.get(topic) or mapping.get(topic.lower()) or "NOVEL"
        entry = {"topic": topic, "label": item.get("label", topic), "pajais_theme": theme}
        if theme == "NOVEL":
            novel.append(entry)
            taxonomy_map["NOVEL"].append(topic)
        else:
            mapped.append(entry)
            if theme in taxonomy_map:
                taxonomy_map[theme].append(topic)
            else:
                taxonomy_map[theme] = [topic]

    total = len(labelled_topics)
    coverage = round(100.0 * len(mapped) / total, 2) if total else 0.0

    return {
        "status": "ok",
        "taxonomy_map": taxonomy_map,
        "mapped": mapped,
        "novel": novel,
        "coverage_pct": coverage,
        "pajais_themes_used": list({m["pajais_theme"] for m in mapped}),
        "novel_count": len(novel),
        "mapped_count": len(mapped),
    }


def save_taxonomy_json(taxonomy_map: dict) -> dict:
    """
    Serialise taxonomy map to a JSON string.
    Returns: {"status", "json_text"}
    """
    try:
        return {"status": "ok", "json_text": json.dumps(taxonomy_map, indent=2)}
    except Exception as e:
        return {"status": "error", "message": str(e), "json_text": "{}"}


# ── Phase 5: Narrative Support ────────────────────────────────────────────────

def build_narrative_context(
    frequency_table: list,
    comparison: list,
    taxonomy_result: dict,
    record_count: int,
) -> dict:
    """
    Assemble a structured context dict to feed into the narrative generation prompt.
    Returns: {"status", "context"}
    """
    top_topics = [
        {"topic": t["topic"], "label": t["label"], "freq": t["frequency"], "pct": t["percentage"]}
        for t in frequency_table[:10]
    ]
    biggest_delta = sorted(comparison, key=lambda x: abs(x.get("delta_pct", 0)), reverse=True)[:5]
    novel_sample = [n["topic"] for n in taxonomy_result.get("novel", [])[:10]]

    context = {
        "record_count": record_count,
        "top_10_topics": top_topics,
        "biggest_title_abstract_deltas": biggest_delta,
        "novel_themes_count": len(taxonomy_result.get("novel", [])),
        "novel_sample": novel_sample,
        "pajais_themes_covered": taxonomy_result.get("pajais_themes_used", []),
        "pajais_coverage_pct": taxonomy_result.get("coverage_pct", 0),
    }
    return {"status": "ok", "context": context}


def validate_narrative(text: str, min_words: int = 450) -> dict:
    """
    Check that the narrative meets the minimum word count.
    Returns: {"status", "valid", "word_count", "message"}
    """
    words = len(text.split())
    valid = words >= min_words
    return {
        "status": "ok",
        "valid": valid,
        "word_count": words,
        "message": "OK" if valid else f"Too short: {words} words (need {min_words}+)",
    }


# ── Utility ───────────────────────────────────────────────────────────────────

def summarise_run(
    record_count: int,
    topic_count: int,
    mapped_count: int,
    novel_count: int,
    coverage_pct: float,
) -> dict:
    """Return a summary dict for display."""
    return {
        "status": "ok",
        "summary": {
            "papers_analysed": record_count,
            "topics_discovered": topic_count,
            "pajais_mapped": mapped_count,
            "novel_themes": novel_count,
            "pajais_coverage_pct": coverage_pct,
        },
    }


# ── Tool registry ─────────────────────────────────────────────────────────────

TOOL_REGISTRY = {
    "parse_csv": parse_csv,
    "extract_text_corpus": extract_text_corpus,
    "chunk_corpus_for_topic_extraction": chunk_corpus_for_topic_extraction,
    "merge_and_deduplicate_topics": merge_and_deduplicate_topics,
    "assign_topic_labels": assign_topic_labels,
    "build_topic_frequency_table": build_topic_frequency_table,
    "split_corpus_by_field": split_corpus_by_field,
    "compare_topic_distributions": compare_topic_distributions,
    "save_comparison_csv": save_comparison_csv,
    "map_topics_to_pajais": map_topics_to_pajais,
    "save_taxonomy_json": save_taxonomy_json,
    "build_narrative_context": build_narrative_context,
    "validate_narrative": validate_narrative,
    "summarise_run": summarise_run,
}

PAJAIS_THEMES_LIST = PAJAIS_THEMES