#!/usr/bin/env python3
"""Count and plot papers whose titles contain the whole word "diffusion".

Sources:
  * CVPR and ICCV: Computer Vision Foundation Open Access Repository
  * ECCV: European Computer Vision Association Open Access Repository

Only main-conference paper titles are searched. Matching is case-insensitive
and uses ``\bdiffusion\b``, so "diffusion" is counted but "diffusive" is not.

Run:
    python plot_vision_conferences_diffusion_papers.py

Outputs:
    vision_conferences_diffusion_papers.csv
    vision_conferences_diffusion_papers.png
    vision_conferences_diffusion_papers.pdf
    vision_conferences_diffusion_papers.svg
"""

from __future__ import annotations

from html.parser import HTMLParser
from pathlib import Path
import csv
import gzip
import re
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urljoin
from urllib.request import Request, urlopen

import matplotlib.pyplot as plt
from matplotlib import font_manager
from matplotlib.ticker import MultipleLocator


ROOT = Path(__file__).resolve().parent
YEARS = list(range(2020, 2027))
WORD_PATTERN = re.compile(r"\bdiffusion\b", flags=re.IGNORECASE)
USER_AGENT = "Mozilla/5.0 (compatible; academic-title-analysis/1.0)"

CVF_BASE = "https://openaccess.thecvf.com/"
ECVA_PAPERS_URL = "https://www.ecva.net/papers.php"

# A missing value means that the conference had no edition in that year, or
# that its official proceedings were not available when the chart was made.
CONFERENCE_YEARS = {
    "CVPR": set(YEARS),
    "ECCV": {2020, 2022, 2024, 2026},
    "ICCV": {2021, 2023, 2025},
}

CONFERENCE_COLORS = {
    "CVPR": "#2F6B9A",  # deep blue
    "ECCV": "#E6A23C",  # warm amber
    "ICCV": "#D95D5D",  # muted coral
}


class PaperTitleParser(HTMLParser):
    """Extract text from ``<dt class="ptitle">`` elements."""

    def __init__(self) -> None:
        super().__init__()
        self._inside_title = False
        self._buffer: list[str] = []
        self.titles: list[str] = []

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        attributes = dict(attrs)
        classes = (attributes.get("class") or "").split()
        if tag == "dt" and "ptitle" in classes:
            self._inside_title = True
            self._buffer = []

    def handle_endtag(self, tag: str) -> None:
        if tag == "dt" and self._inside_title:
            title = " ".join("".join(self._buffer).split())
            if title:
                self.titles.append(title)
            self._inside_title = False

    def handle_data(self, data: str) -> None:
        if self._inside_title:
            self._buffer.append(data)


def fetch_html(url: str, attempts: int = 3) -> str:
    """Download an official proceedings page using only the standard library."""
    request = Request(url, headers={"User-Agent": USER_AGENT})
    for attempt in range(1, attempts + 1):
        try:
            with urlopen(request, timeout=60) as response:
                payload = response.read()
                if response.headers.get("Content-Encoding") == "gzip" or payload.startswith(b"\x1f\x8b"):
                    payload = gzip.decompress(payload)
                return payload.decode("utf-8", errors="replace")
        except (HTTPError, URLError, TimeoutError):
            if attempt == attempts:
                raise
            time.sleep(1.5 * attempt)
    raise RuntimeError(f"Could not download {url}")


def extract_titles(html: str) -> list[str]:
    parser = PaperTitleParser()
    parser.feed(html)
    return parser.titles


def unique_preserving_order(items: list[str]) -> list[str]:
    return list(dict.fromkeys(items))


def fetch_cvpr_titles(year: int) -> list[str]:
    """Fetch CVPR titles, including the special three-page layout used in 2020."""
    if year != 2020:
        page = fetch_html(urljoin(CVF_BASE, f"CVPR{year}?day=all"))
        return unique_preserving_order(extract_titles(page))

    landing_url = urljoin(CVF_BASE, "CVPR2020")
    landing = fetch_html(landing_url)
    day_links = unique_preserving_order(
        re.findall(r'href=["\']([^"\']*CVPR2020\.py\?day=[^"\']+)["\']', landing)
    )
    titles: list[str] = []
    for day_link in day_links:
        titles.extend(extract_titles(fetch_html(urljoin(landing_url, day_link))))
    return unique_preserving_order(titles)


