K.clear_session()

# --- Simulation & Model Setup ---
L = 8
num_chains, num_steps, equil_steps = 128, 500, 50
boundaries = [200, 250, 300, 350, 400, 450]
values = [0.005 * (0.5 ** i) for i in range(len(boundaries) + 1)]
lr_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
    boundaries, values
)
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)

# --- Complex periodic conv2D layer ---
def cELU(z, alpha=1.0):
    return tf.cast(alpha, tf.complex128) * (tf.exp(z) - tf.cast(1.0, tf.complex128))
class ComplexPeriodicConv2D(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size):
        super().__init__()
        self.filters = filters
        self.kernel_size = kernel_size

    def build(self, input_shape):
        in_channels = input_shape[-2]
        kernel_shape = (self.kernel_size, self.kernel_size, in_channels, self.filters)
        self.w_real = self.add_weight(name="w_real", shape=kernel_shape, initializer="glorot_uniform", trainable=True)
        self.w_imag = self.add_weight(name="w_imag", shape=kernel_shape, initializer="glorot_uniform", trainable=True)
        self.b_real = self.add_weight(name="b_real", shape=(self.filters,), initializer="zeros", trainable=True)
        self.b_imag = self.add_weight(name="b_imag", shape=(self.filters,), initializer="zeros", trainable=True)

        
    def periodic_pad(self, x, pad_h, pad_w):
        x = tf.concat([x, x[:, :pad_h, :, :]], axis=1)
        x = tf.concat([x, x[:, :, :pad_w, :]], axis=2)
        return x

    def call(self, inputs):
        # inputs: (N, H, W, in_channels, 2)
        x_real = inputs[..., 0]
        x_imag = inputs[..., 1]

        # Periodic padding
        pad_h, pad_w = self.kernel_size - 1, self.kernel_size - 1
        inputs_padded_real = self.periodic_pad(x_real, pad_h, pad_w)
        inputs_padded_imag = self.periodic_pad(x_imag, pad_h, pad_w)

        conv_rr = tf.nn.conv2d(inputs_padded_real, self.w_real, strides=1, padding="VALID")
        conv_ri = tf.nn.conv2d(inputs_padded_real, self.w_imag, strides=1, padding="VALID")
        conv_ir = tf.nn.conv2d(inputs_padded_imag, self.w_real, strides=1, padding="VALID")
        conv_ii = tf.nn.conv2d(inputs_padded_imag, self.w_imag, strides=1, padding="VALID")

        real_out = conv_rr - conv_ii + self.b_real
        imag_out = conv_ri + conv_ir + self.b_imag

        return tf.stack([real_out, imag_out], axis=-1)

    def get_complex_weights(self):
        w = tf.complex(self.w_real, self.w_imag)
        b = tf.complex(self.b_real, self.b_imag)
        return [w], [b]

    def apply_complex_grads(self, grad_w_list, grad_b_list):
        grad_w, = grad_w_list
        grad_b, = grad_b_list
        self.w_real.assign_add(tf.math.real(grad_w))
        self.w_imag.assign_add(tf.math.imag(grad_w))
        self.b_real.assign_add(tf.math.real(grad_b))
        self.b_imag.assign_add(tf.math.imag(grad_b))

# --- Model ---
def build_model(L):
    inputs = Input(shape=(L, L, 2), dtype=tf.float64)
    x = PreprocessLayerU1()(inputs)
    x = PlaquetteU1(1)(x)
    conv = ComplexPeriodicConv2D(8, 3)
    x = conv(x)
    x = cELU(tf.complex(x[..., 0], x[..., 1]))
    x = tf.stack([tf.math.real(x), tf.math.imag(x)], axis=-1)
    x = tf.reduce_mean(x, axis=[1, 2, 3])
    outputs = tf.stack([x[..., 0], x[..., 1]], axis=-1)

    model = Model(inputs, outputs)
    model.conv_layer = conv
    return model


# --- Gradient wrapper ---
def compute_wirtinger_gradient(model, config):
    with tf.GradientTape(persistent=True) as tape:
        out = model(config)
        log_amp, phase = out[:, 0], out[:, 1]
    grads_log = tape.jacobian(log_amp, model.trainable_variables)
    grads_phase = tape.jacobian(phase, model.trainable_variables)
    del tape

    O_list = []
    for gl, gp in zip(grads_log, grads_phase):
        gl = tf.reshape(gl, [tf.shape(config)[0], -1])
        gp = tf.reshape(gp, [tf.shape(config)[0], -1])
        O_list.append(tf.complex(gl, -gp) * 0.5)  # Wirtinger
    return tf.concat(O_list, axis=1)  # [batch, n_params]

