#!/usr/bin/env python3
"""
Disease-colored CO-SNE for RadIR/RadIR hyperbolic embeddings.

Uses the same hyperbolic CO-SNE pipeline as scripts/plot_cosne.py, then colors
points by disease while keeping modality markers:
  ○ Image,  △ Text
  Color = disease / class

Labeling strategies (--color_by):
  gt         ground-truth exclusive / priority labels (see --label_mode)
  inferred   disease-seeded DBSCAN on high-dim hyp image feats ("反推法")
  cosne2d    same seeded DBSCAN but on 2D CO-SNE image coords
  cluster_csv load precomputed cluster_disease from infer_disease_clusters.py

Usage (from RadIR/RadIR):
  # Ground-truth exclusive (clean)
  python scripts/plot_anat_cosne.py \\
    --load_emb outputs/cosne_emb.npz \\
    --color_by gt --label_mode exclusive --hide_other \\
    --out outputs/cosne_disease_clean.png

  # Reverse-inference clusters on CO-SNE disk
  python scripts/plot_anat_cosne.py \\
    --load_emb outputs/cosne_emb.npz \\
    --load_feats outputs/tsne_feats.npz \\
    --color_by inferred --hide_other \\
    --out outputs/cosne_disease_inferred.png
"""

from __future__ import annotations

import argparse
import os
import sys
from collections import Counter

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SCRIPTS = os.path.dirname(os.path.abspath(__file__))
RAD_IR_ROOT = os.path.join(ROOT, "Rad_IR")
for p in (ROOT, RAD_IR_ROOT, SCRIPTS):
    if p not in sys.path:
        sys.path.insert(0, p)

import infer_disease_clusters as idc  # noqa: E402
import plot_cosne as cosne  # noqa: E402

DISEASE_NAMES = idc.DISEASE_NAMES
CLASS_ORDER = ["Normal"] + DISEASE_NAMES + ["Other"]

CLASS_COLORS = {
    "Normal": "#00C853",
    "Pneumothorax": "#AA00FF",
    "Edema": "#00B8D4",
    "Pneumonia": "#D50000",
    "Effusion": "#2962FF",
    "Cardiomegaly": "#FF6D00",
    "Atelectasis": "#FFD600",
    "Other": "#BDBDBD",
}

MODALITY_MARKER = {"Image": "o", "Text": "^"}


def style_for(cls: str, modality: str) -> tuple[str, str]:
    return CLASS_COLORS.get(cls, "#333333"), MODALITY_MARKER[modality]


def load_label_table(labels_csv: str) -> pd.DataFrame:
    df = pd.read_csv(labels_csv)
    needed = ["dicom_id", "normal_only"] + [f"{d}_pos" for d in DISEASE_NAMES]
    missing = [c for c in needed if c not in df.columns]
    if missing:
        raise KeyError(f"labels CSV missing columns: {missing}")
    return df.drop_duplicates(subset=["dicom_id"], keep="first").reset_index(drop=True)


def build_labels(
    df: pd.DataFrame,
    sample_idx: np.ndarray,
    label_mode: str,
) -> list[str]:
    """
    Map global row indices → class name.

    exclusive: single-label only; else "Other"
    priority:  Normal if normal_only, else first disease hit, else "Other"
    """
    labels: list[str] = []
    for i in sample_idx:
        row = df.iloc[int(i)]
        hits: list[str] = []
        if int(row["normal_only"]) == 1:
            hits.append("Normal")
        for d in DISEASE_NAMES:
            if int(row[f"{d}_pos"]) == 1:
                hits.append(d)

        if label_mode == "exclusive":
            labels.append(hits[0] if len(hits) == 1 else "Other")
        elif label_mode == "priority":
            if int(row["normal_only"]) == 1 and len(hits) == 1:
                labels.append("Normal")
            else:
                assigned = None
                for d in DISEASE_NAMES:
                    if int(row[f"{d}_pos"]) == 1:
                        assigned = d
                        break
                if assigned is None and int(row["normal_only"]) == 1:
                    assigned = "Normal"
                labels.append(assigned if assigned is not None else "Other")
        else:
            raise ValueError(label_mode)
    return labels


