#
# ELM using even or odd parity features
# - enforce symmetry in features
# - don't rely on learning symmetry from a few training points 
#


import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(12345)


fig, (ax1, ax2) = plt.subplots(1,2,figsize=(15,8))


def activation(z):
    return np.maximum(z, 0.0)
#    return np.tanh(z)


def rho_exact(x):
    """Ground-state density of the harmonic oscillator."""
    return np.exp(-x**2) / np.sqrt(np.pi)

# Training data only on the positive half-axis
x_train = np.linspace(0.0, 2.5, 15)
y_train = rho_exact(x_train)


# Prediction points on both sides
x_plot = np.linspace(-3.0, 3.0, 500)

# Random hidden-layer parameters
nhid = 80
w1 = rng.normal(0.0, 1.0, nhid)
b1 = rng.normal(0.0, 1.0, nhid)


def parity_decaying_matrix(x, w, b, parity):
    """
    parity = +1: even
    parity = -1: odd
    """
    x = np.asarray(x)

    z_plus = w[:, None] * x[None, :] + b[:, None]
    z_minus = w[:, None] * (-x)[None, :] + b[:, None]

    projected = 0.5 * (
        activation(z_plus)
        + parity * activation(z_minus)
    )

    # decay or not: choose envelope
    envelope = np.exp(-x**2 / 2)[None, :]
    #envelope = 1.0
        

    return (
        envelope
        * projected
        / np.sqrt(len(w))
    )



# even: ground state
M = parity_decaying_matrix(x_train, w1, b1, parity = +1)

print("M shape:", M.shape)

lam = 1.0e-6

A_aug = np.vstack((
    M.T,
    np.sqrt(lam) * np.eye(nhid)
))


y_aug = np.concatenate((
    y_train,
    np.zeros(nhid)
))

w2 = np.linalg.lstsq(
    A_aug,
    y_aug,
    rcond=None
)[0]

# Note: even-projected features decay automatically for tanh, but not for ReLU
M_plot = parity_decaying_matrix(x_plot, w1, b1, parity = +1)
y_plot = w2 @ M_plot


ax1.plot(
    x_plot,
    rho_exact(x_plot),
    label="exact"
)

ax1.plot(
    x_plot,
    y_plot,
    "--",
    label="even ELM"
)

ax1.plot(
    x_train,
    y_train,
    "o",
    label=r"training data, $x\geq0$"
)

ax1.set_title("Harmonic oscillator ground state density")
ax1.set_xlabel(r"$x$")
ax1.set_ylabel(r"$\rho_0(x)$")
ax1.legend()



x_test = np.linspace(0.0, 3.0, 100)

y_positive = w2 @ parity_decaying_matrix(x_test, w1, b1, parity = +1)

y_negative = w2 @ parity_decaying_matrix(-x_test, w1, b1, parity = +1)

print(
    "maximum parity error:",
    np.max(np.abs(y_positive - y_negative))
)


# 1st excited state 
# =================
def psi_1_exact(x):
    """First-excited-state wavefunction."""
    return (np.sqrt(2)* x * np.exp(-x**2 / 2) / np.pi**0.25 )



# Training data only on the positive half-axis
x_train = np.linspace(0.0, 2.5, 15)
y_train = psi_1_exact(x_train)


# Prediction points on both sides
x_plot = np.linspace(-3.0, 3.0, 500)

# Random hidden-layer parameters
nhid = 80
w1 = rng.normal(0.0, 1.0, nhid)
b1 = rng.normal(0.0, 1.0, nhid)


# Note: odd-projected features *do not* decay automatically in most cases


M = parity_decaying_matrix(x_train, w1, b1, parity=-1)


print("M shape:", M.shape)

lam = 1.0e-6

A_aug = np.vstack((
    M.T,
    np.sqrt(lam) * np.eye(nhid)
))


y_aug = np.concatenate((
    y_train,
    np.zeros(nhid)
))

w2 = np.linalg.lstsq(
    A_aug,
    y_aug,
    rcond=None
)[0]



M_plot = parity_decaying_matrix(x_plot, w1, b1, parity=-1)
y_plot = w2 @ M_plot



ax2.plot(
    x_plot,
    psi_1_exact(x_plot),
    label="exact"
)

ax2.plot(
    x_plot,
    y_plot,
    "--",
    label="odd ELM"
)

ax2.plot(
    x_train,
    y_train,
    "o",
    label=r"training data, $x\geq0$"
)
ax2.set_title(r"Harmonic oscillator first excited state wave function")
ax2.set_xlabel(r"$x$")
ax2.set_ylabel(r"$\psi_1(x)$")
ax2.legend()


x_test = np.linspace(0.0, 3.0, 100)

y_positive = w2 @ parity_decaying_matrix(
    x_test, w1, b1, parity=-1
)

y_negative = w2 @ parity_decaying_matrix(
    -x_test, w1, b1, parity=-1
)

print(
    "maximum parity error:",
    np.max(np.abs(y_positive - y_negative))
)



plt.show()