text:
using StatsBase

# get 2D xy array from pointcloud
"read xy from Cloud"
function xy_array(cloud::Cloud, index::AbstractVector)
    n = length(index)
    XY = zeros(2, n)
    for i = 1:n
        p = positions(cloud)[index[i]]
        XY[1, i] = p[1]
        XY[2, i] = p[2]
    end
    XY
end
xy_array(cloud::Cloud) = xy_array(cloud, 1:length(cloud))

# get arrays of x/y/z coordinates
getx(cloud::Cloud, index::Vector{U}) where U <: Integer = getindex.(positions(cloud)[index], 1)
getx(cloud::Cloud, index::Integer) = positions(cloud)[index][1]
getx(cloud::Cloud) = getindex.(positions(cloud), 1)
gety(cloud::Cloud, index::Vector{U}) where U <: Integer = getindex.(positions(cloud)[index], 2)
gety(cloud::Cloud, index::Integer) = positions(cloud)[index][2]
gety(cloud::Cloud) = getindex.(positions(cloud), 2)
getz(cloud::Cloud, index::Vector{U}) where U <: Integer = getindex.(positions(cloud)[index], 3)
getz(cloud::Cloud, index::Integer) = positions(cloud)[index][3]
getz(cloud::Cloud) = getindex.(positions(cloud), 3)

# set arrays of x/y/z coordinates
setx!(cloud::Cloud, index::Vector{U}, value::Vector{Float64}) where U <: Integer = setindex!.(positions(cloud)[index], value, 1)
sety!(cloud::Cloud, index::Vector{U}, value::Vector{Float64}) where U <: Integer = setindex!.(positions(cloud)[index], value, 2)
setz!(cloud::Cloud, index::Vector{U}, value::Vector{Float64}) where U <: Integer = setindex!.(positions(cloud)[index], value, 3)
setx!(cloud::Cloud, index::Integer, value::Float64) = setindex!(positions(cloud)[index], value, 1)
sety!(cloud::Cloud, index::Integer, value::Float64) = setindex!(positions(cloud)[index], value, 2)
setz!(cloud::Cloud, index::Integer, value::Float64) = setindex!(positions(cloud)[index], value, 3)

# normalize z from given z value
function normalizez!(cloud::Cloud, index::Vector{U}, value::Float64) where U <: Integer
    z = getz(cloud, index)
    setz!(cloud, index, z - value)
end

# statistical functions
reducer_minz(cloud::Cloud, index::Vector{U}) where U <: Integer = minimum(getz(cloud, index))
reducer_maxz(cloud::Cloud, index::Vector{U}) where U <: Integer = maximum(getz(cloud, index))
reducer_medz(cloud::Cloud, index::Vector{U}) where U <: Integer = median(getz(cloud, index))
reducer_meanz(cloud::Cloud, index::Vector{U}) where U <: Integer = mean(getz(cloud, index))
reducer_madz(cloud::Cloud, index::Vector{U}) where U <: Integer = mad(getz(cloud, index), normalize=false)
reducer_stdz(cloud::Cloud, index::Vector{U}) where U <: Integer = std(getz(cloud, index))
reducer_99pz(cloud::Cloud, index::Vector{U}) where U <: Integer = maximum(trim(getz(cloud, index), prop=0.01))

# reducer for vegetation histogram splits
function reducer_veg_below_2m(cloud::Cloud, index::Vector{U}) where U <: Integer
    heights = getz(cloud, index)
    below = heights .< 2.
    return sum(below) / length(heights) * 100.
end
function reducer_veg_below_5m(cloud::Cloud, index::Vector{U}) where U <: Integer
    heights = getz(cloud, index)
    below = heights .< 5.
    return sum(below) / length(heights) * 100.
end
function reducer_veg_above_25m(cloud::Cloud, index::Vector{U}) where U <: Integer
    heights = getz(cloud, index)
    above = heights .> 25.
    return sum(above) / length(heights) * 100.
end
function reducer_veg_above_20m(cloud::Cloud, index::Vector{U}) where U <: Integer
    heights = getz(cloud, index)
    above = heights .> 20.
    return sum(above) / length(heights) * 100.
end

function reducer_ranz(cloud::Cloud, index::Vector{U}) where U <: Integer
    z = getz(cloud, index)
    maximum(z) - minimum(z)
end

function reducer_pctz(cloud::Cloud, index::Vector{U}) where U <: Integer
    quantile(getz(cloud, index), [0.1, 0.25, 0.5, 0.75, 0.9])
end

# reducer function to count number of points per cell
reducer_count(cloud::XYZ.Cloud, index::Vector{U}) where U <: Integer = length(index)
reducer_dens = reducer_count  # backwards compatibility

# reducer functions for any point attribute
function reducer_min_attr(field::Symbol)
    reducer(cloud::XYZ.Cloud, index::Vector{U}) where U <: Integer = minimum(cloud[field][index])
end
function reducer_max_attr(field::Symbol)
    reducer(cloud::XYZ.Cloud, index::Vector{U}) where U <: Integer = maximum(cloud[field][index])
end
function reducer_med_attr(field::Symbol)
    reducer(cloud::XYZ.Cloud, index::Vector{U}) where U <: Integer = median(cloud[field][index])
end
function reducer_min_abs_attr(field::Symbol)
    reducer(cloud::XYZ.Cloud, index::Vector{U}) where U <: Integer = minimum(abs.(cloud[field][index]))
end

# reducer function to get any attribute of the lowest or highest point
function reducer_lowestpoint_attr(field::Symbol)
    reducer(cloud::XYZ.Cloud, index::Vector{U}) where U <: Integer = cloud[field][index[indmin(getz(cloud, index))]]
end
function reducer_highestpoint_attr(field::Symbol)
    reducer(cloud::XYZ.Cloud, index::Vector{U}) where U <: Integer = cloud[field][index[indmax(getz(cloud, index))]]
end

# reduce_index, returns local(!) index of minimum / maximum
reducer_minz_index(cloud::Cloud, index::Vector{U}) where U <: Integer = indmin(getz(cloud, index))
reducer_maxz_index(cloud::Cloud, index::Vector{U}) where U <: Integer = indmax(getz(cloud, index))

const class_description = Dict{Int, String}(
    0 => "Created, never classified",
    1 => "Unclassified",
    2 => "Ground",
    3 => "Low Vegetation",
    4 => "Medium Vegetation",
    5 => "High Vegetation",
    6 => "Building",
    7 => "Low Point (noise)",
    8 => "Model Key-point (mass point)",
    9 => "Water",
    12 => "Overlap Points",

    18 => "High noise" # LAS 1.4
)

"Describe a `Cloud` with a number of statistics"
function describe(io::IO, cloud::Cloud)
    classes = sort!(unique(cloud[:classification]))
    quantvals = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]
    print(io, "quantiles,")
    join(io, quantvals, ",")
    println(io, "")
    for class in classes
        idxs = cloud[:classification] .== class
        println(io, "    class, $class, $(class_description[class])")
        print(io, "        elevation,")
        join(io, quantile(getz(cloud)[idxs], quantvals), ",")
        print(io, "\n        intensity,")
        join(io, quantile(cloud[:intensity][idxs], quantvals), ",")
        print(io, "\n")
    end
end

function describe(filepath::String, cloud::Cloud)
    open(filepath, "w") do io
        describe(io, cloud)
    end
end

"Paint a Cloud efficiently with a Raster index and RGB bands from an orthophoto.
Does not load the entire orthophoto into memory, but one block at a time."
function paint!(pc::Cloud, r::Raster,
        rband::Ptr{GDAL.GDALRasterBandH},
        gband::Ptr{GDAL.GDALRasterBandH},
        bband::Ptr{GDAL.GDALRasterBandH})

    # 257 would be a full mapping, but 256 is in the standard,
    # see https://groups.google.com/d/msg/lasroom/Th2_-ywc2q8/4JRajnxxO0wJ
    const colorscale = 0x0100 # = 256
    # assume all bands have same block sizes and data types
    @assert GDAL.getrasterdatatype(rband) === GDAL.GDT_Byte
    ncolblock, nrowblock = Ref(Cint(-1)), Ref(Cint(-1))
    GDAL.getblocksize(rband, ncolblock, nrowblock) # size of single blocks
    ncolblock, nrowblock = ncolblock[], nrowblock[] # get values from Ref
    ncol = GDAL.getrasterbandxsize(rband)
    nrow = GDAL.getrasterbandysize(rband)
    # number of blocks in the x and y direction
    nblockx = cld(ncol, ncolblock)
    nblocky = cld(nrow, nrowblock)
    # number of rows ans columns in the last block of every row or column
    ncollastblock = rem(ncol, ncolblock) == 0 ? ncolblock : rem(ncol, ncolblock)
    nrowlastblock = rem(nrow, nrowblock) == 0 ? nrowblock : rem(nrow, nrowblock)
    # asserted it is GDAL.GDT_Byte
    # switched dimensions instead of transposing
    rblock = zeros(UInt8, ncolblock, nrowblock)
    gblock = zeros(UInt8, ncolblock, nrowblock)
    bblock = zeros(UInt8, ncolblock, nrowblock)
    # iterate over blocks
    for blockrow = 1:nblocky
        for blockcol = 1:nblockx
            nxblockoff = blockcol-1
            nyblockoff = blockrow-1
            # possible optimization: only read blocks that contain points?
            GDAL.readblock(rband, nxblockoff, nyblockoff, rblock)
            GDAL.readblock(gband, nxblockoff, nyblockoff, gblock)
            GDAL.readblock(bband, nxblockoff, nyblockoff, bblock)
            nrowthisblock = blockrow == nblocky ? nrowlastblock : nrowblock
            ncolthisblock = blockcol == nblockx ? ncollastblock : ncolblock
            # calculate the global indices for Cloud Raster index
            startrow = (blockrow - 1) * nrowblock
            startcol = (blockcol - 1) * ncolblock
            # iterate over block cells
            for i in 1:nrowthisblock
                for j in 1:ncolthisblock
                    row = startrow + i
                    col = startcol + j
                    idxs = r[row,col]
                    isempty(idxs) && continue # big speedup if many pixels don't have points
                    color = RGB{N0f16}(
                        reinterpret(N0f16, rblock[j,i] * colorscale),
                        reinterpret(N0f16, gblock[j,i] * colorscale),
                        reinterpret(N0f16, bblock[j,i] * colorscale))
                    pc[:color][idxs] = color
                end
            end
        end
    end
    nothing
end

function paint!(pc::Cloud, orthofile::AbstractString)
    # GDAL setup
    GDAL.allregister()
    ds = GDAL.open(orthofile, GDAL.GA_ReadOnly)
    geotransform = zeros(6)
    GDAL.getgeotransform(ds, geotransform)
    epsg = wkt2epsg(GDAL.getprojectionref(ds))
    rband = GDAL.getrasterband(ds, 1)
    gband = GDAL.getrasterband(ds, 2)
    bband = GDAL.getrasterband(ds, 3)
    ncol = Int(GDAL.getrasterbandxsize(rband))
    nrow = Int(GDAL.getrasterbandysize(rband))

    # create raster index that is guaranteed to align with the orthophoto pixels
    r = define_raster(pc, nrow, ncol, geotransform, epsg=epsg, pointfilter=nothing)
    paint!(pc, r, rband, gband, bband)
    GDAL.close(ds)
    nothing
end

function create_kdtree(cloud::Cloud;
    dim=2,
    metric=Euclidean(),
    index = 1:length(cloud))

    # create tree
    if dim == 2
        tree = KDTree(xy_array(cloud, index), metric)
    elseif dim == 3
        tree = KDTree(positions(cloud)[index], metric)
    end

    tree
end

function point_density!(cloud::Cloud, spatialindex::NearestNeighbors.NNTree;
                        range=(1.0 / (4/3.0 * pi ))^(1.0/3),
                        pointfilter=nothing) # vol = 1m3
    @assert (spatialindex.metric == Euclidean()) && (length(spatialindex.data[1]) == 3)
    n = length(cloud)
    cloud[:point_density] = zeros(Float32, n)
    vol = 4/3.0 * pi * range^3

    # loop through points and calculate point density
    @showprogress 5 "looping through las points..." for i in 1:n
        (pointfilter != nothing) && (pointfilter(cloud, i) || continue)
        # find number of points in range
        neighbors = inrange(spatialindex, positions(cloud)[i], range)
        # set point source attribute to visualize no. of neighbors
        cloud[:point_density][i] = length(neighbors) / vol
    end

    nothing
end
point_density!(cloud::Cloud; range=(1.0 / (4/3.0 * pi ))^(1.0/3)) = point_density!(cloud, create_kdtree(cloud; dim=3); range=range)

unique_id!(cloud::Cloud) = cloud[:id] = 1:length(cloud);

