text:
###################################################################################
# Antony Schutz 2015, ANR - POLCA - 2016
###################################################################################
function checkPack()

    if( Pkg.installed("OptimPack") == nothing )
        error("OptimPack.jl not installed")
    end

    # NFFT
    if( Pkg.installed("NFFT") == nothing )
        error("NFFT.jl not installed")
    end

    # Wavelets
    if( Pkg.installed("Wavelets") == nothing )
        error("Wavelets.jl not installed")
    end

    # OIFITS
    if( Pkg.installed("OIFITS") == nothing )
        error("OIFITS.jl not installed")
    end

    # HDF5
    if( Pkg.installed("JLD") == nothing )
        error("JLD.jl not installed")
    end

end
# ---------------------------------------------------------------------------------
# Check 3D mask
# ---------------------------------------------------------------------------------
function checkmask(mask3D::Array,nx::Int,nw::Int)
# check if the mask dimensions are valid
# mask3D: Array of dimensions:
# nx*nx*nw: 3D cube mask given by user
# nx^2 *nw: 3D Matrix
# nx*nx: gray image Matrix
# nx^2: gray image Vector
# empty array leads to default mask: no support contraint only positivity cube of ones
#
# if mask3D is a string it must be the path to a fits file the conversion from fits to data is done in painterinit(...)
#
# case: no mask given
    if isempty(mask3D)
        mask3D = ones(Float64, nx, nx, nw)
    end

    dims = ndims(mask3D)

# case, dims=1 mask3D is gray and must be of size (nx*nx)*1
    if dims == 1
        a = length(mask3D)

        if  a == nx * nx
          mask3D = reshape(repmat(reshape(vec(mask3D), nx, nx), 1, nw), nx, nx, nw)

        else
          error(" dims of mask3D =1, mask3D must be nx*nx (nx: Nb pixels)")
        end

# case, dims=3 mask3D is multi color and must be nx*nx*nw
    elseif dims == 3
        (a,b,c) = size(mask3D)

        if(a == nx)&&(b == nx)&&(c == nw)
        else
            error(" dims of mask3D =3, mask3D must be nx*nx*nw (nx: Nb pixels, nw: Nb wavelengths)")
        end

# case, dims=2 mask3D is multi color and must be (nx*nx)*nw
# or it is an image (nx*nx) and must be (nx*nx)*nw
    elseif dims == 2
        (a,b) = size(mask3D)

        if (a == nx)&&(b == nx)
            mask3D = reshape(repmat(mask3D, 1, nw), nx, nx, nw)

        elseif (a == nx*nx)&&(b == nw)
            mask3D = reshape(mask3D, nx, nx, nw)
            println(" Mask3D nx^2 * nw")

        else
            error(" dims of mask3D =2, mask3D is multi color and must be of size (nx*nx)*nw (nx: Nb pixels, nw: Nb wavelengths) or mask3D is an image and must be of size nx*nx (nx: Nb pixels)")
        end

    end
    return mask3D
end
# ---------------------------------------------------------------------------------
# First Estimate
# ---------------------------------------------------------------------------------
# compute the first estimate from xinit, a 3D an array
function inityc(xinit::Array,nb::Int,nx::Int,nw::Int,plan::Array)
# xinit Array given by checkinit(...) several kind of dimensions are allowed or path to a fits file
# nb: number of bases, spatial frequencies per wavelength
# nx: side size of images in pixels
# nw: number of wavelength
# plan: non uniform fft plan
    yc = zeros(Complex{Float64},nb,nw)

    for n in 1:nw
        yc[:,n] = nfft(plan[n], im * 0. + xinit[:, :, n]) / nx
    end

    return yc
end
function checkinit(xinit::Array,nb::Int,nx::Int,nw::Int,plan::Array)
# Check and Compute the first estimate
# xinit: Array of dimensions:
# nx*nx*nw: 3D cube initialization given by user
# nx^2 *nw: 3D Matrix
# nx*nx: gray image Matrix
# nx^2: gray image Vector
# or is complexe visibilities of size nb*nw Matrix or (nb*nx)*1 Vector
#
# nb: number of bases, spatial frequencies per wavelength
# nx: side size of images in pixels
# nw: number of wavelength
# plan: non uniform fft plan
    if isempty(xinit)
        xinit = zeros(nx, nx, nw)
        xinit[round(Int,  (nx / 2) + 1), round(Int, (nx / 2) + 1),:] = 1
        yc = inityc(xinit, nb, nx, nw, plan)
    end

    if ndims(xinit) == 1
        if size(xinit, 1) == nx * nx
            println("Gray vectorized Initial estimate")
            xinit = reshape(repmat(reshape(xinit, nx, nx), 1, nw), nx, nx, nw)
            yc = inityc(xinit, nb, nx, nw, plan)

        elseif size(xinit, 1) == nx * nx * nw
            println("Colored vectorized Initial estimate")
            xinit = reshape(xinit, nx, nx, nw)
            yc = inityc(xinit, nb, nx, nw, plan)

        elseif size(xinit,1) == nb * nw
            println("Input is vectorized Colored complexe visibilities (warm start will be less good)")
            xinit = zeros(nx, nx, nw)
            yc = reshape(xinit, nb, nw) + im * 0.
        end

    elseif ndims(xinit) == 2

        if (size(xinit, 1) == nx)&&(size(xinit, 2) == nx)
            println("Gray images Initial estimate")
            xinit = reshape(repmat(xinit, 1, nw), nx, nx, nw)
            yc = inityc(xinit, nb, nx, nw, plan)

        elseif (size(xinit, 1) == (nx * nx))&&(size(xinit, 2) == nw)
            println("Colored vectorized images as Initial estimate")
            xinit = reshape(xinit, nx, nx, nw)
            yc = inityc(xinit, nb, nx, nw, plan)

        elseif (size(xinit, 1) == nb)&&(size(xinit, 2) == nw)
            println("Input is Colored complexe visibilities  (warm start will be less good)")
            xinit = zeros(nx, nx, nw)
            yc = xinit + im * 0.
        end

    elseif ndims(xinit) == 3

        if (size(xinit, 1) == nx)&&(size(xinit, 2) == nx)&&(size(xinit, 3) == nw)
            println("Colored cube Initial estimate")
            yc = inityc(xinit, nb, nx, nw, plan)
        end

    end

    return yc,xinit

end
function checkinit(xinit::String,nb::Int,nx::Int,nw::Int,plan::Array)
# method 2 read initial estimate from fits and use method
    xfits = read(FITS(xinit))
    return checkinit(xfits, nb, nx, nw, plan)
end
###################################################################################
function mask(nx::Int,side::Int;choice="square")
# create image 2D of mask for mask3D
# choice = square or disk (disk) default: rect
# nx: size of the constructed image
# side: parameter for mask as radius of disk or half side length or square
    if(side > (nx / 2))
        error("side must be less than nx/2")
    end

    if(choice == "square")
        side2 = round(Int, nx / 2) - side
        mask3D = ones(nx, nx)
        mask3D[1:(1 + side2), :] = 0
        mask3D[:, 1:(1 + side2)] = 0
        mask3D[(nx - side2):nx, :] = 0
        mask3D[:, (nx - side2):nx] = 0

    elseif(choice == "disk")
        xg = reshape(repeat(collect(1:nx), outer=[nx]), nx, nx)
        yg = xg'
        rg = sqrt( (xg - .5 - (nx / 2)).^2 + (yg - .5 - (nx / 2)).^2 )
        mask3D = float(abs(rg) .< (side + .5))
    end

    return mask3D
end
###################################################################################
# PAINTER Config
###################################################################################
# Folder        : Folder of fits files
# nx            : size in pixels of image (image of size nx*nx)
# Wvlt          : list of wavelets basis
# lambda__spat  : Spatial Regularization parameter (weight) (Eq.29,31 [1])
# lambda__spec	: Spectral Regularization parameter (weight) (Eq.29,31 [1])
# lambda_L1     : L1 (for star field) constraint, emphasize sparsity of object (Novelty)
# epsilon       : Ridge/tikhonov parameter epilson |X|^2_2 (Eq.29,31 [1])
# rho_y         : ADMM parameter for data fidelity (convergence rate) (Eq.35,50-52 [1])
# rho_spat      : ADMM parameter for Spatial Regularization (convergence rate) (Eq.25,31 [1])
# rho_spec      : ADMM parameter for Spectral Regularization (convergence rate) (Eq.42,55 [1])
# rho_ps        : ADMM parameter for positivity (convergence rate) (Eq.47,54 [1])
# alpha         : Weight for Complexe visibility estimation from V2 (Eq.25,31 [1])
# beta	        : Weight for Complexe visibility estimation from phases difference (Eq.25,31 [1])
# eps1          : Primal Residual stop criterium (3.3 Optimality Conditions and Stopping Criterion , [2])
# eps2          : Dual Residual stop criterium (3.3 Optimality Conditions and Stopping Criterion , [2])
# FOV           : Field Of View User parameter, must be in ArcSecond
# nx            : size of desired output image, in pixel, image is then nx*nx pixels
# mask3D	      : Support constraint: can be a path to a fits file, or to data image (2D or 3D in both case)
# xinit3D       : Initial Estimate: can be a path to a fits file, or to data image (2D or 3D in both case)
###################################################################################
function painterinit(OIDATA::PAINTER_Input,Folder,nx,lambda_spat,lambda_spec,lambda_L1,epsilon,rho_y, rho_y_gamma, rho_y_xi,
  rho_spat,rho_spec,rho_ps,alpha,beta,eps1,eps2,FOV,mask3D,xinit3D,Wvlt,dptype,dpprm,indwvl,PlotFct)
# check if user parameters are valid parameters, correct them if type is not good or replace by default if parameter are not valid

# if only 1 wavelength, no spectral regularization
    OIDATA.rho_spec = !(length(indwvl)==1) * OIDATA.rho_spec

    if(FOV < 0)
        println("FOV must be non negative (default: 0.04 arcsecond)")
        println("initialized to default value")
        FOV = .04
    end

    FOV = FOV + 0.

    if !isinteger(nx)
        println("nx must an integer")
        println("Converted to integer")
    end

    nx = round(Int, nx)

    if(nx < 0)
        println("nx must be non negative (default: 64 pixels)")
        println("initialized to default value")
        nx = 64
    end

    if(lambda_spat < 0)

        println("lambda_spat must be non negative (default: 1.0)")
        println("initialized to default value")
        lambda_spat = 1e-3
    end

    lambda_spat = lambda_spat + 0.

    if(lambda_spec < 0)
        println("lambda_spec must be non negative (default: 1.0)")
        println("initialized to default value")
        lambda_spec = 1e-3
    end

    lambda_spec = lambda_spec + 0.

    if(lambda_L1 < 0)
        println("lambda_L1 must be non negative (default: 1e-20)")
        println("initialized to default value")
        lambda_L1 = 1e-3
    end

    lambda_L1 = lambda_L1 + 0.

    if(epsilon < 0)
        println("epsilon must be non negative (default: 1e-6, Ridge Regularization)")
        println("initialized to default value")
        epsilon = 1e-6
    end

    epsilon = epsilon + 0.

    if(rho_y < 0)
        println("rho_y must be non negative (default: 10.0)")
        println("initialized to default value")
        rho_y = 10.
    end

    rho_y = rho_y + 0.

    if(rho_y_gamma < 0)
        println("rho_y_gamma must be non negative (default: rho_y)")
        println("initialized to default value")
        rho_y_gamma = rho_y
    end

    rho_y_gamma = rho_y_gamma + 0.

    if(rho_y_xi < 0)
        println("rho_y_xi must be non negative (default: rho_y)")
        println("initialized to default value")
        rho_y_xi = rho_y
    end

    rho_y_xi = rho_y_xi + 0.

    if(rho_spat < 0)
        println("rho_spat must be non negative (default: 1.0)")
        println("initialized to default value")
        rho_spat = 1.
    end

    rho_spat = rho_spat + 0.

    if(rho_spec < 0)
        println("rho_spec must be non negative (default: 1.0)")
        println("initialized to default value")
        rho_spec = 4.
    end

    rho_spec = rho_spec + 0.

    if(rho_ps < 0)
        println("rho_ps must be non negative (default: 0.5)")
        println("initialized to default value")
        rho_ps = .5
    end

    rho_ps = rho_ps + 0.

    if(alpha < 0)
        println("alpha must be non negative (default: 1.0)")
        println("initialized to default value")
        alpha = 1.
    end

    alpha = alpha + 0.

    if(beta < 0)
        println("beta must be non negative (default: 1.0)")
        println("initialized to default value")
        beta = 1.
    end

    beta = beta + 0.

    if(eps1 < 0)
        println("eps1 must be non negative (default: 1e-4)")
        println("initialized to default value")
        eps1 = 1e-4
    end

    eps1 = eps1 + 0.

    if(eps2 < 0)
        println("eps2 must be non negative (default: 1e-4)")
        println("initialized to default value")
        eps2 = 1e-4
    end

    eps2 = eps2 + 0.

# Check mask3d type

    if typeof(mask3D) == String
        mask3dprint = "3D Mask Initialization from fits file: $mask3D"
        f = FITS(mask3D)
        mask3D = read(f[1])

    elseif typeof(mask3D) == Array{Float64, 2}
        Sz = size(mask3D)
        mask3dprint = "3D Mask Initialization from 2D data of size $Sz"

    elseif typeof(mask3D) == Array{Float64, 3}
        Sz = size(mask3D)
        mask3dprint = "3D Mask Initialization from 3D data of size $Sz"

    elseif isempty(mask3D)
        mask3dprint = "3D Mask Initialization from default, no constraint"
    end

