using AbstractAlgebra
using Groups
using GroupRings
using PropertyT

using SCS
solver(tol, iterations) =
    SCSSolver(linearsolver=SCS.Direct,
        eps=tol, max_iters=iterations,
        alpha=1.95, acceleration_lookback=1)

include("../main.jl")

using PropertyTGroups

args = Dict("SAut" => 5, "upper-bound" => 50.0, "radius" => 2, "nosymmetry"=>false, "tol"=>1e-12, "iterations" =>200000, "warmstart" => true)

Gr = PropertyTGroups.PropertyTGroup(args)
sett = PropertyT.Settings(Gr, args,
    solver(args["tol"], args["iterations"]))

@show sett

fullpath = PropertyT.fullpath(sett)
isdir(fullpath) || mkpath(fullpath)
# setup_logging(PropertyT.filename(fullpath, :fulllog), :fulllog)

function small_generating_set(RG::GroupRing{AutGroup{N}}, n) where N
    indexing = [(i,j) for i in 1:n for j in 1:n if i≠j]

    rmuls = [Groups.rmul_autsymbol(i,j) for (i,j) in indexing]
    lmuls = [Groups.lmul_autsymbol(i,j) for (i,j) in indexing]
    gen_set = RG.group.([rmuls; lmuls])

    return [gen_set; inv.(gen_set)]
end

function computeX(RG::GroupRing{AutGroup{N}}) where N
    Tn = small_generating_set(RG, N-1)

    ℤ = Int64
    Δn = length(Tn)*one(RG, ℤ) - RG(Tn, ℤ);

    Alt_N = [g for g in elements(PermutationGroup(N)) if parity(g) == 0]

    @time X = sum(σ(Δn)*sum(τ(Δn) for τ ∈ Alt_N if τ ≠ σ) for σ in Alt_N);
    return X
end

function Sq(RG::GroupRing{AutGroup{N}}) where N
    T2 = small_generating_set(RG, 2)
    ℤ = Int64
    Δ₂ = length(T2)*one(RG, ℤ) - RG(T2, ℤ);

    Alt_N = [g for g in elements(PermutationGroup(N)) if parity(g) == 0]
    elt = sum(σ(Δ₂)^2 for σ in Alt_N)
    return elt
end

function Adj(RG::GroupRing{AutGroup{N}}) where N
    T2 = small_generating_set(RG, 2)

    ℤ = Int64
    Δ₂ = length(T2)*one(RG, ℤ) - RG(T2, ℤ);

    Alt_N = [g for g in elements(PermutationGroup(N)) if parity(g) == 0]

    adj(σ::perm, τ::perm, i=1, j=2) = Set([σ[i], σ[j]]) ∩ Set([τ[i], τ[j]])
    adj(σ::perm) = [τ for τ in Alt_N if length(adj(σ, τ)) == 1]

    @time elt = sum(σ(Δ₂)*sum(τ(Δ₂) for τ in adj(σ)) for σ in Alt_N);
    # return RG(elt.coeffs÷factorial(N-2)^2)
    return elt
end

function Op(RG::GroupRing{AutGroup{N}}) where N
    T2 = small_generating_set(RG, 2)

    ℤ = Int64
    Δ₂ = length(T2)*one(RG, ℤ) - RG(T2, ℤ);

    Alt_N = [g for g in elements(PermutationGroup(N)) if parity(g) == 0]

    adj(σ::perm, τ::perm, i=1, j=2) = Set([σ[i], σ[j]]) ∩ Set([τ[i], τ[j]])
    adj(σ::perm) = [τ for τ in Alt_N if length(adj(σ, τ)) == 0]

    @time elt = sum(σ(Δ₂)*sum(τ(Δ₂) for τ in adj(σ)) for σ in Alt_N);
    # return RG(elt.coeffs÷factorial(N-2)^2)
    return elt
end

const ELT_FILE = joinpath(dirname(PropertyT.filename(sett, :Δ)), "SqAdjOp_coeffs.jld")
const WARMSTART_FILE = PropertyT.filename(sett, :warmstart)

