text:
# =============================================================
#		MODULE: reading
# =============================================================
module reading

	using DelimitedFiles

	# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
	#		FUNCTION : psd
	# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
	function PSD(Path)
		# Reading PSD data

		Id =  READ_HEADER(Path, "Id" )
		N_∑Psd = length(Id)
		∑Psd = Array{Float64}(undef, (N_∑Psd, 10))

		∑Psd[:,10] = READ_HEADER(Path, "PSD_2_mm")
		∑Psd[:,9] = READ_HEADER(Path, "PSD_1_mm")
		∑Psd[:,8] = READ_HEADER(Path, "PSD_05_mm")
		∑Psd[:,7] = READ_HEADER(Path, "PSD_025_mm")
		∑Psd[:,6] = READ_HEADER(Path, "PSD_0125_mm")
		∑Psd[:,5] = READ_HEADER(Path, "PSD_0063_mm")
		∑Psd[:,4] = READ_HEADER(Path, "PSD_002_mm")
		∑Psd[:,3] = READ_HEADER(Path, "PSD_001_mm")
		∑Psd[:,2] = READ_HEADER(Path, "PSD_0006_mm")
		∑Psd[:,1] = READ_HEADER(Path, "PSD_0002_mm")

		return ∑Psd, Id,N_∑Psd
	end  # function PSD

	# ============================
	# 		READ_HEADER
	# ============================
	function READ_HEADER(Path, Name)
		# Read data
			Data =  DelimitedFiles.readdlm(Path, ',')
			N_X, N_Y = size(Data) # Size of the array

			# Reading header
			Header = fill("", N_Y)
			for i in 1:N_Y
				Header[i] = Data[1,i]
			end

			# Getting the column which matches the name of the header
			Name = replace(Name, " " => "") # Remove white spaces
			try
				global Data_Output = Data[2:N_X,findfirst(isequal(Name), Header)]
			catch
				error("\n \n HyCatch ERROR: cannot find   $Name  in $Path \n \n")
			end

			N_X -= 1 # To take consideration of the header

			return Data_Output
		end # function READ_HEADER

end  # module reading
# ............................................................

max_stars_repo_name: joseph-pollacco/TextureTriangleNZ

max_issues_repo_name: joseph-pollacco/TextureTriangleNZ

max_forks_repo_name: joseph-pollacco/TextureTriangleNZ


text:
module NestedSampling

# Represents an N-dimensional ellipsoid
immutable Ellipsoid
    ctr::Vector{Float64}  # center coordinates
    cov::Array{Float64, 2}
    icov::Array{Float64, 2}  # inverse of cov
    vol::Float64
end

# Calculate log(exp(x) + exp(y)).
# surprised this isn't in Julia base
function logaddexp(x::FloatingPoint, y::FloatingPoint)
    if x == y
        return x + log(2.)
    else
        tmp = x - y
        if tmp > zero(tmp)
            return x + log1p(exp(-tmp))
        elseif tmp <= 0.
            return y + log1p(exp(tmp))
        else
            return tmp  # NaNs
        end
    end
end

# Draw a random point from within a unit N-ball
function randnball(ndim)
    z = randn(ndim)
    r2 = 0.
    for i=1:ndim
        r2 += z[i]*z[i]
    end
    factor = rand()^(1./ndim) / sqrt(r2)
    for i=1:ndim
        z[i] *= factor
    end
    return z
end

# proportionality constant depending on dimension
# for n even:      (2pi)^(n    /2) / (2 * 4 * ... * n)
# for n odd :  2 * (2pi)^((n-1)/2) / (1 * 3 * ... * n)
function nball_vol_factor(ndim::Int)
    if ndim % 2 == 0
        c = 1.
        for i=2:2:ndim
            c *= 2pi / i
        end
        return c
    else
        c = 2.
        for i = 3:2:ndim
            c *= 2pi / i
        end
        return c
    end
end

function ellipsoid_volume(scaled_cov::Matrix{Float64})
    ndim = size(scaled_cov, 1)
    return nball_vol_factor(ndim) * sqrt(det(scaled_cov))
end

# find the bounding ellipsoid of points x where
function bounding_ellipsoid(x::Matrix{Float64}, enlarge=1.0)

    ndim, npoints = size(x)

    ctr = mean(x, 2)[:, 1]
    delta = x .- ctr
    cov = Base.unscaled_covzm(delta, 2)
    icov = inv(cov)

    # Calculate expansion factor necessary to bound each point.
    # This finds the maximum of (delta_i' * icov * delta_i)
    fmax = -Inf
    for k in 1:npoints
        f = 0.0
        for j=1:ndim
            for i=1:ndim
                f += icov[i, j] * delta[i, k] * delta[j, k]
            end
        end
        fmax = max(fmax, f)
    end

    fmax *= enlarge
    scale!(cov, fmax)
    scale!(icov, 1./fmax)
    vol = ellipsoid_volume(cov)

    return Ellipsoid(ctr, cov, icov, vol)
end

function sample_ellipsoid(ell::Ellipsoid)
    ndim = length(ell.ctr)

    # Get scaled eigenvectors (in columns): vs[:,i] is the i-th eigenvector.
    f = eigfact(ell.cov)
    v, w = f[:vectors], f[:values]
    for j=1:ndim
        tmp = sqrt(w[j])
        for i=1:ndim
            v[i, j] *= tmp
        end
    end

    return v*randnball(ndim) + ell.ctr
end

# nested sampling algorithm to evaluate Bayesian evidence.
function sample(loglikelihood::Function, prior::Function, ndim::Int;
                npoints::Int=100, enlarge::Float64=1.5, maxiter::Int=10000)

    # enlarge is volume enlargement factor
    enlarge_linear = enlarge^(1./ndim)

    # Choose initial points and calculate likelihoods
    points_u = rand(ndim, npoints)  # position of active pts in unit cube
    points_v = zeros(ndim, npoints)  # position of active pts in prior space
    points_logl = zeros(npoints)  # log(likelihood) at each point
    for i=1:npoints
        v = prior(sub(points_u, :, i))
        points_v[:, i] = v
        points_logl[i] = loglikelihood(v)
    end
    ncall = npoints  # number of likelihood calls we just made

    # Initialize values for nested sampling loop.
    samples_v = Float64[]    # stored objects for posterior results
    samples_logl = Float64[]
    samples_logprior = Float64[]
    samples_logwt = Float64[]
    loglstar = 0.  # ln(Likelihood constraint)
    h = 0.         # Information, initially 0.
    logz = -1.e300  # log(evidence Z, initially 0)

    # ln(width in prior mass), outermost width is 1 - e^(-1/n)
    logwidth = log(1. - exp(-1./npoints))  # log(prior volume), outermost
                                           # element is 1 - e^(-1/n)

    # Nested sampling loop.
    ndecl = 0
    logwt_old = -Inf
    niter = 0
    while niter < maxiter
        niter += 1

        if logz > -1.e6
            @printf "\riter=%6d logz=%8f" niter logz
        else
            @printf "\riter=%6d logz=" niter
        end
        flush(STDOUT)

        # find lowest logl in active points
        loglstar, minidx = findmin(points_logl)
        logwt = logwidth + loglstar

        # update evidence and information
        logz_new = logaddexp(logz, logwt)
        h = (exp(logwt-logz_new) * loglstar +
             exp(logz-logz_new) * (h+logz) - logz_new)
        logz = logz_new

        # Add worst object to samples.
        append!(samples_v, sub(points_v, :, minidx))
        push!(samples_logwt, logwt)
        push!(samples_logprior, logwidth)
        push!(samples_logl, loglstar)

        expected_vol = exp(-niter/npoints)

        # calculate the ellipsoid in prior space that contains all the
        # samples (including the worst one).
        ell = bounding_ellipsoid(points_u, enlarge_linear)

        # choose a point from within the ellipse until it has likelihood
        # better than loglstar
        while true
            u = sample_ellipsoid(ell)
            ok = true
            if any(u .< 0.) && any(u .> 1.)
                continue
            end
            v = prior(u)
            logltmp = loglikelihood(v)
            ncall += 1

            # Accept if and only if within likelihood constraint.
            if logltmp > loglstar
                points_u[:, minidx] = u
                points_v[:, minidx] = v
                points_logl[minidx] = logltmp
                break
            end
        end

        # Shrink interval
        logwidth -= 1./npoints

        # stop when logwt has been declining for more than 2*npoints
        # or niter/6 consecutive iterations.
        ndecl = (logwt < logwt_old) ? ndecl+1 : 0
        if (ndecl > 2*npoints) && (ndecl > niter/6)
            break
        end
        logwt_old = logwt
    end

    # Add remaining objects.
    # After N samples have been taken out, the remaining width is e^(-N/nobj)
    # The remaining width for each object is e^(-N/nobj) / nobj
    # The log of this for each object is:
    # log(e^(-N/nobj) / nobj) = -N/nobj - log(nobj)
    nsamples = div(length(samples_v), ndim)
    logwidth = -nsamples/npoints - log(npoints)
    for i in 1:npoints
        logwt = logwidth + points_logl[i]
        logz_new = logaddexp(logz, logwt)
        h = (exp(logwt - logz_new) * points_logl[i] +
             exp(logz - logz_new) * (h + logz) - logz_new)
        logz = logz_new

        append!(samples_v, sub(points_v, :, i))
        push!(samples_logwt, logwt)
        push!(samples_logl, points_logl[i])
        push!(samples_logprior, logwidth)
    end

    nsamples += npoints

    return ["niter" => niter,
            "ncall" => ncall,
            "logz" => logz,
            "logzerr" => sqrt(h/npoints),
            "loglmax" => maximum(points_logl),
            "h" => h,
            "samples" => reshape(samples_v, (nsamples, ndim)),
            "weights" => exp(samples_logwt .- logz),
            "logprior" => samples_logprior,
            "logl" => samples_logl]
end

end # module

ext: jl


text:
# In the 5 by 5 matrix below, the minimal path sum from the top left to the
# bottom right, by only moving to the right and down, is indicated in bold red
# and is equal to 2427.
#
# 131 673 234 103 18
# 201 96  342 965 150
# 630 803 746 422 111
# 537 699 497 121 956
# 805 732 524 37  331
#
# Find the minimal path sum from the top left to the bottom right by only moving
# right and down in matrix.txt (right click and "Save Link/Target As..."), a 31K
# text file containing an 80 by 80 matrix.

using ProjectEulerSolutions

# Simple dynamic programming solution to fill in a matrix starting with the
# corner and then picking the min path to each consecutive cell.
function p081solution(path::String="")::Integer

    if !isfile(path)
        if path != ""
            println("Cannot find path: '", path, "'")
        end
        return -1
    end

    # Read in matrix from file
    matrix = reduce(hcat, [parse.(Int32, split(line, ',')) for line in eachline(path)])
    n, m = size(matrix)

    best_path = zeros(Integer, size(matrix))
    best_path[1,1] = matrix[1,1]

    # Initialize the edges since they only have a single path.  Make use of the
    # built-in cumulative sum function
    best_path[:,1] = cumsum(matrix[:,1])
    best_path[1,:] = cumsum(matrix[1,:])

    # Very simple dynamic programming solution for the rest
    for i = 2:n
        for j = 2:m
            best_path[i,j] = min(best_path[i-1,j], best_path[i,j-1]) + matrix[i,j]
        end
    end

    return best_path[end,end]
end

p081 = Problems.Problem(p081solution)

input_path = "data/p081_matrix.txt"

Problems.benchmark(p081, input_path)

