import argparse
import os
from dataclasses import dataclass

import matplotlib.pyplot as plt
import mpmath
import numpy as np
import torch


DEFAULT_GAMMA = [
    14.134725,
    21.022040,
    25.010858,
    30.424876,
    32.935062,
    37.586178,
    40.918719,
    43.327073,
    48.005150,
    49.773832,
]


@dataclass
class CirclePatch:
    origin: complex
    radius: float
    roots: np.ndarray
    mapped_roots: np.ndarray
    coeffs: np.ndarray
    fit_loss: float


@dataclass
class PatchLayout:
    n_circles: int
    roots_per_circle: int
    patches: list
    recovered: np.ndarray
    mean_error: float
    max_mapped_radius: float
    fit_loss: float
    score: float


def get_zeta_zeros(n_roots):
    n_roots = max(1, int(n_roots))
    if n_roots <= len(DEFAULT_GAMMA):
        gamma = DEFAULT_GAMMA[:n_roots]
        return np.array([0.5 + 1j * g for g in gamma], dtype=np.complex128)

    zeros = []
    for k in range(1, n_roots + 1):
        zero = mpmath.zetazero(k)
        zeros.append(complex(float(mpmath.re(zero)), float(mpmath.im(zero))))
    return np.array(zeros, dtype=np.complex128)


def format_complex(z):
    return f"{z.real:.6f}{z.imag:+.6f}j"


def circle_map(s, origin, radius):
    return (s - origin) / radius


def inv_circle_map(w, origin, radius):
    return origin + radius * w


def fit_circle_cluster_torch(cluster, padding=0.35, steps=400, lr=0.03):
    cluster = np.array(cluster, dtype=np.complex128)
    points = np.column_stack((cluster.real, cluster.imag))
    points_t = torch.tensor(points, dtype=torch.float64)

    center_init = points_t.mean(dim=0)
    radii_init = torch.linalg.norm(points_t - center_init, dim=1)
    radius_init = torch.clamp(radii_init.mean(), min=1.0)

    center = torch.nn.Parameter(center_init.clone())
    log_radius = torch.nn.Parameter(torch.log(radius_init))
    optimizer = torch.optim.Adam([center, log_radius], lr=lr)

    best_loss = None
    best_center = None
    best_radius = None

    for _ in range(steps):
        optimizer.zero_grad()
        radius = torch.exp(log_radius)
        distances = torch.linalg.norm(points_t - center, dim=1)
        radial_residual = distances - radius
        containment_penalty = torch.relu(torch.max(distances) - radius) ** 2
        loss = torch.mean(radial_residual ** 2) + 10.0 * containment_penalty
        loss.backward()
        optimizer.step()

        loss_value = float(loss.detach())
        if best_loss is None or loss_value < best_loss:
            best_loss = loss_value
            best_center = center.detach().clone()
            best_radius = radius.detach().clone()

    fitted_center = complex(float(best_center[0]), float(best_center[1]))
    fitted_radius = float(best_radius) * (1.0 + padding)
    return fitted_center, fitted_radius, float(best_loss)


def build_circle_patches(roots, roots_per_circle=3, padding=0.35):
    patches = []
    roots = np.array(sorted(roots, key=lambda z: z.imag), dtype=np.complex128)

    for start in range(0, len(roots), roots_per_circle):
        cluster = roots[start : start + roots_per_circle]
        if len(cluster) == 0:
            continue

        origin, radius, fit_loss = fit_circle_cluster_torch(cluster, padding=padding)
        mapped_roots = circle_map(cluster, origin, radius)
        coeffs = np.poly(mapped_roots)
        patches.append(
            CirclePatch(
                origin=origin,
                radius=radius,
                roots=cluster,
                mapped_roots=mapped_roots,
                coeffs=coeffs,
                fit_loss=fit_loss,
            )
        )

    return patches


def build_circle_patches_for_n(roots, n_circles, padding=0.35, torch_steps=400, torch_lr=0.03):
    roots = np.array(sorted(roots, key=lambda z: z.imag), dtype=np.complex128)
    n_circles = max(1, min(int(n_circles), len(roots)))
    clusters = [cluster for cluster in np.array_split(roots, n_circles) if len(cluster) > 0]

    patches = []
    for cluster in clusters:
        origin, radius, fit_loss = fit_circle_cluster_torch(
            cluster,
            padding=padding,
            steps=torch_steps,
            lr=torch_lr,
        )
        mapped_roots = circle_map(cluster, origin, radius)
        coeffs = np.poly(mapped_roots)
        patches.append(
            CirclePatch(
                origin=origin,
                radius=radius,
                roots=cluster,
                mapped_roots=mapped_roots,
                coeffs=coeffs,
                fit_loss=fit_loss,
            )
        )

    roots_per_circle = int(np.ceil(len(roots) / n_circles))
    return patches, roots_per_circle


