# Build your own nuclear DFT -- step 6: close the loop.
#
# Steps 3 and 5 are the two halves of a circle: fields make orbitals, orbitals
# make densities, densities make fields. Self-consistency is just running that
# circle until it stops moving:
#
#   guess rho -> build B, U, W -> solve the blocks -> rebuild rho -> MIX -> repeat
#
# Two things are new and both matter.
#
# (1) The kinetic operator must be re-discretized. B(r) now varies with
#     position, and the naive stencil for -d/dr[B du/dr] is NOT symmetric --
#     a non-Hermitian matrix means complex energies. The fix: average B to the
#     half-points B_{i+1/2} and use the SAME value in the diagonal and the
#     off-diagonal it couples. (This is the trick the production code's bug
#     stories orbit around.)
# (2) Mixing. Feeding the new density straight back (mix = 1) overshoots and
#     oscillates; blending 40% new into 60% old converges smoothly. Watch
#     e_hist for the approach.
#
# What we track here is the sum of occupied single-particle energies -- an
# easy convergence monitor, but NOT the binding energy: each eps contains the
# interaction with every other nucleon, so summing them counts every pair
# TWICE. Step 7 computes the energy the right way and the number becomes
# physical. (Still N = Z, still no Coulomb: one species, doubled.)

pi = 3.141592653589793
hbar2_2m = 20.7355
h = 0.1
nr = 140
amass = 16.0

t0 = -1128.75
t1 = 395.0
t2 = -95.0
t3 = 14000.0
x0 = 0.45
w0 = 120.0

r = zeros(nr)
for i to nr { r[i] = (i + 1.0) * h }

def integrate_r2(f, n, hh) {
    s = 0.5 * f[0] * r[0] * r[0] + 0.5 * f[n - 1.0] * r[n - 1.0] * r[n - 1.0]
    for i to n {
        if i > 0 {
            if i < n - 1.0 { s += f[i] * r[i] * r[i] }
        }
    }
    return 4.0 * pi * s * hh
}

def grad1(f, out, n, hh) {
    out[0] = (-1.5 * f[0] + 2.0 * f[1] - 0.5 * f[2]) / hh
    for i to n {
        if i > 0 {
            if i < n - 1.0 { out[i] = (f[i + 1] - f[i - 1]) / (2.0 * hh) }
        }
    }
    out[n - 1.0] = (1.5 * f[n - 1.0] - 2.0 * f[n - 2.0] + 0.5 * f[n - 3.0]) / hh
    return 0.0
}

# ---- block solver, upgraded: variable mass with half-point Hermitization -----
b_field = zeros(nr)
u_field = zeros(nr)
w_field = zeros(nr)
diag = zeros(nr)
off = zeros(nr)
u_new = zeros(nr)
u_cur = zeros(nr)
tc = zeros(nr)
td = zeros(nr)
ev = zeros(1024)
eps_k = zeros(2)

def build_h(l, sdl) {
    for i to nr {
        bm = b_field[0]
        if i > 0 { bm = 0.5 * (b_field[i - 1.0] + b_field[i]) }
        bp = 0.5 * (b_field[i] + hbar2_2m)
        if i < nr - 1.0 { bp = 0.5 * (b_field[i] + b_field[i + 1.0]) }
        diag[i] = (bm + bp) / (h * h) + b_field[i] * l * (l + 1.0) / (r[i] * r[i]) + u_field[i] + sdl * w_field[i] / r[i]
        if i < nr - 1.0 { off[i] = 0.0 - bp / (h * h) }
    }
    return 0.0
}

def lowest_k(n, kk) {
    sigma = diag[0]
    for i to n {
        g = diag[i]
        if i > 0 { g = g - abs(off[i - 1.0]) }
        if i < n - 1.0 { g = g - abs(off[i]) }
        if g < sigma { sigma = g }
    }
    sigma = sigma - 1.0
    for s to kk {
        for i to n { u_new[i] = sin((s + 1.0) * pi * (i + 1.0) / (n + 1.0)) }
        lam = 0.0
        for sweep to 300 {
            for i to n { u_cur[i] = u_new[i] }
            tc[0] = off[0] / (diag[0] - sigma)
            td[0] = u_cur[0] / (diag[0] - sigma)
            for i to n {
                if i > 0 {
                    m = diag[i] - sigma - off[i - 1.0] * tc[i - 1.0]
                    if i < n - 1.0 { tc[i] = off[i] / m }
                    td[i] = (u_cur[i] - off[i - 1.0] * td[i - 1.0]) / m
                }
            }
            u_new[n - 1.0] = td[n - 1.0]
            k = n - 2.0
            while k >= 0.0 {
                u_new[k] = td[k] - tc[k] * u_new[k + 1.0]
                k = k - 1.0
            }
            for t to s {
                dot = 0.0
                for i to n { dot += ev[t * 512.0 + i] * u_new[i] }
                for i to n { u_new[i] = u_new[i] - dot * ev[t * 512.0 + i] }
            }
            nrm = 0.0
            for i to n { nrm += u_new[i] * u_new[i] }
            nrm = sqrt(nrm)
            for i to n { u_new[i] = u_new[i] / nrm }
            newlam = 0.0
            for i to n {
                newlam += diag[i] * u_new[i] * u_new[i]
                if i < n - 1.0 { newlam += 2.0 * off[i] * u_new[i] * u_new[i + 1.0] }
            }
            dl = abs(newlam - lam)
            lam = newlam
            if sweep > 2.0 {
                if dl < 1.0e-12 { break }
            }
        }
        eps_k[s] = lam
        for i to n { ev[s * 512.0 + i] = u_new[i] }
    }
    return 0.0
}