max_stars_repo_name: gnujosh/julia-euler-project

max_issues_repo_name: gnujosh/julia-euler-project

max_forks_repo_name: gnujosh/julia-euler-project


text:
using Plots, Printf
# pyplot()
do_viz = true

@views function heat_1D()
    # Physics
    lx  = 10.0
    λ   = 1.0
    ρCp = 1.0
    nt  = 200
    # Numerics
    nx  = 128
    # Derived numerics
    dx  = lx/nx
    dt  = dx^2/ρCp/λ/2.1
    xc  = LinRange(dx/2, lx-dx/2, nx)
    # Array allocation
    qx  = zeros(nx-1)
    # Initial condition
    T   = exp.(.-(xc.-lx./2.0).^2)
    # Time loop
    for it = 1:nt
        if (it==11) global t0 = Base.time() end
        # TODO add physics
    end
    time_s = (Base.time()-t0)
    @printf("Time = %1.4e s, T_eff = %1.2f GB/s \n", time_s, round((2/1e9*nx*sizeof(lx))/(time_s/(nt-10)), sigdigits=2))
    # Visualise
    if do_viz display(plot(xc, T, legend=false, framestyle=:box, xlabel="lx", ylabel="heat", title="diffusion")) end
    return
end

@time heat_1D()

max_stars_repo_name: luraess/geo-hpc-course

max_issues_repo_name: luraess/geo-hpc-course

max_forks_repo_name: luraess/geo-hpc-course


text:
# AUTO GENERATED FILE - DO NOT EDIT

export 'feffery'_antdaffix

"""
    'feffery'_antdaffix(;kwargs...)
    'feffery'_antdaffix(children::Any;kwargs...)
    'feffery'_antdaffix(children_maker::Function;kwargs...)

An AntdAffix component.

Keyword arguments:
- `children` (a list of or a singular dash component, string or number; optional)
- `id` (String; optional)
- `className` (String; optional)
- `loading_state` (optional): . loading_state has the following type: lists containing elements 'is_loading', 'prop_name', 'component_name'.
Those elements have the following types:
  - `is_loading` (Bool; optional): Determines if the component is loading or not
  - `prop_name` (String; optional): Holds which property is loading
  - `component_name` (String; optional): Holds the name of the component that is loading
- `offsetBottom` (Real; optional)
- `offsetTop` (Real; optional)
- `style` (Dict; optional)
- `target` (String; optional)
"""
function 'feffery'_antdaffix(; kwargs...)
        available_props = Symbol[:children, :id, :className, :loading_state, :offsetBottom, :offsetTop, :style, :target]
        wild_props = Symbol[]
        return Component("'feffery'_antdaffix", "AntdAffix", "feffery_antd_components", available_props, wild_props; kwargs...)
end

'feffery'_antdaffix(children::Any; kwargs...) = 'feffery'_antdaffix(;kwargs..., children = children)
'feffery'_antdaffix(children_maker::Function; kwargs...) = 'feffery'_antdaffix(children_maker(); kwargs...)

max_stars_repo_name: ReilYoung/feffery-antd-components

max_issues_repo_name: ReilYoung/feffery-antd-components

max_forks_repo_name: ReilYoung/feffery-antd-components


text:
function realsphharmbasis(Ω::AngularGrid, q::Vector{Point2{Int}})
    Float64[realsphharm(l,m,θ,ϕ) for (r,θ,ϕ) in cart2sph(Ω), (l,m) in q]
end

# precompute spherical harmonics basis to use in multipole expansion
function realsphharmbasis!(at::Atom, q::Vector{Point2{Int}})
    Ω = at[:g].Ω
    at[:qξ] = q
    at[:ξ]  = reshape(realsphharmbasis(Ω,q), (length(Ω),1,length(q)))
    return
end

realsphharmbasis!(at::Atom, lmax::Int) = realsphharmbasis!(at,enumeratelm(lmax))

function realsphharmbasis!(atoms::Atoms, q::Vector{Point2{Int}})
    for at in atoms
        realsphharmbasis!(at,q)
    end
end

realsphharmbasis!(atoms::Atoms, lmax::Int) = realsphharmbasis!(atoms,enumeratelm(lmax))

function expandmultipole(g    ::AtomicGrid,                         # atomic grid
                         f    ::AbstractArray,                      # function to expand
                         ξ    ::AbstractArray,                      # spherical harmonics basis
                         q    ::Vector{Point2{Int}};                # (l,m) quantum numbers
                         R    ::Point3{Float64} = (0.,0.,0.),       # multipole center
                         nspl ::Int             = 100,              # number of spline points
                         trim ::Bool            = true)             # trim negligible components
    n,m = length(g.r),length(q)
    f_mp = reshape(integrate(g.Ω, f.*ξ), (n,m))
    basis = MultipoleBasis()
    for (i,(l,m)) in enumerate(q)
        if !trim || sumabs2(f_mp[:,i]) > 1e-14
            push!(basis, MultipoleFunc(g.r,f_mp[:,i],l=l,m=m,R=R,npts=nspl))
        end
    end
    basis
end

function expandmultipole(at   ::Atom,                               # atomic grid
                         f    ::AbstractArray;                      # function to expand
                         nspl ::Int  = 100,                         # number of spline points
                         trim ::Bool = true)                        # trim negligible components
    expandmultipole(at[:g],f,at[:ξ],at[:qξ],R=at.R,nspl=nspl,trim=trim)
end

function expandmultipole(at   ::Atom,                               # atomic grid
                         f    ::AbstractVector;                     # function to expand
                         nspl ::Int  = 100,                         # number of spline points
                         trim ::Bool = true)                        # trim negligible components
    expandmultipole(at[:g],reshape(f,size(at[:g])),at[:ξ],at[:qξ],R=at.R,nspl=nspl,trim=trim)
end

function expandmultipole(g     ::MultiGrid,                    # molecular grid
                         rho   ::AbstractArray,                # electronic density
                         atoms ::Atoms;                        # atoms
                         nspl  ::Int = 100,                    # number of spline points
                         trim ::Bool = true)
    rho_mp = MultipoleBasis[]                                  # rho multipole expansion
    for (i,j) in enumerate(eachatom(g))
        rho_at = expandmultipole(atoms[i],g.p[j].*rho[j],nspl=nspl,trim=trim)
        push!(rho_mp, rho_at)
    end
    rho_mp
end

# function expandmultipole(m     ::Molecule;
#                          nspl  ::Int = 100,                    # number of spline points
#                          trim ::Bool = true)
#     expandmultipole(m.g,m[:rho],m.atoms,nspl=nspl,trim=trim)
# end

function poissonmultipole(rh   ::LogarithmicGrid,               # radial grid to solve poisson equation
                          rho  ::MultipoleFuncSpl;              # multipole component of density
                          R    ::Point3f = (0.,0.,0.),          # multipole center
                          nspl ::Int     = 100)                 # number of spline points
    l,m  = rho.l,rho.m
    rcut = rho.u.x[end]                                         # radial extent of splined density
    f  = rho.u(rh)                                              # evaluate radial part of multipole component
    vh = 4π/(2l+1)*radpoisson(rh,f,l)                           # multipole componenet of hartree potential
    q  = 4π/(2l+1)*multipolemoment(rh,f,l)                      # multipole moment
    MultipoleFunc(rh,vh,l=l,m=m,R=R,rcut=rcut,q=q,npts=nspl)
end

function poissonmultipole(rho  ::AbstractArray,
                          at   ::Atom;
                          nh   ::Int = 300,
                          nspl ::Int = 100)
    rho_mp = expandmultipole(at,rho,nspl=nspl)                  # multipole expansion of density
    rh     = loggrid(at[:g].r,nh)                               # radial grid to solve poisson equation
    vh_mp  = MultipoleFuncSpl[poissonmultipole(rh,rho_mp_lm,R=at.R,nspl=nspl) for rho_mp_lm in rho_mp]      # multipole expansion of hartree potential
    rho_mp, vh_mp
end

function poissonmultipole(g     ::MultiGrid,                    # molecular grid
                          rho   ::AbstractArray,                # electronic density
                          atoms ::Atoms;                        # atoms
                          nh    ::Int = 300,                    # number of hartree radial points
                          nspl  ::Int = 100)                    # number of spline points
    rho_mp = MultipoleBasis()                                   # rho multipole expansion for each atom
    vh_mp  = MultipoleBasis()                                   # vh multipole expansion for each atom
    for (i,j) in enumerate(eachatom(g))
        rho_at, vh_at = poissonmultipole(g.p[j].*rho[j],atoms[i],nh=nh,nspl=nspl)
        append!(rho_mp, rho_at)
        append!(vh_mp, vh_at)
    end
    rho_mp, vh_mp
end

max_stars_repo_name: johncsnyder/Element


text:
function init_game(img::Array{T,2}; nof_row_blocks::Int=4, nof_col_blocks::Int=4 ) where T  <: Colorant
    block_size = Int.(size(img) ./ (nof_row_blocks, nof_col_blocks))
    row_slices = [ i*block_size[1] + 1 : (i+1) * block_size[1] for i in 0:nof_row_blocks-1]
    col_slices = [ i*block_size[2] + 1 : (i+1) * block_size[2] for i in 0:nof_col_blocks-1]

    blocks = [(x,y) for y in col_slices for x in row_slices]

    shuf = zeros(eltype(img), size(img))
    shuf_blocks = shuffle(blocks)

    for b in 1:length(blocks)
        shuf[shuf_blocks[b][1], shuf_blocks[b][2]] = img[blocks[b][1], blocks[b][2]]
    end
    shuf
end


text:
using Diversity
using DataFrames
using EcoBase

"""
    hillnumber(proportions, qs)

Calculate the Hill number (or naive diversity) of order q of
population(s) with given relative proportions

# Arguments:

- `proportions`: relative proportions of different individuals / species
                 in population (vector, or matrix where columns are
                 individual populations)

- `qs`: single number or vector of orders of diversity measurement

# Returns:

- Diversity of order qs (single number or vector of diversities)
"""
function hillnumber(proportions, qs)
    hill = subdiv(NormalisedAlpha(Metacommunity(proportions)), qs)
    hill[!,:measure] .= "HillNumber"
    select!(hill, Not(:div_type))
    return hill
end

function hillnumber(asm::EcoBase.AbstractAssemblage, qs)
    hassimilarity(asm) && error("function cannot run with $(typeof(gettypes(asm))) types as contains similarity")
    return hillnumber(occurrences(asm), qs)
end


text:
using Test
using Dorothy

@DorothyAll()

@testset "Atoms" begin

	@testset "Guessing elements" begin
		m1 = readf("$(Dorothy.datapath)/1BTL.gro")
		@test inferelement("CA", "CA") == "Ca"
		@test inferelement("CL", "CL") == "Cl"
		@test inferelement("CA", "ASP") == "C"
		@test inferelement("N", "ASP") == "N"
		@test_throws Exception inferelement("", "ASP")
		inferelements!(m1)
		@test m1.elements[1:4] == ["N", "C", "C", "O"]
	end

	@testset "Guessing masses" begin
		m1 = readf("$(Dorothy.datapath)/1BTL.gro")
		@test infermass("MN1", "") == 0.0
		@test infermass("HG11", "H") == 1.008
		@test_throws Exception infermass("", "")
		infermasses!(m1)
		@test m1.masses[1:4] == [14.00, 12.01, 12.01, 15.99]
	end

end # @testset


text:
module LBLRTM

