module SexSpecificMigrations

using ForwardPopGenSimulations

# reexport from ForwardPopGenSimulations as users use these to specify model parameters.
export Female,
       Male,
       Mother,
       Father,
       Daughter,
       Son,

       Autosome,
       XChromosome,
       YChromosome,
       Mitochondrion,

       BasicData

export ModelParameters,
       LocusList,
       PopulationSizes,
       MutationRates,
       RecombinationRates,
       ProportionsOfMigrants,
       FitnessOfMigrants,
       FitnessOfForeignTraits,
       TraitTransmissionRates,

       simulate

immutable ForeignTrait end
immutable Migrant end

immutable LocusList
    val::Array{Tuple{ChromosomeType, Int}, 1}
end
Base.call(::Type{LocusList}, val::Tuple{ChromosomeType, Int}...) = LocusList([val...])
Base.length(t::LocusList) = length(t.val)
Base.start(t::LocusList) = start(t.val)
Base.next(t::LocusList, s) = next(t.val, s)
Base.done(t::LocusList, s) = done(t.val, s)

abstract Parameter
abstract DemeAndSexSpecificParameter <: Parameter
abstract TwoGensSexSpecificParameter <: Parameter
abstract SexSpecificParameter <: Parameter
Base.getindex(p::Parameter, i::Integer) = p.data[i]

Base.length(i::DemeAndSexSpecificParameter) = length(i.data)

function Base.call{T<:DemeAndSexSpecificParameter, V<:Number}(
    ::Type{T}, vals::Pair{Int, NTuple{2, Pair{SexType, V}}}...)

    d = Dict([val[1] => Dict(val[2]) for val in vals])
    demes = sort(collect(keys(d)))
    T(tuple([(d[deme][Female], d[deme][Male]) for deme in demes]...))
end

function Base.call{T<:TwoGensSexSpecificParameter, V<:Number}(
    ::Type{T}, vals::Pair{SexType, NTuple{2, Pair{SexType, V}}}...)

    d = Dict([val[1] => Dict(val[2]) for val in vals])
    T(((d[Mother][Daughter], d[Mother][Son]), (d[Father][Daughter], d[Father][Son])))
end

function Base.call{T<:SexSpecificParameter, V<:Number}(
    ::Type{T}, vals::Pair{SexType, V}...)
    d = Dict(vals)
    T((d[Female], d[Male]))
end

immutable PopulationSizes <: DemeAndSexSpecificParameter
    # deme and sex-specific
    data::NTuple{2, NTuple{2, Int}}
    PopulationSizes(d::NTuple{2, NTuple{2, Int}}) = new(d)
end

immutable MutationRates <: SexSpecificParameter
    # sex-specific
    data::NTuple{2, Float64}
    MutationRates(d::NTuple{2, Float64}) = new(d)
end

immutable RecombinationRates <: SexSpecificParameter
    # sex-specific
    data::NTuple{2, Float64}
    RecombinationRates(d::NTuple{2, Float64}) = new(d)
end

immutable ProportionsOfMigrants <: DemeAndSexSpecificParameter
    # deme and sex-specific
    data::NTuple{2, NTuple{2, Float64}}
    ProportionsOfMigrants(d::NTuple{2, NTuple{2, Float64}}) = new(d)
end

abstract Fitness <: DemeAndSexSpecificParameter
immutable FitnessOfMigrants <: Fitness
    # deme and sex-specific
    data::NTuple{2, NTuple{2, Float64}}
    FitnessOfMigrants(d::NTuple{2, NTuple{2, Float64}}) = new(d)
end

immutable FitnessOfForeignTraits <: Fitness
    # deme and sex-specific
    data::NTuple{2, NTuple{2, Float64}}
    FitnessOfForeignTraits(d::NTuple{2, NTuple{2, Float64}}) = new(d)
end

immutable TraitTransmissionRates <: TwoGensSexSpecificParameter
    # sex-specific in parents and offspring
    data::NTuple{2, NTuple{2, Float64}}
    TraitTransmissionRates(d::NTuple{2, NTuple{2, Float64}}) = new(d)
end

function trait(momrate, poprate, momtrait, poptrait, deme)
    if deme == momtrait == poptrait
        deme
    elseif deme != momtrait && deme != poptrait
        rand() < 1.0 - (1.0 - momrate) * (1.0 - poprate) ? momtrait : deme
    elseif deme != momtrait
        rand() < momrate ? momtrait : deme
    else
        rand() < poprate ? poptrait : deme
    end
end

immutable ModelParameters
    chromosomes::LocusList
    population_sizes::PopulationSizes
    mutation_rates::MutationRates
    recombination_rates::RecombinationRates
    proportions_of_migrant::ProportionsOfMigrants
    fitness_of_migrants::FitnessOfMigrants
    fitness_of_foreign_traits::FitnessOfForeignTraits
    transmission_rates_of_trait::TraitTransmissionRates
end

immutable SexSpecificParameters
    deme::Int
    population_size::Int
    number_of_loci::Int
    sex::SexType
    mutation_rate::Float64
    recombination_rate::Array{Float64, 1}
    type_of_locus::Array{ChromosomeType, 1}
    proportion_of_migrant::Float64
    fitness_of_migrant::Float64
    fitness_of_foreign_trait::Float64
    transmission_rate_of_trait::NTuple{2, Float64}
end

function SexSpecificParameters(p::ModelParameters, deme::Integer, sex::SexType)
    sexidx = value(sex)

    chrs = p.chromosomes
    recrates = vcat([[i == 1 ? 0.5 : recrate for i in 1:locus[2]] for locus in chrs]...)
    chrtypes = vcat([[locus[1] for _ in 1:locus[2]] for locus in chrs]...)

    length(recrates) == length(chrtypes) || error()

    nloci = length(chrtypes)

    SexSpecificParameters(
        deme,
        p.population_sizes[deme][sexidx],
        nloci,
        sex,
        p.mutation_rates[sexidx],
        recrates,
        chrtypes,
        p.proportions_of_migrant[deme][sexidx],
        p.fitness_of_migrants[deme][sexidx],
        p.fitness_of_foreign_traits[deme][sexidx],
        p.transmission_rates_of_trait[sexidx]
    )
end

type SexSpecificPopulation
    data::Array{Int, 2} # Chromosome is along row, and locus is along column.
    trait::Array{Int, 1}
    params::SexSpecificParameters
end

function SexSpecificPopulation(params::ModelParameters, deme::Int, sex::SexType)
    ps = SexSpecificParameters(params, deme, sex)
    nloci = length(ps.type_of_locus)
    popsize = ps.population_size
    SexSpecificPopulation(zeros(Int, 2, nloci * popsize), fill(deme, popsize), ps)
end

get(p::SexSpecificPopulation, s::Symbol) = getfield(p.params, s)
get(p::SexSpecificPopulation, s::Symbol, sex::SexType) = getfield(p.params, s)[value(sex)]
get(p::SexSpecificPopulation, s::Symbol, i::Integer) = getfield(p.params, s)[i]
Base.setindex!(p::SexSpecificPopulation, val::Int, ::Type{Val{:trait}}, i::Integer) = p.trait[i] = val
Base.getindex(p::SexSpecificPopulation, ::Type{Val{:trait}}, i::Integer) = p.trait[i]
Base.length(p::SexSpecificPopulation) = p.params.population_size
Base.eachindex(p::SexSpecificPopulation) = eachindex(p.data)
Base.getindex(p::SexSpecificPopulation, i::Integer) = p.data[i]
Base.getindex(p::SexSpecificPopulation, c::Integer, l::Integer) = p.data[c, l]
Base.setindex!(p::SexSpecificPopulation, val::Int, i::Integer) = p.data[i] = val
Base.setindex!(p::SexSpecificPopulation, val::Int, c::Integer, l::Integer) = p.data[c, l] = val

type Population
    data::NTuple{2, SexSpecificPopulation}
end
Base.length(p::Population) = length(p.data)
Base.getindex(p::Population, i::Integer) = p.data[i]
Base.start(p::Population) = 0
Base.next(p::Population, state) = (state += 1; (p.data[state], state))
Base.done(p::Population, state) = length(p.data) == state
Base.eachindex(p::Population) = 1:length(p.data)

function createpopulations(params::ModelParameters)
    ndemes = length(params.population_sizes)
    pops = Array{Population}(ndemes)
    sexes = (Female, Male)
    for i = eachindex(pops)

        spops = Array{SexSpecificPopulation}(2)
        for (j, sex) = enumerate(sexes)
            spops[j] = SexSpecificPopulation(params, i, sex)
        end

        pops[i] = Population(tuple(spops...))
    end
    pops
end

function setinitialgenotype2!(core, p, i)
    for j = 1:2
        p[j,i] = insert!(db(core), GeneRecord(0, nextstate!(core)))
    end
    nothing
end

function setinitialgenotype1!(core, p, i, j)
    p[j,i] = insert!(db(core), GeneRecord(0, nextstate!(core)))
    nothing
end

function setinitialgenotype!(
    core::BasicData,
    p::SexSpecificPopulation,
    ind,
    locus,
    ::Type{Val{Female}},
    ::Union(Type{Val{Autosome}}, Type{Val{XChromosome}}))

    setinitialgenotype2!(core, p, toposition(p, ind, locus))
end

function setinitialgenotype!(
    core::BasicData,
    p::SexSpecificPopulation,
    ind,
    locus,
    ::Type{Val{Female}},
    ::Type{Val{YChromosome}})

    nothing
end

function setinitialgenotype!(
    core::BasicData,
    p::SexSpecificPopulation,
    ind,
    locus,
    ::Type{Val{Female}},
    ::Type{Val{Mitochondrion}})

    setinitialgenotype1!(core, p, toposition(p, ind, locus), 1)
end

function setinitialgenotype!(
    core::BasicData,
    p::SexSpecificPopulation,
    ind,
    locus,
    ::Type{Val{Male}},
    ::Type{Val{Autosome}})

    setinitialgenotype2!(core, p, toposition(p, ind, locus))
end

function setinitialgenotype!(
    core::BasicData,
    p::SexSpecificPopulation,
    ind,
    locus,
    ::Type{Val{Male}},
    ::Type{Val{XChromosome}})

    setinitialgenotype1!(core, p, toposition(p, ind, locus), 1)
end

function setinitialgenotype!(
    core::BasicData,
    p::SexSpecificPopulation,
    ind,
    locus,
    ::Type{Val{Male}},
    ::Type{Val{YChromosome}})

    setinitialgenotype1!(core, p, toposition(p, ind, locus), 2)
end

function setinitialgenotype!(
    core::BasicData,
    p::SexSpecificPopulation,
    ind,
    locus,
    ::Type{Val{Male}},
    ::Type{Val{Mitochondrion}})

    nothing
end

function initialize!(core, pops)
    for pop in pops
        for sexpop in pop
            nl = get(sexpop, :number_of_loci)
            sex = get(sexpop, :sex)
            popsize = length(sexpop)
            for j = 1:popsize, i in 1:nl
                chr = get(sexpop, :type_of_locus, i)
                setinitialgenotype!(core, sexpop, j, i, Val{sex}, Val{chr})
            end
        end
    end
end

function reinitialize!(oldcore, pops)
    oldgdb = db(oldcore)
    core = BasicData()
    gdb = db(core)
    # Preemptively map 0 to 0 as this value has a special meaning.
    smap = Dict{Int, Int}(0 => 0)
    smax = 1
    for pop in pops
        for sexpop in pop
            for idx in eachindex(sexpop.data)
                sexpop[idx] == 0 && continue
                state = oldgdb[sexpop[idx]].state
                if !haskey(smap, state)
                    smap[state] = smax
                    smax += 1
                end
                sexpop[idx] = insert!(gdb, GeneRecord(0, smap[state]))
            end
        end
    end
    core.state = maximum(keys(db(core)))
    core
end

function listgenes!(gids::AbstractArray, pops, locus, idx)
    for pop in pops
        for sexpop in pop
            chrtype = Val{get(sexpop, :type_of_locus, locus)}
            sex = get(sexpop, :sex)
            idx = listgenes!(gids, sexpop, locus, idx, chrtype, Val{sex})
        end
    end
    idx
end

listgenes!(gids::AbstractArray, pops, locus) = listgenes!(gids, pops, locus, 1)

listgenes!(gids::AbstractArray, sexpop, locus, idx, chrtype, sex) = idx

