# Build your own nuclear DFT -- step 4: from orbitals back to densities.
#
# Step 3 turned a potential into orbitals. Density functional theory needs the
# reverse map too: turn occupied orbitals into the three local densities the
# Skyrme force feeds on,
#
#   rho(r) = sum (2j+1)/4pi * R^2               number density
#   tau(r) = sum (2j+1)/4pi * (R'^2 + l(l+1)/r^2 R^2)     kinetic density
#   J(r)   = sum (2j+1)/4pi * <sigma.l>/r * R^2           spin-orbit density
#
# where R(r) = u(r)/(r sqrt(h)) recovers the properly normalized radial wave
# function from the unit-norm eigenvector. The sum runs over OCCUPIED orbitals
# only -- for 16O that is 1s1/2, 1p3/2, 1p1/2, filled by 8 neutrons and 8
# protons, which in this Coulomb-free step are identical: solve one species
# and double it.
#
# Sanity checks a real code lives on: integrating rho must count 16 nucleons;
# J should nearly vanish because 16O is spin-saturated (both 1p partners full,
# their <sigma.l> weights cancel: 1*(2*1+2) + (-2)*(2*1) = 0).

pi = 3.141592653589793
hbar2_2m = 20.7355
h = 0.1
nr = 140
amass = 16.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
}

# np.gradient-style derivative: central differences, one-sided 2nd order edges
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
}

v0 = -51.0
rws = 1.27 * pow(amass, 1.0 / 3.0)
aws = 0.67
u_pot = zeros(nr)
wso = zeros(nr)
for i to nr {
    f = 1.0 / (1.0 + exp((r[i] - rws) / aws))
    u_pot[i] = v0 * f
    wso[i] = 22.0 * (0.0 - f * (1.0 - f) / aws)
}

# ---- block solver (step 3, unchanged) ----------------------------------------
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 {
        diag[i] = 2.0 * hbar2_2m / (h * h) + hbar2_2m * l * (l + 1.0) / (r[i] * r[i]) + u_pot[i] + sdl * wso[i] / r[i]
        if i < nr - 1.0 { off[i] = 0.0 - hbar2_2m / (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
}

# ---- NEW: orbital -> densities ------------------------------------------------
rho = zeros(nr)
tau = zeros(nr)
jso = zeros(nr)
R_tmp = zeros(nr)
dR = zeros(nr)

def add_orbital(deg, l, sdl) {
    # take the freshly solved state 0 out of ev, rebuild R = u/(r sqrt h)
    for i to nr { R_tmp[i] = ev[0.0 * 512.0 + i] / sqrt(h) / r[i] }
    grad1(R_tmp, dR, nr, h)
    w = deg / (4.0 * pi)
    for i to nr {
        rho[i] = rho[i] + w * R_tmp[i] * R_tmp[i]
        tau[i] = tau[i] + w * (dR[i] * dR[i] + l * (l + 1.0) / (r[i] * r[i]) * R_tmp[i] * R_tmp[i])
        jso[i] = jso[i] + w * sdl / r[i] * R_tmp[i] * R_tmp[i]
    }
    return 0.0
}

# occupied orbitals of one species (8 nucleons), then double for N = Z
build_h(0.0, 0.0)
lowest_k(nr, 1.0)
add_orbital(2.0, 0.0, 0.0)      # 1s1/2

build_h(1.0, 1.0)
lowest_k(nr, 1.0)
add_orbital(4.0, 1.0, 1.0)      # 1p3/2

build_h(1.0, -2.0)
lowest_k(nr, 1.0)
add_orbital(2.0, 1.0, -2.0)     # 1p1/2

for i to nr {
    rho[i] = 2.0 * rho[i]       # neutrons + protons
    tau[i] = 2.0 * tau[i]
    jso[i] = 2.0 * jso[i]
}

a_check = integrate_r2(rho, nr, h)
show a_check          # 16: the orbitals hold exactly the nucleus

show rho              # compare step 1's guess: same size, more structure
show tau              # kinetic density: where the wavefunctions wiggle
show jso              # nearly zero everywhere: 16O is spin-saturated