def _to_plot_label(name: str) -> str:
    if name in (None, "", "Unassigned", "Other"):
        return "Other"
    return str(name)


def labels_from_cluster_csv(
    cluster_csv: str, labels_csv: str, sample_idx: np.ndarray
) -> list[str]:
    """Map sample_idx → cluster_disease via dicom_id alignment."""
    gt = pd.read_csv(labels_csv).drop_duplicates(subset=["dicom_id"], keep="first")
    cl = pd.read_csv(cluster_csv).drop_duplicates(subset=["dicom_id"], keep="first")
    cl = cl.set_index("dicom_id")["cluster_disease"]
    out: list[str] = []
    for i in sample_idx:
        did = str(gt.iloc[int(i)]["dicom_id"])
        if did in cl.index:
            out.append(_to_plot_label(cl.loc[did]))
        else:
            out.append("Other")
    return out


def labels_inferred_highdim(
    feats_npz: str,
    labels_csv: str,
    sample_idx: np.ndarray,
    seed_mode: str,
    eps: float | None,
    min_samples: int,
    eps_quantile: float,
    min_cluster_size: int,
) -> tuple[list[str], list[dict]]:
    """
    Disease-seeded DBSCAN on full-set high-dim hyp image feats, then slice
    assignments to the CO-SNE subsample.
    """
    data = np.load(feats_npz, allow_pickle=True)
    feats = idc.prepare_image_feats(data, "hyperbolic")
    df = pd.read_csv(labels_csv).drop_duplicates(subset=["dicom_id"], keep="first")
    if len(df) != feats.shape[0]:
        raise RuntimeError(
            f"feats N={feats.shape[0]} vs labels rows={len(df)} "
            "(expect same valid.jsonl order)"
        )
    normal_only = df["normal_only"].astype(int).to_numpy().astype(bool)
    disease_mat = np.stack(
        [df[f"{d}_pos"].astype(int).to_numpy().astype(bool) for d in DISEASE_NAMES],
        axis=1,
    )
    print(f"===== Infer clusters on high-dim hyp feats ({seed_mode}) =====")
    assigned, _uid, stats = idc.infer_disease_clusters(
        feats,
        normal_only,
        disease_mat,
        seed_mode=seed_mode,
        eps=eps,
        min_samples=min_samples,
        eps_quantile=eps_quantile,
        min_cluster_size=min_cluster_size,
    )
    labels = [_to_plot_label(assigned[int(i)]) for i in sample_idx]
    return labels, stats


def labels_inferred_cosne2d(
    emb_img: np.ndarray,
    df: pd.DataFrame,
    sample_idx: np.ndarray,
    seed_mode: str,
    eps: float | None,
    min_samples: int,
    eps_quantile: float,
    min_cluster_size: int,
) -> tuple[list[str], list[dict]]:
    """
    Same disease-seeded DBSCAN, but distance is Euclidean in the 2D CO-SNE
    Poincaré plane (answers: does the disk layout form disease neighborhoods?).
    """
    n = emb_img.shape[0]
    normal_only = np.zeros(n, dtype=bool)
    disease_mat = np.zeros((n, len(DISEASE_NAMES)), dtype=bool)
    for k, i in enumerate(sample_idx):
        row = df.iloc[int(i)]
        normal_only[k] = bool(int(row["normal_only"]))
        for j, d in enumerate(DISEASE_NAMES):
            disease_mat[k, j] = bool(int(row[f"{d}_pos"]))

    # Use Euclidean metric on 2D coords via L2-normalize trick disabled:
    # cluster_subset uses cosine on L2-normalized vectors; for 2D Poincaré
    # points we pass raw coords and temporarily use a Euclidean DBSCAN path.
    print(f"===== Infer clusters on CO-SNE 2D ({seed_mode}) =====")
    assigned, _uid, stats = infer_clusters_euclidean_2d(
        emb_img,
        normal_only,
        disease_mat,
        seed_mode=seed_mode,
        eps=eps,
        min_samples=min_samples,
        eps_quantile=eps_quantile,
        min_cluster_size=min_cluster_size,
    )
    labels = [_to_plot_label(a) for a in assigned]
    return labels, stats