function listgenes!(gids::AbstractArray, sexpop, locus, idx, ::Type{Val{Autosome}}, sex)

    psize = length(sexpop)
    offset = get(sexpop, :number_of_loci)
    for i = 1:psize, chr = 1:2
        gids[idx] = sexpop[chr, (i - 1) * offset + locus]
        idx += 1
    end
    idx
end

function listgenes!(gids::AbstractArray, sexpop, locus, idx,
    ::Type{Val{XChromosome}}, ::Type{Val{Female}})

    psize = length(sexpop)
    offset = get(sexpop, :number_of_loci)
    for i = 1:psize, chr = 1:2
        gids[idx] = sexpop[chr, (i - 1) * offset + locus]
        idx += 1
    end
    idx
end

function listgenes!(gids::AbstractArray, sexpop, locus, idx,
    ::Type{Val{XChromosome}}, ::Type{Val{Male}})

    psize = length(sexpop)
    offset = get(sexpop, :number_of_loci)
    for i = 1:psize
        gids[idx] = sexpop[1, (i - 1) * offset + locus]
        idx += 1
    end
    idx
end

function listgenes!(gids::AbstractArray, sexpop, locus, idx,
    ::Type{Val{YChromosome}}, ::Type{Val{Male}})

    psize = length(sexpop)
    offset = get(sexpop, :number_of_loci)
    for i = 1:psize
        gids[idx] = sexpop[2, (i - 1) * offset + locus]
        idx += 1
    end
    idx
end

function listgenes!(gids::AbstractArray, sexpop, locus, idx,
    ::Type{Val{Mitochondrion}}, ::Type{Val{Female}})

    psize = length(sexpop)
    offset = get(sexpop, :number_of_loci)
    for i = 1:psize
        gids[idx] = sexpop[1, (i - 1) * offset + locus]
        idx += 1
    end
    idx
end

function reproduce!(core, parpop, ppos, chpop, cpos)
    # Maternally derived genes are always stored on the first row, whereas
    # paternally derived genes are on the second row.
    # This function deals with a single offspring, not the entire population of offspring.
    csex = get(chpop, :sex)
    cdeme = get(chpop, :deme)
    pdeme = get(parpop[1], :deme) # both parents are from the same deme.
    for (sexpop, pos) in zip(parpop, ppos)
        psex = get(sexpop, :sex)
        nl = get(sexpop, :number_of_loci)
        cchr = value(psex)
        chr = rand(1:2)
        mut = get(sexpop, :mutation_rate)
        for i = 1:nl
            chrtype = get(sexpop, :type_of_locus, i)
            rrate = get(sexpop, :recombination_rate, i)

            chr, passgene = selectchr(chr, rrate, Val{chrtype}, Val{psex}, Val{csex})
            passgene || continue

            pos = toposition(chpop, cpos, i)
            if rand() < mut
                state = nextstate!(core)
            else
                state = -1
            end
            pgene = sexpop[chr, pos]
            geneid = transmit!(db(core), time(core), pgene, state=state, src=pdeme, dest=cdeme)
            chpop[cchr, pos] = geneid
        end
    end
end

selectchr(chr, rate, ::Type{Val{Autosome}}, psex, csex) = pickchr(chr, rate), true
selectchr(chr, rate, ::Type{Val{XChromosome}}, ::Type{Val{Female}}, csex) =
    pickchr(chr, rate), true
selectchr(chr, rate, ::Type{Val{XChromosome}}, ::Type{Val{Male}}, ::Type{Val{Female}}) = 1, true
selectchr(chr, rate, ::Type{Val{YChromosome}}, ::Type{Val{Male}}, ::Type{Val{Male}}) = 2, true
selectchr(chr, rate, ::Type{Val{Mitochondrion}}, ::Type{Val{Female}}, ::Type{Val{Female}}) =
    1, true
selectchr(chr, rate, chrtype, psex, csex) = chr, false

pickchr(chr, rate) = rand() < rate ? 3 - chr : chr

toposition(pop, ind, locus) = get(pop, :number_of_loci) * (ind - 1) + locus

function evolve!(
    core::BasicData,
    parpops::Array{Population, 1},
    termon::Int,
    tclean::Int)

    info("evolution started on ", now(), ".")

    ndemes = length(parpops) # basically fixed at 2
    nloci = get(parpops[1][1], :number_of_loci)
    loci = get(parpops[1][1], :type_of_locus)

    chpops = deepcopy(parpops)

    ncoals = [0 for _ = 1:nloci]
    lastcoals = [0 for _ = 1:nloci]

    # storage of gene ids. Autosome
    gidstore = Array{Array{Int, 1}, 1}(4)
    # Autosome
    gidstore[1] = zeros(
        Int,
        2 * sum((vcat([[length(pop) for pop in parpop] for parpop in parpops]...)))
    )
    # XChromosome
    gidstore[2] = zeros(
        Int,
        2 * sum([length(parpop[1]) for parpop in parpops]) +
        sum([length(parpop[2]) for parpop in parpops])
    )
    # YChromosome
    gidstore[3] = zeros(
        Int,
        sum([length(parpop[2]) for parpop in parpops])
    )
    # Mitochondrion
    gidstore[4] = zeros(
        Int,
        sum([length(parpop[1]) for parpop in parpops])
    )

    nnewids = sum(map(x -> length(gidstore[value(x)]), loci))

    gen = 0

    # main loop of a simulation
    for gen = core
        for pop in chpops
            for sexpop in pop
                # cdeme and csex indicate the location and sex of offspring (child).
                cdeme = get(sexpop, :deme)
                csex = get(sexpop, :sex)
                # i = toposition(sexpop)
                propmig = get(sexpop, :proportion_of_migrant)
                migfit = get(sexpop, :fitness_of_migrant)
                for org = 1:length(sexpop)
                    while true
                        # decide if an offspring is a native or migrant.
                        pdeme = cdeme
                        if rand() < propmig
                            rand() < migfit || continue
                            pdeme = 3 - cdeme
                        end

                        # If cdeme == pdeme, an offspring is born in pdeme and stay there
                        # for reproduction.
                        # If cdeme != pdeme, an offspring is born in pdeme, and it migrates
                        # to cdeme before reproduction.
                        # Because all offspring are born in pdeme, it is OK to have a parent
                        # to have offspring in both demes.

                        pp = parpops[pdeme]

                        # sample parents
                        mom = rand(1:length(pp[1]))
                        pop = rand(1:length(pp[2]))

                        # sample the trait of a offspring
                        momtrait = pp[1][Val{:trait}, mom]
                        poptrait = pp[2][Val{:trait}, pop]
                        ctrait = trait(
                            get(pp[1], :transmission_rate_of_trait, csex),
                            get(pp[2], :transmission_rate_of_trait, csex),
                            momtrait, poptrait, pdeme)

                        if pdeme != cdeme && ctrait != cdeme
                            tfit = get(sexpop, :fitness_of_foreign_trait)
                            rand() < tfit || continue
                        end

                        reproduce!(core, pp, (mom, pop), sexpop, org)

                        sexpop[Val{:trait}, org] = ctrait
                        break
                    end
                end
            end
        end

        parpops, chpops = chpops, parpops

        for (idx, locus) in enumerate(loci)
            lpos = value(locus)
            listgenes!(gidstore[lpos], parpops, idx)
            anc = mrca(db(core), gidstore[lpos])
            if anc.epoch > lastcoals[idx]
                ncoals[idx] += 1
                lastcoals[idx] = gen
                info("turn over ", ncoals[lpos], " of locus ", idx, " on generation ", gen)
            end
        end
        if termon == minimum(ncoals)
            break
        end
        gen % tclean == 0 && clean!(db(core), nnewids * gen + 1, nnewids * (gen + 1))
    end
    clean!(db(core), nnewids * gen + 1, nnewids * (gen + 1))

    if termon == minimum(ncoals)
        info("evolution terminated by number of turn-overs in gen ", gen, " on ", now(), ".")
    else
        info("evolution terminated by reaching max gen ", gen, " on ", now(), ".")
    end
    parpops, core

end

function simulate(params::ModelParameters, burnin::Int, t::Int, termon::Int, tclean::Int, rep=1)
    info("process started on ", now(), ".")
    core = BasicData()
    pops = createpopulations(params)

    initialize!(core, pops)
    settmax!(core, burnin)
    pops, core  = evolve!(core, pops, -1, tclean)

    datastore = Array{Tuple{typeof(pops), typeof(db(core)), typeof(time(core))}}(0)

    for i = 1:rep
        pops = deepcopy(pops) # clone population that to avoid overwrite.
        core = reinitialize!(core, pops)
        settmax!(core, t)
        pops, core = evolve!(core, pops, termon, tclean)
        push!(datastore, (pops, db(core), time(core)))
    end
    info("process terminated on ", now(), ".")
    datastore
end

end # module


using Potential, Preprocessing, Optimization

function VelocityVerlet(x, v, q, t, a, b, c, pbc, app, potentials)

    dim = Int32(size(x,1))
    inum = Int32(length(t))

    ensemblemode = app.ensemblemode;
    time = app.time;             # initial time
    dt = app.dt
    ntimesteps =  app.ntimesteps;       # # time steps
    globalfreq =  app.globalfreq;
    skin = app.neighskin;
    every = app.neighevery;      # perform neighbor rebuild list for "every" iterations
    delay = app.neighdelay;      # delay neighbor rebuild
    distcheck = app.neighcheck;  # 0 -> neighbor rebuild by delay and every settings,
    outputflag = app.outputflag
    eta = app.eta[:]
    kappa = app.kappa[:]
    atommass = [0; app.atommasses[:]]

    if (lowercase(ensemblemode) == "nve")
        ensemblemodenum = 0;
    elseif (lowercase(ensemblemode) == "nvelimit")
        ensemblemodenum = 1;
    elseif (lowercase(ensemblemode) == "nvt")
        ensemblemodenum = 2;
    else
        ensemblemodenum = -1;
    end
    ensemblemodenum = Int32(ensemblemodenum)

    units = Preprocessing.setunits(app.unitstyle)
    boltz = Int32(units.boltz)
    mvv2e = Int32(units.mvv2e)
    nktv2p = Int32(units.nktv2p)

    dtarray, tarray, vlimitsq,  eta_mass, eta_thermostat, eta_dot, eta_dotdot,
        eta_mass_flag, biasflag, mtchain, nc_tchain = ensemble(app, dim, inum)

    dom = setdomainstruct(pbc, a, b, c)
    boxlo = dom.boxlo
    boxhi = dom.boxhi
    hi_lambda = dom.boxhi_lamda
    lo_lambda = dom.boxlo_lamda
    h = dom.h
    h_inv = dom.h_inv
    h_rate = dom.h_rate
    vdeform = dom.vdeform
    triclinic = dom.triclinic
    inv_volume = 1.0 / (dom.h[1] * dom.h[2] * dom.h[3])

    rcutmax = skin + Potential.maximumcutoff(potentials)
    x, alist, neighlist, neighnum = Potential.fullneighborlist(x, a, b, c, pbc, rcutmax);
    etot, eatom, f, vatom, vtot = Potential.emlpotential(x, q, t, alist, neighlist, neighnum, eta, kappa, potentials)
    pe, temp, pres, ke = thermooutput(v, t, alist, atommass, vtot, etot, mvv2e, boltz, nktv2p, inv_volume, outputflag)
    print("Total number of atoms = $inum\n");
    print("Time step     Temperature        Potential energy        Total energy            Pressure \n");
    printoutput(0, ntimesteps, temp, pe, ke, pres)

    xhold = 1.0*x
    for istep=1:ntimesteps
        time = time + dt
        gnum = Int32(length(alist))
        dtarray[6] = istep
        xstart = 1.0*x

        InitialIntegrate(x, v, f, atommass, dtarray, tarray, eta_mass, eta_thermostat,
            eta_dot, eta_dotdot, vlimitsq, Int32.(t[:]), Int32.(alist.-1), eta_mass_flag, biasflag, mtchain,
            nc_tchain, ensemblemodenum, dim, inum);

        PostInitialIntegration(x, v, f, t, eatom, vatom, atommass, dtarray, app.setvelocity, app.wallreflect)

        build = checkbuild(x, xhold, skin, delay, every, distcheck, istep, dim, inum)
        if build==1
            print("Rebuild neighbor list at timestep $istep...\n");
            if sum(pbc) > 0
                ApplyPBC(x, v, image, boxhi, boxlo, hi_lambda, lo_lambda, h, h_inv,
                    h_rate, pbc, vdeform, triclinic, dim, inum)
            end
            x, alist, neighlist, neighnum = Potential.fullneighborlist(x[:,1:inum], a, b, c, pbc, rcutmax);
            xhold = 1.0*x
        else
            if sum(pbc) > 0
                # calculate the distance between x and xstart: xstart = x - xstart
                xstart = x - xstart
                # shift periodic images by the above distance
                xtm = x[:,(inum+1):end]
                atm = alist[(inum+1):end]
                ArrayPlusAtColumnIndex(xtm, xstart, Int32.(atm.-1), dim, gnum);
                x[:,(inum+1):end] = xtm
            end
        end

        etot, eatom, f, vatom, vtot = Potential.emlpotential(x, q, t, alist, neighlist, neighnum, eta, kappa, potentials)

        PostForceComputation(x, v, f, eatom, vatom, app.setforce, app.lineforce, app.planeforce,
                    app.wallharmonic, app.walllj126, app.wallmorse)

        FinalIntegrate(x, v, f, atommass, dtarray, tarray, eta_mass, eta_thermostat, eta_dot,
            eta_dotdot, vlimitsq, Int32.(t[:]), Int32.(alist.-1), eta_mass_flag, biasflag,
            mtchain, nc_tchain,  ensemblemodenum, dim, inum);

        PostFinalIntegration(x, v, f, t, atommass, dtarray, tarray,
                biasflag, app.setvelocity, app.velocityrescaling)

        if (istep % globalfreq == 0)
            pe, temp, pres, ke = thermooutput(v, t, alist, atommass, vtot, etot,
                mvv2e, boltz, nktv2p, inv_volume, outputflag)
            printoutput(istep, ntimesteps, temp, pe, ke, pres)
        end
    end
