import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split

# -------------------------------
# Synthetic training data (replace with experimental dataset)
# -------------------------------
frequencies = np.linspace(100, 1000, 100)
nodal_counts = 3*np.sin(0.01*frequencies) + 0.5*np.cos(0.005*frequencies) + np.random.normal(0,0.05,len(frequencies))

X = np.vstack([
    frequencies,
    frequencies**2,
    np.sin(frequencies/100),
    np.cos(frequencies/200),
    np.power(1.618, frequencies/500),
    np.sqrt(frequencies),
    np.log1p(frequencies)
]).T

Y = np.vstack([
    nodal_counts + 0.1*np.random.randn(len(frequencies)),
    nodal_counts*2 + 0.2*np.random.randn(len(frequencies)),
    nodal_counts*0.7 + 0.1*np.random.randn(len(frequencies)),
    nodal_counts*1.1 + 0.2*np.random.randn(len(frequencies))
]).T

X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=42)
model = LinearRegression().fit(X_train, Y_train)

# -------------------------------
# φ-hash tuning functions
# -------------------------------
def safe_phi_hash(x, phi=1.6180339887):
    x_mod = np.mod(x, 1000)
    return np.mod(np.power(phi, x_mod), 1.0)

def conditional_phi_tune(y_pred, x, y_true, lam=0.01):
    hash_vec = safe_phi_hash(x)
    hash_mapped = np.ones_like(y_pred) * np.sum(hash_vec)/len(hash_vec)
    y_candidate = y_pred + lam * hash_mapped
    return y_candidate if np.linalg.norm(y_candidate - y_true) < np.linalg.norm(y_pred - y_true) else y_pred

# -------------------------------
# Polar/Cartesian grid setup
# -------------------------------
Nx, Ny = 300, 300
x = np.linspace(0, 1, Nx)
y = np.linspace(0, 1, Ny)
Xg, Yg = np.meshgrid(x, y)
R = np.sqrt(Xg**2 + Yg**2)
Theta = np.arctan2(Yg, Xg)

def morph_coords(Xc, Yc, R, Theta, m):
    Xm = (1 - m) * Xc + m * R * np.cos(Theta)
    Ym = (1 - m) * Yc + m * R * np.sin(Theta)
    return Xm, Ym

# -------------------------------
# Note ↔ Frequency
# -------------------------------
def note_to_freq(n):
    return 440 * 2**((n - 69)/12)

def get_cymatic_params(f_query):
    features = np.array([
        f_query,
        f_query**2,
        np.sin(f_query/100),
        np.cos(f_query/200),
        np.power(1.618, f_query/500),
        np.sqrt(f_query),
        np.log1p(f_query)
    ]).reshape(1,-1)
    params = model.predict(features)[0]
    y_true_ref = Y_test.mean(axis=0)
    params_tuned = conditional_phi_tune(params, features.flatten(), y_true_ref)
    return {
        "alpha": abs(params_tuned[0]),
        "beta": abs(params_tuned[1]),
        "eta": abs(params_tuned[2]),
        "zeta": abs(params_tuned[3])
    }

def cymatic(XY, alpha, beta, eta, zeta):
    return np.sin(alpha * np.pi * XY[0]) * np.sin(beta * np.pi * XY[1]) + eta * np.cos(zeta * np.pi * (XY[0]+XY[1]))

# -------------------------------
# Interactive plot
# -------------------------------
init_note = 69
init_morph = 0.0
f_query = note_to_freq(init_note)
params = get_cymatic_params(f_query)
Xm, Ym = morph_coords(Xg, Yg, R, Theta, init_morph)
Z = cymatic((Xm,Ym), **params)

fig, ax = plt.subplots(figsize=(6,6))
plt.subplots_adjust(bottom=0.3)
contour = ax.contour(Xm, Ym, Z, levels=[0], colors='black')
ax.set_title(f"Note: {init_note}, Freq: {f_query:.2f} Hz")
ax.axis('off')

# Morph slider
ax_morph = plt.axes([0.25, 0.15, 0.5, 0.03])
slider_morph = Slider(ax_morph, 'Morph', 0.0, 1.0, valinit=init_morph)

# Note slider
ax_note = plt.axes([0.25, 0.1, 0.5, 0.03])
slider_note = Slider(ax_note, 'Note', 21, 108, valinit=init_note)

def update(val):
    morph_val = slider_morph.val
    note_val = slider_note.val
    f_query = note_to_freq(note_val)
    params = get_cymatic_params(f_query)
    Xm, Ym = morph_coords(Xg, Yg, R, Theta, morph_val)
    ax.clear()
    Z = cymatic((Xm,Ym), **params)
    ax.contour(Xm, Ym, Z, levels=[0], colors='black')
    ax.set_title(f"Note: {int(note_val)}, Freq: {f_query:.2f} Hz")
    ax.axis('off')
    fig.canvas.draw_idle()

slider_morph.on_changed(update)
slider_note.on_changed(update)
plt.show()
