#!/usr/bin/env python3

import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as opt

def Main():

    ### this N is one less than the N in the sheet for simplicity
    N =
    nbar =

    x = np.arange(0.05, 0.5+1e-8, 1e-2)
    y = []

    ### you will want to calculate and plot other quantities
    w = []

    # Iterate over z*t products
    print(f'zt\tEn\tL_err')
    for zt in x:
        a, err = find_a(N, nbar, zt)
        energy = Energy(a, N, nbar, zt)
        print(f'{zt:.1e}\t{energy:+.2f}\t{err:+.2e}')
        y.append(energy)

    fig, ax = plt.subplots(1, 1)
    ax.plot(x, y, label='$\\varepsilon_t$')
    ax.set_xlabel('$zt$')
    ax.legend()
    plt.show()
    plt.close(fig)
    return

def find_a(N, nbar, zt):
    k = np.arange(N)

    ### x0 is the array of initial parameters
    x0 = np.zeros(N+2)
    x0[:N] =
    ### Initial Lambda parameters
    x0[-2], x0[-1] = -1, -1

    ### root search
    x = opt.root(Lprime, x0=x0, method='lm', args=(N, nbar, zt)).x

    ### Calculate the magnitude of |nabla L| to check convergence
    err =
    return x[:N], err

### Calculate nabla L
def Lprime(x, N, nbar, zt):
    a, λ = x[:N], x[N:]

    ### derivatives in a of L
    L_a = np.zeros(N)

    ### constraints
    g0 =
    g1 =

    return np.append(L_a, [g0, g1])

def Energy(a, N, nbar, zt):
    ### calculate energy per site here
    return

Main()