def infer_clusters_euclidean_2d(
    coords: np.ndarray,
    normal_only: np.ndarray,
    disease_mat: np.ndarray,
    seed_mode: str,
    eps: float | None,
    min_samples: int,
    eps_quantile: float,
    min_cluster_size: int,
) -> tuple[np.ndarray, np.ndarray, list[dict]]:
    """DBSCAN with Euclidean metric on 2D coords (within disease seeds)."""
    from collections import defaultdict

    from sklearn.cluster import DBSCAN
    from sklearn.neighbors import NearestNeighbors

    n = coords.shape[0]
    excl = idc.exclusive_labels(normal_only, disease_mat)
    candidates: list[dict] = []
    class_names = ["Normal"] + DISEASE_NAMES

    for name in class_names:
        if seed_mode == "exclusive":
            member_idx = np.where(excl == name)[0]
        else:
            if name == "Normal":
                member_idx = np.where(normal_only)[0]
            else:
                j = DISEASE_NAMES.index(name)
                member_idx = np.where(disease_mat[:, j])[0]

        print(f"  Seed {name:14s} ({seed_mode}): {member_idx.size} images", end="")
        if member_idx.size < min_cluster_size:
            print(" → 0 clusters")
            continue

        sub = coords[member_idx].astype(np.float64)
        eps_used = eps
        if eps_used is None:
            k = min(min_samples, sub.shape[0] - 1)
            if k < 1:
                print(" → 0 clusters")
                continue
            nn = NearestNeighbors(n_neighbors=k + 1, metric="euclidean").fit(sub)
            dists, _ = nn.kneighbors(sub)
            eps_used = float(np.quantile(dists[:, -1], eps_quantile))

        labels_db = DBSCAN(
            eps=eps_used, min_samples=min_samples, metric="euclidean"
        ).fit_predict(sub)
        n_kept = 0
        for cid in sorted(set(labels_db.tolist()) - {-1}):
            local = np.where(labels_db == cid)[0]
            if local.size < min_cluster_size:
                continue
            global_idx = member_idx[local]
            candidates.append(
                {
                    "disease": name,
                    "global_indices": global_idx.astype(int),
                    "size": int(local.size),
                    "eps_used": float(eps_used),
                }
            )
            n_kept += 1
        print(f" → {n_kept} clusters ≥{min_cluster_size} (eps={eps_used:.4f})")

    claims: dict[int, list[tuple]] = defaultdict(list)
    for cid, c in enumerate(candidates):
        idxs = c["global_indices"]
        sub = coords[idxs]
        centroid = sub.mean(axis=0)
        dist = np.linalg.norm(sub - centroid, axis=1)
        for k, i in enumerate(idxs):
            is_excl = 1 if excl[i] == c["disease"] else 0
            claims[int(i)].append(
                (is_excl, c["size"], -float(dist[k]), cid, c["disease"])
            )

    assigned = np.array(["Unassigned"] * n, dtype=object)
    cluster_uid = np.full(n, -1, dtype=int)
    for i, opts in claims.items():
        opts.sort(reverse=True)
        _, _, _, cid, disease = opts[0]
        assigned[i] = disease
        cluster_uid[i] = cid

    stats = []
    for cid, c in enumerate(candidates):
        idxs = c["global_indices"]
        kept = idxs[assigned[idxs] == c["disease"]]
        stats.append(
            {
                "cluster_uid": cid,
                "disease": c["disease"],
                "size_raw": c["size"],
                "size_kept": int(kept.size),
                "eps_used": c["eps_used"],
            }
        )
    return assigned, cluster_uid, stats


def print_cluster_stats(stats: list[dict], labels: list[str]) -> None:
    print("----- kept clusters -----")
    by_d: dict[str, int] = Counter()
    for s in stats:
        if s.get("size_kept", 0) > 0:
            print(
                f"  uid={s['cluster_uid']:3d}  {s['disease']:14s}  "
                f"kept={s['size_kept']:4d}  raw={s['size_raw']:4d}  "
                f"eps={s['eps_used']:.4f}"
            )
            by_d[s["disease"]] += s["size_kept"]
    print("----- subsample assignment -----")
    counts = Counter(labels)
    for cls in CLASS_ORDER:
        if counts.get(cls, 0):
            print(f"  {cls:14s} {counts[cls]}")
    labeled = sum(v for k, v in counts.items() if k != "Other")
    print(f"  labeled={labeled}/{len(labels)}  Other={counts.get('Other', 0)}")