def fetch_iccv_titles(year: int) -> list[str]:
    page = fetch_html(urljoin(CVF_BASE, f"ICCV{year}?day=all"))
    return unique_preserving_order(extract_titles(page))


def extract_eccv_year_section(all_years_html: str, year: int) -> str:
    """Select one ECCV edition from ECVA's combined proceedings page."""
    start_marker = f"<!-- ECCV {year} -->"
    start = all_years_html.find(start_marker)
    if start < 0:
        raise ValueError(f"ECCV {year} section was not found")

    next_section = re.search(r"<!--\s*ECCV\s+\d{4}\s*-->", all_years_html[start + 1 :])
    end = len(all_years_html)
    if next_section:
        end = start + 1 + next_section.start()
    return all_years_html[start:end]


def count_matching_titles(titles: list[str]) -> int:
    return sum(bool(WORD_PATTERN.search(title)) for title in titles)


def collect_counts() -> dict[str, dict[int, int | None]]:
    counts: dict[str, dict[int, int | None]] = {
        conference: {year: None for year in YEARS}
        for conference in ("CVPR", "ECCV", "ICCV")
    }

    for year in YEARS:
        counts["CVPR"][year] = count_matching_titles(fetch_cvpr_titles(year))

    ecva_html = fetch_html(ECVA_PAPERS_URL)
    for year in sorted(CONFERENCE_YEARS["ECCV"]):
        try:
            section = extract_eccv_year_section(ecva_html, year)
        except ValueError:
            # Keep the value as None until the official proceedings appear.
            continue
        counts["ECCV"][year] = count_matching_titles(extract_titles(section))

    for year in sorted(CONFERENCE_YEARS["ICCV"]):
        counts["ICCV"][year] = count_matching_titles(fetch_iccv_titles(year))

    return counts


def save_csv(counts: dict[str, dict[int, int | None]]) -> None:
    output = ROOT / "vision_conferences_diffusion_papers.csv"
    with output.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["year", "CVPR", "ECCV", "ICCV", "combined_total"])
        for year in YEARS:
            combined_total = sum(
                value for name in ("CVPR", "ECCV", "ICCV")
                if (value := counts[name][year]) is not None
            )
            writer.writerow(
                [year]
                + ["" if counts[name][year] is None else counts[name][year] for name in ("CVPR", "ECCV", "ICCV")]
                + [combined_total]
            )


def choose_font() -> str:
    installed = {font.name for font in font_manager.fontManager.ttflist}
    for candidate in (
        "Aptos",
        "Arial",
        "Helvetica",
        "Liberation Sans",
        "DejaVu Sans",
    ):
        if candidate in installed:
            return candidate
    return "DejaVu Sans"


