using Distributed
using Random
using Logging

workers = 8

if nprocs() <= workers
    addprocs(workers + 1 - nprocs())
end

@everywhere include("models/Vibron.jl")
@everywhere include("modules/ClassicalDynamics.jl")

@everywhere disable_logging(Logging.Info)

""" Calculates Poincaré sections with Lyapunov exponents for various energies
    Parallel calculation, takes usually days to finish
"""

@everywhere function SolveItem(energy, parameters, dimension; file="Vibron.txt", path="", alreadySolved=[], sectionPlane=2)    
    for (E, λ) in alreadySolved
        if isapprox(E, energy) && isapprox(parameters[1], λ)
            println("Skipped $parameters, E=$energy, dim=$dimension (already calculated)")
            return false
        end
    end

    time = @elapsed averageLyapunov, freg, trajectories, lyapunovs = SolveEnergy(energy, parameters, dimension, savePath=path, timeout=7200, showFigures=false, randomize=true)

    chaos = 0
    total = 0
    error = 0
    meanLyapunov = 0
    meanLyapunovChaos = 0

    for x in averageLyapunov
        if x > 0.0
            total += 1
            meanLyapunov += x
        end
        if x == 0.0
            error += 1
        end
        if x > 0.01
            chaos += 1
            meanLyapunovChaos += x
        end        
    end    

    varλ = 0
    if length(lyapunovs) == 0
        maxλ = 0
        meanλ = 0
    else
        maxλ = maximum(lyapunovs)
        meanλ = mean(lyapunovs)
        if length(lyapunovs) > 1
            varλ = var(lyapunovs)
        end
    end

    result = [energy, parameters[1], total, chaos, error, freg, total > 0 ? meanLyapunov / total : 0, chaos > 0 ? meanLyapunovChaos / chaos : 0, length(lyapunovs), maxλ, meanλ, varλ, myid(), time, trajectories]

    open(path * file, "a") do io
        println(io, result)
    end

    return true
end

function RunMapC(; C=0.2, path="", dimension=101, step=0.1, sectionPlane=3)
    path *= "Vibron_"

    file = "Map_dim=$(dimension)_C=$C.txt"
    alreadySolved = ReadMap(path * file)

    input = shuffle([(energy, [0.5 * (1 - A), -A, C], dimension) for energy in -1.2:step:1, A in 0:step:1])

    println("To be calculated $(length(input)).")
    println("Already calculated $(length(alreadySolved)) points.")

    pmap((args)->SolveItem(args...; file=file, path=path, alreadySolved=alreadySolved, sectionPlane=sectionPlane), input)

    return
end

""" Calculates freg for one given A and C """
function RunAC(; C=0.2, A=0.4, path="", dimension=101, step=0.01, sectionPlane=3)
    path *= "Vibron_"

    file = "Energy_dim=$(dimension)_$([A, C]).txt"
    alreadySolved = ReadMap(path * file)

    input = [(energy, [0.5 * (1 - A), -A, C], dimension) for energy in -1.2:step:1]

    println("To be calculated $(length(input)).")
    println("Already calculated $(length(alreadySolved)) points.")

    pmap((args)->SolveItem(args...; file=file, path=path, alreadySolved=alreadySolved, sectionPlane=sectionPlane), input)    

    return
end

function ReadMap(file="")
    Eλ = []

    try
        for line in eachline(file)     
            line = replace(line, "[" => "")   
            line = replace(line, "]" => "")
            line = replace(line, "," => "")
            elements = split(line)

            append!(Eλ, [(parse(Float64, elements[1]), parse(Float64, elements[2]))])
        end
    catch x
    end

    return Eλ
end