using Unitful
using PhysicalConstants.CODATA2018: h, ħ, k_B, c_0, π, g_n
using Printf
export run_lnfl, run_lblrtm, planck, radiance_to_bt

# Constants
include("constants.jl")

# Helper functions
include("helpers.jl")

# LNFL
include("run_lnfl.jl")

# LBLRTM
include("run_lblrtm.jl")

end


text:
"Intermediate representation (IR) for many-body propagators"
module SparseIR

using MultiFloats: Float64x2
using LinearAlgebra: dot, svd, SVD, QRIteration, Factorization, factorize, mul!, ldiv!
import LinearAlgebra: cond
import LinearAlgebra.BLAS: gemm!
using QuadGK: gauss, kronrod, quadgk
using SpecialFunctions: SpecialFunctions

Base.sinh(x::Float64x2) = setprecision(() -> Float64x2(sinh(big(x))), precision(Float64x2))
Base.cosh(x::Float64x2) = setprecision(() -> Float64x2(cosh(big(x))), precision(Float64x2))
Base.Math.hypot(x::Float64x2, y::Float64x2) = Base.Math._hypot(x, y) # FIXME: remove if MultiFloats is fixed

export fermion, boson
export DimensionlessBasis, FiniteTempBasis
export SparsePoleRepresentation, to_IR, from_IR
export overlap
export LegendreBasis, MatsubaraConstBasis
export FiniteTempBasisSet
export LogisticKernel, RegularizedBoseKernel
export CompositeBasis, CompositeBasisFunction, CompositeBasisFunctionFT
export TauSampling, MatsubaraSampling, TauSampling64, MatsubaraSampling64, evaluate, fit,
    evaluate!, fit!, workarrlengthfit

@enum Statistics boson fermion

include("_linalg.jl")
include("_roots.jl")
include("_specfuncs.jl")
using ._LinAlg: tsvd

include("svd.jl")
include("gauss.jl")
include("poly.jl")
include("kernel.jl")
include("basis.jl")
include("sve.jl")
include("augment.jl")
include("composite.jl")
include("sampling.jl")
include("spr.jl")
include("basis_set.jl")

# Precompile
precompile(FiniteTempBasis, (Statistics, Float64, Float64, Float64))
for cls in [:TauSampling, :MatsubaraSampling]
    for func in [:fit, :evaluate]
        for vartype in [:Float64, :ComplexF64]
            for dim in [:1, :2, :3, :4, :5, :6, :7]
                @eval precompile($(func), ($(cls), Array{$(vartype),$(dim)}))
            end
        end
    end
end

end # module


text:
include("magpoly.jl")

using .magpoly
const mp = magpoly
using DelimitedFiles, Random, Distributions, LinearAlgebra
#= This file contains the structures and functions to do carry
   out the inference using synthetic likelihood =#

#=struct summary_stats{T1<:Int, T2<:Float64}
    n_samples::T1
    magnetizations::T2
    spin_feature_overlaps::T2
    lag_p_autocorrelation::T2
end

function summary_stats(n_samples::Int)
    summary_stats(n_samples,zeros(n_samples),zeros(n_samples),zeros(n_samples))
end

#####################################################################################
# This is the structure for the parameter samples generated through synthetic likelihood
struct SL_data{T1<:Int, T2<:Float64}
    n_samples::T1
    n_params::T1
    param_values::T2
end

function SL_data(n_samples::Int, n_params::Int)
    SL_data(n_samples, n_params, zeros(n_params, n_samples))
end

#####################################################################################
#####################################################################################
function lag_p_autocovariance(vect::Array{Int}, p::Int)
    avg = 0.0
    avg_lag = 0.0
    var = 0.0
    var_lag = 0.0
    corr = 0.0
    n = length(vect)
    for i in 1:(n-p)
        avg += vect[i]
        avg_lag += vect[i+p]
        corr += vect[i]*vect[i+p]
    end
    avg = avg/(n-p)
    avg_lag = avg_lag/(n-p)
    corr = corr/(n-p)
    for i in 1:(n-p)
        var += (vect[i]-avg)^2
        var_lag += (vect[i+p]-avg_lag)^2
    end
    var = var/(n-p-1)
    var_lag = var_lag/(n-p-1)
    return (corr - avg*avg_lag)/(sqrt(var*var_lag))
end=#

function compute_summary_stats!(ss::Matrix{Float64}, sp_conf::Matrix{Int}, feats::Matrix{Float64})
    n_s = size(ss,2) #number of samples
    n_m = size(sp_conf,1)
    n_feats = size(feats, 2)
    for i_sample in 1:n_s
        for i_feat in 1:n_feats
            ss[i_feat,i_sample]=0
        end
        for i_m in 1:n_m
            for i_feat in 1:n_feats
                ss[i_feat,i_sample] += sp_conf[i_m,i_sample]*feats[i_m, i_feat]
            end
        end
        for i_feat in 1:n_feats
            ss[i_feat,i_sample] = ss[i_feat,i_sample]/n_m
        end
    end
end

function log_synth_likelihood(mat::Matrix{Float64}, avg::Array{Float64}, data::Array{Float64})
    s = -0.5*(avg .- data)'*inv(mat)*(avg .- data) -0.5*log(abs(det(mat)))
    return s
end

#####################################################################################
#####################################################################################

## for now this is just the backbone of the function to carry out the likelihood-free inference task
function synthetic_likelihood_polymer(n_samples::Int, sample_lag::Int, stride::Int, n_params::Int, delta_w::Float64,initial_weights::Array{Float64}, features_file::String, data_file::String)
    io = open(features_file, "r")
    features = readdlm(io,Float64)
    close(io)
    n_mono = size(features, 1)
    n_feats = size(features, 2)

    io = open(data_file, "r")
    data = readdlm(io,Float64; header=true)[1][1:end,:]
    close(io)
    println(data,"\n\n")
    n_data = size(data, 1)
    n_ss = size(data, 2) # number of summary statistics
    println(n_ss, "\n\n")

    accepted_moves = 0
    n_strides = cld(n_samples*sample_lag, stride) # integer ceiling of the division
    spins_coupling = 0.5
    inv_temps = [1.0, 0.9, 0.8, 0.74, 0.65, 0.62, 0.6, 0.55, 0.5, 0.4]
    n_temps = length(inv_temps)

    spins_configs = zeros(Int, n_mono, n_samples)

    polymers = Array{mp.Magnetic_polymer}(undef, n_temps)
    trajs = Array{mp.MC_data}(undef, n_temps)
    for i_temp in 1:n_temps
        polymers[i_temp] = mp.Magnetic_polymer(n_mono)#, inv_temps[i_temp], spins_coupling, alpha)
        trajs[i_temp] = mp.MC_data(n_strides*stride)
        if isfile("simulation_data/final_config_$(n_mono).txt")
            mp.initialize_poly!(polymers[i_temp],"simulation_data/final_config_$(n_mono).txt")
        else
            mp.initialize_poly!(polymers[i_temp])
        end
    end

    cov_mat = zeros(n_ss,n_ss)
    #inv_cov_mat = zeros(3,3)
    #delta_w = 0.05
    param_series = zeros(n_feats,n_params)
    SL_series = zeros(n_params)
    ss_mean_series = zeros(n_ss, n_params)
    #ss_cov_determinant_series = zeros(n_params)
    summary_stats = zeros(n_ss, n_samples)
    weights = zeros(n_feats)
    trial_weights = zeros(n_feats)
    weights .= initial_weights

    fields = zeros(n_mono)
    for i in 1:n_feats
        fields .+= features[:,i] .* weights[i]
    end

    #mp.set_fields!(polymers, fields)
    for i_eq in 1:20
        mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps, spins_coupling, fields)
    end

    mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps,spins_configs,sample_lag,n_samples, spins_coupling, fields)
    compute_summary_stats!(summary_stats, spins_configs, features)
    ss_mean_series[:,1] .= vec(mean(summary_stats, dims=2))
    cov_mat .= cov(summary_stats, dims=2)

    syn_like = 0.0
    for i_data in 1:n_data
        syn_like += log_synth_likelihood(cov_mat, ss_mean_series[:,1], data[i_data,:])
    end

    param_series[1,1] = weights[1]
    param_series[2,1] = weights[2]
    param_series[3,1] = weights[3]
    SL_series[1] = syn_like
    w_acceptance = 0.0

    for i_param in 2:n_params
        #=trial_weights .= weights
        if rand((1,2))==1
            trial_weights[1] += (2*rand()-1)*delta_w
        else
            trial_weights[2] += (2*rand()-1)*delta_w
        end=#

        trial_weights .= weights .+ (2 .* rand(n_feats) .-1) .*delta_w

        fields = zeros(n_mono)
        for i in 1:n_feats
            fields .+= features[:,i] .* trial_weights[i]
        end
        #mp.set_fields!(polymers, fields)

        mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps, spins_coupling, fields) # This is an short equilibration run so that the expectations are a bit better when changing weights
        mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps,spins_configs,sample_lag,n_samples, spins_coupling, fields)
        compute_summary_stats!(summary_stats, spins_configs, features)

        ss_mean_series[:,i_param] .= vec(mean(summary_stats, dims=2))
        cov_mat .= cov(summary_stats, dims=2)

        trial_syn_like = 0.0
        for i_data in 1:n_data
            trial_syn_like += log_synth_likelihood(cov_mat, ss_mean_series[:,i_param], data[i_data,:])
        end

        delta_syn_like = trial_syn_like -syn_like
        delta_syn_like>=0 ? w_acceptance=1 : w_acceptance=exp(delta_syn_like)
        if i_param%10 == 0
            println(i_param)
            println("w1: ",weights[1]," ---> ",trial_weights[1])
            println("w2: ",weights[2]," ---> ",trial_weights[2])
            println("w3: ",weights[3]," ---> ",trial_weights[3])
            println("delta_syn_like: ", delta_syn_like)
            println("acceptance: ", w_acceptance)
            println("avg acceptance: ", accepted_moves/i_param)
        end
        if w_acceptance > rand()
            weights .= trial_weights
            syn_like = trial_syn_like
            accepted_moves += 1
        end
        param_series[1,i_param] = weights[1]
        param_series[2,i_param] = weights[2]
        param_series[3,i_param] = weights[3]
        SL_series[i_param] = syn_like
    end
    println("Acceptance ratio: ", accepted_moves/n_params)

    !isdir("SL_data") && mkdir("SL_data")
    open("SL_data/weights.txt", "w") do io
        writedlm(io, param_series)
    end
    open("SL_data/syn_likes.txt", "w") do io
        writedlm(io, SL_series)
    end
    mp.write_results(polymers[1],trajs[1])

end

#####################################################################################
#####################################################################################