def run_or_load_cosne(args) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]:
    """
    Returns emb_img [M,2], emb_txt [M,2], sample_idx [M], curv.
    """
    if args.load_emb:
        print(f"Loading CO-SNE coords: {args.load_emb}")
        data = np.load(args.load_emb, allow_pickle=True)
        emb_img = np.asarray(data["emb_img"])
        emb_txt = np.asarray(data["emb_txt"])
        sample_idx = np.asarray(data["sample_idx"], dtype=np.int64)
        curv = float(data["curv"]) if "curv" in data.files else 1.0
        if emb_img.shape[0] != emb_txt.shape[0] or emb_img.shape[0] != len(sample_idx):
            raise ValueError("emb_img / emb_txt / sample_idx length mismatch")
        return emb_img, emb_txt, sample_idx, curv

    if not args.load_feats:
        raise ValueError("Provide --load_emb or --load_feats")

    print(f"Loading hyperbolic feats: {args.load_feats}")
    feats = np.load(args.load_feats, allow_pickle=True)
    if "hyp_img" not in feats.files:
        raise KeyError("Need hyp_img / hyp_txt in feats npz")
    hyp_img = torch.from_numpy(feats["hyp_img"])
    hyp_txt = torch.from_numpy(feats["hyp_txt"])
    curv = float(feats["curv"]) if "curv" in feats.files else 1.0

    max_samples = None if args.max_samples <= 0 else args.max_samples
    hyp_img, hyp_txt, sample_idx = cosne.subsample_pair(
        hyp_img, hyp_txt, max_samples, args.seed
    )
    print(f"  subsampled to {len(sample_idx)} pairs, curv={curv:.4f}")

    with torch.no_grad():
        poi_img = cosne.lorentz_space_to_poincare(hyp_img.float(), curv).numpy()
        poi_txt = cosne.lorentz_space_to_poincare(hyp_txt.float(), curv).numpy()

    X = np.concatenate([poi_img, poi_txt], axis=0)
    model = cosne.COSNE(
        n_components=2,
        perplexity=args.perplexity,
        n_iter=args.n_iter,
        learning_rate=args.learning_rate,
        learning_rate_h=args.learning_rate_h,
        early_exaggeration=args.early_exaggeration,
        student_t_gamma=args.student_t_gamma,
        seed=args.seed,
        device=args.cosne_device,
        verbose=True,
    )
    Y = model.fit_transform(X)
    n = poi_img.shape[0]
    emb_img, emb_txt = Y[:n], Y[n:]

    if args.save_emb:
        os.makedirs(
            os.path.dirname(os.path.abspath(args.save_emb)) or ".", exist_ok=True
        )
        np.savez(
            args.save_emb,
            emb_img=emb_img,
            emb_txt=emb_txt,
            sample_idx=sample_idx,
            curv=curv,
        )
        print(f"Saved CO-SNE coords to {args.save_emb}")

    return emb_img, emb_txt, sample_idx, curv