# Check xinit3d type

    if typeof(xinit3D) == String
        xinit3dprint = "3D Init Initialization from fits file: $xinit3D"
        f = FITS(xinit3D)
        xinit3D = read(f[1])

    elseif typeof(xinit3D) == Array{Float64, 2}
        Sx = size(xinit3D)
        xinit3dprint = "3D Init Initialization from 2D data of size $Sx"

    elseif typeof(xinit3D) == Array{Float64, 3}
        Sx = size(xinit3D)
        xinit3dprint = "3D Init Initialization from 3D data of size $Sx"

    elseif isempty(xinit3D)
        xinit3dprint = "3D Init Initialization from default, centered dirac"
    end

    # Check Differential phases matrix parameters
    if typeof(dptype) == String
        if (dptype!="all") && (dptype!="ref") && (dptype!="frame")&& (dptype!="sliding")&& (dptype!="diag")&& (dptype!="phase")
              println("dptype: all (default), ref, diag, frame, sliding ")
              println("dptype: can be 'phase' for phases of complexe visibilities")
              println("initialized to default value")
              dptype = "all"
        end
    else
        error("dptype must be a string: all (default), ref, frame, sliding, diag, phase ")
    end

    if !isinteger(dpprm)
        println("lambda_ref/horizon must an integer")
        println("Converted to integer")
    end

    dpprm = round(Int,dpprm)

    if(dpprm <= 0)
        println("lambda_ref/horizon must be non negative (default: 1)")
        println("initialized to default value")
        dpprm = 1
    end

# Check data path
    if isempty(Folder)
        tmp = pwd()
        Folder = joinpath(dirname(@__FILE__), "OIFITS")
        println(Folder)
    end

    if (typeof(Folder) == String)||(typeof(Folder) == UTF8String)
        cpath = Folder

    else
        error("data path is not correct")
    end

# if pwd/OIFITS does not exits so, search in user/.julia/vx.x/PAINTER/src/OIFITS
    if !isdir(cpath)
        cptmp = cpath
        cpath = joinpath(dirname(@__FILE__), "OIFITS")
        println("OIFITS files folder does not exists in $cptmp, replaced by default $cpath")
        println("")
    end

    pathprint = cpath

# Check Wavelet input
    if (typeof(Wvlt) <: Array) & all( [typeof(wlt) <: WT.OrthoWaveletClass for wlt in Wvlt])
        Wvltprint = "$Wvlt"
    else
        error("Wavelets list must be a vector of orthogonal wavelets")
    end

    if !(typeof(PlotFct) <: Function)
        error("PlotFct must be a Function, default plot function: painterplotfct() will be used")
        PlotFct = painterplotfct
    end

    println("OIFits path = $pathprint")
    println("lambda_spat = $lambda_spat ")
    println("lambda_spec = $lambda_spec ")
    println("lambda_L1   = $lambda_L1 ")
    println("epsilon     = $epsilon ")
    println("rho_y       = $rho_y ")
    println("rho_spat    = $rho_spat ")
    println("rho_spec    = $rho_spec ")
    println("rho_ps      = $rho_ps ")
    println("alpha       = $alpha ")
    println("beta	    = $beta ")
    println("eps1        = $eps1 ")
    println("eps2        = $eps2 ")
    println("FOV         = $FOV arcsecond")
    println("nx          = $nx pixels")
    println("mask3D      = $mask3dprint")
    println("xinit3d     = $xinit3dprint")
    println("Wavelets    = $Wvltprint")
    println("Plot Func   = $PlotFct")
    println("dptype      = $dptype")
    println("dpprm       = $dpprm")

    OIDATA.PlotFct = PlotFct
    OIDATA.Folder = cpath
    OIDATA.lambda_spat = lambda_spat
    OIDATA.lambda_spec = lambda_spec
    OIDATA.lambda_L1 = lambda_L1
    OIDATA.epsilon = epsilon
    OIDATA.rho_y = rho_y
    OIDATA.rho_spat = rho_spat
    OIDATA.rho_spec = rho_spec
    OIDATA.rho_ps = rho_ps
    OIDATA.alpha = alpha
    OIDATA.beta = beta
    OIDATA.eps1 = eps1
    OIDATA.eps2 = eps2
    OIDATA.FOV = FOV
    OIDATA.nx = nx
    OIDATA.mask3D = mask3D
    OIDATA.xinit3D = xinit3D
    OIDATA.Wvlt = Wvlt
    OIDATA.dptype = dptype
    OIDATA.dpprm = dpprm
    return OIDATA
end
###################################################################################
# Initialise ADMM Array
function painterarrayinit(PDATA::PAINTER_Data,OIDATA::PAINTER_Input)

    nwvlt = length(OIDATA.Wvlt)
    DegRd = 2.0 * pi / 360.0
    RadAs = 3600.0 / DegRd
    coef = OIDATA.FOV / RadAs

# matrices, constants, plan creation
    PDATA.eta = nwvlt * OIDATA.rho_spat + OIDATA.rho_spec + OIDATA.rho_ps + OIDATA.epsilon
    PDATA.plan = planarray_par(OIDATA.U * coef, OIDATA.V * coef, OIDATA.nx, OIDATA.nw)
    PDATA.F3D = nudft3d_par(OIDATA.U * coef, OIDATA.V * coef, OIDATA.nb, OIDATA.nx, OIDATA.nw)
    PDATA.M  = invmat_par(PDATA.F3D, OIDATA.rho_y, PDATA.eta, OIDATA.nw)
    println("process independent Cluster")
    for n in 1:length(OIDATA.baseNb)
        OIDATA.K[n] = ItKappa(OIDATA.K[n])
        PDATA.H[n] = phasetophasediff(OIDATA.orderedCluster[n], OIDATA.nw, length(OIDATA.baseNb[n]), 1, OIDATA.isDP, OIDATA.dptype, OIDATA.dpprm)
    end
# Array Initialization
    PDATA.x = zeros(Float64, (OIDATA.nx, OIDATA.nx, OIDATA.nw))
    PDATA.tau_s = zeros(Float64, OIDATA.nx, OIDATA.nx, OIDATA.nw, nwvlt)
    PDATA.w = zeros(Float64, OIDATA.nx, OIDATA.nx, OIDATA.nw)
    PDATA.tau_w = zeros(Float64, OIDATA.nx, OIDATA.nx, OIDATA.nw)
    PDATA.tau_v = zeros(Float64, OIDATA.nx, OIDATA.nx, OIDATA.nw)
    PDATA.tau_r = zeros(Float64, OIDATA.nx, OIDATA.nx, OIDATA.nw)
    PDATA.Fx = zeros(Complex{Float64}, (OIDATA.nb, OIDATA.nw))
    PDATA.tau_xc = zeros(Complex128, OIDATA.nb, OIDATA.nw)
    PDATA.tau_pwc = zeros(Complex128, OIDATA.nb, OIDATA.nw)
    PDATA.tau_xic = zeros(Complex128, OIDATA.nb, OIDATA.nw)
    PDATA.y_v2 = zeros(Complex128, OIDATA.nb, OIDATA.nw)
    PDATA.y_phi = zeros(Complex128, OIDATA.nb, OIDATA.nw)
    PDATA.yc, OIDATA.xinit3D = checkinit(OIDATA.xinit3D, OIDATA.nb, OIDATA.nx, OIDATA.nw, PDATA.plan)
    PDATA.z = zeros(Float64, OIDATA.nx, OIDATA.nx, OIDATA.nw, nwvlt)
    PDATA.v = zeros(Float64, OIDATA.nx, OIDATA.nx, OIDATA.nw)
    PDATA.r = zeros(Float64, OIDATA.nx, OIDATA.nx, OIDATA.nw)
    return PDATA,OIDATA
end
###################################################################################
# Initialise Data and xinit from V2 and flux
function painterautoparametersinit(PDATA::PAINTER_Data,OIDATA::PAINTER_Input)
    nP = sum(OIDATA.P)
    OIDATA.norm = nP
    OIDATA.P = OIDATA.P / nP
    OIDATA.W = OIDATA.W / nP^2

    # AllK = Float64{}[]
    # for n in 1 : length(OIDATA.K)
    #     for m in 1 : length(OIDATA.K[n])
    #         push!(AllK,OIDATA.K[n][m])
    #     end
    # end
    #
    # OIDATA.alpha = OIDATA.alpha .*maximum( OIDATA.W )
    # OIDATA.beta = OIDATA.beta./maximum(abs( AllK ))

    OIDATA.alpha = 1
    OIDATA.beta = 1

    # normalise lambda: spatial and spectral
    OIDATA.lambda_spat = OIDATA.lambda_spat ./ (OIDATA.nx*OIDATA.nx)
    OIDATA.lambda_spec = OIDATA.lambda_spec ./ (OIDATA.nw)

    # give to initial estimate the good power in each channel
    for n in 1:OIDATA.nw
        factor = sqrt( sum(OIDATA.P[:,n]) ./ sum(abs2.(PDATA.yc[:,n])) )
        OIDATA.xinit3D[:,:,n] = OIDATA.xinit3D[:,:,n]  *  factor
        PDATA.yc[:,n] = PDATA.yc[:,n] * factor
    end

    miniv2, minivp = compute_hessianVP(PDATA, OIDATA)

    OIDATA.rho_y_gamma = max( OIDATA.alpha * (-miniv2)*1.01 , 1e-6)
    OIDATA.rho_y_xi    = max( OIDATA.beta  * (-minivp)*1.01 , 1e-6)

    return PDATA,OIDATA
end
###################################################################################
# Compute Hessian from initial estimate
function compute_hessianVP(PDATA::PAINTER_Data,OIDATA::PAINTER_Input)
    diagofhessV2 = (-4 ./ OIDATA.W) .* ( OIDATA.P - 3* abs2.( PDATA.yc ) )
    miniv2 = minimum(diagofhessV2)

    vp = zeros( length(OIDATA.K) )
    for n in 1 : length(OIDATA.K)
      H = PDATA.H[n]
      dK = diagm( OIDATA.K[n] )
      dX = diagm( cos.( H*vec(angle.(PDATA.yc[OIDATA.baseNb[n],:])) - OIDATA.Xi[n]  ))
      hess = H' * dK * dX * H
      vp[n] = eigs(hess,nev=1,which=:SR)[1][1]
    end
    minivp = minimum(vp)
    return miniv2, minivp
end
###################################################################################
# Initialise Lagrange Multipliers
function painterlagrangemultipliersinit(PDATA::PAINTER_Data,OIDATA::PAINTER_Input)

    nx = OIDATA.nx
    nb = OIDATA.nb
    nw = OIDATA.nw
    wvl = OIDATA.wvl
    Wvlt = OIDATA.Wvlt
    NWvlt = length(Wvlt)

    rho_y = OIDATA.rho_y
    rho_spat = OIDATA.rho_spat
    rho_spec = OIDATA.rho_spec
    rho_ps = OIDATA.rho_ps

    lambda_spat = OIDATA.lambda_spat
    lambda_spec = OIDATA.lambda_spec
    lambda_L1 = OIDATA.lambda_L1

    vHt = convert( SharedArray, zeros(nx, nx, nw))
    Hx = convert( SharedArray, zeros( nx, nx, nw, NWvlt))

    # Cube of flux
    flux_cube = 1. # vec(sum(OIDATA.P,1))

    if rho_spat >0
        tmpx = copy(OIDATA.xinit3D)
        @sync @parallel for ind in 1:nw*NWvlt
            n,b = ind2sub((nw,NWvlt),ind)
            Hx[:, :, n, b] = dwt(tmpx[:, :, n], wavelet(Wvlt[b]))
        end
        # update of z
        PDATA.z = copy(Hx)
        PDATA.z = max.(1 - ((flux_cube.*lambda_spat / rho_spat) ./ abs.(PDATA.z)), 0.) .* PDATA.z
    end

    # update of v
    if rho_spec >0
        PDATA.v = OIDATA.xinit3D  / 2.
        vecv = permutedims(PDATA.v, [3, 1, 2])
        @sync @parallel for ind in 1:nx*nx
            m,n = ind2sub((nx,nx),ind)
            vHt[m, n, :] = dct(vecv[:, m, n] )
        end
        # update of r
        PDATA.r = copy(vHt)
        PDATA.r = max.(1 - (lambda_spec / rho_spec) ./ abs.(PDATA.r), 0) .* PDATA.r
    end

    # update of w
    if rho_ps>0
        PDATA.w = max.(max.(0.0, OIDATA.xinit3D) .* OIDATA.mask3D - flux_cube.*lambda_L1, 0)
        PDATA.tau_w = rho_ps * (OIDATA.xinit3D - PDATA.w)
    end

    if rho_spat >0
        PDATA.tau_s = rho_spat * (copy(Hx) - PDATA.z)
    end

    if rho_spec >0
        PDATA.tau_v = rho_spec * OIDATA.xinit3D /2
        PDATA.tau_r = rho_spec * (copy(vHt) - PDATA.r)
    end

    return PDATA,OIDATA
end


text:
using Printf

include("test_hb.jl")

ext: jl


text:
"""
    Subcommunities(num)

AbstractPartition subtype with multiple subcommunities.

"""
struct Subcommunities <: Diversity.API.AbstractPartition{Nothing}
    num::Int64
    names::Vector{String}

    function Subcommunities(num::Integer)
        num > 0 || error("Too few subcommunities")
        new(num, map(x -> "$x", 1:num))
    end

    function Subcommunities(names::Vector{String})
        num = length(names)
        num > 0 || error("Too few subcommunities")
        new(num, names)
    end
end

import Diversity.API._getsubcommunitynames
function _getsubcommunitynames(sc::Subcommunities)
    return sc.names
end

import Diversity.API._countsubcommunities
function _countsubcommunities(sub::Subcommunities)
    return sub.num
end

"""
    Onecommunity

AbstractPartition subtype containing only one subcommunity.
"""
struct Onecommunity <: Diversity.API.AbstractPartition{Nothing}
    namev::Vector{String}

    function Onecommunity(name::String = "1")
        new([name])
    end
end

function _getsubcommunitynames(oc::Onecommunity)
    return oc.namev
end

function _countsubcommunities(::Onecommunity)
    return 1
end


text:
using Test
using GraphIO.CDF

@testset "CDF" begin
    g = loadgraph(joinpath(testdir, "testdata", "30bus.jlg"))
    read_test(CDFFormat(), g, "graph", joinpath(testdir, "testdata", "30bus.cdf"))
end


text:
# Takes an array of ResponseData and predicts the label based on the model in mlmodel
function predict_linear(mom::StatisticalMoments,  fm::FeatureMap,  ml::ML_model, options)
    x = zeros(1,ml.nfeatures);
    fm.apply(x, 1, mom,options);   # Applied the features map
    x = x_apply_transform(x,ml); #Applies any data transformation we applied to the training set including add the 1 bias term
    #Compute linear prediction implements representar theorem:  $<w^*, \psi(x )> = \sum_{\ell =1} (\alpha^*)^\ell K( x^\ell,x) $
    ys = ml.kernelapply(ml.X,x,ml.kernelparam)*ml.parameters;# println("ys before undo_transform: ",ys)
    #Undo the data transformations applied to y
    ys = y_undo_transform(ys,ml);
    return ys;
