text:
module FedDCD

using LinearAlgebra
using Printf
using SparseArrays
using Random
using StatsBase
using Distributed
using LinearMaps
using IterativeSolvers
using DataFrames
using CSV
using Zygote
using Flux

export FedProxClient, ScaffoldClient, FedDCDClient, AccFedDCDClient, getObjValue
export FedProxServer, ScaffoldServer, FedDCDServer, AccFedDCDServer, getObjValue
export FedDCDClientNN, FedDCDServerNN, AccFedDCDClientNN, AccFedDCDServerNN
export FedProxClientNN, FedProxServerNN, ScaffoldClientNN, ScaffoldServerNN
export softmax, neg_log_loss
export split_data, splitDataByRow, labelTransform, writeToFile, splitDataByClass
export read_libsvm
export fedAvgAndProx, Scaffold, fedDCD, accfedDCD
export obj, accuracy
export getStochasticGrad, getGradient
export ComputeNewtonDirection, SoftmaxNewtonMethod, lineSearch, lineSearch2
export dot_product
export sgd!, newton!, adam!

include("./utils.jl")
include("./LogReg/logReg.jl")
include("./oracle.jl")
include("./Client/client.jl")
include("./Server/server.jl")
include("./training.jl")

end # module


text:
print("matrixops, ")
## this test for typos, mainly, not for correct functioning...

## TODO I am not convinced this does anything related to NamedArrays...

using LinearAlgebra

for t in (Float32, Float64)
    global a = rand(t, 10, 10)
    global b = rand(t, 10, 5)
    global c = rand(t, 10, 10)
    apd = a' * a
    asym = a' + a
    luf = lu(a)
    chf = cholesky(apd, Val(false))
    qrf = qr(a)
    eif = eigen(a)
    eigvals(a)
    hsf = hessenberg(a)
    scf = schur(a)
    scf1 = schur(a, c)
    svdf = svd(a)
    svdvals(a)
    global d = Matrix(Diagonal(a))
    Matrix(Diagonal(d))
    cond(a)
    nullspace(a)
    kron(a, a)
    # linreg(a[:,1], b[:,1])
    lyap(a, c)
    sylvester(a, c, a)
    isposdef(apd)
end

ext: jl


text:
# FeatureServices.jl
#
# Gets features from sources
# whether this be files on disk in certain formats
# or web accessible repositories

using DataFrames
using GZip

include(Pkg.dir("MethylUtils", "src","Features.jl"))
include(Pkg.dir("MethylUtils", "src","Config.jl"))

function load_bis_seq_file(path, fileinfo={}, format="moabs")
    cgs = Array(CytosineCount,0)
    if ( format == "moabs")
        cgs =  load_moabs(path)
    else
        println("$format is currently not supported - (formats supported: moabs")
    end
    return cgs
end

# read_moabs_format_as_sequence_features
# will serialize and read from serialized if it exists
function read_moabs_format_as_sequence_features(path;exclude_unmapped = true)
    path_serialized = join( [path ".jld"] )
    if ( !isfile(path_serialized) )
        cgs = load_moabs_format_as_sequence_features( path, exclude_unmapped = exclude_unmapped)
        iostream = open( path_serialized  ,"w")
        serialize(iostream,cgs)
        close(iostream)
        println("serialised to $path_serialized")
    else
        println("Read from serialized $path_serialized")
        cgs = open( path_serialized ) do file
            deserialize(file)
        end
        println("Done reading from serialized path $path_serialized")
    end
    return cgs
end

# load_moabs_format_as_sequence_features
# reads from file and loads into sequence_feature container
function load_moabs_format_as_sequence_features( path; exclude_unmapped = true)
    uncompressed_stream = open(path)
    #cgs = Array(CytosineCount,0)
    cgs = SequenceFeatures(Dict(),Dict())
    idx = 0
    for line in eachline(uncompressed_stream)
       idx +=1
       # Parse line
       # Get the pos and point, i.e. this will result in two Annotation Points, for start and stop.
       (chr, start, stop, ratio,totalC, methC, strand, next, plus, plustotalC, plusmethC, minus, minustotalC,minusmethC ) = [ "" for i in 1:14 ]
       try
           (chr, start, stop, ratio,totalC, methC, strand, next, plus, plustotalC, plusmethC, minus, minustotalC,minusmethC ) = split(line, "\t")
       catch y
             error( "Error: Offending line ($idx) : $line" )
       end
       methC_int  = parseint(methC)
       totalC_int = parseint(totalC)
       cg = CytosineCount(SequencePos(start, stop) , true, methC_int, totalC_int - methC_int )
       add_sequence_feature(cgs,chr,cg)
       if idx % 1000000 == 0
          println("loaded: ",idx)
       end
   end
   return cgs
end
#####################

# DEPRECATED
function load_moabs(path; exclude_unmapped = true )
    #uncompressed_stream = ismatch(r"gz", path) ? GZip.open(path) : open(path)
    uncompressed_stream = open(path)
    cgs = Array(CytosineCount,0)
    idx = 0
    for line in eachline(uncompressed_stream)
       idx +=1
       # Parse line
       # Get the pos and point, i.e. this will result in two Annotation Points, for start and stop.
       (chr, start, stop, ratio,totalC, methC, strand, next, plus, plustotalC, plusmethC, minus, minustotalC,minusmethC ) = [ "" for i in 1:14 ]
       try
           (chr, start, stop, ratio,totalC, methC, strand, next, plus, plustotalC, plusmethC, minus, minustotalC,minusmethC ) = split(line, "\t")
       catch y
             error( "Error: Offending line ($idx) : $line" )
       end
       methC_int  = parseint(methC)
       totalC_int = parseint(totalC)
       cg = CytosineCount(GenomePos(chr, start, stop) , true, methC_int, totalC_int - methC_int )
       push!(cgs,cg)
       if idx % 1000000 == 0
          println("loaded: ",idx)
       end
   end
   return cgs
end

function ensgene_dataframe(ensgene_path)
    names_str = [ "bin", "name", "chrom", "strand", "txStart", "txEnd", "cdsStart",
              "cdsEnd", "exonCount", "exonStarts", "exonEnds", "score", "name2",
              "cdsStartStat", "cdsEndStat", "exonFrames" ]
    names = map(symbol, names_str)

    #check that ensgene_filename exists
    if !isfile(ensgene_path)
        error("$path does not exist on disk")
    end
    # try catch this and throw out path?
    ensgene_df = readtable(ensgene_path,separator='\t', header=false, names=names)
    return ensgene_df
end

function make_ensgene_gene_feature_table( ensgene_path, organism)
     ensgene_df = ensgene_dataframe( ensgene_path )
     gene_dict = Dict{Any,Any}()
     nrows = size(ensgene_df,1)
     for row = 1:nrows
         gene = ensgene_df[row, :]
         gene_name = gene[:name2][1]
         chr       = gene[:chrom][1]
         start     = gene[:txStart][1]
         stop      = gene[:txEnd][1]
         strand    = gene[:strand][1]
         if haskey(gene_dict, gene_name) == true
             if gene_dict[gene_name][1] > start
                gene_dict[gene_name][1] = start
             end
             if gene_dict[gene_name][2] < stop
                gene_dict[gene_name][2] = stop
             end
         else
             gene_dict[gene_name] = [start, stop, chr, strand]
         end
     end

     println("The length of the gene_dict is: ", length(gene_dict))
     names  = Array(ASCIIString,0)
     starts = Array(Int32,0)
     stops  = Array(Int32,0)
     chrs   = Array(ASCIIString,0)
     strands = Array(Bool,0)
     for key in keys(gene_dict)
        push!(names,key)
        push!(starts,  gene_dict[key][1])
        push!(stops,   gene_dict[key][2])
        push!(chrs, gene_dict[key][3])
        if(gene_dict[key][4] == "+")
            push!(strands, true)
        else
            push!(strands, false)
        end
     end
     println( "lengths of names: ", length(names) )
     D = DataFrame(name=names,chr=chrs,start=starts,stop=stops,strand=strands)
     A = FeatureTable(["organism" => organism],D)
     return A
end

# this creates a file for the start
# and ends of gene regions from the transcripts
# ensgene_dataframe
# DEPRECATED
function ensgene_regions(organism, ensgene_filename="ensGene.txt.gz")
     ensgene_df = ensgene_dataframe(organism,ensgene_filename)
     gene_dict = Dict{Any,Any}()
     nrows = size(ensgene_df,1)
     for row = 1:nrows
         gene = ensgene_df[row, :]
         #println(gene)
         gene_name = gene[:name2][1]
         chr       = gene[:chrom][1]
         start     = gene[:txStart][1]
         stop      = gene[:txEnd][1]
         strand    = gene[:strand][1]
         #println("$gene_name $start $stop")
         if haskey(gene_dict, gene_name) == true
             if gene_dict[gene_name][1] > start
                gene_dict[gene_name][1] = start
             end
             if gene_dict[gene_name][2] < stop
                gene_dict[gene_name][2] = stop
             end
         else
             gene_dict[gene_name] = [start, stop, chr, strand]
         end
     end

     CURRENT_ANNOTATION_CHR_INDEX = current_annotation_chr_index() # from config

     println("The length of the gene_dict is: ", length(gene_dict))
     names  = Array(ASCIIString,0)
     starts = Array(Int32,0)
     stops  = Array(Int32,0)
     chrs   = Array(Int16,0)
     strands = Array(Bool,0)
     for key in keys(gene_dict)
        push!(names,key)
        push!(starts,  gene_dict[key][1])
        push!(stops,   gene_dict[key][2])
        push!(chrs,    CURRENT_ANNOTATION_CHR_INDEX[organism][ gene_dict[key][3] ])
        if(gene_dict[key][4] == "+")
            push!(strands, true)
        else
            push!(strands, false)
        end
     end
     println( "lengths of names: ", length(names) )
     D = DataFrame(name=names,chr=chrs,start=starts,stop=stops,strand=strands)
     A = FeatureTable(["organism" => organism],D)
     return A
end

# DEPRECATED
function ensgene_dataframe(organism, ensgene_filename)
    path = joinpath(LOCAL_STORE_PATH(), organism,ensgene_filename)
    names_str = [ "bin", "name", "chrom", "strand", "txStart", "txEnd", "cdsStart",
              "cdsEnd", "exonCount", "exonStarts", "exonEnds", "score", "name2",
              "cdsStartStat", "cdsEndStat", "exonFrames" ]
    names = map(symbol, names_str)

    #check that ensgene_filename exists
    if !isfile(path)
        error("$path does not exist on disk")
    end
    # try catch this and throw out path?
    ensgene_df = readtable(path,separator='\t', header=false, names=names)
    return ensgene_df
end

#### TO IMPLEMENT ####
# function load_moabs(path,include_intervals::IntervalTree)
#    error(" load_moabs(path,include_intervals) Not implemented")
# end
######################


text:
module GoTMSupport

export completion_choices_plist

include("funcparser.jl")
include("commandcompletion.jl")

end # module


text:
@testset "Tables Interface" begin
    n = 1000
    x, y, z = 1:n, rand(Bool, n), randn(n)

    t = table((x=x, y=y, z=z), pkey=[:x, :y])

    @test Tables.istable(t)
    # @test t == table(Tables.rowtable((x=x,y=y,z=z)))
    @test Tables.istable(columns(t))
    @test Tables.istable(Columns(columns(t)))
end


text:
@testset "JuMP direct model - Linear Constraints - Affine parameters" begin
    optimizer = POI.ParametricOptimizer(GLPK.Optimizer())

    model = direct_model(optimizer)

    @variable(model, x[i = 1:2] >= 0)

    @variable(model, y in POI.Parameter(0))
    @variable(model, w in POI.Parameter(0))
    @variable(model, z in POI.Parameter(0))

    @constraint(model, 2 * x[1] + x[2] + y <= 4)
    @constraint(model, 1 * x[1] + 2 * x[2] + z <= 4)

    @objective(model, Max, 4 * x[1] + 3 * x[2] + w)

    optimize!(model)

    @test isapprox.(value(x[1]), 4.0 / 3.0, atol = ATOL)
    @test isapprox.(value(x[2]), 4.0 / 3.0, atol = ATOL)
    @test isapprox.(value(y), 0, atol = ATOL)

    # ===== Set parameter value =====
    MOI.set(model, POI.ParameterValue(), y, 2.0)
    optimize!(model)

    @test isapprox.(value(x[1]), 0.0, atol = ATOL)
    @test isapprox.(value(x[2]), 2.0, atol = ATOL)
    @test isapprox.(value(y), 2.0, atol = ATOL)
end

@testset "JuMP direct model - Linear Constraints - Parameter x variable" begin
    optimizer = POI.ParametricOptimizer(GLPK.Optimizer())

    model = direct_model(optimizer)

    @variable(model, x[i = 1:2] >= 0)

    @variable(model, y in POI.Parameter(0))
    @variable(model, w in POI.Parameter(0))
    @variable(model, z in POI.Parameter(0))

    @constraint(model, 2 * x[1] + x[2] + y <= 4)
    @constraint(model, (1 + y) * x[1] + 2 * x[2] + z <= 4)

    @objective(model, Max, 4 * x[1] + 3 * x[2] + w)

    optimize!(model)

    @test isapprox.(value(x[1]), 4.0 / 3.0, atol = ATOL)
    @test isapprox.(value(x[2]), 4.0 / 3.0, atol = ATOL)
    @test isapprox.(value(y), 0, atol = ATOL)

    # ===== Set parameter value =====
    MOI.set(model, POI.ParameterValue(), y, 2.0)
    optimize!(model)

    @test isapprox.(value(x[1]), 0.0, atol = ATOL)
    @test isapprox.(value(x[2]), 2.0, atol = ATOL)
    @test isapprox.(value(y), 2.0, atol = ATOL)
end

@testset "JuMP - Linear Constraints - Affine parameters" begin
    model = Model(() -> POI.ParametricOptimizer(GLPK.Optimizer()))

    @variable(model, x[i = 1:2] >= 0)

    @variable(model, y in POI.Parameter(0))
    @variable(model, w in POI.Parameter(0))
    @variable(model, z in POI.Parameter(0))

    @constraint(model, 2 * x[1] + x[2] + y <= 4)
    @constraint(model, 1 * x[1] + 2 * x[2] + z <= 4)

    @objective(model, Max, 4 * x[1] + 3 * x[2] + w)

    optimize!(model)

    @test isapprox.(value(x[1]), 4.0 / 3.0, atol = ATOL)
    @test isapprox.(value(x[2]), 4.0 / 3.0, atol = ATOL)
    @test isapprox.(value(y), 0, atol = ATOL)

    # ===== Set parameter value =====
    MOI.set(model, POI.ParameterValue(), y, 2.0)
    optimize!(model)

    @test isapprox.(value(x[1]), 0.0, atol = ATOL)
    @test isapprox.(value(x[2]), 2.0, atol = ATOL)
    @test isapprox.(value(y), 2.0, atol = ATOL)
