#!/usr/bin/env python3
# winding-step-data.py
# Generate step function data for gnuplot

import math
import numpy as np


def compute_winding_safe(r_val, samples=2000, avoid_roots=0.05):
    """Compute winding number, returning None if r is too close to a root."""
    pi = math.pi
    t = np.linspace(0, 2 * pi, samples)
    roots = [1.0, 2.0, 3.0]

    for root_mag in roots:
        if abs(r_val - root_mag) < avoid_roots:
            return None

    x = (
        r_val**5 * np.cos(5 * t)
        - 7 * r_val**4 * np.cos(4 * t)
        + 19 * r_val**3 * np.cos(3 * t)
        - 37 * r_val**2 * np.cos(2 * t)
        + 60 * r_val * np.cos(t)
        - 36
    )
    y = (
        r_val**5 * np.sin(5 * t)
        - 7 * r_val**4 * np.sin(4 * t)
        + 19 * r_val**3 * np.sin(3 * t)
        - 37 * r_val**2 * np.sin(2 * t)
        + 60 * r_val * np.sin(t)
    )

    arg_raw = np.arctan2(y, x)
    arg_unwrapped = np.unwrap(arg_raw)
    winding = (arg_unwrapped[-1] - arg_unwrapped[0]) / (2 * pi)

    return round(winding)


# Generate r values across full range
r_values = np.linspace(0.1, 4.0, 500)

# Collect data points
w_values = []
r_plot = []

for r in r_values:
    w = compute_winding_safe(r, avoid_roots=0.08)
    if w is not None:
        w_values.append(w)
        r_plot.append(r)

# Write to single data file
with open("winding-step-function.txt", "w") as f:
    f.write("# r       W(r)\n")
    for r, w in zip(r_plot, w_values):
        f.write(f"{r:.4f}  {w}\n")

print("✓ Data written to: winding-step-function.txt")
print(f"  {len(r_plot)} points")