plot_compared_normalized_frequencies.py

import argparse
import csv
import itertools
import json
from pathlib import Path
from typing import Dict, List, Tuple

import matplotlib.pyplot as plt
import numpy as np


def load_counts_from_json(path: Path) -> Dict[str, float]:
    """
    Load token counts from a JSON file.

    Expected JSON format:
        {token: count, ...}
    """
    with path.open("r", encoding="utf-8") as f:
        data = json.load(f)

    if not isinstance(data, dict):
        raise ValueError(f"Unsupported JSON format in {path}")

    return {str(token): float(count) for token, count in data.items()}


def normalized_freqs(counts: Dict[str, float]) -> Dict[str, float]:
    """
    Convert raw counts to normalized frequencies.
    """
    total = float(sum(counts.values()))

    if total <= 0:
        return {}

    return {token: count / total for token, count in counts.items()}


def top_k_terms(counts: Dict[str, float], k: int) -> List[str]:
    """
    Return the top-k tokens by raw count.
    """
    return [
        token
        for token, _ in sorted(
            counts.items(),
            key=lambda item: item[1],
            reverse=True,
        )[:k]
    ]


def format_number_for_filename(value: float) -> str:
    """
    Format numeric values for cleaner filenames.

    Examples:
        2.0 -> "2"
        0.5 -> "0p5"
    """
    if float(value).is_integer():
        return str(int(value))

    return str(value).replace(".", "p")


def get_top_k_union(
    counts1: Dict[str, float],
    counts2: Dict[str, float],
    top_k: int,
) -> List[str]:
    """
    Return the union of the top-k tokens from two datasets.
    """
    return sorted(set(top_k_terms(counts1, top_k)) | set(top_k_terms(counts2, top_k)))