def local_patch_polynomial(s, patch):
    w = circle_map(s, patch.origin, patch.radius)
    return np.polyval(patch.coeffs, w)


def multi_circle_root_model(s, patches):
    value = 1.0 + 0.0j
    for patch in patches:
        value *= local_patch_polynomial(s, patch)
    return value


def zeta_product(s, roots):
    prod = 1.0 + 0.0j
    for rho in roots:
        prod *= (1 - s / rho) * np.exp(s / rho)
    return prod


def zeta_model(s, roots):
    s_mp = mpmath.mpc(s.real, s.imag)
    factor = (mpmath.pi ** (s_mp / 2)) / (
        2 * (s_mp - 1) * mpmath.gamma(s_mp / 2 + 1)
    )
    prod = zeta_product(s, roots)
    value = factor * mpmath.mpc(prod.real, prod.imag)
    return complex(value)


def recover_roots_from_patches(patches):
    recovered = []
    for patch in patches:
        local_roots = np.roots(patch.coeffs)
        recovered.extend(inv_circle_map(w, patch.origin, patch.radius) for w in local_roots)
    recovered = np.array(sorted(recovered, key=lambda z: z.imag), dtype=np.complex128)
    return recovered


def mean_root_error(reference_roots, estimated_roots):
    reference = np.array(sorted(reference_roots, key=lambda z: z.imag), dtype=np.complex128)
    estimated = np.array(sorted(estimated_roots, key=lambda z: z.imag), dtype=np.complex128)
    return float(np.mean(np.abs(reference - estimated)))


def layout_score(patches, recovered, reference_roots):
    mean_error = mean_root_error(reference_roots, recovered)
    max_mapped_radius = max(
        float(np.max(np.abs(patch.mapped_roots))) for patch in patches
    )
    fit_loss = float(np.mean([patch.fit_loss for patch in patches]))
    score = mean_error + 1e-6 * max_mapped_radius + fit_loss
    return mean_error, max_mapped_radius, fit_loss, score


def evaluate_layout(roots, n_circles, padding=0.35, torch_steps=400, torch_lr=0.03):
    patches, roots_per_circle = build_circle_patches_for_n(
        roots,
        n_circles,
        padding=padding,
        torch_steps=torch_steps,
        torch_lr=torch_lr,
    )
    recovered = recover_roots_from_patches(patches)
    mean_error, max_mapped_radius, fit_loss, score = layout_score(patches, recovered, roots)
    return PatchLayout(
        n_circles=n_circles,
        roots_per_circle=roots_per_circle,
        patches=patches,
        recovered=recovered,
        mean_error=mean_error,
        max_mapped_radius=max_mapped_radius,
        fit_loss=fit_loss,
        score=score,
    )


def search_circle_layouts(roots, n_min, n_max, padding=0.35, torch_steps=400, torch_lr=0.03):
    layouts = []
    for n_circles in range(n_min, n_max + 1):
        layouts.append(
            evaluate_layout(
                roots,
                n_circles=n_circles,
                padding=padding,
                torch_steps=torch_steps,
                torch_lr=torch_lr,
            )
        )
    best = min(layouts, key=lambda layout: (layout.score, layout.n_circles))
    return best, layouts


def plot_circle_patches(patches, zeros):
    theta = np.linspace(0, 2 * np.pi, 240)
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    axes[0].scatter(zeros.real, zeros.imag, c="red", label="Known zeros")
    for idx, patch in enumerate(patches, start=1):
        circle = patch.origin + patch.radius * np.exp(1j * theta)
        axes[0].plot(circle.real, circle.imag, label=f"Patch {idx}")
        axes[0].scatter(patch.roots.real, patch.roots.imag, s=35)
        axes[0].text(patch.origin.real, patch.origin.imag, str(idx), fontsize=9)
    axes[0].set_title("Multiple circles in the s-plane")
    axes[0].set_xlabel("Re(s)")
    axes[0].set_ylabel("Im(s)")
    axes[0].grid(alpha=0.3)
    axes[0].legend(loc="best", fontsize=8)

    unit_circle = np.exp(1j * theta)
    axes[1].plot(unit_circle.real, unit_circle.imag, "k--", label="Unit circle")
    for idx, patch in enumerate(patches, start=1):
        axes[1].scatter(
            patch.mapped_roots.real,
            patch.mapped_roots.imag,
            label=f"Patch {idx} mapped roots",
        )
    axes[1].set_title("Patch-local normalized root maps")
    axes[1].set_aspect("equal", adjustable="box")
    axes[1].grid(alpha=0.3)
    axes[1].legend(loc="best", fontsize=8)

    plt.tight_layout()