end

@testset "JuMP - Linear Constraints - Parameter x variable" begin
    model = Model(() -> POI.ParametricOptimizer(GLPK.Optimizer()))

    @variable(model, x[i = 1:2] >= 0)

    @variable(model, y in POI.Parameter(0))
    @variable(model, w in POI.Parameter(0))
    @variable(model, z in POI.Parameter(0))

    @constraint(model, 2 * x[1] + x[2] + y <= 4)
    @constraint(model, (1 + y) * x[1] + 2 * x[2] + z <= 4)

    @objective(model, Max, 4 * x[1] + 3 * x[2] + w)

    optimize!(model)

    @test isapprox.(value(x[1]), 4.0 / 3.0, atol = ATOL)
    @test isapprox.(value(x[2]), 4.0 / 3.0, atol = ATOL)
    @test isapprox.(value(y), 0, atol = ATOL)

    # ===== Set parameter value =====
    MOI.set(model, POI.ParameterValue(), y, 2.0)
    optimize!(model)

    @test isapprox.(value(x[1]), 0.0, atol = ATOL)
    @test isapprox.(value(x[2]), 2.0, atol = ATOL)
    @test isapprox.(value(y), 2.0, atol = ATOL)
end


text:
import Base
export TypeInfo, InstTo, NoProp
export TypeInfoKind, InstToKind, NoPropKind

@data TypeInfo begin
    InstTo(HMT)
    NoProp
end

@data TypeInfoKind begin
    InstToKind
    NoPropKind
end

(i::TypeInfoKind)(x :: HMT) =
    @match i begin
        InstToKind => InstTo(x)
        NoPropKind => NoProp
    end


text:
using GeoMakie, CairoMakie, AbstractPlotting, ImageMagick, NASAEarthObservations, Glob
using Test

source = LonLat()
dest = WinkelTripel()

imgdir = observations("rgb/MYDAL2_M_AER_RA")

imgpaths = sort(Glob.glob(joinpath(relpath(imgdir), "*.PNG"))) # change this to your requirements

img = ImageMagick.load(imgpaths[1])

re = GeoMakie.date_regex("MYDAL2_M_AER_RA", "PNG")
titletext = Node(join(match(re, basename(imgpaths[1])).captures, '-'))

lons = LinRange(-179.5, 179.5, size(img)[2])
lats = LinRange(89.5, -89.5, size(img)[1])

xs = [lon for lat in lats, lon in lons]
ys = [lat for lat in lats, lon in lons]

points, faces = GeoMakie.triangulated_grid(lons, lats)
tpoints = Proj4.transform.(source, dest, points)

scene = mesh(tpoints, faces; color = img[:], shading = false, show_axis = false, scale_plot = false)

geoaxis!(scene, -180, 180, -90, 90; crs = (src = source, dest = dest,));

titletext = Node("07/2016")

fullsc = title(scene, titletext; fontsize = 40);

record(fullsc, "particulates.mp4", filter!(x -> uppercase(splitext(x)[2]) == ".TIFF", sort(readdir(imgdir))); framerate = 10) do img

    year, month = match(re, img).captures

    scene.plots[1].color = ImageMagick.load(joinpath(imgdir,img))
    titletext[] = "$month/$year"
end


text:
isdefined(Base, :__precompile__) && __precompile__(true)

module Glove

import Base: insert!, haskey, getindex, length
import DataStructures

export Model, LookupTable, Token,
       CooccurenceDict, CooccurenceVector,
       weightedsums, adagrad!

# Not sure if I'll need this yet.
using Compat
using LinearAlgebra

include("lookuptable.jl")
include("cooccurence.jl")
include("model.jl")

end


text:
export metadata, info, sampleinfo, genodata, locusinfo, samplenames
export locationdata, loci, samples, populations
export genotype, genotypes

"""
    info(::PopData)
Display the metadata (PopDataInfo) of a PopData object.
"""
metadata(data::PopData) = data.metadata
info(data::PopData) = data.metadata

"""
    sampleinfo(::PopData)
Show the sample information found within the metadata of a `PopData` object. Returns a view of a dataframe
"""
sampleinfo(data::PopData) = data.sampleinfo

"""
    genodata(::PopData)
Method to show the genotype information of a `PopData` object. Returns a view of a dataframe.
"""
genodata(data::PopData) = data.genodata

"""
    locusinfo(::PopData)
Show the locus information found within the metadata of a `PopData` object. Returns a view of a dataframe
"""
locusinfo(data::PopData) = data.locusinfo

"""
    locationdata(data::PopData)
View the longitude and latitude data in a `PopData` object. Returns a table
derived from the PopData. Changes made to this table will not alter the source
`PopData` object.

Use `locations!` to add spatial data to a `PopData` object.
"""
function locationdata(data::PopData)
    if :longitude ∉ propertynames(data.sampleinfo) && :latitude ∉ propertynames(data.sampleinfo)
        throw(ArgumentError(":longitude and :latitude columns not present in metadata."))
    else
        @view data.sampleinfo[!, [:longitude, :latitude]]
    end
end

"""
    loci(data::PopData)
Returns an array of strings of the loci names in a `PopData` object.
"""
function loci(data::PopData)
    data.locusinfo.locus
end

"""
    samplenames(data::PopData)
View individual/sample names in a `PopData`
"""
function samplenames(data::PopData)
    data.sampleinfo.name
end

"""
    genotypes(data::PopData, samplelocus::String)
Return a vector of all the genotypes of a sample (or locus) in a `PopData` object.
```
cats = @nancycats
genotypes(cats, "N115")
genotypes(cats, "fca8")

```
"""
function genotypes(data::PopData, samplelocus::AbstractString)
    if samplelocus ∈ samplenames(data)
        @view data.genodata[data.genodata.name .== samplelocus, :genotype]
    elseif samplelocus ∈ loci(data)
        @view data.genodata[data.genodata.locus .== samplelocus, :genotype]
    else
        throw(ArgumentError("$samplelocus is not found in either samples or loci."))
    end
end

"""
    genotypes(data::PopData, samplelocus::Pair{String, String}) ::DataFrame
    genotypes(data::PopData, samplelocus::Pair{Vector{String}, String}) ::DataFrame
    genotypes(data::PopData, samplelocus::Pair{String, Vector{String}}) ::DataFrame

Return a genotype or dataframe of genotypes for one or more samples/loci
in a `PopData` object. Uses the `Pair` notation of `samples => loci`.
### Examples
```
cats = @nancycats;
genotypes(cats, "N115" => "fca8")
genotypes(cats, ["N115", "N7"] => "fca8")
genotypes(cats, "N115" => ["fca8", "fca37"])
genotypes(cats, ["N1", "N2"] => ["fca8", "fca37"])
```
"""
function genotypes(data::PopData, samplocus::Pair{T,U})  where T<:AbstractString where U<:AbstractString
    sample, locus = samplocus
    @views data.genodata[(data.genodata.name .== sample) .& (data.genodata.locus .== locus), :genotype][1]
end

function genotypes(data::PopData, samplelocus::Pair{U, Vector{T}}) where T<:AbstractString where U<:AbstractString
    sample,locus = samplelocus
    @view data.genodata[(data.genodata.name .== sample) .& (data.genodata.locus .∈ Ref(locus)), :]
end

function genotypes(data::PopData, samplelocus::Pair{Vector{U}, Vector{T}}) where T<:AbstractString where U<:AbstractString
    sample,locus = samplelocus
    @view data.genodata[(data.genodata.name .∈ Ref(sample)) .& (data.genodata.locus .∈ Ref(locus)), :]
end

function genotypes(data::PopData, samplelocus::Pair{Vector{U}, T}) where T<:AbstractString where U<:AbstractString
    sample,locus = samplelocus
    @view data.genodata[(data.genodata.name .∈ Ref(sample)) .& (data.genodata.locus .== locus), :]
end

"""
    populations(data::PopData; counts::Bool = false)
View unique population ID's and/or their counts in `PopData`.

- `counts` returns a dataframe of samples per `population` instead (default = `false`)
"""
@inline function populations(data::PopData; counts::Bool = false)
    if all(ismissing.(data.sampleinfo.population)) == true
        println("no population data present in PopData")
        return
    end
    uniq_pops = unique(data.sampleinfo.population)
    if counts == false
        return uniq_pops
    else
        pops = countmap(data.sampleinfo.population)
        return DataFrame(:population => uniq_pops, :count => [pops[i] for i in uniq_pops])
    end
end


text:
# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# EAGO
# A development environment for robust and global optimization
# See https://github.com/PSORLab/EAGO.jl
#############################################################################
# src/eago_optimizer/functions/nonlinear/quadratic.jl
# Defines buffered structures to store quadratic functions:
# BufferedQuadraticIneq, BufferedQuadraticEq, as well as the
# lower_interval_bound, interval_bound, and eliminate_fixed_variables!
# functions associated with each structure.
#############################################################################

###
### Structure definitions
###

"""
$(TYPEDEF)

Stores a general quadratic inequality constraint with a buffer.
"""
mutable struct BufferedQuadraticIneq <: AbstractEAGOConstraint
    func::SQF
    buffer::Dict{Int, Float64}
    saf::SAF
    len::Int
end

"""
$(TYPEDEF)

Stores a general quadratic equality constraint with a buffer.
"""
mutable struct BufferedQuadraticEq <: AbstractEAGOConstraint
    func::SQF
    minus_func::SQF
    buffer::Dict{Int, Float64}
    saf::SAF
    len::Int
end

#=
mutable struct BufferedConvexQuadratic <: AbstractEAGOConstraint
    func::SQF
    buffer::Dict{Int, Float64}
    saf::SAF
    len::Int
end
=#

###
### Constructor definitions
###

function create_buffer_dict(func::SQF)

    buffer = Dict{Int, Float64}()

    for term in func.quadratic_terms
        buffer[term.variable_index_1.value] = 0.0
        buffer[term.variable_index_2.value] = 0.0
    end

    for term in func.affine_terms
        buffer[term.variable_index.value] = 0.0
    end

    return buffer
end

BufferedQuadraticIneq() = BufferedQuadraticIneq(SQF(SQT[], SAT[], 0.0), Dict{Int, Float64}(), SAF(SAT[], 0.0), 0)

function BufferedQuadraticIneq(func::SQF, set::LT)

    buffer = create_buffer_dict(func)
    saf = SAF([SAT(0.0, VI(k)) for k in keys(buffer)], 0.0)
    len = length(buffer)
    cfunc = copy(func)
    cfunc.constant -= set.upper

    return BufferedQuadraticIneq(cfunc, buffer, saf, len)
end

function BufferedQuadraticIneq(func::SQF, set::GT)

    buffer = create_buffer_dict(func)
    saf = SAF([SAT(0.0, VI(k)) for k in keys(buffer)], 0.0)
    len = length(buffer)
    cfunc = MOIU.operate(-, Float64, func)
    cfunc.constant += set.lower

    return BufferedQuadraticIneq(cfunc, buffer, saf, len)
end

BufferedQuadraticEq() = BufferedQuadraticEq(SQF(SQT[], SAT[], 0.0), SQF(SQT[], SAT[], 0.0), Dict{Int, Float64}(), SAF(SAT[], 0.0), 0)

function BufferedQuadraticEq(func::SQF, set::ET)

    buffer = create_buffer_dict(func)
    saf = SAF([SAT(0.0, VI(k)) for k in keys(buffer)], 0.0)
    len = length(buffer)
    cfunc1 = copy(func)
    cfunc1.constant -= set.value
    cfunc2 = MOIU.operate(-, Float64, func)
    cfunc2.constant += set.value

    return BufferedQuadraticEq(cfunc1, cfunc2, buffer, saf, len)
end

#=
function BufferedConvexQuadratic(f::BufferedQuadraticIneq)
    BufferedConvexQuadratic(copy(f.func), copy(f.buffer), copy(f.saf), f.len)
end
=#

###
### Parsing definitions
###

function eliminate_fixed_variables!(f::T, v::Vector{VariableInfo}) where T <: Union{BufferedQuadraticIneq,
                                                                                BufferedQuadraticIneq}
    deleted_count = 0
    index = 1
    while i + deleted_count <= f.len
        term = @inbounds f.sqf.terms[i]
        variable_info_1 = @inbounds v[term.variable_index_1.value]
        variable_info_2 = @inbounds v[term.variable_index_2.value]
        if variable_info_1.is_fixed && variable_index_2.is_fixed
            f.sqf.constant += coeff*variable_info_1.lower_bound*variable_index_2.lower_bound
            deleteat!(f.sqf.terms, i)
            deleted_count += 1
        else
            i += 1
        end
    end
    f.len -= deleted_count

    return nothing
end


text:
# Bicycle domain ported from rlpy https://github.com/rlpy/rlpy/blob/master/rlpy/Domains/Bicycle.py
# original paper Learning to Drive a Bicycle using Reinforcement Learning and Shaping, Jette Randlov, Preben Alstrom, 1998. https://www.semanticscholar.org/paper/Learning-to-Drive-a-Bicycle-Using-Reinforcement-and-Randl%C3%B8v-Alstr%C3%B8m/10bad197f1c1115005a56973b8326e5f7fc1031c

# Bicycle domain might not be correct.
struct BicycleParams{T} <:Any where {T<:Real}
	g::T    	# gravity 9.82
	v::T		# velocity of the bicycle (10 km/h paper value) default is 10. / 3.6 (10kph appox 6.2mph and 10kph is 2.77778 meters per second or 10/3)
	d_CM::T 	# vertical distance between the CM for the bicycle and for the cyclist 0.3m
	c::T    	# horizontal distance between the point where the front wheel touches the ground and the CM 0.66 m
	h::T		# height of the CM over the ground  0.94 cm
	M_c::T  	# mass bicycle  9kg ≈ 19.8lbs 15kg ≈ 33lbs (default is 15kg)
	M_d::T		# mass tire 1.0kg ≈ 2.2lbs 1.5kg ≈ 3.3lbs 3.0kg ≈ 4.4lbs (default is 1.7)
	M_p::T		# mass cyclist M_p = 45.4kg ≈ 100lbs 60.0kg ≈ 132.0lbs 90.7kg ≈ 200lbs (default is 60kg)
	M::T   		# mass of cyclist and bike M_p + M_c
	r::T	   	# radius of tire 0.34m
	dsigma::T  	# angular velcity of a tire (v/r)
	I::T    	# moment of inertia for bicycle and cyclist 13/3 * M_c * h^2 + M_p * (h + d_CM)^2
	I_dc::T 	# moment of inertia for tire M_d * r^2
	I_dv::T 	# moment of inertia for tire 3/2 * M_d * r^2
	I_dl::T 	# moment of inertia for tire 1/2 * M_d * r^2
	l::T    	# disance between the front and tire and the back tire at the point where they both touch the ground 1.11m

    function BicycleParams(T::Type=Float64)
		g = T(9.82)
		v = T(10. / 3.6)
		d_CM = T(0.3)
		c = T(0.66)
		h = T(0.94)
		M_c = T(15.)
		M_d = T(1.7)
		M_p = T(60.)
		M = M_c + M_p
		r = T(0.34)
		dsigma = T(v / r)
		I = T(13. / 3. * M_c * h^2 + M_p * (h+d_CM)^2)
		I_dc = T(M_d * r^2)
		I_dv = T(3. / 2.)
		I_dl = T(M_d / 2. * r^2)
		l = T(1.11)
		new{T}(g, v, d_CM, c, h, M_c, M_d, M_p, M, r, dsigma, I, I_dc, I_dv, I_dl, l)
	end

	function BicycleParams(::Type{T}, g, M_c, M_d, M_p, v) where {T}
		d_CM = T(0.3)
		c = T(0.66)
		h = T(0.94)
		M = M_c + M_p
		r = T(0.34)
		dsigma = T(v / r)
		I = T(13. / 3. * M_c * h^2 + M_p * (h+d_CM)^2)
		I_dc = T(M_d * r^2)
		I_dv = T(3. / 2. * M_d * r^2)
		I_dl = T(M_d / 2. * r^2)
		l = T(1.11)
		new{T}(g, v, d_CM, c, h, M_c, M_d, M_p, M, r, dsigma, I, I_dc, I_dv, I_dl, l)
	end