function amhi_polymer(n_samples::Int, sample_lag::Int, stride::Int, n_params::Int, delta_theta::Float64,initial_theta::Array{Float64}, features_file::String, data_file::String)
    # Read features from file
    io = open(features_file, "r")
    features = readdlm(io,Float64)
    close(io)
    n_mono = size(features, 1)
    n_feats = size(features, 2)

    # Read generated data. In this algo they're not summary stats but the full data!
    io = open(data_file,"r")
    data_spins = readdlm(io, Int64; header=true)[1][1:end,:]
    close(io)
    n_data = size(data_spins, 1)
    println(n_mono)
    println(n_data)
    println(typeof(data_spins))

    accepted_moves = 0
    n_strides = cld(n_samples*sample_lag, stride) # integer ceiling of the division
    #spins_coupling = 0.5

    inv_temps = [1.0, 0.9, 0.8, 0.74, 0.65, 0.62, 0.6, 0.55, 0.5, 0.4]
    n_temps = length(inv_temps)

    spins_configs = zeros(Int, n_mono, n_samples) # where you store the generated data
    ising_energies = zeros(n_samples)

    ########################################################### Initialize some quantitites in this section
    polymers = Array{mp.Magnetic_polymer}(undef, n_temps)
    trajs = Array{mp.MC_data}(undef, n_temps)
    for i_temp in 1:n_temps
        polymers[i_temp] = mp.Magnetic_polymer(n_mono)
        trajs[i_temp] = mp.MC_data(n_strides*stride)
        if isfile("simulation_data/final_config_$(n_mono).txt")
            mp.initialize_poly!(polymers[i_temp],"simulation_data/final_config_$(n_mono).txt")
        else
            mp.initialize_poly!(polymers[i_temp])
        end
    end

    param_series = zeros(n_feats+1,n_params)
    overlaps = zeros(n_feats+1, n_samples)
    avg_overlaps = zeros(n_feats+1)
    #cov_overlaps = zeros(n_feats+1, n_feats+1)

    theta = zeros(n_feats+1)    #from theta[1] to theta[n_feats] we have the regression weights, theta[n_feats+1] is the spins coupling
    trial_theta = zeros(n_feats+1)
    theta .= initial_theta
    if theta[n_feats+1] <= 0 || theta[n_feats+1] > 1
        theta[n_feats+1] = 0.5
    end
    ###
    #theta[n_feats+1] = 0.5
    ###

    fields = zeros(n_mono)
    for i in 1:n_feats
        fields .+= features[:,i] .* theta[i]
    end
    #############################################################

    for i_eq in 1:20
        mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps, theta[n_feats+1], fields) # theta[n_feats+1 is the coupling]
    end

    param_series[:,1] .= theta
    acceptance = 0.0

    energy = 0
    for i_data in 1:n_data
        for i_mono in 1:n_mono
            energy -= fields[i_mono]*data_spins[i_data,i_mono]
        end
    end

    energy_correction = 0.0

    ##
    theta_variation = zeros(n_feats+1)
    ##

    for i_param in 2:n_params

        theta_variation .= (2 .* rand(n_feats+1) .-1) .*delta_theta
        ###
        #theta_variation[n_feats+1] = 0.0
        ###
        trial_theta .= theta .+ theta_variation

        if trial_theta[n_feats+1]>0 && trial_theta[n_feats+1]<= 1.0
            trial_fields = zeros(n_mono)
            for i in 1:n_feats
                trial_fields .+= features[:,i] .* trial_theta[i]
            end

            trial_energy = 0
            for i_data in 1:n_data
                for i_mono in 1:n_mono
                    trial_energy -= trial_fields[i_mono]*data_spins[i_data,i_mono]
                end
            end

            ##################################################

            mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps, theta[n_feats+1], fields) # This is an short equilibration run so that the expectations are a bit better when changing weights
            mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps,spins_configs,ising_energies,sample_lag,n_samples, theta[n_feats+1], fields)

            fill!(overlaps, 0.0)
            for i_sample in 1:n_samples
                for i_feat in 1:n_feats
                    for i_mono in 1:n_mono
                        overlaps[i_feat, i_sample] += spins_configs[i_mono, i_sample]*features[i_mono, i_feat]
                    end
                end
                overlaps[n_feats+1, i_sample] = ising_energies[i_sample]
            end

            avg_overlaps = vec(mean(overlaps, dims=2))

            energy_correction = 0
            #linear correction
            for j in 1:n_feats+1
                energy_correction += theta_variation[j]*avg_overlaps[j]
            end
            energy_correction = energy_correction * n_data
            ##################################################

            energy_correction_J = 0.0
            for i_data in 1:n_data
                mp.set_spins!(polymers, data_spins[i_data,:])
                # here I run the MC's with the extra 'true' argument to quench the spins
                mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps, theta[n_feats+1], fields, true) # This is an short equilibration run so that the expectations are a bit better when changing weights
                mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps,spins_configs,ising_energies,sample_lag,n_samples, theta[n_feats+1], fields, true)
                energy_correction_J += mean(ising_energies) * theta_variation[n_feats+1]
            end

            ##################################################
            delta_acc = -(trial_energy - energy) - energy_correction + energy_correction_J
            delta_acc>=0 ? acceptance=1 : acceptance=exp(delta_acc)
            if i_param%10 == 0
                println(i_param)
                println("w1: ",theta[1]," ---> ",trial_theta[1])
                println("w2: ",theta[2]," ---> ",trial_theta[2])
                #println("w3: ",theta[3]," ---> ",trial_theta[3])
                println("J:  ",theta[3]," ---> ",trial_theta[3])
                println("delta_energy: ", -(trial_energy-energy))
                println("delta_acc: ", delta_acc)
                println("order of error: ", - sum(theta_variation.^2) *n_data)
                println("acceptance: ", acceptance)
                println("avg acceptance: ", accepted_moves/i_param)
            end
            if acceptance > rand()
                theta .= trial_theta
                energy = trial_energy
                #resample_needed = true
                fields .= trial_fields
                accepted_moves += 1
            end
        end

        param_series[1,i_param] = theta[1]
        param_series[2,i_param] = theta[2]
        param_series[3,i_param] = theta[3]
        #param_series[4,i_param] = theta[4]
    end
    println("Acceptance ratio: ", accepted_moves/n_params)

    !isdir("AMHI_data") && mkdir("AMHI_data")
    open("AMHI_data/thetas_$(n_data)amh$(n_samples)_$(delta_theta).txt", "w") do io
        writedlm(io, param_series)
    end

    mp.write_results(polymers[1],trajs[1])

end

#####################################################################################

# QAMHI: Approximate metropolis hastings inference with Quadratic corrections
function Qamhi_polymer(n_samples::Int, sample_lag::Int, stride::Int, n_params::Int, delta_theta::Float64,initial_theta::Array{Float64}, features_file::String, data_file::String)
    # Read features from file
    println("gaussian")
    io = open(features_file, "r")
    features = readdlm(io,Float64)
    close(io)
    n_mono = size(features, 1)
    n_feats = size(features, 2)
    println("Number of features: " ,n_feats)

    # Read generated data. In this algo they're not summary stats but the full data!
    io = open(data_file,"r")
    data_spins = readdlm(io, Int64; header=true)[1][1:end,:]
    close(io)
    n_data = size(data_spins, 1)
    println(n_mono)
    println(n_data)
    println(typeof(data_spins))

    accepted_moves = 0
    n_strides = cld(n_samples*sample_lag, stride) # integer ceiling of the division
    #spins_coupling = 0.5

    #inv_temps = [1.0, 0.9, 0.8, 0.74, 0.65, 0.62, 0.6, 0.55, 0.5, 0.4]
    inv_temps = [1.0, 0.96, 0.92, 0.89, 0.85, 0.83, 0.82, 0.79, 0.75, 0.67, 0.57, 0.5, 0.4]
    n_temps = length(inv_temps)

    spins_configs = zeros(Int, n_mono, n_samples) # where you store the generated data
    ising_energies = zeros(n_samples)

    ########################################################### Initialize some quantitites in this section
    polymers = Array{mp.Magnetic_polymer}(undef, n_temps)
    trajs = Array{mp.MC_data}(undef, n_temps)
    for i_temp in 1:n_temps
        polymers[i_temp] = mp.Magnetic_polymer(n_mono)#, inv_temps[i_temp], spins_coupling, alpha)
        trajs[i_temp] = mp.MC_data(n_strides*stride)
        #if isfile("simulation_data/final_config_$(n_mono).txt")
        #    mp.initialize_poly!(polymers[i_temp],"simulation_data/final_config_$(n_mono).txt")
        #else
        mp.initialize_poly!(polymers[i_temp])
        #end
    end

    polymers_2 = Array{mp.Magnetic_polymer}(undef, n_temps)
    trajs_2 = Array{mp.MC_data}(undef, n_temps)
    println("ok, different polymers")
    for i_temp in 1:n_temps
        polymers_2[i_temp] = mp.Magnetic_polymer(n_mono)#, inv_temps[i_temp], spins_coupling, alpha)
        trajs_2[i_temp] = mp.MC_data(n_strides*stride)
        #if isfile("simulation_data/final_config_$(n_mono).txt")
        #    mp.initialize_poly!(polymers_2[i_temp],"simulation_data/final_config_$(n_mono).txt")
        #else
        println("scratch initialization")
        mp.initialize_poly!(polymers_2[i_temp])
        #end
    end

    param_series = zeros(n_feats+1,n_params)
    #avg_spins = zeros(n_mono, n_samples)
    overlaps = zeros(n_feats+1, n_samples)
    avg_overlaps = zeros(n_feats+1)
    cov_overlaps = zeros(n_feats+1, n_feats+1)

    theta = zeros(n_feats+1)
    trial_theta = zeros(n_feats+1)
    theta .= initial_theta
    if theta[n_feats+1] <= 0
        theta[n_feats+1] = 0.5
    end

    fields = zeros(n_mono)
    for i in 1:n_feats
        fields .+= features[:,i] .* theta[i]
    end
    #############################################################

    mp.set_spins!(polymers_2, data_spins[1,:])
    for i_eq in 1:500
        mp.MMC_run!(polymers,  trajs,  n_strides,stride,inv_temps, theta[n_feats+1], fields)  #theta[n_feats+1] is the spin coupling
        mp.MMC_run!(polymers_2,trajs_2,n_strides,stride,inv_temps, theta[n_feats+1], fields, true)  #theta[n_feats+1] is the spin coupling
        if i_eq%10 == 0
            println(i_eq)
        end
    end

    param_series[:,1] .= theta
    acceptance = 0.0

    energy = 0
    for i_data in 1:n_data
        for i_mono in 1:n_mono
            energy -= fields[i_mono]*data_spins[i_data,i_mono]
        end
    end

    energy_correction = 0.0

    ##
    theta_variation = zeros(n_feats+1)
    ##

    for i_param in 2:n_params

        #theta_variation .= (2 .* rand(n_feats+1) .-1) .*delta_theta
        theta_variation .= randn(n_feats+1) .* delta_theta
        trial_theta .= theta .+ theta_variation

        if trial_theta[n_feats+1] > 0 #&& trial_theta[n_feats+1] <= 1.0
            trial_fields = zeros(n_mono)
            for i in 1:n_feats
                trial_fields .+= features[:,i] .* trial_theta[i]
            end

            trial_energy = 0
            for i_data in 1:n_data
                for i_mono in 1:n_mono
                    trial_energy -= trial_fields[i_mono]*data_spins[i_data,i_mono]
                end
            end

            mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps, theta[n_feats+1], fields) # This is an short equilibration run so that the expectations are a bit better when changing weights
            mp.MMC_run!(polymers,trajs,n_strides,stride,inv_temps,spins_configs,ising_energies, sample_lag,n_samples, theta[n_feats+1], fields)

            fill!(overlaps, 0.0)
            for i_sample in 1:n_samples
                for i_feat in 1:n_feats
                    for i_mono in 1:n_mono
                        overlaps[i_feat, i_sample] += spins_configs[i_mono, i_sample]*features[i_mono, i_feat]
                    end
                end
                overlaps[n_feats+1, i_sample] = ising_energies[i_sample]
            end

            #overlaps = features'*spins_configs ## shorter but maybe less efficient way to write stuff in the previous nested for loops
            avg_overlaps = vec(mean(overlaps, dims=2))
            cov_overlaps = cov(overlaps, dims=2)

            energy_correction = 0
            #linear correction
            for j in 1:n_feats+1
                energy_correction += theta_variation[j]*avg_overlaps[j]
            end
            #quadratic correction
            quadratic_correction = 0
            for j1 in 1:n_feats+1
                for j2 in 1:n_feats+1
                    quadratic_correction += theta_variation[j1]*cov_overlaps[j1,j2]*theta_variation[j2] * 0.5
                end
            end
            energy_correction += quadratic_correction
            energy_correction = energy_correction * n_data

            ##################################################

            energy_correction_J = 0.0
            for i_data in 1:n_data
                mp.set_spins!(polymers_2, data_spins[i_data,:])
                # here I run the MC's with the extra 'true' argument to quench the spins
                mp.MMC_run!(polymers_2,trajs_2,n_strides,stride,inv_temps, theta[n_feats+1], fields, true) # This is an short equilibration run so that the expectations are a bit better when changing weights
                mp.MMC_run!(polymers_2,trajs_2,n_strides,stride,inv_temps,spins_configs,ising_energies,sample_lag,n_samples, theta[n_feats+1], fields, true)
                energy_correction_J += mean(ising_energies) * theta_variation[n_feats+1] + 0.5 * theta_variation[n_feats+1]^2 * var(ising_energies)
            end

            ##################################################

            delta_acc = -(trial_energy - energy) - energy_correction + energy_correction_J
            delta_acc>=0 ? acceptance=1 : acceptance=exp(delta_acc)
            if i_param%10 == 0
                println(i_param)
                for i_feat in 1:n_feats
                    println("w$(i_feat): ",theta[i_feat]," ---> ",trial_theta[i_feat])
                end
                println("J : ",theta[n_feats+1]," ---> ",trial_theta[n_feats+1])
                println("delta_energy: ", -(trial_energy-energy))
                println("delta_acc: ", delta_acc)
                println("order of error: ", sum(theta_variation.^2)^1.5 *n_data)
                println("ratio quadratic_correction/tot_corr: ", quadratic_correction/energy_correction)
                println("ratio ising_correction/normal_correction: ", energy_correction_J/energy_correction)
                println("acceptance: ", acceptance)
                println("avg acceptance: ", accepted_moves/i_param)
            end
            if acceptance > rand()
                theta .= trial_theta
                energy = trial_energy
                fields .= trial_fields
                accepted_moves += 1
            end
        end

        for i_feat in 1:n_feats
            param_series[i_feat, i_param] = theta[i_feat]
        end
        param_series[n_feats+1, i_param] = theta[n_feats+1]

        if i_param%1000 == 0
            !isdir("AMHI_data") && mkdir("AMHI_data")
            open("AMHI_data/thetas_$(n_data)Qamh$(n_samples)_$(delta_theta).txt", "w") do io
                writedlm(io, param_series)
            end
        end
    end
    println("Acceptance ratio: ", accepted_moves/n_params)

    !isdir("AMHI_data") && mkdir("AMHI_data")
    open("AMHI_data/thetas_$(n_data)Qamh$(n_samples)_$(delta_theta).txt", "w") do io
        writedlm(io, param_series)
    end

    mp.write_results(polymers[1],trajs[1])