def plot_disease_poincare(
    emb_img: np.ndarray,
    emb_txt: np.ndarray,
    labels: list[str],
    out_path: str,
    title: str,
    dpi: int,
    hide_other: bool,
    max_pair_lines: int,
) -> None:
    n = emb_img.shape[0]
    assert len(labels) == n

    fig, ax = plt.subplots(figsize=(9, 9))
    ax.add_patch(plt.Circle((0, 0), 1.0, color="black", fill=False, linewidth=1.2))
    for r in (0.25, 0.5, 0.75):
        ax.add_patch(
            plt.Circle((0, 0), r, color="gray", fill=False, linewidth=0.4, alpha=0.35)
        )

    # optional pair connectors (only for non-Other if hide_other)
    n_lines = min(max_pair_lines, n)
    drawn = 0
    for i in range(n):
        if drawn >= n_lines:
            break
        if hide_other and labels[i] == "Other":
            continue
        ax.plot(
            [emb_img[i, 0], emb_txt[i, 0]],
            [emb_img[i, 1], emb_txt[i, 1]],
            color="gray",
            alpha=0.12,
            linewidth=0.5,
            zorder=0,
        )
        drawn += 1

    # draw Other first (background), then each disease
    draw_order = ["Other"] + [c for c in CLASS_ORDER if c != "Other"]
    for cls in draw_order:
        if hide_other and cls == "Other":
            continue
        mask = np.array([l == cls for l in labels], dtype=bool)
        if not mask.any():
            continue
        for modality, emb in (("Image", emb_img), ("Text", emb_txt)):
            color, marker = style_for(cls, modality)
            is_other = cls == "Other"
            ax.scatter(
                emb[mask, 0],
                emb[mask, 1],
                c=color,
                marker=marker,
                s=18 if is_other else (32 if modality == "Image" else 42),
                alpha=0.35 if is_other else 0.88,
                edgecolors="none" if is_other else "black",
                linewidths=0.0 if is_other else 0.3,
                label=f"{cls}-{modality} (n={int(mask.sum())})",
                zorder=1 if is_other else (2 if modality == "Image" else 3),
            )

    # radius summary on labeled (non-Other) points
    labeled_mask = np.array([l != "Other" for l in labels], dtype=bool)
    if labeled_mask.any():
        r_img = np.linalg.norm(emb_img[labeled_mask], axis=1).mean()
        r_txt = np.linalg.norm(emb_txt[labeled_mask], axis=1).mean()
        subtitle = f"labeled mean radius  Image={r_img:.3f}  Text={r_txt:.3f}"
    else:
        subtitle = "no exclusive labels in subsample"

    ax.set_title(f"{title}\n{subtitle}")
    ax.set_xlim(-1.05, 1.05)
    ax.set_ylim(-1.05, 1.05)
    ax.set_aspect("equal")
    ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1.0), fontsize=8, frameon=True)
    ax.axis("off")
    fig.tight_layout()
    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
    fig.savefig(out_path, dpi=dpi, bbox_inches="tight")
    if out_path.lower().endswith(".png"):
        fig.savefig(out_path[:-4] + ".pdf", dpi=dpi, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved disease CO-SNE to {out_path}")


def parse_args():
    p = argparse.ArgumentParser(
        description="Disease-colored CO-SNE (hyperbolic) for RadIR"
    )
    p.add_argument(
        "--load_emb",
        default=None,
        help="Cached CO-SNE coords from plot_cosne.py (--save_emb)",
    )
    p.add_argument(
        "--load_feats",
        default=None,
        help="Hyperbolic feats npz (required for --color_by inferred)",
    )
    p.add_argument(
        "--labels_csv",
        default=os.path.join(ROOT, "outputs/tsne_disease_labels.csv"),
    )
    p.add_argument(
        "--cluster_csv",
        default=os.path.join(ROOT, "outputs/tsne_cluster_disease.csv"),
        help="Precomputed cluster_disease CSV for --color_by cluster_csv",
    )
    p.add_argument(
        "--jsonl",
        default=os.path.join(ROOT, "dataset/valid.jsonl"),
        help="Only used to sanity-check length vs labels (optional)",
    )
    p.add_argument(
        "--color_by",
        choices=["gt", "inferred", "cosne2d", "cluster_csv"],
        default="gt",
        help="gt=ground-truth; inferred=反推 on high-dim; "
        "cosne2d=反推 on CO-SNE plane; cluster_csv=load CSV",
    )
    p.add_argument(
        "--label_mode",
        choices=["exclusive", "priority"],
        default="exclusive",
        help="For --color_by gt only",
    )
    p.add_argument(
        "--seed_mode",
        choices=["exclusive", "inclusive"],
        default="exclusive",
        help="Disease seed mode for inferred / cosne2d",
    )
    p.add_argument("--eps", type=float, default=None, help="DBSCAN eps (auto if None)")
    p.add_argument("--min_samples", type=int, default=5)
    p.add_argument("--eps_quantile", type=float, default=0.5)
    p.add_argument("--min_cluster_size", type=int, default=8)
    p.add_argument(
        "--hide_other",
        action="store_true",
        help="Do not draw Other/Unassigned points",
    )
    p.add_argument("--max_samples", type=int, default=500)
    p.add_argument("--perplexity", type=float, default=30.0)
    p.add_argument("--n_iter", type=int, default=800)
    p.add_argument("--learning_rate", type=float, default=0.05)
    p.add_argument("--learning_rate_h", type=float, default=5.0)
    p.add_argument("--early_exaggeration", type=float, default=4.0)
    p.add_argument("--student_t_gamma", type=float, default=0.1)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--cosne_device", default="cpu")
    p.add_argument("--max_pair_lines", type=int, default=40)
    p.add_argument(
        "--out",
        default=os.path.join(ROOT, "outputs/cosne_disease.png"),
    )
    p.add_argument("--save_emb", default=None)
    p.add_argument("--dpi", type=int, default=150)
    return p.parse_args()


def main():
    args = parse_args()
    if args.load_emb is None and args.load_feats is None:
        default_emb = os.path.join(ROOT, "outputs/cosne_emb.npz")
        default_feats = os.path.join(ROOT, "outputs/tsne_feats.npz")
        if os.path.isfile(default_emb):
            args.load_emb = default_emb
        elif os.path.isfile(default_feats):
            args.load_feats = default_feats
        else:
            raise FileNotFoundError(
                "Need --load_emb or --load_feats (or outputs/cosne_emb.npz)"
            )

    emb_img, emb_txt, sample_idx, curv = run_or_load_cosne(args)
    df = load_label_table(args.labels_csv)
    if len(df) < int(sample_idx.max()) + 1:
        raise RuntimeError(
            f"labels CSV has {len(df)} rows but sample_idx max={sample_idx.max()}"
        )

    stats: list[dict] = []
    if args.color_by == "gt":
        labels = build_labels(df, sample_idx, args.label_mode)
        print(f"===== Label mode: gt/{args.label_mode} (n={len(labels)}) =====")
        for cls in CLASS_ORDER:
            c = Counter(labels).get(cls, 0)
            if c:
                print(f"  {cls:14s} {c}")
        title = (
            f"CO-SNE + GT ({args.label_mode}) | hyperbolic | "
            f"n={len(labels)} | curv={curv:.3f}"
        )
    elif args.color_by == "cluster_csv":
        labels = labels_from_cluster_csv(args.cluster_csv, args.labels_csv, sample_idx)
        print(f"===== Color by cluster CSV: {args.cluster_csv} =====")
        for cls in CLASS_ORDER:
            c = Counter(labels).get(cls, 0)
            if c:
                print(f"  {cls:14s} {c}")
        title = (
            f"CO-SNE + inferred clusters (CSV) | hyperbolic | "
            f"n={len(labels)} | curv={curv:.3f}"
        )
    elif args.color_by == "inferred":
        feats_path = args.load_feats or os.path.join(ROOT, "outputs/tsne_feats.npz")
        if not os.path.isfile(feats_path):
            raise FileNotFoundError(f"--color_by inferred needs feats: {feats_path}")
        labels, stats = labels_inferred_highdim(
            feats_path,
            args.labels_csv,
            sample_idx,
            seed_mode=args.seed_mode,
            eps=args.eps,
            min_samples=args.min_samples,
            eps_quantile=args.eps_quantile,
            min_cluster_size=args.min_cluster_size,
        )
        print_cluster_stats(stats, labels)
        title = (
            f"CO-SNE + inferred (high-dim {args.seed_mode}) | hyperbolic | "
            f"n={len(labels)} | curv={curv:.3f}"
        )
    elif args.color_by == "cosne2d":
        labels, stats = labels_inferred_cosne2d(
            emb_img,
            df,
            sample_idx,
            seed_mode=args.seed_mode,
            eps=args.eps,
            min_samples=args.min_samples,
            eps_quantile=args.eps_quantile,
            min_cluster_size=args.min_cluster_size,
        )
        print_cluster_stats(stats, labels)
        title = (
            f"CO-SNE + inferred (2D disk {args.seed_mode}) | hyperbolic | "
            f"n={len(labels)} | curv={curv:.3f}"
        )
    else:
        raise ValueError(args.color_by)

    plot_disease_poincare(
        emb_img,
        emb_txt,
        labels,
        args.out,
        title=title,
        dpi=args.dpi,
        hide_other=args.hide_other,
        max_pair_lines=args.max_pair_lines,
    )


if __name__ == "__main__":
    main()