end


text:
using Doubles
using Base.Test

# courtesy: https://github.com/simonbyrne/DoubleDouble.jl

x = sqrt(2.0)
bx = big(x)
sx = x
dx = Double(x)

y = 0.1
by = big(y)
sy = y
dy = Double(y)

@test x == sx == dx
@test y == sy == dy

dxy = dx*dy
bxy = bx*by
@test sx*sy == dxy
@test x*y == Float64(dxy)
# @test dxy == Double(bxy)

@test x+y == Float64(dx+dy)
# @test dx+dy == Double(bx+by)

@test x-y == Float64(dx-dy)
# @test dx-dy == Double(bx-by)

@test x/y == Float64(dx/dy)
# @test dx/dy == Double(bx/by)

@test sqrt(y) == Float64(sqrt(dy))
# @test sqrt(dy) == Double(sqrt(by))


text:
if VERSION >= v"0.7.0-DEV-"
    using Test
else
    using Base.Test
end

@test 1==1


text:
function get_eop_history(obj::TEMEProvider)
    return jcall(obj, "getEOPHistory", EOPHistory, ())
end

function get_non_interpolating_provider(obj::TEMEProvider)
    return jcall(obj, "getNonInterpolatingProvider", TEMEProvider, ())
end

function get_transform(obj::TEMEProvider, arg0::AbsoluteDate)
    return jcall(obj, "getTransform", Transform, (AbsoluteDate,), arg0)
end

function get_transform(obj::TEMEProvider, arg0::FieldAbsoluteDate)
    return jcall(obj, "getTransform", FieldTransform, (FieldAbsoluteDate,), arg0)
end


text:
export AbstractSystem,
       AbstractSystemCache,
       SystemNullCache,
       cache,
       evaluate,
       evaluate!,
       jacobian,
       jacobian!,
       differentiate_parameters!,
       differentiate_parameters,
       evaluate_and_jacobian,
       evaluate_and_jacobian!

"""
    SystemNullCache

An empty cache if no cache is necessary.
"""
struct SystemNullCache <: AbstractSystemCache end

"""
    cache(F::AbstractSystem, x)::AbstractSystemCache

Create a cache for the evaluation (incl. Jacobian) of `F` with elements of the type
of `x`.

    cache(F::AbstractSystem, x, p)::AbstractSystemCache

Create a cache for the evaluation (incl. Jacobian) of `F` with elements of the type
of `x` and parameters `p`.
"""
cache(F::AbstractSystem, args...) = throw(MethodError(cache, tuple(F, args...)))

"""
    evaluate!(u, F::AbstractSystem, x, cache::AbstractSystemCache)

Evaluate the system `F` at `x` and store the result in `u`.

    evaluate!(u, F::AbstractSystem, x, p, cache::AbstractSystemCache)

Evaluate the system `F` at `x` and parameters `p` and store the result in `u`.
"""
evaluate!(u, F::AbstractSystem, args...) =
    error(MethodError(evaluate!, tuple(u, F, args...)))

"""
    evaluate(F::AbstractSystem, x::AbstractVector, cache=cache(F, x))

Evaluate the system `F` at `x`.

    evaluate(F::AbstractSystem, x::AbstractVector, p, cache=cache(F, x))

Evaluate the system `F` at `x` and parameters `p`.
"""
evaluate(F::AbstractSystem, x, c::AbstractSystemCache = cache(F, x)) = evaluate(F, x, c)
evaluate(F::AbstractSystem, x, p, c::AbstractSystemCache = cache(F, x, p)) =
    evaluate(F, x, p, c)

"""
    jacobian!(u, F::AbstractSystem, x , cache::AbstractSystemCache)

Evaluate the Jacobian of the system `F` at `x` and store the result in `u`.

    jacobian!(u, F::AbstractSystem, x , p, cache::AbstractSystemCache)

Evaluate the Jacobian of the system `F` at `x` and parameters `p` and store the result in `u`.
"""
jacobian!(u, F::AbstractSystem, args...) =
    error(MethodError(jacobian!, tuple(u, F, args...)))

"""
    jacobian(F::AbstractSystem, x, cache=cache(F, x))

Evaluate the Jacobian of the system `F` at `x`.

    jacobian(F::AbstractSystem, x , p, cache=cache(F, x))

Evaluate the Jacobian of the system `F` at `x` and parameters `p`.
"""
jacobian(F::AbstractSystem, x, c::AbstractSystemCache = cache(F, x)) = jacobian(F, x, c)
jacobian(F::AbstractSystem, x, p, c::AbstractSystemCache = cache(F, x, p)) =
    jacobian(F, x, p, c)

"""
    differentiate_parameters!(u, F::AbstractSystem, x, p, cache::AbstractSystemCache)

Evaluate the Jacobian of the system `F` at `x` and parameters `p` w.r.t. the parameters
and store the result in `u`.
"""
function differentiate_parameters! end

"""
    differentiate_parameters(F::AbstractSystem, x, p, cache=cache(F, x))

Evaluate the Jacobian of the system `F` at `x` and parameters `p` w.r.t. the parameters
"""
differentiate_parameters(F::AbstractSystem, x, p, c::AbstractSystemCache = cache(F, x, p)) =
    differentiate_parameters(F, x, p, c)

# Optional
"""
    evaluate_and_jacobian!(u, U, F, x , cache::AbstractSystemCache)

Evaluate the system `F` and its Jacobian at `x` and store the results in `u` (evalution)
and `U` (Jacobian).
"""
function evaluate_and_jacobian!(u, U, F::AbstractSystem, x, cache::AbstractSystemCache)
    evaluate!(u, F, x, cache)
    jacobian!(U, F, x, cache)
    nothing
end

"""
    evaluate_and_jacobian!(u, U, F, x, p, cache::AbstractSystemCache)

Evaluate the system `F` and its Jacobian at `x` and parameters `p` and store the results in `u` (evalution)
and `U` (Jacobian).
"""
function evaluate_and_jacobian!(u, U, F::AbstractSystem, x, p, cache::AbstractSystemCache)
    evaluate!(u, F, x, p, cache)
    jacobian!(U, F, x, p, cache)
    nothing
end

"""
    evaluate_and_jacobian(F::AbstractSystem, x, cache=cache(F, x))

Evaluate the system `F` and its Jacobian at `x`.
"""
function evaluate_and_jacobian(F::AbstractSystem, x, c::AbstractSystemCache = cache(F, x))
    u = evaluate(F, x, c)
    U = jacobian(F, x, c)
    u, U
end
"""
    evaluate_and_jacobian(F::AbstractSystem, x, p, cache=cache(F, x))

Evaluate the system `F` and its Jacobian at `x` and parameters `p`.
"""
function evaluate_and_jacobian(
    F::AbstractSystem,
    x,
    p,
    c::AbstractSystemCache = cache(F, x),
)
    u = evaluate(F, x, p, c)
    U = jacobian(F, x, p, c)
    u, U
end

"""
    Base.size(F::AbstractSystem)

Returns a tuple `(m, n)` indicating that `F` is a system of `m` polynomials `m` in `n` variables.
"""
Base.size(::AbstractSystem) = error("Mandatory to define `Base.size` for `AbstractSystem`s")
Base.size(F::AbstractSystem, i::Integer) = size(F)[i]
Base.length(F::AbstractSystem) = size(F, 1)
npolynomials(F::AbstractSystem) = length(F)

"""
    degrees(F::AbstractSystem)

Returns a the degrees of the polynomials.
"""
degrees(F::AbstractSystem) = nothing

include("systems/sp_system.jl")
include("systems/fixed_homotopy.jl")
include("systems/fp_system.jl")
include("systems/totaldegree_system.jl")
include("systems/multihom_totaldegree_system.jl")
include("systems/composition_system.jl")
include("systems/fixed_parameter_system.jl")
include("systems/squared_up_system.jl")
include("systems/trace_test_system.jl")
include("systems/trace_test_pencil.jl")


text:
module Preprocessing
using Base.Threads
using Images

function convert_to_float32(zstack::AbstractArray)
    return Float32.(reinterpret(Images.N0f16, zstack))
end

function equalize_zstack(zstack)
    # equalize zstack with equalization separately performed for each z slice
    output = similar(zstack)
    equalizer = Images.Equalization(nbins = 256, minval = 0.0, maxval = 0.7)
    @threads for I in CartesianIndices(axes(zstack)[1:2])
        output[I, :, :] = Images.adjust_histogram(zstack[I, :, :], equalizer)
    end
    return output
end

export convert_to_float32
end


text:
"""
    refine_triangulation(triang_vertices, triang_simplex_indices, split_indices, k)

Refine a triangulation with explicitly providing the splitting rules. In this function,
these are computed internally, using a splitting factor of k.
"""
function refine_triangulation(triang_vertices, triang_simplex_indices, split_indices, k)

    if length(split_indices) == 0
        untouched_indices = collect(1:size(triang_simplex_indices, 1))
        return triang_vertices, triang_simplex_indices, untouched_indices
    end

    # The number of simplices to split
    n_split_simplices = length(split_indices)

    # The dimension of the space
    E = size(triang_vertices, 2)

    # Generic rules for splitting a simplex
    splitting_rules = simplicial_subdivision(k, E)

    # Rules for forming the strictly new vertices of the subtriangulation
    rules = splitting_rules[1]

    # Array where each row represents one of the new simplices in the splitted simplex
    subtriangulation = splitting_rules[2]

    # How many new vertices are created each split?
    n_newvertices_eachsplit = size(rules, 1)

    # We need an array that can accomodate all of them. Each row in this
    # array will be a new vertex. Stacks of n_newvertices_eachsplit * E arrays.
    # We have as many as we have simplices to split.
    new_vertices = zeros(n_newvertices_eachsplit * n_split_simplices, E)

    # Fill the array by looping over all simplices that we want to split
    for i = 1:n_split_simplices
        # Figure out what the row indices corresponding to the ith simplex
        # must be. Marks the beginning of each of the simplex stacks in new_vertices
        ind = n_newvertices_eachsplit * (i -1)

        # Index of the simplex we need to split
        simplex_idx = split_indices[i]

        # Get the vertices of the simplex currently being splitted. Each of the
        # n_newvertices_eachsplit new vertices will be a linear combination
        # of these vertices. Each row is a vertex.
        vertices = triang_vertices[triang_simplex_indices[simplex_idx, :], :]

        # Generate the strictly new vertices for each sub
        for j = 1:n_newvertices_eachsplit
            # The index for a particular new vertex of the ith new simplex
            ind_newvertex = ind + j

            # Compute the jth vertex of the subtriangulation of the ith simplex

            # Go the jth new subsimplex of the ith splitted simplex. The entries of this
            # vector
            jth_subsimplex = rules[j, :]

            # Pick the corresponding original vertices with indices contained in rules[j, :]
            original_vertices = vertices[rules[j, :], :]

            new_vertices[ind_newvertex, :] = sum(original_vertices, dims=1) ./ k

        end
    end

    # Find the unique new vertices
    new_vertices_noreps = unique(new_vertices, dims=1)
    Ind = Vector{Int}(undef, size(new_vertices, 1))
    for i = 1:size(new_vertices_noreps, 1)
        for j = 1:size(new_vertices, 1)
            if new_vertices_noreps[i, :] == new_vertices[j, :]
                Ind[j] = i
            end
        end
    end

    # Combine old and newly introduced vertices
    num_vertices_beforesplit = size(triang_vertices, 1)
    triang_vertices = vcat(triang_vertices, new_vertices_noreps)

    # Update the Ind array, so that we start at the new vertices
    Ind = Ind .+ num_vertices_beforesplit
    num_simplices_each_split = size(subtriangulation, 1)

    # The subsimplices formed by the splitting. Each row contains E + 1 indices referencing
    # the vertices furnishing that particular subsimplex (now found in the updated
    # triang_vertices array).
    newtriangulation = Array{Float64, 2}(undef, num_simplices_each_split * n_split_simplices, E + 1)

    # For each simplex that we need to split,
    for i = 1:n_split_simplices
        # The beginning of the stack we need to fill.
        index = num_simplices_each_split * (i - 1)

        # Figure out what the row indices corresponding to the ith simplex
        # must be. Marks the beginning of each of the simplex stacks in new_vertices
        ind = n_newvertices_eachsplit * (i -1)

        # Index of the simplex we need to split
        simplex_idx = split_indices[i]

        # Pick the indices of the original vertices. Should be a column vector.
        inds_original_vertices = triang_simplex_indices[simplex_idx, :]

        # Indices of the new vertices. Should be a column vector
        inds_new_vertices = Ind[(ind + 1):(ind + n_newvertices_eachsplit)]
        inds_all_vertices = vcat(inds_original_vertices, inds_new_vertices)

        # Go through each of the subsimplices formed by splitting that simplex
        for j = 1:num_simplices_each_split
            newtriangulation[index + j, :] = inds_all_vertices[subtriangulation[j, :]]
        end
    end

    # Indices of the simplices that are not split
    untouched_indices = setdiff(1:size(triang_simplex_indices, 1), split_indices)

    triang_simplex_indices = round.(Int, vcat(triang_simplex_indices[untouched_indices, :],
                                  newtriangulation))

    return triang_vertices, triang_simplex_indices, untouched_indices
end


text:
using HorizonSideRobots
include("BackPath.jl")
include("Structs.jl")

function chess_putmarker!(robot::Robot)
    back_path = BackPath(robot)
    r = ChessRobot(robot)
    if sum(back_path.path) % 2 == r.flag
        r.flag = (r.flag - 1) * (-1)
    else
        putmarker!(robot)
    end

    snake(r)
    to_sud_west_corner = BackPath(robot)
    back!(robot, back_path)
end

max_stars_repo_name: Mananhina/mirea-progs

max_issues_repo_name: Mananhina/mirea-progs

max_forks_repo_name: Mananhina/mirea-progs


text:
module Receive

include("patch.jl")

struct Break end

haslock(ch::Channel) = current_task() == ch.cond_put.lock.locked_by
isflushed(ch::Channel) = haslock(ch) && !isready(ch)