# ---- densities ---------------------------------------------------------------
rho = zeros(nr)
tau = zeros(nr)
jso = zeros(nr)
nrho = zeros(nr)
ntau = zeros(nr)
njso = zeros(nr)
R_tmp = zeros(nr)
dR = zeros(nr)

def add_orbital(deg, l, sdl) {
    for i to nr { R_tmp[i] = ev[i] / sqrt(h) / r[i] }
    grad1(R_tmp, dR, nr, h)
    w = deg / (4.0 * pi)
    for i to nr {
        nrho[i] = nrho[i] + w * R_tmp[i] * R_tmp[i]
        ntau[i] = ntau[i] + w * (dR[i] * dR[i] + l * (l + 1.0) / (r[i] * r[i]) * R_tmp[i] * R_tmp[i])
        njso[i] = njso[i] + w * sdl / r[i] * R_tmp[i] * R_tmp[i]
    }
    return 0.0
}

# Woods-Saxon guess for the total density (tau, J start at zero)
r0ws = 1.2 * pow(amass, 1.0 / 3.0)
for i to nr { rho[i] = 1.0 / (1.0 + exp((r[i] - r0ws) / 0.6)) }
norm = integrate_r2(rho, nr, h)
for i to nr { rho[i] = rho[i] * amass / norm }

# ---- the loop ------------------------------------------------------------------
mix = 0.4
niter = 60
e_hist = zeros(niter)
drho = zeros(nr)
d1 = zeros(nr)
d2 = zeros(nr)
lap = zeros(nr)
djso = zeros(nr)
e_sum = 0.0

for iter to niter {
    # densities -> fields (N = Z symmetric: rho_q = rho/2 etc., no Coulomb)
    grad1(rho, drho, nr, h)
    grad1(rho, d1, nr, h)
    grad1(d1, d2, nr, h)
    for i to nr { lap[i] = d2[i] + 2.0 / r[i] * d1[i] }
    grad1(jso, djso, nr, h)
    for i to nr {
        b_field[i] = hbar2_2m + 0.125 * (2.0 * t1 + 2.0 * t2) * rho[i] + 0.125 * (t2 - t1) * 0.5 * rho[i]
        u_field[i] = t0 * ((1.0 + x0 / 2.0) * rho[i] - (x0 + 0.5) * 0.5 * rho[i])
        u_field[i] = u_field[i] + 0.25 * t1 * ((tau[i] - 1.5 * lap[i]) - 0.5 * (0.5 * tau[i] - 0.75 * lap[i]))
        u_field[i] = u_field[i] + 0.25 * t2 * ((tau[i] + 0.5 * lap[i]) + 0.5 * (0.5 * tau[i] + 0.25 * lap[i]))
        u_field[i] = u_field[i] + 0.25 * t3 * 0.5 * rho[i] * (rho[i] + 0.5 * rho[i])
        u_field[i] = u_field[i] - 0.5 * w0 * (1.5 * jso[i] / r[i] + 0.75 * djso[i])
        w_field[i] = 0.5 * w0 * 1.5 * drho[i] + 0.125 * (t1 - t2) * 0.5 * jso[i]
    }

    # fields -> orbitals -> new densities
    for i to nr {
        nrho[i] = 0.0
        ntau[i] = 0.0
        njso[i] = 0.0
    }
    e_sum = 0.0
    build_h(0.0, 0.0)
    lowest_k(nr, 1.0)
    add_orbital(2.0, 0.0, 0.0)
    e_sum += 2.0 * eps_k[0]
    build_h(1.0, 1.0)
    lowest_k(nr, 1.0)
    add_orbital(4.0, 1.0, 1.0)
    e_sum += 4.0 * eps_k[0]
    build_h(1.0, -2.0)
    lowest_k(nr, 1.0)
    add_orbital(2.0, 1.0, -2.0)
    e_sum += 2.0 * eps_k[0]
    e_sum = 2.0 * e_sum                 # both species

    # mix: 40% new density into 60% old
    for i to nr {
        rho[i] = (1.0 - mix) * rho[i] + mix * 2.0 * nrho[i]
        tau[i] = (1.0 - mix) * tau[i] + mix * 2.0 * ntau[i]
        jso[i] = (1.0 - mix) * jso[i] + mix * 2.0 * njso[i]
    }
    e_hist[iter] = e_sum
}

show e_hist          # watch self-consistency arrive: big swing, then flat
show e_sum           # sum of occupied eps -- NOT the energy (pairs counted twice)
show rho             # the self-consistent density: no Woods-Saxon was harmed