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.05 * (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-1):
    N = tf.shape(config)[0]
    two_pi = tf.constant(2.0 * np.pi, tf.float64)
    Δ = tf.constant(delta, tf.float64)
    epsilon = tf.constant(1e-10, tf.complex128) # Stabilization constant

    θ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 + epsilon) # Added epsilon
    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, step, g=1.0, diag_shift=1e-3, max_iter=50):
    with tf.GradientTape(persistent=True) as tape:
        out = model(config)
        # Unpack the model's output
        lnψ, φ = out[:, 0], out[:, 1]
        # Calculate the local energy for each sample in the batch
        E_loc_complex = local_energy_u1(config, model, g=g)

    # The physical energy is real. SR is based on minimizing its expectation value.
    E_loc = tf.math.real(E_loc_complex)
    E_mean = tf.reduce_mean(E_loc)
    E_centered = E_loc - E_mean

    # The trainable variables are the real and imaginary parts of the weights and biases.
    # The order defined in the ComplexPeriodicConv2D layer is: [w_real, w_imag, b_real, b_imag]
    real_trainable_vars = model.conv_layer.trainable_variables
    
    # Calculate the Jacobians (per-sample gradients) of the log-amplitude and phase
    # with respect to each of the real trainable variables.
    jac_lnψ_list = tape.jacobian(lnψ, real_trainable_vars)
    jac_φ_list = tape.jacobian(φ, real_trainable_vars)
    del tape

    # O_list will store the log-derivatives with respect to the *complex* parameters.
    O_list = []
    # We iterate through the parameters in pairs (real, imag) to build the complex derivatives.
    for i in range(0, len(real_trainable_vars), 2):
        # Derivatives with respect to the real part of the parameter (e.g., w_r)
        d_lnψ_dr = jac_lnψ_list[i]
        d_φ_dr   = jac_φ_list[i]
        
        # Derivatives with respect to the imaginary part of the parameter (e.g., w_i)
        d_lnψ_di = jac_lnψ_list[i+1]
        d_φ_di   = jac_φ_list[i+1]
        
        # Using the correctly derived formula for O_w:
        # O_w = 0.5 * [ (d(lnψ)/dr + d(φ)/di) + i * (d(φ)/dr - d(lnψ)/di) ]
        O_w_real = 0.5 * (d_lnψ_dr + d_φ_di)
        O_w_imag = 0.5 * (d_φ_dr - d_lnψ_di)
        O_w = tf.complex(O_w_real, O_w_imag)
        
        # Reshape to [batch_size, num_elements_in_param] and append
        O_list.append(tf.reshape(O_w, [tf.shape(config)[0], -1]))

    # Concatenate into one large matrix O of shape [batch_size, total_num_complex_params]
    O = tf.concat(O_list, axis=1)
    O = tf.cast(O, tf.complex128)

    # Center the log-derivatives: O_k' = O_k - ⟨O_k⟩
    O_mean = tf.reduce_mean(O, axis=0, keepdims=True)
    O_centered = O - O_mean

    # 1. Calculate the force vector F_k = ⟨(E_loc - ⟨E⟩) * O_k'⟩. Note the conjugate on O.
    F = tf.reduce_mean(tf.expand_dims(tf.cast(E_centered, tf.complex128), -1) * tf.math.conj(O_centered), axis=0)
    F = tf.reshape(F, [-1, 1])

    # 2. Define the S-matrix multiplication S * x
    def matvec(x):
        batch_size = tf.cast(tf.shape(O_centered)[0], tf.complex128)
        # S_ij = ⟨O_i'* O_j'⟩
        Ox = tf.matmul(O_centered, x)
        O_conj_Ox = tf.matmul(O_centered, Ox, adjoint_a=True)
        Sx = O_conj_Ox / batch_size
        return Sx + diag_shift * x

    # 3. Solve the linear system S * δ = F using Conjugate Gradient for the update vector δ
    x = tf.zeros_like(F, dtype=tf.complex128)
    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) + 1e-14)
        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 + 1e-14)
        p_new = r_new + beta * p
        return i+1, r_new, p_new, x_new, rdotr_new

    i = tf.constant(0)
    cond = lambda i, *_: tf.logical_and(i < max_iter, tf.math.real(rdotr) > 1e-12)
    _, _, _, delta, _ = tf.while_loop(cond, body, [i, r, p, x, rdotr])

    # 4. Apply the update w' = w - lr * δ.
    lr = optimizer.learning_rate(step)
    lr_c = tf.cast(lr, tf.complex128)
    update_vector = -lr_c * delta

    # Unflatten the update vector to match the shapes of the complex parameters (weights and biases)
    w_shape = model.conv_layer.w_real.shape
    b_shape = model.conv_layer.b_real.shape
    w_size = tf.size(model.conv_layer.w_real)
    
    update_w = tf.reshape(update_vector[:w_size], w_shape)
    update_b = tf.reshape(update_vector[w_size:], b_shape)

    # Your custom layer method correctly applies the real and imag parts of the complex update
    model.conv_layer.apply_complex_grads([update_w], [update_b])
    
    return tf.reduce_mean(E_loc), tf.math.reduce_variance(E_loc)



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

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)
    
    # Pass the current step to the training function for the LR scheduler
    e, var = train_step_sr_cg(current_config, tf.constant(step, dtype=tf.int64))

    if step % 20 == 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 got the energy as below, but it should be ~50. Correct the wrong part.
Starting U(1) VMC: L=8, chains=128
E_init: 64.140968859492659
Step   0 | E 63.771777 | Var 28.8090 | T 4.241s
Step  20 | E 39.473553 | Var 21.6121 | T 1.289s
Step  40 | E 30.124487 | Var 9.5140 | T 1.306s
Step  60 | E 25.805975 | Var 11.2211 | T 1.319s