end

function VelocityVerlet(config, app, potentials, ci=1)

    natom = [0; cumsum(config.natom[:])];

    x = config.x[:,(natom[ci]+1):natom[ci+1]] # atom positions of configuration ci
    v = config.v[:,(natom[ci]+1):natom[ci+1]] # atom positions of configuration ci
    t = config.t[:,(natom[ci]+1):natom[ci+1]] # atom types of configuration ci
    if length(config.q)>0
        q = config.q[:,(natom[ci]+1):natom[ci+1]]
    else
        q = []
    end
    if size(config.a,2) > 1
        a = config.a[:,ci]
        b = config.b[:,ci]
        c = config.c[:,ci]
    else
        a = config.a[:,1]
        b = config.b[:,1]
        c = config.c[:,1]
    end
    if size(config.pbc,2) > 1
        pbc = config.pbc[:,ci]
    else
        pbc = config.pbc[:,1]
    end

    VelocityVerlet(x, v, q, t, a, b, c, pbc, app, potentials)
end

function printoutput(istep, ntimesteps, temp, pe, ke, pres)
    n = length(digits(ntimesteps))
    s0 = repeat(' ', n)
    s1 = replacestring(s0, string(istep)) * "        "
    s2 = num2string(temp, 16) * "      "
    s3 = num2string(pe, 16) * "      "
    s4 = num2string(pe+ke, 16) * "      "
    s5 = num2string(pres, 16) * "      "
    print(s1 * s2 * s3 * s4 * s5 * "\n");
end

function ensemble(app, dim, inum)

    ensemblemode = app.ensemblemode
    dtarray = zeros(10)
    tarray = zeros(10)
    eta_mass = zeros(10)
    eta_thermostat = zeros(10)
    eta_dot = zeros(10)
    eta_dotdot = zeros(10)

    units = Preprocessing.setunits(app.unitstyle)
    boltz = units.boltz
    mvv2e = units.mvv2e
    ftm2v = units.ftm2v
    tdof = (inum-1)*dim;

    dt = app.dt
    dtarray[1] = dt;     # dt
    dtarray[2] = 0.5*dt*ftm2v; # dtf
    dtarray[3] = dt;     # dtv
    dtarray[4] = 0;             # beginstep
    dtarray[5] = app.ntimesteps; # endstep
    vlimitsq = 0.0;
    eta_mass_flag = 1;
    biasflag = 0;
    mtchain = 1;
    nc_tchain = 1;
    if ensemblemode != "nve"
        if (ensemblemode == "nvelimit")
            nveparam = app.nveparam
            vlimitsq = (nveparam[1]/dt) * (nveparam[1]/dt);
        end
        if (ensemblemode == "nvt")  # NVT ensemble
            nvtparam = app.nvtparam
            tarray[1] = nvtparam[0]; # temp start
            tarray[2] = nvtparam[1]; # temp stop
            tarray[3] = 1/nvtparam[2]; # temp frequency
            #delta = (ntimestep - beginstep)/(endstep - beginstep)
            #tarray[4] = t_start + delta * (t_stop - t_start); # t_target
            #t_current = cpuComputeTempScalar(v, mass, tfactor, type, ilist, dim, inum);
            #tarray[5] = t_current;
            tarray[6] = tdof;
            tarray[7] = boltz;
            tarray[8] = (length(nvtparam)>3) ? nvtparam[3] : 0.0; # drag factor added to barostat/thermostat
            tarray[9] = mvv2e;
            mtchain   = (length(nvtparam)>4) ? Int32(nvtparam[4]) : 3;
            nc_tchain = (length(nvtparam)>5) ? Int32(nvtparam[5]) : 1;
        end
    end

    eta_mass_flag = Int32(eta_mass_flag)
    biasflag = Int32(biasflag)
    mtchain = Int32(mtchain)
    nc_tchain = Int32(nc_tchain)
    return dtarray, tarray, vlimitsq,  eta_mass, eta_thermostat, eta_dot, eta_dotdot, eta_mass_flag, biasflag, mtchain, nc_tchain
end

function thermooutput(v, t, ilist, atommass, vtot, etot, mvv2e, boltz, nktv2p, inv_volume, flag=0)

    dim, inum = size(v)
    tdof = (inum-1)*dim;
    tfactor = mvv2e/(tdof * boltz)

    # function ComputeKEAtom(ke::Vector{Float64}, mass::Vector{Float64}, v::Vector{Float64},
    #     mvv2e::Float64, type::Vector{Int32}, ilist::Vector{Int32}, dim::Int32, inum::Int32)

    keatom = zeros(inum)
    ComputeKEAtom(keatom, atommass, v[:], 2.0, Int32.(t[:]), Int32.(ilist .- 1), Int32(dim), Int32(inum));
    kesum = sum(keatom)

    pe = etot
    temp = tfactor*kesum
    ke = 0.5 * tdof * boltz * temp
    pres = (tdof * boltz * temp + vtot[1] + vtot[2] + vtot[3]) / 3.0 * inv_volume * nktv2p

    if (flag == 0)
        pe = pe/inum
        ke = ke/inum
    end

    return pe, temp, pres, ke
end

