fine-tune.jl

using Optim
using Roots
using LinearAlgebra
using LineSearches

using Pkg

evaluations = 0

local_path = joinpath(@__DIR__, "..", "AstroUtils.jl")

if isdir(local_path)
    Pkg.develop(path=local_path)
else
    Pkg.add(url="https://github.com/lamont-granquist/AstroUtils.jl.git")
end

using AstroUtils

# Earth grav param
const μ0 = 398600435436096
# Moon grav param
const μ1 = 4902800066163.8
# Moon soi
const soi1 = 66167158.6569544

# Spacecraft initial conditions
const r0 = [-6419931.35599855, -1598504.22548976, 1968486.96029449]
const v0 = [-747.608332256415, -9974.07871417746, -3695.67552943825]

# Moon initial conditions
const r1 = [4781917.4840911, -344412903.013029, -122488006.928327]
const v1 = [996.891435624119, 120.171877824752, -367.753988404139]

# Periapsis target
#const perT = 2_000_000
# Inclination target
#const incT = deg2rad(90)

#const perT = 4.473096171177541e7
#const incT = 2.629887001719956
const perT = 1e4
#const incT = deg2rad(38.2090635905795)
#const incT = 2.6
const incT = deg2rad(90)

#
# Scaling
#

const r_scale0 = sqrt(norm(r0)*norm(r1))
const v_scale0 = sqrt(μ0/r_scale0)
const t_scale0 = r_scale0 / v_scale0

const r_scale1 = sqrt(perT*soi1)
const v_scale1 = sqrt(μ1/r_scale1)
const t_scale1 = r_scale1 / v_scale1

const soi1s = soi1 / r_scale0
const r0s = r0 / r_scale0
const v0s = v0 / v_scale0
const r1s = r1 / r_scale0
const v1s = v1 / v_scale0

const perTs = perT / r_scale1

function distance(mu, dt, r0s, v0s, r1s, v1s)
    r0sf, v0sf = twobody(mu, dt, r0s, v0s)
    r1sf, v1sf = twobody(mu, dt, r1s, v1s)
    return norm(r1sf - r0sf)
end

# XXX: many edge conditions in this algorithm: type 2 transfers, multirotation, etc
function time_for_min_distance(mu, r0s, v0s, r1s, v1s)
    dtmax = if sma_from_state_vectors(mu, r0s, v0s) > 0
        period_from_state_vectors(mu, r0s, v0s)
    else
        time_to_next_radius(mu, r0s, v0s, 2*norm(r1s))
    end
    result = optimize(dt->distance(mu, dt, r0s, v0s, r1s, v1s), 0, dtmax)
    dt = Optim.minimizer(result)
    return dt
end

function residuals(dv)
    maxdt = time_for_min_distance(1.0, r0s, v0s + dv, r1s, v1s)
    minr = distance(1.0, maxdt, r0s, v0s + dv, r1s, v1s)
    dt = if minr <= soi1s
        find_zero(dt->distance(1.0, dt, r0s, v0s + dv, r1s, v1s) - soi1s, (0, maxdt), Roots.Brent())  # FIXME: failures?
    else
        maxdt
    end
    r0sf, v0sf = twobody(1.0, dt, r0s, v0s + dv)
    r1sf, v1sf = twobody(1.0, dt, r1s, v1s)
    rsoi = r0sf - r1sf
    vsoi = v0sf - v1sf
    rsoi = rsoi * r_scale0/r_scale1
    vsoi = vsoi * v_scale0/v_scale1
    per = if minr <= soi1s
        periapsis_from_state_vectors(1.0, rsoi, vsoi)
    else
        minr * r_scale0/r_scale1
    end
    h1 = normalize(cross(rsoi, vsoi))[3] - cos(incT)
    h2 = per - perTs
    return [0, h2]
    return [h1, h2]
end

function objective(dv)
    return 0.5 * dot(dv, dv)
end