end

#####################################################################################
#####################################################################################

function generate_data(features_file::String, n_strides::Int, weights::Array{Float64})
    io = open(features_file, "r")
    features = readdlm(io,Float64)
    close(io)
    n_mono = size(features, 1)
    n_feats = size(features, 2)

    spins_coupling = 1.25
    #inv_temps = [10000.0, 0.9, 0.8, 0.74, 0.65, 0.62, 0.6, 0.55, 0.5, 0.4]
    #inv_temps = [0.01, 0.01]
    inv_temps = [1.0, 0.96, 0.92, 0.89, 0.85, 0.83, 0.82, 0.79, 0.75, 0.67, 0.57, 0.5, 0.4]
    n_temps = length(inv_temps)
    stride = 100

    polymers = Array{mp.Magnetic_polymer}(undef, n_temps)
    trajs = Array{mp.MC_data}(undef, n_temps)
    for i_temp in 1:n_temps
        polymers[i_temp] = mp.Magnetic_polymer(n_mono)#, inv_temps[i_temp], spins_coupling, alpha)
        trajs[i_temp] = mp.MC_data(n_strides*stride)
        #if isfile("simulation_data/final_config_$(n_mono).txt")
        #    mp.initialize_poly!(polymers[i_temp],"simulation_data/final_config_$(n_mono).txt")
        #    print("yes")
        #else
        mp.initialize_poly!(polymers[i_temp])
        #end
    end

    fields = zeros(n_mono)
    for i in 1:n_feats
        fields .+= features[:,i] .* weights[i]
    end
    #mp.set_fields!(polymers, fields)

    #here I run the simulation for a while in order to equilibrate the chain
    for i_burnin in 1:50
        println("burnin number: ", i_burnin)
        mp.MMC_run!(polymers, trajs, n_strides, stride, inv_temps, spins_coupling, fields)
    end

    n_data = 500
    summary_stats = zeros(n_feats,n_data)
    spins_conf = zeros(Int,n_mono,n_data)

    poly_confs = zeros(Int,n_data*3,n_mono)

    for i_data in 1:n_data
        println("Number of data point:  ", i_data)
        mp.MMC_run!(polymers, trajs, n_strides, stride, inv_temps, spins_coupling, fields)
        #mp.write_results(polymers[1],trajs[1])
        spins_conf[:,i_data] .= polymers[1].spins
        for j in 1:3
            poly_confs[(i_data-1)*3+j,:] .= polymers[1].coord[j,:]
        end
        if i_data%10 == 0
            open("saw_conf_data_prova.txt","w") do io
                writedlm(io,transpose([weights; spins_coupling]))
                writedlm(io,transpose(spins_conf))
            end
            open("poly_confs_prova.txt","w") do io
                writedlm(io,transpose([weights; spins_coupling]))
                writedlm(io, poly_confs)
            end
        end
    end
    #compute_summary_stats!(summary_stats,spins_conf,features)
    #println(summary_stats)
    #println(spins_conf)
    open("saw_conf_data_prova.txt","w") do io
        writedlm(io,transpose([weights; spins_coupling]))
        writedlm(io,transpose(spins_conf))
    end
    #=open("data_file.txt","w") do io
        writedlm(io,transpose(weights))
        writedlm(io,transpose(summary_stats))
    end=#
    open("poly_confs_prova.txt","w") do io
        writedlm(io,transpose([weights; spins_coupling]))
        writedlm(io, poly_confs)
    end

    mp.write_results(polymers[1])
end


text:
#demos for SplineModel.jl module
include("../src/DemoGeometries.jl")
plotly()
#pyplot()
# ###### demo for nrbtestsrf()
# Plots.surface(reuse = false)
# srf = nrbtestsrf()
# plt = nrbplot(srf, [20, 30])
# display(plt)
#
# ##### demo for nrbtestcrv()
# Plots.surface(reuse = false)
# crv = nrbtestcrv()
# #plt = nrbplot(crv, [100])
# plt = nrbctrlplot(crv)
# display(plt)
#
# icrv = nrbkntins(crv, [[0.125, 0.375, 0.625, 0.875]])
# Plots.surface(reuse = false)
# plt =nrbctrlplot(icrv)
# display(plt)
#
# icrv = nrbdegelev(crv, [1])
# Plots.surface(reuse = false)
# plt =nrbctrlplot(icrv)
# display(plt)
#
#
# ###### demo for nrbtestcrv3d()
# Plots.surface(reuse = false)
# crv3d = nrbtestcrv3d()
# #plt = nrbplot(crv3d, [100])
# plt = nrbctrlplot(crv3d)
# display(plt)
#
# ##### demo for nrbtestvol()
# Plots.surface(reuse = false)
# horseshoe = nrbtestvol()
# plt = nrbplot(horseshoe, [6, 6, 50])
# display(plt)
#
# ##### demo for nrbline()
# Plots.surface(reuse = false)
# crv = nrbline([0.0 0.0 0.0]',[5.0 4.0 2.0]')
# #crv2 = nrbline()
# plt = nrbplot(crv, [1])
# #plt = nrbplot(crv2, [1])
# display(plt)
#
# ##### demo for nrbrect(w,h)
# Plots.surface(reuse = false)
# crv = nrbrect(2,1)
# plt = nrbplot(crv, [4])
# display(plt)
#
#
# ##### demo for nrbcirc()
# Plots.surface(reuse = false)
# for r=1:9
#     global crv = nrbcirc(r, [], 45*pi/180, 315*pi/180)
#     #global plt=nrbplot(crv, [50])
#     global plt=nrbkntplot(crv);
# end
# display(plt)
#
##### demo for nrbkntplot()
Plots.surface(reuse = false)
crv = nrbtestcrv()
plt = nrbkntplot(crv)
display(plt)

##### demo for nrbkntplot()
Plots.surface(reuse = false)
crv = nrbtestcrv3d()
plt = nrbkntplot(crv)
display(plt)

###### demo for nrbtestsrf()
Plots.surface(reuse = false)
srf = nrbtestsrf()
#plt = nrbkntplot(srf)
plt = nrbctrlplot(srf)
display(plt)
#
# ##### demo for nrbtestvol()
# Plots.surface(reuse = false)
# horseshoe = nrbtestvol()
# #plt = nrbkntplot(horseshoe)
# plt = nrbctrlplot(horseshoe)
# display(plt)
#
#
# ######### demo for nrbruled()
# Plots.surface(reuse = false)
# crv1 = nrbtestcrv()
# crv2 = nrbtform(nrbcirc(4,[4.5,0],pi,0.0), vectrans([0.0, 4.0, -4.0]));
# srf = nrbruled(crv1, crv2);
# plt = nrbplot(srf,[40 20]);
# #plt = nrbctrlplot(srf);
# display(plt)
# # title ('Ruled surface construction from two NURBS curves.');
# # hold off
#
# ######### demo for nrb4surf()
# Plots.surface(reuse = false)
# srf = nrb4surf([0.0, 0.0, 0.5],[1.0, 0.0, -0.5],[0.0, 1.0, -0.5],[1.0, 1.0, 0.5]);
# plt = nrbplot(srf,[10,10]);
# display(plt)
#
#
# ##### demo for nrbtransp()
# Plots.surface(reuse = false)
# srf = nrb4surf([0 0 0], [1 0 1], [0 1 1], [1 1 2]);
# plt = nrbplot(srf,[20 5]);
# coefs = srf.coefs
# coefs[3,:,:] = coefs[3,:,:] .+ 10
# srf = nrbmak(coefs, srf.knots);
# srf = nrbtransp(srf);
# plt = nrbplot(srf,[20 5]);
# display(plt)
#
# ##### demo for nrbcoons()
# Plots.surface(reuse = false)
# pnts = [ 0.0  3.0  4.5  6.5 8.0 10.0;
#          0.0  0.0  0.0  0.0 0.0  0.0;
#          2.0  2.0  7.0  4.0 7.0  9.0];
# crv1 = nrbmak(pnts, [[0, 0, 0, 1/3, 0.5, 2/3, 1, 1, 1]]);
#
# pnts= [ 0.0  3.0  5.0  8.0 10.0;
#         10.0 10.0 10.0 10.0 10.0;
#         3.0  5.0  8.0  6.0 10.0];
# crv2 = nrbmak(pnts, [[0, 0, 0, 1/3, 2/3, 1, 1, 1]]);
#
# pnts= [ 0.0 0.0 0.0 0.0;
#         0.0 3.0 8.0 10.0;
#         2.0 0.0 5.0 3.0];
# crv3 = nrbmak(pnts, [[0, 0, 0, 0.5, 1, 1, 1]]);
#
# pnts= [ 10.0 10.0 10.0 10.0 10.0;
#         0.0   3.0  5.0  8.0 10.0;
#         9.0   7.0  7.0 10.0 10.0];
# crv4 = nrbmak(pnts, [[0, 0, 0, 0.25, 0.75, 1, 1, 1]]);
# srf = nrbcoons(crv1, crv2, crv3, crv4);
#
# plt = nrbplot(srf,[20 20]);
# display(plt)

