# modulus.py
# Andrew Davison, ad@coe.psu.ac.th, Nov. 2025

# Modulus plot of a cubic.

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

import numpy as np   # used for meshgrid() and array()
import cubicUtils


class CyclicNormalize(mcolors.Normalize):
  def __init__(self, vmin=None, vmax=None, period=1.0, clip=False):
    self.period = period
    super().__init__(vmin, vmax, clip)
  
  def __call__(self, value, clip=None):
    # Map values to [0, 1] with cycling
    result = ((value - self.vmin) % self.period) / self.period
    return np.ma.array(result, mask=np.ma.getmask(value))


def getMagnitudes(realVs, imagVs):
  mags = []
  for im in imagVs:
    rowMag = []
    for re in realVs:
      z = complex(re, im)
      try:
        w = cubicUtils.cubic(z, a, b, c, d)
      except (ValueError, ZeroDivisionError):
        w = complex(float('nan'), float('nan'))
      rowMag.append( abs(w))
    mags.append(rowMag)
  return mags



# ---------------------------------
steps = 150

a, b, c, d = map(float, input("Enter cubic coefs (a b c d): ").split())
size = int(input("Range (or 3): ") or "3")

realMin, realMax, imagMin, imagMax = -size, size, -size, size
print(f"Real range: [{realMin}, {realMax}]")
print(f"Imaginary range: [{imagMin}, {imagMax}]")

reals = cubicUtils.linspace(realMin, realMax, steps)
imags = cubicUtils.linspace(imagMin, imagMax, steps)
mags = getMagnitudes(reals, imags)

realGrid, imagGrid = np.meshgrid(reals, imags)
Z = np.array(mags)

fig = plt.figure(figsize=(10, 7))
ax = fig.add_subplot(111, projection='3d')

cycle = CyclicNormalize(vmin=0, period=10)
surf = ax.plot_surface(realGrid, imagGrid, Z, 
            cmap='hsv', norm=cycle, 
            antialiased=True)

ax.set_xlabel('Re(z)')
ax.set_ylabel('Im(z)')
ax.set_zlabel('|f(z)|')
ax.set_title('Modulus Plot for ' + 
               cubicUtils.toString(a, b, c, d))

ax.set_zlim(0, max(max(row) for row in mags))

plt.tight_layout()
plt.show()