"""copy classification based on unique id,
assumes the field 'id' in cloud copy_from has the index of the same point in cloud copy_to"""
function copy_attributes!(cloud_copy_to::Cloud, cloud_copy_from::Cloud; field=:classification)
    @assert haskey(cloud_copy_from.attributes, :id) "both point clouds should have an 'id' field"
    @assert haskey(cloud_copy_from.attributes, field) "the 'copy_from' point cloud does not have a field $(field)"

    n = length(cloud_copy_to)
    if !haskey(cloud_copy_to.attributes, field)
        cloud_copy_to[field] = zeros(eltype(cloud_copy_from[field]), n)
    end

    # copy attribute valeu back to original cloud
    for (idx, v) in zip(cloud_copy_from[:id], cloud_copy_from[field])
        cloud_copy_to[field][idx] = v
    end
    nothing
end

"""project the values of grid to point in pointcloud based on a raster index"""
function grid2cloud_attribute!(cloud::Cloud, grid::Array{T,2}, r::Raster, name::Symbol) where T <: Real
    @assert size(grid) == (r.nrow, r.ncol) "raster definition 'r' should have same size as 'grid'"
    cloud[name] = zeros(T,length(cloud))
    for icell in 1:length(r)
        v = grid[icell]
        for i in r[icell]
            cloud[name][i] = v
        end
    end
    nothing
end


text:
using GroupedErrors
using Test
using IndexedTables
using ShiftedArrays
using TextParse
using Statistics
import Lazy

function loadtable(path)
    cols, names = TextParse.csvread(path)
    table(cols...; names = Symbol.(names))
end

function check_equality(stored_table, test_table, atol)
    for j in colnames(stored_table)
        if eltype(columns(stored_table, j)) == Float64
            @test columns(stored_table, j) ≈ columns(test_table, j) atol = atol
        else
            @test columns(stored_table, j) == columns(test_table, j)
        end
    end
end

include("build_test_tables.jl")

for i in 1:length(tables)
    println(i)
    test_table = tables[i]
    println(test_table)
    stored_table = loadtable(GroupedErrors.exampletable("t$i.csv"))
    println(stored_table)
    atol = i == 8 ? 1e-1 : 2e-3
    check_equality(stored_table, test_table, atol)
end

split_vars = [:Minrty, :Sx]

processed_table = @> school begin
    @splitby(Tuple(getfield(_, s) for s in split_vars))
    @x(_.MAch, :continuous)
    @y(_.SSS)
    ProcessedTable
end

processed_table_col = Lazy.@as x GroupedErrors.ColumnSelector(school) begin
    GroupedErrors._splitby(x, split_vars)
    GroupedErrors._x(x, :MAch, :continuous)
    GroupedErrors._y(x, :locreg, :SSS)
    ProcessedTable
end

check_equality(processed_table.table, processed_table_col.table, 1e-6)

processed_table = @> school begin
    @splitby (_.Minrty, _.Sx)
    @across _.School
    @set_attr :linestyle _[1] == "Yes" ? :solid : :dash
    @set_attr :color _[2] == "Male" ? :black : :blue
    @x _.CSES
    @y :density bandwidth = 0.2
    ProcessedTable
end

@test [t[1] for t in processed_table.kw[:plot_kwargs]] == [:linestyle, :color]

na_school_vec = convert(Vector{GroupedErrors.DataValue{Float64}}, columns(school, :MAch))
na_school_vec[1] = GroupedErrors.DataValue{Float64}()

table2process_nafree = @> setcol(school, :MAch, na_school_vec) begin
    @x(_.MAch, :continuous)
    @y(_.SSS)
    GroupedErrors.Table2Process
end

table2process = @> school begin
    @x(_.MAch, :continuous)
    @y(_.SSS)
    GroupedErrors.Table2Process
end

check_equality(table(table2process.table...)[2:end], table(table2process_nafree.table...), 1e-8)

v = [1, 2, 3, 4, 5, 6, 7, 8]

t = table([1, 2], ShiftedArray.((v,), [-3, -7]), names = [:x, :y])

res = @> t begin
    @across _.x
    @x -1:1 :discrete
    @y _.y
    ProcessedTable
end
expected_res = table((s1 = fill("y1", 3), x = -1:1, y = [4., 5., 6.], err = fill(2., 3)),
    pkey = :s1)

check_equality(expected_res, res.table, 1e-8)

cs = Lazy.@as x GroupedErrors.ColumnSelector(school) begin
    GroupedErrors._across(x, :School)
    GroupedErrors._x(x, :MAch)
    GroupedErrors._y(x, :MAch)
    GroupedErrors._compare(x, :Sx)
    ProcessedTable(x)
end

s = @> school begin
    @across _.School
    @xy _.MAch
    @compare _.Sx
    ProcessedTable
end

check_equality(cs.table, s.table, 1e-8)

ext: jl


text:
module ComoniconGUI

include("codegen/codegen.jl")
using .CodeGen

include("parse/parse.jl")
using .Parse

export @main, @cast

end


text:
module HansenWrapper

using JavaCall

const Class = @jimport java.lang.Class
const Field = @jimport org.hipparchus.Field
const FieldHansenTesseralLinear = @jimport org.orekit.propagation.semianalytical.dsst.utilities.hansen.FieldHansenTesseralLinear
const FieldHansenThirdBodyLinear = @jimport org.orekit.propagation.semianalytical.dsst.utilities.hansen.FieldHansenThirdBodyLinear
const FieldHansenZonalLinear = @jimport org.orekit.propagation.semianalytical.dsst.utilities.hansen.FieldHansenZonalLinear
const HansenTesseralLinear = @jimport org.orekit.propagation.semianalytical.dsst.utilities.hansen.HansenTesseralLinear
const HansenThirdBodyLinear = @jimport org.orekit.propagation.semianalytical.dsst.utilities.hansen.HansenThirdBodyLinear
const HansenUtilities = @jimport org.orekit.propagation.semianalytical.dsst.utilities.hansen.HansenUtilities
const HansenZonalLinear = @jimport org.orekit.propagation.semianalytical.dsst.utilities.hansen.HansenZonalLinear
const JString = @jimport java.lang.String
const Object = @jimport java.lang.Object
const PolynomialFunction = @jimport org.hipparchus.analysis.polynomials.PolynomialFunction
const PolynomialFunctionMatrix = @jimport org.orekit.propagation.semianalytical.dsst.utilities.hansen.PolynomialFunctionMatrix
const RealFieldElement = @jimport org.hipparchus.RealFieldElement

export FieldHansenTesseralLinear
export FieldHansenThirdBodyLinear
export FieldHansenZonalLinear
export HansenTesseralLinear
export HansenThirdBodyLinear
export HansenUtilities
export HansenZonalLinear
export PolynomialFunctionMatrix
export add
export build_identity_matrix2
export build_identity_matrix4
export build_zero_matrix2
export build_zero_matrix4
export compute_init_values
export equals
export get_class
export get_derivative
export get_elem
export get_matrix_line
export get_value
export hash_code
export multiply
export notify
export notify_all
export set_elem
export set_matrix
export set_matrix_line
export to_string
export wait

include("field_hansen_tesseral_linear.jl")
include("field_hansen_third_body_linear.jl")
include("field_hansen_zonal_linear.jl")
include("hansen_tesseral_linear.jl")
include("hansen_third_body_linear.jl")
include("hansen_utilities.jl")
include("hansen_zonal_linear.jl")
include("polynomial_function_matrix.jl")

end


text:
include("/home/wanglab/Documents/WhiskerTask2/WhiskerTask.jl")

myt=Task_CameraTask()
myfpga=FPGA(1,[0],usb3=true)
myparams=Algorithm[DetectAbs(),ClusterTemplate(49),AlignProm(),FeatureTime(),ReductionNone(),ThresholdMeanN()];
(myrhd,ss,myfpgas)=makeRHD([myfpga],params=myparams,single_channel_mode=true);

handles = makegui(myrhd,ss,myt,myfpgas);

if !isinteractive()
    c = Condition()
    signal_connect(handles.win, :destroy) do widget
        notify(c)
    end
    wait(c)
end

max_stars_repo_name: paulmthompson/WhiskerTask

max_issues_repo_name: paulmthompson/WhiskerTask

max_forks_repo_name: paulmthompson/WhiskerTask


text:
#Code written by Skipper af Rosenborg and Esben Bager
using JuMP
using Gurobi
using StatsBase
using DataFrames
using CSV
include("SupportFunction.jl")
include("DataLoad.jl")
println("Leeeeroooy Jenkins")

### Set path ###
#Esben's path
#cd("$(homedir())/Documents/GitHub/Thesis/Data")
#path = "$(homedir())/Documents/GitHub/Thesis/Data"

#Skipper's path
#path = "/Users/SkipperAfRosenborg/Google Drive/DTU/10. Semester/Thesis/GitHubCode/Thesis/Data/IndexDataDiff"

#HPC path
path = "/zhome/9f/d/88706/SpecialeCode/Thesis/Data/IndexDataDiff/"

### Load data ###
#mainData = loadIndexDataNoDur(path)
mainData = loadIndexDataLOGReturn("NoDur", path)
path = "/zhome/9f/d/88706/SpecialeCode/Results/IndexData/SpeedTest/"
#path = "/Users/SkipperAfRosenborg/Google Drive/DTU/10. Semester/Thesis/GitHubCode/Results/IndexData/SpeedTest/"

#Reset HPC path
#path = "/zhome/9f/d/88706/SpecialeCode/Thesis/ML"
#cd(path)

# Split data
dateAndReseccion = Array(mainData[:,end-1:end])
mainDataArr = Array(mainData[:,1:end-2])

colNames = names(mainData)

nRows = size(mainDataArr)[1]
nCols = size(mainDataArr)[2]

mainXarr = mainDataArr[:,1:nCols-1]
mainYarr = mainDataArr[:, nCols:nCols]

# Transform with time elements
mainXarr = expandWithTime3612(mainXarr)

# Transform with non linear expansions and techincal indicators
mainXarr = expandWithTransformations(mainXarr)
mainXarr = expandWithMAandMomentum(mainXarr, mainYarr, (nCols-1))

# Standardize data
standX = zScoreByColumn(mainXarr)
standY = mainDataArr[:, nCols:nCols]
allData = hcat(standX, standY)
bCols = size(standX)[2]
nRows = size(standX)[1]

# Set prediction parameters
trainingSize = 12
predictions = 1
testRuns = nRows-trainingSize-predictions

Xtrain = allData[:, 1:bCols]
trainingData = Xtrain[:,1:nCols-1]

### STAGE 1 ###
println("STAGE 1 INITIATED")
#Define solve
m = JuMP.Model(solver = GurobiSolver(OutputFlag =0 ))

#Add binary variables variables
@variable(m, 0 <= z[1:nCols-1] <= 1, Bin )

#Calculate highly correlated matix
HC = cor(trainingData)

#Define objective function
@objective(m, Max, sum(z[i] for i=1:length(z)))

#Define max correlation and constraints
rho = 0.8
for i=1:length(z)
    for j=1:length(z)
        if i != j
            if HC[i,j] >= rho
                @constraint(m,z[i]+z[j] <= 1)
            end
        end
    end
end

#Get solution status
status = solve(m);

#Get objective value
println("Objective value kMax: ", getobjectivevalue(m))
kmax = getobjectivevalue(m)

println("STAGE 1 DONE")
# Initialise arrays
amountOfGammas = 5
realArr = zeros((nRows-trainingSize-predictions), 1)
ISRSquared = zeros((nRows-trainingSize-predictions), kmax*amountOfGammas)
bSolMatrix = zeros((nRows-trainingSize-predictions), kmax*amountOfGammas, bCols)

### STAGE 2 ###
println("STAGE 2 INITIATED")
#Initialise values for check later
bSolved = []
bestBeta = []

bigM = 10
#Spaced between 0 and half the SSTO since this would then get SSTO*absSumOfBeta which would force everything to 0
gammaArray = log10.(logspace(0.001, 5, amountOfGammas))

