"""
simulate_grid.py — Sweep a 10×10 grid across the workspace and analyze
tension differential, belt lengths, and worst-case positions.

Outputs:
  - workspace_heatmap.png  (matplotlib heatmap of tension differential)
  - workspace_heatmap.csv  (fallback if no matplotlib, also written as data log)
"""

from __future__ import annotations

import csv
import math
import os
import sys

import numpy as np

from kinematics import (
    belt_lengths,
    solve_forward,
    BELT_NAMES,
    TEST_GRID,
)
from tension_analysis import analyze_tension, resting_lengths

# Output directory
OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))

# Workspace bounds
X_MIN, X_MAX = -0.6, 0.6
Y_MIN, Y_MAX = -1.2, 1.2


def sweep_grid(nx: int = 10, ny: int = 10,
               z: float = 0.0) -> tuple[np.ndarray, np.ndarray, np.ndarray, tuple]:
    """Sweep an nx × ny grid across the workspace at height z.

    Returns:
      xs, ys: 1D arrays of X and Y grid lines
      diff_map: (ny, nx) array of max-min tension differential at each point
      worst: ((x, y), max_diff) — the point with greatest tension differential
    """
    xs = np.linspace(X_MIN, X_MAX, nx)
    ys = np.linspace(Y_MIN, Y_MAX, ny)

    diff_map = np.zeros((ny, nx))
    worst_diff = 0.0
    worst_xy = (0.0, 0.0)

    rest = resting_lengths(0.0, 0.0, 0.0)

    for i, x in enumerate(xs):
        for j, y in enumerate(ys):
            result = analyze_tension(x, y, z, rest_lengths=rest)
            tmin = np.min(result.tension_multipliers)
            tmax = np.max(result.tension_multipliers)
            diff = tmax - tmin
            diff_map[j, i] = diff
            if diff > worst_diff:
                worst_diff = diff
                worst_xy = (x, y)

    return xs, ys, diff_map, (worst_xy, worst_diff)


def write_csv(xs: np.ndarray, ys: np.ndarray,
              diff_map: np.ndarray, path: str) -> None:
    """Write the grid data as a CSV file."""
    with open(path, "w", newline="") as f:
        writer = csv.writer(f)
        # Header: first cell empty, then X coordinates
        header = [""] + [f"{x:.6f}" for x in xs]
        writer.writerow(header)
        for j, y in enumerate(ys):
            row = [f"{y:.6f}"] + [f"{diff_map[j, i]:.6f}" for i in range(len(xs))]
            writer.writerow(row)
    print(f"  Wrote CSV: {path}")


def plot_heatmap(xs: np.ndarray, ys: np.ndarray,
                 diff_map: np.ndarray, worst_xy: tuple,
                 worst_diff: float,
                 path: str) -> bool:
    """Generate and save a heatmap using matplotlib.

    Returns True on success, False if matplotlib is unavailable.
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        return False

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

    X, Y = np.meshgrid(xs, ys)
    levels = 50
    cf = ax.contourf(X, Y, diff_map, levels=levels, cmap="plasma")
    cbar = fig.colorbar(cf, ax=ax, label="Tension Differential (multiplier range)")

    # Mark worst point
    wx, wy = worst_xy
    ax.plot(wx, wy, marker="*", color="white", markersize=14,
            markeredgecolor="black", markeredgewidth=1.0)
    ax.annotate(f"Worst: ({wx:.3f}, {wy:.3f})\nDiff = {worst_diff:.3f}",
                xy=(wx, wy), xytext=(wx + 0.12, wy + 0.08),
                color="white", fontsize=9,
                arrowprops=dict(arrowstyle="->", color="white", lw=1.2),
                bbox=dict(boxstyle="round,pad=0.3", facecolor="black",
                          edgecolor="white", alpha=0.7))

    # Mark the 9 test points
    for name, (tx, ty, tz) in TEST_GRID:
        ax.plot(tx, ty, marker="o", color="cyan", markersize=4, alpha=0.8)

    ax.set_xlabel("X (m)")
    ax.set_ylabel("Y (m)")
    ax.set_title("Gordix 8-Belt — Tension Differential Across Workspace\n"
                 "(10×10 grid, Z=0)")
    ax.set_aspect("equal")
    ax.grid(True, alpha=0.3)

    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print(f"  Saved heatmap: {path}")
    return True


def _run_sweep():
    print("=" * 70)
    print("  Grid Sweep — Gordix 8-Belt Workspace Analysis")
    print("=" * 70)

    xs, ys, diff_map, (worst_xy, worst_diff) = sweep_grid(10, 10, 0.0)

    print(f"\n  Grid: 10 × 10 = 100 points")
    print(f"  Workspace: X=[{X_MIN:.2f}, {X_MAX:.2f}]  Y=[{Y_MIN:.2f}, {Y_MAX:.2f}]")
    print(f"\n  Worst-case tension differential:")
    print(f"    Point:  ({worst_xy[0]:.4f}, {worst_xy[1]:.4f}) m")
    print(f"    Differential: {worst_diff:.4f} (tension multiplier range)")

    # Also report raw belt length range
    rest = resting_lengths(0.0, 0.0, 0.0)
    bl = belt_lengths(worst_xy[0], worst_xy[1], 0.0)
    vals = list(bl.values())
    print(f"    Belt lengths: {min(vals):.6f} – {max(vals):.6f} m")
    print(f"    ΔL from rest:")
    for name in BELT_NAMES:
        delta = (bl[name] - rest[name]) * 1000
        print(f"      {name:<12}: {delta:+8.4f} mm")

    # Save CSV always
    csv_path = os.path.join(OUTPUT_DIR, "workspace_heatmap.csv")
    write_csv(xs, ys, diff_map, csv_path)

    # Save PNG if possible
    png_path = os.path.join(OUTPUT_DIR, "workspace_heatmap.png")
    ok = plot_heatmap(xs, ys, diff_map, worst_xy, worst_diff, png_path)
    if not ok:
        print("  [matplotlib not available — skipped PNG, CSV saved]")

    print()
    return worst_xy, worst_diff


if __name__ == "__main__":
    _run_sweep()