end

function random_bike_params()
	g = rand(Uniform(8., 10.))  		# gravity
	v = rand(Uniform(8., 15.)) / 3.6	# velocity of the bicycle mps
	M_c = rand(Uniform(9., 15.))  		# mass bicycle  9kg ≈ 19.8lbs 15kg ≈ 33lbs (default is 15kg)
	M_d = rand(Uniform(1., 3.))	    	# mass tire 1.0kg ≈ 2.2lbs 1.5kg ≈ 3.3lbs 3.0kg ≈ 4.4lbs (default is 1.7)
	M_p = rand(Uniform(45.4, 90.7))		# mass cyclist M_p = 45.4kg ≈ 100lbs 60.0kg ≈ 132.0lbs 90.7kg ≈ 200lbs (default is 60kg)

	params = BicycleParams(Float64, g, M_c, M_d, M_p, v)
	return params
end

function bicycle_balance(;randomize=false, maxT=30.0)
	if randomize
		params = random_bike_params()
	else
		params = BicycleParams()
	end
	m = create_bicycle_balance(params, maxT=maxT, dt=0.01)
end

function bicycle_goal(;randomize=false, maxT=30.0)
	if randomize
		params = random_bike_params()
	else
		params = BicycleParams()
	end
	m = create_bicycle_goal(params, maxT=maxT, dt=0.01)
end

function bicycle_spaces(maxT)
	X1 = zeros((5,2))
    X1[1,:] .= [-π * 12. / 180., π * 12 / 180]
	X1[2,:] .= [-π, π]
	X1[3,:] .= [-π * 80. / 180., π * 80. / 180.]
	X1[4,:] .= [-π, π]
	X1[5,:] .= [-π, π]
	X2 = zeros(2,2)
	X2[1,:] .= [-Inf, Inf]
	X2[2,:] .= [-Inf, Inf]

	S = ([0.0, maxT], X1, X2)
	A = 1:9
	return S, A
end

function create_bicycle_balance(params; maxT=60.0, dt=0.01)
	S, A = bicycle_spaces(maxT)
	X = S[2]

	p = (s,a)->bicycle_balance_step!(s, a, params, dt, maxT)
	d0 = bicycle_balance_initial

	meta = Dict{Symbol,Any}()
    # meta[:minreward] = -5.0
    # meta[:maxreward] = 5
    # meta[:minreturn] = -5 * ceil(maxT / (dt * 20))
    # meta[:maxreturn] =  5 * ceil(maxT / (dt * 20))
    meta[:stochastic] = true
    meta[:minhorizon] = 1  # not sure the the true minimum is. This seems like a good lower bound
    meta[:maxhorizon] = ceil(Int, maxT / dt)
    meta[:discounted] = false
	meta[:episodes] = 2000
	bplot = BicyclePlotData()
    render = (state,clearplot=false)->bicycleplot(bplot,state,clearplot)
    m = SequentialProblem(S,X,A,p,d0,meta,render)
	return m
end

function create_bicycle_goal(params; maxT=600.0, dt=0.01)
	S, A = bicycle_spaces(maxT)
	X = S[1]

	p = (s,a)->bicycle_goal_step!(s, a, params, dt, maxT)
	d0 = bicycle_goal_initial

	meta = Dict{Symbol,Any}()
    # meta[:minreward] = -5.0
    # meta[:maxreward] = 5
    # meta[:minreturn] = -5 * ceil(maxT / (dt * 20))
    # meta[:maxreturn] =  5 * ceil(maxT / (dt * 20))
    meta[:stochastic] = true
    meta[:minhorizon] = 1  # not sure the the true minimum is. This seems like a good lower bound
    meta[:maxhorizon] = ceil(Int, maxT / dt)
    meta[:discounted] = false
	meta[:episodes] = 3000
	bplot = BicyclePlotData()
    render = (state,clearplot=false)->bicycleplot(bplot,state,clearplot)
    m = SequentialProblem(S,X,A,p,d0,meta,render)
	return m
end

function bicycle_balance_initial()
	t = 0.0
	x = zeros(5)
	pos = zeros(2)

	return (t,x,pos), x
end

function bicycle_goal_initial()
	t = 0.0
	x = zeros(5)
	pos = zeros(2)
	pos[1] = -500.0  # 20 * randn()
	pos[2] = 0.0 # 20 * randn()
	return (t,x,pos), x
end

function bicycle_sim!(state, pos, constants::BicycleParams{T}, u::T, d::T, dt::T) where {T<:Real}
	omega, omegaDot, theta, thetaDot, psi = state
	g = constants.g
	v = constants.v
	d_CM = constants.d_CM
	c = constants.c
	h = constants.h
	M_c = constants.M_c
	M_d = constants.M_d
	M_p = constants.M_p
	M = constants.M
	r = constants.r
	dsigma = constants.dsigma
	I = constants.I
	I_dc = constants.I_dc
	I_dv = constants.I_dv
	I_dl = constants.I_dl
	l = constants.l

	w = 0.02 * (2.0 * rand() - 1.0)

	phi = omega + atan(d+w) / h
	invr_f = abs(sin(theta)) / l
	invr_b = abs(tan(theta)) / l
	if theta != 0.
		invr_CM = ((l-c)^2 + invr_b^(-2))^(-0.5)
	else
		invr_CM = 0.
	end

	nomega = omega + dt * omegaDot
	nomegaDot = omegaDot + dt * (M * h * g * sin(phi) - cos(phi) * (I_dc * dsigma * thetaDot + sign(theta) * v^2 * (M_d * r * (invr_f + invr_b) + M * h * invr_CM))) / I
	out = theta + dt * thetaDot

	if abs(out) < (π * 80. / 180.)
		ntheta = out
		nthetaDot = thetaDot + dt * (u - I_dv * dsigma * omegaDot) / I_dl
	else
		ntheta = sign(out) * (π * 80. / 180.)
		nthetaDot = 0.
	end
	# ntheta = mod(ntheta + π, 2. * π) - π

	npsi = psi + dt * sign(theta) * v * invr_b

	npsi = npsi % (2 * π)
	if npsi > π
		npsi -= 2 * π
	end

	# dposxf = v * dt * -sin(psi + theta + sign(psi + theta) * asin(v*dt*0.5*invr_f))
	# dposyf = v * dt * cos(psi + theta + sign(psi + theta) * asin(v*dt*0.5*invr_f))
	dposxb = v * dt * -sin(psi + sign(psi) * asin(v*dt*0.5*invr_b))
	dposyb = v * dt * cos(psi + sign(psi) * asin(v*dt*0.5*invr_b))

	state[1] = nomega
	state[2] = omegaDot = nomegaDot
	state[3] = ntheta
	state[4] = nthetaDot
	state[5] = npsi
	pos[1] += dposxb
	pos[2] += dposyb

	return npsi - psi

end

function compute_actions(action::Int)
	if action <= 0 || action > 10
        error("Action needs to be an integer in [1, 9]")
    end
    u = 0.
	d = 0.
	uaction = floor(Int, (action-1) / 3) + 1
	daction = action % 3

	if uaction == 1
        u = -2.
    elseif uaction == 2
    	u =  0.
	else
		u = 2.
    end

	if daction == 1
		d = -0.02
	elseif daction == 2
		d =  0.
	else
		d = 0.02
	end
	return u,d
end

function bicycle_balance_step!(state, action::Int, params, dt, maxT)
    u,d = compute_actions(action)
	t,x,pos = state
	psi_diff = bicycle_sim!(x, pos, params, u, d, dt)
	t += dt

	done = bike_fell(x) || t ≥ maxT

	reward = 1

	γ = 1.0
	if done
		γ = 0.0
	end

    return (t, x, pos), x, reward, γ
end

function bicycle_goal_step!(state, action::Int, params, dt, maxT)
    u,d = compute_actions(action)
	t,x,pos = state
	psi_diff = bicycle_sim!(x, pos, params, u, d, dt)
	t += dt

	if bike_fell(x)
		reward = -1.0
		γ = 0.0
	elseif bike_at_goal(pos)
		reward = 0.01
		γ = 0.0
	elseif t ≥ maxT
		reward = (4.0 - atan(pos...)^2) * 0.00004
		γ = 0.0
	else
		reward = (4.0 - atan(pos...)^2) * 0.00004
		γ = 1.0
	end

    return (t, x, pos), x, reward, γ
end

function bike_fell(x)
	omega = x[1]
	fallcond = abs(omega) > (π * 12. / 180.) # 12 degrees in rads
	return fallcond
end

function bike_at_goal(pos)
	d = √sum(pos.^2)
	return d ≤ 10  # within 10 meters
end

mutable struct BicyclePlotData <: Any
	ts::Vector{Float64}
	ωs::Vector{Float64}
	xs::Vector{Float64}
	ys::Vector{Float64}

	function BicyclePlotData()
		new(Vector{Float64}(),Vector{Float64}(),Vector{Float64}(),Vector{Float64}())
	end
end

@userplot BicyclePlot
@recipe function f(ap::BicyclePlot)
	data, state, clearplot = ap.args
	t, x, pos = state
	if clearplot
		empty!(data.ts)
		empty!(data.ωs)
		empty!(data.xs)
		empty!(data.ys)
	end

	push!(data.ts, t)
	push!(data.ωs, x[1])
	push!(data.xs, pos[1])
	push!(data.ys, pos[2])

	legend := false
	# xlims := (0., 1.)
	# ylims := (0., 1.)
	grid := true
	ticks := true
	layout := (2,1)
	foreground_color := :white
	# aspect_ratio := 1.

	@series begin
		seriestype := :path
		seriescolor := :blue
		subplot := 1
		# xlims := (0., 30.)
		fall = π * 12.0 / 180.0
		ylims := (-fall, fall)
		data.ts, data.ωs
	end

	@series begin
		seriestype := :path
		seriescolor := :black
		subplot := 2
		# xlims := (-5, 5)
		# ylims := (-3, 10)
		data.xs, data.ys
	end

	# goal
	@series begin
		seriestype := :shape
		linecolor := nothing
		seriescolor := :green
		aspect_ratio := 1.
		fillalpha := 0.4
		subplot := 2
		circle_shape(0.0, 0.0, 10)
	end

	# current pos
	@series begin
		seriestype := :shape
		linecolor := nothing
		seriescolor := :red
		aspect_ratio := 1.

		circle_shape(pos[1], pos[2], 1.0)
	end

end


text:
module Estim

export fixPoint!, setupMu!, stdError, share!, jacobDelta!, linearIV!, linearIV
export gmmWeight, optimIV!, prediction!, updateGMMweight!
using TypeDef, Optim, Helper
using Printf, LinearAlgebra
using LoopVectorization

function fixPoint!(eδ::Vector{Vector{Float64}}, ∂δ∂θ::Vector{Matrix{Float64}},
  eμ::Vector{Matrix{Float64}}, dat::datType, dim::dimType, mat::matType,
  indx::indxType, pred::predType, diags::diagsType, θ₂::Vector{Float64})

  for t=1:dim.market
    if any(isnan,eδ[t])
      return
    end
    eδt  = eδ[t]
    eμt  = eμ[t]
    numProdt = dim.prodt[t]
    ms1 = pred.ms[t]  # ms1 = zeros(Float64,numProdt,)
    num = pred.num[t] # num = Array{Float64,1}(numProdt)
    eμ₀ = pred.eμ₀[t]
    ms0 = mat.ms[t] #view(mat.ms,mask)

    iter::Int64 = 0; dist::Float64 = 1.0
    while (dist>diags.tolBLP) & (iter<1e5)
      share!(ms1,num,eμ₀,eδt,eμt,numProdt,dim)
      dist = 0.0
      for j=1:numProdt
        Δ = ms0[j]/ms1[j]
        dist = max(dist, abs(log(Δ)))
        eδt[j] *= Δ
      end
      iter += 1
    end
    # if iter>=1e5
    #   eδt[:] = NaN
    #   println("Fixed point iteration exceeding the limit.")
    #   return
    # end
    jacobDelta!(∂δ∂θ[t],pred.∂S∂θ[t],mat.dμdL[t],eδt,eμt,num,eμ₀,
                pred.s[t],ms1,pred.mss[t],numProdt,dim,indx,θ₂)
    diags.iter += iter
    #@printf "Converged in %d iterations at markt %d\n" iter t
  end
  diags.iterBLP += 1
end

function jacobDelta!(
  ∂δ∂θ::Matrix{Float64},
  ∂S∂θ::Matrix{Float64},
  dμdL::Array{Float64,3},
  eδt::Array{Float64,1},
  eμt::Matrix{Float64},
  num::Vector{Float64},
  eμ₀::Vector{Float64},
  s::Vector{Float64},
  ms::Vector{Float64},
  mss::Matrix{Float64},
  numProdt::Int64,
  dim::dimType,
  indx::indxType,
  θ₂::Vector{Float64})

  fill!(mss,0.0)
  fill!(∂S∂θ,0.0)
  for i=1:dim.person
    den = eμ₀[i]
    for j=1:numProdt
      num[j] = eδt[j]*eμt[j,i]
      den += num[j]
    end
    for j=1:numProdt
      s[j] = num[j]/den
    end
    for j=1:numProdt
      for k=1:numProdt
        mss[k,j] += s[j]*s[k]
      end
    end
    for l=1:dim.θ₂
      dμdLᵢ = view(dμdL,:,l,i)
      ΣdμdLS = 0.0
      for k=1:numProdt
        ΣdμdLS += dμdLᵢ[k]*s[k]
      end
      for j=1:numProdt
        ∂S∂θ[j,l] += s[j]*(dμdLᵢ[j]-ΣdμdLS)
      end
    end
  end
  ∂S∂θ /= dim.person
  ∂S∂δ = diagm(0=>ms) - mss/dim.person
  # if θ₂[1]>0
  #   ∂S∂θ[:] /= 2*sqrt(θ₂)
  # end
  for l in indx.expθ
    ∂S∂θ[:,l] *= exp(θ₂[l])
  end
  ∂δ∂θ[:] = -inv(∂S∂δ)*∂S∂θ

