#!/usr/bin/env python3
"""Build auditable landscape counts and evidence candidates for Agent Memory."""

from __future__ import annotations

import argparse
import csv
import json
import re
from collections import Counter, defaultdict
from datetime import date
from pathlib import Path


ROOT = Path(__file__).resolve().parents[2]

THEMES = {
    "retrieval_access": r"retriev|recall|search|query|evidence access|memory navigation",
    "representation_organization": r"knowledge graph|graph memory|hierarch|structure|representation|index|database|topic document|ontology",
    "update_forgetting": r"update|forget|consolidat|supersed|stale|conflict|write path|write-path|memory admission|retention",
    "execution_state": r"execution state|task state|workflow state|belief state|long-horizon|trajectory|subgoal|planning",
    "experience_skill_learning": r"experience|procedural|skill|reflect|self-evolv|test-time learn|continual learn|reinforcement learn|\brl\b",
    "evaluation_diagnosis": r"benchmark|evaluat|diagnos|metric|ablation|empirical study|controlled",
    "security_privacy_trust": r"attack|poison|security|privacy|trust|provenance|govern|leak|integrity|adversarial",
    "multi_agent_shared": r"multi-agent|multiagent|shared memory|collective memory|collaborative memory",
    "multimodal_embodied_gui": r"multimodal|visual|video|robot|embodied|gui|mobile|web agent|navigation",
    "systems_efficiency": r"token|latency|cost|efficient|compression|serving|storage|throughput|budget",
}

# Title-only labels describe the paper's declared contribution more conservatively
# than abstract keywords such as "evaluate", which occur in almost every paper.
TITLE_THEMES = {
    "benchmark_diagnosis": r"benchmark|evaluat|diagnos|empirical|survey|taxonomy|anatomy|rethinking|what .* need|how far|study of|analysis",
    "retrieval_representation": r"retriev|recall|search|graph|hierarch|structure|index|database|ontology|organization|representation|rag\b|topic document",
    "update_forgetting": r"update|forget|consolidat|stale|supersed|decay|write memory|not to write|promotion|temporal conflict|belief revision|lifecycle",
    "experience_skill_policy": r"experience|skill|procedural|learn|evolv|adapt|reinforcement|\brl\b|credit assignment|distill",
    "execution_multimodal": r"execution state|task-driving state|task state|long-horizon|gui|embodied|robot|navigation|web agent|multimodal|video|spatial",
    "security_governance": r"attack|poison|security|privacy|trust|govern|provenance|integrity|leak|hijack|vulnerab|unlearn|false promotion|contamination",
    "systems_efficiency": r"efficient|cost|token|latency|serving|on-device|edge|compress|scalable|storage|runtime|budget|wire format",
    "personal_shared": r"personali|user-aware|user-centric|multi-agent|multiagent|shared memory|collaborative|collective",
}

WINDOWS = {
    "previous_six_months": ("2025-07-09", "2026-01-10"),
    "recent_six_months": ("2026-01-10", "2026-07-11"),
}

CENTRAL_TITLE = re.compile(r"memory|memori|remember|forget|stateful", re.I)
AGENT_CONTEXT = re.compile(r"agent|large language model|\bllm\b|foundation model", re.I)
HARDWARE_ONLY = re.compile(
    r"gpu memory|memory bandwidth|memory footprint|memory-efficient|memory efficient|bounded memory agents on",
    re.I,
)
NUMERIC_RESULT = re.compile(r"(?:\d+(?:\.\d+)?\s*%|\+?\d+(?:\.\d+)?\s*(?:pp|points?|×|x))", re.I)


def period(published: str) -> str:
    if published < "2025-07-01":
        return "2025-H1"
    if published < "2026-02-01":
        return "baseline-2025H2-Jan2026"
    if published < "2026-04-01":
        return "transition-2026-Feb-Mar"
    return "recent-2026-Apr-Jul"


def central_by_text(record: dict) -> bool:
    title = record["title"]
    abstract = record["abstract"]
    return bool(
        CENTRAL_TITLE.search(title)
        and AGENT_CONTEXT.search(f"{title} {abstract}")
        and not HARDWARE_ONLY.search(title)
    )


def theme_labels(record: dict) -> list[str]:
    text = f"{record['title']} {record['abstract']}".lower()
    return [name for name, pattern in THEMES.items() if re.search(pattern, text, re.I)]


def title_theme_labels(record: dict) -> list[str]:
    title = record["title"].lower()
    return [name for name, pattern in TITLE_THEMES.items() if re.search(pattern, title, re.I)]


def window_name(published: str) -> str | None:
    for name, (start, end) in WINDOWS.items():
        if start <= published < end:
            return name
    return None


def evidence_score(record: dict) -> int:
    text = record["abstract"].lower()
    score = 0
    score += 2 if re.search(r"benchmark|evaluate|evaluation|experiment", text) else 0
    score += 2 if re.search(r"baseline|compare|outperform|versus|\bvs\.?\b", text) else 0
    score += 2 if NUMERIC_RESULT.search(text) else 0
    score += 1 if re.search(r"ablation|controlled|factorial|paired|significan", text) else 0
    score += 1 if re.search(r"code is|github|we release|open-source", text) else 0
    score += 1 if re.search(r"limitation|does not|not always|far from|fails?", text) else 0
    return score