#### demo for nrbrevolve()
Plots.surface(reuse = false)
sphere = nrbrevolve(nrbcirc(1,[],0.0,pi),[0.0, 0.0, 0.0],[1.0, 0.0, 0.0]);
plt = nrbplot(sphere,[40, 40]);
torus = nrbrevolve(nrbcirc(0.2,[0.9, 1.0]),[0.0, 0.0, 0.0],[1.0, 0.0, 0.0]);
plt = nrbplot(torus,[40, 40]);
plt = nrbplot(nrbtform(torus,vectrans([-1.8])),[20, 10]);
display(plt)

#### demo for nrbrevolve()
Plots.surface(reuse = false)
pnts = [3.0 5.5 5.5 1.5 1.5 4.0 4.5;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.5 1.5 4.5 3.0 7.5 6.0 8.5]
crv = nrbmak(pnts,[[0, 0, 0, 1/4, 1/2, 3/4, 3/4, 1, 1, 1]])
xx = vecrotz(25*pi/180)*vecroty(15*pi/180)*vecrotx(20*pi/180)
nrb = nrbtform(crv,vectrans([5, 5])*xx)
pnt = [5, 5, 0]
vect = xx*[0, 0, 1, 1]
srf = nrbrevolve(nrb,pnt,vect[1:3])
plt = nrbplot(srf, [40, 40])
display(plt)

#### demo for nrbrevolve()
Plots.surface(reuse = false)
crv1 = nrbcirc(1,[0, 0],0, pi/2);
crv2 = nrbcirc(2,[0, 0],0, pi/2);
srf = nrbruled(crv1, crv2);
srf = nrbtform(srf, [1 0 0 0; 0 1 0 1; 0 0 1 0; 0 0 0 1]);
vol = nrbrevolve(srf, [0, 0, 0], [1, 0, 0], pi/2);
plt = nrbplot(vol, [30, 30, 30])
display(plt)


text:
using QuantLib.Time

mutable struct FixedRateCoupon{DC <: DayCount, IR <: InterestRate} <: Coupon
  couponMixin::CouponMixin{DC}
  paymentDate::Date
  nominal::Float64
  rate::IR
end

FixedRateCoupon(paymentDate::Date,
                faceAmount::Float64,
                rate::IR,
                accrual_start::Date,
                accrual_end::Date,
                ref_start::Date,
                ref_end::Date,
                dc::DC,
                accrual::Float64) where {DC <: DayCount, IR <: InterestRate} =
                FixedRateCoupon{DC, IR}(CouponMixin{DC}(accrual_start, accrual_end, ref_start, ref_end, dc, accrual), paymentDate, faceAmount, rate)

## COUPON METHODS ##
amount(coup::FixedRateCoupon) =
        coup.nominal * (compound_factor(coup.rate, accrual_start_date(coup), accrual_end_date(coup), ref_period_start(coup), ref_period_end(coup)) - 1)

calc_rate(coup::FixedRateCoupon) = coup.rate.rate

mutable struct FixedRateLeg{FRC <: FixedRateCoupon} <: Leg
  coupons::Vector{FRC}
  redemption::Union{SimpleCashFlow, Nothing}
end

function FixedRateLeg(schedule::Schedule,
                      faceAmount::Float64,
                      rate::Float64,
                      calendar::BusinessCalendar,
                      paymentConvention::BusinessDayConvention,
                      dc::DayCount;
                      add_redemption::Bool = true)
  # n = add_redemption ? length(schedule.dates) : length(schedule.dates) - 1
  # coups = Vector{Union{FixedRateCoupon, SimpleCashFlow}}(n)
  #
  # start_date = schedule.dates[1]
  # end_date = schedule.dates[2]
  # payment_date = adjust(calendar, paymentConvention, end_date)
  # #TODO: setup payment adjustments and the like
  # coups[1] = FixedRateCoupon(payment_date, faceAmount, InterestRate(rate, dc, SimpleCompounding(), schedule.tenor.freq), start_date, end_date, start_date, end_date, dc, -1.0)
  #
  # # build coupons
  # count = 2
  # start_date = end_date
  # end_date = count == length(schedule.dates) ? schedule.dates[end] : schedule.dates[count + 1]
  # payment_date = adjust(calendar, paymentConvention, end_date)
  # while start_date < schedule.dates[end]
  #   @inbounds coups[count] = FixedRateCoupon(payment_date, faceAmount, InterestRate(rate, dc, SimpleCompounding(), schedule.tenor.freq), start_date, end_date, start_date, end_date, dc, -1.0)
  #
  #   count += 1
  #   start_date = end_date
  #   end_date = count == length(schedule.dates) ? schedule.dates[end] : schedule.dates[count + 1]
  #   payment_date = adjust(calendar, paymentConvention, end_date)
  # end
  #
  # if add_redemption
  #   @inbounds coups[end] = SimpleCashFlow(faceAmount, end_date)
  # end

  ratesLen = length(schedule.dates) - 1

  ratesVec = fill(rate, ratesLen)

  FixedRateLeg(schedule, faceAmount, ratesVec, calendar, paymentConvention, dc; add_redemption = add_redemption)
end

function FixedRateLeg(schedule::Schedule,
                      faceAmount::Float64,
                      rates::Vector{Float64},
                      calendar::BusinessCalendar,
                      paymentConvention::BusinessDayConvention,
                      dc::DC;
                      add_redemption::Bool = false) where {DC <: DayCount}
  n = length(schedule.dates) - 1
  length(rates) == length(schedule.dates) - 1 || error("mismatch in coupon rates")

  coup_type = FixedRateCoupon{DC, InterestRate{DC, SimpleCompounding, typeof(schedule.tenor.freq)}}
  coups = Vector{coup_type}(undef, n)

  start_date = schedule.dates[1]
  end_date = schedule.dates[2]
  payment_date = adjust(calendar, paymentConvention, end_date)
  #TODO: setup payment adjustments and the like
  coups[1] = FixedRateCoupon(payment_date, faceAmount, InterestRate(rates[1], dc, SimpleCompounding(), schedule.tenor.freq), start_date, end_date, start_date, end_date, dc, -1.0)

  # build coupons
  count = 2
  start_date = end_date
  end_date = count == length(schedule.dates) ? schedule.dates[end] : schedule.dates[count + 1]
  payment_date = adjust(calendar, paymentConvention, end_date)
  while start_date < schedule.dates[end]
    @inbounds coups[count] = FixedRateCoupon(payment_date, faceAmount, InterestRate(rates[count], dc, SimpleCompounding(), schedule.tenor.freq), start_date, end_date, start_date, end_date, dc, -1.0)

    count += 1
    start_date = end_date
    end_date = count == length(schedule.dates) ? schedule.dates[end] : schedule.dates[count + 1]
    payment_date = adjust(calendar, paymentConvention, end_date)
  end

  if add_redemption
    redempt = SimpleCashFlow(faceAmount, end_date)
  else
    redempt = nothing
  end

  return FixedRateLeg{coup_type}(coups, redempt)
end

# Coupon methods
get_pay_dates(coups::Vector{F}) where {F <: FixedRateCoupon} = Date[date(coup) for coup in coups]
get_reset_dates(coups::Vector{F}) where {F <: FixedRateCoupon} = Date[accrual_start_date(coup) for coup in coups]

function accrued_amount(coup::FixedRateCoupon, settlement_date::Date)
  if settlement_date <= accrual_start_date(coup) || settlement_date > coup.paymentDate
    return 0.0
  end

  return coup.nominal *
      (compound_factor(coup.rate, accrual_start_date(coup), min(settlement_date, accrual_end_date(coup)), ref_period_start(coup), ref_period_end(coup)) - 1.0)
end


text:
module TestFolds

import LoadAllPackages
using Test
using Distributed: addprocs, nworkers
using Test

if get(ENV, "CI", "false") == "true"
    addprocs(3)

    # Tests in `PerformanceTestTools.@include_foreach` might cause
    # pre-compilation errors as two processes try to compile packages
    # at the same time.  This can happen when the tests are run via
    # `Pkg.test`.  Doing this after `addprocs` to workaround a quirk
    # in Distributed.jl.
    LoadAllPackages.loadall()
end
@info "Testing with:" nworkers()

@testset "$file" for file in sort([
    file for file in readdir(@__DIR__) if match(r"^test_.*\.jl$", file) !== nothing
])
    if file == "test_doctest.jl"
        if lowercase(get(ENV, "JULIA_PKGEVAL", "false")) == "true"
            @info "Skipping doctests on PkgEval."
            continue
        elseif VERSION >= v"1.6-"
            @info "Skipping doctests on Julia $VERSION."
            continue
        end
    end

    include(file)
end

end  # module


text:
function position_periodic(x::Float64, L::Float64)
	if x > L
		x = x - L
	elseif x < 0.0
		x = x + L
	end

	return x
end


text:
include("calc_line.jl")

function test_CN_parallel()

  start_dir = pwd()
  @testset "---- Crank-Nicolson Convergence Tests, PETSc + CS Jacobian -----" begin

    cd(dirname(@__FILE__))
    cd("./m1")
    fname = "input_vals1.jl"
    mesh, sbp, eqn, opts = solvePDE(fname)

    cd("../m2")
    fname = "input_vals1.jl"
    mesh, sbp, eqn, opts = solvePDE(fname)

#    cd("../m3")
#    fname = "input_vals1.jl"
#    mesh, sbp, eqn, opts = solvePDE(fname)

    cd("..")
#    include(joinpath(pwd(), "calc_line.jl"))

    MPI.Barrier(mesh.comm)
    slope = calc_line()
    # println("slope = ", slope)

    MPI.Barrier(mesh.comm)

    data = readdlm("err_data.dat")
    err_vals = data[:, 2]
    #println("err_vals = ", err_vals)

    slope_val = 2.00
    slope_margin = 0.1

    @test  slope  > slope_val - slope_margin
    @test  slope  < slope_val + slope_margin

    err_val = 0.09095728504176116
    slope_fac = 1.25
    # println("err_vals[1] = ", err_vals[1])
    @test  err_vals[1]  > err_val/slope_fac
    @test  err_vals[1]  < err_val*slope_fac

  end

  @testset "---- Testing restart -----" begin
    cd("./m1")
    mesh, sbp, eqn, opts = solvePDE("input_vals_restart")
    MPI.Barrier(mesh.comm)
    data = readdlm("error_calc.dat")
    datas = readdlm("../../crank_nicolson_PETSc_serial/m1/error_calc.dat")

    @test isapprox( data[1], datas[1]) atol=1e-13
    @test isapprox( data[2], datas[2]) atol=1e-13

    cd("..")

  end

  @testset "--- Testing matrix-free Petsc -----" begin
    cd("./m1")
    opts = read_input("input_vals1.jl")

    # clear out the old options