function PostInitialIntegration(x, v, f, t, eatom, vatom, mass,
            dtarray, setvelocity=nothing, wallreflect=nothing)

    dim, anum = size(v)
    alist = Int32.(0:(anum-1))
    eflag_atom = 0;
    vflag_atom = 0;
    dtf = dtarray[2]
    dtv = dtarray[3]

    x = x[:]
    v = v[:]
    f = f[:]
    if (setvelocity !== nothing)
        nsetvelocity = length(setvelocity)
        for  i = 1:nsetvelocity
            fparam = setvelocity[i].fparam
            iparam = setvelocity[i].iparam
            ind =  findall(t == setvelocity[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            SetVelocityInitialIntegrate(x, v, f, mass, fparam, dtf,
                    dtv, t, ilist, iparam, dim, inum);
        end
    end

    if (wallreflect !== nothing)
        nwallreflect = length(wallreflect)
        for i = 1:nwallreflect
            fparam = wallreflect[i].fparam
            iparam = wallreflect[i].iparam
            ind =  findall(t == wallreflect[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            FixWallReflect(x, v, f, eatom, vatom, fparam,
                iparam, ilist, eflag_atom, vflag_atom, dim, inum);
        end
    end
end

function checkbuild(x, xhold, skin, delay, every, distcheck, istep, dim, inum)
    build = 0;
    ago = istep + 1;

    if ((ago >= delay) && (ago % every == 0))
        if (distcheck == 0)
            build = 1;
        else
            ds = zeros(inum)
            ArrayDistSquareSum(ds, x[:], xhold[:], dim, inum);
            maxdistsquare = maximum(ds)
            if (maxdistsquare > 0.25*skin*skin)
                build = 1;
            end
        end
    end

    return build
end

function PostForceComputation(x, v, f, t, eatom, vatom, setforce=nothing, lineforce=nothing,
    planeforce=nothing, wallharmonic=nothing, walllj126=nothing, wallmorse=nothing)

    dim, anum = size(v)
    alist = Int32.(0:(anum-1))

    eflag_atom = 0;
    vflag_atom = 0;

    if (setforce !== nothing)
        nsetforce = length(setforce)
        for i = 1:nsetforce
            fparam = setforce[i].fparam
            iparam = setforce[i].iparam
            ind =  findall(t == setforce[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            FixSetForce(x, v, f, eatom, vatom, fparam,
                iparam, ilist, eflag_atom, vflag_atom, dim, inum);
        end
    end

    if (lineforce !== nothing)
        nlineforce = length(lineforce)
        for i = 1:nlineforce
            fparam = lineforce[i].fparam
            iparam = lineforce[i].iparam
            ind =  findall(t == lineforce[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            FixLineForce(x, v, f, eatom, vatom, fparam,
                iparam, ilist, eflag_atom, vflag_atom, dim, inum);
        end
    end

    if (planeforce !== nothing)
        nplaneforce = length(planeforce)
        for i = 1:nplaneforce
            fparam = planeforce[i].fparam
            iparam = planeforce[i].iparam
            ind =  findall(t == planeforce[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            FixPlaneForce(x, v, f, eatom, vatom, fparam,
                iparam, ilist, eflag_atom, vflag_atom, dim, inum);
        end
    end

    if (wallharmonic !== nothing)
        nwallharmonic = length(wallharmonic)
        for i = 1:nwallharmonic
            fparam = wallharmonic[i].fparam
            iparam = wallharmonic[i].iparam
            ind =  findall(t == wallharmonic[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            FixWallHarmonic(x, v, f, eatom, vatom, fparam,
                iparam, ilist, eflag_atom, vflag_atom, dim, inum);
        end
    end

    if (walllj126 !== nothing)
        nwalllj126 = length(walllj126)
        for i = 1:nwalllj126
            fparam = walllj126[i].fparam
            iparam = walllj126[i].iparam
            ind =  findall(t == walllj126[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            FixWallLJ126(x, v, f, eatom, vatom, fparam,
                iparam, ilist, eflag_atom, vflag_atom, dim, inum);
        end
    end

    if (wallmorse !== nothing)
        nwallmorse = length(wallmorse)
        for i = 1:nwallmorse
            fparam = wallmorse[i].fparam
            iparam = wallmorse[i].iparam
            ind =  findall(t == wallmorse[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            FixWallMorse(x, v, f, eatom, vatom, fparam,
                iparam, ilist, eflag_atom, vflag_atom, dim, inum);
        end
    end
end

function PostFinalIntegration(x, v, f, t, mass, dtarray, tarray,
        biasflag, setvelocity=nothing, velocityrescaling=nothing)

    dim, anum = size(v)
    alist = Int32.(0:(anum-1))
    energy = 0.0
    dtf = dtarray[2]
    #dtv = dtarray[3]

    if (setvelocity !== nothing)
        nsetvelocity = length(setvelocity)
        for  i = 1:nsetvelocity
            #fparam = setvelocity[i].fparam
            iparam = setvelocity[i].iparam
            ind =  findall(t == setvelocity[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            SetVelocityFinalIntegrate(x, v, f, mass, dtf, t, ilist, iparam, dim, inum);
        end
    end

    if (velocityrescaling !== nothing)
        nvelocityrescaling = length(velocityrescaling)
        second = zeros(1)
        seed = Int32.(ones(1))
        save = Int32.(zeros(1))
        for  i = 1:nvelocityrescaling
            fparam = velocityrescaling[i].fparam
            iparam = velocityrescaling[i].iparam
            energy = fparam[1];
            mode = iparam[1];
            ind =  findall(t == velocityrescaling[i].groupnum)
            ilist = alist[ind]
            inum = length(ilist)
            energy = VelocityRescalingThermostat(v, mass, dtarray, tarray, second, energy,
                    t, ilist, seed, save, biasflag, mode, dim, inum);
        end
    end

    return energy
end


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Attributes #
# ========================== #
"""
    AbstractLShapedAttribute

Abstract supertype for attribute objects specific to the L-shaped algorithm.
"""
abstract type AbstractLShapedAttribute <: AbstractStructuredOptimizerAttribute end
"""
    FeasibilityStrategy

An optimizer attribute for specifying a strategy for dealing with second-stage feasibility the L-shaped algorithm. Options are:

- [`IgnoreFeasibility`](@ref) (default)
- [`FeasibilityCuts`](@ref)
"""
struct FeasibilityStrategy <: AbstractLShapedAttribute end
"""
    IntegerStrategy

An optimizer attribute for specifying a strategy for dealing with integers the L-shaped algorithm. Options are:

- [`IgnoreIntegers`](@ref) (default)
- [`CombinatorialCuts`](@ref)
- [`Convexification`](@ref)
"""
struct IntegerStrategy <: AbstractLShapedAttribute end
"""
    Regularizer

An optimizer attribute for specifying a regularization procedure to be used in the L-shaped algorithm. Options are:

- [`NoRegularization`](@ref):  L-shaped algorithm (default)
- [`RegularizedDecomposition`](@ref):  Regularized decomposition ?RegularizedDecomposition for parameter descriptions.
- [`TrustRegion`](@ref):  Trust-region ?TrustRegion for parameter descriptions.
- [`LevelSet`](@ref):  Level-set ?LevelSet for parameter descriptions.
"""
struct Regularizer <: AbstractLShapedAttribute end
"""
    Aggregator

An optimizer attribute for specifying an aggregation procedure to be used in the L-shaped algorithm. Options are:

- [`NoAggregation`](@ref):  Multi-cut L-shaped algorithm (default)
- [`PartialAggregation`](@ref):  ?PartialAggregation for parameter descriptions.
- [`FullAggregation`](@ref):  ?FullAggregation for parameter descriptions.
- [`DynamicAggregation`](@ref):  ?DynamicAggregation for parameter descriptions.
- [`ClusterAggregation`](@ref):  ?ClusterAggregation for parameter descriptions.
- [`HybridAggregation`](@ref):  ?HybridAggregation for parameter descriptions.
"""
struct Aggregator <: AbstractLShapedAttribute end
"""
    Consolidator

An optimizer attribute for specifying a consolidation procedure to be used in the L-shaped algorithm. Options are:

- [`NoConsolidation`](@ref) (default)
- [`Consolidation`](@ref)
"""
struct Consolidator <: AbstractLShapedAttribute end
"""
    IntegerParameter

Abstract supertype for integer-specific attributes.
"""
abstract type IntegerParameter <: AbstractLShapedAttribute end
"""
    RegularizationParameter

Abstract supertype for regularization-specific attributes.
"""
abstract type RegularizationParameter <: AbstractLShapedAttribute end
"""
    AggregationParameter

Abstract supertype for aggregation-specific attributes.
"""
abstract type AggregationParameter <: AbstractLShapedAttribute end
"""
    ConsolidationParameter

Abstract supertype for consolidation-specific attributes.
"""
abstract type ConsolidationParameter <: AbstractLShapedAttribute end


@testset "Test NLP Evals" begin

    x0 = zeros(2)

    nlp = ADNLPModel(x -> norm(x.^2), x0)
    max_nlp = init_max_counters(obj = 2)
    nlp_stop_evals = NLPStopping(nlp, max_cntrs = max_nlp)

    slack = SlackModel(nlp)
    max_slack = init_max_counters(obj = 2)
    slack_stop_evals = NLPStopping(nlp, max_cntrs = max_slack)

    nls = ADNLSModel(x -> x.^2, x0, 2)
    max_nls = init_max_counters_NLS(obj = 2, residual = 1)
    nls_stop_evals = NLPStopping(nls, max_cntrs = max_nls)

    @test typeof(nlp_stop_evals.pb.counters) == Counters
    @test nlp_stop_evals.meta.max_cntrs[:neval_obj] == 2
    @test typeof(nls_stop_evals.pb.counters) == NLSCounters
    @test nls_stop_evals.meta.max_cntrs[:neval_obj] == 2
    @test nls_stop_evals.meta.max_cntrs[:neval_residual] == 1

    Stopping._resources_check!(nlp_stop_evals, x0)
    @test nlp_stop_evals.meta.resources == false
    Stopping._resources_check!(slack_stop_evals, x0)
    @test slack_stop_evals.meta.resources == false
    Stopping._resources_check!(nls_stop_evals, x0)
    @test nls_stop_evals.meta.resources == false

    nlp_stop_evals.meta.max_cntrs[:neval_obj] = 0
    Stopping._resources_check!(nlp_stop_evals, x0)
    @test nlp_stop_evals.meta.resources == false
    nls_stop_evals.meta.max_cntrs[:neval_residual] = -1
    Stopping._resources_check!(nls_stop_evals, x0)
    @test nls_stop_evals.meta.resources == true

    max_evals!(nlp_stop_evals, 10)
    @test nlp_stop_evals.meta.max_cntrs[:neval_obj] == 10
    @test nlp_stop_evals.meta.max_cntrs[:neval_sum] == typemax(Int)

    max_evals!(nlp_stop_evals, allevals = 10, obj = 2)
    @test nlp_stop_evals.meta.max_cntrs[:neval_cons] == 10
    @test nlp_stop_evals.meta.max_cntrs[:neval_sum]  == typemax(Int)
    @test nlp_stop_evals.meta.max_cntrs[:neval_obj] == 2

end


using SeisIO, FileIO, Dates, DSP, Statistics, DataFrames, CSV, Trapz, Distributed, Plots, PlotSeis
@everywhere using SeisPGVPGA

#---output project directory (all outputs are dumped into here)---#
fodir = "output"

#---parameters for request stations---#
fi_stations = "BP_gmap-stations.txt" #"BP_gmap-stations.txt"
locchan = Dict(
            "BP" => [(*, "BP*"), (*, "SP*")]
            )

#---parameters for processsing---#
starttime = DateTime(2004, 10, 1)
endtime  = DateTime(2004, 10, 2)
download_time_unit = 86400.0 #[s]
sampling_frequency = 20.0 #[Hz]
get_data_opt = [true, true, true, true, true]#: [unscale, demean, detrend, taper, ungap]
download_margin = 300.0 # [s] download margin
save_seisdata = true # True if save seisdata, otherwise save only pgvpga output.
bp_fmin = 0.1;
bp_fmax = 5.0;
drms_fmin = 0.9;
drms_fmax = 2.0;

#---make output directories---#
InputDict=Dict()
mkpath(fodir)
InputDict["pgvpga_dir"] = joinpath(fodir, "pgvpga")
InputDict["seisdata_dir"] = joinpath(fodir, "seisdata")
mkpath(InputDict["pgvpga_dir"])
mkpath(InputDict["seisdata_dir"])

#---make stationlist---#
stationlist_fopath = joinpath(fodir, "request_stations")
InputDict["requeststations"] = get_requeststations(fi_stations, locchan, stationlist_fopath);

starttimelist = get_starttimelist(starttime, endtime, download_time_unit)
InputDict["starttimelist"] = starttimelist
InputDict["download_time_unit"] = download_time_unit
InputDict["fodir"] = fodir
InputDict["sampling_frequency"] = sampling_frequency
InputDict["get_data_opt"] = get_data_opt
InputDict["download_margin"] = download_margin
InputDict["save_seisdata"] = save_seisdata

if (drms_fmin < bp_fmin) || (bp_fmax<drms_fmax) || (drms_fmax < drms_fmin)
    error("The order of frequency band is incorrect. Please order like bp_fmin < drms_fmin< drms_fmax< bp_fmax.")
end

InputDict["bp_fmin"] = bp_fmin
InputDict["bp_fmax"] = bp_fmax
InputDict["drms_fmin"] = drms_fmin
InputDict["drms_fmax"] = drms_fmax

#----Restrict number of processors------#
#NEVER CHANGE THIS THRESHOLD OTHERWISE IT OVERLOADS THE DATA SERVER
np = nprocs()
if np > 64 throw(DomainError(np, "np must be smaller than 64.")) end
#---------------------------------------#

# Start downloading data
println("-------START DOWNLOADING AND PROCESSING PGVPGA--------")
t_download = 0
t_download = @elapsed pmap(x -> map_seisPGVPGA(x, InputDict), 1:length(starttimelist))

println("-------START CONVERTING OUTPUT INTO CSV--------")

files_all = String[]

for (root, dirs, files) in walkdir(InputDict["pgvpga_dir"])
   for file in files
       fi = joinpath(root, file)
        (split(fi, ".")[end] == "dat") && push!(files_all, fi)# filter if it is .seisio
   end
end
t_convert = 0
t_convert = @elapsed df_mapped = pmap(x -> map_readPGVPGA(x), files_all)

df_all = DataFrame(id=String[], starttime=DateTime[], endtime=DateTime[], midtime=DateTime[], PGV=Float64[],
                    PGA=Float64[], drms_mean=Float64[])

for df in df_mapped
    (!isnothing(df) && !isnan(df.PGV[1])) && append!(df_all, df)
end

csvname = joinpath(InputDict["fodir"], "pgv_pga_all.csv")
CSV.write(csvname, df_all)

printstyled("---Summary---\n"; color = :cyan, bold = true)
println("Total time for download and process data = $(t_download)[s]")
println("Total time for gather data = $(t_convert)[s]")
printstyled("SeisPGVPGA has been successfully done.\n"; color = :cyan, bold = true)


module GRASS

# import external modules
using CSV
using FITSIO
using Random
using StatsBase
using DataFrames
using Statistics
using SharedArrays
using Interpolations
import Glob.glob
import Dates.DateTime

# abbreviations for commonly used types
const AA = AbstractArray
const AF = AbstractFloat

# configure directories
include("config.jl")

# ancillary functions + constants
include("utils.jl")
include("constants.jl")

# user-defined composite types
include("structures.jl")

# star geometry + thermal/RT physics
include("stargeometry.jl")
include("starphysics.jl")

# data read-in + calculations
include("lineIO.jl")
include("bisectors.jl")

# star simulation
include("trim.jl")
include("synthesize.jl")
include("disk_sim.jl")

# processing spectra
include("velocities.jl")

# simulating observations
include("observing/convolutions.jl")
include("observing/signaltonoise.jl")
include("observing/ObservationPlan.jl")

# export some stuff
export SpecParams, DiskParams, synthesize_spectra, calc_ccf, calc_rvs_from_ccf, calc_rms

end # module


# include(joinpath(WATERFALL_DIR,"src","includes.jl"))
import Base
import Combinatorics
import CSV
import DataFrames
import LinearAlgebra, LinearAlgebra.I
import Printf
import SparseArrays
import Statistics
import StatsPlots

# DATA_DIR = joinpath(WATERFALL_DIR,"data","pvrd")
# df = CSV.read(joinpath(DATA_DIR,"pvrd2-investment-metrics.csv"), DataFrames.DataFrame)

DATA_DIR = joinpath(WATERFALL_DIR,"data")
df = CSV.read(joinpath(DATA_DIR,"viz","investment.csv"), DataFrames.DataFrame)
[df[!,ii] .= coalesce.(df[:,ii],missing,"") for ii in 1:size(df,2)]
# df = df[[1:11;15],:]

# df[!,:Process] .= replace.(df[:,:Process], "interconnection"=>"intercon")
# df[!,:Label] .= string.(titlecase.(df[:,:Process])," (",df[:,:Step],")")
# df[!,:Units] .= "Efficiency (%)"

# Define global arguments.
distribution=:normal
fuzziness=(0.01,0.3)
interactivity=(0.2,0.4)
sample=:Sample
value=:Value
label=:Step
sublabel=:Process
nsample=5

kwargs = (
    label=label,
    sublabel=sublabel,
    distribution=distribution,
    fuzziness=fuzziness,
    interactivity=interactivity,
    # vmin=12.0,
    # ncor=ncor,
    # nsample=nsample,
)


#functionality in this file implements writing .vtu files
#complient to https://vtk.org/wp-content/uploads/2015/04/file-formats.pdf
import Base64
#import CodecZlib
##
function vtk_write1(file_name::String,mesh::Mesh,data=[],binary=true,compressed=false)
   open(file_name*".vtu","w") do fid
      if compressed
         compressor=" compressor=\"vtkZLibDataCompressor\""
      else
         compressor=""
      end
      write(fid,"<VTKFile type=\"UnstructuredGrid\" version=\"0.1\" byte_order=\"LittleEndian\" header_type=\"UInt64\"$compressor>\n")
      write(fid,"\t<UnstructuredGrid>\n")
      write(fid,"\t\t<Piece NumberOfPoints=\"$(size(mesh.points,2))\" NumberOfCells=\"$(length(mesh.tetrahedra))\">\n")
      write(fid,"\t\t\t<Points>\n")
      #write(fid,"\t\t\t\t<DataArray type=\"Float64\" NumberOfComponents=\"3\" format=\"ascii\">")
      #for pnt in mesh.points
      #   write(fid,string(pnt)*" ")
      #end
      #write(fid,"</DataArray>")
      write_data_array(fid,Dict("Coordinates"=>mesh.points),binary,compressed)
      write(fid,"\t\t\t</Points>\n")
      write(fid,"\t\t\t<Cells>\n")
      # write(fid,"\t\t\t\t<DataArray type=\"Int32\" Name=\"connectivity\" format=\"ascii\">")
      # for tet in mesh.tetrahedra
      #    for pnt in tet
      #       write(fid,string(Int32(pnt-1))*" ")
      #    end
      # end
      # write(fid,"</DataArray>\n")
      N_points=size(mesh.points,2)
      #connectivioty must be Int32. Paraview is serious about this in its latest
      #versions! https://discourse.paraview.org/t/cannot-open-vtu-files-with-paraview-5-8/3759
      connectivity=Array{Int32}(undef,length(mesh.tetrahedra)*4)
      idx=1
      for tet in mesh.tetrahedra
         for pnt_idx in tet
            connectivity[idx]=Int32(pnt_idx-1)
            idx+=1
         end
      end

      write_data_array(fid,Dict("connectivity"=>connectivity),binary,compressed)

      write(fid,"\t\t\t\t<DataArray type=\"Int32\" Name=\"offsets\" format=\"ascii\">")
      off=0
      for tet in mesh.tetrahedra
            off+=length(tet)
            write(fid,string(off)*" ")
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t\t<DataArray type=\"UInt8\" Name=\"types\" format=\"ascii\">")
      for tet in mesh.tetrahedra
         for pnt in tet
            write(fid,"10 ")
         end
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t</Cells>\n")

      write(fid,"\t\t\t<PointData>\n")
      write_data_array(fid,data,binary,compressed)
      write(fid,"\t\t\t</PointData>\n")
      write(fid,"\t\t</Piece>\n")
      write(fid,"\t</UnstructuredGrid>\n")
      write(fid,"</VTKFile>\n")
   end
   return nothing
end
##
function vtk_write0(file_name::String,mesh::Mesh,data=[],binary=true,compressed=false)
   open(file_name*".vtu","w") do fid
      if compressed
         compressor=" compressor=\"vtkZLibDataCompressor\""
      else
         compressor=""
      end
      write(fid,"<VTKFile type=\"UnstructuredGrid\" version=\"0.1\" byte_order=\"LittleEndian\" header_type=\"UInt64\"$compressor>\n")
      write(fid,"\t<UnstructuredGrid>\n")
      write(fid,"\t\t<Piece NumberOfPoints=\"$(size(mesh.points,2))\" NumberOfCells=\"$(length(mesh.tetrahedra))\">\n")
      write(fid,"\t\t\t<Points>\n")
      # write(fid,"\t\t\t\t<DataArray type=\"Float64\" NumberOfComponents=\"3\" format=\"ascii\">")
      # for pnt in mesh.points
      #    write(fid,string(pnt)*" ")
      # end
      # write(fid,"</DataArray>")
      write_data_array(fid,Dict("Coordinates"=>mesh.points),binary,compressed)
      write(fid,"\t\t\t</Points>\n")
      write(fid,"\t\t\t<Cells>\n")
      write(fid,"\t\t\t\t<DataArray type=\"Int32\" Name=\"connectivity\" format=\"ascii\">")
      for tet in mesh.tetrahedra
         for pnt in tet
            write(fid,string(Int32(pnt-1))*" ")
         end
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t\t<DataArray type=\"Int32\" Name=\"offsets\" format=\"ascii\">")
      off=0
      for tet in mesh.tetrahedra
            off+=length(tet)
            write(fid,string(off)*" ")
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t\t<DataArray type=\"UInt8\" Name=\"types\" format=\"ascii\">")
      for tet in mesh.tetrahedra
         for pnt in tet
            write(fid,"10 ")
         end
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t</Cells>\n")

      write(fid,"\t\t\t<CellData>\n")
      write_data_array(fid,data,binary,compressed)
      write(fid,"\t\t\t</CellData>\n")
      write(fid,"\t\t</Piece>\n")
      write(fid,"\t</UnstructuredGrid>\n")
      write(fid,"</VTKFile>\n")
   end
   return nothing
end

function vtk_write2(file_name::String,mesh::Mesh,data=[],binary=true,compressed=false)
   open(file_name*".vtu","w") do fid
      if compressed
         compressor=" compressor=\"vtkZLibDataCompressor\""
      else
         compressor=""
      end
      write(fid,"<VTKFile type=\"UnstructuredGrid\" version=\"0.1\" byte_order=\"LittleEndian\" header_type=\"UInt64\"$compressor>\n")
      write(fid,"\t<UnstructuredGrid>\n")
      write(fid,"\t\t<Piece NumberOfPoints=\"$(size(mesh.points,2)+length(mesh.lines))\" NumberOfCells=\"$(length(mesh.tetrahedra))\">\n")
      write(fid,"\t\t\t<Points>\n")
      #write(fid,"\t\t\t\t<DataArray type=\"Float64\" NumberOfComponents=\"3\" format=\"ascii\">")
      # for pnt in mesh.points
      #    write(fid,string(pnt)*" ")
      # end
      # for line in mesh.lines
      #    pnts=sum(mesh.points[:,line],dims=2)/2
      #    for pnt in pnts
      #       write(fid,string(pnt)*" ")
      #    end
      # end
      #write(fid,"</DataArray>")
      N_points=size(mesh.points,2)
      points=Array{Float64}(undef,3,N_points+length(mesh.lines))
      points[:,1:N_points]=mesh.points
      for (idx,line) in enumerate(mesh.lines)
         points[:,N_points+idx]=sum(mesh.points[:,line],dims=2)/2
      end
      write_data_array(fid,Dict("coordinates"=>points),binary,compressed)
      write(fid,"\t\t\t</Points>\n")
      write(fid,"\t\t\t<Cells>\n")
      write(fid,"\t\t\t\t<DataArray type=\"Int32\" Name=\"connectivity\" format=\"ascii\">")
      Npoints=size(mesh.points,2)
      for tet in mesh.tetrahedra
         for pnt in tet
            write(fid,string(Int32(pnt-1))*" ")
         end
         pnt=get_line_idx(mesh,[tet[1]; tet[2]])+Npoints
         write(fid,string(Int32(pnt-1))*" ")
         pnt=get_line_idx(mesh,[tet[2]; tet[3]])+Npoints
         write(fid,string(Int32(pnt-1))*" ")
         pnt=get_line_idx(mesh,[tet[3]; tet[1]])+Npoints
         write(fid,string(Int32(pnt-1))*" ")
         pnt=get_line_idx(mesh,[tet[1]; tet[4]])+Npoints
         write(fid,string(Int32(pnt-1))*" ")
         pnt=get_line_idx(mesh,[tet[2]; tet[4]])+Npoints
         write(fid,string(Int32(pnt-1))*" ")
         pnt=get_line_idx(mesh,[tet[3]; tet[4]])+Npoints
         write(fid,string(Int32(pnt-1))*" ")
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t\t<DataArray type=\"Int32\" Name=\"offsets\" format=\"ascii\">")
      off=0
      for tet in mesh.tetrahedra
            off+=10#length(tet)
            write(fid,string(off)*" ")
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t\t<DataArray type=\"UInt8\" Name=\"types\" format=\"ascii\">")
      for tet in mesh.tetrahedra
         #for pnt in tet
            write(fid,"24 ")
         #end
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t</Cells>\n")

      write(fid,"\t\t\t<PointData>\n")
      write_data_array(fid,data,binary,compressed)
      write(fid,"\t\t\t</PointData>\n")
      write(fid,"\t\t</Piece>\n")
      write(fid,"\t</UnstructuredGrid>\n")
      write(fid,"</VTKFile>\n")
   end
   return nothing
end
##
function vtk_write_tri(file_name::String,mesh::Mesh,data=[],binary=true,compressed=false)
   open(file_name*".vtu","w") do fid
      if compressed
         compressor=" compressor=\"vtkZLibDataCompressor\""
      else
         compressor=""
      end
      write(fid,"<VTKFile type=\"UnstructuredGrid\" version=\"0.1\" byte_order=\"LittleEndian\" header_type=\"UInt64\"$compressor>\n")
      write(fid,"\t<UnstructuredGrid>\n")
      write(fid,"\t\t<Piece NumberOfPoints=\"$(size(mesh.points,2))\" NumberOfCells=\"$(length(mesh.triangles))\">\n")
      write(fid,"\t\t\t<Points>\n")
      # write(fid,"\t\t\t\t<DataArray type=\"Float64\" NumberOfComponents=\"3\" format=\"ascii\">")
      # for pnt in mesh.points
      #    write(fid,string(pnt)*" ")
      # end
      # write(fid,"</DataArray>")
      write_data_array(fid,Dict("Coordinates"=>mesh.points),binary,compressed)
      write(fid,"\t\t\t</Points>\n")
      write(fid,"\t\t\t<Cells>\n")
      write(fid,"\t\t\t\t<DataArray type=\"Int32\" Name=\"connectivity\" format=\"ascii\">")
      for tri in mesh.triangles
         for pnt in tri
            write(fid,string(Int32(pnt-1))*" ")
         end
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t\t<DataArray type=\"Int32\" Name=\"offsets\" format=\"ascii\">")
      off=0
      for tri in mesh.triangles
            off+=length(tri)
            write(fid,string(off)*" ")
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t\t<DataArray type=\"UInt8\" Name=\"types\" format=\"ascii\">")
      for tri in mesh.triangles
         for pnt in tri
            write(fid,"5 ")
         end
      end
      write(fid,"</DataArray>\n")
      write(fid,"\t\t\t</Cells>\n")

      write(fid,"\t\t\t<CellData>\n")
      write_data_array(fid,data,binary,compressed)
      write(fid,"\t\t\t</CellData>\n")
      write(fid,"\t\t</Piece>\n")
      write(fid,"\t</UnstructuredGrid>\n")
      write(fid,"</VTKFile>\n")
   end
   return nothing
end

##

"""
    vtk_write(file_name, mesh, data)

Write vtk-file containing the datavalues `data` given on the usntructured grid `mesh`.

# Arguments
- `file_name::String`: name given to the written files. The name will be preceeded by a specifier. See Notes below.
- `mesh::Mesh`: mesh associated with the data.
- `data::Dict`: Dictionary containing the data.

# Notes
The routine automatically sorts the data according to its type and writes it in up to three diffrent files.
Data that is constant on a tetrahedron goes into `"\$(filename)_const.vtu"`,
data that is linearly interpolated on a tetrahedron goes into `"\$(filename)_lin.vtu"`,
and data that is quadratically interpolated on a tetrahedron goes into `"\$(filename)_quad.vtu"`
"""
function vtk_write(file_name::String,mesh::Mesh,data::Dict,binary=false,compressed=false) #TODO: write into one file
   data0=Dict()
   data1=Dict()
   data2=Dict()
   data_tri=Dict()
   for (key,val) in data
      if maximum(size(val))==length(mesh.tetrahedra)
         data0[key]=val
      elseif maximum(size(val))==size(mesh.points,2) #TODO: check for rare case of same number of points and tets
         data1[key]=val
      elseif maximum(size(val))==size(mesh.points,2)+length(mesh.lines)
         data2[key]=val
      elseif maximum(size(val))==length(mesh.triangles)
         data_tri[key]=val
      else
         println("WARNING: data length of '$key' does not fit to any mesh dimension.")
      end
   end

   if length(data0)!=0
      vtk_write0(file_name*"_const",mesh,data0,binary,compressed)
   end
   if length(data1)!=0
      vtk_write1(file_name*"_lin",mesh,data1,binary,compressed)
   end
   if length(data2)!=0
      vtk_write2(file_name*"_quad",mesh,data2,binary,compressed)
   end
   if length(data_tri)!=0
      vtk_write_tri(file_name*"_tri",mesh,data_tri,binary,compressed)
   end
   return
end
## bloch utilities
# function bloch_expand(mesh::Mesh,sol,b=:b)
#     naxis=mesh.dos.naxis
#     nxsector=mesh.dos.nxsector
#     DOS=mesh.dos.DOS
#     v=zeros(ComplexF64,naxis+nxsector*DOS)
#     v[1:naxis]=sol.v[1:naxis]
#     B=sol.params[b]
#     for s=0:DOS-1
#         v[naxis+1+s*nxsector:naxis+(s+1)*nxsector]=sol.v[naxis+1:naxis+nxsector].*exp(+2.0im*pi/DOS*B*s)
#     end
#     return v
# end
# function bloch_expand(mesh::Mesh,vec::Array,B::Real=0)
#     naxis=mesh.dos.naxis
#     nxsector=mesh.dos.nxsector
#     DOS=mesh.dos.DOS
#     v=zeros(ComplexF64,naxis+nxsector*DOS)
#     v[1:naxis]=sol.v[1:naxis]
#     for s=0:DOS-1
#         v[naxis+1+s*nxsector:naxis+(s+1)*nxsector]=vec[naxis+1:naxis+nxsector].*exp(+2.0im*pi/DOS*B*s)
#     end
#     return v
# end

## helper functions for writing vtk data

function write_data_array(fid,data,binary=true,compressed=false)
   if binary
      if compressed
         iob = IOBuffer()
         #io.append=true
         #iob = Base64.Base64EncodePipe(iob);
      else
         iob = Base64.Base64EncodePipe(fid);
      end
      format="\"binary\""
   else
      format="\"ascii\""
   end

   for dataname in keys(data)
      if size(data[dataname],2)>1 #write vector output #TODO: Sanity checks for three components
         datatype="$(typeof(data[dataname][1]))"
         bytesize=sizeof(data[dataname][1])
         write(fid,"\t\t\t\t<DataArray type=\"$datatype\" NumberOfComponents=\"3\" Name=\"$dataname\" format=$format>")
         write(fid,"\n\t\t\t\t\t")
         if binary
            len=length(data[dataname])
            if !compressed
               num_of_bytes=UInt64(length(data[dataname])*bytesize)
               write(iob, num_of_bytes)
            end
            for vec in data[dataname]
               for num in vec
                  write(iob, num)#write number to base64 coded stream
               end
            end
         else
            for vec in data[dataname]
               write(fid,string(vec)*" ")
            end
         end
      else
         datatype="$(typeof(data[dataname][1]))"
         bytesize=sizeof(data[dataname][1])
         #datatype="Float64"
         write(fid,"\t\t\t\t<DataArray type=\"$datatype\" Name=\"$dataname\" format=$format>")
         write(fid,"\n\t\t\t\t\t")
         if binary
               #num_of_bytes=UInt64(length(data[dataname])*bytesize)
            if !compressed
               num_of_bytes=UInt64(length(data[dataname])*bytesize)
               write(iob, num_of_bytes)
            end
            for datum in data[dataname]
               idxx=0
               #println("$idxx : $dataname")
               for (idxc,num) in enumerate(datum)
                  idxx+=1
                  #println("$idxx : $num")
                  write(iob, num)#write number to base64 coded stream
               end
            end
         else
            for datum in data[dataname]
                  write(fid,string(datum)*" ")
            end
         end
      end
      if binary
         #close(iob) #close io stream
         #zip compress data
         if compressed
            block_zcompress!(fid,take!(iob))
         end
         #write(fid,String(take!(io))) #write encoded data to file
      end
      write(fid,"\n\t\t\t\t</DataArray>\n")
   end
   if compressed
      close(iob)
   end
end

function block_zcompress!(fid,data,block_length=2^10)
    iob64_encode = Base64.Base64EncodePipe(fid);
    len=length(data)
    number_of_blocks=len÷block_length
    length_of_last_partial_block=len%block_length
    if length_of_last_partial_block!=0
        number_of_blocks+=1
    end
    compressed_block_size = Array{UInt64}(undef,number_of_blocks)
    compressed_data=[]
    for idx=1:number_of_blocks
        if idx==number_of_blocks
            block=data[((idx-1)*block_length+1):end]
        else
            block=data[((idx-1)*block_length+1):(idx*block_length)]
        end
        compressed_block = CodecZlib.transcode(CodecZlib.GzipCompressor, block)
        compressed_block_size[idx]=sizeof(compressed_block)
        append!(compressed_data,[compressed_block])
    end

    #write compressed binary entry
    #io = IOBuffer()
    #header
    write(iob64_encode, UInt64(number_of_blocks))
    write(iob64_encode, UInt64(block_length))
    write(iob64_encode, UInt64(length_of_last_partial_block))
    for idx=1:number_of_blocks
      write(iob64_encode, UInt64(compressed_block_size[idx]))
    end
    #data
    for idx=1:number_of_blocks
        for num in compressed_data[idx]
            write(iob64_encode, num)
        end
    end

    return nothing
end


function simulate!(s::Simulator{T}; verbose=false) where T
    status = false

    p = s.policy
	w = s.dist
	traj = s.traj
	N = length(traj.u)

	# last_computed_control = p.ref_traj.u[1] / N_sample
	buffer_time = 0.0
    for t = 1:N
		# println("timestep $t / $N, time $time")
		policy_time = @elapsed traj.u[t] = exec_policy(p, traj, t, buffer_time)
		s.opts.record && (s.stats.policy_time[t] = policy_time)
		(buffer_time < 0.0) && (buffer_time = policy_time)
		buffer_time -= s.h # decrement the buffer time of the sim step

		# traj.u[t] .= last_computed_control
		# if buffer_time <= 0.0
			# policy
			# buffer_time = @elapsed last_computed_control = exec_policy(p, traj, t, buffer_time)
			# s.opts.record && (s.stats.policy_time[t] = buffer_time)
        # end
		# buffer_time -= s.h # decrement the buffer time of the sim step

        # disturbances
        traj.w[t] .= disturbances(w, traj.q[t+1], t)

        # step
        status = RoboDojo.step!(s, t, verbose=verbose)
        !status && break
    end

    return status
end

function exec_policy(p::CIMPC{T,NQ,NU,NW,NC}, traj::Trajectory{T}, t::Int, buffer_time::T) where {T,NQ,NU,NW,NC}
	# reset
	if t == 1
		p.cnt[1] = p.N_sample
		p.q0 .= p.ref_traj.q[1]
		p.altitude .= 0.0
		set_trajectory!(p.traj, p.ref_traj)
		set_implicit_trajectory!(p.im_traj, p.im_traj_cache)
		reset_window!(p.window)
	end

    if p.cnt[1] == p.N_sample
		# altitude_update
		(p.opts.altitude_update && t > 1) && (update_altitude!(p.altitude, p.ϕ, p.s,
									traj, t, NC, p.N_sample,
									threshold = p.opts.altitude_impact_threshold,
									verbose = p.opts.altitude_verbose))
		set_altitude!(p.im_traj, p.altitude)

		# # # optimize
		# q1 = traj.q[t+1]
		# newton_solve!(p.newton, p.s, p.q0, q1,
		# 	p.window, p.im_traj, p.traj, warm_start = t > 1)

		update!(p.im_traj, p.traj, p.s, p.altitude, p.κ[1], p.traj.H)

		# visualize
		p.opts.live_plotting && live_plotting(p.s.model, p.traj, traj, p.newton, p.q0, traj.q[t+1], t)

		# shift trajectory
		rot_n_stride!(p.traj, p.traj_cache, p.stride, p.window)

		# update
		update_window!(p.window, p.ref_traj.H)
		# p.q0 .= q1

		# reset count
		p.cnt[1] = 0
    end
	p.cnt[1] += 1

	if buffer_time <= 0.0
		@show t
		# optimize
		q0 = traj.q[max(1,t+1-N_sample)]
		q1 = traj.q[t+1]
		newton_solve!(p.newton, p.s, q0, q1,
			p.window, p.im_traj, p.traj, warm_start = t > 1)

		update!(p.im_traj, p.traj, p.s, p.altitude, p.κ[1], p.traj.H)
	end

	# scale control
	if p.newton_mode == :direct
		p.u .= p.newton.traj.u[1]
		p.u ./= p.N_sample
	elseif p.newton_mode == :structure
		p.u .= p.newton.u[1]
		p.u ./= p.N_sample
	else
		println("newton mode specified not available")
	end

	return p.u
end


using Random
using Plots
#=
Julia Func B:
 This is a sequence of numbers I found on the OEIS and expanded to large porportions.
 I was shocked to find the minimum would always be -199.

 Perhaps I'll come back and redo this one later. i got a little frustrated, so this looks pretty bad
 with how the variables are laid out.
=#

numMax = 200

nums = ( function ()
  global lpMax
  maxiNum::UInt64 = (2^(63)-1)
  arrPos = 1
  numArr = zeros(Int64, numMax)
  itrcnt = 3
  chck = true
  i = 1
  num = 0
  str = 0
  while i <= numMax

    if (chck == true)
     str = string(-(i))
     num = parse(Int64, str[1:2])
     chck = false
   else
     num += parse(Int8, str[itrcnt])
     itrcnt += 1
   end

   if (itrcnt > length(str))
   num ^= 2
   num += -(i)
     if ((string(num))[1] != '-' && num != 0 ) numArr[arrPos] = -(i)
     println(arrPos, " ", -(i))
     arrPos += 1
     end
     chck = true
     itrcnt = 3
     i += 1
   end
   @assert (num >= maxiNum)  ? false : true (string( "Input too high; Num: ", num))

  end
return numArr
end)()

function Plot(maxNums, nums)
  println("Would you like to see the plot?")
  if (readline() == "y")
    println("Enter Character in Console to Exit")
    gui(plot(1:maxNums,nums))
  else
    exit()
  end
  readline()
end

Plot(numMax, nums)
exit()


module MCMC

using Celeste: Model, Transform, SensitiveFloats, MCMC
using Distributions, StatsBase, DataFrames, StaticArrays, WCS

# TODO move these to model/log_prob.jl
star_param_names = ["lnr", "cug", "cgr", "cri", "ciz", "ra", "dec"]
galaxy_param_names = [star_param_names; ["gdev", "gaxis", "gangle", "gscale"]]

# functions to create star/gal log likelihoods
include("mcmc/mcmc_functions.jl")

# slice sampling function
include("mcmc/slicesample.jl")

# misc MCMC functions
include("mcmc/mcmc_misc.jl")

end


module MPFR

export
    exp10,
    MPFRFloat,
    prec,
    get_default_precision,
    set_default_precision,
    with_precision

import
    Base.(*),
    Base.+,
    Base.-,
    Base./,
    Base.<,
    Base.<=,
    Base.==,
    Base.>,
    Base.>=,
    Base.^,
    Base.besselj,
    Base.besselj0,
    Base.besselj1,
    Base.bessely,
    Base.bessely0,
    Base.bessely1,
    Base.ceil,
    Base.cmp,
    Base.convert,
    Base.copysign,
    Base.exp,
    Base.exp2,
    Base.exponent,
    Base.factorial,
    Base.floor,
    Base.integer_valued,
    Base.iround,
    Base.isfinite,
    Base.isinf,
    Base.isnan,
    Base.log,
    Base.log2,
    Base.log10,
    Base.max,
    Base.min,
    Base.mod,
    Base.modf,
    Base.nextfloat,
    Base.prevfloat,
    Base.promote_rule,
    Base.rem,
    Base.round,
    Base.show,
    Base.showcompact,
    Base.sum,
    Base.sqrt,
    Base.string,
    Base.trunc,
    # import trigonometric functions
    Base.sin,
    Base.cos,
    Base.tan,
    Base.sec,
    Base.csc,
    Base.cot,
    Base.acos,
    Base.asin,
    Base.atan,
    Base.cosh,
    Base.sinh,
    Base.tanh,
    Base.sech,
    Base.csch,
    Base.coth,
    Base.acosh,
    Base.asinh,
    Base.atanh

const ROUNDING_MODE = [0]
const DEFAULT_PRECISION = [53]

# Basic type and initialization definitions

immutable MPFRFloat{N} <: FloatingPoint
    mpfr::Vector{Int32}
    function MPFRFloat()
        if N < 2
            error("Invalid precision")
        end
        z = Array(Int32, 5)
        ccall((:mpfr_init2,:libmpfr), Void, (Ptr{Void}, Int), z, N)
        b = new(z)
        finalizer(b.mpfr, MPFR_clear)
        return b
    end
end

MPFR_clear(mpfr::Vector{Int32}) = ccall((:mpfr_clear, :libmpfr), Void, (Ptr{Void},), mpfr)

function MPFRFloat(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_set, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

for (fJ, fC) in ((:si,:Int), (:ui,:Uint), (:d,:Float64))
    @eval begin
        function MPFRFloat(x::($fC))
            z = MPFRFloat{DEFAULT_PRECISION[end]}()
            ccall(($(string(:mpfr_set_,fJ)), :libmpfr), Int32, (Ptr{Void}, ($fC), Int32), z.mpfr, x, ROUNDING_MODE[end])
            return z
        end
    end
end

function MPFRFloat(x::BigInt)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_set_z, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpz, ROUNDING_MODE[end])
    return z
end

function MPFRFloat(x::BigFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_set_f, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpf, ROUNDING_MODE[end])
    return z
end

function MPFRFloat(x::String, base::Int)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    err = ccall((:mpfr_set_str, :libmpfr), Int32, (Ptr{Void}, Ptr{Uint8}, Int32, Int32), z.mpfr, x, base, ROUNDING_MODE[end])
    if err != 0; error("Invalid input"); end
    return z
end
MPFRFloat(x::String) = MPFRFloat(x, 10)

MPFRFloat(x::Bool) = MPFRFloat(uint(x))
MPFRFloat(x::Signed) = MPFRFloat(int(x))
MPFRFloat(x::Unsigned) = MPFRFloat(uint(x))
if WORD_SIZE == 32
    MPFRFloat(x::Int64) = MPFRFloat(string(x))
    MPFRFloat(x::Uint64) = MPFRFloat(BigInt(x))
end
MPFRFloat(x::Float32) = MPFRFloat(float64(x))
MPFRFloat(x::Rational) = MPFRFloat(num(x)) / MPFRFloat(den(x))

# TODO: fix the precision support here
convert{N}(::Type{MPFRFloat{N}}, x::Rational) = MPFRFloat(x) # to resolve ambiguity
convert{N}(::Type{MPFRFloat{N}}, x::Real) = MPFRFloat(x)
convert(::Type{MPFRFloat}, x::Real) = MPFRFloat(x)
convert{N,T}(::Type{MPFRFloat{N}}, x::MPFRFloat{T}) = MPFRFloat(x)

convert(::Type{Int64}, x::MPFRFloat) = integer_valued(x) ?
    ccall((:mpfr_get_si,:libmpfr), Int64, (Ptr{Void}, Int32), x.mpfr, ROUNDING_MODE[end]) :
    throw(InexactError())
convert(::Type{Int32}, x::MPFRFloat) = int32(convert(Int64, x))
convert(::Type{Uint64}, x::MPFRFloat) = integer_valued(x) ?
    ccall((:mpfr_get_ui,:libmpfr), Uint64, (Ptr{Void}, Int32), x.mpfr, ROUNDING_MODE[end]) :
    throw(InexactError())
function convert(::Type{BigInt}, x::MPFRFloat)
    if integer_valued(x)
        z = BigInt()
        ccall((:mpfr_get_z,:libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpz, x.mpfr, ROUNDING_MODE[end])
        return z
    else
        throw(InexactError())
    end
end
convert(::Type{Uint32}, x::MPFRFloat) = uint32(convert(Int64, x))
convert(::Type{Float64}, x::MPFRFloat) = ccall((:mpfr_get_d,:libmpfr), Float64, (Ptr{Void},), x.mpfr)
convert(::Type{Float32}, x::MPFRFloat) = ccall((:mpfr_get_flt,:libmpfr), Float32, (Ptr{Void},), x.mpfr)

# If two different precisions given, promote to the default
promote_rule{T,S}(::Type{MPFRFloat{T}}, ::Type{MPFRFloat{S}}) =
    MPFRFloat{DEFAULT_PRECISION[end]}

promote_rule{T<:Real,N}(::Type{MPFRFloat{N}}, ::Type{T}) = MPFRFloat{N}
promote_rule{T<:Real}(::Type{MPFRFloat}, ::Type{T}) = MPFRFloat

# TODO: Decide if overwriting the default BigFloat rule is good
#promote_rule{T<:FloatingPoint}(::Type{BigInt},::Type{T}) = MPFRFloat{DEFAULT_PRECISION[end]}
#promote_rule{T<:FloatingPoint}(::Type{BigFloat},::Type{T}) = MPFRFloat{DEFAULT_PRECISION[end]}

# Basic operations

for (fJ, fC) in ((:+,:add), (:-,:sub), (:*,:mul), (:/,:div), (:^, :pow))
    @eval begin
        function ($fJ)(x::MPFRFloat, y::MPFRFloat)
            z = MPFRFloat{DEFAULT_PRECISION[end]}()
            ccall(($(string(:mpfr_,fC)),:libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, y.mpfr, ROUNDING_MODE[end])
            return z
        end
    end
end

function -(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_neg, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function cmp(x::MPFRFloat, y::MPFRFloat)
    ccall((:mpfr_cmp, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}), x.mpfr, y.mpfr)
end

function sqrt(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_sqrt, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    if isnan(z)
        throw(DomainError())
    end
    return z
end

for f in (:ceil, :floor, :trunc)
    @eval begin
        function ($f)(x::MPFRFloat)
            z = MPFRFloat{DEFAULT_PRECISION[end]}()
            ccall(($(string(:mpfr_,f)), :libmpfr), Int32, (Ptr{Void}, Ptr{Void}), z.mpfr, x.mpfr)
            return z
        end
    end
end

function ^(x::MPFRFloat, y::Uint)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_pow_ui, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Uint, Int32), z.mpfr, x.mpfr, y, ROUNDING_MODE[end])
    return z
end

function ^(x::MPFRFloat, y::Int)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_pow_si, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int, Int32), z.mpfr, x.mpfr, y, ROUNDING_MODE[end])
    return z
end

function ^(x::MPFRFloat, y::BigInt)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_pow_z, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, y.mpz, ROUNDING_MODE[end])
    return z
end

function exp(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_exp, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function exp2(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_exp2, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function exp10(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_exp10, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function besselj0(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_j0, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function besselj1(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_j1, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function besselj(nu::Integer, x::MPFRFloat)
    n = int64(nu)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_jn, :libmpfr), Int32, (Ptr{Void}, Int64, Ptr{Void}, Int32), z.mpfr, n, x.mpfr, ROUNDING_MODE[end])
    return z
end

function bessely0(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_y0, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function bessely1(x::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_y1, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function bessely(nu::Integer, x::MPFRFloat)
    n = int64(nu)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_yn, :libmpfr), Int32, (Ptr{Void}, Int64, Ptr{Void}, Int32), z.mpfr, n, x.mpfr, ROUNDING_MODE[end])
    return z
end

function factorial(x::MPFRFloat)
    if x < 0 || !integer_valued(x)
        throw(DomainError())
    end
    ui = uint64(x)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_fac_ui, :libmpfr), Int32, (Ptr{Void}, Uint64, Int32), z.mpfr, ui, ROUNDING_MODE[end])
    return z
end

function log(x::MPFRFloat)
    if x < 0
        throw(DomainError())
    end
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_log, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function log2(x::MPFRFloat)
    if x < 0
        throw(DomainError())
    end
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_log2, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function log10(x::MPFRFloat)
    if x < 0
        throw(DomainError())
    end
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_log10, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
    return z
end

function max(x::MPFRFloat, y::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_max, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, y.mpfr, ROUNDING_MODE[end])
    return z
end

function min(x::MPFRFloat, y::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_min, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, y.mpfr, ROUNDING_MODE[end])
    return z
end

function modf(x::MPFRFloat)
    if isinf(x)
        return (MPFRFloat(NaN), x)
    end
    zint = MPFRFloat{DEFAULT_PRECISION[end]}()
    zfloat = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_modf, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Ptr{Void}, Int32), zint.mpfr, zfloat.mpfr, x.mpfr, ROUNDING_MODE[end])
    return (zfloat, zint)
end

function rem(x::MPFRFloat, y::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_remainder, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, y.mpfr, ROUNDING_MODE[end])
    return z
end

function sum{T<:MPFRFloat}(arr::AbstractArray{T})
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    n = length(arr)
    ptrarr = [pointer(x.mpfr) for x in arr]
    ccall((:mpfr_sum, :libmpfr), Int32,
        (Ptr{Void}, Ptr{Void}, Uint, Int32),
        z.mpfr, ptrarr, n, ROUNDING_MODE[1])
    return z
end

# Trigonometric functions
# Every NaN is thrown as an error, and it follows somewhat closely
# the Base functions behavior
for f in (:sin,:cos,:tan,:sec,:csc,:cot,:acos,:asin,:atan,
        :cosh,:sinh,:tanh,:sech,:csch,:coth,:acosh,:asinh,:atanh)
    @eval begin
        function ($f)(x::MPFRFloat)
            z = MPFRFloat{DEFAULT_PRECISION[end]}()
            ccall(($(string(:mpfr_,f)), :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, ROUNDING_MODE[end])
            if isnan(z)
                throw(DomainError())
            end
            return z
        end
    end
end

# Utility functions
==(x::MPFRFloat, y::MPFRFloat) = ccall((:mpfr_equal_p, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}), x.mpfr, y.mpfr) != 0
<=(x::MPFRFloat, y::MPFRFloat) = ccall((:mpfr_lessequal_p, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}), x.mpfr, y.mpfr) != 0
>=(x::MPFRFloat, y::MPFRFloat) = ccall((:mpfr_greaterequal_p, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}), x.mpfr, y.mpfr) != 0
<(x::MPFRFloat, y::MPFRFloat) = ccall((:mpfr_less_p, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}), x.mpfr, y.mpfr) != 0
>(x::MPFRFloat, y::MPFRFloat) = ccall((:mpfr_greater_p, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}), x.mpfr, y.mpfr) != 0

function prec(x::MPFRFloat)
    return ccall((:mpfr_get_prec, :libmpfr), Int, (Ptr{Void},), x.mpfr)
end

get_default_precision() = DEFAULT_PRECISION[end]
function set_default_precision(x::Int)
    if x < 2
        throw(DomainError())
    end
    DEFAULT_PRECISION[end] = x
end

function copysign(x::MPFRFloat, y::MPFRFloat)
    z = MPFRFloat{DEFAULT_PRECISION[end]}()
    ccall((:mpfr_copysign, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Ptr{Void}, Int32), z.mpfr, x.mpfr, y.mpfr, ROUNDING_MODE[end])
    return z
end

function exponent(x::MPFRFloat)
    if x == 0 || !isfinite(x)
        throw(DomainError())
    end
    # The '- 1' is to make it work as Base.exponent
    return ccall((:mpfr_get_exp, :libmpfr), Int, (Ptr{Void},), x.mpfr) - 1
end

function integer_valued(x::MPFRFloat)
    return ccall((:mpfr_integer_p, :libmpfr), Int32, (Ptr{Void},), x.mpfr) != 0
end

function iround(x::MPFRFloat)
    fits = ccall((:mpfr_fits_slong_p, :libmpfr), Int32, (Ptr{Void}, Int32), x.mpfr, ROUNDING_MODE[end])
    if fits != 0
        return ccall((:mpfr_get_si, :libmpfr), Int64, (Ptr{Void}, Int32), x.mpfr, ROUNDING_MODE[end])
    end
    z = BigInt()
    ccall((:mpfr_get_z, :libmpfr), Int32, (Ptr{Void}, Ptr{Void}, Int32), z.mpz, x.mpfr, ROUNDING_MODE[end])
    return z
end

for (f,t) in ((:uint8,:Uint8), (:uint16,:Uint16), (:uint32,:Uint32),
              (:int64,:Int64), (:uint64,:Uint64),
              # requires int128/uint128(::Type{BigInt}) support
              # (:int128,:Int128), (:uint128,:Uint128),
              (:unsigned,:Uint), (:uint,:Uint))
    @eval iround(::Type{$t}, x::MPFRFloat) = ($f)(iround(x))
end

isfinite(x::MPFRFloat) = !isinf(x)
function isinf(x::MPFRFloat)
    return ccall((:mpfr_inf_p, :libmpfr), Int32, (Ptr{Void},), x.mpfr) != 0
end

function isnan(x::MPFRFloat)
    return ccall((:mpfr_nan_p, :libmpfr), Int32, (Ptr{Void},), x.mpfr) != 0
end

function nextfloat(x::MPFRFloat)
    z = MPFRFloat(x)
   ccall((:mpfr_nextabove, :libmpfr), Int32, (Ptr{Void},), z.mpfr) != 0
   return z
end

function prevfloat(x::MPFRFloat)
    z = MPFRFloat(x)
   ccall((:mpfr_nextbelow, :libmpfr), Int32, (Ptr{Void},), z.mpfr) != 0
   return z
end

function with_precision(f::Function, precision::Integer)
    old_precision = get_default_precision()
    set_default_precision(precision)
    ret = f()
    set_default_precision(old_precision)
    return ret
end

function round(x::MPFRFloat, prec::Int)
    if prec < 1
        throw(DomainError())
    end
    prec = int(ceil(log2(10^prec)))
    z = MPFRFloat(x)
    ccall((:mpfr_prec_round, :libmpfr), Int32, (Ptr{Void}, Int, Int32), z.mpfr, prec, ROUNDING_MODE[end])
    return z
end

function string(x::MPFRFloat)
    lng = 128
    for i = 1:2
        z = Array(Uint8, lng)
        lng = ccall((:mpfr_snprintf,:libmpfr), Int32, (Ptr{Uint8}, Uint, Ptr{Uint8}, Ptr{Void}...), z, lng, "%.Re", x.mpfr)
        if lng < 128 || i == 2
            return bytestring(convert(Ptr{Uint8}, z[1:lng]))
        end
    end
end

show(io::IO, b::MPFRFloat) = print(io, string(b) * " with $(prec(b)) bits of precision")
showcompact(io::IO, b::MPFRFloat) = print(io, string(b))

end #module


module Funds

    using Dates
    import ..YearMonths
    import ..HTTP

    struct FundData
        fund_type::Union{String, Nothing}
        cnpj::String
        date::Date
        total_value::Float64
        quota_value::Float64
        net_worth::Float64
        deposits::Float64
        withdraws::Float64
        number_quota_holders::Int64
    end

    function get_month_quotes(ym::YearMonths.YearMonth)::Vector{FundData}
        @assert ym >= YearMonths.YearMonth(2017, 1)

        raw = _get_data(ym)

        out = Vector{FundData}()

        header = split(String(raw[1]), ";")

        for i in 2:length(raw)

            line = raw[i]

            if line == ""
                continue
            end

            s = split(String(line), ";")

            if ym <= YearMonths.YearMonth(2020,3)

                @assert length(s) == 8 "bad CVM data, this line should have 8 fields: $s"
                fund_type = nothing
                cnpj = String(s[1])
                date = Date(String(s[2]))
                total_value = parse(Float64, String(s[3]))
                quota_value = parse(Float64, String(s[4]))
                net_worth = parse(Float64, String(s[5]))
                deposits = parse(Float64, String(s[6]))
                withdraws = parse(Float64, String(s[7]))
                number_quota_holders = parse(Int64, String(s[8]))

            else
                # after 2020-03
                @assert length(s) == 9 "bad CVM data, this line should have 9 fields: $s"
                fund_type = String(s[1])
                cnpj = String(s[2])
                date = Date(String(s[3]))
                total_value = parse(Float64, String(s[4]))
                quota_value = parse(Float64, String(s[5]))
                net_worth = parse(Float64, String(s[6]))
                deposits = parse(Float64, String(s[7]))
                withdraws = parse(Float64, String(s[8]))
                number_quota_holders = parse(Int64, String(s[9]))

            end
            d = FundData(
                fund_type, cnpj, date,
                total_value, quota_value, net_worth,
                deposits, withdraws, number_quota_holders
            )
            push!(out, d)

        end

        return(out)
    end

    function _get_url(ym::YearMonths.YearMonth)::String
        str_year = "$(YearMonths.year(ym))"
        str_month = lpad("$(YearMonths.month(ym))", 2, "0")
        url = "http://dados.cvm.gov.br/dados/FI/DOC/INF_DIARIO/DADOS/inf_diario_fi_$str_year$str_month.csv"
        return url
    end

    function _get_data(ym::YearMonths.YearMonth)
        url = _get_url(ym)
        r = HTTP.get(url, status_exception = false)
        if r.status == 200
            return split(String(r.body), "\r\n")
        elseif r.status == 404
            throw(BoundsError("CVM fund data not found for $ym"))
        else
            error("error retrieving CVM fund data: $r")
        end
    end

end


module KnuthBendix

using ProgressMeter

export Word, Alphabet, RewritingSystem, LenLex, knuthbendix

include("abstract_words.jl")
include("words.jl")
include("bufferwords.jl")
include("alphabets.jl")
include("orderings.jl")
include("rules.jl")

include("rewriting.jl")
include("automata.jl")
include("helper_structures.jl")
include("derive_rule.jl")
include("force_confluence.jl")
include("kbs.jl")

include("parsing.jl")
end


module STREAMBenchmark

using BenchmarkTools
using Statistics, Downloads
using Base.Threads: nthreads, @threads
using LoopVectorization

include("kernels_allthreads.jl")
include("kernels_nothreads.jl")
include("kernels_nthreads.jl")
include("benchmarks.jl")
include("original.jl")

export memory_bandwidth, benchmark, scaling_benchmark

end


module HypergraphsEpidemics

using CSV
using DataFrames
using Query
using Tables
using Dates
using DataStructures
using Distributions
using SimpleHypergraphs
using Statistics
using Combinatorics
using Pipe: @pipe
using Random
using JSON
using JSON3
using Serialization

export Person
export Household
export Company
export Contact_simulation_options

export generate_dataset
export start_simulation

include("common/Constants.jl")
include("common/EventWriter.jl")
include("common/utils.jl")
include("common/structs.jl")
include("RoutineAutomaton.jl")
include("DatasetGeneration.jl")
include("Simulation.jl")

include("../deps/HGEpidemics/src/HGEpidemics.jl")

using .HGEpidemics

end


export weight_tuple

"""
    trainable_first(net)

Returns the first array in trainable(net), by recursively applying first to it.
"""
trainable_first(net) = trainable_first(first(trainable(net)))
trainable_first(x::AbstractArray)   = x

"""
    trainable_length

"""
trainable_length(net::NeuralNetwork) = _tlen(net)
function _tlen(t)
    i = 0
    for el=trainable(t)
        i+=_tlen(el)
    end
    return i
end
_tlen(t::AbstractArray) = length(t)

"""
    weight_tuple(net)

Returns a named tuple holding all the fields in the network, and an
extra field named `tuple_all_weights` who has the same type
"""
function weight_tuple(obj, vec=similar(trainable_first(obj), 0),
                      start=1)
    x = trainable(obj)
    i = 0
    if x isa Tuple
        d=Vector()
        for f=propertynames(x)
            di, val = weight_tuple(getproperty(x,f), vec, start+i)
            i += di; push!(d, val)
        end
        i, Tuple(d)
    else
        d=Dict{Symbol, Any}()
        for f=propertynames(x)
            di, val = weight_tuple(getproperty(x,f), vec, start+i)
            i += di; push!(d, f=>val)
        end
        i, (;d...)
    end
end

function weight_tuple(obj::Tuple, vec::AbstractVector, start)
    x = trainable(obj)
    i = 0
    d=Vector()
    for f=propertynames(x)
        di, val = weight_tuple(getfield(x,f), vec, start+i)
        i += di; push!(d, val)
    end
    i, Tuple(d)
end

function weight_tuple(x::AbstractArray{<:Number}, vec::AbstractVector, start)
    length(vec) < start+length(x)-1 && resize!(vec, start+length(x)-1)
    @views data_vec = vec[start:start+length(x)-1]
    if size(x) == size(data_vec)
        reshpd_params = data_vec
    else
        reshpd_params = reshape(data_vec, size(x))
    end
    return length(x), reshpd_params
end

## For Batches
function batched_weight_tuple(grad_tup, bsz=1)
    @info "useless"
    all_weights = grad_tup.tuple_all_weights
    all_weghts_new = [similar(w, length(w), bsz) for w=all_weights][1]
    return batched_weight_tuple(grad_tup, all_weghts_new)[2]
end

function batched_weight_tuple(obj, vec::AbstractMatrix, start=1)
    x = trainable(obj)
    i = 0
    if x isa Tuple
        d=Vector()
        for f=propertynames(x)
            f==:tuple_all_weights && continue
            di, val = batched_weight_tuple(getfield(x,f), vec, start+i)
            i += di; push!(d, val)
        end
        i, Tuple(d)
    else
        d=Dict{Symbol, Any}()
        for f=propertynames(x)
            f==:tuple_all_weights && continue
            di, val = batched_weight_tuple(getfield(x,f), vec, start+i)
            i += di; push!(d, f=>val)
        end
        i, (;d...)
    end
end

function batched_weight_tuple(x::AbstractArray{<:Number}, vec::AbstractMatrix, start)
    @assert length(vec) >= start+length(x)-1
    bsz = size(vec, 2)

    @views data_vec = vec[start:start+length(x)-1, :]
    if size(data_vec) != (size(x)..., bsz)
        reshpd_params = reshape(data_vec, size(x)..., bsz)
        if x isa Array
            reshpd_params = StridedView(reshpd_params)
        end
    else
        reshpd_params = data_vec
    end
    return length(x), reshpd_params
end # ? stridedView?

"""
    apply_recurse!(out::tuple, orig::Zygote.Grads, weights, op=identity)

recurses into all fields of `out`, which is assumed to have the structure
generated by `trainable(weights)`. When it hits abstractarrays, it copies
the corresponding field `orig[weights[key]]` to `out[key]`, by applying the
elementwise unary operation `op`.

This convoluted way is because zygote returns gradient results that do not respect
the trainable interface, and we need to reconstruct it.
"""
function apply_recurse!(out, orig, weights, op=identity)
    weights = trainable(weights)
    for key=keys(out)
        apply_recurse!(out[key], orig, weights[key], op)
    end
    return out
end

apply_recurse!(out::AbstractArray, orig, weight::AbstractArray, op) =
    op(copyto!(out, orig[weight]))

"""
    accum_recurse!(out::tuple, orig::Zygote.Grads, weights, op=identity; init=true)

Similar to `apply_recurse!`, but accumulates and only sets the out to 0 if
`init==true`.
"""
function accum_recurse!(out, orig, weights, op=identity; init=true)
    weights = trainable(weights)
    for key=keys(out)
        accum_recurse!(out[key], orig, weights[key], op; init=init)
    end
    return out
end

function accum_recurse!(out::AbstractArray, orig, weight::AbstractArray, op; init=true)
    if init
        out .= 0
    end

    if length(size(out)) == 1 && length(size(orig[weight])) == 2
        out .+= op.(view(orig[weight],:))
    else
        out .+= op.(orig[weight])
    end
end


#
# Combinator-based embedded query language.
#

module DataKnots

export
    @query,
    Collect,
    Count,
    DataKnot,
    Drop,
    Each,
    Exists,
    Filter,
    First,
    Get,
    Given,
    Group,
    Is,
    Is0to1,
    Is0toN,
    Is1to1,
    Is1toN,
    It,
    Join,
    Keep,
    Label,
    Last,
    Let,
    Lift,
    Max,
    Min,
    Mix,
    Nth,
    Record,
    Sort,
    Sum,
    Tag,
    Take,
    Unique,
    unitknot

include("layouts.jl")
include("vectors.jl")
include("shapes.jl")
include("knots.jl")
include("pipelines.jl")
include("rewrites.jl")
include("queries.jl")

end


using Distributed

using CSV
using DataFrames
@everywhere using FileIO
using JLD2
using ProgressMeter
using Statistics

@assert length(ARGS) >= 1 "no experiment path given"

prefix_path = ARGS[1]
num_trajectories = length(ARGS)==2 ? parse(Int, ARGS[2]) : 1000

df = CSV.read(joinpath(prefix_path, "parameters_map.csv"), DataFrame)
output_dirs = map( x-> joinpath(prefix_path, x, "output"), df.path)
nans = fill(NaN, num_trajectories)

results = @showprogress pmap(output_dirs) do path
  try
    load(joinpath(path, "summary.jld2"), "last_infections", "num_infections", "peak_daily_detections", "peak_daily_infections")
  catch e
    @warn "could not load summary" path
    nans, nans, nans, nans
  end
end

last_infections = getindex.(results, 1)
num_infections = getindex.(results, 2)
peak_daily_detections = getindex.(results, 3)
peak_daily_infections = getindex.(results, 4)

df.mean_last_infection_time = last_infections .|> mean
df.mean_num_infected = num_infections .|> mean
df.peak_daily_detections = peak_daily_detections .|> mean
df.peak_daily_infections = peak_daily_infections .|> mean

save(joinpath(prefix_path,"summary.jld2"),
  "summary", df,
  "last_infections", last_infections,
  "num_infections", num_infections,
  "peak_daily_detections", peak_daily_detections,
  "peak_daily_infections", peak_daily_infections,
  compress=true
)

save(joinpath(prefix_path,"summary.csv"), df)