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

Unlike scripts/plot_tsne.py (log_map0 → Euclidean t-SNE), this keeps points in
hyperbolic geometry: Lorentz space components are mapped to the Poincaré disk,
then CO-SNE (CVPR 2022) reduces them to 2D while preserving:
  - local hyperbolic neighborhoods
  - each point's geodesic distance to the origin (hierarchy / radius)

Reference: Guo et al., "CO-SNE: Dimensionality Reduction and Visualization for
Hyperbolic Data", CVPR 2022.  https://github.com/yunhuiguo/CO-SNE

This file ships a self-contained CO-SNE core (modern PyTorch) so we do not
depend on the official repo's old sklearn / hyptorch pins.

Usage (from RadIR/RadIR):
  python scripts/plot_cosne.py \\
    --load_feats outputs/tsne_feats.npz \\
    --out outputs/cosne_radir.png

  # Smaller / faster smoke run:
  python scripts/plot_cosne.py --load_feats outputs/tsne_feats.npz \\
    --max_samples 200 --n_iter 500 --out outputs/cosne_radir_test.png
"""

from __future__ import annotations

import argparse
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from transformers import AutoModel, AutoTokenizer

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

from radir import lorentz as L
from radir.RadIR import RADIR, l2norm
from radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box
from transformer_maskgit.transformer_maskgit.ctvit import CTViT

EPS = 1e-8
BALL_EPS = 1e-5


# ---------------------------------------------------------------------------
# Lorentz (RadIR) → Poincaré ball (CO-SNE native space)
# ---------------------------------------------------------------------------

def lorentz_space_to_poincare(z_space: torch.Tensor, curv: float) -> torch.Tensor:
    """
    Map Lorentz *space* components (RadIR/MERU) to the Poincaré ball of
    curvature -curv, then rescale into the unit disk for visualization.

    z_space: [N, D]
    returns: [N, D] with ||u|| < 1
    """
    c = float(max(curv, EPS))
    sqrt_c = c**0.5
    # time component: t = sqrt(1/c + ||z||^2)
    t = torch.sqrt(1.0 / c + z_space.pow(2).sum(dim=-1, keepdim=True))
    # Poincaré coords in ball of radius 1/sqrt(c):
    #   u = z / (t + 1/sqrt(c))
    u = z_space / (t + 1.0 / sqrt_c)
    # rescale to unit disk: û = sqrt(c) * u
    u = sqrt_c * u
    return project_to_ball(u)


def project_to_ball(u: torch.Tensor, eps: float = BALL_EPS) -> torch.Tensor:
    """Project points into the open unit disk."""
    norm = u.norm(dim=-1, keepdim=True).clamp_min(EPS)
    max_norm = 1.0 - eps
    scale = torch.clamp(max_norm / norm, max=1.0)
    return u * scale


def poincare_dist(x: torch.Tensor, y: torch.Tensor, eps: float = EPS) -> torch.Tensor:
    """
    Pairwise geodesic distance on the unit Poincaré disk (curvature -1).
    x: [N, D], y: [M, D] → [N, M]
    """
    x2 = x.pow(2).sum(dim=-1, keepdim=True)  # [N,1]
    y2 = y.pow(2).sum(dim=-1, keepdim=True).T  # [1,M]
    xy = x @ y.T
    num = (x2 + y2 - 2.0 * xy).clamp_min(0.0)
    den = ((1.0 - x2).clamp_min(eps) * (1.0 - y2).clamp_min(eps))
    z = 1.0 + 2.0 * num / den
    return torch.acosh(z.clamp_min(1.0 + eps))


def poincare_dist_to_origin(x: torch.Tensor, eps: float = EPS) -> torch.Tensor:
    """Geodesic distance from each point to the origin on the unit disk."""
    r2 = x.pow(2).sum(dim=-1).clamp(max=1.0 - eps)
    return torch.atanh(torch.sqrt(r2).clamp(max=1.0 - eps))


# ---------------------------------------------------------------------------
# CO-SNE core (self-contained)
# ---------------------------------------------------------------------------

def _binary_search_perplexity(
    dist_row: np.ndarray,
    perplexity: float,
    tol: float = 1e-5,
    max_iter: int = 50,
) -> np.ndarray:
    """Given distances to other points, find Gaussian bandwidth for target perplexity."""
    n = dist_row.shape[0]
    beta = 1.0  # = 1/(2σ^2)
    betamin, betamax = -np.inf, np.inf
    log_perp = np.log(perplexity)

    for _ in range(max_iter):
        p = np.exp(-dist_row * beta)
        sum_p = np.maximum(p.sum(), EPS)
        p = p / sum_p
        entropy = -np.sum(p * np.log(np.maximum(p, EPS)))
        diff = entropy - log_perp
        if abs(diff) < tol:
            break
        if diff > 0:
            betamin = beta
            beta = beta * 2.0 if np.isinf(betamax) else (beta + betamax) / 2.0
        else:
            betamax = beta
            beta = beta / 2.0 if np.isinf(betamin) else (beta + betamin) / 2.0
    p = np.exp(-dist_row * beta)
    p = p / np.maximum(p.sum(), EPS)
    return p.astype(np.float64)


def joint_probabilities_from_distances(
    dist: np.ndarray, perplexity: float
) -> np.ndarray:
    """
    Build symmetric joint probabilities P from a hyperbolic distance matrix,
    matching the t-SNE / CO-SNE high-dim affinity construction.
    """
    n = dist.shape[0]
    P = np.zeros((n, n), dtype=np.float64)
    for i in range(n):
        d = dist[i].copy()
        d[i] = np.inf
        P[i] = _binary_search_perplexity(d, perplexity)
    P = P + P.T
    P = np.maximum(P / np.maximum(P.sum(), EPS), EPS)
    np.fill_diagonal(P, 0.0)
    return P


class COSNE:
    """
    CO-SNE: hyperbolic t-SNE with origin-distance preservation.

    High-dim affinities from hyperbolic geodesic distances.
    Low-dim affinities: hyperbolic Cauchy  Q_ij ∝ (1 + (d_ij/γ)^2)^-1
    Extra loss: || d_H(x_i,0) - d_H(y_i,0) ||^2  (hierarchy / radius)
    """

    def __init__(
        self,
        n_components: int = 2,
        perplexity: float = 30.0,
        n_iter: int = 1000,
        learning_rate: float = 1.0,
        learning_rate_h: float = 0.1,
        early_exaggeration: float = 4.0,
        student_t_gamma: float = 0.1,
        exaggeration_iters: int = 250,
        seed: int = 42,
        device: str = "cpu",
        verbose: bool = True,
    ):
        self.n_components = n_components
        self.perplexity = perplexity
        self.n_iter = n_iter
        self.learning_rate = learning_rate
        self.learning_rate_h = learning_rate_h
        self.early_exaggeration = early_exaggeration
        self.student_t_gamma = student_t_gamma
        self.exaggeration_iters = exaggeration_iters
        self.seed = seed
        self.device = device
        self.verbose = verbose

    def fit_transform(self, X_poincare: np.ndarray) -> np.ndarray:
        """
        X_poincare: [N, D] points already on the unit Poincaré disk.
        Returns Y: [N, 2] on the unit disk.
        """
        X = torch.as_tensor(X_poincare, dtype=torch.float64)
        X = project_to_ball(X)
        n, d = X.shape
        device = torch.device(self.device)

        if self.verbose:
            print(f"[CO-SNE] n={n}, dim={d}, perplexity={self.perplexity}")

        # --- high-dim hyperbolic distances & P ---
        with torch.no_grad():
            dist_h = poincare_dist(X, X).cpu().numpy()
        np.fill_diagonal(dist_h, 0.0)
        max_perp = max(2.0, (n - 1) / 3.0)
        perp = min(self.perplexity, max_perp)
        P = joint_probabilities_from_distances(dist_h, perp)
        P = torch.as_tensor(P, dtype=torch.float64, device=device)

        # origin distances in high-dim (hierarchy target)
        with torch.no_grad():
            d0_high = poincare_dist_to_origin(X).to(device)

        # --- init low-dim embedding from first Poincaré coords + radius match ---
        rng = np.random.default_rng(self.seed)
        if d >= 2:
            Y0 = X[:, :2].cpu().numpy() + rng.normal(0, 1e-4, size=(n, 2))
        else:
            Y0 = rng.normal(0, 1e-3, size=(n, self.n_components))
        r_high = np.tanh(d0_high.cpu().numpy())  # (0,1) proxy for hierarchy radius
        r_y = np.linalg.norm(Y0, axis=1) + EPS
        Y0 = Y0 * (r_high / r_y)[:, None]
        Y = torch.as_tensor(Y0, dtype=torch.float64, device=device)
        Y = project_to_ball(Y)
        Y = torch.nn.Parameter(Y)

        opt = torch.optim.Adam([Y], lr=self.learning_rate)
        # after early exaggeration, decay LR a bit for radius stability
        sched = torch.optim.lr_scheduler.StepLR(opt, step_size=max(self.exaggeration_iters, 1), gamma=0.5)
        gamma = float(self.student_t_gamma)

        for it in range(self.n_iter):
            opt.zero_grad()
            Y_proj = project_to_ball(Y)

            # hyperbolic distances in 2D
            dist_y = poincare_dist(Y_proj, Y_proj)
            # hyperbolic Cauchy affinities (CO-SNE / HTSNE style)
            q_num = 1.0 / (1.0 + (dist_y / (gamma**2 + EPS)).pow(2))
            q_num = q_num.fill_diagonal_(0.0)
            Q = q_num / q_num.sum().clamp_min(EPS)
            Q = Q.clamp_min(EPS)

            # early exaggeration on P
            if it < self.exaggeration_iters:
                P_use = P * self.early_exaggeration
                P_use = P_use / P_use.sum()
            else:
                P_use = P

            kl = (P_use * (torch.log(P_use.clamp_min(EPS)) - torch.log(Q))).sum()

            # hierarchy loss: match Poincaré *radii* (more stable than raw d_H
            # when high-dim radii are small, as with RadIR's near-constant image norms)
            r_low = Y_proj.norm(dim=-1).clamp(max=1.0 - BALL_EPS)
            r_high = torch.tanh(d0_high).clamp(max=1.0 - BALL_EPS)
            h_loss = F.mse_loss(r_low, r_high)

            loss = kl + self.learning_rate_h * h_loss
            loss.backward()
            # clip grads — Poincaré optimization is sensitive near the boundary
            torch.nn.utils.clip_grad_norm_([Y], max_norm=1.0)
            opt.step()

            with torch.no_grad():
                Y.data.copy_(project_to_ball(Y.data))
            sched.step()

            if self.verbose and (it % 100 == 0 or it == self.n_iter - 1):
                r_mean = Y.data.norm(dim=-1).mean().item()
                print(
                    f"[CO-SNE] iter {it:4d}/{self.n_iter}  "
                    f"KL={kl.item():.4f}  H={h_loss.item():.6f}  "
                    f"loss={loss.item():.4f}  mean_r={r_mean:.4f}"
                )

        return project_to_ball(Y.detach()).cpu().numpy()


# ---------------------------------------------------------------------------
# Feature extraction (hyperbolic only; aligned with Evaluate / plot_tsne)
# ---------------------------------------------------------------------------

def preprocess_image_like_forward(image: torch.Tensor) -> torch.Tensor:
    if image.ndim == 3:
        image = image.unsqueeze(0)
    if image.ndim == 4:
        if image.shape[1] != 1:
            image = image.mean(dim=1, keepdim=True)
    elif image.ndim == 5:
        if image.shape[1] != 1:
            image = image.mean(dim=1, keepdim=True)
    else:
        raise ValueError(f"Unexpected image ndim={image.ndim}")
    return image


@torch.no_grad()
def extract_hyperbolic_features(model, dataloader, device, modal_embedding=False):
    """Only hyperbolic features (project_img / project_txt)."""
    all_hyp_img, all_hyp_txt = [], []
    model.eval()
    for batch in dataloader:
        imgs = preprocess_image_like_forward(
            batch["imgs"].to(device, non_blocking=True)
        )
        enc = model.visual_transformer(
            imgs,
            return_encoded_tokens=True,
            modal_embedding=modal_embedding,
            modal_indexs=None,
            is_condition=False,
        )
        if enc.dim() == 2:
            image_global = l2norm(enc)
        elif enc.dim() == 3:
            image_global = l2norm(enc[:, 0, :])
        else:
            raise ValueError(f"Unexpected visual shape: {enc.shape}")

        text_out = model.text_transformer(
            input_ids=batch["caption_ids"].to(device, non_blocking=True),
            token_type_ids=batch["token_type_ids"].to(device, non_blocking=True),
            attention_mask=batch["attention_mask"].to(device, non_blocking=True),
        )
        text_global = text_out[0][:, 0, :]

        all_hyp_img.append(model.project_img(image_global).cpu())
        all_hyp_txt.append(model.project_txt(text_global).cpu())

    return torch.cat(all_hyp_img, dim=0), torch.cat(all_hyp_txt, dim=0)


def build_model(args, device: str) -> RADIR:
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path, trust_remote_code=True, local_files_only=True
    )
    text_model = AutoModel.from_pretrained(
        args.model_path, trust_remote_code=True, local_files_only=True
    )
    image_encoder = CTViT(
        dim=768,
        codebook_size=8192,
        image_size=args.imsize,
        patch_size=16,
        temporal_patch_size=10,
        spatial_depth=8,
        temporal_depth=6,
        cls_depth=4,
        dim_head=32,
        heads=8,
        channels=1,
    )
    model = RADIR(
        tokenizer=tokenizer,
        image_encoder=image_encoder,
        text_encoder=text_model,
        dim_text=768,
        dim_image=768,
        dim_latent=768,
        use_mlm=False,
        use_all_token_embeds=False,
    )
    print(f"Loading checkpoint: {args.radir_ckpt}")
    model.load(args.radir_ckpt)
    model.to(device)
    return model


def subsample_pair(
    hyp_img: torch.Tensor,
    hyp_txt: torch.Tensor,
    max_samples: int | None,
    seed: int,
):
    n = hyp_img.shape[0]
    if max_samples is None or max_samples <= 0 or max_samples >= n:
        return hyp_img, hyp_txt, np.arange(n)
    rng = np.random.default_rng(seed)
    idx = np.sort(rng.choice(n, size=max_samples, replace=False))
    return hyp_img[idx], hyp_txt[idx], idx


# ---------------------------------------------------------------------------
# Plotting
# ---------------------------------------------------------------------------

def plot_poincare_disk(
    emb_img: np.ndarray,
    emb_txt: np.ndarray,
    out_path: str,
    title: str,
    dpi: int,
    max_pair_lines: int = 40,
) -> None:
    fig, ax = plt.subplots(figsize=(8, 8))
    circle = plt.Circle((0, 0), 1.0, color="black", fill=False, linewidth=1.2)
    ax.add_patch(circle)

    # light radial rings to emphasize hierarchy
    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.4)
        )

    n = emb_img.shape[0]
    n_lines = min(max_pair_lines, n)
    for i in range(n_lines):
        ax.plot(
            [emb_img[i, 0], emb_txt[i, 0]],
            [emb_img[i, 1], emb_txt[i, 1]],
            color="gray",
            alpha=0.15,
            linewidth=0.5,
            zorder=0,
        )

    ax.scatter(
        emb_img[:, 0],
        emb_img[:, 1],
        c="#2962FF",
        marker="o",
        s=28,
        alpha=0.85,
        edgecolors="black",
        linewidths=0.25,
        label="Image",
        zorder=2,
    )
    ax.scatter(
        emb_txt[:, 0],
        emb_txt[:, 1],
        c="#D50000",
        marker="^",
        s=36,
        alpha=0.85,
        edgecolors="black",
        linewidths=0.25,
        label="Text",
        zorder=3,
    )

    # report mean radius (hierarchy signal)
    r_img = np.linalg.norm(emb_img, axis=1).mean()
    r_txt = np.linalg.norm(emb_txt, axis=1).mean()
    ax.set_title(
        f"{title}\n"
        f"mean Poincaré radius  Image={r_img:.3f}  Text={r_txt:.3f}"
    )
    ax.set_xlim(-1.05, 1.05)
    ax.set_ylim(-1.05, 1.05)
    ax.set_aspect("equal")
    ax.legend(loc="upper right", 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")
    # also save pdf sibling if png
    if out_path.lower().endswith(".png"):
        fig.savefig(out_path[:-4] + ".pdf", dpi=dpi, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved CO-SNE plot to {out_path}")
    print(f"  mean radius Image={r_img:.4f}, Text={r_txt:.4f}  "
          f"(Text closer to origin ⇒ hierarchy)")


def print_radius_stats(name: str, poincare: np.ndarray) -> None:
    r = np.linalg.norm(poincare, axis=1)
    d0 = np.arctanh(np.clip(r, 0, 1 - 1e-6))
    print(
        f"  {name:6s}  radius mean={r.mean():.4f} std={r.std():.4f}  "
        f"|  d_H(·,0) mean={d0.mean():.4f} std={d0.std():.4f}"
    )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(
        description="CO-SNE visualization of RadIR hyperbolic embeddings"
    )
    p.add_argument("--device", default="cuda")
    p.add_argument(
        "--model_path",
        default=os.path.join(
            ROOT, "hf_models/microsoft/BiomedVLP-CXR-BERT-specialized"
        ),
    )
    p.add_argument(
        "--radir_ckpt",
        default=os.path.join(ROOT, "outputs/radir_final.pt"),
    )
    p.add_argument("--dataset_root", default=os.path.join(ROOT, "dataset"))
    p.add_argument("--caption_file", default="valid.jsonl")
    p.add_argument(
        "--cig_matched_dir",
        default=os.path.join(ROOT, "dataset/processed/matched"),
    )
    p.add_argument("--split", default="valid")
    p.add_argument("--batch_size", type=int, default=16)
    p.add_argument("--num_workers", type=int, default=8)
    p.add_argument("--imsize", type=int, default=224)
    p.add_argument("--max_words", type=int, default=128)
    p.add_argument("--modal_embedding", action="store_true")
    p.add_argument(
        "--load_feats",
        default=None,
        help="Load outputs/tsne_feats.npz (uses hyp_img / hyp_txt / curv only)",
    )
    p.add_argument("--save_feats", default=None)
    p.add_argument(
        "--max_samples",
        type=int,
        default=3120,
        help="Subsample pairs for CO-SNE speed; 0 = all (slow, O(N^2))",
    )
    p.add_argument("--perplexity", type=float, default=30.0)
    p.add_argument("--n_iter", type=int, default=1000)
    p.add_argument("--learning_rate", type=float, default=0.05)
    p.add_argument(
        "--learning_rate_h",
        type=float,
        default=5.0,
        help="Weight for origin-radius (hierarchy) loss; keep high so radii are preserved",
    )
    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("--max_pair_lines", type=int, default=40)
    p.add_argument(
        "--out",
        default=os.path.join(ROOT, "outputs/cosne_radir.png"),
    )
    p.add_argument(
        "--save_emb",
        default=None,
        help="Optional .npz to cache 2D CO-SNE Poincaré coords",
    )
    p.add_argument("--dpi", type=int, default=150)
    p.add_argument(
        "--cosne_device",
        default="cpu",
        help="Device for CO-SNE optimization (cpu is usually fine for N~1k)",
    )
    return p.parse_args()


def main():
    args = parse_args()
    device = args.device

    if args.load_feats:
        print(f"Loading cached features: {args.load_feats}")
        data = np.load(args.load_feats, allow_pickle=True)
        if "hyp_img" not in data.files or "hyp_txt" not in data.files:
            raise KeyError(
                f"{args.load_feats} must contain hyp_img / hyp_txt "
                "(hyperbolic features). Re-extract with plot_tsne.py --save_feats."
            )
        hyp_img = torch.from_numpy(data["hyp_img"])
        hyp_txt = torch.from_numpy(data["hyp_txt"])
        curv = float(data["curv"]) if "curv" in data.files else 1.0
        print(f"  hyp_img={tuple(hyp_img.shape)}, curv={curv:.6f}")
    else:
        model = build_model(args, device)
        curv = float(model.get_curv().detach().cpu().item())
        print(f"Curvature = {curv:.6f}")
        dataset = HyCoClipMimicBoxDataset_JPG(
            split=args.split,
            tokenizer=model.tokenizer,
            dataset_root=args.dataset_root,
            cig_matched_dir=args.cig_matched_dir,
            caption_file=args.caption_file,
            imsize=args.imsize,
            max_words=args.max_words,
            transform=None,
        )
        dataloader = DataLoader(
            dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            pin_memory=True,
            collate_fn=collate_hmclip_box,
        )
        print(f"Extracting hyperbolic features from {len(dataset)} samples ...")
        hyp_img, hyp_txt = extract_hyperbolic_features(
            model, dataloader, device, modal_embedding=args.modal_embedding
        )
        if args.save_feats:
            os.makedirs(
                os.path.dirname(os.path.abspath(args.save_feats)) or ".",
                exist_ok=True,
            )
            np.savez(
                args.save_feats,
                hyp_img=hyp_img.numpy(),
                hyp_txt=hyp_txt.numpy(),
                curv=curv,
            )
            print(f"Saved hyperbolic features to {args.save_feats}")

    max_samples = None if args.max_samples <= 0 else args.max_samples
    hyp_img, hyp_txt, idx = subsample_pair(
        hyp_img, hyp_txt, max_samples, args.seed
    )
    print(f"Using {hyp_img.shape[0]} pairs for CO-SNE")

    # Lorentz space → Poincaré (unit disk)
    with torch.no_grad():
        poi_img = lorentz_space_to_poincare(hyp_img.float(), curv).numpy()
        poi_txt = lorentz_space_to_poincare(hyp_txt.float(), curv).numpy()

    print("High-dim Poincaré radius / origin-distance stats:")
    print_radius_stats("Image", poi_img)
    print_radius_stats("Text", poi_txt)

    # Joint CO-SNE on Image + Text (preserves cross-modal + hierarchy)
    X = np.concatenate([poi_img, poi_txt], axis=0)
    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 = cosne.fit_transform(X)
    n = poi_img.shape[0]
    emb_img, emb_txt = Y[:n], Y[n:]

    print("Low-dim (CO-SNE) Poincaré radius / origin-distance stats:")
    print_radius_stats("Image", emb_img)
    print_radius_stats("Text", emb_txt)

    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=idx,
            curv=curv,
        )
        print(f"Saved 2D embeddings to {args.save_emb}")

    ckpt = os.path.basename(args.radir_ckpt) if not args.load_feats else "cached"
    title = f"CO-SNE (hyperbolic) | RadIR | {ckpt} | n={n}"
    plot_poincare_disk(
        emb_img,
        emb_txt,
        args.out,
        title=title,
        dpi=args.dpi,
        max_pair_lines=args.max_pair_lines,
    )


if __name__ == "__main__":
    main()