"""
    _receive_nonblocking(f::Function, ch_vec::Vector{Channel})

Acquires lock for each channel in nonblocking way with `trylock` and
for each unsuccessful attempt takes value from the channel. When lock
is acquired channel is flushed until `isready(ch)==false`. To prevent
starvation the order of channels are permuted like 123 231 312 for
three channels.
"""
function _receive_nonblocking(f::Function, ch_vec::Vector{Channel})

    N = length(ch_vec)

    k = 1
    j = 1

    terminated = false
    bingo = false

    while true

        if k == j
            j = mod(j, N) + 1
            k = j

            bingo==true && break
            bingo = true
        end

        k = mod(k, N) + 1

        ch = ch_vec[k]

        # If lock is not owned this part shall be executed only if
        # either lock can not be acquired. If lock is acquired then I
        # need to check whether there are available values which I can
        # take.

        #if (!(haslock(ch) || trylock(ch)) && isready(ch) ) || (haslock(ch) && isready(ch))

        haslock(ch) || trylock(ch)

        if isready(ch)

            value = take!(ch)
            res = f(ch, value)

            if res isa Break
                terminated = true
                break
            end

            bingo = false

        end

    end

    return terminated

end

function unlock_channels(chvec::Vector{Channel})

    for ch in chvec
        haslock(ch) || unlock(ch)
    end

    return nothing
end

function set_conditions!(ch::Channel, cond_wait::Threads.Condition)

    lk = cond_wait.lock

    ch.cond_wait = cond_wait
    ch.cond_put = Threads.Condition(lk)
    ch.cond_take = Threads.Condition(lk)

    return nothing
end

function restore_conditions!(ch::Channel)

    lk = ReentrantLock()

    old_cond_put = ch.cond_put

    ch.cond_wait = Threads.Condition(lk)
    ch.cond_put = Threads.Condition(lk)
    ch.cond_take = Threads.Condition(lk)

    if length(old_cond_put.waitq) > 0

        @info "Unblock producer from `wait(ch.cond_put)`"

        el = popfirst!(old_cond_put.waitq)
        push!(ch.cond_put.waitq, el)
        notify(old_cond_put)
    end

    return nothing
end

"""
    _receive_blocking(f::Function, ch_vec::Vector{Channel})

Takes values from channels in a ready state. If none is ready waits
for `cond_wait` to be triggered. It is assumed that the thread owns a
lock of for the condition and that all channels have the same
`cond_wait`.
"""
function _receive_blocking(f::Function, ch_vec::Vector{Channel})

    c = ch_vec[1].cond_wait

    N = length(ch_vec)

    k = 1
    j = 1

    bingo = false

    while true

        if k == j
            j = mod(j, N) + 1
            k = j

            # It seems like an unfortunate race condition here :(
            if bingo
                wait(c)
            end

            bingo = true
        end

        k = mod(k, N) + 1

        ch = ch_vec[k]

        if isready(ch)

            value = take!(ch)

            res = f(ch, value)

            if res isa Break
                break
            end

            bingo = false
        end

    end

    return nothing

end

receive(f::Function, ch1::Channel, ch2::Channel) = receive(f, Channel[ch1, ch2])

function receive(f::Function, ch_vec::Vector{Channel})

    # Step 1: Acquire locks for channels and meanwhile take the values
    # from channels in nonblocking way. To prevent starvation the order
    # of channels are permuted like 123 231 312 for three channels.

    terminated = false
    try
        terminated = _receive_nonblocking(f, ch_vec)
    catch err
        terminated = true
        rethrow(err)
    finally
        if terminated
            unlock_channels(ch_vec)
            return nothing
        end
    end

    # Step 2: Replace every channel conditions with new ones having
    # the same lock and for cond_wait the same condition.

    lk = ReentrantLock()
    c = Threads.Condition(lk)

    lock(lk)

    for ch in ch_vec
        olk = ch.cond_wait.lock
        set_conditions!(ch, c)
        unlock(olk)
    end

    # Step 3: Takes values from channels in a ready state. If none is
    # ready waits for `cond_wait` to be triggered. To prevent
    # starvation the order of channels are permuted like 123 231 312
    # for three channels.

    try

        _receive_blocking(f, ch_vec)

    finally

        # Step 4: Cleanup.

        for ch in ch_vec
            restore_conditions!(ch)
        end

        unlock(lk)

    end

    return nothing

end

export receive

end # module


text:
part1(input) = first(part1and2(input))
part2(input) = last(part1and2(input))

function part1and2(io)
    return check_lines(io.data)
end

const lut = Int8[-3, 0, 3, -4, 4, -1, 0, 1, 0, 0, 0, 0, 0, 0, -2, 0, 2]
const corrupted_scores = [25137, 1197, 57, 3]

function check_lines(data)
    prevs = Vector{UInt16}(undef, length(data))
    last_opener = 0
    i = 1
    score1 = 0
    scores2 = Int[]
    while i <= length(data)
        c = lut[data[i] % 18]
        if c < 0
            prevs[i] = last_opener | ((-(c + 1)) << 14)
            last_opener = i
        elseif c > 0
            x = prevs[last_opener]
            if x >> 14 != c - 1
                score1 += corrupted_scores[c]
                while data[i] != UInt8('\n')
                    i += 1
                end
                last_opener = 0
            else
                last_opener = x & 0x3FFF
            end
        else
            score2 = 0
            while last_opener != 0
                x = prevs[last_opener]
                score2 = 5 * score2 + 4 - (x >> 14)
                last_opener = x & 0x3FFF
            end
            push!(scores2, score2)
        end
        i += 1
    end
    return score1, partialsort!(scores2, 1 + length(scores2) ÷ 2)
end


text:
using Distributions
using RmathDist

import Distributions: zval

distlist = [
            LogNormal(0,1),
            LogNormal(1e10,1e10),
            LogNormal(1e-10,1e-10)
            ]

for d in distlist
    for x in [logspace(-300,300,1_000_000)]
        z = zval(Normal(d.meanlog,d.sdlog),log(x))

        if x <= 1e298 # rmath underflow
            test_rmath(pdf,d,x,1e4)
            test_rmath(logpdf,d,x,1e5)
        end
        if !(37 < abs(z) < 39) # rmath underflow
            test_rmath(cdf,d,x,1e5)
            test_rmath(ccdf,d,x,1e5)
        end

        test_rmath(logcdf,d,x,1e5)
        test_rmath(logccdf,d,x,1e5)
    end

    for p in logspace(-300,0,1_000_000)
        test_rmath(quantile,d,p,1e5)
        test_rmath(cquantile,d,p,1e5)
    end

    for lp in -logspace(-300,300,1_000_000)
        test_rmath(invlogcdf,d,lp)
        test_rmath(invlogccdf,d,lp)
    end
end


text:
using DifferentialEvolutionMCMC, Test, Random, Turing, Distributions
Random.seed!(973536)

bounds = ((-Inf,Inf),(0.0,Inf))

data = rand(Normal(0.0, 1.0), 100)

function prior_loglike(μ, σ)
    LL = 0.0
    LL += logpdf(Normal(0, 10), μ)
    LL += logpdf(truncated(Cauchy(0, 1), 0, Inf), σ)
    return LL
end

function sample_prior()
    μ = rand(Normal(0, 10))
    σ = rand(truncated(Cauchy(0, 1), 0, Inf))
    return [μ,σ]
end

bounds = ((-Inf,Inf),(0.0,Inf))

data = rand(Normal(0.0, 1.0), 50)

function loglike(data, μ, σ)
    return sum(logpdf.(Normal(μ, σ), data))
end

names = (:μ,:σ)

model = DEModel(;
    sample_prior,
    prior_loglike,
    loglike,
    data,
    names
)

de = DE(;sample_prior, bounds, burnin=1500, Np=6)
n_iter = 3000
chains = sample(model, de, n_iter, progress=true)
μ_de = describe(chains)[1][:,:mean]
σ_de = describe(chains)[1][:,:std]
rhat = describe(chains)[1][:,:rhat]

@model model1(data) = begin
    μ ~ Normal(0, 10)
    σ ~ truncated(Cauchy(0, 1), 0.0, Inf)
    for i in 1:length(data)
        data[i] ~ Normal(μ, σ)
    end
end
chn = sample(model1(data), NUTS(1000, .85), 2000)
μ_nuts = describe(chn)[1][:,:mean]
σ_nuts = describe(chn)[1][:,:std]

@test all(isapprox.(rhat, fill(1.0, 2), atol=.05))
@test all(isapprox.(μ_nuts, μ_de, atol=.01))
@test all(isapprox.(σ_nuts, σ_de, atol=.01))


text:
using ZigZagBoomerang
using LinearAlgebra
using Random
using SparseArrays
using Distributions
using Statistics
using ZigZagBoomerang: idot, idot_moving!
using ReverseDiff
using Test
using FileIO
using GLM

println("Sparse logistic regression")

# Design matrix
Random.seed!(2)
sparsity(A, d = 3) = round(nnz(A)/length(A), digits=d)
include("sparsedesign.jl")
# create mock design matrix with 2 categorical explanatory variables
# and their interaction effects and 2 continuous explanatory variable
A = sparse_design([20,20], 2, 4*20)
n, p = size(A)
@show n, p
Γ = A'*A
@show sparsity(A), sparsity(Γ)

const m = 1 # m Bernoulli samples for each set of covariates
println("Av. number of regressors per column ", mean(sum(A .!= 0, dims=1)), ", row ", mean(sum(A .!= 0, dims=2)))