end

function share!(ms::Vector{Float64}, num::Vector{Float64}, eμ₀::Vector{Float64},
  eδt::Array{Float64,1}, eμt::Matrix{Float64},
  numProdt::Int64, dim::dimType)

  fill!(ms,0.0)
  @inbounds for i=1:dim.person
    den = eμ₀[i]
    @avx for j=1:numProdt
      num[j] = eδt[j]*eμt[j,i]
      den += num[j]
    end
    for j=1:numProdt
      ms[j] += num[j]/den/dim.person
    end
  end
  # @inbounds for j=1:numProdt
  #   ms[j] /= dim.person
  # end
end

function setupMu!(eμ::Array{Matrix{Float64},1}, eμ₀::Array{Vector{Float64},1},
  adj::Array{Vector{Float64},1},
  θ₂::Vector{Float64}, dim::dimType, indx::indxType, L::Matrix{Float64},
  Lν::Matrix{Float64}, Xᵢ::Matrix{Float64}, ν::Matrix{Float64})

  for k=1:dim.θ₂
    i = indx.L[k,1]
    j = indx.L[k,2]
    L[i,j] = θ₂[k]
  end
  Lν[:] = L*ν

  for t=1:dim.market
    eμt = eμ[t]
    mask = view(indx.mkt,:,t)
    Xᵢt = view(Xᵢ,mask,:)
    for i=1:dim.person
      adj[t][i] = 0.0
      for j=1:dim.prodt[t]
        μ::Float64 = 0.0
        for k=1:dim.Xᵢ
          μ += Xᵢt[j,k]*Lν[k,i]
        end
        (μ > max(adj[t][i],6.91)) ? (adj[t][i]=μ) : nothing
        eμt[j,i] = μ
      end # end of j loop
      for j=1:dim.prodt[t]
        eμt[j,i] = exp(eμt[j,i]-adj[t][i])
      end
      eμ₀[t][i] = exp(-adj[t][i])
    end # end of i loop
  end
end
function setupMu!(
  eμ::Array{Matrix{Float64},1},
  eμ₀::Array{Vector{Float64},1},
  adj::Array{Vector{Float64},1},
  θ₂::Vector{Float64},
  Xᵢ,
  Dᵢ,
  νᵢ,
  dim::dimType,
  indx::indxType)

  θ̃₂ = copy(θ₂)
  for k in indx.expθ
    θ̃₂[k] = exp(θ₂[k])
  end
  for t=1:dim.market
    eμt = eμ[t]
    mask = view(indx.mkt,:,t)
    Xᵢt = view(Xᵢ,mask,:)
    for i=1:dim.person
      adj[t][i] = 0.0
      for j=1:dim.prodt[t]
        μ = 0.0
        for k=1:dim.Dᵢ
          μ += θ̃₂[k] * Xᵢt[j,k] * Dᵢ[t][k,j,i]
        end
        for k=(dim.Dᵢ+1):(dim.Dᵢ+dim.νᵢ)
          μ += θ̃₂[k] * Xᵢt[j,k] * νᵢ[k-dim.Dᵢ,i]
        end
        (μ < min(adj[t][i],0.0)) ? (adj[t][i]=μ) : nothing
        eμt[j,i] = μ
      end # end of j loop
      for j=1:dim.prodt[t]
        eμt[j,i] = exp(eμt[j,i]-adj[t][i])
      end
      eμ₀[t][i] = exp(-adj[t][i])
    end # end of i loop
  end
end