# --- Vectorized local energy ---
def metropolis_step_u1(config, model, step_size=0.3):
    proposal = config + tf.random.uniform(tf.shape(config), minval=-step_size, maxval=step_size, dtype=tf.float64)
    proposal = tf.math.floormod(proposal, 2.0 * np.pi)

    out0 = model(config, training=False)
    out1 = model(proposal, training=False)
    ψ0 = tf.exp(out0[:, 0])
    ψ1 = tf.exp(out1[:, 0])
    accept_prob = tf.square(ψ1 / ψ0)
    accept = tf.random.uniform(tf.shape(accept_prob), dtype=tf.float64) < accept_prob
    accept = tf.reshape(accept, [-1, 1, 1, 1])
    return tf.where(accept, proposal, config)


# ====== Local Energy ======
@tf.function
def local_energy_u1(config, model, g=1., delta=1e-2):
    N = tf.shape(config)[0]
    two_pi = tf.constant(2.0 * np.pi, tf.float64)
    Δ = tf.constant(delta, tf.float64)

    θx, θy = config[..., 0], config[..., 1]
    plaq = θx + tf.roll(θy, -1, 1) - tf.roll(θx, -1, 2) - θy
    V_mag = (1.0 / (g**2)) * tf.reduce_sum(1. - tf.cos(plaq), axis=[1, 2])

    out0 = model(config, training=False)
    lnψ0, φ0 = out0[:, 0], out0[:, 1]
    ψ0 = tf.exp(tf.complex(lnψ0, φ0))

    num_links = L * L * 2
    cfgs_rep = tf.repeat(config, repeats=num_links, axis=0)

    ii = tf.range(L, dtype=tf.int32)
    jj = tf.range(L, dtype=tf.int32)
    μs = tf.range(2, dtype=tf.int32)
    I, J, M = tf.meshgrid(ii, jj, μs, indexing='ij')
    link_idx = tf.stack([tf.reshape(I, [-1]), tf.reshape(J, [-1]), tf.reshape(M, [-1])], axis=1)

    chain_idx = tf.repeat(tf.range(N, dtype=tf.int32)[:, None], num_links, axis=1)
    chain_idx = tf.reshape(chain_idx, [-1, 1])
    link_rep = tf.tile(link_idx[None], [N, 1, 1])
    link_rep = tf.reshape(link_rep, [-1, 3])
    idx = tf.concat([chain_idx, link_rep], axis=1)

    ups = tf.fill([N * num_links], Δ)
    dns = -ups
    cfg_p = tf.tensor_scatter_nd_add(cfgs_rep, idx, ups)
    cfg_p = tf.math.floormod(cfg_p, two_pi)
    cfg_m = tf.tensor_scatter_nd_add(cfgs_rep, idx, dns)
    cfg_m = tf.math.floormod(cfg_m, two_pi)

    out_p = model(cfg_p, training=False)
    out_m = model(cfg_m, training=False)
    ψ_p = tf.exp(tf.complex(out_p[:, 0], out_p[:, 1]))
    ψ_m = tf.exp(tf.complex(out_m[:, 0], out_m[:, 1]))

    ψ_p = tf.reshape(ψ_p, [N, num_links])
    ψ_m = tf.reshape(ψ_m, [N, num_links])
    kin_sum = tf.reduce_sum(ψ_p + ψ_m - 2.0 * ψ0[:, None], axis=1)

    prefac = tf.cast(g**2 / (2.0 * Δ * Δ), tf.complex128)
    lap_over_ψ0 = -tf.cast(kin_sum, tf.complex128) / ψ0
    T_el = prefac * lap_over_ψ0

    return tf.cast(V_mag, tf.complex128) + T_el


# ====== SR Step with CG and Complex Gradients ======