function solveAndLogForAllK(kmax, standXVali, standYVali, r, warmstart)
    # Initialise values
    best3Beta = zeros(3,bCols+3)
    gammaArray = log10.(logspace(5, 0.01, 3))
    println("Remember to check gammaArray line 121")
    bSolved = zeros(bCols,1)
    bestObjtives = zeros(kmax, length(gammaArray))
    println("Warmstart = ",warmstart)

    HC = cor(standX)
	for i in 1:kmax
        if warmstart == true
            bSolved, s = getBetaAndGamma(standX,standY,i)
        end
		for g in 1:length(gammaArray)
			gamma = gammaArray[g]

        	HCPairCounter = 0

            #println("Building model")
        	#Define parameters and model
        	stage2Model = JuMP.Model(solver = GurobiSolver(TimeLimit = 30, OutputFlag = 0, Threads = 2, PreMIQCPForm=0, MIPFocus=1, ImproveStartTime=40));

        	#Define variables
        	@variable(stage2Model, b[1:bCols], start = 0) #Beta values

        	#Define binary variable (5b)
        	@variable(stage2Model, 0 <= z[1:bCols] <= 1, Bin, start = 0)

        	@variable(stage2Model, v[1:bCols], start = 0) # auxiliary variables for abs

            # Set warmstart values
            if warmstart == true || g>1
                for j in find(bSolved)
    				setvalue(b[j], bSolved[j])
    			    setvalue(v[j], abs(bSolved[j]))
    				setvalue(z[j], 1)
    			end
            end

        	@variable(stage2Model, T) #First objective term
        	@variable(stage2Model, G) #Second objective term

        	#Define objective function (5a)
        	@objective(stage2Model, Min, T+G)

            @constraint(stage2Model, soc, norm( [1-T;2*(standX*b-standY)] ) <= 1+T)

        	#Define constraints (5c)
        	@constraint(stage2Model, conBigMN, -1*b .<= bigM*z) #from 0 to bCols -1
        	@constraint(stage2Model, conBigMP,  1*b .<= bigM*z) #from bCols to 2bCols-1

        	#Second objective term
        	@constraint(stage2Model, 1*b .<= v) #from 2bCols to 3bCols-1
        	@constraint(stage2Model, -1*b .<= v) #from 3bCols to 4bCols-1

        	#gamma[g]*oneNorm <= G ---> -G <= -gamma[g]*oneNorm --> G >= gamma[g]*oneNorm
            @constraint(stage2Model, gammaConstr, gamma*ones(bCols)'*v <= G) #4bCols

            #oneNorm = sum(v[i] for i=1:bCols)
            #@constraint(stage2Model, gammaConstr, gamma*oneNorm <= G) #4bCols

        	#Define kmax constraint (5d)
        	@constraint(stage2Model, kMaxConstr, sum(z[j] for j=1:bCols) <= i) #4bCols+1

        	#Constraint 5f - can only select one of a pair of highly correlated features
        	rho = 0.8
        	for k=1:bCols
        		for j=1:bCols
        			if k != j
        				if HC[k,j] >= rho
        					HCPairCounter += 1
        					@constraint(stage2Model,z[k]+z[j] <= 1) #from 4bCols+1 to (4bCols+1+HCPairCounter)
        				end
        			end
        		end
        	end

        	#Constraint (5g) - only one transformation allowed (x, x^2, log(x) or sqrt(x))
        	for j=1:(nCols-1)
        		@constraint(stage2Model, z[j]+z[j+(nCols-1)]+z[j+2*(nCols-1)]+z[j+3*(nCols-1)] <= 1) #from (4bCols+1+HCPairCounter) to (4bCols+1+HCPairCounter+nCols)
        	end

        	#Adding empty cuts to possibly contain cuts. Can generate a maximum of 6 cuts for
        	#each stage 3 and default value is 3 iterations between stage 2 and stage 3,
        	#so 18 empty constraints are added + 1 for fixing the solution
        	nEmptyCuts = 19
        	zOnes = ones(bCols)
        	for c = 1:nEmptyCuts
        		@constraint(stage2Model, sum(zOnes[i]*z[i] for i=1:bCols) <= bCols) #(4bCols+1+HCPairCounter+nCols) to (4bCols+1+HCPairCounter+nCols+18)
        	end

			#Solve Stage 2 model
            solve(stage2Model)
            bSolved = getvalue(b)
            bSolMatrix[r,Int64((i-1)*amountOfGammas+g),:] = bSolved

            SSres = sum((standY[i]-(standX[i,:]'*bSolved))^2 for i=1:length(standY))
            ISRSquared[r,Int64((i-1)*amountOfGammas+g)] = 1 - SSres/SSTO
			#Get solution and calculate R^2
			#bSolved = sol[1:bCols]
			#zSolved = sol[1+bCols:2*bCols]

            obj = Gurobi.getobjval(internalmodel(stage2Model))
            bound = Gurobi.getobjbound(internalmodel(stage2Model))
            bestObjtives[Int64(i),Int64(g)] = obj
            gap = abs.(-1+bound/obj)

            prediction = standXVali*bSolved

			println("Prediction =$prediction\t Real =$standYVali\t kMax =$i \t gamma =$g \t obj = $obj, bound = $bound, gap = $gap")
            realArr[r,1] = standYVali[1]

		end
		#printNonZeroValues(bSolved)
	end
    println(warmstart)
    return bestObjtives
	#return #best3Beta, model
end

function solveAndLogForAllKHeuristics(kmax, standXVali, standYVali, r)
    # Initialise values
    best3Beta = zeros(3,bCols+3)
    gammaArray = log10.(logspace(5, 0.01, 3))
    bSolved = zeros(bCols,1)
    bestObjtives = zeros(kmax, length(gammaArray))
    #println("Warmstart heuristics")

    HC = cor(standX)

    #Define an array of variabels that can't be picked to gether
    HCArray = zeros(length(find(x -> x>=0.8, HC)),2)
    rho = 0.8
    itCounter = 1
    for k=1:size(standX)[2]
        for j=1:size(standX)[2]
            if k != j
                if HC[k,j] >= rho
                    HCArray[itCounter,1] = k
                    HCArray[itCounter,2] = j
                    itCounter += 1
                end
            end
        end
    end
    HCArray = HCArray[1:find(x -> x==0,HCArray[:,1])[1]-1,:]

	for i = 1:kmax
        bSolved = gradDecent(standX, zscore(standY), 100000, 1e-3, i, HC, 0, HCArray)
        tempX = find(bSolved)
        bSolved = zeros(bSolved)
        tempbSolved = inv(standX[:,tempX]'*standX[:,tempX])*standX[:,tempX]'*standY
        bSolved[tempX] = tempbSolved

        for g in 1:length(gammaArray)
			gamma = gammaArray[g]

            find(bSolved)

        	HCPairCounter = 0

            #println("Building model")
        	#Define parameters and model
        	stage2Model = JuMP.Model(solver = GurobiSolver(TimeLimit = 30, OutputFlag = 1, Threads = 2, PreMIQCPForm=0, MIPFocus=1, ImproveStartTime=40));

        	#Define variables
        	@variable(stage2Model, b[1:bCols], start = 0) #Beta values

        	#Define binary variable (5b)
        	@variable(stage2Model, 0 <= z[1:bCols] <= 1, Bin, start = 0)

        	@variable(stage2Model, v[1:bCols], start = 0) # auxiliary variables for abs

            # Set warmstart values
            for j in find(bSolved)
				setvalue(b[j], bSolved[j])
			    setvalue(v[j], abs(bSolved[j]))
				setvalue(z[j], 1)
			end

        	@variable(stage2Model, T) #First objective term
        	@variable(stage2Model, G) #Second objective term

        	#Define objective function (5a)
        	@objective(stage2Model, Min, T+G)

            @constraint(stage2Model, soc, norm( [1-T;2*(standX*b-standY)] ) <= 1+T)

        	#Define constraints (5c)
        	@constraint(stage2Model, conBigMN, -1*b .<= bigM*z) #from 0 to bCols -1
        	@constraint(stage2Model, conBigMP,  1*b .<= bigM*z) #from bCols to 2bCols-1

        	#Second objective term
        	@constraint(stage2Model, 1*b .<= v) #from 2bCols to 3bCols-1
        	@constraint(stage2Model, -1*b .<= v) #from 3bCols to 4bCols-1

        	#gamma[g]*oneNorm <= G ---> -G <= -gamma[g]*oneNorm --> G >= gamma[g]*oneNorm
            @constraint(stage2Model, gammaConstr, gamma*ones(bCols)'*v <= G) #4bCols

            #oneNorm = sum(v[i] for i=1:bCols)
            #@constraint(stage2Model, gammaConstr, gamma*oneNorm <= G) #4bCols

        	#Define kmax constraint (5d)
        	@constraint(stage2Model, kMaxConstr, sum(z[j] for j=1:bCols) <= i) #4bCols+1

        	#Constraint 5f - can only select one of a pair of highly correlated features
        	rho = 0.8
        	for k=1:bCols
        		for j=1:bCols
        			if k != j
        				if HC[k,j] >= rho
        					HCPairCounter += 1
        					@constraint(stage2Model,z[k]+z[j] <= 1) #from 4bCols+1 to (4bCols+1+HCPairCounter)
        				end
        			end
        		end
        	end

        	#Constraint (5g) - only one transformation allowed (x, x^2, log(x) or sqrt(x))
        	for j=1:(nCols-1)
        		@constraint(stage2Model, z[j]+z[j+(nCols-1)]+z[j+2*(nCols-1)]+z[j+3*(nCols-1)] <= 1) #from (4bCols+1+HCPairCounter) to (4bCols+1+HCPairCounter+nCols)
        	end

        	#Adding empty cuts to possibly contain cuts. Can generate a maximum of 6 cuts for
        	#each stage 3 and default value is 3 iterations between stage 2 and stage 3,
        	#so 18 empty constraints are added + 1 for fixing the solution
        	nEmptyCuts = 19
        	zOnes = ones(bCols)
        	for c = 1:nEmptyCuts
        		@constraint(stage2Model, sum(zOnes[i]*z[i] for i=1:bCols) <= bCols) #(4bCols+1+HCPairCounter+nCols) to (4bCols+1+HCPairCounter+nCols+18)
        	end

			#Solve Stage 2 model
            solve(stage2Model)
            bSolved = getvalue(b)
            bSolMatrix[r,Int64((i-1)*amountOfGammas+g),:] = bSolved

            SSres = sum((standY[i]-(standX[i,:]'*bSolved))^2 for i=1:length(standY))
            ISRSquared[r,Int64((i-1)*amountOfGammas+g)] = 1 - SSres/SSTO
			#Get solution and calculate R^2
			#bSolved = sol[1:bCols]
			#zSolved = sol[1+bCols:2*bCols]

            obj = Gurobi.getobjval(internalmodel(stage2Model))
            bound = Gurobi.getobjbound(internalmodel(stage2Model))
            bestObjtives[Int64(i),Int64(g)] = obj
            gap = abs.(-1+bound/obj)

            prediction = standXVali*bSolved

			println("Prediction =$prediction\t Real =$standYVali\t kMax =$i \t gamma =$g \t obj = $obj, bound = $bound, gap = $gap")
            realArr[r,1] = standYVali[1]
		end
		#printNonZeroValues(bSolved)
	end
    println(warmstart)
    return bestObjtives
	#return #best3Beta, model
end

"""
Function that builds and solve the LASSO regression
"""
function solveLasso(Xtrain, Ytrain, gamma)
	bCols = size(Xtrain)[2]

    #Define model
	M = JuMP.Model(solver = GurobiSolver(OutputFlag = 0))
	@variables M begin
	        b[1:bCols]
	        t[1:bCols]
	        w
	end
	@objective(M,Min,0.5*w+gamma*ones(bCols)'*t)
	@constraint(M, soc, norm( [1-w;2*(Xtrain*b-Ytrain)] ) <= 1+w)
	@constraint(M,  b .<= t)
	@constraint(M, -t .<= b)

	solve(M)
	bSolved = getvalue(b)

    #Shrink values to 0
	for i in 1:length(bSolved)
		if bSolved[i] <= 1e-6
			if bSolved[i] >= -1e-6
				bSolved[i] = 0
			end
		end
	end

    #Count non zero prediction terms
	k = countnz(bSolved)
	return bSolved, k
end

"""
Function to search for lasso solution that uses a specific number of non zero variables k
"""
function getBetaAndGamma(standX,standY, kmax)
    # Initialise parameters
    bSolved=zeros(bCols)
	k=1000 #
	iter = 250 #start gamma index
    iterCounter = 1
	gammaArray = log10.(logspace(0, 10, 500)) #Searchs space
    distance = 125 # distance iter is moved in 1st iteration

    #Trying to find k=kmax within 25 iterations
	while (k!=kmax && iterCounter <= 25)
        if iter>500 || iter <1
            break
        end

        # solve lasse at current gamma and get k non zeros variables
		bSolved, k = solveLasso(standX,standY,gammaArray[iter])
        #println("k= $k \t distance = $distance, iter = $iter")
        if k > kmax
            iter += distance
        else
            iter -= distance
        end
        distance = Int64(ceil(distance/2)) #Half distance to narrow search space
        iterCounter += 1
		#println("Iteration: ",iter," gamma=",gammaArray[iter])
	end

    # If k is not found in the search, bruce force is used to find k
    if k != kmax
        iter = 1
        while (k!=kmax && iter<500)
    		bSolved, k = solveLasso(standX,standY,gammaArray[iter])
            #println("k= $k \t distance = $distance, iter = $iter")
            iter += 1
    		#println("Iteration: ",iter," gamma=",gammaArray[iter])
    	end
    end

	k = countnz(bSolved)
	return bSolved, k
end

"""
Build and solve function for stage 2 used in math heuristic
"""
function buildAndSolveProblem(kmax, standX, standY, r, gamma, leftBranchMatrix, rightBranchMatrix, searchSize, HC, rightBranchBool)
    #Define parameters and model
    stage2Model = JuMP.Model(solver = GurobiSolver(TimeLimit = 5, OutputFlag = 0, Threads = 2, PreMIQCPForm=0, MIPFocus=1))

    #Define variables
    @variable(stage2Model, b[1:bCols], start = 0) #beta variables
    @variable(stage2Model, 0 <= z[1:bCols] <= 1, Bin, start = 0) #binary variables
    @variable(stage2Model, v[1:bCols], start = 0) # auxiliary variables for abs(beta)
    @variable(stage2Model, T) #First objective term
    @variable(stage2Model, G) #Second objective term

    #Set warmstart solution
    for j in find(leftBranchMatrix)
        setvalue(b[j], leftBranchMatrix[j])
        setvalue(v[j], abs(leftBranchMatrix[j]))
        setvalue(z[j], 1)
    end

    #Define objective function (5a)
    @objective(stage2Model, Min, T+G)

    #Define SOC constraint
    @constraint(stage2Model, soc, norm( [1-T;2*(standX*b-standY)] ) <= 1+T)

    #Define constraints (5c)
    @constraint(stage2Model, conBigMN, -1*b .<= bigM*z) #from 0 to bCols -1
    @constraint(stage2Model, conBigMP,  1*b .<= bigM*z) #from bCols to 2bCols-1

    #Second objective term
    @constraint(stage2Model, 1*b .<= v) #from 2bCols to 3bCols-1
    @constraint(stage2Model, -1*b .<= v) #from 3bCols to 4bCols-1
    #gamma[g]*oneNorm <= G ---> -G <= -gamma[g]*oneNorm --> G >= gamma[g]*oneNorm
    @constraint(stage2Model, gammaConstr, gamma*ones(bCols)'*v <= G) #4bCols

    #Define kmax constraint (5d)
    @constraint(stage2Model, kMaxConstr, sum(z[j] for j=1:bCols) <= kmax) #4bCols+1

    #Constraint 5f - can only select one of a pair of highly correlated features
    rho = 0.8
    for k=1:bCols
        for j=1:bCols
            if k != j
                if HC[k,j] >= rho
                    @constraint(stage2Model,z[k]+z[j] <= 1) #from 4bCols+1 to (4bCols+1+HCPairCounter)
                end
            end
        end
    end

    #Constraint (5g) - only one transformation allowed (x, x^2, log(x) or sqrt(x))
    for j=1:(nCols-1)
        @constraint(stage2Model, z[j]+z[j+(nCols-1)]+z[j+2*(nCols-1)]+z[j+3*(nCols-1)] <= 1) #from (4bCols+1+HCPairCounter) to (4bCols+1+HCPairCounter+nCols)
    end

    #=
    Adding empty cuts to possibly contain cuts. Can generate a maximum of 6 cuts for
    each stage 3 and default value is 3 iterations between stage 2 and stage 3,
    so 18 empty constraints are added + 1 for fixing the solution
    =#
    nEmptyCuts = 19
    zOnes = ones(bCols)
    for c = 1:nEmptyCuts
        @constraint(stage2Model, sum(zOnes[j]*z[j] for j=1:bCols) <= bCols) #(4bCols+1+HCPairCounter+nCols) to (4bCols+1+HCPairCounter+nCols+18)
    end

    #Adding local branching
    #sum of binary values selected (1-z_j) + sum(z_j not selected) <= kmax
    # Loop over number of bsolved in branchingmatrix
    if leftBranchMatrix != 0
        @constraint(stage2Model, sum(1-z[j] for j=find(leftBranchMatrix))
            + sum(z[j] for j=find(x -> x == 0, leftBranchMatrix)) <= searchSize)
    end

    if rightBranchBool
        println("Added right branch")
        if leftBranchMatrix != rightBranchMatrix
            for row in 2:size(rightBranchMatrix)[2]
                @constraint(stage2Model, sum(1-z[j] for j=find(rightBranchMatrix[:,row]))
                    + sum(z[j] for j=find(x -> x == 0, rightBranchMatrix[:,row])) >= searchSize + 1)
            end
        end
    end

    #Solve Stage 2 model
    status = solve(stage2Model)

    if status != :InfeasibleOrUnbounded
        bSolved = getvalue(b)
        obj = Gurobi.getobjval(internalmodel(stage2Model))
        bound = Gurobi.getobjbound(internalmodel(stage2Model))
        gap = abs.(-1+bound/obj)
        #println("Objective is: ", obj, "\tGap is: ",gap)
        return bSolved, status, obj, gap
    else # if model is InfeasibleOrUnbounded all beta = 0 and obj and gap is set to a high value
        #println("\nWriting Problem!\n")
        #writeLP(stage2Model, "TestLP.LP"; genericnames=false)
        bSolved = zeros(bCols)
        obj = 9999
        gap = 9999
        return bSolved, status, obj, gap
    end

    return bSolved, status, obj, gap
end

function solveAndLogForAllKMath(kmax, standXVali, standYVali, r)
	status = :Optimal
    HC = cor(standX)
    gammaArray = log10.(logspace(5, 0.01, 3))
    bestObjtives = zeros(kmax, length(gammaArray))
	for i in 1:kmax
		lassoBsolved, k = getBetaAndGamma(standX,standY, i)   #Find warmstarted solution
		prediction = 0
        bSolved = zeros(bCols)
		for g in 1:length(gammaArray)
            bestbSolved = zeros(bCols)
            #println("\nNew Gamma \n")
            if g == 1
                leftBranchMatrix = lassoBsolved
                rightBranchMatrix = lassoBsolved
            else
                leftBranchMatrix = bSolved
                rightBranchMatrix = bSolved
            end
            #println("kMax = ",i," g = ", g)

            #Initialise parameters
            searchSize = 2
			oldObj = 1000.0
            obj = 1000.0
            bestObj = 9999
			gamma = gammaArray[g]
            iter = 1
            rightBranchBool = false
			while (true) #Continues the while loop is broken by a condition within the loop
                #println("Iteration $iter")
                oldObj = obj # Update old objective value

                bSolved, status, obj, gap = buildAndSolveProblem(i, standX, standY, r, gamma,
                    leftBranchMatrix, rightBranchMatrix, searchSize, HC, rightBranchBool) #Solve local problem
                #println("Kmax = $i and bsolved = ", countnz(bSolved))

                # InfeasibleOrUnbounded case
                if status == :InfeasibleOrUnbounded
                    #println("Went into ",status)
                    # End loop as last node can't be solved and the last optimal solve was with a worse objective
                    if (typeof(leftBranchMatrix) == Int64 || typeof(leftBranchMatrix) == Float64) && iter != 1
                        break
                    end

                    # Convert leftBranchMatrix into rightBranchMatrix with RHS >= searchSize
                    rightBranchMatrix = hcat(rightBranchMatrix, leftBranchMatrix)

                    # Delete leftBranchMatrix, bSolve = 0 if status = InfeasibleOrUnbounded
                    leftBranchMatrix = bSolved

                    # Weak diversification --> Increase searchSize + 1
                    if searchSize >= i+1 # Ensure that there is a k <= kmax constraint / limit
                        break
                    else
                        if i >= 8 && searchSize < i-2 # go quicker through small search areas when kMax is big.
                            searchSize += 2
                        else
                            searchSize += 1
                        end
                    end

                    iter += 1
                    # Rerun algorithm
                    continue
                end

                # Returns optimal with better solution
                if status == :Optimal && obj+1e-3 < bestObj
                    #println("Went into ",status, ", with better objective")
                    bestObj = obj
                    bestbSolved = bSolved

                    # Add constraint to rightBranchMatrix Delta(x,x_old) >= k + 1
                    # Diversify due to an optimale worse solution previously found or due to solving the last node to optimality
                    if typeof(leftBranchMatrix) == Int64 || typeof(leftBranchMatrix) == Float64
                        if searchSize < i+1
                            if i >= 8 && searchSize < i-2
                                searchSize += 2
                            else
                                searchSize += 1
                            end
                            iter += 1
                            continue
                        else
                            break
                        end
                    end

                    #println("leftBranchMatrix = ", leftBranchMatrix)
                    rightBranchBool = true
                    rightBranchMatrix = hcat(rightBranchMatrix, leftBranchMatrix)

                    # Add constraints to leftBranchMatrix Delta(x,x_new) <= k
                    leftBranchMatrix = bSolved

                    iter += 1
                    # Rerun algorithm
                    #println("Status = ", status,"\tObjective = ", obj, "\tgap = ", gap)
                    continue
                end

                # Return optimal with same or worse solution
                if status == :Optimal && obj+1e-3 >= bestObj
                    println("Went into ",status, ", with worse objective")

                    # Diversify due to an optimale worse solution previously found or due to solving the last node to optimality
                    if typeof(leftBranchMatrix) == Int64 || typeof(leftBranchMatrix) == Float64
                        if searchSize < i+1
                            if i >= 8 && searchSize < i-2
                                searchSize += 2
                            else
                                searchSize += 1
                            end
                        else
                            #Break loop as we can't diversify any more.
                            break
                        end
                    end

                    # Diversify due to an optimale worse solution previously found
                    if typeof(leftBranchMatrix) == Int64 || typeof(leftBranchMatrix) == Float64
                        leftBranchMatrix = 0
                        iter += 1
                        continue
                    else #First time we see optimal with worse solution -- Add leftbranch constraint to right branch and tries to solve rootnode
                        rightBranchBool = true
                        rightBranchMatrix = hcat(rightBranchMatrix, leftBranchMatrix)
                    end
                    leftBranchMatrix = 0 #Set leftBranchMatrix to an integer/float value

                    # Add constraint to rightBranchMatrix Delta(x,x_old) >= k + 1
                    # Solve to depth / timeLimit reached

                    iter += 1
                    continue
                end

                # Return TimeLimit / UserLimit with better solution
                if status == :UserLimit && obj+1e-3 < bestObj
                    #println("Went into ",status, ", with better objective")
                    bestObj = obj
                    bestbSolved = bSolved

                    # Update constraints to leftBranchMatrix Delta(x,x_new) <= k
                    leftBranchMatrix = bSolved

                    # Rerun algorith
                    iter += 1
                    continue
                end

                # Return TimeLimit / UserLimit with worse solution
                if status == :UserLimit && obj+1e-3 >= bestObj
                    #println("Went into ",status, ", with worse objective")
                    # End algorithm under the argument that the stanadard implementation would have solved it by now
                    iter += 1
                    break
                end
			end #While loops ends

            bestObjtives[Int64(i), Int64(g)] = bestObj

			#println("Number of non zero actually: ", countnz(bSolved))

			#println("Prediction =$prediction\t Real =$standYVali\t kMax =$i \t gamma =$g")
			#realArr[r,1] = standYVali[1]
            prediction = standXVali*bestbSolved

			println("Prediction =$prediction\t Real =$standYVali\t kMax =$i \t gamma =$g \t obj = $(bestObjtives[Int64(i),Int64(g)])")
		end

	end
    return bestObjtives
end

function AALR_Time_Run(standX, standY, standXVali, standYVali, trainingData, r)
    nRows = size(standX)[1]

    bSample = []
    allCuts = []
    signifBoolean = zeros(3)

    #mathheuristic
    warmstart = false
	bestObjtivesMath = @time(solveAndLogForAllKMath(kmax, standXVali, standYVali, r))
    println("This was time for mathheuristic")

    #Nonwarmstart
    warmstart = false
    bestObjtivesNormal = @time(solveAndLogForAllK(kmax, standXVali, standYVali, r, warmstart))
    print("This was time for non warmstart")
    println("")

    #lasso warmstart
    warmstart = true
    bestObjtivesWarmstart = @time(solveAndLogForAllK(kmax, standXVali, standYVali, r, warmstart))
    println("This was time for warmstart")

    #Heuristic warmstart
    println("Fucking start here!")
    bestObjtivesHeuristic = @time(solveAndLogForAllKHeuristics(kmax, standXVali, standYVali, r))
    println("This was time for heuristc")

    return bestObjtivesNormal, bestObjtivesWarmstart, bestObjtivesMath, bestObjtivesHeuristic
end

SSTO = sum((standY[i]-mean(standY[:]))^2 for i=1:length(standY))
#gammaArray = log10.(logspace(0.01, SSTO, 3))
gammaArray = log10.(logspace(0.01, 5, 3))

bestObjtivesNormal = 1
bestObjtivesWarmstart = 1
bestObjtivesMath = 1

#writedlm(path*"/bestNorm.csv",bestObjtivesNormal)
#writedlm(path*"/bestWarm.csv",bestObjtivesWarmstart)
#writedlm(path*"/bestMath.csv",bestObjtivesMath)

r=1
for r = 1:100:501#(nRows-trainingSize-predictions)
	standX = allData[r:(trainingSize+r), 1:bCols]
    nRows = size(Xtrain)[1]
    trainingData = Xtrain[:,1:nCols-1]
	standY = allData[r:(trainingSize+r), bCols+1]
	standXVali  = allData[(trainingSize+r+1):(trainingSize+r+predictions), 1:bCols]
	standYVali  = allData[(trainingSize+r+1):(trainingSize+r+predictions), bCols+1]
	bestObjtivesNormal, bestObjtivesWarmstart, bestObjtivesMath, bestObjtivesHeuristic  = AALR_Time_Run(standX, standY, standXVali, standYVali, trainingData, r)

    writedlm(path*string(r)*"bestNorm.csv",bestObjtivesNormal)
    writedlm(path*string(r)*"bestWarm.csv",bestObjtivesWarmstart)
    writedlm(path*string(r)*"bestMath.csv",bestObjtivesMath)
    writedlm(path*string(r)*"bestHeuristic.csv",bestObjtivesHeuristic)
end

max_stars_repo_name: skipperafrosenborg/Thesis

max_issues_repo_name: skipperafrosenborg/Thesis

max_forks_repo_name: skipperafrosenborg/Thesis


text:
#------------------------------------------------------------------------------------------------------------
#   Code to solve elastic net-ish version of covariance estimation
#   min_{X \in \R^{d x r}} F(X) :=  E_{a,b} | <XX^T, aa^T> - b | + (<XX^T, aa^T> - b )^2
#   for a bunch of trials
#------------------------------------------------------------------------------------------------------------
using Random  # for setting the random seed
using PyPlot
using LaTeXStrings, Printf  # for plot labels

include("solve_enet.jl");

#-------------------------------------
#   Set parameters
#-------------------------------------
r = 2;  # rank
dims = [100];  # ambient dimension
dims_colors = ["#1f78b4"]; #, "#33a02c"];
stoch_err = 0.1;  # standard deviation of errors in stochastic measurements b

maxIter = 2000;
dims_steps_subgrad = [1e-6];
dims_steps_mirror = [0.1];

Random.seed!(123);  # for reproducibility

#-------------------------------------
#   Initialize plot
#-------------------------------------
dist_fig = figure(figsize=[10,6]);
xlabel(L"Iteration $k$");
ylabel(L"$\min_U \; \frac{\| X_k U - X_{true} \|_2^2 }{\| X_{true} \|_2^2}$");
title(@sprintf("Distance to solution for elastic net covariance estimation (r=%i)", r));
xlim(0,maxIter)

funval_fig = figure(figsize=[10,6]);
xlabel(L"Iteration $k$");
ylabel("Empirical Function Value");
title(@sprintf("Function values for elastic net covariance estimation (r=%i)", r));
xlim(0,maxIter)

#-----------------------------------------------------------------------------------------
#   Run stochastic subgradient method for each dimension
#-----------------------------------------------------------------------------------------
for i in 1:length(dims)
    d = dims[i];

    # Generate true matrix we are searching for
    Xtrue = randn(d,r);

    # Generate initial point
    radius = 1.0;  # 2-norm measure of relative deviation between Xtrue and Xinit
    pert = randn(d,r);
    Xinit = Xtrue + radius * (norm(Xtrue, 2) / norm(pert, 2)) * pert;

    # Specify step sizes
    steps_subgrad = fill(dims_steps_subgrad[i], maxIter);
    steps_mirror = fill(dims_steps_mirror[i], maxIter);

    # Run the two methods
    (err_hist_subgrad, fun_hist_subgrad) = solve_enet(Xinit, Xtrue, stoch_err, maxIter,
                                                                                steps_subgrad,
                                                                                method="subgradient")
    (err_hist_mirror, fun_hist_mirror) = solve_enet(Xinit, Xtrue, stoch_err, maxIter,
                                                                                steps_mirror,
                                                                                method="mirror")

    # Add distances to appropriate plot
    plt[:figure](dist_fig[:number])
    semilogy(err_hist_subgrad, linestyle="--", color=dims_colors[i],
                    label=@sprintf("SGD, d=%i", d));
    semilogy(err_hist_mirror, color=dims_colors[i], label=@sprintf("SMD, d=%i", d));

    # Add function values to appropriate plot
    plt[:figure](funval_fig[:number])
    semilogy(fun_hist_subgrad, linestyle="--", color=dims_colors[i],
                    label=@sprintf("SGD, d=%i", d));
    semilogy(fun_hist_mirror, color=dims_colors[i], label=@sprintf("SMD, d=%i", d));
end

# Add legends to plots and save final results
plt[:figure](dist_fig[:number])
legend(loc="lower right")
savefig("plots/enet_distances.pdf");

plt[:figure](funval_fig[:number])
legend(loc="lower right")
savefig("plots/enet_values.pdf");

max_stars_repo_name: kelliemac/stochastic-model-based

max_issues_repo_name: kelliemac/stochastic-model-based

max_forks_repo_name: kelliemac/stochastic-model-based


text:
using SafeTestsets

@safetestset "Element tests" begin include("element_test.jl") end

@safetestset "CartesianMesh tests" begin include("cartesianmesh_test.jl") end

max_stars_repo_name: WaveProp/WavePropBase

max_issues_repo_name: WaveProp/WavePropBase

max_forks_repo_name: WaveProp/WavePropBase


text:
#Visualization Functions
export visualize

using Plots

gr()

function _plot(c::confusion_matrix; func, type, title, labels = nothing)
    x = nothing
    if type == "condition-positive"; x = func(labels, condition_positive(c) , title = title, labels = permutedims(labels))
    elseif type == "condition-negative"; x = func(labels, condition_negative(c) , title = title, labels = permutedims(labels))
    elseif type == "predicted-positive"; x = func(labels, predicted_positive(c) , title = title, labels = permutedims(labels))
    elseif type == "predicted-negative"; x = func(labels, predicted_positive(c) , title = title, labels = permutedims(labels))
    elseif type == "correctly-classified"; x = func(labels, correctly_classified(c) , title = title, labels = permutedims(labels))
    elseif type == "incorrectly-classified"; x = func(labels, incorrectly_classified(c) , title = title, labels = permutedims(labels))
    elseif type == "sensitivity-score"; x = func(labels, sensitivity_score(c) , title = title, labels = permutedims(labels))
    elseif type == "recall-score"; x = func(labels, recall_score(c), title = title, labels = permutedims(labels))
    elseif type == "specificity-score";x =  func(labels, specificity_score(c) ,title = title, labels = permutedims(labels))
    elseif type == "precision-score";x =  func(labels, precision_score(c), title = title, labels = permutedims(labels))
    elseif type == "positive-predictive-value"; x = func(labels, positive_predictive_value(c) , title = title, labels = permutedims(labels))
    elseif type == "accuracy-score"; x = func(labels, accuracy_score(c), title = title, labels = permutedims(labels))
    elseif type == "balanced-accuracy-score"; x = func(labels, balanced_accuracy_score(c), title = title, labels = permutedims(labels))
    elseif type == "negative-predictive-value"; x = func(labels, negative_predictive_value(c), title = title, labels = permutedims(labels))
    elseif type == "false-negative-rate"; x = func(labels, false_negative_rate(c), title = title, labels = permutedims(labels))
    elseif type == "false-positive-rate"; x = func(labels, false_positive_rate(c), title = title, labels = permutedims(labels))
    elseif type == "false-discovery-rate"; x = func(labels, false_discovery_rate(c), title = title, labels = permutedims(labels))
    elseif type == "false-omission-rate"; x = func(labels, false_omission_rate(c), title = title, labels = permutedims(labels))
    elseif type == "f1-score"; x = func(labels, f1_score(c), title = title, labels = permutedims(labels))
    elseif type == "prevalence-threshold"; x = func(labels, prevalence_threshold(c), title = title, labels = permutedims(labels))
    elseif type == "threat-score"; x = func(labels, threat_score(c), title = title, labels = permutedims(labels))
    elseif type == "matthews-correlation-coeff"; x = func(labels, matthews_correlation_coeff(c), title = title, labels = permutedims(labels))
    elseif type == "fowlkes-mallows-index"; x = func(labels,fowlkes-mallows-index(c), title = title, labels = permutedims(labels))
    elseif type == "informedness"; x = func(labels, informedness(c), title = title, labels = permutedims(labels))
    elseif type == "markedness"; x = func(labels, markedness(c), title = title, labels = permutedims(labels))
    elseif type == "cohen-kappa-score"; x = func(labels, cohen_kappa_score(c), title = title, labels = permutedims(labels))
    elseif type == "hamming-loss"; x = func(labels, hamming_loss(c), title = title, labels = permutedims(labels))
    elseif type == "jaccard-score"; x = func(labels, jaccard_score(c), title = title, labels = permutedims(labels))
    end
    return x
end

"""
```visualize(c::confusion_matrix, <keywords>)```

Visualize the given properties of the confusion matrix as specified

## Keywords
**`mode`** : String, Array of String; represents properties given below of the confusion matrix, can be either an array containing different properties
or a single string.\n
    Supported Modes: \n
    - `matrix`\n
    - `condition-positive`\n
    - `condition-negative`\n
    - `predicted-positive\n
    - `predicted-negative`\n
    - `correctly-classified`\n
    - `incorrectly-classified`\n
    - `sensitivity-score`\n
    - `recall-score`\n
    - `specificity-score`\n
    - `precision-score`\n
    - `positive-predictive-value`\n
    - `accuracy-score`\n
    - `balanced-accuracy-score`\n
    - `negative-predictive-value`\n
    - `false-negative-rate`\n
    - `false-positive-rate`\n
    - `false-discovery-rate`\n
    - `false-omission-rate`\n
    - `f1-score`\n
    - `prevalence-threshold`\n
    - `threat-score`\n
    - `matthews-correlation-coeff`\n
    - `fowlkes-mallows-index`\n
    - `informedness`\n
    - `markedness`\n
    - `cohen-kappa-score`\n
    - `hamming-loss`\n
    - `jaccard-score`\n

**`seriestype`** : String, denotes which visualization function will be used; default: heatmap
    Supported visualization functions:
        - `heatmap`
        - `bar`
        - `histogram`
        - `scatter`
        - `line`

**`title`** : String, denotes the title that will displayed above the drawn plot, default: nothing

**`labels`** : Vector, denotes the labels that will be used for the plot. If equals to nothing, labels of the given confusion matrix will be used.

"""
function visualize(c::confusion_matrix; mode = "matrix", seriestype::String = "heatmap", title= nothing, labels = nothing)
    @assert seriestype in ["scatter", "heatmap", "line", "histogram", "bar"] "Unknown visualization format"
    labels = labels != nothing ? labels : convert(Array{typeof(c.Labels[1])}, c.Labels)
    if title == nothing; title = mode isa Array ? mode : String(Base.copymutable(mode)); end
    title = title isa Array ? title : [title]
    mode = mode isa Array ? mode : [mode]
    plt = []
    for i in 1:length(mode)
        @assert mode[i] in ["matrix", "condition-positive", "condition-negative", "predicted-positive","predicted-negative", "correctly-classified", "incorrectly-classified", "sensitivity-score", "recall-score", "specificity-score", "precision-score", "positive-predictive-value", "accuracy-score", "balanced-accuracy-score", "negative-predictive-value", "false-negative-rate", "false-positive-rate", "false-discovery-rate",
         "false-omission-rate", "f1-score", "prevalence-threshold", "threat-score", "matthews-correlation-coeff", "fowlkes-mallows-index",
         "informedness", "markedness", "cohen-kappa-score", "hamming-loss", "jaccard-score"] "Unknown visualization mode"
        if mode[i] != "matrix"; @assert seriestype in ["scatter", "line", "histogram", "bar"] "The given mode does not support this visualization format"; end
        x = nothing
        if mode[i] == "matrix"
            if seriestype == "histogram"; x = histogram(labels, c.matrix, labels = permutedims(labels), title = title[i])
            elseif seriestype == "scatter"; x = scatter(labels, c.matrix, labels = permutedims(labels), title = title[i])
            elseif seriestype == "line"; x = plot(labels, c.matrix, labels = permutedims(labels), title = title[i])
            elseif seriestype == "bar"; x = bar(labels, c.matrix, labels = permutedims(labels), title = title[i])
            elseif seriestype == "heatmap"; x = heatmap(labels, labels, c.matrix, labels = permutedims(labels), title = title[i])
            end
        else
            if seriestype == "histogram"; x = _plot(c; func = histogram, type = mode[i], title = title[i], labels = labels)
            elseif seriestype == "scatter"; x = _plot(c; func = scatter, type = mode[i], title = title[i], labels = labels)
            elseif seriestype == "line"; x =  _plot(c; func = plot, type = mode[i], title = title[i], labels = labels)
            elseif seriestype == "bar"; x =  _plot(c; func = bar, type = mode[i], title = title[i], labels = labels)
            elseif seriestype == "heatmap"; x =  _plot(c; func = heatmap, type = mode[i], title = title[i], labels = labels)
            end
        end
        push!(plt, x)
    end
    plot(plt..., layout = (length(plt), 1))
end


text:
using ThreeBodyDecay
using Parameters
using StaticArrays
using PartialWaveFunctions
using Plots
using Plots.PlotMeasures: mm
using LaTeXStrings

struct AngularBusiness
    k::Int
    two_js::Union{ThreeBodySpins, SVector{4, Int}}
    two_ls::Union{NamedTuple,Tuple}
    two_LS::Union{NamedTuple,Tuple}
    two_j::Int
end

function angular_ampitude(z, two_λs, ab::AngularBusiness)
    k = ab.k
    #
    @unpack k, two_ls, two_LS, two_js, two_j = ab
    i,j = ij_from_k(k)
    #
    two_l, two_s = ab.two_ls
    two_L, two_S = ab.two_LS
    #
    hR = jls_coupling(two_js[i], two_λs[i], two_js[j], two_λs[j], two_j, two_l, two_s)
    h0 = jls_coupling(two_j, two_λs[4]+two_λs[k], two_js[k], two_λs[k], two_js[4], two_L, two_S)
    #
    return h0*wignerd_doublearg(two_j,two_λs[4]+two_λs[k],two_λs[i]-two_λs[j], z)*hR
end

angular_intensity(z,ab::AngularBusiness) =
    sum(abs2, angular_ampitude(z,two_λs,ab) for two_λs in itr(ab.two_js))
#
#
function explore(k, jp::jp, jps::Vector)
    i,j = ij_from_k(k)
    lsv = possible_ls(jps[i], jps[j]; jp = jp)
    LSv = possible_ls(jp,  jps[k]; jp = jps[4])
    ABs = [AngularBusiness(k, SVector{4}(two_j.(jps)), x2.(ls), x2.(LS), two_j(jp))
        for ls in lsv, LS in LSv]
    return ABs
end
#
#
function vector_of_vector_of_plots(vov)
    l = max(length.(vov)...)
    plot(layout=grid(3,l), size=(300*l,300*3))
    for (i,ab_all) in enumerate(vov)
        for k in 1:l
            sp = (i-1)*l + k
            if k > length(ab_all)
                plot!(sp=sp, xaxis=false, yaxis=false, grid=false)
                continue
            end
            ab = ab_all[k]
            plot!(sp=sp, ab)
        end
    end
    plot!()
end

@recipe function f(vov::Vector{Vector{T}} where T)
    l = max(length.(vov)...)
    #
    size --> (300*l,300*length(vov))
    layout := grid(length(vov), l)
    #
    for (i,objects) in enumerate(vov)
        for k in 1:l
            sp = (i-1)*l + k
            if k > length(objects)
                @series begin
                    label := ""
                    subplot := sp
                    xaxis := false
                    yaxis := false
                    grid := false
                    ()
                end
                continue
            end
            ab = objects[k]
            @series begin
                subplot := sp
                ab
            end
        end
    end
end

@recipe function f(ab::AngularBusiness)
    xv = range(-1,1, length=100)
    calv = angular_intensity.(xv, Ref(ab))

    framestyle --> :origin
    title --> "l=$(over2(ab.two_ls[1])), L=$(over2(ab.two_LS[1]))"

    xv, calv
end

# Ωb⁻ → Ξc⁺π⁻K⁻
let
    jps_pc = [jp"1/2+", jp"0-", jp"0-", jp"1/2+"]
    jps_pv = [jp"1/2+", jp"0-", jp"0-", jp"1/2-"]

    jpRv = [jp"1/2+", jp"3/2+", jp"5/2+"]

    ABs_all = []
    for jp in jpRv
        ABm = explore.(3, Ref(jp), [jps_pc, jps_pv])
        push!(ABs_all, hcat(ABm...))
    end
    plot([vcat(v...) for v in ABs_all], lab="", xlab=L"\cos\,\theta")
end

# B⁻ → Λc⁻Λc⁺K⁻
let
    jps_pc = [jp"1/2-", jp"1/2+", jp"0-", jp"0+"]
    jps_pv = [jp"1/2-", jp"1/2+", jp"0-", jp"0-"]

    jpRv = [jp"1/2+", jp"3/2+", jp"5/2+"]

    ABs_all = []
    for jp in jpRv
        ABm = explore.(1, Ref(jp), [jps_pc, jps_pv])
        push!(ABs_all, hcat(ABm...))
    end
    plot([vcat(v...) for v in ABs_all], lab="", xlab=L"\cos\,\theta")
end

# Λb⁰ → J/ψ p K⁻
let
    jps_pc = [jp"1-", jp"1/2+", jp"0-", jp"1/2+"]
    jps_pv = [jp"1-", jp"1/2+", jp"0-", jp"1/2-"]

    jpRv = [jp"1/2+", jp"3/2+", jp"5/2+"]

    ABs_all = []
    for jp in jpRv
        ABm = explore.(3, Ref(jp), [jps_pc, jps_pv])
        push!(ABs_all, hcat(ABm...))
    end
    plot([vcat(v...) for v in ABs_all], lab="", xlab=L"\cos\,\theta")
end

# Ωb⁻ → Ωc⁰π⁻π⁰
let
    jps_pc = [jp"1/2+", jp"0-", jp"0-", jp"1/2+"]
    jps_pv = [jp"1/2+", jp"0-", jp"0-", jp"1/2-"]

    jpRv = [jp"1-"]

    ABs_all = []
    for jp in jpRv
        ABm = explore.(1, Ref(jp), [jps_pc, jps_pv])
        push!(ABs_all, hcat(ABm...))
    end
    plot([vcat(v...) for v in ABs_all], lab="", xlab=L"\cos\,\theta", bottom_margin=5mm)
end

max_stars_repo_name: mmikhasenko/ThreeBodyDecay

max_issues_repo_name: mmikhasenko/ThreeBodyDecay

max_forks_repo_name: mmikhasenko/ThreeBodyDecay


text:
# # # Higher level debug info support
struct DebugSections{T<:ObjectHandle, S}
    oh::T
    debug_abbrev::Nullable{S}
    debug_aranges::Nullable{S}
    debug_frame::Nullable{S}
    debug_info::Nullable{S}
    debug_line::Nullable{S}
    debug_loc::Nullable{S}
    debug_macinfo::Nullable{S}
    debug_pubnames::Nullable{S}
    debug_ranges::Nullable{S}
    debug_str::Nullable{S}
    debug_types::Nullable{S}
end
handle(dbgs::DebugSections) = dbgs.oh

function DebugSections(oh::H; debug_abbrev = nothing, debug_aranges = nothing,
        debug_frame = nothing, debug_info = nothing, debug_line = nothing,
        debug_macinfo = nothing, debug_pubnames = nothing, debug_loc= nothing,
        debug_ranges = nothing, debug_str = nothing, debug_types = nothing) where {H <: ObjectHandle}
    # Build Section typeunion
    S = Union{map(typeof, [debug_abbrev, debug_aranges, debug_frame, debug_info,
        debug_line, debug_loc, debug_macinfo, debug_pubnames, debug_ranges,
        debug_str, debug_types])...}
    DebugSections{T,S}(oh, debug_abbrev, debug_aranges, debug_frame, debug_info,
        debug_line, debug_loc, debug_macinfo, debug_pubnames, debug_ranges,
        debug_str, debug_types)
end

function DebugSections(oh::H, sections::Dict) where {H <: ObjectHandle}
    DebugSections(oh,
        debug_abbrev = get(sections, "debug_abbrev", nothing),
        debug_aranges = get(sections, "debug_aranges", nothing),
        debug_frame = get(sections, "debug_frame", nothing),
        debug_info = get(sections, "debug_info", nothing),
        debug_line = get(sections, "debug_line", nothing),
        debug_loc = get(sections, "debug_loc", nothing),
        debug_macinfo = get(sections, "debug_macinfo", nothing),
        debug_pubnames = get(sections, "debug_pubnames", nothing),
        debug_ranges = get(sections, "debug_ranges", nothing),
        debug_str = get(sections, "debug_str", nothing),
        debug_types = get(sections, "debug_types", nothing))
end

const DEBUG_SECTIONS = [
    "debug_abbrev",
    "debug_aranges",
    "debug_frame",
    "debug_info",
    "debug_line",
    "debug_loc",
    "debug_macinfo",
    "debug_pubnames",
    "debug_ranges",
    "debug_str",
    "debug_types"]

function show(io::IO, dsect::DebugSections)
    println(io, "Debug Sections for $(dsect.oh)")
    println(io, "========================= debug_abbrev =========================")
    println(io, dsect.debug_abbrev)
    println(io, "======================== debug_aranges =========================")
    println(io,  dsect.debug_aranges)
    println(io, "========================= debug_frame ==========================")
    println(io, dsect.debug_frame)
    println(io, "========================= debug_info ===========================")
    println(io, dsect.debug_info)
    println(io, "========================= debug_line ===========================")
    println(io, dsect.debug_line)
    println(io, "========================== debug_loc ===========================")
    println(io, dsect.debug_loc)
    println(io, "======================== debug_macinfo =========================")
    println(io, dsect.debug_macinfo)
    println(io, "======================= debug_pubnames =========================")
    println(io, dsect.debug_pubnames)
    println(io, "======================== debug_ranges ==========================")
    println(io, dsect.debug_ranges)
    println(io, "=========================== debug_str ==========================")
    println(io, dsect.debug_str)
    println(io, "========================= debug_types ==========================")
    println(io, dsect.debug_types)
end

"""
    debugsections(oh::ObjectHandle)

Return list of sections needed for debug info within given `ObjectHandle`
"""
@mustimplement debugsections(oh::ObjectHandle)


text:
# TODO: Document additional trick functionality
"""
    gbtranspose!(C::GBMatrix, A::GBMatrix; kwargs...)::Nothing

Eagerly evaluated matrix transpose, storing the output in `C`.

# Arguments
- `C::GBMatrix`: output matrix.
- `A::GBMatrix`: input matrix.

# Keywords
- `mask::Union{Ptr{Nothing}, GBMatrix} = C_NULL`: optional mask.
- `accum::Union{Ptr{Nothing}, AbstractBinaryOp} = C_NULL`: binary accumulator operation
    where `C[i,j] = accum(C[i,j], T[i,j])` where T is the result of this function before accum is applied.
- `desc::Descriptor = DEFAULTDESC`
"""
function gbtranspose!(
    C::GBMatrix, A::GBMatOrTranspose;
    mask = nothing, accum = nothing, desc = nothing
)
    mask, accum = _handlenothings(mask, accum)
    desc === nothing && (desc = DEFAULTDESC)
    if A isa Transpose && desc.input1 == TRANSPOSE
        throw(ArgumentError("Cannot have A isa Transpose and desc.input1 = TRANSPOSE."))
    elseif A isa Transpose
        A = A.parent
        desc = desc + T0
    end
    accum = getaccum(accum, eltype(C))
    libgb.GrB_transpose(C, mask, accum, A, desc)
    return C
end

"""
    gbtranspose(A::GBMatrix; kwargs...)::GBMatrix

Eagerly evaluated matrix transpose which returns the transposed matrix.

# Keywords
- `mask::Union{Ptr{Nothing}, GBMatrix} = C_NULL`: optional mask.
- `accum::Union{Ptr{Nothing}, AbstractBinaryOp} = C_NULL`: binary accumulator operation
    where `C[i,j] = accum(C[i,j], T[i,j])` where T is the result of this function before accum is applied.
- `desc::Descriptor = DEFAULTDESC`

# Returns
- `C::GBMatrix`: output matrix.
"""
function gbtranspose(
    A::GBMatOrTranspose;
    mask = C_NULL, accum = nothing, desc::Descriptor = DEFAULTDESC
)
    C = similar(A, size(A,2), size(A, 1))
    gbtranspose!(C, A; mask, accum, desc)
    return C
end

function LinearAlgebra.transpose(A::GBMatrix)
    return Transpose(A)
end

#TODO: This should be lazy
function LinearAlgebra.transpose(v::GBVector)
    A = GBMatrix{eltype(v)}(size(v, 2), size(v, 1))
    nz = findnz(v)
    for i ∈ 1:length(nz[1])
        A[1, nz[1][i]] = nz[2][i]
    end
    return A
end

function Base.copy!(
    C::GBMatrix, A::LinearAlgebra.Transpose{<:Any, <:GBMatrix};
    mask = C_NULL, accum = nothing, desc::Descriptor = C_NULL
)
    return gbtranspose!(C, A.parent; mask, accum, desc)
end

function Base.copy(
    A::LinearAlgebra.Transpose{<:Any, <:GBMatrix};
    mask = C_NULL, accum = nothing, desc::Descriptor = DEFAULTDESC
)
    return gbtranspose(A.parent; mask, accum, desc)
end

function _handletranspose(
    A::Union{GBArray, Nothing} = nothing,
    desc::Union{Descriptor, Nothing, Ptr{Nothing}} = nothing,
    B::Union{GBArray, Nothing} = nothing
)
    if desc == C_NULL
        desc = DEFAULTDESC
    end
    if A isa Transpose
        desc = desc + T0
        A = A.parent
    end
    if B isa Transpose
        desc = desc + T1
        B = B.parent
    end
    return A, desc, B
end

#This is ok per the GraphBLAS Slack channel. Should change its effect on Complex input.
LinearAlgebra.adjoint(A::GBMatrix) = transpose(A)

LinearAlgebra.adjoint(v::GBVector) = transpose(v)

#arrrrgh, type piracy.
LinearAlgebra.transpose(::Nothing) = nothing


text:
abstract ModelGenerator

function next_model(gen::ModelGenerator, rmdp::RobustMDP, s, a, snode::RobustStateNode, anode::RobustActionNode)
    next_model(gen, rmdp, s, a)
end

type RandomModelGenerator <: ModelGenerator
    rng::AbstractRNG
end
RandomModelGenerator() = RandomModelGenerator(MersenneTwister())

function next_model(gen::RandomModelGenerator, rmdp::RobustMDP, s, a, snode::RobustStateNode, anode::RobustActionNode)
    next_model(gen, rmdp, s, a)
end

next_model(gen::RandomModelGenerator, rmdp::RobustMDP, s, a) = rand_model(rmdp, gen.rng)


text:
using ImageMagick
using FileIO
using Documenter
using Highlights
using Markdown
using Random
using GLMakie
using AbstractPlotting
AbstractPlotting.inline!(true)
import AbstractPlotting: to_string

# Pause renderloop for slow software rendering.
# This way, we only render if we actualy save e.g. an image
GLMakie.set_window_config!(;
    framerate = 15.0,
    pause_rendering = true
)
# ImageIO seems broken on 1.6 ... and there doesn't
# seem to be a clean way anymore to force not to use a loader library?
filter!(x-> x !== :ImageIO, FileIO.sym2saver[:PNG])
filter!(x-> x !== :ImageIO, FileIO.sym2loader[:PNG])

"""
    print_table(io::IO, dict::Dict)

Print a Markdown-formatted table with the entries from `dict` to specified `io`.
"""
function print_table(io::IO, dict::Dict)
    # get max length of the keys
    k = string.("`", collect(keys(dict)), "`")
    maxlen_k = max(length.(k)...)

    # get max length of the values
    v = string.(collect(values(dict)))
    maxlen_v = max(length.(v)...)

    j = sort(collect(dict), by = x -> x[1])

    # column labels
    labels = ["Symbol", "Description"]

    # print top header
    print(io, "|")
    print(io, "$(labels[1])")
    print(io, " "^(maxlen_k - length(labels[1])))
    print(io, "|")
    print(io, "$(labels[2])")
    print(io, " "^(maxlen_v - length(labels[2])))
    print(io, "|")
    print(io, "\n")

    # print second line (toprule)
    print(io, "|")
    print(io, "-"^maxlen_k)
    print(io, "|")
    print(io, "-"^maxlen_v)
    print(io, "|")
    print(io, "\n")

    for (idx, entry) in enumerate(j)
        print(io, "|")
        print(io, "`$(entry[1])`")
        print(io, " "^(maxlen_k - length(string(entry[1])) - 2))
        print(io, "|")
        print(io, "$(entry[2])")
        print(io, " "^(maxlen_v - length(entry[2])))
        print(io, "|")
        print(io, "\n")
    end
end

################################################################################
#                              Utility functions                               #
################################################################################

################################################################################
#                                    Setup                                     #
################################################################################

pathroot  = normpath(@__DIR__, "..")
docspath  = joinpath(pathroot, "docs")
srcpath   = joinpath(docspath, "src")
srcgenpath   = joinpath(docspath, "src_generation")
buildpath = joinpath(docspath, "build")
genpath   = joinpath(srcpath, "generated")

mkpath(genpath)

################################################################################
#                          Syntax highlighting theme                           #
################################################################################

@info("Writing highlighting stylesheet")

open(joinpath(srcpath, "assets", "syntaxtheme.css"), "w") do io
    Highlights.stylesheet(io, MIME("text/css"), Highlights.Themes.DefaultTheme)
end

################################################################################
#                      Automatic Markdown page generation                      #
################################################################################

########################################
#       Plot attributes overview       #
########################################

# automatically generate an overview of the plot attributes (keyword arguments), using a source md file
@info("Generating attributes page")
# axis_attr_list = []
# for a in (Axis3D,)
#     attr = keys(default_theme(nothing, a))
#     push!(axis_attr_list, attr...)
# end
# axis_attr_list = string.(sort!(unique(axis_attr_list)))

# dict = default_theme(nothing, Axis3D)
# print_rec(STDOUT, dict)

const Axis3D_attr_desc = Dict(
    :frame => "See the detailed descriptions for `frame` attributes.",
    :names => "See the detailed descriptions for `names` attributes.",
    :scale => "NTuple{3,Float}. Specifies the scaling for the axes.",
    :showaxis => "NTuple{3,Bool}. Specifies whether to show the axes.",
    :showgrid => "NTuple{3,Bool}. Specifies whether to show the axis grids.",
    :showticks => "NTuple{3,Bool}. Specifies whether to show the axis ticks.",
    :ticks => "See the detailed descriptions for `ticks` attributes."
)

# frame
const Axis3D_attr_frame = Dict(
    :axiscolor => "Symbol or Colorant. Specifies the color of the axes. Can be a color symbol/string like :red, or a Colorant.",
    :linewidth => "Number. Width of the axes grid lines.",
    :linecolor => "Symbol or Colorant. Specifies the color of the grid lines. Can be a color symbol/string like :red, or a Colorant."
)

# names
const Axis3D_attr_names = Dict(
    :align => "`NTuple{3,(:pos, :pos)}`. Specify the text alignment for the axis labels, where `:pos` can be `:left`, `:center`, or `:right`.",
    :axisnames => "NTuple{3,String}. Specifies the axis labels.",
    :font => "NTuple{3,String}. Specifies the font for the axis labels, and can choose any font available on the system.",
    :gap => "Number. Specifies the gap (in pixels) between the axis labels and the axes themselves.",
    :rotation => "NTuple{3,Quaternion{Float32}}. Specifies the rotations for each axis's label, in radians.",
    :textcolor => "NTuple{3,Symbol or Colorant}. Specifies the color of the axes labels. Can be a color symbol/string like :red, or a Colorant.",
    :textsize => "NTuple{3,Int}. Font pointsize for axes labels."
)

# ticks
const Axis3D_attr_ticks = Dict(
    :align => "`NTuple{3,(:pos, :pos)}`. Specify the text alignment for the axis ticks, where `:pos` can be `:left`, `:center`, or `:right`.",
    :font => "NTuple{3,String}. Specifies the font for the axis ticks, and can choose any font available on the system.",
    :gap => "Number. Specifies the gap (in pixels) between the axis ticks and the axes themselves.",
    :rotation => "NTuple{3,Quaternion{Float32}}. Specifies the rotations for each axis's ticks, in radians.",
    :textcolor => "NTuple{3,Symbol or Colorant}. Specifies the color of the axes ticks. Can be a color symbol/string like :red, or a Colorant.",
    :textsize => "Integer. Font pointsize for text."
)

const Axis3D_attr_groups = Dict(
    :frame => Axis3D_attr_frame,
    :names => Axis3D_attr_names,
    :ticks => Axis3D_attr_ticks
)
# dict = default_theme(nothing, Axis2D)
# print_rec(STDOUT, dict)

const Axis2D_attr_desc = Dict(
    :frame => "See the detailed descriptions for `frame` attributes.",
    :grid => "See the detailed descriptions for `grid` attributes.",
    :names => "See the detailed descriptions for `names` attributes.",
    :ticks => "See the detailed descriptions for `ticks` attributes."
)

# frame
const Axis2D_attr_frame = Dict(
    :arrow_size => "Number. Size of the axes arrows.",
    :axis_position => "",
    :axis_arrow => "Bool. Toggles the axes arrows.",
    :frames => "NTuple{2,NTuple{2,Bool}}.",
    :linecolor => "Symbol or Colorant. Specifies the color of the grid lines. Can be a color symbol/string like :red, or a Colorant.",
    :linestyle => "",
    :linewidth => "Number. Widths of the axes frame lines."
)

# grid
const Axis2D_attr_grid = Dict(
    :linecolor => "Symbol or Colorant. Specifies the color of the grid lines. Can be a color symbol/string like :red, or a Colorant.",
    :linestyle => "",
    :linewidth => "NTuple{2, Number}. Width of the x and y grids."
)

# names
const Axis2D_attr_names = Dict(
    :align => "`(:pos, :pos)`. Specify the text alignment, where `:pos` can be `:left`, `:center`, or `:right`.",
    :axisnames => "NTuple{2,String}. Specifies the text labels for the axes.",
    :font => "NTuple{2,String}. Specifies the font and can name any font available on the system.",
    :rotation => "NTuple{3,Float32}. Specifies the rotations for each axis's label, in radians.",
    :textcolor => "NTuple{2,Symbol or Colorant}. Specifies the color of the axes labels. Can be a color symbol/string like :red, or a Colorant.",
    :textsize => "Integer. Font pointsize for text."
)

# ticks
const Axis2D_attr_ticks = Dict(
    :align => "`NTuple{2,(:pos, :pos)}`. Specify the text alignment for the axis ticks, where `:pos` can be `:left`, `:center`, or `:right`",
    :font => "NTuple{2,String}. Specifies the font and can name any font available on the system.",
    :gap => "Number. Specifies the gap (in pixels) between the axis tick labels and the axes themselves.",
    :linecolor => "NTuple{2,Symbol or Colorant}. Specifies the color of the grid lines. Can be a color symbol/string like :red, or a Colorant.",
    :linestyle => "",
    :linewidth => "NTuple{2,Number}. Width of the axes ticks.",
    :rotation => "NTuple{3,Float32}. Specifies the rotations for each axis's ticks, in radians.",
    :textcolor => "NTuple{2,Symbol or Colorant}. Specifies the color of the axes ticks. Can be a color symbol/string like :red, or a Colorant.",
    :textsize => "NTuple{2,Int}. Font pointsize for tick labels.",
    :title_gap => "Number. Specifies the gap (in pixels) between the axis titles and the axis tick labels."
)

const Axis2D_attr_groups = Dict(
    :frame => Axis2D_attr_frame,
    :grid => Axis2D_attr_grid,
    :names => Axis2D_attr_names,
    :ticks => Axis2D_attr_ticks
)

const plot_attr_desc = Dict(
    :absorption => "Float32. Sets the absorption value for `volume` plots.",
    :algorithm => "Algorithm to be used for `volume` plots. Can be one of `:iso`, `:absorption`, `:mip`, `:absorptionrgba`, or `:indexedabsorption`.",
    :align => "`(:pos, :pos)`. Specify the text alignment, where `:pos` can be `:left`, `:center`, or `:right`.",
    :alpha => "Float in [0,1]. The alpha value (transparency).",
    :color => "Symbol or Colorant. The color of the main plot element (markers, lines, etc.). Can be a color symbol/string like :red, or a Colorant.  Can also be an array or matrix of 'z-values' that are converted into colors by the colormap automatically.",
    :colormap => "The color map of the main plot. Call `available_gradients()` to see what gradients are available. Can also be used with any Vector{<: Colorant}, or e.g. [:red, :black], or `ColorSchemes.jl` colormaps (by `colormap = ColorSchemes.<colorscheme name>.colors`).",
    :colorrange => "A tuple `(min, max)` where `min` and `max` specify the data range to be used for indexing the colormap. E.g. color = [-2, 4] with colorrange = (-2, 4) will map to the lowest and highest color value of the colormap.",
    :fillrange => "Bool. Toggles range filling in `contour` plots.",
    :font => "String. Specifies the font, and can choose any font available on the system.",
    :glowcolor => "Color Type. Color of the marker glow (outside the border) in `scatter` plots.",
    :glowwidth => "Number. Width of the marker glow in `scatter` plots.",
    :image => "The image to be plotted on the plot.",
    :interpolate => "Bool. For `heatmap` and `images`. Toggles color interpolation between nearby pixels.",
    :isorange => "Float32. Sets the isorange for `volume` plots.",
    :isovalue => "Float32. Sets the isovalue for `volume` plots.",
    :levels => "Integer. Number of levels for a `contour`-type plot.",
    :linestyle => "Symbol. Style of the line (for `line` and `linesegments` plots). Available styles are `:dash`, `:dot`, `:dashdot`, and `:dashdotdot`. You can also supply an array describing the length of each gap/fill.",
    :linewidth => "Number. Width of the line in `line` and `linesegments` plots.",
    :marker => "Symbol, Char, Shape, or AbstractVector. Call `available_marker_symbols`() to see which ones.",
    :marker_offset => "Array of `GeometryBasics.Point`'s. Specifies the offset coordinates for the markers. See the [Marker offset](https://simondanisch.github.io/ReferenceImages/gallery/marker_offset/index.html) example.",
    :markersize => "Number or AbstractVector. Specifies size (radius pixels) of the markers.",
    :position => "NTuple{2,Float}, `(x, y)`. Specify the coordinates to position text at.",
    :rotation => "Float32. Specifies the rotation in radians.",
    :rotations => "AbstractVector{Float32}. Similar to `:rotation`, except it specifies the rotations for each element in the plot.",
    :shading => "Bool. Specifies if shading should be on or not (for meshes).",
    :strokecolor => "Color Type. Color of the marker stroke (border).",
    :strokewidth => "Number. Width of the marker stroke (in pixels).",
    :textsize => "Integer. Font pointsize for text.",
    :transformation => "`(:plane, location)`. Transforms the `:plane` to the specified location. Possible `:plane`'s are `:xy`, `:yz`, and `:xz`.",
    :visible => "Bool. Toggle visibility of plot."
)

path = joinpath(genpath, "plot-attributes.md")
srcdocpath = joinpath(srcgenpath, "src-plot-attributes.md")
open(path, "w") do io
    !ispath(srcdocpath) && error("source document doesn't exist!")
    println(io, "# Plot attributes")
    src = read(srcdocpath, String)
    println(io, src)
    print(io, "\n")
    println(io, "## List of attributes")
    print_table(io, plot_attr_desc)
end

########################################
#       OldAxis attributes overview       #
########################################

# automatically generate an overview of the axis attributes, using a source md file
@info("Generating axis page")
path = joinpath(genpath, "axis.md")
srcdocpath = joinpath(srcgenpath, "src-axis.md")

open(path, "w") do io
    !ispath(srcdocpath) && error("source document doesn't exist!")
    src = read(srcdocpath, String)
    println(io, src)
    print(io)
    # Axis2D section
    println(io, "## `Axis2D`")
    println(io, "### `Axis2D` attributes groups")
    print_table(io, Axis2D_attr_desc)
    print(io)
    for (k, v) in Axis2D_attr_groups
        println(io, "#### `:$k`\n")
        print_table(io, v)
        println(io)
    end
    # Axis3D section
    println(io, "## `Axis3D`")
    println(io, "### `Axis3D` attributes groups")
    print_table(io, Axis3D_attr_desc)
    print(io)
    for (k, v) in Axis3D_attr_groups
        println(io, "#### `:$k`\n")
        print_table(io, v)
        println(io)
    end

end

########################################
#          Colormap reference          #
########################################

@info "Generating colormap reference"

include(joinpath(srcgenpath, "colormap_generation.jl"))

generate_colorschemes_markdown(
    joinpath(srcgenpath, "src-colors.md"),
    joinpath(genpath, "colors.md")
)

################################################################################
#                 Building HTML documentation with Documenter                  #
################################################################################

@info("Running `makedocs` with Documenter.")

makedocs(
    doctest = false, clean = true,
    repo = "https://github.com/JuliaPlots/AbstractPlotting.jl/blob/{commit}{path}#{line}",
    format = Documenter.HTML(
        prettyurls = false,
        assets = [
            "assets/favicon.ico",
            "assets/syntaxtheme.css"
        ],
    ),
    sitename = "Makie Plotting Ecosystem",
    pages = Any[
        "Home" => "index.md",
        "Basics" => [
            "Basic Tutorial" => "basic-tutorial.md",
            "Layout Tutorial" => "makielayout/tutorial.md",
            "animation.md",
            "interaction.md",
            "Plotting Functions" =>
                joinpath.(
                    "plotting_functions",
                    filter(
                        endswith(".md"),
                        readdir(joinpath(srcpath, "plotting_functions"),
                            sort = true)
                    )
                ),
            "theming.md",
        ],
        "Documentation" => [
            "plot_method_signatures.md",
            "Figure" => "figure.md",
            "Layoutables & Widgets" => [
                "makielayout/layoutables.md",
                "makielayout/axis.md",
                "makielayout/axis3.md",
                "makielayout/box.md",
                "makielayout/button.md",
                "makielayout/colorbar.md",
                "makielayout/gridlayout.md",
                "makielayout/intervalslider.md",
                "makielayout/label.md",
                "makielayout/legend.md",
                "makielayout/lscene.md",
                "makielayout/menu.md",
                "makielayout/slider.md",
                "makielayout/toggle.md",

            ],
            "makielayout/layouting.md",
            "generated/colors.md",
            "generated/plot-attributes.md",
            "recipes.md",
            "backends_and_output.md",
            "scenes.md",
            "lighting.md",
            "cameras.md",
            "faq.md",
            "API Reference AbstractPlotting" => "abstractplotting_api.md",
            "API Reference MakieLayout" => "makielayout/reference.md",
            "generated/axis.md",
        ],
    ],
    strict = true, # experimental kwarg, so that the docs fail if there are any errors encountered
    # this way the docs serve better as another test case, because nobody looks at warnings
)

################################################################################
#                           Deploying documentation                            #
################################################################################

# for github actions, documenter checks that GITHUB_REPOSITORY matches the repo
# keyword, but since we want to push to a different repo, we need to override the
# env variable, which is JuliaPlots/AbstractPlotting.jl by default
ENV["GITHUB_REPOSITORY"] = "JuliaPlots/MakieDocumentation"

deploydocs(
    repo = "github.com/JuliaPlots/MakieDocumentation",
    push_preview = true
)


text:
#==============================================================================#
# Abstract types
#==============================================================================#

"""
    abstract type AbstractSystem

Supertype for a field-theoretic system representing a physical model.
"""
abstract type AbstractSystem end

"""
    abstract type AbstractConstraint

Supertype for all physical constraints applied to a system.
"""
abstract type AbstractConstraint end

"""
    abstract type AbstractSpecies

Supertype for all molecular species in a field-theoretic simulation.
"""
abstract type AbstractSpecies end

"""
    abstract type AbstractPropagator

Supertype for algorithms to evaluate species propagators under a given external field.
"""
abstract type AbstractPropagator end

"""
    abstract type AbstractInteraction

Supertype for all excess interactions in a field-theoretic simulation.
"""
abstract type AbstractInteraction end

"""
    abstract type AbstractFieldUpdater

Supertype for iteration algorithms to update fields in a field-theoretic simulation.
"""
abstract type AbstractFieldUpdater end

#==============================================================================#
# Type definitions
#==============================================================================#

# 3D array/shape alias
const FieldGrid{T} = Array{T, 3}

# Nullable union holder
const Option{T} = Union{Nothing,T}

# 3D vectors and matrices
const SVec3D{T} = SVector{3,T}
const MVec3D{T} = MVector{3,T}
const Vec3D{T} = Union{SVec3D{T}, MVec3D{T}}

const SMat3D{T} = SMatrix{3,3,T,9}
const MMat3D{T} = MMatrix{3,3,T,9}
const Mat3D{T} = Union{SMat3D{T}, MMat3D{T}}


text:
"""
NETALIGNBP
----------
    solve the network alignment problem with Belief Propagation
"""

function netalignbp(S::SparseMatrixCSC{T,Int64},w::Vector{F},
                            a::R,b::K,li::Vector{Int64},
                            lj::Vector{Int64},gam::Float64,
                            dtype::Int64,maxiter::Int64,verbose::Bool) where {T,F,R,K}
  nedges = length(li)
  nsquares = nnz(S)/2
  m = maximum(li)
  n = maximum(lj)
  SU = findall(x->x!=0,triu(S))#only consider each square once.

  #compute a vector that allows us to transpose data between squares.
  #Recall the BP algorithm requires edge(i,j) to send messages to edge(r,s)
  #when (i,j) and (r,s) form a square.  However, edge(i,j) needs information
  #the information that (r,s) sent it from the previous iteraton.
  #If we imagine that each non-zero of S has a value, we want to tranpose
  #these values.

  SI = spzeros(Int,size(S,1),size(S,1))
  SI[SU] = 1:length(SU)
  # this is faster than the old version of ind2sub, now SU is an array of CartesianIndex type
  # SI = sparse(sui,suj,1:length(sui),size(S,1),size(S,2)) #assign indices

  SI = SI + SI' #SI now has symmetric indices
  si = SI.rowval
  sind = SI.nzval

  # faster than:
  # (si, sj) = ind2sub(size(SI),find(SI))
  # sind = nonzeros(SI)

  SP = sparse(si,sind,true,size(S,1),nsquares); # each column in SP has 2 nz

  sij_cartesian = findall(x->x!=0,SP)
  sij = [x.I[1] for x in sij_cartesian]
  sijrs = [x.I[2] for x in sij_cartesian]

  spair = sind
  len_spair = length(spair)
  spair[1:2:end] = 2:2:len_spair
  spair[2:2:end] = 1:2:len_spair

  # sij, sijrs maps between rows and squares
  # spair is now an indexing vector that accomplishes what we need

  # Initialize the messages
  ma = zeros(Float64,nedges)
  mb = zeros(Float64,nedges)
  ms = zeros(Float64,nnz(S))
  sums = zeros(Float64,nedges)
  damping = gam
  curdamp = 1
  iter = 1
  alpha = a
  beta = b

  # Initialize history
  hista = zeros(Float64,maxiter,4) # history of messages from ei->a vertices
  histb = zeros(Float64,maxiter,4) # history of messages from ei->b vertices
  fbest = 0
  fbestiter = 0
  if verbose # print the header
    @printf("%4s   %4s   %7s %7s %7s %7s   %7s %7s %7s %7s\n",
        "best", "iter", "obj_ma", "wght_ma", "card_ma", "over_ma",
            "obj_mb", "wght_mb", "card_mb", "over_mb")
  end
  # setup the matching problem once
  Matching_Setup = bipartite_matching_setup(w,li,lj,m,n)
  tripi = Matching_Setup.tripi
  matm = Matching_Setup.m
  matn = Matching_Setup.n
  rp = Matching_Setup.rp
  ci = Matching_Setup.ci
  mperm = tripi[tripi.>0]

  while iter <= maxiter
    prevma = copy(ma)
    prevmb = copy(mb)
    prevms = copy(ms)
    prevsums = copy(sums)
    curdamp = damping * curdamp

    omaxb = max.((othermaxplus(2,li,lj,mb,m,n)),0)
    omaxa = max.((othermaxplus(1,li,lj,ma,m,n)),0)

    msflip = ms[spair] # swap ij->ijrs to rs->ijrs

    mymsflip = msflip .+ beta
    mymsflip = min.(beta,mymsflip)
    mymsflip = max.(0,mymsflip)

    sums = zeros(Float64,nedges)
    for i = 1:length(sij)
        sums[sij[i]] += mymsflip[i]
    end

    ma = alpha*w .- omaxb .+ sums
    mb = alpha*w .- omaxa .+ sums

    ms = alpha*w[sij]-(omaxb[sij] + omaxa[sij])
    ms = ms .+ othersum(sij,mymsflip,nedges)
    # Original updates
    if dtype == 1
      ma = curdamp*(ma) .+ (1-curdamp)*(prevma)
      mb = curdamp*(mb) .+ (1-curdamp)*(prevmb)
      ms = curdamp*(ms) .+ (1-curdamp)*(prevms)
    elseif dtype == 2
      ma = ma .+ (1-curdamp)*(prevma.+prevmb.-alpha*w+prevsums)
      mb = mb .+ (1-curdamp)*(prevmb.+prevma.-alpha*w+prevsums)
      ms = ms .+ (1-curdamp)*(prevms.+prevms[spair].-beta)
    elseif dtype == 3
      ma = curdamp*ma .+ (1-curdamp)*(prevma.+prevmb.-alpha*w.+prevsums)
      mb = curdamp*mb .+ (1-curdamp)*(prevmb.+prevma.-alpha*w.+prevsums)
      ms = curdamp*ms .+ (1-curdamp)*(prevms.+prevms[spair].-beta)
    end

    # now compute the matchings
    hista[iter,:] = round_messages(ma,S,w,alpha,beta,rp,ci,tripi,matn,matm,mperm,li,lj)
    histb[iter,:] = round_messages(mb,S,w,alpha,beta,rp,ci,tripi,matn,matm,mperm,li,lj)
    if hista[iter,1] > fbest
      fbestiter = iter
      mbest = ma
      fbest = hista[iter,1]
    end
    if histb[iter,1] > fbest
      fbestiter = -iter
      mbest = mb
      fbest = histb[iter,1]
    end
    if verbose
      if fbestiter == iter
        bestchar = "*a"
      elseif fbestiter == -iter
        bestchar = "*b"
      else
        bestchar = ""
      end
      @printf("%4s   %4i    %5.2f %5.2f %4d %7d %14.2f %5.2f %4d %7d\n",
            bestchar, iter, hista[iter,1], hista[iter,2], hista[iter,3], hista[iter,4],
              histb[iter,1], histb[iter,2], histb[iter,3], histb[iter,4]);

    end
    iter += 1
  end
  return (ma,S,w,alpha,beta,rp,ci,tripi,matn,matm,mperm,li,lj)

end