def write_json(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=Path,
        default=ROOT / "data" / "research" / "agent-memory-expanded.json",
    )
    parser.add_argument(
        "--screen",
        type=Path,
        default=ROOT / "data" / "research" / "agent-memory-screened.json",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=ROOT / "data" / "research" / "memory-landscape.json",
    )
    parser.add_argument(
        "--csv",
        type=Path,
        default=ROOT / "data" / "research" / "memory-paper-ledger.csv",
    )
    args = parser.parse_args()

    source = json.loads(args.input.read_text(encoding="utf-8"))
    screen = {}
    if args.screen.exists():
        screened = json.loads(args.screen.read_text(encoding="utf-8"))
        screen = {item["arxiv_id"]: item.get("screen") or {} for item in screened.get("papers", [])}

    records = []
    for paper in source["papers"]:
        screened = screen.get(paper["arxiv_id"])
        records.append(
            {
                **paper,
                "period": period(paper["published"]),
                "central_by_title": central_by_text(paper),
                "themes": theme_labels(paper),
                "title_themes": title_theme_labels(paper),
                "evidence_score": evidence_score(paper),
                "screen": screened,
            }
        )

    period_counts: defaultdict[str, Counter] = defaultdict(Counter)
    theme_counts: defaultdict[str, Counter] = defaultdict(Counter)
    window_counts: defaultdict[str, Counter] = defaultdict(Counter)
    title_theme_counts: defaultdict[str, Counter] = defaultdict(Counter)
    screen_counts: defaultdict[str, Counter] = defaultdict(Counter)
    for record in records:
        bucket = record["period"]
        period_counts[bucket]["candidates"] += 1
        if record["central_by_title"]:
            period_counts[bucket]["title_central"] += 1
        relevance = (record.get("screen") or {}).get("relevance")
        if relevance:
            period_counts[bucket][f"screen_{relevance}"] += 1
        for theme in record["themes"]:
            if record["central_by_title"]:
                theme_counts[bucket][theme] += 1
        window = window_name(record["published"])
        if window:
            window_counts[window]["candidates"] += 1
            if record["central_by_title"]:
                window_counts[window]["title_central"] += 1
                for theme in record["title_themes"]:
                    title_theme_counts[window][theme] += 1
        screened = record.get("screen") or {}
        if screened:
            screen_counts["relevance"][screened.get("relevance", "unknown")] += 1
            screen_counts["problem"][screened.get("problem", "unknown")] += 1
            screen_counts["source"][screened.get("source", "unknown")] += 1

    top_evidence = {}
    for theme in THEMES:
        candidates = [
            record
            for record in records
            if record["central_by_title"] and theme in record["themes"]
        ]
        candidates.sort(key=lambda item: (item["evidence_score"], item["published"]), reverse=True)
        top_evidence[theme] = [
            {
                "arxiv_id": item["arxiv_id"],
                "published": item["published"],
                "title": item["title"],
                "score": item["evidence_score"],
                "url": item["url"],
            }
            for item in candidates[:20]
        ]

    output = {
        "generated_at": date.today().isoformat(),
        "source_candidates": len(records),
        "title_central": sum(record["central_by_title"] for record in records),
        "screened": len(screen),
        "windows": {
            name: {
                "from_inclusive": start,
                "to_exclusive": end,
                **dict(window_counts[name]),
                "title_theme_counts": dict(title_theme_counts[name]),
            }
            for name, (start, end) in WINDOWS.items()
        },
        "period_counts": {key: dict(value) for key, value in period_counts.items()},
        "theme_counts_title_central": {key: dict(value) for key, value in theme_counts.items()},
        "screen_counts": {key: dict(value) for key, value in screen_counts.items()},
        "evidence_score_distribution": dict(
            Counter(record["evidence_score"] for record in records if record["central_by_title"])
        ),
        "top_evidence_candidates": top_evidence,
    }
    write_json(args.output, output)

    args.csv.parent.mkdir(parents=True, exist_ok=True)
    with args.csv.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=[
                "arxiv_id",
                "published",
                "period",
                "title",
                "url",
                "title_central",
                "screen_relevance",
                "screen_problem",
                "screen_source",
                "themes",
                "title_themes",
                "evidence_score",
                "matched_queries",
            ],
        )
        writer.writeheader()
        for record in records:
            screened = record.get("screen") or {}
            writer.writerow(
                {
                    "arxiv_id": record["arxiv_id"],
                    "published": record["published"],
                    "period": record["period"],
                    "title": record["title"],
                    "url": record["url"],
                    "title_central": record["central_by_title"],
                    "screen_relevance": screened.get("relevance", ""),
                    "screen_problem": screened.get("problem", ""),
                    "screen_source": screened.get("source", ""),
                    "themes": ";".join(record["themes"]),
                    "title_themes": ";".join(record["title_themes"]),
                    "evidence_score": record["evidence_score"],
                    "matched_queries": ";".join(record["matched_queries"]),
                }
            )
    print(
        json.dumps(
            {
                "source_candidates": output["source_candidates"],
                "title_central": output["title_central"],
                "screened": output["screened"],
                "windows": output["windows"],
                "screen_counts": output["screen_counts"],
            },
            ensure_ascii=False,
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