def create_chart(counts: dict[str, dict[int, int | None]]) -> None:
    font = choose_font()
    plt.rcParams.update(
        {
            "font.family": font,
            "font.size": 15,
            "axes.labelcolor": "#22313F",
            "xtick.color": "#4B5563",
            "ytick.color": "#4B5563",
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
            "svg.fonttype": "none",
        }
    )

    fig, ax = plt.subplots(figsize=(15, 8.8), facecolor="white")
    fig.subplots_adjust(left=0.09, right=0.98, top=0.76, bottom=0.20)

    x_positions = list(range(len(YEARS)))
    combined_totals = [
        sum(
            value for name in ("CVPR", "ECCV", "ICCV")
            if (value := counts[name][year]) is not None
        )
        for year in YEARS
    ]

    bottoms = [0] * len(YEARS)
    for conference in ("CVPR", "ECCV", "ICCV"):
        values = [counts[conference][year] or 0 for year in YEARS]
        bars = ax.bar(
            x_positions,
            values,
            width=0.68,
            bottom=bottoms,
            color=CONFERENCE_COLORS[conference],
            edgecolor="white",
            linewidth=1.0,
            label=conference,
            zorder=3,
        )

        for bar, value, bottom in zip(bars, values, bottoms):
            if value >= 45:
                ax.text(
                    bar.get_x() + bar.get_width() / 2,
                    bottom + value / 2,
                    f"{conference}  {value:,}",
                    ha="center",
                    va="center",
                    fontsize=12.5,
                    fontweight="bold",
                    color="white",
                    zorder=4,
                )
        bottoms = [bottom + value for bottom, value in zip(bottoms, values)]

    # Small early-year segments are too short for in-bar labels, so place their
    # conference-colored labels just above the corresponding bar.
    small_label_counts: list[int] = []
    for x_position, year, total in zip(x_positions, YEARS, combined_totals):
        small_segments = [
            (conference, counts[conference][year])
            for conference in ("CVPR", "ECCV", "ICCV")
            if counts[conference][year] is not None and 0 < int(counts[conference][year]) < 45
        ]
        small_label_counts.append(len(small_segments))
        for line_number, (conference, value) in enumerate(small_segments):
            ax.annotate(
                f"{conference}  {int(value):,}",
                xy=(x_position, total),
                xytext=(0, 8 + line_number * 15),
                textcoords="offset points",
                ha="center",
                va="bottom",
                fontsize=11.5,
                fontweight="bold",
                color=CONFERENCE_COLORS[conference],
                clip_on=False,
            )

    for x_position, value, small_count in zip(x_positions, combined_totals, small_label_counts):
        ax.annotate(
            f"Total  {value:,}",
            xy=(x_position, value),
            xytext=(0, 10 + small_count * 15),
            textcoords="offset points",
            ha="center",
            va="bottom",
            fontsize=13.5,
            fontweight="bold",
            color="#15324A",
            clip_on=False,
        )

    ax.set_xlim(-0.65, len(YEARS) - 0.35)
    ax.set_ylim(0, 725)
    ax.set_xticks(x_positions, [str(year) for year in YEARS], fontsize=16)
    ax.yaxis.set_major_locator(MultipleLocator(100))
    ax.tick_params(axis="y", labelsize=15, length=0, pad=8)
    ax.tick_params(axis="x", length=0, pad=10)
    ax.set_xlabel("Publication year", fontsize=17, fontweight="bold", labelpad=14)
    ax.set_ylabel("Number of papers", fontsize=17, fontweight="bold", labelpad=14)

    ax.grid(axis="y", color="#DDE5EC", linewidth=1.0, alpha=0.9, zorder=0)
    ax.grid(axis="x", visible=False)
    for side in ("top", "right", "left"):
        ax.spines[side].set_visible(False)
    ax.spines["bottom"].set_color("#AAB4C0")
    ax.spines["bottom"].set_linewidth(1.0)

    legend = ax.legend(
        loc="lower left",
        bbox_to_anchor=(0.0, 1.025),
        ncol=3,
        frameon=False,
        fontsize=14,
        handlelength=1.4,
        columnspacing=2.0,
    )
    for text in legend.get_texts():
        text.set_fontweight("bold")

    fig.text(
        0.09,
        0.925,
        "Diffusion Papers at CVPR + ICCV + ECCV",
        fontsize=28,
        fontweight="bold",
        color="#102A43",
        ha="left",
    )
    fig.text(
        0.09,
        0.87,
        "Annual totals with each conference’s contribution shown as a colored segment",
        fontsize=16,
        color="#52606D",
        ha="left",
    )
    fig.text(
        0.09,
        0.06,
        "Sources: CVF Open Access (CVPR, ICCV) and ECVA (ECCV)  •  "
        "Exact case-insensitive title match; ‘diffusive’ and workshops excluded",
        fontsize=11,
        color="#6B7280",
        ha="left",
    )
    fig.text(
        0.09,
        0.027,
        "ECCV 2026 proceedings are not yet available; therefore, the 2026 bar currently contains CVPR only.",
        fontsize=11,
        color="#6B7280",
        ha="left",
    )

    for extension in ("png", "pdf", "svg"):
        output = ROOT / f"vision_conferences_diffusion_papers.{extension}"
        kwargs: dict[str, object] = {"bbox_inches": "tight", "facecolor": "white"}
        if extension == "png":
            kwargs["dpi"] = 300
        fig.savefig(output, **kwargs)
    plt.close(fig)


def main() -> None:
    counts = collect_counts()
    save_csv(counts)
    create_chart(counts)
    for year in YEARS:
        values = [counts[name][year] for name in ("CVPR", "ECCV", "ICCV")]
        print(year, *["N/A" if value is None else value for value in values])


if __name__ == "__main__":
    main()