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" => 10.0, "radius" => 2, "nosymmetry"=>false, "tol"=>1e-10, "iterations" =>200000, "repetitions"=>5)

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 (p::perm)(A::GroupRingElem)
    RG = parent(A)
    T = eltype(A.coeffs)
    result = zero(RG, T)

    for (idx, c) in enumerate(A.coeffs)
        if c!= zero(T)
            result[p(RG.basis[idx])] = c
        end
    end
    return result
end

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+3op;

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

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

info("Reconstructing P...")
@time P = PropertyT.reconstruct(Ps, orbit_data);
save(PropertyT.filename(sett, :solution), "λ", λ, "P", P)

addprocs(4)
@everywhere using PropertyT

@time const Q = real(sqrtm(P));
const EOI = elt - λ*Δ;
@show λ - 2^3*check_SOS_precision(Q, EOI);

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

EOIint = GroupRingElem([@interval(c) for c in EOI.coeffs], parent(Δ));
b_int = check_SOS_precision(Qint, EOIint)
@show @interval(λ) - 2^3*b_int;