#    PetscClearOptions(opts["petsc_options"])
    opts["petsc_options"] = Dict{AbstractString, AbstractString}(
                                "-pc_type" => "none",
                                "-malloc" => "",
                                "-malloc_debug" => "",
                                "-ksp_monitor" => "",
                                "-ksp_gmres_modifiedgramschmidt" => "",
                                "-ksp_gmres_restart" => "30",
                                )

    mesh, sbp, eqn, opts = solvePDE("input_vals_restart")
    MPI.Barrier(mesh.comm)
    data = readdlm("error_calc.dat")
    datas = readdlm("../../crank_nicolson_PETSc_serial/m1/error_calc.dat")

    @test isapprox( data[1], datas[1]) atol=1e-13
    @test isapprox( data[2], datas[2]) atol=1e-13

    cd("..")

  end

  cd(start_dir)

  return nothing
end

add_func1!(AdvectionTests, test_CN_parallel, [TAG_SHORTTEST])


text:
using Test
using MathOptInterface
import Cbc

const MOI = MathOptInterface

const OPTIMIZER = Cbc.Optimizer()
MOI.set(OPTIMIZER, MOI.Silent(), true)

@testset "SolverName" begin
    @test MOI.get(OPTIMIZER, MOI.SolverName()) == "COIN Branch-and-Cut (Cbc)"
end

@testset "supports_default_copy_to" begin
    @test !MOI.Utilities.supports_allocate_load(OPTIMIZER, false)
    @test !MOI.Utilities.supports_allocate_load(OPTIMIZER, true)
    @test !MOI.Utilities.supports_default_copy_to(OPTIMIZER, false)
    @test !MOI.Utilities.supports_default_copy_to(OPTIMIZER, true)
end

const CACHED = MOI.Utilities.CachingOptimizer(
    MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
    OPTIMIZER
)
const BRIDGED = MOI.Bridges.full_bridge_optimizer(CACHED, Float64)

const CONFIG = MOI.Test.TestConfig(duals = false, infeas_certificates = false)

@testset "basic_constraint_tests" begin
    MOI.Test.basic_constraint_tests(CACHED, CONFIG)
end

@testset "Unit" begin
    MOI.Test.unittest(BRIDGED, CONFIG, [
        # TODO(odow): implement attributes:
        "number_threads",

        # INFEASIBLE_OR_UNBOUNDED instead of DUAL_INFEASIBLE
        "solve_unbounded_model",

        # No quadratics
        "delete_soc_variables",
        "solve_qcp_edge_cases",
        "solve_qp_edge_cases",
    ])
end

@testset "ModelLike" begin
    @test MOI.get(CACHED, MOI.SolverName()) == "COIN Branch-and-Cut (Cbc)"
    @testset "default_objective_test" begin
         MOI.Test.default_objective_test(CACHED)
     end
     @testset "default_status_test" begin
         MOI.Test.default_status_test(CACHED)
     end
    @testset "nametest" begin
        MOI.Test.nametest(CACHED)
    end
    @testset "validtest" begin
        MOI.Test.validtest(CACHED)
    end
    @testset "emptytest" begin
        MOI.Test.emptytest(BRIDGED)
    end
    @testset "orderedindicestest" begin
        MOI.Test.orderedindicestest(CACHED)
    end
end

@testset "Continuous Linear" begin
    MOI.Test.contlineartest(BRIDGED, CONFIG)
end

@testset "Integer Linear" begin
    MOI.Test.intlineartest(BRIDGED, CONFIG, [
        # Cbc does not support indicator constraints.
        "indicator1", "indicator2", "indicator3", "indicator4",
        # TODO(odow): needs MOI at least 0.9.14.
        "semiconttest", "semiinttest"
    ])
end

@testset "Test params" begin
    # Note: we generate a non-trivial problem to ensure that Cbc struggles to
    # find a solution at the root node.
    knapsack_model = MOI.Utilities.Model{Float64}()
    N = 100
    x = MOI.add_variables(knapsack_model, N)
    MOI.add_constraint.(knapsack_model, MOI.SingleVariable.(x), MOI.ZeroOne())
    MOI.add_constraint(
        knapsack_model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1 + sin(i) for i = 1:N], x),
            0.0
        ),
        MOI.LessThan(10.0),
    )
    MOI.set(
        knapsack_model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([cos(i) for i = 1:N], x),
            0.0
        )
    )
    model = Cbc.Optimizer(
        maxSol = 1, presolve = "off", cuts = "off", heur = "off", logLevel = 0
    )
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.SOLUTION_LIMIT
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    MOI.empty!(model)
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.SOLUTION_LIMIT
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
end

@testset "Test PrimalStatus" begin
    model = MOI.Utilities.Model{Float64}()
    x = MOI.add_variable(model)
    MOI.add_constraint(model, x, MOI.GreaterThan(1.0))
    MOI.add_constraint(model, x, MOI.LessThan(0.0))
    cbc = Cbc.Optimizer()
    MOI.copy_to(cbc, model)
    MOI.optimize!(cbc)
    MOI.get(cbc, MOI.PrimalStatus()) == MOI.NO_SOLUTION
end


text:
using FastRef
using LazySets
import FastRef: forward_network, forward_affine_map

center = fill(1.0, 10)
radius = fill(1.0, 10)
inputSet = Hyperrectangle(center, radius)

nnet = read_nnet("nnet/test10.nnet")
solver1 = MaxSens(2.0, true)
solver2 = FastGrid(2.0)

(W, b) = (nnet.layers[1].weights, nnet.layers[1].bias)

outputSet1 = forward_network(solver1, nnet, inputSet)

z = forward_affine_map(solver2, W, b, inputSet)
outputSet2 = forward_network(solver2, nnet, z)

println(outputSet1)
println(outputSet2)


text:
#################################################################################
# Comparison of Zig-Zag for diffusion bridges with tailored Poisson rates       #
# and with adapted Poisson rate. Reference: https://arxiv.org/abs/2001.05889.   #
#################################################################################

using ZigZagBoomerang, SparseArrays, LinearAlgebra
#using CairoMakie
include("../../src/faberschauder.jl")
include("../../src/asvar.jl")
include("../../src/pdmp.jl")

const ZZB = ZigZagBoomerang
using CSV
using DataFrames
using Statistics
# Drift
b(x) = α * sin(x)
# First derivative
b′(x) = α * cos(x)
# Second derivative
b″(x) = -α * sin(x)

####################################################################
# Overloading Poisson times in order to have tighter upperbounds   #
####################################################################
struct MyBound
    c::Float64
end
function ZZB.adapt!(b::Vector{MyBound}, i, x)
    b[i] = MyBound(b[i].c * x)
end

"""
    poisson_time(a, b, c, u)
Obtaining waiting time for inhomogeneous Poisson Process
with rate of the form λ(t) = a + (b + c*t)^+, where `c`,`a`> 0 ,`b` ∈ R, `u` uniform random variable
"""
function ZZB.poisson_time((a, b, c)::NTuple{3}, u = rand()) # formula (22)
    if b > 0
        return (-(b + a) + sqrt((b + a)^2 - 2.0 * c * log(u))) / c # positive solution of quadratic equation c*0.5 x^2 + (b + a) x + log(u) = 0
    elseif a * b / c <= log(u)
        return -log(u) / a
    else
        return (-(a + b) + sqrt((a + b)^2 - 2.0 * c * (b * b * 0.5 / c + log(u)))) / c    # positive solution of quadratic equation c*0.5 x^2 + (b + a) x + log(u) + b*b*0.5/c = 0
    end
end

ZZB.sλ̄((a, b, c)::NTuple{3}, Δt) = a + ZZB.pos(b + c * Δt)
"""
    abc(G, i, x, θ, c, Flow)
Returns the constant term `a` and linear term `b` when computing the Poisson times
from the upper upper bounding rates λᵢ(t) = max(a + b*t)^2. The factors `a` and `b`
can be function of the current position `x`, velocity `θ`, tuning parameter `c` and
the Graph `G`
"""
function ZZB.ab(G, i, x, θ, c::Vector{MyBound}, F::ZigZag)
    if i == 1
        a = c[i].c + T^(1.5)*0.5*(α^2 + α) * abs(θ[i])  # initial point
        b1 = θ[i]*(x[i] - x[end])
        b2 = θ[i]*(θ[i] - θ[end])
    elseif i == (2 << L) + 1
        a = c[i].c + T^(1.5)*0.5*(α^2 + α) * abs(θ[i])  # final point
        b1 = θ[i]*(x[i] - x[1])
        b2 = θ[i]*(θ[i] - θ[1])
    else
        l = lvl(i, L)
        a = c[i].c + T^(1.5) / 2^((L - l) * 1.5 + 2) * (α^2 + α) * abs(θ[i]) # formula (22)
        b1 = x[i] * θ[i]
        b2 = θ[i] * θ[i]
    end
    return a, b1, b2
end

function ∇ϕmoving_var1(t, ξ, θ, i, t′, F, L, T) # formula (17)
    if i == (2 << L) + 1    # final point
        s = T * (rand())
        x = dotψmoving(t, ξ, θ, t′, s, F, L, T)
        return -b(ξ[end]*sqrt(T))*sqrt(T) + 0.5 * sqrt(T) * s * (2b(x) * b′(x) + b″(x)) + ξ[i] - ξ[1]
    elseif i == 1   # initial point
        s = T * (rand())
        x = dotψmoving(t, ξ, θ, t′, s, F, L, T)
        return 0.5 * T^(1.5) * (1 - s / T) * (2b(x) * b′(x) + b″(x)) + ξ[1]
    else
        l = lvl(i, L)
        k = (i - 1) ÷ (2 << l)
        δ = T / (1 << (L - l))
        res = 0.0
        ave_fact = 1/(l+1)
        for _ in 1:l+1
            s = δ * (k + rand())
            x = dotψmoving(t, ξ, θ, t′, s, F, L, T)
           res += ave_fact * 0.5 * δ * Λ(s, L - l, T) * (2b(x) * b′(x) + b″(x))
       end
       return res + ξ[i]
    end
end

function ∇ϕmoving_var2(t, ξ, θ, i, t′, F, L, T) # formula (17)
    if i == (2 << L) + 1    # final point
        s = T * (rand())
        x = dotψmoving(t, ξ, θ, t′, s, F, L, T)
        return -b(ξ[end]*sqrt(T))*sqrt(T) + 0.5 * sqrt(T) * s * (2b(x) * b′(x) + b″(x)) + ξ[i] - ξ[1]
    elseif i == 1   # initial point
        s = T * (rand())
        x = dotψmoving(t, ξ, θ, t′, s, F, L, T)
        return 0.5 * T^(1.5) * (1 - s / T) * (2b(x) * b′(x) + b″(x)) + ξ[1]
    else
        l = lvl(i, L)
        Δi = T / (1 << L) # subintervals
        k = (i - 1) ÷ (2 << l)
        δ = T / (1 << (L - l))
        res = 0.0
        ave_fact = 1/(l+1)
        for jj in 1:l+1 # number of points proportional to the length
            s = δ * k + Δi*(jj - 1 + rand()) #
            @assert δ * k < s < δ * (k+1)
            x = dotψmoving(t, ξ, θ, t′, s, F, L, T)
           res += ave_fact * 0.5 * δ * Λ(s, L - l, T) * (2b(x) * b′(x) + b″(x))
       end
       return res + ξ[i]
    end
end