function auglag(x, λ, ρ)
    # Augmented Lagrangian: f(x) + (ρ/2) Σ (hᵢ + λᵢ/ρ)²
    h = residuals(x)
    L = objective(x)
    for i in eachindex(h)
        s = h[i] + λ[i] / ρ
        L += 0.5 * ρ * s^2
    end
    return L
end

function auglag_minimize(x0;
        τ       = 0.5,     # ICM must shrink by this factor or ρ increases
        γ       = 10.0,    # factor by which ρ grows when ICM stalls
        λ_max   = 1e20,    # clamp magnitude for multiplier estimates
        xtol    = 1e-8,    # convergence: x stopped changing
        ftol    = 1e-8,    # convergence: objective stopped changing
        htol    = 1e-8,    # convergence: all constraints within this tolerance
        maxiter = 1000     # maximum outer (augmented Lagrangian) iterations
    )

    # ── Initialise working state ──

    x = copy(x0)
    h = residuals(x)
    pp = length(h)              # total number of scalar equality constraints
    λ = zeros(pp)               # multiplier estimates, initialised to zero

    # Initial ρ: balance objective scale against constraint violation scale
    # (heuristic from Birgin & Martinez)
    con2 = sum(h .^ 2)
    f0 = objective(x)
    ρ = con2 > 0 ? clamp(2 * abs(f0) / con2, 1e-6, 10.0) : 10.0

    # Track the best solution found so far
    best_x = copy(x)
    best_f = f0
    best_h = copy(h)
    best_penalty = sum(abs.(h))
    best_feasible = all(abs.(h) .<= htol)

    ICM = Inf   # infeasibility measure (worst constraint violation)

    # ── Main outer loop ──

    for iter in 1:maxiter
        prev_ICM = ICM

        # ── Inner solve: minimise augmented Lagrangian with BFGS ──

        result = optimize(
                          z -> auglag(z, λ, ρ), x,
                          BFGS(
                               linesearch = BackTracking(order=3),
                               alphaguess = InitialStatic(alpha=0.1, scaled=true)
                              ),
                          Optim.Options(
                                        g_tol = 1e-8,
                                        f_reltol = 1e-12,
                                        x_abstol = 1e-10,
                                        iterations = 200
                                       )
                         )
        x = Optim.minimizer(result)

        # ── Re-evaluate true objective and constraints ──

        fcur = objective(x)
        h = residuals(x)

        # ── Update multipliers and compute infeasibility ──

        ICM = 0.0
        penalty = 0.0
        feasible = true

        for i in eachindex(h)
            hi = h[i]

            # First-order multiplier update: λ ← λ + ρ h(x)
            λ[i] = clamp(λ[i] + ρ * hi, -λ_max, λ_max)

            penalty += abs(hi)
            feasible = feasible && abs(hi) <= htol
            ICM = max(ICM, abs(hi))
        end

        # ── Increase penalty if infeasibility didn't shrink enough ──

        if ICM > τ * prev_ICM
            ρ *= γ
        end

        if ρ == Inf
            break
        end

        # ── Update best-known solution ──
        #
        # Accept if: (a) feasible and improves on previous best, or
        #            (b) reduces total constraint violation when no
        #                feasible point has been found yet.

        if (feasible && (!best_feasible || penalty < best_penalty || fcur < best_f)) ||
            (!best_feasible && penalty < best_penalty)

            prev_best_f = best_f
            prev_best_x = best_x

            best_x = copy(x)
            best_f = fcur
            best_h = copy(h)
            best_penalty = penalty
            best_feasible = feasible

            # ── Check convergence (only meaningful once feasible) ──

            if feasible
                if abs(fcur - prev_best_f) < ftol * (1 + abs(fcur))
                    break
                end
                if maximum(abs.(x .- prev_best_x)) < xtol * (1 + maximum(abs.(x)))
                    break
                end
            end
        end

        # Perfect feasibility — nothing left for the outer loop to do
        if ICM == 0
            break
        end
    end

    return best_x, best_f, best_h
end

x0 = [0.0,0.0,0.0]

(x, f, h) = auglag_minimize(x0)

display(x * v_scale0)
display(h)