def train_step_sr_cg(config, g=1.0, diag_shift=1e-2, max_iter=20):
    with tf.GradientTape(persistent=True) as tape:
        out = model(config)
        lnψ, φ = out[:, 0], out[:, 1]
        ψ = tf.exp(tf.complex(lnψ, φ))
        E_loc = local_energy_u1(config, model, g=g)
    
    # Get complex weights
    vars_w, vars_b = model.conv_layer.get_complex_weights()
    params = vars_w + vars_b
    grads = tape.gradient(tf.math.real(tf.reduce_mean(E_loc * tf.math.conj(ψ) / ψ)), params)
    
    # Compute log(ψ) gradients
    O_list = []
    for param in params:
        grad_lnψ = tape.gradient(lnψ, param)
        grad_φ = tape.gradient(φ, param)

        if grad_lnψ is None:
            grad_lnψ = tf.zeros_like(param, dtype=tf.float64)
        if grad_φ is None:
            grad_φ = tf.zeros_like(param, dtype=tf.float64)

        N = tf.shape(config)[0]  # batch size

        grad_lnψ = tf.convert_to_tensor(grad_lnψ)
        grad_φ = tf.convert_to_tensor(grad_φ)

        if tf.rank(grad_lnψ) == 1:
            grad_lnψ = tf.broadcast_to(grad_lnψ[None, :], [N, tf.shape(grad_lnψ)[0]])
        if tf.rank(grad_φ) == 1:
            grad_φ = tf.broadcast_to(grad_φ[None, :], [N, tf.shape(grad_φ)[0]])

        # Handle case where shape might be (filters,) or (batch_size, filters,...)
        if tf.shape(grad_lnψ)[0] != N:
            grad_lnψ = tf.broadcast_to(grad_lnψ, [N] + grad_lnψ.shape.as_list())
        if tf.shape(grad_φ)[0] != N:
            grad_φ = tf.broadcast_to(grad_φ, [N] + grad_φ.shape.as_list())

        grad = tf.complex(grad_lnψ, -grad_φ) * 0.5
        grad = tf.reshape(grad, [N, -1])
        O_list.append(grad)

    O = tf.concat(O_list, axis=1)
    O = tf.cast(O, tf.complex128)

    O_mean = tf.reduce_mean(O, axis=0, keepdims=True)
    O_centered = O - O_mean
    F = tf.reduce_mean(tf.expand_dims(tf.math.conj(E_loc), -1) * O_centered, axis=0)
    F = tf.reshape(F, [-1, 1])

    def matvec(x):
        batch_size = tf.cast(tf.shape(O_centered)[0], tf.complex128)
        Sx = tf.matmul(tf.matmul(O_centered, O_centered, adjoint_a=True), x) / batch_size
        return Sx + diag_shift * x

    x = tf.zeros_like(F)
    r = F
    p = F
    rdotr = tf.reduce_sum(tf.math.conj(r) * r)

    def body(i, r, p, x, rdotr):
        Ap = matvec(p)
        alpha = rdotr / tf.reduce_sum(tf.math.conj(p) * Ap)
        x_new = x + alpha * p
        r_new = r - alpha * Ap
        rdotr_new = tf.reduce_sum(tf.math.conj(r_new) * r_new)
        beta = rdotr_new / rdotr
        p_new = r_new + beta * p
        return i+1, r_new, p_new, x_new, rdotr_new

    i = tf.constant(0)
    _, _, _, delta, _ = tf.while_loop(lambda i, *_: i < max_iter, body, [i, r, p, x, rdotr])

    # Apply -lr * delta to complex weights
    sizes = [tf.size(p) for p in params]
    indices = tf.concat([[0], tf.cumsum(sizes)], axis=0)
    updates = []
    for k in range(len(params)):
        updates.append(tf.reshape(delta[indices[k]:indices[k+1]], tf.shape(params[k])))

    lr = optimizer.learning_rate(step)
    lr_c = tf.cast(lr, tf.complex128)
    neg_updates = [-lr_c * u for u in updates]
    model.conv_layer.apply_complex_grads(neg_updates[:1], neg_updates[1:])

    return tf.reduce_mean(tf.math.real(E_loc)), tf.math.reduce_variance(tf.math.real(E_loc))


# ====== Run VMC Simulation ======
model = build_model(L)
num_chains = 128
num_steps = 300
equil_steps = 5

current_config = tf.random.uniform([num_chains, L, L, 2], 0.0, 2.0 * np.pi, tf.float64)
print(f"Starting U(1) VMC: L={L}, chains={num_chains}")
tf.print("E_init:", tf.reduce_mean(tf.math.real(local_energy_u1(current_config, model))))

for step in range(num_steps):
    t0 = time.time()
    for _ in range(equil_steps):
        current_config = metropolis_step_u1(current_config, model)
    e, var = train_step_sr_cg(current_config)

    #if step % 50 == 0 or step == num_steps - 1:
    tf.print(f"Step {step:3d} | E {e:.6f} | Var {var:.4f} | T {time.time()-t0:.3f}s")

I complemented VMC for complex wavefunction and complex parameters. But get the result below. What's wrong?
Starting U(1) VMC: L=8, chains=128
E_init: 64.811250383702088
Step   0 | E 64.774303 | Var 33.7131 | T 2.499s
Step   1 | E nan | Var nan | T 0.326s
Step   2 | E nan | Var nan | T 0.321s
Step   3 | E nan | Var nan | T 0.322s
Step   4 | E nan | Var nan | T 0.322s
Step   5 | E nan | Var nan | T 0.320s
Step   6 | E nan | Var nan | T 0.326s
Step   7 | E nan | Var nan | T 0.328s
Step   8 | E nan | Var nan | T 0.335s
Step   9 | E nan | Var nan | T 0.328s
Step  10 | E nan | Var nan | T 0.329s