At = SparseMatrixCSC(A')
γ0 = 0.01
# Data from the model
xtrue = 5*randn(p).*(rand(p) .< 0.1)
sigmoid(x) = inv(one(x) + exp(-x))
lsigmoid(x) = -log(one(x) + exp(-x))
y_ = Matrix(rand(n, m) .< sigmoid.(A*xtrue))
y = sum(y_, dims=2)[:]
ny = m .- y
# prior and potential
# prior precision
ϕ(x, A, y::Vector) = γ0*dot(x,x)/2 - sum(y .* lsigmoid.(A*x) + (m .- y) .* lsigmoid.(-A*x))

# Sparse gradient
# helper functions for sparse gradient
function fdot(A::SparseMatrixCSC, At::SparseMatrixCSC, f, j, x, y)
   rows = rowvals(A)
   vals = nonzeros(A)
   s = zero(eltype(A))
   @inbounds for i in nzrange(A, j)
       s += vals[i]*y[rows[i]]*f(idot(At, rows[i], x))
   end
   s
end

sigmoidn(x) = sigmoid(-x)
nsigmoid(x) = -sigmoid(x)
# helper function for sparse gradient estimate through subsampling
function fdotr(A::SparseMatrixCSC, At::SparseMatrixCSC, j, x, μ, y, ny, k)
   rows = rowvals(A)
   vals = nonzeros(A)
   s = zero(eltype(A))
   r = nzrange(A, j)
   sampler = Random.SamplerRangeNDL(r)
   l = length(r)
   @inbounds for _ in 1:k
       i = rand(sampler)
       u = idot(At, rows[i], x)
       s += l/k*vals[i]*y[rows[i]]*sigmoidn(u)
       s += l/k*vals[i]*ny[rows[i]]*nsigmoid(u)
       u0 = idot(At, rows[i], μ)
       s -= l/k*vals[i]*y[rows[i]]*sigmoidn(u0)
       s -= l/k*vals[i]*ny[rows[i]]*nsigmoid(u0)   end
   s
end

# the same, but advancing only the coordinates necessary with `idot_moving!`
function fdot_moving(A::SparseMatrixCSC, At::SparseMatrixCSC, j, t, x, θ, t′, F, μ, y, ny, k)
   rows = rowvals(A)
   vals = nonzeros(A)
   s = zero(eltype(A))
   r = nzrange(A, j)
   sampler = Random.SamplerRangeNDL(r)
   l = length(r)
   @inbounds for _ in 1:k
       i = rand(sampler)
       u = idot_moving!(At, rows[i], t, x, θ, t′, F)
       s += l/k*vals[i]*y[rows[i]]*sigmoidn(u)
       s += l/k*vals[i]*ny[rows[i]]*nsigmoid(u)
       u0 = idot(At, rows[i], μ)
       s -= l/k*vals[i]*y[rows[i]]*sigmoidn(u0)
       s -= l/k*vals[i]*ny[rows[i]]*nsigmoid(u0)
   end
   s
end

# Gradient (found using http://www.matrixcalculus.org/ S. Laue, M. Mitterreiter, and J. Giesen. A Simple and Efficient Tensor Calculus, AAAI 2020.)
∇ϕ(x, A, At, y) = γ0*x - A'*(y .* sigmoidn.(A*x)) - A'*((m .- y).*nsigmoid.(A*x))

# Element i of the gradient exploiting sparsity
∇ϕ(x, i, A, At, y, ny = m .- y) = γ0*x[i] - fdot(A, At, sigmoidn, i, x, y) - fdot(A, At, nsigmoid, i, x, ny)
# Element i of the gradient exploiting sparsity and random subsampling
∇ϕr(x, i, A, At, μ, y, ny = m .- y, k = 5) = γ0*x[i] - fdotr(A, At, i, x, μ, y, ny, k)
# The same, but exploiting that only dependencies in subsamples need to be evaluated
∇ϕmoving(t, x, θ, i, t′, F, A, At, μ, y, ny = m .- y, k = 5) = γ0*x[i] - fdot_moving(A, At, i, t, x, θ, t′, F, μ, y, ny, k)
# Tests, to be sure
#@test norm(ReverseDiff.gradient(x->ϕ(x, A, y), xtrue) - ∇ϕ(xtrue, A, At, y)) < 1e-7
#@test norm(ReverseDiff.gradient(x->ϕ(x, A, y), xtrue) - [∇ϕ(xtrue, i, A, At, y) for i in 1:p]) < 1e-7

#res = glm(A, y/m, Binomial(m), LogitLink())
#x0 = coef(res)
#norm(xtrue - x0)

# Some Newton steps towards the mode as starting point
println("Newton steps")
x0 = 0.1rand(p)
@time for i in 1:30
    global x0
    H = hcat((sparse(ReverseDiff.gradient(x -> ∇ϕ(x, i, A, At, y, ny), x0)) for i in 1:p)...)
    x0 = x0 - (Symmetric(0.00I + H))\[∇ϕ(x0, i, A, At, y, ny) for i in 1:p]
end
norm(xtrue - x0)

# Note that the A'A has the same sparsity structure as the Hessian of ϕ
# We can use it to tune the doubly local ZigZig with spdmp sparsifying the
# precision matrix estimate
Γ0 = A'*A
nonzeros(Γ0) .= 1

Γ = hcat((sparse(ReverseDiff.gradient(x -> ∇ϕ(x, i, A, At, y), x0)) for i in 1:p)...)
μ = copy(x0)

Γdrop = droptol!(copy(Γ), 1e-2)
@assert nnz(diag(Γdrop)) == p
println("Precision Γ is ", 100*nnz(Γ)/length(Γ), "% filled")
println("Dropped precision is ", 100*nnz(Γdrop)/length(Γdrop), "% filled")

# Random initial values
t0 = 0.0

# Define ZigZag
c = 0.01*ones(p) # Rejection bounds
c0 = copy(c)
σ = sqrt.(diag(inv(Matrix(Γ))))
#σ = (Vector(diag(Γ))).^(-0.5) # cheaper

Z = ZigZag(Γ, μ, σ, ρ=0.5, λref=0.00)
Zdiag = ZigZag(sparse(Diagonal(Γ)), μ, σ, ρ=0.5, λref=0.00)
Zdrop = ZigZag(Γdrop, μ, σ, ρ=0.5, λref=0.00)
T = 2000.0

θ0 = rand([-1.0,1.0], p) .* σ

# Run sparse ZigZag for T time units and collect trajectory
println("Distance starting point")
@show norm(x0 - xtrue)

println("Run spdmp")
trace0, u0, (acc0,num0), c0 = @time spdmp(∇ϕmoving, t0, x0, θ0, T, c, Zdrop, SelfMoving(), A, At, μ, y, m .- y, 10; structured = true, adapt=true, factor=5)

c = 0.01*ones(p)
κ = (γ0/sqrt(2π))/(1/0.9 -1)*ones(p)
trace, u, (acc,num), c = @time sspdmp(∇ϕmoving, t0, x0, θ0, T, c, Zdrop, κ, SelfMoving(), A, At, μ, y, m .- y, 10; structured = true, adapt=true, factor=5)

using Pkg
Pkg.activate(@__DIR__)
cd(@__DIR__)
using Makie, AbstractPlotting
using ColorSchemes, GeometryBasics

@time traj0 = collect(discretize(trace0, 0.1))
@time traj = collect(discretize(trace, 0.1))

xhat0 = [median(getindex.(last.(traj0), i)) for i in eachindex(μ)]
xhat = [median(getindex.(last.(traj), i)) for i in eachindex(μ)]
xcol = xtrue.*(xtrue .!= 0.0) .- (xtrue .== 0.0)*16.0
fi0 = linesegments(repeat(1:p, inner = 2), vec([xhat0 xtrue]'), linewidth = 0.5, resolution = (1200,900))
scatter!(1:p, xhat0, color = xcol, colormap=:berlin, markersize=7, strokewidth = 0.5,
            marker = map(x -> x != 0 ? :circle : :x, xtrue))
fi0
fi = linesegments(repeat(1:p, inner = 2), vec([xhat xtrue]'), linewidth = 0.5, resolution = (1200,900))
scatter!(1:p, xhat, color = xcol, colormap=:berlin, markersize=7,
            strokewidth = 0.5,marker = map(x -> x != 0 ? :circle : :x, xtrue))
fi

save("logistic.png", fi0)
save("logisticsticky.png", fi)

@show acc/num
@show extrema(c ./ c0)

function confusion(xhat)
    Dict(
    "false negative" => sum((xhat .== 0) .& (xtrue .!= 0)),
    "true positive" => sum((xhat .!= 0) .& (xtrue .!= 0)),
    "true negative" => sum((xhat .== 0) .& (xtrue .== 0)),
    "false positive" => sum((xhat .!= 0) .& (xtrue .== 0)),
    )
end

confusion(xhat)

confusion(abs.(xhat0) .> 1.9)

dt = T/4000
x̂ = mean(x for (t,x) in discretize(trace, dt))
X = Float64[]
ts = Float64[]
for (t,x) in discretize(trace, dt)
    append!(ts, t)
    append!(X, x)
end
X = reshape(X, p, length(X)÷p)

[xtrue x̂ abs.(xtrue - x̂)]

@show norm(xtrue - μ)
@show norm(xtrue - x̂)

println("Plot")
using Makie
using Colors
using GoldenSequences

ps = [4, 22, 50, 378, p-2, p-1]
p0 = length(ps)
cs = map(x->RGB(x...), (Iterators.take(GoldenSequence(3), p0)))
pis = []
fig = Figure(resolution = (1200, 1200))
for i in 1:p0
    global p_1
    p_i = fig[i, 1] = Axis(fig)
    if i == 1
        p_1 = p_i
    end
    lines!(p_i, ts, X[ps[i], :], color=cs[i])
    lines!(p_i, [0, T], [xtrue[ps[i]], xtrue[ps[i]]], color=cs[i], linewidth=2.0)
    p_i.ylabel = "var$(ps[i])"
    p_i.xlabel = "t"
    i > 0 && linkaxes!(p_1, p_i)
end
#fig[0, :] = Label(fig, "Sparse design logistic regression n=$(m*n), p=$p", textsize = 30)
fig
save("logistic$(typeof(Z).name).png", fig)

#p2 = title(vbox(hbox([text("$p", show_axis=false, textsize=10) for p in ps]...),
#    [hbox([i==j ? lines(ts, X[i,:]) : lines(X[i, :], X[j, :]) for i in ps]...) for j in ps]...), "$ps")

p1


text:
import Base: quantile, convert
import Distributions.pnormalize!

abstract InferenceAlgorithm
abstract SamplingAlgorithm <: InferenceAlgorithm
abstract MCMCAlgorithm <: SamplingAlgorithm

## Types of algorithm
# - Preimage Partition: pre_bfs
# - Exact Preimage Sample: ___
# - Approximate Preimage Sample: pre_tlmh
# - Exact Conditional Abstract Sample: __
# - Approximate Conditional Abstract Sample: ___
# - Exact Conditional Point Sample
# - Approximate Conditional Point Sample

for finame in ["prob.jl",
               "rand.jl",
               "model.jl"]
    include(joinpath("query", finame))
end


text:
#
# Main Factor Code
#

# THE MOST BASIC ASSUMPTION IS THAT ALL VARIABLES ARE CATEGORICAL AND THEREFORE
# Base.OneTo WORTHY. IF THAT IS VIOLATED, NOTHING WILL WORK

"""
    Factor(dims, potential)

Create a Factor corresponding to the potential.
"""
mutable struct Factor
    dimensions::NodeNames
    potential::Array{Float64} # Unnormalized probability
                              # In most cases this will be a probability

    function Factor(dims::NodeNameUnion, potential::Array{Float64})
        dims = convert(NodeNames, dims)
        _ckeck_dims_unique(dims)

        (length(dims) != ndims(potential)) &&
            throw(DimensionMismatch("`potential` must have as many " *
                        "dimensions as dims"))

        (:potential in dims) &&
            @warn("Having a dimension called `potential` will cause problems")

        return new(dims, potential)
    end
end

"""
    Factor(dims, lengths, fill_value=0)

Create a factor with dimensions `dims`, each with lengths corresponding to
`lengths`. `fill_value` will fill the potential array with that value.
To keep uninitialized, use `fill_value=nothing`.
"""
Factor(dims::NodeNames, lengths::Vector{Int}, ::Nothing) =
    Factor(dims, Array{Float64}(undef, lengths...))

Factor(dims::NodeNames, lengths::Vector{Int}, fill_value::Number=0) =
    Factor(dims, fill(Float64(fill_value), lengths...))

Factor(dim::NodeName, length::Int, ::Nothing) =
    Factor([dims], Array{Float64}(undef, length))

Factor(dim::NodeName, length::Int, fill_value::Number=0) =
    Factor(dims, fill(Float64(fill_value), length))

"""
    convert(DiscreteCPD, cpd)

Construct a Factor from a DiscreteCPD.
"""
function Base.convert(::Type{Factor}, cpd::DiscreteCPD)
    dims = vcat(name(cpd), parents(cpd))
    lengths = tuple(ncategories(cpd), cpd.parental_ncategories...)
    p = Array{Float64}(undef,lengths)
    p[:] = vcat([d.p for d in cpd.distributions]...)
    return Factor(dims, p)
end
Base.showable(::MIME"text/html", ϕ::DiscreteCPD) = true
Base.show(io::IO, cpd::DiscreteCPD) = show(io, convert(Factor, cpd))
Base.show(io::IO, a::MIME"text/html", cpd::DiscreteCPD) = show(io, a, convert(DataFrame, convert(Factor, cpd)))

"""
    Factor(bn, name, evidence=Assignment())

Create a factor for a node, given some evidence.
"""
function Factor(bn::DiscreteBayesNet, name::NodeName, evidence::Assignment=Assignment())
    cpd = get(bn, name)
    ϕ = convert(Factor, cpd)
    return ϕ[evidence]
end

###############################################################################
#                   Methods

"""
    similar(ϕ)

Return a factor similar to `ϕ` with unitialized values
"""
Base.similar(ϕ::Factor) = Factor(ϕ.dimensions, similar(ϕ.potential))

"""
Returns Float64
"""
Base.eltype(ϕ::Factor) = Float64

"""
Names of each dimension
"""
Base.names(ϕ::Factor) = ϕ.dimensions

Base.ndims(ϕ::Factor) = ndims(ϕ.potential)

"""
    size(ϕ, [dims...])

Returns a tuple of the dimensions of `ϕ`
"""
Base.size(ϕ::Factor) = size(ϕ.potential)
Base.size(ϕ::Factor, dim::NodeName) = size(ϕ.potential, indexin(dim, ϕ))
Base.size(ϕ::Factor, dims::Vararg{NodeName, N}) where {N} =
    ntuple(k -> size(ϕ, dims[k]), Val(N))

"""
Total number of elements in Factor (potential)
"""
Base.length(ϕ::Factor) = length(ϕ.potential)

"""
    in(dim, ϕ) -> Bool

Return true if `dim` is in the Factor `ϕ`
"""
Base.in(dim::NodeName, ϕ::Factor) = dim in names(ϕ)

"""
    indexin(dims, ϕ)

Return the index of dimension `dim` in `ϕ`, or 0 if not in `ϕ`.
"""
Base.indexin(dim::NodeName, ϕ::Factor) = findnext(isequal(dim), ϕ.dimensions, 1)
Base.indexin(dims::NodeNames, ϕ::Factor) = indexin(dims, names(ϕ))

"""
    rand!(ϕ)

Fill with random values
"""
Random.rand!(ϕ::Factor) = Random.rand!(ϕ.potential)

"""
Appends a new dimension to a Factor
"""
@inline function Base.push!(ϕ::Factor, dim::NodeName, length::Int)
    if dim in names(ϕ)
        error("Dimension $(dim) already exists")
    end

    p = duplicate(ϕ.potential, (length, ))
    ϕ.dimensions = push!(ϕ.dimensions, dim)
    ϕ.potential = p

    return ϕ
end

@inline function Base.permutedims!(ϕ::Factor, perm)
    ϕ.potential = permutedims(ϕ.potential, perm)
    ϕ.dimensions = ϕ.dimensions[perm]
    return ϕ
end

Base.permutedims(ϕ::Factor, perm) = permutedims!(deepcopy(ϕ), perm)

"""
    pattern(ϕ, [dims])

Return an array with the pattern of each dimension's state for all possible
instances
"""
function pattern(ϕ::Factor, dims)
    inds = indexin(dims, ϕ)

    zero_loc = findfirst(isequal(0), inds)
    if zero_loc != nothing
        not_in_factor_error(dims[zero_loc])
    end

    lens = [size(ϕ)...]

    inners = vcat(1, cumprod(lens[1:(end-1)]))
    outers = Int[(length(ϕ) ./ lens[inds] ./ inners[inds])...]

    hcat([repeat(collect(1:l), inner=i, outer=o) for (l, i, o) in
        zip(lens[inds], inners[inds], outers)]...)
end

function pattern(ϕ::Factor)
    lens = [size(ϕ)...]

    inners = vcat(1, cumprod(lens[1:(end-1)]))
    outers = Int[(length(ϕ) ./ lens ./ inners)...]

    hcat([repeat(collect(1:l), inner=i, outer=o) for (l, i, o) in
            zip(lens, inners, outers)]...)
end


text:
# Thrift SASL transport
#
# Thrift SASL: https://github.com/apache/thrift/blob/master/doc/specs/thrift-sasl-spec.txt
# SASL PLAIN mechanism: https://tools.ietf.org/html/rfc4616

# protocol status codes (also represent SASL negotiation state)
const SASL_ZERO     = 0x00
const SASL_START    = 0x01
const SASL_OK       = 0x02
const SASL_BAD      = 0x03
const SASL_ERROR    = 0x04
const SASL_COMPLETE = 0x05

# sasl mechanism names
const SASL_MECH_PLAIN   = "PLAIN"
const SASL_MECH_KERB    = "KERBEROS"
const SASL_MECH_LDAP    = "LDAP"
const SASL_MECHANISMS   = (SASL_MECH_PLAIN, SASL_MECH_KERB, SASL_MECH_LDAP)

# sasl exception types
const SASL_ERR_UNKNOWN      = 0
const SASL_ERR_UNSUPPORTED  = 1
const SASL_ERR_NEGOTIATION  = 2
const SASL_ERR_INVALID      = 3

type SASLException <: Exception
    code::Int
    message::AbstractString

    SASLException(code=SASL_ERR_UNKNOWN, message::AbstractString="") = new(code, message)
end

validate_sasl_mech(mech) = (mech in SASL_MECHANISMS) || throw(SASLException(SASL_ERR_UNSUPPORTED, "Unsupported SASL mechanism \"$mech\""))

validate_sasl_status(status::UInt8, message::Vector{UInt8}, okstatus::Tuple) = validate_sasl_status(status, Compat.String(message), okstatus)
function validate_sasl_status(status::UInt8, message::AbstractString, okstatus::Tuple)
    (status != SASL_BAD) && (status != SASL_ERROR) && (status in okstatus) && return

    if status == SASL_BAD
        msg = "SASL mechanism not supported"
        err = SASL_ERR_UNSUPPORTED
    elseif status == SASL_ERROR
        msg = "SASL negotiation failed"
        err = SASL_ERR_NEGOTIATION
    else
        msg = "SASL negotiation failed: unexpected status code $status received from server"
        err = SASL_ERR_NEGOTIATION
    end

    if !isempty(message)
        msg = string(msg, ": ", message)
    end
    throw(SASLException(err, msg))
end

# SASL messages are of the form:
# | 1-byte status code | 4-byte payload length | variable-length payload |
sasl_write(io::IO, status::UInt8, payload::AbstractString) = sasl_write(io, status, convert(Vector{UInt8}, payload))
function sasl_write(io::IO, status::UInt8, payload::Vector{UInt8}=UInt8[])
    len = length(payload)
    iob = IOBuffer()
    write(iob, status)
    _write_fixed(iob, UInt32(len), true)
    (len > 0) && write(iob, payload)
    write(io, takebuf_array(iob))
end

function sasl_read(io::IO)
    status = read(io, UInt8)
    @logmsg("read_sasl read status $status")
    len = _read_fixed(io, UInt32(0), 4, true)
    @logmsg("read_sasl read len $len")
    data = read(io, UInt8, len)
    @logmsg("read_sasl read data $data")
    (status, len, data)
end

function sasl_negotiate_plain(io::IO, callback::Function)
    # prepare credentials
    authzid = callback(:authzid)
    authcid = callback(:authcid)
    password = callback(:passwd)
    (isempty(authcid) || isempty(password)) && throw(SASLException(SASL_ERR_INVALID, "Invalid or empty credentials"))
    (max(length(authcid), length(password), length(authzid)) > 255) && throw(SASLException(SASL_ERR_INVALID, "Credentials too large"))

    creds = IOBuffer()
    isempty(authzid) || creds.write(convert(Vector{UInt8}, authzid))
    write(creds, 0x00)
    write(creds, convert(Vector{UInt8}, authcid))
    write(creds, 0x00)
    write(creds, convert(Vector{UInt8}, password))

    # start negotiation, indicate protocol
    nbyt = sasl_write(io, SASL_START, SASL_MECH_PLAIN)
    @logmsg("negotiate_sasl wrote $nbyt bytes")
    # send credentials
    nbyt = sasl_write(io, SASL_OK, takebuf_array(creds))
    @logmsg("negotiate_sasl wrote $nbyt bytes")

    (status, len, data) = sasl_read(io)
    @logmsg("negotiate_sasl read $status, $len, $data")
    validate_sasl_status(status, data, (SASL_COMPLETE,))

    # send COMPLETE
    #nbyt = sasl_write(io, SASL_COMPLETE, "")
    #@logmsg("negotiate_sasl wrote $nbyt bytes")

    nothing
end

function sasl_negotiate(io::IO, mech, callback::Function)
    if mech == SASL_MECH_PLAIN
        return sasl_negotiate_plain(io, callback)
    end
    throw(SASLException(SASL_ERR_UNSUPPORTED, "Unsupported mechanism $mech"))
end

function sasl_callback_default(part::Symbol)
    (part == :authcid) && (return get(ENV, "USER", ""))
    (part == :passwd) && (return "password")
    (part == :show) && (return get(ENV, "USER", ""))
    (part == :mechanism) && (return "SASL-Plain")

    return "" # for authzid
end


text:
"""
    SeisRadonFreqFor(in, nt; <keyword arguments>)

Transform a tau-p gather to a time-offset gather using a frequency domain
forward parabolic or linear Radon operator.

# Arguments
- `in::Array{Float64,2}`: 2D Radon panel, `in[1:ntau,1:np]`, where `ntau` is the
number of intercept times and `np` the number of curvatures or ray parameters.
- `nt::Int`: number of time samples.

# Keyword arguments
- `order="parab"`: `"parab"` for parabolic transform, `"linear"`
for linear transform.
- `dt=0.004`: time sampling interval in seconds.
- `h=collect(0.0:20.0:1000.0)`: offset vector; `h[1:nh]`.
- `href=0.0`: reference offset for parabolic Radon Transform. If the
defautl value `href=0.0` is given, `href` is set to `max(abs(h))`.
- `p=collect(-0.05:0.01:2.2)`: `p[1:np]`. If `order="parab"`, `p`
is a vector of residual moveout ("curvatures") at reference offset `href` in
seconds; if `order=linear`, `p` is a vector of ray parameters in s/m.
- `flow=0.0`: minimum frequency in the data in Hz.
- `fhigh=125.0`: maximum frequency in the data in Hz.

# Output
- `d`: data synthetized via forward Radon modeling, `d[1:nt, 1:nh]`.

# References
*  Hampson, D., 1986, Inverse velocity stacking for multiple elimination:
Canadian Journal of Exploration Geophysics, 22, 44-55.
*  Sacchi, M.D. and Ulrych, T.J., 1995, High-resolution velocity gathers and
offset space reconstruction: Geophysics, 60, 1169-1177.
"""
function SeisRadonFreqFor(m::Array{Tm,2}, nt::Int; order="parab",
                            dt=0.004, href=0.0,
                            h=collect(0.0:20.0:1000.0),
                            p=collect(-0.05:0.01:2.2),
                            flow=0.0, fhigh=125.0)where{Tm<:Real}

    if order=="parab"
        ind = 2
        href == 0 && (href = maximum(abs.(h)))
    elseif order=="linear"
        ind = 1
        href = 1.0
    else
        error("Order should be equal to \"parab\" or \"linear\"")
    end
    ntau = size(m, 1)
    np = length(p)
    np == size(m, 2) || error("lenght(p) must be equal to size(m, 2)")
    nh = length(h)
    nw = 2*nextpow(2,ntau)
    m = cat( m, zeros(nw-ntau, np),dims=1)
    M = fft(m, 1)
    iw_low = round(Int, flow*dt*nw+1)
    iw_high = round(Int, fhigh*dt*nw+1)
    D = zeros(ComplexF64, nw, nh)
    for iw = iw_low:iw_high
        w  = 2.0*pi*(iw-1)/(nw*dt)
        L = zeros(ComplexF64, nh, np)
        for ip = 1:np
            for ih = 1:nh
                phi = w*p[ip]*(h[ih]/href)^ind
      	        L[ih, ip] = exp(-im*phi)
            end
        end
        D[iw, :] = L*M[iw,:]
    end
    for iw = round(Int, nw/2)+2:nw
        D[iw, :] = conj(D[nw-iw+2, :])
    end
    d = real(ifft(D, 1))
    d = d[1:nt, :]
    return d
end


text:
# Julia wrapper for header: /usr/include/scip/pub_implics.h
# Automatically generated using Clang.jl wrap_c

function SCIPcliqueSearchVar(clique, var, value)
    ccall((:SCIPcliqueSearchVar, libscip), Cint, (Ptr{SCIP_CLIQUE}, Ptr{SCIP_VAR}, UInt32), clique, var, value)
end

function SCIPcliqueHasVar(clique, var, value)
    ccall((:SCIPcliqueHasVar, libscip), UInt32, (Ptr{SCIP_CLIQUE}, Ptr{SCIP_VAR}, UInt32), clique, var, value)
end

function SCIPcliqueGetNVars(clique)
    ccall((:SCIPcliqueGetNVars, libscip), Cint, (Ptr{SCIP_CLIQUE},), clique)
end

function SCIPcliqueGetVars(clique)
    ccall((:SCIPcliqueGetVars, libscip), Ptr{Ptr{SCIP_VAR}}, (Ptr{SCIP_CLIQUE},), clique)
end

function SCIPcliqueGetValues(clique)
    ccall((:SCIPcliqueGetValues, libscip), Ptr{UInt32}, (Ptr{SCIP_CLIQUE},), clique)
end

function SCIPcliqueGetId(clique)
    ccall((:SCIPcliqueGetId, libscip), UInt32, (Ptr{SCIP_CLIQUE},), clique)
end

function SCIPcliqueGetIndex(clique)
    ccall((:SCIPcliqueGetIndex, libscip), Cint, (Ptr{SCIP_CLIQUE},), clique)
end

function SCIPcliqueIsCleanedUp(clique)
    ccall((:SCIPcliqueIsCleanedUp, libscip), UInt32, (Ptr{SCIP_CLIQUE},), clique)
end

function SCIPcliqueIsEquation(clique)
    ccall((:SCIPcliqueIsEquation, libscip), UInt32, (Ptr{SCIP_CLIQUE},), clique)
end


text:
VERSION < v"0.7.0-beta2.199" && __precompile__()

module MDCT

using LinearAlgebra
import AbstractFFTs
export mdct, imdct, plan_mdct, plan_imdct
import Base: *, size
import LinearAlgebra: mul!

using FFTW
import FFTW: fftwNumber, r2r, r2r!, REDFT11, plan_r2r!, plan_r2r

fftwsimilar(X::AbstractArray{T}, sz) where {T<:fftwNumber} = Array{T}(undef, sz...)
fftwsimilar(X::AbstractArray{T}, sz) where {T<:Real} = Array{Float64}(undef, sz...)
fftwsimilar(X::AbstractArray{T}, sz) where {T<:Complex} = Array{Complex128}(undef, sz...)

# The following two routines compute the MDCT and IMDCT via
# a type-IV DCT (FFTW's REDFT11 r2r transform).  For a review
# of this relationship, see the notes I posted on Wikipedia:
#     http://en.wikipedia.org/wiki/Modified_discrete_cosine_transform

function mdct(X::AbstractVector{T}) where {T<:Number}
    sz = length(X)
    if isodd(sz)
        throw(ArgumentError("mdct requires an even-length vector"))
    end
    N = div(sz, 2);
    Y = fftwsimilar(X, N)
    if isodd(N)
        throw(ArgumentError("mdct requires a multiple-of-4 vector length"))
        # FIXME: handle odd case via DCT-III?
    else
        N2 = div(N,2)
        for i = 1:N2
            Y[i] = -0.5 * (X[(3*N2+1)-i] + X[3*N2+i])
            Y[N2+i] = -0.5 * (X[(2*N2+1)-i] - X[i])
        end
        r2r!(Y, REDFT11)
    end
    return Y
end

function imdct(X::StridedVector{T}) where {T<:Number}
    N = length(X)
    Z = fftwsimilar(X, 2*N)
    if isodd(N)
        throw(ArgumentError("imdct requires an even vector length"))
        # FIXME: handle odd case via DCT-II?
    else
        Y = r2r(X, REDFT11)
        N2 = div(N,2)
        s = 0.5 / N
        for i = 1:N2
            Z[N+1-i] = -(Z[i] = Y[N2+i]*s)
            Z[N+N2+1-i] = (Z[N+N2+i] = -Y[i]*s)
        end
    end
    return Z
end

imdct(X::AbstractVector{T}) where {T<:Number} =
    imdct(copy!(fftwsimilar(X, size(X)), X))

mutable struct Plan{T<:fftwNumber, N, inv} <: AbstractFFTs.Plan{T}
    plan::FFTW.r2rFFTWPlan{T} # plan_r2r! REDFT11 plan
    pinv::AbstractFFTs.Plan{T}
    Plan{T, N, inv}(plan) where {T, N, inv} = new(plan)
end

size(p::Plan{T, N, true}) where {T<:Number, N} = (div(N,2), N)
size(p::Plan{T, N, false}) where {T<:Number, N} = (2*N, N)

function mul!(Y::StridedVector{T}, p::Plan{T, N, false}, X::AbstractVector{T}) where {T<:Number, N}
    @boundscheck (length(X) == 2*N && length(Y) == N) || throw(DimensionMismatch())
    N2 = div(N,2)
    @inbounds for i = 1:N2
        Y[i] = -0.5 * (X[(3*N2+1)-i] + X[3*N2+i])
        Y[N2+i] = -0.5 * (X[(2*N2+1)-i] - X[i])
    end
    return p.plan * Y
end

function mul!(Z::StridedVector{T}, p::Plan{T, N, true}, X::StridedVector{T}) where {T<:Number, N}
    @boundscheck length(Z) == N || throw(DimensionMismatch())
    Y = p.plan*X
    N1 = div(N,2)
    N2 = div(N1,2)
    s = 0.5 / N1
    @inbounds for i = 1:N2
        Z[N1+1-i] = -(Z[i] = Y[N2+i]*s)
        Z[N1+N2+1-i] = (Z[N1+N2+i] = -Y[i]*s)
    end
    return Z
end

function *(p::Plan{T, N}, X::AbstractVector{T}) where {T<:Number, N}
    Y = fftwsimilar(X, N)
    return mul!(Y, p, X)
end

function plan_mdct(X::AbstractVector{T}, flags::Integer=FFTW.ESTIMATE, timelimit::Real=FFTW.NO_TIMELIMIT) where {T<:Number}
    sz = length(X)
    if mod(sz, 4) != 0
        throw(ArgumentError("mdct requires a multiple-of-4 vector length"))
        # FIXME: handle odd case via DCT-III?
    end
    N = div(sz, 2);
    return Plan{T, N, false}(plan_r2r!(fftwsimilar(X, N), REDFT11; flags=flags, timelimit=timelimit))
end

function plan_imdct(X::AbstractVector{T}, flags::Integer=FFTW.ESTIMATE, timelimit::Real=FFTW.NO_TIMELIMIT) where {T<:Number}
    sz = length(X)
    if isodd(sz)
        throw(ArgumentError("imdct requires an even vector length"))
        # FIXME: handle odd case via DCT-II?
    end
    return Plan{T, 2sz, true}(plan_r2r(fftwsimilar(X, sz), REDFT11; flags=flags, timelimit=timelimit))
end

end # MDCT


text:
using GeometricAlgebra

@testset "trig" begin
	t, x, y, z = basis((-1, 1, 1, 1))

	@testset "scalar square" begin
		@test exp(+x) ≈ cosh(1) + sinh(1)x
		@test exp(0y) ≈ 1
		@test exp(-z) ≈ cosh(1) - sinh(1)z

		@test exp(x*y) ≈ cos(1) + sin(1)x*y
	end

	@testset "identities: $label" for (label, objects) ∈ [
		"scalar square" => (x, 2y*x, x + y, x*y + 3y*z),
		# convergence of series can be spooky, so for these
		# proof-of-concept tests, only try for small multivectors
		"non-scalar square" => (t*x - y*z, 1 + 2x)./5,
	]
		for a ∈ objects
			@test exp(a) ≈ cosh(a) + sinh(a)

			@test tan(a) ≈ sin(a)/cos(a)
			@test tanh(a) ≈ sinh(a)/cosh(a)
			@test cot(a) ≈ cos(a)/sin(a)
			@test coth(a) ≈ cosh(a)/sinh(a)

			@test cos(a)^2 + sin(a)^2 ≈ 1
			@test 1 + tan(a)^2 ≈ sec(a)^2
			@test cot(a)^2 + 1 ≈ csc(a)^2

			@test cosh(a)^2 - sinh(a)^2 ≈ 1
			@test 1 - tanh(a)^2 ≈ sech(a)^2
			@test coth(a)^2 - 1 ≈ csch(a)^2
		end
	end

end


text:
using CSV
using TimeSeries
using DataFrames

###### John Hopkins University data loader #################

_load_data(source::Val{:JHU}, region::Union{AbstractString,Regex}, path::AbstractString) = _load_data(source, (r".*", region, "global"), path)

_load_data(source::Val{:JHU}, region::NTuple{2,Union{AbstractString,Regex}}, path::AbstractString) = _load_data(source, (region..., "global"), path)

function _load_data(source::Val{:JHU}, region::Tuple{Union{AbstractString,Regex},Union{AbstractString,Regex},AbstractString}, path::AbstractString)

    basepath = joinpath(path,"COVID-19/csse_covid_19_data/csse_covid_19_time_series")
    if region[3] == "global"
        fns = Dict(:Confirmed => "time_series_covid19_confirmed_global.csv",
                      :Deaths => "time_series_covid19_deaths_global.csv")
        region_cols = map(Symbol,("Province/State","Country/Region"))
        date_cols_s = 5
    elseif region[3] == "US"
        fns = Dict(:Confirmed => "time_series_covid19_confirmed_US.csv",
                      :Deaths => "time_series_covid19_deaths_US.csv")
        region_cols = map(Symbol,("Admin2","Province_State"))
        date_cols_s = 13
    else
        error("region $region unknown")
    end

    ta = Dict{Symbol,TimeArray}()
    for (key,fn) in fns
        fp = joinpath(basepath,fn)
        df = CSV.read(fp,copycols=true)
        let col = region_cols[1]
            rows = ismissing.(df[!,col]) .| (df[!,col] .== "Unassigned")
            df[rows,col] .= ""
        end

        dates = Date.(String.(names(df)[date_cols_s:end]),"m/d/y") + Dates.Year(2000)
        sub = trues(size(df,1))
        for (i,col) in enumerate(region_cols)
            sub = sub .& occursin.(region[i],df[!,col])
        end
        if any(sub)
            ta[key] = TimeArray(dates,sum(Array{Int64}(df[sub,date_cols_s:end]),dims=1)[1,:],[key,])
        else
            error("no entry in `$fp' for $region")
        end
    end

    merge(values(ta)...)
end


text:
# epsxc only version
function XC_x_pbe_E( rho, grho )

    third = 1.0/3.0
    k = 0.804
    mu = 0.2195149727645171
    c1 = 0.75/pi
    c2 = 3.093667726280136
    c5 = 4.0*third

    agrho = sqrt(grho)
    kf = c2 * rho^third
    dsg = 0.5/kf
    s1 = agrho*dsg / rho
    s2 = s1 * s1
    ds = -c5 * s1

    # Energy
    f1 = s2 * mu/k
    f2 = 1.0 + f1
    f3 = k / f2
    fx = k - f3
    exunif = -c1 * kf
    sx = exunif * fx * rho

    return sx
end


text:
# Load -
include("Include.jl")

# Load the data dictionary -
data_dictionary = DataDictionary(0.0,0.0,0.0);

# Load the parameter array -
parameter_array = readdlm("./parameter_best_v1.dat.1")

# Update data dictionary to match new parameters before calculating obj
parameter_mapping_array = data_dictionary["parameter_name_mapping_array"]
for index = 1:length(parameter_mapping_array)
  if index <= data_dictionary["number_of_binding"]
    data_dictionary["binding_parameter_dictionary"][parameter_mapping_array[index]] = parameter_array[index]
  else
    data_dictionary["control_parameter_dictionary"][parameter_mapping_array[index]] = parameter_array[index]
  end
end

# Setup the time -
time_start = 0.0
time_stop = 120.0
time_step_size = 0.01

# Run the simulation -
(T,X) = Simulation(time_start,time_stop,time_step_size,data_dictionary);


text:
using GroupsCore
import Random

include(joinpath(pathof(GroupsCore), "..", "..", "test", "cyclic.jl"))
StarAlgebras.star(g::GroupElement) = inv(g)

@testset "Arithmetic" begin
    G = CyclicGroup(6)
    b = StarAlgebras.Basis{UInt8}(collect(G))
    l = length(b)
    RG = StarAlgebra(G, b, (l, l))

    @test sprint(show, RG) == "*-algebra of Group of residues modulo 6"

    @testset "Module structure" begin
        a = AlgebraElement(ones(Int, order(G)), RG)

        @test -a isa AlgebraElement
        @test coeffs(-a) == -coeffs(a)

        @test 2 * a isa AlgebraElement
        @test eltype(2 * a) == typeof(2)
        @test coeffs(2 * a) == 2coeffs(a)

        @test 2.0 * a isa AlgebraElement
        @test eltype(2.0 * a) == typeof(2.0)
        @test coeffs(2.0 * a) == 2.0 * coeffs(a)

        @test coeffs(a / 2) == coeffs(a) / 2
        b = a / 2
        @test b isa AlgebraElement
        @test eltype(b) == typeof(1 / 2)
        @test coeffs(b / 2) == 0.25 * coeffs(a)

        c = a // 1

        @test eltype(c) == Rational{Int}
        @test c // 4 isa AlgebraElement
        @test c // big(4) isa AlgebraElement
        @test eltype(c // (big(4) // 1)) == Rational{BigInt}

        @test (1.0a) * 1 // 2 == (0.5a) == c // 2
    end

    @testset "Additive structure" begin
        a = AlgebraElement(ones(Int, order(G)), RG)
        b = sum((-1)^isodd(g.residual) * RG(g) for g in G)

        @test a == AlgebraElement(ones(Int, order(G)), RG) == sum(RG(g) for g in G)

        @test 1 / 2 * (a + b).coeffs == [1.0, 0.0, 1.0, 0.0, 1.0, 0.0]

        g = CyclicGroupElement(2, G)
        h = CyclicGroupElement(3, G)
        k = CyclicGroupElement(5, G)

        a = RG(1) + RG(g) + RG(h)
        b = RG(1) - RG(k) - RG(h)

        @test a - b == RG(g) + RG(k) + 2RG(h)

        @test a + b - 2a == b - a

        @test 1 // 2 * 2a == a
        @test a + 2a == (3 // 1) * a
        @test 2a - (1 // 1) * a == a
    end

    @testset "Multiplicative structure" begin

        for g in G, h in G
            a = RG(g)
            b = RG(h)
            @test a * b == RG(g * h)
            @test (a + b) * (a + b) == a * a + a * b + b * a + b * b
        end

        for g in G
            @test star(RG(g)) == RG(inv(g))
            @test (one(RG) - RG(g)) * star(one(RG) - RG(g)) ==
                  2 * one(RG) - RG(g) - RG(inv(g))
            @test aug(one(RG) - RG(g)) == 0
        end

        g = CyclicGroupElement(2, G)
        h = CyclicGroupElement(3, G)
        k = CyclicGroupElement(5, G)

        a = RG(1) + RG(g) + RG(h)
        b = RG(1) - RG(k) - RG(h)

        # a = RG(1) + RG(perm"(2,3)") + RG(perm"(1,2,3)")
        # b = RG(1) - RG(perm"(1,2)(3)") - RG(perm"(1,2,3)")

        @test a * b == StarAlgebras.mul!(a, a, b)

        @test aug(a) == 3
        @test aug(b) == -1
        @test aug(a) * aug(b) == aug(a * b) == aug(b * a)

        z = sum((one(RG) - RG(g)) * star(one(RG) - RG(g)) for g in G)
        @test aug(z) == 0

        @test supp(z) == basis(parent(z))
        @test supp(RG(1) + RG(g)) == [one(G), g]
        @test supp(a) == [one(G), g, h]

        if false
            @testset "Projections in Symm(3)" begin
                G = SymmetricGroup(3)
                b = StarAlgebras.Basis{UInt8}(collect(G))
                l = length(b)

                RG = StarAlgebra(G, b)
                @test RG isa StarAlgebra

                P = sum(RG(g) for g in b) // l
                @test P * P == P

                P3 = 2 * sum(RG(g) for g in b if sign(g) > 0) // l
                @test P3 * P3 == P3

                PAlt = sum(sign(g) * RG(g) for g in b) // l
                @test PAlt * PAlt == PAlt

                @test P3 * PAlt == PAlt * P3

                P2 = (RG(1) + RG(b[2])) // 2
                @test P2 * P2 == P2

                @test P2 * P3 == P3 * P2 == P

                P2m = (RG(1) - RG(b[2])) // 2
                @test P2m * P2m == P2m

                @test P2m * P3 == P3 * P2m == PAlt
                @test iszero(P2m * P2)
            end
        end
    end

    @testset "Mutable API and trivial mstructure" begin

        A = [:a, :b, :c]
        b = StarAlgebras.Basis{UInt16}(words(A, radius = 8))
        l = findfirst(w -> length(w) > 4, b) - 1

        RG = StarAlgebra(one(first(b)), b)

        @test basis(RG) === b
        @test basis(RG.mstructure) === basis(RG)

        RGc = StarAlgebra(one(first(b)), b, (l, l))
        @test basis(RGc) === b
        @test basis(RGc.mstructure) === basis(RGc)

        @test all(RG.mstructure[1:121, 1:121] .== RGc.mstructure)

        Z = zero(RG)
        W = zero(RGc)

        let g = b[rand(1:121)]
            X = RG(g)
            Y = -RG(star(g))
            for i in 1:3
                X[b[rand(1:121)]] += rand(-3:3)
                Y[b[rand(1:121)]] -= rand(3:3)
            end

            Xc = AlgebraElement(coeffs(X), RGc)
            Yc = AlgebraElement(coeffs(Y), RGc)

            @test coeffs(X * Y) == coeffs(Xc * Yc) == coeffs(StarAlgebras.mul!(Z, X, Y))

            @test coeffs(X^2) == coeffs(Xc^2) == coeffs(X * X)
            @test coeffs(Y^2) == coeffs(Yc^2) == coeffs(Y * Y)

            @test coeffs(Z) ==
                  StarAlgebras.mul!(coeffs(W), coeffs(X), coeffs(Y), RG.mstructure)
            @test coeffs(Z) == coeffs(W)
            @test coeffs(Z) ==
                  StarAlgebras.mul!(coeffs(W), coeffs(X), coeffs(Y), RGc.mstructure)
            @test coeffs(Z) == coeffs(W)

            StarAlgebras.zero!(W)
            StarAlgebras.fmac!(coeffs(W), coeffs(X), coeffs(Y), RG.mstructure)

            @test coeffs(2 * X * Y) == coeffs(StarAlgebras.mul!(W, W, 2))
        end
    end
end

@testset "Group Algebra caching" begin
    A = [:a, :b, :c]
    b = StarAlgebras.Basis{UInt8}(words(A, radius = 4))
    k = findfirst(w -> length(w) == 3, b) - 1

    RG = StarAlgebra(Word(A, Int[]), b, (k, k))
    @test RG isa StarAlgebra

    D = sum(RG(b[i]) for i in 1:k)
    @test D isa AlgebraElement
    g = one(RG)
    @test isone(g)

    @test one(RG) == g
    @test iszero(zero(RG))
    @test 0 * g == zero(RG)
    @test iszero(0 * g)

    h = RG(b[3])

    @test D * one(RG) == D
    @test one(RG) * D == D

    @test supp(D) == b[1:k]

    @test_throws StarAlgebras.ProductNotDefined StarAlgebras._check(RG.mstructure)

    @test D * D isa AlgebraElement

    @test StarAlgebras._check(RG.mstructure)

    @test all(!iszero, RG.mstructure.table)

    RG = StarAlgebra(Word(A, Int[]), b, (k, k), precompute=true)
    @test all(!iszero, RG.mstructure.table)
end


text:
using Distributed
addprocs(2) # add 2 workers

@everywhere using Pkg
@everywhere Pkg.activate(joinpath(@__DIR__, "..", "examples"))

@everywhere using ProgressiveHedging
@everywhere import JuMP
@everywhere import Ipopt

@everywhere function two_stage_model(scenario_id::ScenarioID)

    model = JuMP.Model(()->Ipopt.Optimizer())
    JuMP.set_optimizer_attribute(model, "print_level", 0)
    JuMP.set_optimizer_attribute(model, "tol", 1e-12)
    JuMP.set_optimizer_attribute(model, "acceptable_tol", 1e-12)

    scen = value(scenario_id)

    ref = JuMP.@variable(model, x >= 0.0)
    stage1 = [ref]

    ref = JuMP.@variable(model, y >= 0.0)
    stage2 = [ref]

    b_s = scen == 0 ? 11.0 : 4.0
	c_s = scen == 0 ? 0.5 : 10.0

    JuMP.@constraint(model, x + y == b_s)

    JuMP.@objective(model, Min, 1.0*x + c_s*y)

    return JuMPSubproblem(model,
                          scenario_id,
                          Dict(stid(1) => stage1,
                               stid(2) => stage2)
                          )
end

scen_tree = two_stage_tree(2)

(niter, abs_res, rel_res, obj, soln_df, phd) = solve(scen_tree,
                                                     two_stage_model,
                                                     ScalarPenaltyParameter(1.0)
                                                     )
@show niter
@show abs_res
@show rel_res
@show obj
@show soln_df


text:
tests = ["test/HopfieldNets.jl", "test/storkey.jl"]

for path in tests
	include(path)
end


text:
# p = (1, 0, 0)
p = zeros(4)
p[2] = 1.0
p[3] = 0.0
p[4] = 0.0

# v = (0, 0, 1)
v = zeros(3)
v[3] = 1.0

# theta
# 右手系と逆になる
# Mz > 0 ならば負の方向に回転
theta = -pi/2

# q = (cos(t/2),  vx*sin(t/2),  vy*sin(t/2),  vz*sin(t/2))
q = zeros(4)

q[1] = cos(theta/2)
q[2] = v[1]*sin(theta/2)
q[3] = v[2]*sin(theta/2)
q[4] = v[3]*sin(theta/2)

A = zeros(3,3)

A[1,1] = 1 - 2*q[3]^2 - 2*q[4]^2
A[1,2] = 2*q[2]*q[3] + 2*q[1]*q[4]
A[1,3] = 2*q[2]*q[4] - 2*q[1]*q[3]

A[2,1] = 2*q[2]*q[3] - 2*q[1]*q[4]
A[2,2] = 1 - 2*q[4]^2 - 2*q[2]^2
A[2,3] = 2*q[3]*q[4] + 2*q[1]*q[2]

A[3,1] = 2*q[2]*q[4] + 2*q[1]*q[3]
A[3,2] = 2*q[3]*q[4] - 2*q[1]*q[2]
A[3,3] = 1 - 2*q[2]^2 - 2*q[3]^2

newx = A[1,1]*p[2] + A[1,2]*p[3] + A[1,3]*p[4]
newy = A[2,1]*p[2] + A[2,2]*p[3] + A[2,3]*p[4]
newz = A[3,1]*p[2] + A[3,2]*p[3] + A[3,3]*p[4]

println(newx)
println(newy)
println(newz)

p[2] = newx
p[3] = newy
p[4] = newz

newx = A[1,1]*p[2] + A[1,2]*p[3] + A[1,3]*p[4]
newy = A[2,1]*p[2] + A[2,2]*p[3] + A[2,3]*p[4]
newz = A[3,1]*p[2] + A[3,2]*p[3] + A[3,3]*p[4]

println(newx)
println(newy)
println(newz)


text:
"""
    results_dict(suite, submission, solution)
    results_dict(testcase, submission, solution)

Convert the test suite or case into a `Dict` form.

This is used to create a `json`able representation of a test suite/case.
"""
function results_dict(
    suite::TestSuite,
    submission::TestSuiteResult,
    solution::TestSuiteResult,
)
    test_cases = [
        results_dict(
            case,
            submission.results[i],
            solution.results[i],
        )
        for (i, case) in enumerate(suite.cases)
    ]
    return Dict(
        "passed" => all(case -> case["passed"], test_cases),
        "testCases" => test_cases,
    )
end

function results_dict(
    case::TestCase,
    submission::TestCaseResult,
    solution::TestCaseResult,
)
    out, err = truncatestring.([submission.stdout, submission.stderr])
    if solution.exception !== nothing
        return Dict(
            "passed" => false,
            "error" => "Solution code threw an exception on the input.",
            "errorDescription" => "$(solution.exception)\n$(solution.backtrace)",
        )
    end
    if submission.exception !== nothing
        return Dict(
            "passed" => false,
            "error" => "Submission code threw an exception on the input.",
            "errorDescription" => "$(submission.exception)\n$(submission.backtrace)",
            "stdout" => out,
            "stderr" => err,
        )
    end

    delta = δ(solution.result, submission.result)
    return Dict(
        "passed" => delta < case.tolerance,
        "delta" => delta,
        "result" => smallrepr(submission.result),
        "stdout" => out,
        "stderr" => err,
    )
end


text:
module ConditionalThreads

using Base.Threads
using Base.Threads: threadid, threading_run

"""Conditionally apply multi-threading to `for` loops.

This is a variation on `Base.Threads.@threads` that adds a run-time boolean
flag to enable or disable threading. It is intended for *internal use* in
packages building on `QuantumControlBase`.

Usage:

```julia
using QuantumControlBase.ConditionalThreads: @threadsif

function optimize(objectives; use_threads=true)
    @threadsif use_threads for k = 1:length(objectives)
    # ...
    end
end
```
"""
macro threadsif(cond, loop)
    if !(isa(loop, Expr) && loop.head === :for)
        throw(ArgumentError("@threadsif requires a `for` loop expression"))
    end
    if !(loop.args[1] isa Expr && loop.args[1].head === :(=))
        throw(ArgumentError("nested outer loops are not currently supported by @threadsif"))
    end
    quote
        if $(esc(cond))
            $(Threads._threadsfor(loop.args[1], loop.args[2], :static))
        else
            $(esc(loop))
        end
    end
end

end


text:
export
    AtiyahBottFormula,
    AtiyahBottFormulaForGraph

const  number_trees = [1, 1, 2, 3, 6, 11, 23, 47, 106, 235, 551, 1301, 3159]
#the number of non-isomorphic graphs with given number of vertices (starting from 2)

# ProgressData contains the data necessary to keep the progress bar up-to-date
mutable struct ProgressData
    progress_bar::Progress
    current_graph::Int64
    top_aut::Int64
    threshold::Int64
end

# create trivial progress data in case the user does not want progress bar
function EmptyProgressData()::ProgressData
    P::Progress = Progress(0)
    P.enabled = false
    return ProgressData(P, 0, 0, 0)
end

# the following function performs the computation of the Atiyah-Bott formula
# for a particular graph.
"""
    AtiyahBottFormulaForGraph(g, pruf_str, aut, n, deg, n_marks, P, s)

Apply the Atiyah-Bott formula to all colorations of a specific graph. It is useful for splitting the computation in multiple parts, to be computed in single threads.
"""
function AtiyahBottFormulaForGraph(g::SimpleGraph, pruf_str::String,
    aut::Int64, n::Int64, deg::Int64, n_marks::Int64, P, s::Vector{fmpq},
    progress_data::ProgressData = EmptyProgressData())::Vector{fmpq}

    local weights::Vector{Vector{Int64}} = get_weights(nv(g)-1, deg) #the array of array of weights

    local n_results::Int64 = length(P)   #this store the number of final results of our computation
    local result::Vector{fmpq} = zeros(fmpq, n_results) #the array of final results

    #for N in 1:n

        from_file, file_name = exists_file_with_colorings( pruf_str, n ) #n <- N

        if from_file
            cols = graph_coloring_from_file( file_name )
        else
            #println( "graph not found!!!" )
            cols = graph_coloring( g, UInt8(n+1) )   #n <- N
        end

        for c in cols   #we run among all colorations of g
            if from_file
                aut = cols.current_aut  #we are read the number of automorphisms of the colorated graph from the file
            #else
                #N+1 in cols.current_color || continue
            end

            for m in Marks(nv(g), n_marks)    #we run among all marks of g, if n_marks==0 we have only the empty mark
                for w in weights          #we run among all weights of g
                    try
                        local Euler::fmpq = Euler_inv(g,c,w,s,m)//(aut*prod(w)) #the contribuition of the Euler class in the Atiyah-Bott formula
                        for res in 1:n_results      #compute each term of the array P
                            result[res] += P[res](g,c,w,s,m)*Euler    #apply Atiyah-Bott
                        end
                    catch err
                        if isa(err, DivideError)
                            error("Some division by zero occurred. Try again")
                        end
                        println(err)
                        error("Some error occurred")
                        return zeros(fmpq, n_results)
                    end
                end

                if progress_data.progress_bar.enabled

                    progress_data.current_graph += progress_data.top_aut÷aut
                    #update the progress bar
                    update!(progress_data.progress_bar,
                            progress_data.current_graph,
                            showvalues = [(:"Total number of graphs",progress_data.threshold),
                            (:"Current graph",progress_data.current_graph)])
                end
            end
        end
    #end
    return result
end

"""
    AtiyahBottFormula(n, d, m, P, do_check, show_bar, down_col)

Apply the Atiyah-Bott residue formula to the class `P`, in the moduli space of rational marked stable maps to the projective space of dimension `n` of degree `d` with `m` marks.
# Arguments
- `n::Int64`: the dimension of the projective space, it must be between ``1`` and ``254``.
- `d::Int64`: the degree of the stable maps, it must be between ``1`` and ``13``.
- `m::Int64`: the number of marks.
- `P`: the equivariant class.
- `do_check::Bool`: if `true`, checks if `P` is a well defined zero cycle, and stops the computation if this is not true. If `false`, the computation may have an unexpected behaviour. By default is `true`.
- `show_bar::Bool`: hide the progress bar if and only if this condition is `false`. By default is `true`.
- `down_col::Bool`: check if all colorations needed in the computation are in the folder Data, and download them otherwise. Once downloaded, those files can be used for future computations. By default is `true`.

The general construction of `P` is the following:
```julia-repl
julia> P = (g,c,w,s,m) ->
```
After `->`, one has to write an expression in the equivariant classes. All such equivariant classes are functions starting with `(g,c,w,s)` or `(g,c,w,s,m)`. At the end, they can have more arguments. The expression is a polynomial combination of the equivariant classes. We compute the degree of `P` by

```julia-repl
julia> AtiyahBottFormula(n,d,m,P);
```

# Example
```julia-repl
julia> P = (g,c,w,s,m) -> Hypersurface(g,c,w,s,5);
julia> AtiyahBottFormula(3,1,0,P);
Warning: the class is not a 0-cycle.
julia> AtiyahBottFormula(4,1,0,P);
Result: 2875//1
julia> AtiyahBottFormula(4,1,0,P,false);             #same as before, but without the preliminary check on `P`
julia> AtiyahBottFormula(4,1,0,P,false,false);       #same as before, but without showing the progress bar
julia> AtiyahBottFormula(4,1,0,P,false,false,false); #same as before, but without checking for the colorations
```

The function returns an array of the same dimension of `P` (non-vectorized classes are assumed as 1-dimensional arrays). The Julia notation for accessing to array is `name_of_array[i]` where `i` is an index starting from 1.

# Example
```julia-repl
julia> P = (g,c,w,s,m) -> Incidency(g,c,w,s,2)*Hypersurface(g,c,w,s,3);
julia> x = AtiyahBottFormula(3,2,0,P)[1];
Result: 81//1
julia> x
81//1
```

The class `P` supports parameters.
```julia-repl
julia> P = (g,c,w,s,m) -> Hypersurface(g,c,w,s,3)*(Incidency(g,c,w,s,2)//3)^(d-1);
julia> d = 2;
julia> AtiyahBottFormula(3,d,0,P);
Result: 27//1
julia> d = 3;
julia> AtiyahBottFormula(3,d,0,P);
Result: 84//1
```

More examples are available in the support of the equivariant classes. It is enough to type `?` and then the name of the class. Currently, the supported classes are:

* `O1_i`         (Euler class of ``\\mathrm{ev}_i^*\\mathcal{O}_{\\mathbb{P}^n}(1)``)
* `O1`           (product of all `O1_i`)
* `Psi`          (cycle of ``\\psi``-classes)
* `Jet`          (Euler class of the jet bundle ``J^p`` of ``\\mathrm{ev}^*\\mathcal{O}_{\\mathbb{P}^n}(q)``)
* `Hypersurface` (Euler class of the direct image of ``\\mathrm{ev}^*\\mathcal{O}_{\\mathbb{P}^n}(b)``)
* `Incidency`    (cycle parameterizing curves meeting a linear subspace)
* `Contact`      (cycle parameterizing contact curves)
* `R1`           (first derived functor of direct image of ``\\mathrm{ev}^*\\mathcal{O}_{\\mathbb{P}^n}(-k)``)

To add more classes, please contact the authors.
"""
function AtiyahBottFormula(n::Int64, deg::Int64, n_marks::Int64, P, do_check::Bool = true, show_bar::Bool = true, down_col::Bool = true)::Vector{fmpq}

    if n < 1
        printstyled("ERROR: ", bold=true, color=:red)
        println("n must be positive, correct ", n)
        return [fmpq(0)]
    end
    if deg > 13 || deg < 1
        printstyled("ERROR: ", bold=true, color=:red)
        println("d must be between 1 and 13, correct ", deg)
        return [fmpq(0)]
    end
    if n_marks < 0
        printstyled("ERROR: ", bold=true, color=:red)
        println("m must be non negative, correct ", n_marks)
        return [fmpq(0)]
    end

    if !isa(P, Array)  #we want that P is an array, possibly with only one element
        P = [P]
    end

    if do_check && !is_zero_cycle(n, deg, n_marks, P)
        return [fmpq(0)]
    end

    if down_col && !fill_Data(n, deg)
        printstyled("ERROR: ", bold=true, color=:red)
        println("I was unable to download the colorations from the repository. Check your internet connection or execute:")
        printstyled("julia> ", bold=true, color=:light_green)
        println("AtiyahBottFormula($n,$deg,$n_marks,P,$do_check,$show_bar,false);")
        return [fmpq(0)]
    end

    local n_results::Int64 = length(P)   #this store the number of final results of our computation
    local result::Vector{fmpq} = zeros(fmpq, n_results) #the array of final results
    local max_col::Int64 = n+1   #the colors are number from 1 to n+1

    list_g::IOStream = open(current_dir*"/list_trees.txt")
    #open the file containing the list of Prufer sequences of graphs

    #s::Vector{Rational{BigInt}} = convert(Vector{Rational{BigInt}},rand(-1000*max_col*deg:1000*max_col*deg,max_col))
    s::Vector{Rational{BigInt}} = Rational{BigInt}.(rand(UInt32,max_col))
    #set up progress data
    threshold::Int64 = sum(v -> number_trees[v-1]*max_col*(n^(v-1))*(v^n_marks), 2:deg+1)

    progress_data = ProgressData(
        Progress(threshold, barglyphs=BarGlyphs("[=> ]"), color=:green, enabled=show_bar), #progress_bar
        0, #current_graph
        0, #top_aut
        threshold) #threshold

    for v in 2:(deg+1) #run the computation among all graphs with fixed number of vertices

        for _ in 1:number_trees[v - 1]  #run the computation for a fixed graph

            str = readline(list_g) #read a new line, we expect a Prufer seq plus the number of automorphisms
            local (g, aut) = get_graph(str)  #g is the graph, aut is the number of automorphisms
            progress_data.top_aut = aut
            #local top_aut::Int64 = copy(aut) #the number of automorphisms of g without any decoration

            pruf_str = string(split(str, ',')[1])
            result += AtiyahBottFormulaForGraph(g, pruf_str, aut, n, deg,
                    n_marks, P, s,
                    progress_data)

        end
    end

    close(list_g)     #close the file with Prufer sequences
    if n_results == 1
        println("Result: ", result[1])
    else
        for res in 1:n_results
            println("Result number ", res, ": ", result[res])
        end
    end
    return result
end