def compute_pair_frequencies(
    counts1: Dict[str, float],
    counts2: Dict[str, float],
    terms: List[str],
    pseudocount: float = 0.5,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute normalized frequencies for the selected terms.

    A pseudocount can be added so that terms missing from one dataset can still
    be displayed on log-scaled axes.
    """
    x_counts = np.array(
        [counts1.get(term, 0.0) + pseudocount for term in terms],
        dtype=float,
    )
    y_counts = np.array(
        [counts2.get(term, 0.0) + pseudocount for term in terms],
        dtype=float,
    )

    x_freqs = x_counts / x_counts.sum()
    y_freqs = y_counts / y_counts.sum()

    return x_freqs, y_freqs


def plot_pair_scatter(
    dataset1: str,
    dataset2: str,
    counts1: Dict[str, float],
    counts2: Dict[str, float],
    output_path: Path,
    *,
    top_k: int,
    factor_band: float = 2.0,
    log_axes: bool = True,
    xlim: Tuple[float, float] = (1e-7, 5e-3),
    ylim: Tuple[float, float] = (1e-7, 5e-3),
    min_freq: float = 1e-7,
    pseudocount: float = 0.5,
) -> None:
    """
    Plot normalized token frequencies for a pair of datasets.

    The plot compares the union of each dataset's top-k terms.
    Points inside the factor band are considered similarly frequent.
    Points outside the band are more frequent in one dataset than the other.
    """
    terms = get_top_k_union(counts1, counts2, top_k)
    x_freqs, y_freqs = compute_pair_frequencies(
        counts1=counts1,
        counts2=counts2,
        terms=terms,
        pseudocount=pseudocount,
    )

    keep = (x_freqs >= min_freq) | (y_freqs >= min_freq)

    terms = [term for term, should_keep in zip(terms, keep) if should_keep]
    x_freqs = x_freqs[keep]
    y_freqs = y_freqs[keep]

    print(
        f"{dataset1} vs {dataset2}: kept {len(terms)} terms "
        f"after min_freq={min_freq:g}"
    )

    fig, ax = plt.subplots(figsize=(10, 9))

    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    xmin, xmax = xlim
    ymin, ymax = ylim

    factor = float(factor_band)
    xs = np.array([xmin, xmax], dtype=float)

    y_lower = xs / factor
    y_upper = xs * factor

    ax.fill_between(xs, y_lower, y_upper, alpha=0.12, zorder=0)
    ax.fill_between(xs, ymin, y_lower, alpha=0.10, zorder=0)
    ax.fill_between(xs, y_upper, ymax, alpha=0.10, zorder=0)

    ratio = x_freqs / y_freqs

    within_band = (ratio >= 1.0 / factor) & (ratio <= factor)
    dataset1_higher = ratio > factor
    dataset2_higher = ratio < 1.0 / factor

    colors = np.empty(len(x_freqs), dtype=object)
    colors[within_band] = "#0E19BD"
    colors[dataset1_higher] = "#BD2716"
    colors[dataset2_higher] = "#1F8A70"

    ax.scatter(
        x_freqs,
        y_freqs,
        s=10,
        alpha=0.55,
        c=colors,
        linewidths=0,
    )

    ax.plot([xmin, xmax], [xmin, xmax], linewidth=1)
    ax.plot(
        [xmin, xmax],
        [factor * xmin, factor * xmax],
        linestyle="--",
        linewidth=1.5,
        alpha=0.9,
        color="red",
    )
    ax.plot(
        [xmin, xmax],
        [xmin / factor, xmax / factor],
        linestyle="--",
        linewidth=1.5,
        alpha=0.9,
        color="red",
    )

    if log_axes:
        ax.set_xscale("log")
        ax.set_yscale("log")

    ax.set_xlabel(f"{dataset1} normalized frequency", fontsize=24, fontweight="bold")
    ax.set_ylabel(f"{dataset2} normalized frequency", fontsize=24, fontweight="bold")

    ax.tick_params(axis="both", labelsize=20)
    ax.grid(True, which="both", alpha=0.25)

    plt.tight_layout()

    output_path.parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(output_path, dpi=200, bbox_inches="tight")
    plt.close(fig)


def pair_common_fraction(
    counts1: Dict[str, float],
    counts2: Dict[str, float],
    *,
    top_k: int,
    factor_band: float = 2.0,
    eps: float = 1e-12,
) -> Tuple[float, int, int]:
    """
    Compute the fraction of top-k union terms whose normalized frequencies are
    within a specified multiplicative factor of each other.
    """
    terms = get_top_k_union(counts1, counts2, top_k)
    total_terms = len(terms)

    if total_terms == 0:
        return float("nan"), 0, 0

    freqs1 = normalized_freqs(counts1)
    freqs2 = normalized_freqs(counts2)

    x_freqs = np.array([freqs1.get(term, 0.0) for term in terms], dtype=float)
    y_freqs = np.array([freqs2.get(term, 0.0) for term in terms], dtype=float)

    ratio = (x_freqs + eps) / (y_freqs + eps)

    factor = float(factor_band)
    within_band = (ratio >= 1.0 / factor) & (ratio <= factor)

    terms_within_band = int(within_band.sum())
    common_fraction = terms_within_band / total_terms

    return common_fraction, terms_within_band, total_terms


def save_summary_csv(
    results: List[Dict[str, object]],
    output_path: Path,
) -> None:
    """
    Save pairwise comparison results to a CSV file.
    """
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with output_path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "dataset1",
                "dataset2",
                "common_fraction",
                "within_band",
                "total_terms",
            ],
        )
        writer.writeheader()
        writer.writerows(results)


def run_all_pairs(
    datasets: List[str],
    dataset_to_json: Dict[str, Path],
    output_dir: Path,
    *,
    top_k: int = 2000,
    factor_band: float = 2.0,
    eps: float = 1e-12,
    log_axes: bool = True,
    make_plots: bool = True,
    part_of_speech: str,
    tokenization: str,
    min_freq: float = 1e-7,
    pseudocount: float = 0.5,
) -> List[Dict[str, object]]:
    """
    Run pairwise normalized-frequency comparisons across datasets.
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    counts_by_dataset = {
        dataset: load_counts_from_json(dataset_to_json[dataset])
        for dataset in datasets
    }

    results: List[Dict[str, object]] = []

    factor_label = format_number_for_filename(factor_band)

    for dataset1, dataset2 in itertools.combinations(datasets, 2):
        counts1 = counts_by_dataset[dataset1]
        counts2 = counts_by_dataset[dataset2]

        common_fraction, within_band, total_terms = pair_common_fraction(
            counts1,
            counts2,
            top_k=top_k,
            factor_band=factor_band,
            eps=eps,
        )

        results.append(
            {
                "dataset1": dataset1,
                "dataset2": dataset2,
                "common_fraction": round(common_fraction, 6),
                "within_band": within_band,
                "total_terms": total_terms,
            }
        )

        print(
            f"{dataset1} vs {dataset2}: "
            f"within={within_band}/{total_terms} "
            f"fraction={common_fraction:.3f}"
        )

        if make_plots:
            output_path = output_dir / (
                f"{dataset1}_vs_{dataset2}"
                f"_scatter_top{top_k}"
                f"_factor{factor_label}"
                f"_{part_of_speech}_{tokenization}.jpg"
            )

            plot_pair_scatter(
                dataset1=dataset1,
                dataset2=dataset2,
                counts1=counts1,
                counts2=counts2,
                output_path=output_path,
                top_k=top_k,
                factor_band=factor_band,
                log_axes=log_axes,
                min_freq=min_freq,
                pseudocount=pseudocount,
            )

    summary_path = output_dir / (
        f"pair_common_fraction"
        f"_top{top_k}"
        f"_factor{factor_label}"
        f"_{part_of_speech}_{tokenization}.csv"
    )

    save_summary_csv(results, summary_path)

    print(f"[OK] Saved summary: {summary_path}")

    return results


def build_dataset_paths(
    base_dir: Path,
    datasets: List[str],
    part_of_speech: str,
    tokenization: str,
) -> Dict[str, Path]:
    """
    Build paths to count JSON files for each dataset.
    """
    return {
        dataset: base_dir / dataset / f"counts_{part_of_speech}_{tokenization}.json"
        for dataset in datasets
    }


def validate_input_files(dataset_to_json: Dict[str, Path]) -> None:
    """
    Check that all expected input files exist before running comparisons.
    """
    missing_files = [
        str(path)
        for path in dataset_to_json.values()
        if not path.exists()
    ]

    if missing_files:
        missing_text = "\n".join(missing_files)
        raise FileNotFoundError(f"Missing count JSON file(s):\n{missing_text}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Compare normalized token frequencies across datasets using "
            "pairwise scatter plots and summary statistics."
        )
    )

    parser.add_argument(
        "--base-dir",
        type=Path,
        required=True,
        help="Directory containing one subdirectory per dataset.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        required=True,
        help="Directory where plots and CSV summaries will be saved.",
    )
    parser.add_argument(
        "--datasets",
        nargs="+",
        required=True,
        help="Dataset names to compare.",
    )
    parser.add_argument(
        "--part-of-speech",
        default="all",
        help="Part-of-speech label used in the count filenames.",
    )
    parser.add_argument(
        "--tokenization",
        default="lemma",
        help="Tokenization label used in the count filenames.",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=50000,
        help="Number of top terms to use from each dataset.",
    )
    parser.add_argument(
        "--factor-band",
        type=float,
        default=2.0,
        help="Frequency ratio band used to define similarly frequent terms.",
    )
    parser.add_argument(
        "--min-freq",
        type=float,
        default=1e-7,
        help="Minimum plotted normalized frequency.",
    )
    parser.add_argument(
        "--pseudocount",
        type=float,
        default=0.5,
        help="Pseudocount added before plotting frequencies.",
    )
    parser.add_argument(
        "--eps",
        type=float,
        default=1e-12,
        help="Small value added when computing frequency ratios.",
    )
    parser.add_argument(
        "--no-log-axes",
        action="store_true",
        help="Disable log-scaled plot axes.",
    )
    parser.add_argument(
        "--no-plots",
        action="store_true",
        help="Only save the CSV summary without scatter plots.",
    )

    return parser.parse_args()


def main() -> None:
    args = parse_args()

    dataset_to_json = build_dataset_paths(
        base_dir=args.base_dir,
        datasets=args.datasets,
        part_of_speech=args.part_of_speech,
        tokenization=args.tokenization,
    )

    validate_input_files(dataset_to_json)

    run_all_pairs(
        datasets=args.datasets,
        dataset_to_json=dataset_to_json,
        output_dir=args.output_dir,
        top_k=args.top_k,
        factor_band=args.factor_band,
        eps=args.eps,
        log_axes=not args.no_log_axes,
        make_plots=not args.no_plots,
        part_of_speech=args.part_of_speech,
        tokenization=args.tokenization,
        min_freq=args.min_freq,
        pseudocount=args.pseudocount,
    )


if __name__ == "__main__":
    main()