if isfile(PropertyT.filename(sett,:Δ)) && isfile(ELT_FILE) &&
    isfile(PropertyT.filename(sett, :OrbitData))
    # cached
    Δ = PropertyT.loadGRElem(PropertyT.filename(sett,:Δ), sett.G)
    RG = parent(Δ)
    orbit_data = load(PropertyT.filename(sett, :OrbitData), "OrbitData")
    sq_c, adj_c, op_c = load(ELT_FILE, "Sq", "Adj", "Op")
    # elt = ELT_FILE, sett.G)
    sq = GroupRingElem(sq_c, RG)
    adj = GroupRingElem(adj_c, RG)
    op = GroupRingElem(op_c, RG);
else
    info("Compute Laplacian")
    Δ = PropertyT.Laplacian(sett.S, sett.radius)
    RG = parent(Δ)

    info("Compute Sq, Adj, Op")
    @time sq, adj, op = Sq(RG), Adj(RG), Op(RG)

    PropertyT.saveGRElem(PropertyT.filename(sett, :Δ), Δ)
    save(ELT_FILE, "Sq", sq.coeffs, "Adj", adj.coeffs, "Op", op.coeffs)

    info("Compute OrbitData")
    if !isfile(PropertyT.filename(sett, :OrbitData))
        orbit_data = PropertyT.OrbitData(parent(Y), sett.autS)
        save(PropertyT.filename(sett, :OrbitData), "OrbitData", orbit_data)
    else
        orbit_data = load(PropertyT.filename(sett, :OrbitData), "OrbitData")
    end
end;

orbit_data = PropertyT.decimate(orbit_data);

elt = adj+2op;

const SOLUTION_FILE = PropertyT.filename(sett, :solution)

if !isfile(SOLUTION_FILE)
    
    SDP_problem, varλ, varP = PropertyT.SOS_problem(elt, Δ, orbit_data; upper_bound=sett.upper_bound)

    begin
        using SCS
        scs_solver = SCS.SCSSolver(linearsolver=SCS.Direct,
            eps=sett.tol,
            max_iters=args["iterations"],
            alpha=1.95,
            acceleration_lookback=1)

        JuMP.setsolver(SDP_problem, scs_solver)
    end

    λ = Ps  = nothing
    ws = PropertyT.warmstart(sett)

    # using ProgressMeter

    # @showprogress "Running SDP optimization step... " for i in 1:args["repetitions"]
    while true
        λ, Ps, ws = PropertyT.solve(PropertyT.filename(sett, :solverlog),
        SDP_problem, varλ, varP, ws);

        if all((!isnan).(ws[1]))
            save(WARMSTART_FILE, "warmstart", ws, "λ", λ, "Ps", Ps)
            save(WARMSTART_FILE[1:end-4]*"_$(now()).jld", "warmstart", ws, "λ", λ, "Ps", Ps)
        else
            warn("No valid solution was saved!")
        end
    end
    
    info("Reconstructing P...")
    @time P = PropertyT.reconstruct(Ps, orbit_data);
    save(SOLUTION_FILE, "λ", λ, "P", P)
end

P, λ = load(SOLUTION_FILE, "P", "λ")
@show λ;

@time const Q = real(sqrtm(P));

function SOS_residual(eoi::GroupRingElem, Q::Matrix)
    RG = parent(eoi)
    @time sos = PropertyT.compute_SOS(RG, Q);
    return eoi - sos
end

info("Floating Point arithmetic:")
EOI = elt - λ*Δ
b = SOS_residual(EOI, Q)
@show norm(b, 1);

info("Interval arithmetic:")
using IntervalArithmetic
Qint = PropertyT.augIdproj(Q);
@assert all([zero(eltype(Q)) in sum(view(Qint, :, i)) for i in 1:size(Qint, 2)])

EOI_int = elt - @interval(λ)*Δ;
Q_int = PropertyT.augIdproj(Q);
@assert all([zero(eltype(Q)) in sum(view(Q_int, :, i)) for i in 1:size(Q_int, 2)])
b_int = SOS_residual(EOI_int, Q_int)
@show norm(b_int, 1);

info("λ is certified to be > ", (@interval(λ) - 2^2*norm(b_int,1)).lo)