function stdError(
  θ₂     ::Vector{Float64},
  model  ::Int64,
  stage  ::Int64,
  iv     ::String,
  nPerson::Int64,
  impF   ::String,
  update ::Bool,
  dat    ::datType,
  dim    ::dimType,
  fxp    ::fxpType,
  mat    ::matType,
  set    ::setType,
  indx   ::indxType,
  pred   ::predType,
  diags  ::diagsType,
  mcid1st::Int64=0)

  if iv=="opt"
    if (impF=="optIVtest") | (impF=="optIV1") | (impF=="optIV2")
      θ̂₂ = importθ̂₂(model,stage,mcid1st,iv,nPerson,impF,dim,update)
      setupMu!(fxp.eμ,pred.eμ₀,pred.adj, θ̂₂,mat.Xᵢ,mat.Dᵢ,mat.νᵢ,dim,indx)
      fixPoint!(fxp.eδ,fxp.∂δ∂θ, fxp.eμ,dat,dim,mat,indx,pred,diags,θ̂₂)
      θ₁ = optimIV!(mat,θ̂₂,dat,dim,fxp,indx,pred,diags)
      println("θ₁: ",round.(θ₁,digits=4))
      setupMu!(fxp.eμ,pred.eμ₀,pred.adj, θ₂,mat.Xᵢ,mat.Dᵢ,mat.νᵢ,dim,indx)
      fixPoint!(fxp.eδ,fxp.∂δ∂θ, fxp.eμ,dat,dim,mat,indx,pred,diags,θ₂)
    else
      setupMu!(fxp.eμ,pred.eμ₀,pred.adj, θ₂,mat.Xᵢ,mat.Dᵢ,mat.νᵢ,dim,indx)
      fixPoint!(fxp.eδ,fxp.∂δ∂θ, fxp.eμ,dat,dim,mat,indx,pred,diags,θ₂)
      optimIV!(mat,θ₂,dat,dim,fxp,indx,pred,diags)
    end
  else
    setupMu!(fxp.eμ,pred.eμ₀,pred.adj, θ₂,mat.Xᵢ,mat.Dᵢ,mat.νᵢ,dim,indx)
    fixPoint!(fxp.eδ,fxp.∂δ∂θ, fxp.eμ,dat,dim,mat,indx,pred,diags,θ₂)
  end
  θ₁ = linearIV!(mat.ξ, fxp.eδ,fxp.eδ1d, dim,mat,indx;output=2)

  ∂δ∂θ = unpackStackedArray(fxp.∂δ∂θ,dim.obs,dim.θ₂,dim.market,indx.mkt)
  # ∂m∂θ = view(mat.Z,mat.mask,:)'*view([∂δ∂θ -mat.X],mat.mask,:)
  ∂m∂θ = mat.Z'*[∂δ∂θ -mat.X]
  invPhi = gmmWeight(mat,3,dat,dim,set,indx)
  covθ = inv(∂m∂θ'*invPhi*∂m∂θ)
  varθ = diag(covθ)
  seθ = try
    sqrt.(varθ)
  catch
    sqrt.(complex.(varθ))
  end
  for k=1:dim.θ₂
    if imag(seθ[k])!=0
      seθ[k] = NaN
    end
  end
  seθ = real(seθ)

  return seθ, θ₁, covθ
end

function linearIV!(ξ::Matrix{Float64},
  eδ::Vector{Vector{Float64}}, eδ1d::Vector{Float64},
  dim::dimType, mat::matType, indx::indxType; output::Int64=1)

  unpackStackedArray!(eδ1d, eδ, dim.market, indx.mkt)
  δ = log.(eδ1d)
  θ₁ = mat.XZPhiZXZ*δ
  if output==1
    return θ₁
  elseif output==2
    mat.ξ[:] = δ - mat.X*θ₁
    return θ₁
  elseif output==3
    mat.ξ[:] = δ - mat.X*θ₁
    return
  else
    error("Invalid output in linearIV()")
  end
end

function linearIV(
  eδ  ::Vector{Vector{Float64}},
  eδ1d::Vector{Float64},
  X   ::Matrix{Float64},
  Z   ::Matrix{Float64},
  dim ::dimType,
  indx::indxType)

  unpackStackedArray!(eδ1d, eδ, dim.market, indx.mkt)
  δ = log.(eδ1d)
  XZ = X'*Z
  invPhi = inv(Z'*Z)
  XZPhiZXZ = (XZ*invPhi*XZ')\(XZ*invPhi*Z')
  θ₁ = XZPhiZXZ*δ
  return θ₁
end

function gmmWeight(mat::matType, flag::Int64,
  dat::datType, dim::dimType, set::setType, indx::indxType)

  ξ = mat.ξ[:,1]
  nZ = size(mat.Z,2)
  if flag==1
    Phi = mat.Z'*diagm(0=>ξ.^2)*mat.Z
  elseif flag==3
    Phi = zeros(Float64,nZ,nZ)
    for gid=1:dim.region
      g = set.region[gid]
      for jid=1:dim.product
        b = set.product[jid,1]
        p = set.product[jid,2]
        mask = (dat.region.==g) .& (dat.brand.==b) .& (dat.prepost.==p)
        id   = findall(mask)
        for j=1:length(id)
          i   = id[j]
          z   = view(mat.Z,i,:)
          Phi = Phi + (z*z')*ξ[i]*ξ[i]
          iLag = findall(mask .& (dat.qtr.==dat.qtr[i]-1))
          if length(iLag)==1
            iL = iLag[1]
            zl = view(mat.Z,iL,:)
            z0zl = z*zl'
            Phi = Phi + 0.5*(z0zl + z0zl')*ξ[i]*ξ[iL]
          elseif length(iLag)>0
            error("Multiple lag obs at Region: %d, Brand: %d, Prepost: %d",g,b,p)
          end
        end # end of j
      end # end of jid
    end # end of gid
  end # end of if

  # Replace matrices
  return inv(Phi)
end

function optimIV!(
  mat::matType,
  θ̂₂::Vector{Float64},
  dat::datType,
  dim::dimType,
  fxp::fxpType,
  indx::indxType,
  pred::predType,
  diags::diagsType)

  # θ̂₁ from regressing δ on Z = [p̂,Xexo]
  # predict ŝ from σ̂,p̂
  # Invert ŝ to obtain δ̂(ŝ)
  # Z = [∂δ̂∂σ̂,Xexo]
  # ξ = δ - δ̂(Z,θ̂), θ̂₁ from Z
  θ̂₁ = linearIV(fxp.eδ,fxp.eδ1d, mat.X,[-mat.p̂ mat.Xexo],dim,indx)
  setupMu!(fxp.eμ̂,pred.eμ₀,pred.adj, θ̂₂,mat.X̂ᵢ,mat.Dᵢ,mat.νᵢ,dim,indx)

  # E[∂δ∂θ|Z]
  iter = 0
  dist = 1.0
  while (dist>diags.tolOptIV) & (iter<1e5)
    predictions!(fxp.∂δ̂∂θ,fxp.eδ̂, θ̂₁,θ̂₂,fxp.eμ̂, dim,mat,indx,pred;intNum=1)
    unpackStackedArray!(mat.∂δ̂∂θ,fxp.∂δ̂∂θ,dim.market,indx.mkt)
    updateIVmatrix!(mat)
    θ₁ = linearIV!(mat.ξ, fxp.eδ,fxp.eδ1d, dim,mat,indx;output=2)
    dist = maximum(abs.(θ̂₁-θ₁))
    #@printf("Iter: %4d alpha: %.4f Dist: %.15f\n",iter,θ₁[1],dist)
    iter += 1
    (rem(iter,1000)==0) ? @printf("Iter: %6d Dist: %.14f\n",iter,dist) : nothing
    # if rem(iter,100)==0
    #   @printf("Iter: %5d Dist: %.14f\n",iter,dist)
    # end
    θ̂₁[:] = copy(θ₁)
  end
  if iter>=1e5
    println("Routine optimIV() failed to converge.")
    fxp.eδ[1][:] = NaN
    return
  end
  return θ̂₁
end

function predictions!(
  ∂δ̂∂θ::Vector{Matrix{Float64}},
  eδ̂::Vector{Vector{Float64}},
  θ̂₁::Vector{Float64},
  θ̂₂::Vector{Float64},
  eμ̂::Vector{Matrix{Float64}},
  dim::dimType,
  mat::matType,
  indx::indxType,
  pred::predType;
  intNum::Int64=1)

  if intNum==1
    for t=1:dim.market
      eδ̂[t] .= exp.(mat.X̂[t]*θ̂₁)
      share!(pred.ms[t],pred.num[t],pred.eμ₀[t],eδ̂[t],eμ̂[t],dim.prodt[t],dim)
      jacobDelta!(∂δ̂∂θ[t],pred.∂S∂θ[t],mat.dμ̂dL[t],eδ̂[t],eμ̂[t],
        pred.num[t],pred.eμ₀[t],pred.s[t],pred.ms[t],pred.mss[t],
        dim.prodt[t],dim,indx,θ̂₂)
    end
  elseif intNum==2
    mat.Lξ .= mat.ξ[indx.idLag,1]
    x = view(mat.Lξ,mat.mask,1)
    y = view(mat.ξ,mat.mask,1)
    ρ = inv(x'*x)*(x'*y)
    @. mat.Eξ[mat.mask,1]  = x*ρ
    for i=1:dim.brand
      mask = view(mat.maskBrand,:,i)
      # mat.Eξ[mask,1] = mean(view(mat.ξ,mask,1))
      # mat.Eξ[mask,1] = mean(view(mat.ξ,.~(mat.mask),1))
      # println(mean(view(mat.ξ,mask,1)))
    end
    @. mat.Eξ[~(mat.mask),1] = mat.ξ[~(mat.mask),1]
    for t=1:dim.market
      mask = view(indx.mkt,:,t)
      eδ̂[t] .= exp.(mat.X̂[t]*θ̂₁ + view(mat.Eξ,mask,1))
      share!(pred.ms[t],pred.num[t],pred.eμ₀[t],eδ̂[t],eμ̂[t],dim.prodt[t],dim)
      jacobDelta!(∂δ̂∂θ[t],pred.∂S∂θ[t],mat.dμ̂dL[t],eδ̂[t],eμ̂[t],
        pred.num[t],pred.eμ₀[t],pred.s[t],pred.ms[t],pred.mss[t],
        dim.prodt[t],dim,indx,θ̂₂)
    end
  end
end

function updateGMMweight!(
  mat::matType,
  θ₂,
  dat::datType,
  dim::dimType,
  fxp::fxpType,
  set::setType,
  indx::indxType,
  pred::predType,
  diags::diagsType)

  setupMu!(fxp.eμ,pred.eμ₀,pred.adj, θ₂,mat.Xᵢ,mat.Dᵢ,mat.νᵢ,dim,indx)
  fixPoint!(fxp.eδ,fxp.∂δ∂θ, fxp.eμ,dat,dim,mat,indx,pred,diags,θ₂)
  θ₁ = linearIV!(mat.ξ, fxp.eδ,fxp.eδ1d, dim,mat,indx;output=2)
  XZ = mat.X' * mat.Z
  mat.invPhi   .= gmmWeight(mat,3,dat,dim,set,indx)
  mat.XZPhiZXZ .= (XZ*mat.invPhi*XZ')\(XZ*mat.invPhi*mat.Z')
end

end # end of module

max_stars_repo_name: yutec/FightBrand

max_issues_repo_name: yutec/FightBrand

max_forks_repo_name: yutec/FightBrand


text:
bessel_j(ν, z) = SpecialFunctions.besselj(ν, z)
bessel_y(ν, z) = SpecialFunctions.bessely(ν, z)


text:
module Schemalizers

using Random
using Dates
using DataFrames
using JuliaDB: ML, table

using AutoMLPipeline.AbsTypes
using AutoMLPipeline.Utils
import AutoMLPipeline.AbsTypes: fit!, transform!

export fit!,transform!
export Schemalizer, ML, table

"""
    Schemalizer(
       Dict(
           :schema => ML.Schema()
       )
    )

Automatically converts continuous features into zscore and categorical values into hot-bin encoding.

Example:

    function generatedf()
      Random.seed!(123)
      gdate = DateTime(2015,1,1):Dates.Minute(30):DateTime(2015,1,2) |> collect
      len = length(gdate)
      sex = rand(['m','f'],len)
      x1=rand(1:100,len)
      x2=rand(1:100,len)
      x3=rand(1:1000,len)
      DataFrame(date=gdate,sex=sex,f1=x1,f2=x2,f3=x3)
    end

    df = generatedf()
    m = Schemalizer(Dict(:schema=>Dict(:sex => ML.Categorical)))
    fit!(m,df)
    res=transform!(m,df)
    @assert isapprox(sum(Matrix(res.^2)),193,atol=1e-2)
    m = Schemalizer(Dict())
    fit!(m,df)
    res = transform!(m,df)
    @assert isapprox(sum(Matrix(res.^2)),144,atol=1e-2)

Implements: `fit!`, `transform!`
"""
mutable struct Schemalizer <: Transformer
  model
  args
  function Schemalizer(args=Dict())
    default_args = Dict(
        :schema => ML.Schema(),
    )
    margs=mergedict(default_args, args)
    new(nothing,margs)
  end
end

"""
    fit!(sch::Schemalizer, features::T, labels::Vector=[]) where {T<:Union{Vector,Matrix,DataFrame}}

Check validity of `features`: Date,Val data or just Vals Matrix
"""
function fit!(sch::Schemalizer, features::DataFrame, labels::Vector=[])
  isempty(features) && error("Schemalizer.fit: data format not recognized.")
  if isempty(sch.args[:schema])
    sch.args[:schema] = ML.schema(table(features))
  else
    sch.args[:schema] = ML.schema(table(features),hints=sch.args[:schema])
  end
  sch.model = sch.args
end

"""
    transform!(sch::Schemalizer, features::T) where {T<:Union{Vector,Matrix,DataFrame}}

Normalized continous features and hot-bit encode categorical features
"""
function transform!(sch::Schemalizer, features::DataFrame)
  isempty(features) && (return DataFrame())
  ML.featuremat(sch.args[:schema],table(features))' |> DataFrame
end

end


text:
struct TimestampResolver{T₁ <: AbstractString, T₂ <: Number, T₃ <: Number} <: Function
    timezone::T₁
    timestamp::T₂
    intersample_duration::T₃
    nested_interval::NestedInterval
end

function get_timezone(resolver::TimestampResolver)
    resolver.timezone
end

function get_timestamp(resolver::TimestampResolver)
    resolver.timestamp
end

function get_intersample_duration(resolver::TimestampResolver)
    resolver.intersample_duration
end

function get_nested_interval(resolver::TimestampResolver)
    resolver.nested_interval
end

function (resolve_timestamp::TimestampResolver)(x::Number)
    # All times are in milliseconds
    timezone = resolve_timestamp.timezone
    nested_interval = resolve_timestamp.nested_interval
    t₀ = resolve_timestamp.timestamp
    tₛ = get_start(nested_interval) * resolve_timestamp.intersample_duration
    tₙ = x * resolve_timestamp.intersample_duration
    # Display timestamp.
    timestamp₀ = Dates.unix2datetime(Dates.value(Second(t₀) + round(Millisecond(tₛ + tₙ), Dates.Second)))
    # TODO Make the timezone a parameter of the TimestampResolver
    timestamp₁ = ZonedDateTime(timestamp₀, TimeZone(timezone, TimeZones.Class(:ALL)); from_utc = true)
    timestr = Dates.format(timestamp₁ , dateformat"HH:MM:SS")
end


text:
# https://www.nature.com/articles/s41598-017-16178-8.pdf?proof=t
using MultidimensionalSpectroscopy, PyPlot, QuantumOptics, LinearAlgebra, FFTW, Colors,
        Printf, DelimitedFiles, Random, Combinatorics, Distributions

QuantumOpticsBase.set_printing(standard_order=false,rounding_tol=1e-3)

# make sure to set script directory as pwd()
cd(@__DIR__)

mod_name = "A"
fn_log =  "logs/" * mod_name * ".log"

# erase .log file
try
    rm("logs/" * mod_name * ".log")
    rm(fn_log)
catch
end

### -- select model parameters --- ###
mod_name = "parameters/" * mod_name
include( mod_name * ".params")

logger(String(read(mod_name * ".params")), fn_log)

pygui(true)

calc_2d = true

logger("\nDo 2D calculation: ", calc_2d, fn_log)

cmp = create_colormap("bright");

# use 2 level system to approximate Frenkel excitons ... seems to coarse, but well
Nlev = 3

b = NLevelBasis(Nlev)

# energy of TDBC Frenkel exciton bright state #TODO: add some details, such as J angles, etc.
Emon = E_monomer
Emonb = 2 * E_monomer + 0.1
# not important here, but in general for TDBC aggregates
# Δ = .299                          # shift monomer -> aggregate
# J = -Δ / (2 * cos(pi /(N+1)))     # nearest-neighbour coupling constant

# transitions and monomer Hamiltonian
σ⁺ = transition(b,2,1)
σ31 = transition(b,3,1)
σ32 = transition(b,3,2)
σ⁻ = transition(b,1,2)
σ13 = transition(b,1,3)
σ23 = transition(b,2,3)
H_mon = Emon * σ⁺ * σ⁻ + Emonb * σ31*σ13

σ⁺ = σ⁺ + σ32
σ⁻ = σ⁻ + σ23

logger("\nMonomer Hamiltonian:\n\n", dense(H_mon), fn_log)

# use N uncoupled TLSs
N = N_monomers      # from .params file
B = b^N             # basis for matter part of H

# random (normal dist) list of energy shifts (factors) scaled arbitrarily by scale. Used to account for static disorder of exciton transition energy
dE_n = energy_disorder  # from .params file

# keep only factors that increase E_mon, since Frenkel exciton absorption spectrum is not Gaussian but favors high energy side. This could be explained by
# a maximum delocalisation length resulting in the lowest energy bright state beyond which no longer delocalisation is possible. Shorter delocalisation lengths,
# due to structural breaks in the exciton or other effecsts, lead to the lowest/bright transition being at higher energies
# idea1
#dE_n[dE_n.<=1] = 2 .- dE_n[dE_n.<=1] #BUG #CHECK This shifts polariton spectrum to higher E
# idea2, better since keeps all values
#dE_n = 1 .+ abs.(dE_n .- 1)

# can deactive random energy shifts
dE_n = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] .+ 1

# functions/methods to build matter H and transition operators, by embedding monomer H and σ at site N
H_n(n)  = embed(B,n,H_mon * dE_n[n]) # Hamiltonian at site n => Hₙ with random, normally distributed energy offset dE
Σ⁺_n(n) = embed(B,n,σ⁺)              # raising operator at site n
Σ⁻_n(n) = embed(B,n,σ⁻)              # lowering operator at site n

# need this workaroud to deal with a single TLS... #BUG
if N != 1
    Σ⁺   = sum(Σ⁺_n.(1:N))
    Σ⁻   = sum(Σ⁻_n.(1:N))
    Σ⁺n  = Σ⁺_n.(1:N)
    Σ⁻n  = Σ⁻_n.(1:N)
    Hexc = sum(H_n.(1:N))                   # Hamiltonian of disorder TLSs ... NOT coupled to each other
else
    Σ⁺   = σ⁺
    Σ⁻   = σ⁻
    Σ⁺n   = σ⁺
    Σ⁻n   = σ⁻
    Hexc = H_mon                             # Hamiltonian of disorder TLSs ... NOT coupled to each other
end

logger("\nMatter Hamiltonian:\n\n", dense(Hexc), fn_log)

# create cavity basis
bcav = FockBasis(N_photonStates)  #FACT: #CHECK need second FockState for ESA !?
#bcav = FockBasis(1)
#BUG#TODO: Having 2 photon states messes up transforming between eigen and site basis in subspace
#IDEA MAYBE I need two 1 photon Fock states ?
Ecav = E_cavity         # slightly red-tuned from excitonic transition, as main system in paper
a    = destroy(bcav) # ⊗ one(bcav) + one(bcav) ⊗ destroy(bcav)
at   = create(bcav)  #⊗ one(bcav) + one(bcav) ⊗ create(bcav)

# embed operators into full basis
A  = a ⊗ one(Hexc)
At = at ⊗ one(Hexc)
Σ⁺ = one(a) ⊗ Σ⁺
Σ⁻ = one(a) ⊗ Σ⁻
Σ⁺n = [one(a) ⊗ Σ⁺n[i] for i in 1:N]
Σ⁻n = [one(a) ⊗ Σ⁻n[i] for i in 1:N]

# create cavity H
Hcav = Ecav * at * a

logger("\nCavity Hamiltonian:\n\n", dense(Hcav), fn_log)

# create basis of complete system
Bfull = bcav ⊗ B

# create Jaynes/Tavis-Cummings Hamiltonian for light-matter coupling g
g = 0.15 / sqrt(N)                               # should make splliting independent of N
#g = 0
H = embed(Bfull,1,Hcav) + one(Hcav) ⊗ Hexc + g * (Σ⁺ * A + Σ⁻ * At)

logger("\nFull Hamiltonian:\n\n", dense(H), fn_log)

# sort by increasing values of Hamiltonian (rough #CHECK ) ... sort into excitation sectors

idx = sortperm(real(diag((H).data)))
H.data = H.data[idx,idx]
A.data = A.data[idx,idx]
At.data = At.data[idx,idx]
Σ⁺.data = Σ⁺.data[idx,idx]
Σ⁻.data = Σ⁻.data[idx,idx]
for i in 1:N
    Σ⁺n[i].data = Σ⁺n[i].data[idx,idx]
    Σ⁻n[i].data = Σ⁻n[i].data[idx,idx]
end

logger("\nOrdered Hamiltonian:\n\n", dense(H), fn_log)

# transition operator
#μ = Σ⁺ + Σ⁻    # transitions between excitonic states only, #FACT: not possible induced by ext. EM field in cavity, every via optical mode
μ = At + A #+ At^2 + A^2                                     #FACT: transitions induced by external field

# density matrices
rho0 = dm(fockstate(bcav,0) ⊗ tensor([nlevelstate(b,1) for i in 1:N]))
rho1 = dm(fockstate(bcav,0) ⊗ tensor([nlevelstate(b,2) for i in 1:N]))

# transition operator from ground to first ES
μ12 = rho0 * μ + μ * rho0

# sort (see above)

μ12.data  = μ12.data[idx,idx]
rho0.data = rho0.data[idx,idx]
rho1.data = rho1.data[idx,idx]

# transition from first ES to higher ES #FACT: should connect excitation on exciton site and excitation of exciton + cavity
μ23 = μ - μ12

# normalization of μ12 and μ23 ... seems not really necessary #DELETE
rho1 = μ12 * rho0 * μ12
#μ12 = μ12 / sqrt(tr(rho1))
#rho1 = μ12 * rho0 * μ12
rho2 = μ23 * rho1 * μ23
#μ23 = μ23 / sqrt(tr(rho2))
#μ23 = μ23
rho2 = μ23 * rho1 * μ23

# Lindblad dissipation operators
L1 = A           # FACT: #CHECK only the relaxation channel from single exc. sector to GS is relevant ...
L2 = Σ⁻          # FACT: actually, having the elements that connect single and double exc. sector seems to ...
                        # FACT: broaden the ESA peaks (the only signal that considers single and double exc. sectors)
L3 = .5 * ((At*A) - (A*At)) # not sure if this one is relevant/meaningful/... #CHECK

#L3 = A  - L1
#L4 = Σ⁻ - L2

# rates and list of Lindblad dissipation operators
Γ =      [.25, .000001]         # Γ[1] was 0.2 before
L = Γ .* [L1,  L2]

# plot energy levels
figure(); title("Energy level diagram")
plot_levels(H,0)
plot_levels(Hcav,-1)
plot_levels(Hexc,1)
plot_levels(H_mon,2)
xticks([-1, 0, 1, 2], ["cavity", "full", "matter", "monomer"])

##

# create subspace (single excitation, double excitation, ...) in eigenbasis/polariton basis
H_si, transf_op_si, P_si, L_si, rho0_si, μ12_si, μ23_si, Σ⁺_si, Σ⁻_si, At_si, A_si = create_subspace([H],"si", L, rho0, μ12, μ23, Σ⁺, Σ⁻, At, A)
H_si = H_si[1]
H, transf_op, P, L, rho0, μ12, μ23, Σ⁺, Σ⁻, At, A = create_subspace([H],"bi", L, rho0, μ12, μ23, Σ⁺, Σ⁻, At, A)
H = H[1]
#FACT: Using "bi_lowest" does not give ESA that eliminates above diagonal cross-peak
#FACT: Neither does using "bi_polariton", which takes lowest and highest levels of double excitation sector
#FACT: only using "bi" adds ESA signal at above diagonal cross-peak that cancels out the GSB above diagonal cross-peak after some time T

Σ⁺n_si = [P_si * Σ⁺n[i] * P_si' for i in 1:N]
Σ⁻n_si = [P_si * Σ⁻n[i] * P_si' for i in 1:N]

Σ⁺n = [P * Σ⁺n[i] * P' for i in 1:N]
Σ⁻n = [P * Σ⁻n[i] * P' for i in 1:N]

logger("\nCHANGE TO SUBSPACE", fn_log)
logger("\n0th, 1st, 2nd Manifold", fn_log)
logger("\nCoupled (eigen) basis", fn_log)
logger("\nHamiltonian:\n\n", dense(H), fn_log)
logger("\nTransition operator gs <-> es:\n\n", dense(μ12), fn_log)
logger("\nTransition operator es <-> d-es:\n\n", dense(μ23), fn_log)
logger("\n0th, 1st, 2nd Manifold", fn_log)
logger("\nUncoupled (site) basis", fn_log)
logger("\nHamiltonian:\n\n", transf_op * H * transf_op', fn_log)
logger("\nTransition operator gs <-> es:\n\n", transf_op * μ12 * transf_op', fn_log)
logger("\nTransition operator es <-> d-es:\n\n", transf_op * μ23 * transf_op', fn_log)
logger("\n0th, 1st Manifold", fn_log)
logger("\nCoupled (eigen) basis", fn_log)
logger("\nHamiltonian:\n\n", dense(H_si), fn_log)
logger("\nTransition operator gs <-> es:\n\n", μ12_si, fn_log)
logger("\nTransition operator es <-> d-es:\n\n", μ23_si, fn_log)
logger("\nUncoupled (site) basis", fn_log)
logger("\nHamiltonian (o, 1, 2 Manifold; site):\n\n", transf_op_si * H_si * transf_op_si', fn_log)
logger("\nTransition operator gs <-> es:\n\n", transf_op_si * μ12_si * transf_op_si', fn_log)
logger("\nTransition operator es <-> d-es:\n\n", transf_op_si * μ23_si * transf_op_si', fn_log)
plot_levels(H_si,-0.05,col="r",ls="dashed")
plot_levels(H   , 0.05,col="b",ls="dotted")

# transf_op can be used to convert between polariton and site basis #FACT: required prior sorting of operators into blocks of excitation sectors
# H_site    = transf_op     * H     * transf_op'
# H_site_si = transf_op_si  * H_si  * transf_op_si'

#TODO: vary g instead of E_mon ... try to see effect
g = g .* dE_n
#Hk(i) = (Ecav * At_si * A_si) + (Emon * Σ⁺_si * Σ⁻_si) + g[1] * (Σ⁺_si * A_si + Σ⁻_si * At_si)

# time list
tlist = [0:1.1:200;]

logger("\nTime list:\n\n", tlist, fn_log)

# Lindblad time evolution
corr    = timecorrelations.correlation(tlist, rho0_si, H_si, L_si, μ12_si, μ12_si)
# #TODO: test with varying g ...
#corr = sum([timecorrelations.correlation(tlist, rho0_si, Hk(i), L_si, μ12_si, μ12_si) for i in 1:1]) ./ 1

# zeropadding for smoother data only
zp = 10
corr    = zeropad(corr,zp)
tnew, ~ = interpt(tlist,zp)

# calculate and plot spectrum
ω, spec = timecorrelations.correlation2spectrum(tnew, corr; normalize_spec=true);
figure()
subplot(211)
plot(ω,spec)

# load and plot experimental absorption spectrum
TDBC_cav   = readdlm("TDBC_cav.dat")
TDBC_cav_E = TDBC_cav[:,2] ./ 8065      # convert from cm^-1 to eV
TDBC_cav_I = TDBC_cav[:,3] ./ 0.176     # normalize to 1 (in the visible)

plot(-TDBC_cav_E,TDBC_cav_I,"r:")

# TODO: spectral density ... so far only trial and error ... Seems like it needs to bridge the E gap between P+ and P- for relaxation to occur
width = .005
#ωq_TDBC     = [0.12 0.336]
#ampl_TDBC   = [0.12 3] / (N)

# from ... http://www.rsc.org/suppdata/c8/sc/c8sc00171e/c8sc00171e1.pdf
#ωq_TDBC = [40, 80, 120, 150, 185, 197] ./ 8065 #.* 0.02998 # cm-1 in Hz #./ 8065 # in eV instead
# => [0.00496 0.00992 0.01488 0.01860  0.02294 0.02443]
#ampl_TDBC = [14, 18,  25,  43,  42, 67, 60]

#ωq_TDBC     = [0.0084, 0.0120, 0.1204, 0.1289, 0.1304, 0.1424, 0.303] #FACT: relaxation from UP to LP when ωq_TDBC bridges energy gap between UP and LP. However, no relaxation to dark states ... DUE TO LACK OF COUPLING ????
ωq_TDBC     = ωq        # from .params file
ampl_TDBC   = Aq        # from .params file

logger("Spectral frequencies: ", ωq_TDBC, fn_log)
logger("Spectral amplitudes: ", ampl_TDBC, fn_log)
logger("Spetral peak width: ", width, fn_log)

try     # need to delete method first when working on it
    local m = @which spectral_density(1)
    Base.delete_method(m)
catch
end
function spectral_density(ω) # power spectral density, thermal bath noise spectrum
    #w = [(ω .- ωq_TDBC[i]) ./ (.5 * width) for i in 1:length(ωq_TDBC)]
    b = 1 / 0.0259 # eV k_B*T at T=300 K
    f = .5 # 0 ... 1: relative importance of non-radiation vs. dephasing ... f = 0 > no dephasing
    γ = 0.7 # from absorption linewidth γ = γ_non-rad + γ_dephasing
    η = 1
    η = f * γ / (2 * pi * 0.0259)
    ω_cut = .1
    if ω == 0
        Jω = .0
    #elseif !isless(real(ω),0)
    #else
    #    w = [(ω .- ωq_TDBC[i]) ./ (.5 * width) for i in 1:length(ωq_TDBC)]
    #    Jω = sum([ampl_TDBC[i] ./ (1 .+ w[i].^2) for i in 1:length(ωq_TDBC)])  # probably most realistic ... ???
    #    #Jω = 1
    end
    if  !isless(real(ω),0) && ω != 0
        w = [(ω .- ωq_TDBC[i]) ./ (.5 * width) for i in 1:length(ωq_TDBC)]
        Jω = sum([ampl_TDBC[i] ./ (1 .+ w[i].^2) for i in 1:length(ωq_TDBC)])
        Jω = η * ω * exp(-(ω/ω_cut)^2) + Jω
        Jω = Jω * (1/exp(b *  ω - 1) + 1)
    elseif isless(real(ω),0)
        w = [(ω .- ωq_TDBC[i]) ./ (.5 * width) for i in 1:length(ωq_TDBC)]
        Jωa = sum([ampl_TDBC[i] ./ (1 .+ w[i].^2) for i in 1:length(ωq_TDBC)])
        Jω = η * -ω * exp(-(-ω/ω_cut)^2)
        Jω = Jω * (1/exp(b * -ω - 1)) + Jωa
    end
    return Jω

end

#plot spectral density
#figure()
subplot(212)
plot(collect(-1:.001:1),(spectral_density.(collect(-1:.001:1))))

# Rate and operator for redfield process (GS -> single exc sector)
Γ_R = .23

# this means collective mode of all TLSs dephases ...
a_ops_si = [Γ_R * Σ⁺_si * Σ⁻_si, spectral_density]
# ... and this: each TLS dephases individually
a_ops_si = vcat([[Σ⁺n_si[i] * Σ⁻n_si[i], spectral_density] for i in 1:N]...)

R_si, ekets = timeevolution.bloch_redfield_tensor(H_si, [a_ops_si]; J=1 .* L_si)

# time evolution, Redfield
tout, rhot = timeevolution.master_bloch_redfield(tlist,μ12_si*rho0_si,R_si,H_si)

corr = expect(μ12_si,rhot)
# calculate and plot spectrum
corr    = zeropad(corr,zp)
tnew, ~ = interpt(tlist,zp)
ω, spec = timecorrelations.correlation2spectrum(tnew, corr; normalize_spec=true);
subplot(211)
plot(ω,spec)
#FACT: J(0) > 0 induces pure dephasing and leads to broadening of ALL transitions
#FACT: J(ω) > 0 for ω > 0 leads to relaxation between states with Ei - Ef = ω ...
#CHECK what role does the operator play ?

figure(figsize=(14,4));
subplot(131)
plot(tout,real(corr[1:length(tout)]))
tout, rhot = timeevolution.master_bloch_redfield(tlist,μ12_si*rho0_si*μ12_si,R_si,H_si)
rho1test = copy(rho0_si)
rho1test.data[1,1] = 0
rho1test.data[2,2] = 1
#tout, rhot = timeevolution.master_bloch_redfield(tlist,rho1test,R_si,H_si)
#corr_gs = expect(Σ⁺_si*Σ⁻_si,rhot)
corr_gs = expect(rho0_si,rhot)
#corr_es = expect(Σ⁻_si*Σ⁺_si,rhot)
#corr_es = expect(μ12_si * rho0_si * μ12_si,rhot)
corr_es_mat = expect(Σ⁺_si * rho0_si * Σ⁻_si / N,rhot)
corr_es_cav = expect(At_si * rho0_si * A_si, rhot)
plot(tout,real(corr_gs))
#plot(tout,corr_es) # == corr_es_cav
plot(tout,real(corr_es_cav))
plot(tout,real(corr_es_mat))
legend(["corr. function", "ground state", "cavity", "matter"])

rhot_site = [transf_op_si * i  * transf_op_si' for i in rhot]
corr_es1_mat = expect(rho0_si, rhot_site)
#plot(tout,corr_es1_mat)

vecs   = eigvecs(dense(H_si).data)
eivecs = [Ket(H_si.basis_l,vecs[:,i]) for i in 1:length(vecs[1,:])]

subplot(132);
title("eigen/coupled basis") # levels: 1 -> GS, 2 -> LP, 3 - N-1 -> ..., N -> UP
for i in 1:H_si.basis_l.shape[1]
    local es_n = real(expect(dm(eivecs[i]),rhot))
    plot(tout,es_n,linestyle="dashed", label=string(i)); #append!(leg,[string(i)])
    ylim((-0.1, 1))
end
legend()

es_mat = zeros(length(tout))
subplot(133); title("uncoupled basis") # levels: 1 -> GS, 2 -> cav., 3 - N -> TLSs
for i in 1:H_si.basis_l.shape[1]
    local es_n = real(expect(dm(eivecs[i]),rhot_site))
    plot(tout,es_n,linestyle="dashed", label=string(i)); #append!(leg,[string(i)])
    if i > 2
        global es_mat = es_mat .+ es_n
    end
end
plot(tout, es_mat, label="Σ(TLSs)")
legend()

# Rate and operator for redfield process (GS up to double exc sector)
a_ops = [Γ_R * Σ⁺ * Σ⁻, spectral_density]
a_ops = vcat([[Σ⁺n[i] * Σ⁻n[i], spectral_density] for i in 1:N]...)

R, ekets = timeevolution.bloch_redfield_tensor(H, [a_ops]; J=1 .* L)

# choose method to use for simulation of 2D spectra
method = "redfield"

# change parameters accordingly #TODO: simplify a bit
if method == "lindblad"
    F = L
    F_si = L_si
    use_sub = true
    H = H[1]
elseif method == "redfield"
    F = R
    F_si = R_si
    use_sub = true
end

if calc_2d
    ## calculate (complex) 3rd order corr function (with T=0)
    zp = 11 # zeropad up to 2^zp

    ## calculate 2D spectra at
    T = [0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 100, 200, 300] # energy eV -> t => [hbar / eV] ~ 0.66 fs
    T = [0, 30]
    spectra2d = Array{out2d}(undef, length(T))

    # cannot plot inside cmds.jl when multithreading
    #for i = 1:length(T)
    # multithreading (run several T steps in parallel)!
    Threads.@threads for i = 1:length(T)
        spectra2d[i] = make2Dspectra(tlist,[rho0, rho0_si],[H, H_si],[F, F_si],[μ12, μ12_si],[μ23, μ23_si],T[i],
                                            method;debug=true,use_sub=use_sub,
                                                t2coh="kin",zp=zp)
    end

    ## crop 2D data and increase dw
    spectra2d = crop2d.(spectra2d,1.5;w_max=2.6,step=1)
    spectra2d = round2d.(spectra2d)
end

if calc_2d
    ## plot 2D spectra for each(?) T
    # what to plot
    rep="abst"
    scal="lin"

    ## make  subplot layout
    nplots = length(T);
    ncols = Int32(ceil(sqrt(nplots)));
    nrows = Int32(ceil(nplots / ncols));

    # determine maximum value in dataset out2d[:].full2d[:,:]
    maxi = Float64(maximum([maximum(real(spectra2d[i].full2d)) for i in 1:length(spectra2d)]))

    # plot 2D spectra
    fig, ax = subplots(nrows,ncols,sharex=true,sharey=true,figsize=(ncols*3.5,nrows*3))
        if nplots == 1
                ax = [ax]
        end
    fig.suptitle(rep * " 2D spectrum (" * scal * ". scaling)")
        k = 0
        for i = 1:nrows
            for j = 1:ncols
                global k += 1
                if k > nplots
                    continue
                end
                sca(ax[j,i])
                ax[j,i].set_aspect="equal"
                plot2d(spectra2d[k].ω,round.(spectra2d[k].full2d,digits=1);repr=rep,scaling=scal,norm=maxi)
                title("2D spectrum at $(T[k]) fs")
            end
        end
        tight_layout()
        subplots_adjust(top=0.825)

    ## plot TA (summed 2D spectrum)
    figure()
    ta = [sum(spectra2d[i].full2d,dims=1) for i in 1:length(spectra2d)]
    #[plot(out2d[1].ω,ta[i]') for i in 1:length(ta)]
    [plot(3e8 * 4.136e-15 ./ spectra2d[1].ω ./ 1e-9,-ta[i]') for i in 1:length(ta)]
    xlabel("Detection (ω₃)")
    legend([string(i) * " fs" for i in T])
    tight_layout()
end

# to plot time traces
# cmds.plot_timeTrace([real(i.full2d) for i in out2d],T,out2d[1].ω,[1.925, 1.925, 2.24, 2.24],[1.925, 2.24, 2.24, 1.925])

max_stars_repo_name: lbmnky/CMDS

max_issues_repo_name: lbmnky/CMDS

max_forks_repo_name: lbmnky/CMDS


text:
module YOLO
include("common.jl")
!isfile(joinpath(pretrained_dir,"v2_tiny","voc2007","v2_tiny_voc.weights")) && @error "YOLO has build errors. Re-run `]build YOLO`"
!isfile(joinpath(pretrained_dir,"v2","voc2007","v2_voc.weights")) && @error "YOLO has build errors. Re-run `]build YOLO`"

@static Sys.isapple() && using QuartzImageIO
using FileIO, ImageMagick
using Knet: Knet, progress, progress!, gpu, KnetArray, relu, minibatch, conv4, pool, softmax
using Random, DelimitedFiles, OffsetArrays
using ImageFiltering, ImageTransformations, Colors, ImageCore
using LightXML
import ProgressMeter
using GeometryTypes
using Requires

const GPU = Knet.gpu()
const xtype = (GPU >= 0 ? Knet.KnetArray{Float32} : Array{Float32})#if gpu exists run on gpu

Base.@kwdef mutable struct BBOX
    x::Float32 #Left hand edge in scaled image width unnits (0-1)
    y::Float32 #Right hand edge in scaled image width unnits (0-1)
    w::Float32 #Width in scaled image width unnits (0-1)
    h::Float32 #Height in scaled image height unnits (0-1)
end

Base.@kwdef mutable struct TruthLabel
    bbox::BBOX
    class::Int #Object class
    difficult::Int32 #Difficult flag
end

Base.@kwdef mutable struct PredictLabel
    bbox::BBOX
    class::Int
    conf::Float32
end

Base.@kwdef mutable struct LabelledImageDataset
    name::String
    objects::Dict{Int,String}
    objectcounts::Vector{Int32}
    image_size_lims::Tuple{Int,Int} = (-1, -1)
    images_dir::String
    labels_dir::String
    image_paths::Array{String} = String[]
    label_paths::Array{String} = String[]
    labels::Vector{Vector{TruthLabel}} = Vector{Vector{TruthLabel}}(undef, 0)
end

Base.@kwdef mutable struct Settings
    dataset_description::String = ""
    source::String = ""
    weights_filepath::String = ""
    image_shape::Tuple{Int,Int}
    image_channels::Int
    namesdic::Dict{String,Int64}
    numsdic::Dict{Int64,String}
    anchors::Array{Tuple{Float64,Float64}}
    num_classes::Int
    minibatch_size::Int = 1
    grid_x::Int
    grid_y::Int
    cell_bboxes::Int
end

Base.@kwdef mutable struct LoadedDataset
    imstack_mat::xtype                       #4D image stack of type Float32 (w,h,colorchannels,numimages)
    paddings::Vector{Array{Int}}
    labels::Vector{Vector{TruthLabel}}
    #label #labels as tupple of arrays. tupples are designed as (ImageWidth, ImageHeight,[x,y,objectWidth,objectHeight],[x,y,objectWidth,objectHeight]..)
end

include("datasets.jl")
include("pretrained.jl")
include("models.jl")
include("preprocess.jl")
include("postprocess.jl")

function __init__()
    @require Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" include("makierendering.jl")
end

# function batchaccuracy()
#     #=User guide
#     inputandlabelsdir => takes Voc labels and inputs folder location respectively and returns 2 arrays
#     images directories and their labels' directories.
#
#     prepareinputlabels => takes images and labels directories and returns 3 arrays
#     input as 416*416*3*totalImage and
#     labels as tupple of arrays.
#     tupples are designed as (ImageWidth, ImageHeight,[x,y,objectWidth,objectHeight],[x,y,objectWidth,objectHeight]..)
#     images are padded version of given images.
#
#     inputdir => takes the directory of input for saving output and returns array of directories of the images
#     prepareinput => takes the array of directories of the images and returns 416*416*3*totalImage. =#
#     #prepare data for accuracy
#     images,labels = collectTruthLabelImagePairs(dataset_labels_dir,dataset_images_dir)
#     inp,out,imgs = prepareinputlabels(images,labels)
#     print("input for accuracy:  ")
#     println(summary(inp))
#     #Minibatching process
#     accdata = minibatch(inp,out,minibatch_size;xtype = xtype)
#     AP = accuracy(model,accdata,0.0,o[:iou],o[:iouth])
#     display(AP)
#     print("Mean average precision: ")
#     println(calculatemean(AP))
#     if o[:record] == true
#         drawdata = minibatch(inp,imgs,minibatch_size; xtype = xtype)
#         #output of Voc dataset.
#         #return output as [ [(x,y,width,height,classNumber,confidenceScore),(x,y,width,height,classNumber,confidenceScore)...] ,[(x,y,width,height,classNumber,confidenceScore),(x,y,width,height,classNumber,confidenceScore)..],...]
#         #save the output images into given location
#         result = saveoutput(model,drawdata,o[:confth],o[:iou]; record = true, location = "VocResult")
#     end
# end
end # module


text:
using Makie: RectanglePacker

function to_meshcolor(color::TOrSignal{Vector{T}}) where T <: Colorant
    TextureBuffer(color)
end

function to_meshcolor(color::TOrSignal{Matrix{T}}) where T <: Colorant
    Texture(color)
end
function to_meshcolor(color)
    color
end

using Makie: get_texture_atlas

vec2quaternion(rotation::StaticVector{4}) = rotation

function vec2quaternion(r::StaticVector{2})
    vec2quaternion(Vec3f(r[1], r[2], 0))
end
function vec2quaternion(rotation::StaticVector{3})
    Makie.rotation_between(Vec3f(0, 0, 1), Vec3f(rotation))
end

vec2quaternion(rotation::Vec4f) = rotation
vec2quaternion(rotation::VectorTypes) = const_lift(x-> vec2quaternion.(x), rotation)
vec2quaternion(rotation::Observable) = lift(vec2quaternion, rotation)
vec2quaternion(rotation::Makie.Quaternion)= Vec4f(rotation.data)
vec2quaternion(rotation)= vec2quaternion(to_rotation(rotation))
GLAbstraction.gl_convert(rotation::Makie.Quaternion)= Vec4f(rotation.data)
to_pointsize(x::Number) = Float32(x)
to_pointsize(x) = Float32(x[1])
struct PointSizeRender
    size::Observable
end
(x::PointSizeRender)() = glPointSize(to_pointsize(x.size[]))

"""
returns the Shape for the distancefield algorithm
"""
primitive_shape(::Union{AbstractString, Char}) = DISTANCEFIELD
primitive_shape(x::X) where {X} = primitive_shape(X)
primitive_shape(::Type{T}) where {T <: Circle} = CIRCLE
primitive_shape(::Type{T}) where {T <: Rect2} = RECTANGLE
primitive_shape(x::Shape) = x

"""
Extracts the scale from a primitive.
"""
primitive_scale(prim::GeometryPrimitive) = Vec2f(widths(prim))
primitive_scale(::Union{Shape, Char}) = Vec2f(40)
primitive_scale(c) = Vec2f(0.1)

"""
Extracts the offset from a primitive.
"""
primitive_offset(x, scale::Nothing) = Vec2f(0) # default offset
primitive_offset(x, scale) = const_lift(/, scale, -2f0)  # default offset

"""
Extracts the uv offset and width from a primitive.
"""
primitive_uv_offset_width(c::Char) = glyph_uv_width!(c)
primitive_uv_offset_width(str::AbstractString) = map(glyph_uv_width!, collect(str))
primitive_uv_offset_width(x) = Vec4f(0,0,1,1)

"""
Gets the texture atlas if primitive is a char.
"""
primitive_distancefield(x) = nothing
primitive_distancefield(::Union{AbstractString, Char}) = get_texture!(get_texture_atlas())
primitive_distancefield(x::Observable) = primitive_distancefield(x[])

function char_scale_factor(char, font)
    # uv * size(ta.data) / Makie.PIXELSIZE_IN_ATLAS[] is the padded glyph size
    # normalized to the size the glyph was generated as.
    ta = Makie.get_texture_atlas()
    lbrt = glyph_uv_width!(ta, char, font)
    width = Vec(lbrt[3] - lbrt[1], lbrt[4] - lbrt[2])
    width * Vec2f(size(ta.data)) / Makie.PIXELSIZE_IN_ATLAS[]
end

# This works the same for x being widths and offsets
rescale_glyph(char::Char, font, x) = x * char_scale_factor(char, font)
function rescale_glyph(char::Char, font, xs::Vector)
    f = char_scale_factor(char, font)
    map(x -> f * x, xs)
end
function rescale_glyph(str::String, font, x)
    [x * char_scale_factor(char, font) for char in collect(str)]
end
function rescale_glyph(str::String, font, xs::Vector)
    map((char, x) -> x * char_scale_factor(char, font), collect(str), xs)
end

@nospecialize
"""
This is the main function to assemble particles with a GLNormalMesh as a primitive
"""
function draw_mesh_particle(p, data)
    rot = get!(data, :rotation, Vec4f(0, 0, 0, 1))
    rot = vec2quaternion(rot)
    delete!(data, :rotation)
    to_opengl_mesh!(data, p[1])
    @gen_defaults! data begin
        position = p[2] => TextureBuffer
        scale = Vec3f(1) => TextureBuffer
        rotation = rot => TextureBuffer
        texturecoordinates = nothing
        shading = true
    end

    @gen_defaults! data begin
        color_map = nothing => Texture
        color_norm = nothing
        intensity = nothing
        image = nothing
        color = nothing => to_meshcolor
        vertex_color = Vec4f(1)
        matcap = nothing => Texture
        fetch_pixel = false
        uv_scale = Vec2f(1)

        instances = const_lift(length, position)
        shading = true
        transparency = false
        shader = GLVisualizeShader(
            "util.vert", "particles.vert", "mesh.frag", "fragment_output.frag",
            view = Dict(
                "position_calc" => position_calc(position, nothing, nothing, nothing, TextureBuffer),
                "light_calc" => light_calc(shading),
                "buffers" => output_buffers(to_value(transparency)),
                "buffer_writes" => output_buffer_writes(to_value(transparency))
            )
        )
    end
    if !isnothing(Makie.to_value(intensity))
        data[:intensity] = intensity_convert_tex(intensity, position)
        data[:len] = const_lift(length, position)
    end
    return assemble_shader(data)
end

"""
This is the most primitive particle system, which uses simple points as primitives.
This is supposed to be the fastest way of displaying particles!
"""
function draw_pixel_scatter(position::VectorTypes, data::Dict)
    @gen_defaults! data begin
        vertex       = position => GLBuffer
        color_map    = nothing  => Texture
        color        = (color_map === nothing ? default(RGBA{Float32}, s) : nothing) => GLBuffer
        color_norm   = nothing
        scale        = 2f0
        transparency = false
        shader       = GLVisualizeShader(
            "fragment_output.frag", "dots.vert", "dots.frag",
            view = Dict(
                "buffers" => output_buffers(to_value(transparency)),
                "buffer_writes" => output_buffer_writes(to_value(transparency))
            )
        )
        gl_primitive = GL_POINTS
    end
    data[:prerender] = PointSizeRender(data[:scale])
    return assemble_shader(data)
end

function draw_scatter(
        p::Tuple{TOrSignal{Matrix{C}}, VectorTypes{P}}, data::Dict
    ) where {C <: Colorant, P <: Point}
    data[:image] = p[1] # we don't want this to be overwritten by user
    @gen_defaults! data begin
        scale = lift(x-> Vec2f(size(x)), p[1])
        offset = Vec2f(0)
    end
    draw_scatter((RECTANGLE, p[2]), data)
end

function draw_scatter(
        p::Tuple{VectorTypes{Matrix{C}}, VectorTypes{P}}, data::Dict
    ) where {C <: Colorant, P <: Point}
    images = to_value(p[1])
    isempty(images) && error("Can not display empty vector of images as primitive")
    sizes = map(size, images)
    if !all(x-> x == sizes[1], sizes) # if differently sized
        # create texture atlas
        maxdims = sum(map(Vec{2, Int}, sizes))
        rectangles = map(x->Rect2(0, 0, x...), sizes)
        rpack = RectanglePacker(Rect2(0, 0, maxdims...))
        uv_coordinates = [push!(rpack, rect).area for rect in rectangles]
        max_xy = maximum(maximum.(uv_coordinates))
        texture_atlas = Texture(C, (max_xy...,))
        for (area, img) in zip(uv_coordinates, images)
            texture_atlas[area] = img #transfer to texture atlas
        end
        data[:uv_offset_width] = map(uv_coordinates) do uv
            m = max_xy .- 1
            mini = reverse((minimum(uv)) ./ m)
            maxi = reverse((maximum(uv) .- 1) ./ m)
            return Vec4f(mini..., maxi...)
        end
        images = texture_atlas
    end
    data[:image] = images # we don't want this to be overwritten by user
    @gen_defaults! data begin
        shape = RECTANGLE
        offset = Vec2f(0)
    end
    return draw_scatter(p, data)
end

"""
Main assemble functions for scatter particles.
Sprites are anything like distance fields, images and simple geometries
"""
function draw_scatter((marker, position), data)
    rot = get!(data, :rotation, Vec4f(0, 0, 0, 1))
    rot = vec2quaternion(rot)
    delete!(data, :rotation)
    # Rescale to include glyph padding and shape
    if isa(to_value(marker), Union{AbstractString, Char})
        scale = data[:scale]
        font = get(data, :font, Observable(Makie.defaultfont()))
        offset = get(data, :offset, Observable(Vec2f(0)))

        # The same scaling that needs to be applied to scale also needs to apply
        # to offset.
        data[:offset] = map(rescale_glyph, marker, font, offset)
        data[:scale] = map(rescale_glyph, marker, font, scale)
    end

    @gen_defaults! data begin
        shape       = const_lift(x-> Int32(primitive_shape(x)), marker)
        position    = position => GLBuffer
        scale       = const_lift(primitive_scale, marker) => GLBuffer
        rotation    = rot => GLBuffer
        image       = nothing => Texture
    end

    @gen_defaults! data begin
        offset          = primitive_offset(marker, scale) => GLBuffer
        intensity       = nothing => GLBuffer
        color_map       = nothing => Texture
        color_norm      = nothing
        color           = nothing => GLBuffer

        glow_color      = RGBA{Float32}(0,0,0,0) => GLBuffer
        stroke_color    = RGBA{Float32}(0,0,0,0) => GLBuffer
        stroke_width    = 0f0
        glow_width      = 0f0
        uv_offset_width = const_lift(primitive_uv_offset_width, marker) => GLBuffer

        distancefield   = primitive_distancefield(marker) => Texture
        indices         = const_lift(length, position) => to_index_buffer
        # rotation and billboard don't go along
        billboard        = rotation == Vec4f(0,0,0,1) => "if `billboard` == true, particles will always face camera"
        fxaa             = false
        transparency     = false
        shader           = GLVisualizeShader(
            "fragment_output.frag", "util.vert", "sprites.geom",
            "sprites.vert", "distance_shape.frag",
            view = Dict(
                "position_calc" => position_calc(position, nothing, nothing, nothing, GLBuffer),
                "buffers" => output_buffers(to_value(transparency)),
                "buffer_writes" => output_buffer_writes(to_value(transparency))
            )
        )
        scale_primitive = true
        gl_primitive = GL_POINTS
    end
    # Exception for intensity, to make it possible to handle intensity with a
    # different length compared to position. Intensities will be interpolated in that case
    data[:intensity] = intensity_convert(intensity, position)
    data[:len] = const_lift(length, position)
    return assemble_shader(data)
end

@specialize


text:
# This file was generated, do not modify it. # hide
by(iris, :Species, :PetalLength => mean)


text:
#= Frobenius Dot product
   ---------------------

1st method of calculation

a) element-wise multiplication between 2 matrices (which have the same size)

b) Sum all elements
=#
Frb1(A,B) = reduce(+, (A .* B))

#= 2nd method
   ----------
a) Vectorise both matrices
b) Compute vector dot product

Vectorising a Matrix
--------------------

Create a vector from a matrix

e.g. [a c e; b d f]

goes to [a b c d e f]' (column-wise concatenation)
=#
function Frb2(A,B)
  va = A[:]
  vb = B[:]

  return va .* vb
end

#=
3rd way of Frobenius (most efficient)
-------------------------------------

take the trace of A' * B
=#
Frb3(A,B) = tr(A' * B)

#=
Frobenius Norm
--------------

Frobenius dot product of matrix with itself
=#

Fnorm(A) = sqrt(Frb3(A,A))

# Implementation in code
# ----------------------
using LinearAlgebra

m = 9; n = 4;

A = randn(m,n);
B = randn(m,n);

Av = A[:] # reshapes into a (mxn) x 1 column vector

Bv = B[:] # same for B

## Ways of calculating Frobenius Dot Product

# 1) Sum of Hadamard multiplication
Frb1(A,B) = reduce(+, (A .* B))

# 2) Dot product of vectorised matrices
Frb2(A,B) = A[:] .* B[:]

# 3) Trace of transposed A matrix multiplied by B matrix
Frb3(A,B) = tr(A' * B)

# Frobenius Norm aka Matrix Norm:
# 1) square root of trace of (A' * A)
Fnorm1(A) = sqrt(Frb3(A,A))

# 2) another way of doing Frobenius norm in Julia
Fnorm2(A) = norm(A, 2)

max_stars_repo_path: linear-algebra-code-challenges/Lecture 59 - Frobenius Dot Product.jl

max_stars_repo_name: jawuku/julia

max_issues_repo_name: jawuku/julia

max_forks_repo_name: jawuku/julia


text:
# This script runs the KS PDE over ensembles to estimate the time-averaged QoI
include("kuramoto-sivashinsky.jl")
using .kuramoto_sivashinsky

# define integrand functions needed for time averages
function ident(val)
    return val
end
function square(val)
    return val*val
end

# set the parameters
tp = Float64             # floating point type to use
num_nodes = 1023         # number of (interior) nodes to use
dt = 0.005
Spin = convert(tp, 1000) # spin-up period
Time = convert(tp, 1000) # statistics gathering period for simulation
#num_spin = 200000        # number of steps to use during spin up
#num_steps = 200000       # number of time steps to use during simulation
num_spin = convert(Int, Spin/dt)
num_steps = convert(Int, Time/dt)
Lx = convert(tp, 128)    # domain size
order = 8                # order of accuracy
num_samples = 10         # number of ensemble samples to use
num_part = 10            # number of partitions that time is broken into
dx = Lx/(num_nodes+1)
x = Array{tp}(dx:dx:Lx-dx)

# loop over the ensemble ensemble samples
u_avg_vals = zeros(num_samples)
u2_avg_vals = zeros(num_samples)
for n = 1:num_samples
    # set initial condition and set-up the KSData structure
    u = 2*rand(num_nodes) .- 1
    ks = kuramoto_sivashinsky.buildKSData(order, num_nodes, tp)

    # spin simulation
    println("Starting spin-up portion of simulation")
    sol = kuramoto_sivashinsky.solveUsingMidpoint(ks, Spin, num_spin, u)

    for p = 1:num_part
        # to avoid running out of memory, we partition the total time
        u = sol[2:end-1,end]
        #clear!(:sol)
        sol = nothing
        println("Starting statistics gathering simulation #",p)
        cputime = @elapsed begin
            sol = kuramoto_sivashinsky.solveUsingMidpoint(ks, Time,
                                                          num_steps, u)
        end
        u_avg_vals[n] += kuramoto_sivashinsky.
            calcSolutionAverage(order, sol, ident)
        u2_avg_vals[n] += kuramoto_sivashinsky.
            calcSolutionAverage(order, sol, square)
    end
    u_avg_vals[n] /= num_part
    u2_avg_vals[n] /= num_part

    println("u averaged = ", u_avg_vals[n])
    println("u^2 averaged = ", u2_avg_vals[n])
end

# display ensemble values
println("ensemble value of average u   = ", sum(u_avg_vals)/num_samples)
println("ensemble value of average u^2 = ", sum(u2_avg_vals)/num_samples)

# write the values to file
file = open("benchmark_ks_data.dat", "w")
for i = 1:num_samples
    print(file, u_avg_vals[i], " ")
end
println(file)
for i = 1:num_samples
    print(file, u2_avg_vals[i], " ")
end
println(file)
close(file)

max_stars_repo_name: jehicken/high-order-and-chaos

max_issues_repo_name: jehicken/high-order-and-chaos

max_forks_repo_name: jehicken/high-order-and-chaos


text:
x0 = ones(6)
c(x) = [sum(x)]
G(x) = [x[1];x[3]]
H(x) = [x[2];x[4]]
admpcc = ADMPCCModel(rosenbrock,  x0,
                     lvar = fill(-10.0,size(x0)), uvar = fill(10.0,size(x0)),
                     y0 = [0.0], c = c, lcon = [-Inf], ucon = [6.],
                     G = G, H = H, lccG = zeros(2), lccH = zeros(2))

@test obj(admpcc, admpcc.meta.x0)  == 0.0
@test grad(admpcc, admpcc.meta.x0) == zeros(6)
@test issparse(hess(admpcc, admpcc.meta.x0))
@test issparse(jac(admpcc, admpcc.meta.x0))
@test cons(admpcc, admpcc.meta.x0) == [6.0, 1.0, 1.0, 1.0, 1.0]
@test viol(admpcc, admpcc.meta.x0) == [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0]

@test hessG(admpcc, admpcc.meta.x0, zeros(2)) == zeros(6,6)
@test hGprod(admpcc, admpcc.meta.x0, zeros(2), ones(6)) == zeros(6)
@test issparse(jacG(admpcc, admpcc.meta.x0))
@test consG(admpcc, admpcc.meta.x0) == ones(2)

@test hessH(admpcc, admpcc.meta.x0, zeros(2)) == zeros(6,6)
@test hHprod(admpcc, admpcc.meta.x0, zeros(2), ones(6)) == zeros(6)
@test sum(jacH(admpcc, admpcc.meta.x0)) == 2
@test consH(admpcc, admpcc.meta.x0) == ones(2)

@test jprod(admpcc, admpcc.meta.x0, zeros(6)) == zeros(5)
@test jtprod(admpcc, admpcc.meta.x0, zeros(5)) == zeros(6)

@test size(hess(admpcc, admpcc.meta.x0)) == (6,6)
@test hess(admpcc, admpcc.meta.x0, zeros(5), obj_weight = 0.0) == zeros(6,6)
@test hprod(admpcc, admpcc.meta.x0, zeros(6)) == zeros(6)
@test hprod(admpcc, admpcc.meta.x0, zeros(5), ones(6), obj_weight = 0.0) == zeros(6)

@test hessH_structure(admpcc) == hessG_structure(admpcc)
@test length(hessH_coord(admpcc, admpcc.meta.x0, zeros(2))) == admpcc.meta.nvar * (admpcc.meta.nvar + 1) / 2
@test length(hessG_coord(admpcc, admpcc.meta.x0, zeros(2))) == admpcc.meta.nvar * (admpcc.meta.nvar + 1) / 2


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""ReplicaSetSpec is the specification of a ReplicaSet.

    IoK8sApiExtensionsV1beta1ReplicaSetSpec(;
        minReadySeconds=nothing,
        replicas=nothing,
        selector=nothing,
        template=nothing,
    )

    - minReadySeconds::Int32 : Minimum number of seconds for which a newly created pod should be ready without any of its container crashing, for it to be considered available. Defaults to 0 (pod will be considered available as soon as it is ready)
    - replicas::Int32 : Replicas is the number of desired replicas. This is a pointer to distinguish between explicit zero and unspecified. Defaults to 1. More info: https://kubernetes.io/docs/concepts/workloads/controllers/replicationcontroller/#what-is-a-replicationcontroller
    - selector::IoK8sApimachineryPkgApisMetaV1LabelSelector : Selector is a label query over pods that should match the replica count. If the selector is empty, it is defaulted to the labels present on the pod template. Label keys and values that must match in order to be controlled by this replica set. More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#label-selectors
    - template::IoK8sApiCoreV1PodTemplateSpec : Template is the object that describes the pod that will be created if insufficient replicas are detected. More info: https://kubernetes.io/docs/concepts/workloads/controllers/replicationcontroller#pod-template
"""
mutable struct IoK8sApiExtensionsV1beta1ReplicaSetSpec <: SwaggerModel
    minReadySeconds::Any # spec type: Union{ Nothing, Int32 } # spec name: minReadySeconds
    replicas::Any # spec type: Union{ Nothing, Int32 } # spec name: replicas
    selector::Any # spec type: Union{ Nothing, IoK8sApimachineryPkgApisMetaV1LabelSelector } # spec name: selector
    template::Any # spec type: Union{ Nothing, IoK8sApiCoreV1PodTemplateSpec } # spec name: template

    function IoK8sApiExtensionsV1beta1ReplicaSetSpec(;minReadySeconds=nothing, replicas=nothing, selector=nothing, template=nothing)
        o = new()
        validate_property(IoK8sApiExtensionsV1beta1ReplicaSetSpec, Symbol("minReadySeconds"), minReadySeconds)
        setfield!(o, Symbol("minReadySeconds"), minReadySeconds)
        validate_property(IoK8sApiExtensionsV1beta1ReplicaSetSpec, Symbol("replicas"), replicas)
        setfield!(o, Symbol("replicas"), replicas)
        validate_property(IoK8sApiExtensionsV1beta1ReplicaSetSpec, Symbol("selector"), selector)
        setfield!(o, Symbol("selector"), selector)
        validate_property(IoK8sApiExtensionsV1beta1ReplicaSetSpec, Symbol("template"), template)
        setfield!(o, Symbol("template"), template)
        o
    end
end # type IoK8sApiExtensionsV1beta1ReplicaSetSpec

const _property_map_IoK8sApiExtensionsV1beta1ReplicaSetSpec = Dict{Symbol,Symbol}(Symbol("minReadySeconds")=>Symbol("minReadySeconds"), Symbol("replicas")=>Symbol("replicas"), Symbol("selector")=>Symbol("selector"), Symbol("template")=>Symbol("template"))
const _property_types_IoK8sApiExtensionsV1beta1ReplicaSetSpec = Dict{Symbol,String}(Symbol("minReadySeconds")=>"Int32", Symbol("replicas")=>"Int32", Symbol("selector")=>"IoK8sApimachineryPkgApisMetaV1LabelSelector", Symbol("template")=>"IoK8sApiCoreV1PodTemplateSpec")
Base.propertynames(::Type{ IoK8sApiExtensionsV1beta1ReplicaSetSpec }) = collect(keys(_property_map_IoK8sApiExtensionsV1beta1ReplicaSetSpec))
Swagger.property_type(::Type{ IoK8sApiExtensionsV1beta1ReplicaSetSpec }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiExtensionsV1beta1ReplicaSetSpec[name]))}
Swagger.field_name(::Type{ IoK8sApiExtensionsV1beta1ReplicaSetSpec }, property_name::Symbol) =  _property_map_IoK8sApiExtensionsV1beta1ReplicaSetSpec[property_name]

function check_required(o::IoK8sApiExtensionsV1beta1ReplicaSetSpec)
    true
end

function validate_property(::Type{ IoK8sApiExtensionsV1beta1ReplicaSetSpec }, name::Symbol, val)
end


text:
mutable struct ModPar
  γ :: Float64 # RiskAversion
  β :: Float64 # Discount
  η :: Float64 # Weight on housing
  q :: Float64 # Part. cost
  r :: Float64 # Risk free rate
  σ :: Float64 # Return std.dev
  μ :: Float64 # Return mean
  ψ :: Float64 # Utility shifter KKK
  xstar :: Float64 # Shifter threshold KKK
  ms :: Float64 # Sales cost

  function ModPar(;
    γ = 2.0,
    β = 1.0,
    η = 0.5,
    q = 0.0,
    r = 0.0,
    σ = 0.16,
    μ = 0.02,
    ψ = 0.0,
    xstar = 0.0,
    ms = 0.0,
  )

    new(γ, β, η, q, r, σ, μ, ψ, xstar, ms)
  end
end

struct NumPar
  nx :: Int
  nh :: Int
  na :: Int
  nq :: Int # How many points to use in quadrature
  nα :: Int # α-choice grid density
  nsav :: Int # savings-choice grid density

  xmax :: Float64
  hmax :: Float64

  xgrd :: StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}
  hgrd :: Vector{Float64}
  ygrd :: Vector{Float64}

  πrs :: Vector{Float64}
  rsgrd :: Vector{Float64}

  p1cons :: Bool

  function NumPar(mp::ModPar;
      nx=51,
      nh=2,
      na=2,
      nq=11,
      nα=21,
      nsav=51,
      xmax=8.0,
      hmax=1.0,
      ygrd=[1.0,1.0],
      p1cons=true,
      )

    xgrd = range(0.1,stop=xmax,length=nx)
    if nh == 1
      hgrd = [1.0]
    elseif nh == 2
      hgrd = [1.0, 2.0]
    else
      throw("No valid housing grid")
    end

    # Use expectation package
    E = expectation(Normal(mp.μ,mp.σ),n=nq)
    rsgrd=exp.(nodes(E)).-1 # Return grid
    πrs=weights(E) # PMF
    @assert sum(πrs)≈1 # Check that PMF sums to one

    new(nx, nh, na, nq, nα, nsav, xmax, hmax, xgrd, hgrd, ygrd, πrs, rsgrd, p1cons)
  end
end

mutable struct ModelSolution
  V :: Array{Float64,3}
  α :: Array{Float64,2}
  b :: Array{Float64,2}
  s :: Array{Float64,2}
  h :: Array{Float64,3}
  c :: Array{Float64,3}
  mp :: ModPar
  np :: NumPar

  function ModelSolution(mp::ModPar,np::NumPar)
    V = fill(-Inf64,(np.nx,np.nh,np.na))
    α = fill(-Inf64,(np.nx,np.nh))
    b = fill(-Inf64,(np.nx,np.nh))
    s = fill(-Inf64,(np.nx,np.nh))
    h = fill(-Inf64,(np.nx,np.nh,np.na))
    c = fill(-Inf64,(np.nx,np.nh,np.na))

    new(V, α, b, s, h, c, mp, np)
  end
end


text:
using Test

@info "Starting tests..."
include("Utils.jl")
include("Frames.jl")
include("Signals.jl")
include("Messages.jl")
include("Databases.jl")
include("Decode.jl")