##
function run_zz(df, T′)
    println("zigzag var 1")
    n = (2 << L) + 1
    ξ0 = 0randn(n)
    u, v = 0.0, 0.0  # initial and fianl point
    ξ0[1] = u / sqrt(T)
    ξ0[end] = v / sqrt(T)
    θ0 = rand((-1.0, 1.0), n)
    θ0[end] = θ0[1] = 0.0 # fix final point
    Γ = sparse(1.0I, n, n)
    c = [MyBound(0.0) for i in 1:n]
    #Improved algorithm
    zz_time = @elapsed((trace, (t, ξ, θ), (acc, num), c) = spdmp(∇ϕmoving_var1, 0.0, ξ0, θ0, T′, c, ZigZag(Γ, ξ0 * 0),
        SelfMoving(), L, T, adapt = false))
    t_trace = getindex.(trace.events, 1)
    t_trace = [0.0, t_trace...]
    x_trace, v_trace = resize(trace)
    burn1 = 100.0
    i1 = findfirst(x -> x > burn1, t_trace)
    xT_2 = sqrt(T)/2*getindex.(x_trace[i1:end], Int((length(ξ0)+1)/2))
    ave_xT_2_var = cumsum(((xT_2[1:end-1] + xT_2[2:end])/2).*diff(t_trace[i1:end]))./(t_trace[i1+1:end] .- burn1)
    for batches in [50]
        ess = ess_pdmp_components(t_trace[i1:end] .- burn1, x_trace[i1:end], v_trace[i1:end], n_batches = batches)
        push!(df, Dict(:sampler => "ZZ_var", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_mean", :y => sum(ess[2:end-1])/(length(ess)-2), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ_var", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_median", :y => median(ess[2:end-1]), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ_var", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_min", :y => minimum(ess[2:end-1]), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ_var", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_x_T2", :y => ess[Int((length(ξ0)+1)/2)], :runtime => zz_time))
    end
    println("standard zigzag")
    n = (2 << L) + 1
    ξ0 = 0randn(n)
    u, v = 0.0, 0.0  # initial and fianl point
    ξ0[1] = u / sqrt(T)
    ξ0[end] = v / sqrt(T)
    θ0 = rand((-1.0, 1.0), n)
    θ0[end] = θ0[1] = 0.0 # fix final point
    Γ = sparse(1.0I, n, n)
    c = [MyBound(0.0) for i in 1:n]
    #Improved algorithm
    zz_time = @elapsed((trace, (t, ξ, θ), (acc, num), c) = spdmp(∇ϕmoving, 0.0, ξ0, θ0, T′, c, ZigZag(Γ, ξ0 * 0),
        SelfMoving(), L, T, adapt = false))
    t_trace1 = getindex.(trace.events, 1)
    t_trace1 = [0.0, t_trace1...]
    x_trace, v_trace = resize(trace)
    i1 = findfirst(x -> x > burn1, t_trace1)
    xT_2 = sqrt(T)/2*getindex.(x_trace[i1:end], Int((length(ξ0)+1)/2))
    ave_xT_2 = cumsum(((xT_2[1:end-1] + xT_2[2:end])/2).*diff(t_trace1[i1:end]))./(t_trace1[i1+1:end] .- burn1)
    for batches in [50]
        ess = ess_pdmp_components(t_trace1[i1:end] .- burn1, x_trace[i1:end], v_trace[i1:end], n_batches = batches)
        push!(df, Dict(:sampler => "ZZ", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_mean", :y => sum(ess[2:end-1])/(length(ess)-2), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_median", :y => median(ess[2:end-1]), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_min", :y => minimum(ess[2:end-1]), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_x_T2", :y => ess[Int((length(ξ0)+1)/2)], :runtime => zz_time))
    end
    println("zigzag var 2")
    n = (2 << L) + 1
    ξ0 = 0randn(n)
    u, v = 0.0, 0.0  # initial and fianl point
    ξ0[1] = u / sqrt(T)
    ξ0[end] = v / sqrt(T)
    θ0 = rand((-1.0, 1.0), n)
    θ0[end] = θ0[1] = 0.0 # fix final point
    Γ = sparse(1.0I, n, n)
    c = [MyBound(0.0) for i in 1:n]
    #Improved algorithm
    zz_time = @elapsed((trace, (t, ξ, θ), (acc, num), c) = spdmp(∇ϕmoving_var2, 0.0, ξ0, θ0, T′, c, ZigZag(Γ, ξ0 * 0),
        SelfMoving(), L, T, adapt = false))
    t_trace3 = getindex.(trace.events, 1)
    t_trace3 = [0.0, t_trace3...]
    x_trace3, v_trace3 = resize(trace)
    i1 = findfirst(x -> x > burn1, t_trace3)
    xT_3 = sqrt(T)/2*getindex.(x_trace3[i1:end], Int((length(ξ0)+1)/2))
    ave_xT_3 = cumsum(((xT_3[1:end-1] + xT_3[2:end])/2).*diff(t_trace3[i1:end]))./(t_trace3[i1+1:end] .- burn1)
    for batches in [50]
        ess = ess_pdmp_components(t_trace3[i1:end] .- burn1, x_trace3[i1:end], v_trace3[i1:end], n_batches = batches)
        push!(df, Dict(:sampler => "ZZ_var2", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_mean", :y => sum(ess[2:end-1])/(length(ess)-2), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ_var2", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_median", :y => median(ess[2:end-1]), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ_var2", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_min", :y => minimum(ess[2:end-1]), :runtime => zz_time))
        push!(df, Dict(:sampler => "ZZ_var2", :alpha => α, :T => T′, :nbatches => batches,
            :stat => "ess_x_T2", :y => ess[Int((length(ξ0)+1)/2)], :runtime => zz_time))
    end
    return df, (ave_xT_2, t_trace1), (ave_xT_2_var, t_trace), (ave_xT_3, t_trace3)
end

function resize(zz_trace)
    x0 = zz_trace.x0
    v0 = zz_trace.θ0
    x_trace = [deepcopy(x0)]
    v_trace = [deepcopy(v0)]
    for i in eachindex(zz_trace.events)
        if i == 1
            x0 = x0 + v0*zz_trace.events[i][1]
        else
            x0 = x0 + v0*(zz_trace.events[i][1] - zz_trace.events[i-1][1])
        end
        v0[zz_trace.events[i][2]] = zz_trace.events[i][4]
        x0[zz_trace.events[i][2]] = zz_trace.events[i][3]
        # if !(v0[zz_trace.events[i][2]] == zz_trace.events[i][4])
        #     println("vel: $(v0[zz_trace.events[i][2]]) vs $(zz_trace.events[i][4])")
        #     error("")
        # end
        # if !(|x0[zz_trace.events[i][2]]| zz_trace.events[i][4])
        #     println("pos: $(x0[zz_trace.events[i][2]]) vs $(zz_trace.events[i][3])")
        #     error("")
        # end
        push!(x_trace, deepcopy(x0))
        push!(v_trace, deepcopy(v0))
    end
    x_trace, v_trace
end

using DataFrames, Random
df = DataFrame(sampler = String[], alpha = Float64[], T = Float64[],
            nbatches = Int64[], stat = String[], y = Float64[], runtime = Float64[])
function data_collection(df)
    Random.seed!(0)
    T′ = 25000.0
    zz_var = []
    zz_var1 = []
    zz = []
    for α′ in [0.1, 0.3, 0.5]
        global α = α′
        global L = 6
        global T = 50.0
        df, x1, x2, x3 = run_zz(df, T′)
        push!(zz, x1)
        push!(zz_var, x2)
        push!(zz_var1, x3)
    end
    return df, zz, zz_var, zz_var1
end
using CSV, JLD2
df, zz, zz_var1, zz_var2 = data_collection(df)
#CSV.write("./scripts/zz_diff_bridges/compare/benchamrk_zz_final.csv", df)
@save "./scripts/compare/covenrgence_zz.jld2" zz
@save "./scripts/compare/covenrgence_zz_var.jld2" zz_var1
@save "./scripts/compare/covenrgence_zz_var1.jld2" zz_var2
error("")

#
# using Makie
# using CairoMakie
#
# T′ = 1000
# α =0.3
# n = (2 << L) + 1
# ξ0 = 0randn(n)
# u, v = 0.0, 0.0  # initial and fianl point
# ξ0[1] = u / sqrt(T)
# ξ0[end] = v / sqrt(T)
# θ0 = rand((-1.0, 1.0), n)
# θ0[end] = θ0[1] = 0.0 # fix final point
# Γ = sparse(1.0I, n, n)
# c = [MyBound(0.0) for i in 1:n]
# zz_time = @elapsed((trace, (t, ξ, θ), (acc, num), c) = spdmp(∇ϕmoving, 0.0, ξ0, θ0, T′, c, ZigZag(Γ, ξ0 * 0),
#     SelfMoving(), L, T, adapt = false))
# t_trace = getindex.(trace.events, 1)
# t_trace = [0.0, t_trace...]
# x_trace, v_trace = resize(trace)
# ess = ess_pdmp_components(t_trace, x_trace, v_trace, n_batches = 50)
# #plot
# ts, ξs = splitpairs(discretize(trace, T′/n))
# S = T*(0:n)/(n+1)
# p1 = lines(S, [dotψ(ξ, s, L, T) for s in S], linewidth=0.3)
# for ξ in ξs[1:5:end]
#     lines!(p1, S, [dotψ(ξ, s, L, T) for s in S], linewidth=0.3)
# end
# display(p1)

max_stars_repo_name: SebaGraz/ZZDiffusionBridge

max_issues_repo_name: SebaGraz/ZZDiffusionBridge

max_forks_repo_name: SebaGraz/ZZDiffusionBridge


text:
using Random
using JuMP
using Gurobi
using DecisionProgramming
using RobustDecisionProgramming
using Base.Iterators: product

include("brute.jl")

rng = MersenneTwister(22)

S = States([2, 2, 3])
C = [ChanceNode(1, Node[]), ChanceNode(3, Node[2])]
D = [DecisionNode(2, Node[1])]
V = [ValueNode(4, Node[2, 3])]
X = [Probabilities(rng, c, S) for c in C]
Y = [Consequences(rng, v, S) for v in V]

validate_influence_diagram(S, C, D, V)
sort!.((C, D, V, X, Y), by = x -> x.j)

ks = [1, 2]
ϵs = [0.1, 0.07]
rs = [[x] for x in X]
for (k, ϵ) in zip(ks, ϵs)
    rs[k] = [Probabilities(X[k].j, d) for d in uncertainty_sets(X[k].data, ϵ)]
end
Xs = vec(map(x -> [x...], product(rs...)))

U = DefaultPathUtility(V, Y)
U⁺ = PositivePathUtility(S, U)

model = Model()
z = DecisionVariables(model, S, D)
Ps = Vector{DefaultPathProbability}()
EVs = Vector{GenericAffExpr}()
for X′ in Xs
    P = DefaultPathProbability(C, X′)
    π_s = PathProbabilityVariables(model, z, S, P; hard_lower_bound=false)
    EV = expected_value(model, π_s, U⁺)
    push!(Ps, P)
    push!(EVs, EV)
end
f = worst_case(model, EVs)
@objective(model, Max, f)

optimizer = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "IntFeasTol" => 1e-9,
)
set_optimizer(model, optimizer)
optimize!(model)

@info("Extracting results.")
Z = DecisionStrategy(z)

@info("Printing decision strategy:")
print_decision_strategy(S, Z)

max_stars_repo_name: jaantollander/RobustDecisionProgramming

max_issues_repo_name: jaantollander/RobustDecisionProgramming

max_forks_repo_name: jaantollander/RobustDecisionProgramming