def maybe_show_plot():
    headless = not os.environ.get("DISPLAY") and not os.environ.get("WAYLAND_DISPLAY")
    backend = plt.get_backend().lower()
    non_interactive_backend = "agg" in backend or "pdf" in backend or "svg" in backend
    if headless or non_interactive_backend:
        plt.close("all")
    else:
        plt.show()


def main():
    parser = argparse.ArgumentParser(description="Map zeta zeros with an n-circle patch search.")
    parser.add_argument("--n-roots", type=int, default=10, help="Number of nontrivial zeta zeros to model.")
    parser.add_argument("--n-circles", type=int, default=4, help="Number of circles to use when search is disabled.")
    parser.add_argument("--search", action="store_true", help="Search over a range of circle counts and choose the best layout.")
    parser.add_argument("--n-min", type=int, default=2, help="Minimum circles to try during search.")
    parser.add_argument("--n-max", type=int, default=6, help="Maximum circles to try during search.")
    parser.add_argument("--padding", type=float, default=0.35, help="Per-circle radius padding factor.")
    parser.add_argument("--torch-steps", type=int, default=400, help="Gradient steps for each torch circle fit.")
    parser.add_argument("--torch-lr", type=float, default=0.03, help="Learning rate for each torch circle fit.")
    args = parser.parse_args()

    zeros = get_zeta_zeros(args.n_roots)

    if args.search:
        layout, layouts = search_circle_layouts(
            zeros,
            n_min=args.n_min,
            n_max=args.n_max,
            padding=args.padding,
            torch_steps=args.torch_steps,
            torch_lr=args.torch_lr,
        )
        print("Circle search results:")
        for candidate in layouts:
            print(
                f"  n={candidate.n_circles}: roots/circle≈{candidate.roots_per_circle}, "
                f"mean_error={candidate.mean_error:.6e}, "
                f"max|w|={candidate.max_mapped_radius:.6f}, "
                f"fit_loss={candidate.fit_loss:.6e}, score={candidate.score:.6e}"
            )
    else:
        layout = evaluate_layout(
            zeros,
            n_circles=args.n_circles,
            padding=args.padding,
            torch_steps=args.torch_steps,
            torch_lr=args.torch_lr,
        )

    patches = layout.patches
    recovered = layout.recovered

    print(f"Built {len(patches)} circle patches")
    for idx, patch in enumerate(patches, start=1):
        print(
            f"Patch {idx}: origin={format_complex(patch.origin)}, "
            f"radius={patch.radius:.6f}, roots={len(patch.roots)}, fit_loss={patch.fit_loss:.6e}"
        )

    global_coeffs = np.poly(circle_map(zeros, 0.5 + 0.0j, 1.0))
    recovered_global = 0.5 + np.roots(global_coeffs)

    print(f"Mean root recovery error, multi-circle model: {layout.mean_error:.6e}")
    print(f"Mean root recovery error, single global circle: {mean_root_error(zeros, recovered_global):.6e}")
    print(f"Max mapped patch radius: {layout.max_mapped_radius:.6f}")
    print(f"Mean torch fit loss: {layout.fit_loss:.6e}")

    s_test = 0.5 + 1j * 14.0
    zeta_true = mpmath.zeta(s_test)
    zeta_approx = zeta_model(s_test, zeros)
    root_model_value = multi_circle_root_model(s_test, patches)
    print(
        f"ζ({s_test}) ≈ {format_complex(zeta_approx)} "
        f"(true: {float(mpmath.re(zeta_true)):.6f}{float(mpmath.im(zeta_true)):+.6f}j)"
    )
    print(f"Multi-circle root model at {s_test}: {format_complex(root_model_value)}")

    plot_circle_patches(patches, zeros)
    maybe_show_plot()


if __name__ == "__main__":
    main()