using SeisIO, JLD2, Distributed, Sockets

export sac2jld2, sac2jld2_par
#

# Converting sac data to JLD2 format. There are options to save to JLD2 format
# as single file or separate files for each sac file.
"""
Key features:
1. Can be used to convert individual SAC file or a group/list of SAC files.
2. Handles the situation when multiple data files were download for the same station
    and channel. Those files will be merged into one single SeisData object with
    multiple SeisChannel entries.
3. All SAC headers are kept and saved in :misc filed of the SeisData object.

"""

"""

    sac2jld2(infile)

Saves individual sac file (in) to JLD2 file.
The output file name will be the same as the input, with new extension *.jld2.

"""

function sac2jld2(infile::String)
    print("Same file, new extension of jld2.\n")
    print("Converting SAC file: [ ",infile," ] ... \n")
    outfile = join([infile[1:end-3],"jld2"]); #trim the file name to remove .sac
    sacin = SeisIO.read_data("sac",infile,full=true);
    sacin.misc[1]["dlerror"] = 0

    stationlist = String[];
    push!(stationlist,sacin.id[1])  #pass the 1st element of sacin.id[] to the station list
    # stationinfo = Dict(["stationlist" => stationlist,"stationmethod" => "sac2jld2","stationsrc" => "converted"])
    stimetemp = sacin.t[1][3]*1e-6; #convert mseconds to seconds;
    saclength = sacin.t[1][2]/sacin.fs; #convert from number of samples to time duration.
    starttime = u2d(stimetemp[1]); #convert union time to calendar time;
    endtime = u2d(stimetemp[1] + saclength[1]);
    DLtimestamplist = String[];
    push!(DLtimestamplist,replace(string(starttime),"-" => "."));

    #save to jld2
    stemp = joinpath(DLtimestamplist[1],stationlist[1]);
    file = jldopen(outfile, "w");
    file["info/DLtimestamplist"] = DLtimestamplist;
    file["info/stationlist"] = stationlist;
    file["info/starttime"]   = string(starttime)
    file["info/endtime"]     = string(endtime)
    file["info/DL_time_unit"]= saclength[1]; #length of the data in time (duration), not number of samples

    #save the waveform data as SeisData format.
    file[stemp] = SeisData(sacin);

    JLD2.close(file);

    print("Saved JLD2 file to: ", outfile, "\n");
end

"""

    sac2jld2(infile,outfile)

Saves individual sac file (in) to JLD2 file (out), which is explicitly specified.

"""

function sac2jld2(infile::String,outfile::String)
    # print("I convert SAC file to JLD2 file.\n")
    print("Converting SAC file: [ ",infile," ] ... \n")
    sacin = SeisIO.read_data("sac",infile,full=true);
    sacin.misc[1]["dlerror"] = 0

    stationlist = String[];
    push!(stationlist,sacin.id[1])  #pass the 1st element of sacin.id[] to the station list
    # stationinfo = Dict(["stationlist" => stationlist,"stationmethod" => "sac2jld2","stationsrc" => "converted"])
    stimetemp = sacin.t[1][3]*1e-6; #convert mseconds to seconds;
    saclength = sacin.t[1][2]/sacin.fs; #convert from number of samples to time duration.
    starttime = u2d(stimetemp[1]); #convert union time to calendar time;
    endtime = u2d(stimetemp[1] + saclength[1]);
    DLtimestamplist = String[];
    push!(DLtimestamplist,replace(string(starttime),"-" => "."));

    #save to jld2
    stemp = joinpath(DLtimestamplist[1],stationlist[1]);
    file = jldopen(outfile, "w");
    file["info/DLtimestamplist"] = DLtimestamplist;
    file["info/stationlist"] = stationlist;
    file["info/starttime"]   = string(starttime)
    file["info/endtime"]     = string(endtime)
    file["info/DL_time_unit"]= saclength[1]; #length of the data in time (duration), not number of samples

    #save the waveform data as SeisData format.
    file[stemp] = SeisData(sacin);

    JLD2.close(file);

    print("Saved JLD2 file to: ", outfile, "\n");
end

"""

    sac2jld2(filelist,timestamp,outfile,verbose=false)

Saves sac files based on a filelist generated using ls() in Julia to
JLD2 file (out), which is explicitly specified.

    Arguments:
    filelist: An String Array including all sacfiles with the same tiemstamp, e.g., same day
    tiemstamp: this must be specified for the sac fiels, assuming they all share
                the same time stamp. Suggested format: 2019.06.27T00.00.00
    outfile: output file name for the JLD2 file. All data fiels will be readin and
                saved to a single JLD2 file.
    verbose: if true, the code will print out more messages. Default is false.
"""

function sac2jld2(filelist::Array{String,1},timestamp::String,outfile::String,verbose::Bool=false)
    # print("I convert and pack all SAC files to a JLD2 file.\n")

    #read the filenames from the infilelist into a strinng array
    # filelist = readfilelist(infilelist)

    stationlist = String[];
    DLtimestamplist = String[];
    push!(DLtimestamplist,timestamp); #use the specified timestamp for the group.

    #get some intial information from the first sac file.
    sacin0 = SeisIO.read_data("sac",filelist[1]);
    stimemin = sacin0.t[1][3]*1e-6; #initial starttime in seconds;
    ftemp = sacin0.t[1][2]/sacin0.fs; #convert from number of samples to time duration.
    saclengthmax = ftemp[1]
    etimemax = stimemin + saclengthmax;

    file = jldopen(outfile, "w");
    file["info/DLtimestamplist"] = DLtimestamplist;
    stemp_pre = "-"
    count = 0
    # stationlisttemp = String[];
    #To deal with multiple channel issue, I am doing the loop twice.
    # 1. The first loop with collect multiple channels into one SeisData object.
    # 2. The second loop will work through the actual number of SeisData objects, counted by `count`
    Sall = SeisData[];
    Str_all = String[];

    #in the loop for each sac file. we will find the minimum starttime, maximum endtime, and maximum saclength.
    for infile = filelist
        println("Converting SAC file: [ ",infile," ]")

        sacin = SeisIO.read_data("sac",infile,full=true);
        sacin.misc[1]["dlerror"] = 0
        push!(stationlist,sacin.id[1])  #pass the 1st element of sacin.id[] to the station list
        # stationinfo = Dict(["stationlist" => stationlist,"stationmethod" => "sac2jld2","stationsrc" => "converted"])
        stimetemp = sacin.t[1][3]*1e-6; #convert mseconds to seconds;
        if stimetemp < stimemin
            stimemin = stimetemp
        end
        saclength = sacin.t[1][2]/sacin.fs; #convert from number of samples to time duration.
        if saclength[1] > saclengthmax
            saclengthmax = saclength[1]
        end
        etimetemp = stimetemp + saclength[1]
        if etimetemp > etimemax
            etimemax = etimetemp
        end

        #save to jld2
        stemp = joinpath(timestamp,sacin.id[1]);
        if stemp == stemp_pre
            println(stemp," exists. Appending to form multiple channels.")
            append!(Sall[count],SeisData(sacin))
        else
            push!(Sall,SeisData(sacin))
            push!(Str_all,stemp)
            count += 1
            # file[stemp] = SeisData(sacin);
        end
        stemp_pre = stemp
    end
    starttime = u2d(stimemin); #convert union time to calendar time;
    endtime = u2d(etimemax);
    #the following info dat is saved after running through the whole group.
    file["info/stationlist"] = unique(stationlist);
    file["info/starttime"]   = string(starttime)
    file["info/endtime"]     = string(endtime)
    file["info/DL_time_unit"]= saclengthmax; #length of the data in time (duration), not number of samples

    #2. start to save the data to jld2
    for i = 1:count
        file[Str_all[i]] = Sall[i]
    end #end of loop for all sac files within one group/directory

    JLD2.close(file);

    print("Saved JLD2 file to: ", outfile, "\n");
end

"""

    sac2jld2(filelist,timestamplist,outfile,verbose=false)

Saves sac files based on a filelist generated using ls() in Julia to
JLD2 file (out), which is explicitly specified. This function assumes each file in
filelist is assigned a timestamp in timestamplist.

    Arguments:
    filelist: An String Array including all sacfiles with the same tiemstamp, e.g., same day
    tiemstamplist: this must be specified for all the sac files in filelist. This list has to be
                the same length as filelist.
    outfile: output file name for the JLD2 file. All data fiels will be readin and
                saved to a single JLD2 file.
    verbose: if true, the code will print out more messages. Default is false.
"""

function sac2jld2(filelist::Array{String,1};timestamplist::Array{String,1},outfile::String,verbose::Bool=false)
    if length(filelist) != length(timestamplist)
        error("filelist and timestamplist must be the same length!")
    end
    stationlist = String[];

    #get some intial information from the first sac file.
    sacin0 = SeisIO.read_data("sac",filelist[1]);
    stimemin = sacin0.t[1][3]*1e-6; #initial starttime in seconds;
    ftemp = sacin0.t[1][2]/sacin0.fs; #convert from number of samples to time duration.
    saclengthmax = ftemp[1]
    etimemax = stimemin + saclengthmax;

    file = jldopen(outfile, "w");
    file["info/DLtimestamplist"] = unique(timestamplist);
    stemp_pre = "-"
    count = 0
    # stationlisttemp = String[];
    #To deal with multiple channel issue, I am doing the loop twice.
    # 1. The first loop with collect multiple channels into one SeisData object.
    # 2. The second loop will work through the actual number of SeisData objects, counted by `count`
    Sall = SeisData[];
    Str_all = String[];

    #in the loop for each sac file. we will find the minimum starttime, maximum endtime, and maximum saclength.
    for (infile,ts) in zip(filelist,timestamplist)
        println("Converting SAC file: [ ",infile," ]")

        sacin = SeisIO.read_data("sac",infile,full=true);
        sacin.misc[1]["dlerror"] = 0
        push!(stationlist,sacin.id[1])  #pass the 1st element of sacin.id[] to the station list
        # stationinfo = Dict(["stationlist" => stationlist,"stationmethod" => "sac2jld2","stationsrc" => "converted"])
        stimetemp = sacin.t[1][3]*1e-6; #convert mseconds to seconds;
        if stimetemp < stimemin
            stimemin = stimetemp
        end
        saclength = sacin.t[1][2]/sacin.fs; #convert from number of samples to time duration.
        if saclength[1] > saclengthmax
            saclengthmax = saclength[1]
        end
        etimetemp = stimetemp + saclength[1]
        if etimetemp > etimemax
            etimemax = etimetemp
        end

        #save to jld2
        stemp = joinpath(ts,sacin.id[1]);
        if stemp == stemp_pre
            println(stemp," exists. Appending to form multiple channels.")
            append!(Sall[count],SeisData(sacin))
        else
            push!(Sall,SeisData(sacin))
            push!(Str_all,stemp)
            count += 1
            # file[stemp] = SeisData(sacin);
        end
        stemp_pre = stemp
    end
    starttime = u2d(stimemin); #convert union time to calendar time;
    endtime = u2d(etimemax);
    #the following info dat is saved after running through the whole group.
    file["info/stationlist"] = unique(stationlist);
    file["info/starttime"]   = string(starttime)
    file["info/endtime"]     = string(endtime)
    file["info/DL_time_unit"]= saclengthmax; #length of the data in time (duration), not number of samples

    #2. start to save the data to jld2
    for i = 1:count
        file[Str_all[i]] = Sall[i]
    end #end of loop for all sac files within one group/directory

    JLD2.close(file);

    print("Saved JLD2 file to: ", outfile, "\n");
end

"""

    sac2jld2(fhandle::JLDfile, sacdir::String,timestamp::String,verbose::Bool=false)

Saves sac files based in sacdir to JLD2 file handle (fhandle), tagged as timestamp.

    Arguments:
    fhandle: A JLD2 file handle, generated using jldopen(), with write permission.
    sacdir: directory containing the sac files need to convert.
    tiemstamp: this must be specified for the sac fiels, assuming they all share
                the same time stamp. Suggested format: 2019.06.27T00.00.00
    verbose: if true, the code will print out more messages. Default is false.
"""

function sac2jld2(fhandle::JLD2.JLDFile, sacdir::String,timestamp::String,verbose::Bool=false)
    # print("I convert and pack all SAC files to a JLD2 file.\n")

    filelist=ls(joinpath(sacdir,"*.sac"));
    println("Converting SAC dir: ",sacdir,"=> ",length(filelist)," files.")

    stemp_pre = "-"
    count = 0
    #To deal with multiple channel issue, I am doing the loop twice.
    # 1. The first loop with collect multiple channels into one SeisData object.
    # 2. The second loop will work through the actual number of SeisData objects, counted by `count`
    Sall = SeisData[];
    Str_all = String[];
    #in the loop for each sac file. we will find the minimum starttime, maximum endtime, and maximum saclength.
    for infile = filelist
        if verbose == true
            print("Converting SAC file: [ ",infile," ] ...\n")
        end

        sacin = SeisIO.read_data("sac",infile,full=true);
        sacin.misc[1]["dlerror"] = 0

        #save to jld2
        stemp = joinpath(timestamp,sacin.id[1]);
        println(stemp)
        if stemp == stemp_pre
            if verbose == true
                println(stemp," exists. Appending to form multiple channels.")
            end
            append!(Sall[count],SeisData(sacin))
        else
            push!(Sall,SeisData(sacin))
            push!(Str_all,stemp)
            count += 1
            # file[stemp] = SeisData(sacin);
        end
        stemp_pre = stemp
    end

    #2. start to save the data to jld2
    for i = 1:count
        fhandle[Str_all[i]] = Sall[i]
    end #end of loop for all sac files within one group/directory
end

"""

    sac2jld2(sacdirlist,timestamplist,outfile,verbose=false)

Searches for all sac files scanning throgh the sacdirlist (Array{String,1}). The sac files
    in each sac directory are assigned to one time stamp based on timestamplist (Array{String,1}).
    The two Arrays must be equal in size. The list can be generated from a file list (using ls in shell),
    using a utility function in this package: readfilelist(infilelist::String).

All sac files will be saved into one single JLD2 file (out), which is explicitly specified as outfile.

    Arguments:
    sacdirlist: An array of strings storing the directories grouped by time stamp, e.g., one day.
    timestamplist: An array of strings storing the time stamps assigning to each SAC file group.
    outfile: output file name for the JLD2 file. All data fiels will be readin and
                saved to a single JLD2 file.
    verbose: if true, the code will print out more messages. Default is false.
"""

function sac2jld2(sacdirlist::Array{String,1},timestamplist::Array{String,1},outfile::String,verbose::Bool=false)
    if length(sacdirlist) != length(timestamplist)
        error("sacdirlist and timestamplist must be the same length!")
    end

    stationlist = String[];

    # saclengthmax = 0; #initiate the length as zero, will find the maximum value after looping through the files.

    file = jldopen(outfile, "w");
    file["info/DLtimestamplist"] = timestamplist;
    file["info/starttime"] = timestamplist[1];
    #here we use the first value in timestamplist as the starttime for a group of multiple timestamp data.
    file["info/endtime"] = timestamplist[end]; #similar to starttime, here we use the last value in timestamplist as the endtime.

    #in the loop for each sac file in each SAC directory. we will find the maximum saclength.
    # countdir=1;
    t1all=0
    t2all=0
    for (sacdir,ts) in zip(sacdirlist,timestamplist)

        t0=time()
        filelist=ls(joinpath(sacdir,"*.sac"));

        print("Converting for directory: [ ",sacdir," ] ... \n")

        count = 0
        # stationlisttemp = String[];
        #To deal with multiple channel issue, I am doing the loop twice.
        # 1. The first loop with collect multiple channels into one SeisData object.
        # 2. The second loop will work through the actual number of SeisData objects, counted by `count`
        Sall = SeisData[];
        Str_all = String[];
        stemp_pre = "-"
        for infile = filelist
            # t1=@elapsed sacin = SeisIO.read_data("sac",infile,full=true);
            sacin = SeisIO.read_data("sac",infile,full=true);
            sacin.misc[1]["dlerror"] = 0

            if sacdir == sacdirlist[1] && infile == filelist[1]
                # println(sacdir)
                file["info/DL_time_unit"]= sacin.t[1][2]/sacin.fs; #length of the data in time (duration), not number of samples;
                #convert from number of samples to time duration.
            end

            stemp = joinpath(ts,sacin.id[1]);
            # if verbose == true
            #     println("00 current: ",stemp)
            #     println("11 previous: ",stemp_pre)
            # end
            # t2=@elapsed if stemp == stemp_pre
            if stemp == stemp_pre
                if verbose == true
                    println(stemp," exists. Append to form multiple channels.")
                end
                append!(Sall[count],SeisData(sacin))
            else
                push!(Sall,SeisData(sacin))
                push!(Str_all,stemp)
                count += 1
                # file[stemp] = SeisData(sacin);
            end

            stemp_pre = stemp;

        end #end of loop for all sac files within one group/directory

        #2. start to save the data to jld2
        for i = 1:count
            file[Str_all[i]] = Sall[i]
        end #end of loop for all sac files within one group/directory

        stationlist = unique(vcat(stationlist,keys(file[ts])))
        if verbose == true
            print("                ----------------> ",length(filelist)," sac files, time used: ",time() - t0,"\n")
        end

    end #end of loop for all SAC directories.

    # println([t1all, t2all])
    #the following info dat is saved after running through the whole group.
    file["info/stationlist"] = stationlist;

    JLD2.close(file);

    print("Saved JLD2 file to: ", outfile, "\n");
end

"""
    sac2jld2_par(sacdirlist,timestamplist,outfile,verbose=false)

Searches for all sac files scanning throgh the sacdirlist (Array{String,1}). The sac files
    in each sac directory are assigned to one time stamp based on timestamplist (Array{String,1}).
    The two Arrays must be equal in size. The list can be generated from a file list (using ls in shell),
    using a utility function in this package: readfilelist(infilelist::String).

All sac files will be saved into one single JLD2 file (out), which is explicitly specified as outfile.

    Arguments:
    sacdirlist: An array of strings storing the directories grouped by time stamp, e.g., one day.
    timestamplist: An array of strings storing the time stamps assigning to each SAC file group.
    outfile: output file name for the JLD2 file. All data fiels will be readin and
                saved to a single JLD2 file.
    verbose: if true, the code will print out more messages. Default is false.
"""

# For parallel processing, the "info/DL_time_unit" is based on ONLY the first sac file.
function sac2jld2_par(sacdirlist::Array{String,1},timestamplist::Array{String,1},outfile::String,verbose::Bool=false)
    if length(sacdirlist) != length(timestamplist)
        error("sacdirlist and timestamplist must be the same length!")
    end

    stationlist = String[];
    #length of the data in time (duration), not number of samples
    file = jldopen(outfile,"w")

    file["info/DLtimestamplist"] = timestamplist;
    file["info/starttime"] = timestamplist[1];
    #here we use the first value in timestamplist as the starttime for a group of multiple timestamp data.
    file["info/endtime"] = timestamplist[end]; #similar to starttime, here we use the last value in timestamplist as the endtime.

    # function testread(x, filelist)
    #     #println(myid())
    #     SeisIO.read_data("sac",filelist[x],full=true)
    # end

    t1all = 0
    t2all = 0
    t3all = 0

    for (sacdir,ts) in zip(sacdirlist,timestamplist)
        t0=time()
        filelist=ls(joinpath(sacdir,"*.sac"));

        print("Converting for directory: [ ",sacdir," ] ... \n")

        S=pmap(x->SeisIO.read_data("sac",x,full=true),filelist)

        if ts == timestamplist[1]
            file["info/DL_time_unit"]= S[1].t[1][2]/S[1].fs;
        end
        count = 0
        # stationlisttemp = String[];
        #To deal with multiple channel issue, I am doing the loop twice.
        # 1. The first loop with collect multiple channels into one SeisData object.
        # 2. The second loop will work through the actual number of SeisData objects, counted by `count`
        Sall = SeisData[];
        Str_all = String[];

        stemp_pre = "-"
        for Sdata = S
            Sdata.misc[1]["dlerror"] = 0
            stemp = joinpath(ts,Sdata.id[1])
            if stemp == stemp_pre
                println(stemp," exists. Append to form multiple channels.")
                append!(Sall[count],SeisData(Sdata))
            else
                push!(Sall,SeisData(Sdata))
                push!(Str_all,stemp)
                count += 1
                # file[stemp] = SeisData(sacin);
            end
            stemp_pre = stemp;
        end

        #2. start to save the data to jld2
        for i = 1:count
            file[Str_all[i]] = Sall[i]
        end #end of loop for all sac files within one group/directory

        stationlist = unique(vcat(stationlist,keys(file[ts])))
        if verbose == true
            print("                ----------------> ",length(filelist)," sac files, time used: ",time() - t0,"\n")
        end

        # t1all += t1
    end #end of loop for all SAC directories.

    # println([t1all, t2all, t3all])

    #the following info dat is saved after running through the whole group.
    file["info/stationlist"] = stationlist;

    JLD2.close(file);

    print("Saved JLD2 file to: ", outfile, "\n");
    rmprocs(workers())
end


using PGFPlotsX
import PGFPlotsX: Plot, PlotInc, Options, Axis

for F in (:Expansion,)
    @eval begin
        Axis(F::$F, trailing...; opts...) =
            @pgf Axis({ymode="log"}, Plot(F, trailing...; opts...))
    end
end

for plot in (:Plot, :PlotInc)
    _plot = Meta.parse("_"*string(plot))
    for EXP in (:Expansion,)
        @eval begin
            $(plot)(F::$EXP; opts...) =
                $(plot)(Options(), F; opts...)
            $(plot)(f::Function, F::$EXP; opts...) =
                $(plot)(Options(), f, F; opts...)
            $(plot)(F::$EXP, f::Function; opts...) =
                $(plot)(Options(), f, F)
            $(plot)(F::$EXP, f::$EXP; opts...) =
                $(plot)(Options(), f, F; opts...)
        end
    end
    @eval begin
        function $(plot)(options::Options, F::Expansion; n=200, plot_complex=false, opts...)
            grid = plotgrid(dictionary(F), n)
            if plot_complex
                vals = F(grid)
                options = @pgf {options..., no_markers}
                $(plot)(options, Table([grid, BasisFunctions.postprocess(dictionary(F), grid, real.(vals))])),
                    $(plot)(options, Table([grid, BasisFunctions.postprocess(dictionary(F), grid, imag.(vals))]))
            else
                vals = real.(F(grid))
                options = @pgf {options..., no_markers}
                $(plot)(options, Table([grid, BasisFunctions.postprocess(dictionary(F), grid, vals)]))
            end
        end

        $(plot)(options::Options, F::Expansion, f::Function; opts...) =
            $(_plot)(options, f, F; opts...)

        $(plot)(options::Options, f::Function, F::Expansion; opts...) =
            $(_plot)(options, f, F; opts...)
        $(plot)(options::Options, f::Expansion, F::Expansion; opts...) =
            $(_plot)(options, f, F; opts...)

        function $(_plot)(options::Options, f, F::Expansion; plot_complex=false, n=200, opts...)
            grid = plotgrid(dictionary(F), n)
            vals = abs.(f.(grid) - F.(grid))
            options = @pgf {options..., no_markers}
            $(plot)(options, Table([grid, BasisFunctions.postprocess(dictionary(F), grid, vals)]))
        end
    end
end


abstract type FluxLayer end

layertype(l) = NaiveNASlib.shapetrait(l)
layer(l) = l

# Types for layers with parameters (e.g. weights and biases) and with similar handling
# w.r.t what shape of parameters means in terms of number of inputs and number of outputs
abstract type FluxParLayer <: FluxLayer end

abstract type Flux2D <: FluxParLayer end
struct GenericFlux2D <: Flux2D end
struct FluxDense <: Flux2D end
NaiveNASlib.shapetrait(::Dense) = FluxDense()

# Might be a Flux2D, but not exactly the same. Also want to leave the door open to when/if they accept 3D input too
# Type hierarchies are hard :/
abstract type FluxRecurrent <:FluxParLayer end
struct GenericFluxRecurrent <: FluxRecurrent end
struct FluxRnn <: FluxRecurrent end
struct FluxLstm <: FluxRecurrent end
struct FluxGru <: FluxRecurrent end

NaiveNASlib.shapetrait(l::Flux.Recur) = NaiveNASlib.shapetrait(l.cell)
NaiveNASlib.shapetrait(::Flux.RNNCell) = FluxRnn()
NaiveNASlib.shapetrait(::Flux.LSTMCell) = FluxLstm()
NaiveNASlib.shapetrait(::Flux.GRUCell) = FluxGru()

abstract type FluxConvolutional{N} <: FluxParLayer end
struct GenericFluxConvolutional{N} <: FluxConvolutional{N} end
struct FluxConv{N} <: FluxConvolutional{N} end
struct FluxConvTranspose{N}  <: FluxConvolutional{N} end
struct FluxDepthwiseConv{N} <: FluxConvolutional{N}  end
struct FluxCrossCor{N} <: FluxConvolutional{N}  end
NaiveNASlib.shapetrait(::Conv{N}) where N = FluxConv{N}()
NaiveNASlib.shapetrait(::ConvTranspose{N}) where N = FluxConvTranspose{N}()
NaiveNASlib.shapetrait(::DepthwiseConv{N}) where N = FluxDepthwiseConv{N}()
NaiveNASlib.shapetrait(::CrossCor{N}) where N = FluxCrossCor{N}()

abstract type FluxTransparentLayer <: FluxLayer end
# Invariant layers with parameters, i.e nin == nout always and parameter selection must
# be performed
abstract type FluxParInvLayer <: FluxTransparentLayer end
struct FluxDiagonal <: FluxParInvLayer end
struct FluxLayerNorm <: FluxParInvLayer end
abstract type FluxParNorm <: FluxParInvLayer end
struct FluxBatchNorm <: FluxParNorm end
struct FluxInstanceNorm <: FluxParNorm end
struct FluxGroupNorm <: FluxParNorm end

NaiveNASlib.shapetrait(::Flux.Diagonal) = FluxDiagonal()
NaiveNASlib.shapetrait(::LayerNorm) = FluxLayerNorm()
NaiveNASlib.shapetrait(::BatchNorm) = FluxBatchNorm()
NaiveNASlib.shapetrait(::InstanceNorm) = FluxInstanceNorm()
NaiveNASlib.shapetrait(::GroupNorm) = FluxGroupNorm()

# Transparent layers, i.e nin == nout always and there are no parameters
abstract type FluxNoParLayer <: FluxTransparentLayer end
struct FluxPoolLayer <: FluxNoParLayer end
struct FluxDropOut <: FluxNoParLayer end

layertype(l::MaxPool) = FluxPoolLayer()
layertype(l::MeanPool) = FluxPoolLayer()
layertype(l::Dropout) = FluxDropOut()
layertype(l::AlphaDropout) = FluxDropOut()
layertype(l::GlobalMaxPool) = FluxPoolLayer()
layertype(l::GlobalMeanPool) = FluxPoolLayer()

# Compositions? Might not have any common methods...
# MaxOut, Chain?


# Helper data for CIE observer functions
include("cie_data.jl")

# Linear interpolation in [a, b] where x is in [0,1],
# or coerced to be if not.
function lerp(x, a, b)
    a + (b - a) * max(min(x, one(x)), zero(x))
end

"""
    hex(c)

Print a color as a RGB hex triple, or a transparent paint as an ARGB
hex quadruplet.
"""
function hex(c::RGB)
    @sprintf("%02X%02X%02X",
             round(Int, lerp(c.r, 0.0, 255.0)),
             round(Int, lerp(c.g, 0.0, 255.0)),
             round(Int, lerp(c.b, 0.0, 255.0)))
end
function hex(c::ARGB)
    @sprintf("%02X%02X%02X%02X",
             round(Int, lerp(alpha(c), 0.0, 255.0)),
             round(Int, lerp(red(c), 0.0, 255.0)),
             round(Int, lerp(green(c), 0.0, 255.0)),
             round(Int, lerp(blue(c), 0.0, 255.0)))
end

hex(c::Color) = hex(convert(RGB, c))
hex(c::Colorant) = hex(convert(ARGB, c))

"""
    weighted_color_mean(w1, c1, c2)

Returns the color `w1*c1 + (1-w1)*c2` that is the weighted mean of `c1` and
`c2`, where `c1` has a weight 0 ≤ `w1` ≤ 1.
"""
function weighted_color_mean(w1::Real, c1::Colorant, c2::Colorant)
    weight1 = convert(promote_type(eltype(c1), eltype(c2)),w1)
    weight2 = weight1 >= 0 && weight1 <= 1 ? oftype(weight1,1-weight1) : throw(DomainError())
    mapc((x,y)->weight1*x+weight2*y, c1, c2)
end
function weighted_color_mean(w1::Real, c1::Gray{Bool}, c2::Gray{Bool})
    # weighting of two Gray{Bool} would return different color type and therefore omitted
    throw(DomainError())
end

"""
    linspace(c1::Color, c2::Color, n=100)

Generates `n`>2 colors in a linearly interpolated ramp from `c1` to`c2`,
inclusive, returning an `Array` of colors.
"""
function linspace{T<:Colorant}(c1::T, c2::T, n::Integer=100)
    return T[weighted_color_mean(w1, c1, c2) for w1 in linspace(1.0,0.0,n)]
end

#Double quadratic Bezier curve
function Bezier{T<:Real}(t::T, p0::T, p2::T, q0::T, q1::T, q2::T)
    B(t,a,b,c)=a*(1.0-t)^2.0+2.0b*(1.0-t)*t+c*t^2.0
    if t <= 0.5
        return B(2.0t, p0, q0, q1)
    else #t > 0.5
        return B(2.0(t-0.5), q1, q2, p2)
    end
end

#Inverse double quadratic Bezier curve
function invBezier{T<:Real}(t::T, p0::T, p2::T, q0::T, q1::T, q2::T)
    invB(t,a,b,c)=(a-b+sqrt(b^2.0-a*c+(a-2.0b+c)*t))/(a-2.0b+c)
    if t < q1
        return 0.5*invB(t,p0,q0,q1)
    else #t >= q1
        return 0.5*invB(t,q1,q2,p2)+0.5
    end
end


function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:BasicDQN},
    ::Val{:EmptyRoom},
    ::Nothing;
    seed = 123,
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_BasicDQN_EmptyRoom$(t)")
    end
    log_dir = joinpath(save_dir, "tb_log")
    lg = TBLogger(log_dir, min_level = Logging.Info)
    rng = StableRNG(seed)

    inner_env = GridWorlds.EmptyRoomDirected(rng = rng)
    action_space_mapping = x -> Base.OneTo(length(RLBase.action_space(inner_env)))
    action_mapping = i -> RLBase.action_space(inner_env)[i]
    env = RLEnvs.ActionTransformedEnv(
        inner_env,
        action_space_mapping = action_space_mapping,
        action_mapping = action_mapping,
    )
    env = RLEnvs.StateOverriddenEnv(env, x -> vec(Float32.(x)))
    env = RewardOverriddenEnv(env, x -> x - convert(typeof(x), 0.01))
    env = MaxTimeoutEnv(env, 240)

    ns, na = length(state(env)), length(action_space(env))
    agent = Agent(
        policy = QBasedPolicy(
            learner = BasicDQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                batch_size = 32,
                min_replay_history = 100,
                loss_func = huber_loss,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularArraySARTTrajectory(
            capacity = 1000,
            state = Vector{Float32} => (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] log_step_increment =
                    0
            end
        end,
    )

    description = """
    This experiment uses three dense layers to approximate the Q value.
    The testing environment is EmptyRoom.

    You can view the runtime logs with `tensorboard --logdir $log_dir`.
    Some useful statistics are stored in the `hook` field of this experiment.
    """

    Experiment(agent, env, stop_condition, hook, description)
end


export samplingFreq, roundFreq, numSamplesPerPeriod, phaseShift

function samplingFreq(rp::RedPitaya,dec::Integer)
  baseFreq = 125e6
  freq = div(baseFreq,dec)
  return freq
end

function roundFreq(rp::RedPitaya,dec::Integer,freq)
  numPeriods = round(Int64,samplingFreq(rp,dec) / freq)
  return div(samplingFreq(rp,dec), numPeriods)
end

function numSamplesPerPeriod(rp::RedPitaya,dec::Integer,freqR)
  return Int(div(samplingFreq(rp,dec),freqR))
end

"""
Return the number of time points the signal u has the be shifted
in order to have phase pi/2 (i.e. be a sine wave)
"""
function phaseShift(u, numPeriods)
  numSampPerPeriod = div(length(u),numPeriods)
  complexSine = exp(-2*pi*im*numPeriods*(0:length(u)-1)/length(u))
  fourierCoeff = dot(complexSine,u)
  return round(Int,numSampPerPeriod*(angle(fourierCoeff) -pi/2) / (2pi))
end

function test2(rp)
  a=zeros(30)
  for l=1:30
    wpos = parse(Int32,query(rp,"ACQ:WPOS?"))
    u = receiveAnalogSignalLowLevel(rp,1,0,16384,"STA",true,true)
    a[l] = wpos
  end
  return a
end


function check_open_orders_klines(config, open_orders, closed_orders, log, dt, x, high, low)

end

function check_open_orders(config, open_orders, closed_orders, log, current)

    if length(keys(open_orders)) > 0
        extra_order_count = parse(Int64, config["extra_order_count"])
        extra_order_exec  = log["extra_order_exec"]
        order_volume = config["order_volume"]
        extra_order_martingale = config["extra_order_martingale"]
        precision = config["precision"]
        extra_order_step = config["extra_order_step"]
        sell_precision = config["sell_precision"]
        stop_loss = parse(Float64, config["stop_loss"])
        fee = config["fee"]
        take_profit = config["take_profit"]

        if config["dry_mode"] == "true"

            if current >= parse(Float64, open_orders["SELL"]["price"]) ### TAKE PROFIT
#                 @info("Take profit was executed")
                open_orders["SELL"]["status"] = "FILLED"
                open_orders["SELL"]["updateTime"] = datetime2unix(now())*1000
                open_orders["SELL"]["executedQty"] = open_orders["SELL"]["origQty"]
                closed_orders[string(now())] = copy(open_orders["SELL"])
                open_orders["BUY"]["status"] == "CANCELLED"
                closed_orders[string(now() - Second(1))] = open_orders["BUY"]
                log["extra_order_exec"] = 0
                log["trade_ongoing"] = false

                open_orders = Dict()

                open("closed_orders.json", "w") do f
                    JSON.print(f, closed_orders, 4)
                end
                open("open_orders.json", "w") do f
                    JSON.print(f, open_orders, 4)
                end
                open("log.json", "w") do f
                    JSON.print(f, log, 4)
                end
                log = Pinguu.daily_log(closed_orders, config)
                #send_notification(config, "take profit", log);

            elseif extra_order_exec < extra_order_count
                if current <= parse(Float64, open_orders["BUY"]["price"]) ### EXTRA ORDER
#                 @info("Extra order was executed")
                    log["extra_order_exec"] += 1
                    closed_orders = JSON.parsefile("closed_orders.json");
                    open_orders = JSON.parsefile("open_orders.json");
                    executed_order = deepcopy(open_orders["BUY"])
                    executed_order["executedQty"] = executed_order["origQty"]
                    executed_order["status"] = "FILLED"
                    executed_order["updateTime"] = datetime2unix(now())*1000
                    failed_order = deepcopy(open_orders["SELL"])
                    failed_order["status"] = "CANCELLED"
                    closed_orders[string(unix2datetime(executed_order["updateTime"]/1000))] = executed_order
                    closed_orders[string(now() - Second(1))] = failed_order
                    open_orders["BUY"]["time"] = datetime2unix(now())*1000
                    open_orders["BUY"]["updateTime"] = datetime2unix(now())*1000
                    open_orders["BUY"]["orderId"] = rand(1:1:1e7)
                    old_quant = parse(Float64, open_orders["BUY"]["origQty"])
                    old_price = parse(Float64, open_orders["BUY"]["price"])
                    new_quant = floor(order_volume * (extra_order_martingale)^extra_order_exec / parse(Float64, open_orders["BUY"]["price"]), digits=precision)
                    open_orders["BUY"]["origQty"] = string(new_quant)
                    open_orders["BUY"]["executedQty"] = "0.0"
                    open_orders["BUY"]["status"] = "NEW"
                    if extra_order_exec != extra_order_count
                        open_orders["BUY"]["price"] = string(round(old_price * (1 - extra_order_step), digits=sell_precision))
                    else
                        open_orders["BUY"]["status"] = "STOPLOSSLIMIT"
                        open_orders["BUY"]["price"] = string(round(parse(Float64, open_orders["SELL"]["price"]) * (100 - stop_loss)/100, digits=sell_precision))
                    end

                    oriqQty_sell = parse(Float64, open_orders["SELL"]["origQty"])
                    newQty_sell  = floor(oriqQty_sell  + old_quant, digits=precision)
                    origPrice_sell = parse(Float64, open_orders["SELL"]["price"])
                    qty_old = 0
                    for i in 0:1:(extra_order_exec-1)
                        qty_old += extra_order_martingale^i
                    end
                    qty_new = extra_order_martingale^extra_order_exec
                    new_price = (qty_old * origPrice_sell + qty_new * old_price * (1 + fee + take_profit/100)) / (qty_old + qty_new)

                    open_orders["SELL"]["time"] = datetime2unix(now())*1000
                    open_orders["SELL"]["updateTime"] = datetime2unix(now())*1000
                    open_orders["SELL"]["orderId"] = rand(1:1:1e7)
                    open_orders["SELL"]["origQty"] = string(newQty_sell)
                    open_orders["SELL"]["price"] = string(round(new_price, digits=sell_precision))

                    open("closed_orders.json", "w") do f
                        JSON.print(f, closed_orders, 4)
                    end
                    open("open_orders.json", "w") do f
                        JSON.print(f, open_orders, 4)
                    end
                    open("log.json", "w") do f
                        JSON.print(f, log, 4)
                    end
                end
#                 send_notification(config, "extra order", log);
            elseif extra_order_exec == extra_order_count ### STOP LOSS
                if current <= parse(Float64, open_orders["SELL"]["price"]) * (100 - parse(Float64, stop_loss))/100
                    open_orders["SELL"]["status"] = "FILLED"
                    open_orders["SELL"]["updateTime"] = datetime2unix(now())*1000
                    open_orders["SELL"]["executedQty"] = open_orders["SELL"]["origQty"]
                    open_orders["SELL"]["price"] = string(round(current, digits=sell_precision))
                    closed_orders[string(now())] = copy(open_orders["SELL"])
                    log["extra_order_exec"] = 0
                    log["trade_ongoing"] = false

                    open_orders = Dict()

                    open("closed_orders.json", "w") do f
                        JSON.print(f, closed_orders, 4)
                    end
                    open("open_orders.json", "w") do f
                        JSON.print(f, open_orders, 4)
                    end
                    open("log.json", "w") do f
                        JSON.print(f, log, 4)
                    end
                    log = Pinguu.daily_log(closed_orders, config)
                    #send_notification(config, "stop loss", log);
                end
            end
        else
            order = client.get_order(
                symbol=config["symbol"],
                orderId=open_orders["SELL"]["orderId"])
            if order["status"] == "FILLED"
#                 @info("Take profit was executed")
                closed_orders[string(unix2datetime(order["time"]/1000))] = order
                result = client.cancel_order(
                    symbol=config["symbol"],
                    orderId=open_orders["BUY"]["orderId"]);
                open_orders = Dict()
            end
            #
            #
            #
            #
            #
            #
            #
            #
            #
            #
            #
            #
            #
            #
            #
            #
            #
        end
    end
    return open_orders, closed_orders, log
end


"""
Driver to run TRON on GPU. This should be called from a kernel.
"""
@inline function tron_kernel(n::Int, shift::Int, max_feval::Int, max_minor::Int, gtol::Float64, scale::Float64, use_polar::Bool,
                     x::CuDeviceArray{Float64,1}, xl::CuDeviceArray{Float64,1},
                     xu::CuDeviceArray{Float64,1},
                     param::CuDeviceArray{Float64,2},
                     YffR::Float64, YffI::Float64,
                     YftR::Float64, YftI::Float64,
                     YttR::Float64, YttI::Float64,
                     YtfR::Float64, YtfI::Float64)
    tx = threadIdx().x

    g = @cuDynamicSharedMem(Float64, n, (3*n)*sizeof(Float64))
    xc = @cuDynamicSharedMem(Float64, n, (4*n)*sizeof(Float64))
    s = @cuDynamicSharedMem(Float64, n, (5*n)*sizeof(Float64))
    wa = @cuDynamicSharedMem(Float64, n, (6*n)*sizeof(Float64))
    wa1 = @cuDynamicSharedMem(Float64, n, (7*n)*sizeof(Float64))
    wa2 = @cuDynamicSharedMem(Float64, n, (8*n)*sizeof(Float64))
    wa3 = @cuDynamicSharedMem(Float64, n, (9*n)*sizeof(Float64))
    wa4 = @cuDynamicSharedMem(Float64, n, (10*n)*sizeof(Float64))
    wa5 = @cuDynamicSharedMem(Float64, n, (11*n)*sizeof(Float64))
    gfree = @cuDynamicSharedMem(Float64, n, (12*n)*sizeof(Float64))
    dsave = @cuDynamicSharedMem(Float64, n, (13*n)*sizeof(Float64))
    indfree = @cuDynamicSharedMem(Int, n, (14*n)*sizeof(Float64))
    iwa = @cuDynamicSharedMem(Int, 2*n, n*sizeof(Int) + (14*n)*sizeof(Float64))
    isave = @cuDynamicSharedMem(Int, n, (3*n)*sizeof(Int) + (14*n)*sizeof(Float64))

    A = @cuDynamicSharedMem(Float64, (n,n), (14*n)*sizeof(Float64)+(4*n)*sizeof(Int))
    B = @cuDynamicSharedMem(Float64, (n,n), (14*n+n^2)*sizeof(Float64)+(4*n)*sizeof(Int))
    L = @cuDynamicSharedMem(Float64, (n,n), (14*n+2*n^2)*sizeof(Float64)+(4*n)*sizeof(Int))

    if tx <= n
        @inbounds for j=1:n
            A[tx,j] = 0.0
            B[tx,j] = 0.0
            L[tx,j] = 0.0
        end
    end
    CUDA.sync_threads()

    task = 0
    status = 0

    delta = 0.0
    fatol = 0.0
    frtol = 1e-12
    fmin = -1e32
    cgtol = 0.1
    cg_itermax = n

    f = 0.0
    nfev = 0
    ngev = 0
    nhev = 0
    minor_iter = 0
    search = true

    while search

        # [0|1]: Evaluate function.

        if task == 0 || task == 1
            if use_polar
                f = eval_f_polar_kernel(n, shift, scale, x, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            else
                f = eval_f_kernel(n, scale, x, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            end
            nfev += 1
            if nfev >= max_feval
                search = false
            end

        end

        # [2] G or H: Evaluate gradient and Hessian.

        if task == 0 || task == 2
            if use_polar
                eval_grad_f_polar_kernel(n, shift, scale, x, g, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
                eval_h_polar_kernel(n, shift, scale, x, A, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            else
                eval_grad_f_kernel(n, scale, x, g, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
                eval_h_kernel(n, scale, x, A, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            end
            ngev += 1
            nhev += 1
            minor_iter += 1
        end

        # Initialize the trust region bound.

        if task == 0
            gnorm0 = dnrm2(n, g, 1)
            delta = gnorm0
        end

        # Call Tron.

        if search
            delta, task = dtron(n, x, xl, xu, f, g, A, frtol, fatol, fmin, cgtol,
                                cg_itermax, delta, task, B, L, xc, s, indfree, gfree,
                                isave, dsave, wa, iwa, wa1, wa2, wa3, wa4, wa5)
        end

        # [3] NEWX: a new point was computed.

        if task == 3
            gnorm_inf = dgpnorm(n, x, xl, xu, g)

            if gnorm_inf <= gtol
                task = 4
            end

            if minor_iter >= max_minor
                status = 1
                search = false
            end
        end

        # [4] CONV: convergence was achieved.
        # [10] : warning fval is less than fmin

        if task == 4 || task == 10
            search = false
        end
    end

    CUDA.sync_threads()

    return status, minor_iter
end

@inline function tron_qp_kernel(n::Int, max_feval::Int, max_minor::Int, gtol::Float64, scale::Float64,
    x::CuDeviceArray{Float64,1}, xl::CuDeviceArray{Float64,1}, xu::CuDeviceArray{Float64,1},
    A::CuDeviceArray{Float64,2}, c::CuDeviceArray{Float64,1})

    tx = threadIdx().x
    I = blockIdx().x

    g = @cuDynamicSharedMem(Float64, n, (3*n)*sizeof(Float64))
    xc = @cuDynamicSharedMem(Float64, n, (4*n)*sizeof(Float64))
    s = @cuDynamicSharedMem(Float64, n, (5*n)*sizeof(Float64))
    wa = @cuDynamicSharedMem(Float64, n, (6*n)*sizeof(Float64))
    wa1 = @cuDynamicSharedMem(Float64, n, (7*n)*sizeof(Float64))
    wa2 = @cuDynamicSharedMem(Float64, n, (8*n)*sizeof(Float64))
    wa3 = @cuDynamicSharedMem(Float64, n, (9*n)*sizeof(Float64))
    wa4 = @cuDynamicSharedMem(Float64, n, (10*n)*sizeof(Float64))
    wa5 = @cuDynamicSharedMem(Float64, n, (11*n)*sizeof(Float64))
    gfree = @cuDynamicSharedMem(Float64, n, (12*n)*sizeof(Float64))
    dsave = @cuDynamicSharedMem(Float64, 3, (13*n)*sizeof(Float64))
    indfree = @cuDynamicSharedMem(Int, n, (13*n+3)*sizeof(Float64))
    iwa = @cuDynamicSharedMem(Int, 2*n, n*sizeof(Int) + (13*n+3)*sizeof(Float64))
    isave = @cuDynamicSharedMem(Int, 3, (3*n)*sizeof(Int) + (13*n+3)*sizeof(Float64))

    B = @cuDynamicSharedMem(Float64, (n,n), (13*n+3+n^2)*sizeof(Float64)+(3*n+3)*sizeof(Int))
    L = @cuDynamicSharedMem(Float64, (n,n), (13*n+3+2*n^2)*sizeof(Float64)+(3*n+3)*sizeof(Int))

    if tx <= n
        @inbounds begin
            for j=1:n
                B[tx,j] = 0.0
                L[tx,j] = 0.0
            end
        end
    end
    CUDA.sync_threads()

    task = 0
    status = 0

    delta = 0.0
    fatol = 0.0
    frtol = 1e-12
    fmin = -1e32
    cgtol = 0.1
    cg_itermax = n

    f = 0.0
    nfev = 0
    ngev = 0
    nhev = 0
    minor_iter = 0
    search = true

    while search

        # [0|1]: Evaluate function.

        if task == 0 || task == 1
            f = eval_qp_f_kernel(n, x, A, c)
            nfev += 1
            if nfev >= max_feval
                search = false
            end
        end

        # [2] G or H: Evaluate gradient and Hessian.

        if task == 0 || task == 2
            eval_qp_grad_f_kernel(n, x, g, A, c)
            # We do not have to evaluate Hessian since A does not change.
            ngev += 1
            nhev += 1
            minor_iter += 1
        end

        # Initialize the trust region bound.

        if task == 0
            gnorm0 = dnrm2(n, g, 1)
            delta = gnorm0
        end

        # Call Tron.

        if search
            delta, task = dtron(n, x, xl, xu, f, g, A, frtol, fatol, fmin, cgtol,
                                cg_itermax, delta, task, B, L, xc, s, indfree, gfree,
                                isave, dsave, wa, iwa, wa1, wa2, wa3, wa4, wa5)
        end

        # [3] NEWX: a new point was computed.

        if task == 3
            gnorm_inf = dgpnorm(n, x, xl, xu, g)

            if gnorm_inf <= gtol
                task = 4
            end

            if minor_iter >= max_minor
                status = 1
                search = false
            end
        end

        # [4] CONV: convergence was achieved.
        # [10] : warning fval is less than fmin

        if task == 4 || task == 10
            search = false
        end
    end

    CUDA.sync_threads()

    return status, minor_iter
end


@graphs begin
	@node phase_3 = Start() -> p3_strat
	@node p3_strat = CheckStrategy("military") -> [n_true=p3_mili_1, n_false=p3_corr_1]

	# Military
	@node p3_mili_1 = BinaryCondition("The Fellowship is on the Mordor track.") -> [n_true = p3_mili_1_yes, n_false = p3_mili_2]
	@node p3_mili_1_yes = PerformAction("Assign the maximum allowed number of dice to the hunt pool.") -> p3_mili_end_phase

	@node p3_mili_2 = BinaryCondition("The Fellowship's progress is greater than 5.") -> [n_true = p3_mili_2_yes, n_false = p3_mili_3]
	@node p3_mili_2_yes = PerformAction("Assign 2 dice to the hunt pool.") -> p3_mili_end_phase

	@node p3_mili_3 = BinaryCondition("The Fellowship is on the starting position and its progress is 0.") -> [n_true = p3_mili_3_yes, n_false = p3_mili_3_no]
	@node p3_mili_3_yes = PerformAction("Assign 0 dice to the hunt pool.") -> p3_mili_end_phase
	@node p3_mili_3_no = PerformAction("Assign 1 dice to the hunt pool.") -> p3_mili_end_phase

	@node p3_mili_end_phase = End("End of Phase") -> []

	# Corruption
	@node p3_corr_1 = BinaryCondition("The Fellowship is on the starting position and its progess is 0.") -> [n_true = p3_corr_1_yes, n_false = p3_corr_2]
	@node p3_corr_1_yes = PerformAction("Roll a d6. On 4+, assign 1 die to the hunt pool, otherwise do nothing.") -> p3_corr_end_phase

	@node p3_corr_2 = BinaryCondition( " The Fellowship is on the Mordor track.  ") -> [n_true = p3_corr_2_yes, n_false = p3_corr_3]
	@node p3_corr_2_yes = PerformAction( " Assign the maximum allowed number of dice to the hunt pool.  ") -> p3_corr_end_phase

	@node p3_corr_3 = BinaryCondition("""
									  A *mobile* army is adjacent to its *target* which provides enough victory points to win the game.
									  Or, the Shadow have 7 dice.
									  """
									  ) -> [n_true = p3_corr_3_yes, n_false = p3_corr_4]
	@node p3_corr_3_yes = PerformAction( " Assign 1 die to the hunt pool.  ") -> p3_corr_end_phase

	@node p3_corr_4 = BinaryCondition( " The Fellowship's progress is greater than 4.  ") -> [n_true = p3_corr_4_yes, n_false = p3_corr_5]
	@node p3_corr_4_yes = PerformAction( " Assign 2 dice to the hunt pool.  ") -> p3_corr_end_phase

	@node p3_corr_5 = BinaryCondition("""
									  The Fellowship's shortest path to Mordor leads via a Shadow stronghold and the progress allows them to pass it or be within 2 steps from it.
									  """) -> [n_true = p3_corr_5_yes, n_false = p3_corr_5_no]
	@node p3_corr_5_yes = PerformAction("Assign 2 dice to the hunt pool.") -> p3_corr_end_phase
	@node p3_corr_5_no = PerformAction("Assign 1 dice to the hunt pool.") -> p3_corr_end_phase

	@node p3_corr_end_phase = End("End of Phase") -> []
end


# Simulations parameters
include("./decision_making.jl")
include("./reaction_time_analysis_functions.jl")
include("./distributions.jl")

param = Dict("T_rest"=>10000,
"mu0"=>30,
"noise_amp"=>0.02,
"I0E1"=>0.3255,
"I0E2"=>0.3255,
"JN11"=>0.2609,
"JN22"=>0.2609,
"JN12"=>0.0497,
"JN21"=>0.0497,
"JAext" => 0.00052,
"Tnmda" => 100,
"Tampa" => 2,
"gamma" => 0.641,
"Tstim"=>5000,
"coh"=>5,
"false_time"=>200,
"Threshold"=>20
)

fixed_param=Dict("a"=>270,"b"=>108,"d"=>0.1540)
#Parameters of the simulation for EUler methods
simu_param=Dict("dt"=>0.5,"time_wind"=>4,"slide_wind"=>4) #4=2/dt

using FileIO
using DataFrames
using CSV

#a=2
Tr=500
println(Tr)
# de 0.02 à 0.025 déjàa sauvegardé (voir 0.02,0.3 50)
df=DataFrame()
df[:Ir]=linspace(0.03,0.1,30)
for j=2:10
#CSV.write("Irsmall.csv",df)
 for i=1:25

for Ir in [0.035,0.047]
println(Ir)
Cohlist=[2,10,20]

 #Need to do the quantile analysis after with r s etc ...

for c in Cohlist[1:end]
c_min=-c
 c_max = c
 number = 2

   distri = make_distributions(c_min,c_max,number,3000)

 param["T_rest"] = Tr
 If=0

  result,r2= sim_mf_decisionmaking_pattern(distri,Ir,If,
    param,fixed_param,simu_param)

 CSV.write("./datatempAll/V$j-Coh$c-Final200Exp$i-T$Tr-Ir$Ir.csv",result)
 CSV.write("./datatempAll/V$j-NondynCoh$c-Final200Exp$i-T$Tr-Ir$Ir.csv",r2)

 end
end

end
end


using ShaderToy, FileIO, GLAbstraction

shadertoy("clouds.frag", Dict{Symbol, Any}(
	:iChannel0 => Texture(load("noise.png").data, x_repeat=:repeat, minfilter=:linear),
))


# using LibSerialPort, COBS,
using Statistics
using DataStructures
using AbstractPlotting, WGLMakie, JSServe, Markdown
using AbstractPlotting.MakieLayout
using JSServe: Slider
markdown_css = JSServe.Asset(JSServe.dependency_path("markdown.css"))

top_rpm = 11_500 + 1_150
t4 = 60e6/4
shortest_t = t4/1.1top_rpm
longest_t = 100_000
fps = 30
baudrate = 115200
history = 50
t₀ = time()

gettime() = time() - t₀

function toint(msg)
    y = zero(UInt32)
    for c in msg
        y <<= 8
        y += c
    end
    return y
end

function getrpm(t)
    if iszero(t)
        0.0
    elseif 0 < t ≤ shortest_t
        nothing
    elseif shortest_t < t ≤ longest_t
        t4/t
    else
        0.0
    end
end

last2 = last ∘ last

gety(line, ::Nothing) = last2(line)
gety(_, y) = y

struct Arduino
    c::ReentrantLock
    sp#::SerialPort
    line::Observable{CircularBuffer{Point2f0}}
    function Arduino(sp)
        line = Node(CircularBuffer{Point2f0}(history))
        for i in 1:history
            push!(line[], Point2f0(i/fps, 0.0))
        end
        c = ReentrantLock()
        new(c, sp, line)
    end
end

function sample!(a::Arduino)
    ts = lock(a.c) do
        # sp_flush(a.sp, SP_BUF_INPUT)
        # encode(a.sp, UInt8(0))
        rand(UInt32)
    end
    t = toint(ts)
    rpm = getrpm(t)
    y = gety(a.line[], rpm)
    xy = Point2f0(gettime(), y)
    push!(a.line[], xy)
    a.line[] = a.line[]
end

# port = only(get_port_list())
# sp = LibSerialPort.open(port, baudrate)

a = Arduino(nothing)

reading = @async while true#isopen(sp)
    cond = Condition()
    Timer(x->notify(cond), 1/fps)
    t = @async sample!(a)
    wait(cond)
    wait(t)
end

pwm = Node(1)
on(pwm) do i
    lock(a.c) do
        println(i)
    end
end

function handler(session, request)
    scene, layout = layoutscene(0)
    ax = layout[1,1] = LAxis(scene, xlabel = "Time (s)", ylabel = "RPM", title = "Fan 1")
    lines!(ax, a.line, color = :blue)
    on(a.line) do xys
        xmin = first(xys[1])
        xmax = first(xys[end]) + 1
        ax.targetlimits[] = FRect2D(xmin, 0, xmax - xmin, top_rpm)
    end
    rpm = lift(a.line) do xys
        round(Int, last(last(xys)))
    end
    slider_s = Slider(1:255, pwm)
    dom = md"""
    $scene
    Speed setting: $pwm $slider_s RPM: $rpm
    """
    return JSServe.DOM.div(markdown_css, dom)
end

app = JSServe.Application(handler, "0.0.0.0", 8000)


"""
Runs the Carleo & Troyer (2017) algorithm.
"""

function main(NetSettings)
    @unpack mfree, mc_trials, dim_rbm, n, repetitions, calc_stat, γ_init, γ_decay,
    save_figures, modelname, use_meter = NetSettings

    HamiltWeights = init_weights(WEIGHTS, modelname, n)
    NetParams = init_params(NETPARAMS, NetSettings)
    CGParams = init_cg_params(CGPARAMS, NetSettings)
    VarDervs = init_vardervs(VARDERIVS, NetSettings)

    if mfree
        OptParams = init_optim_mfree(OPTPARAMSMFREE, NetSettings)
    else
        OptParams = init_optim(OPTPARAMS, NetSettings)
    end

    elocs = SharedArray(zeros(Complex{Float64}, mc_trials, nworkers()))
    # Stores the energy samples of each worker in a column.
    forces_w = SharedArray(zeros(Complex{Float64}, dim_rbm, nworkers()))
    # Stores the forces of each worker in a column.
    warm_states = zeros(Int, n, nworkers())
    # Saves last visited state in a chain.
    mean_e = zeros(repetitions)
    γ = γ_init

    progress_gui = Progress(repetitions)  # Progress bar
    for i=1:repetitions
        @sync for (id, worker) in enumerate(workers())
            @async warm_states[:, id] = remotecall_fetch(mc_sampling!,
                worker,
                NetParams,
                VarDervs.o_tot,
                elocs,
                HamiltWeights,
                NetSettings,
                id,
                warm_state=warm_states[:, id]
                )
            end

        mean_e[i] = mean(real(elocs))
        cg_iters = optimize!(CGParams, VarDervs, OptParams, elocs, forces_w, γ, NetSettings)
        update_netparams!(NetParams, OptParams, NetSettings)
        γ *= γ_decay
        if use_meter
            ProgressMeter.next!(progress_gui; showvalues = [(:iter,i),
                (:E,mean_e[i]), (:CG_iters, cg_iters)])
        end
    end
    @everywhere GC.gc()
    eloc_and_stats = [0.0, init_stat(STATISTICS, n)]
    if calc_stat
        eloc_and_stats = stats_calculator(NetParams, HamiltWeights, warm_states, NetSettings)
    end
    return mean_e, NetParams, eloc_and_stats[1], eloc_and_stats[2]
end

"""
    run_NQS(NetSettings)

# Arguments
    - `NetSettings`: The settings struct containing all possible network settings.
"""
function run_NQS(NetSettings)
    @unpack modelname, n, test_iter, calc_stat, α, mag0, writetofile, save_figures = NetSettings
    init_print(NetSettings)

    energy = zeros(Float64, 0)
    for i=1:test_iter
        println("RUN $i")
        energy, NetParams, final_eloc, GroundStats = main(NetSettings)

        if writetofile
            if calc_stat
                namebase = "$modelname-n=$n-mag0=$mag0-alpha=$α-run=$i-"
                for k in fieldnames(STATISTICS)
                    writedlm(namebase * "statistics-$k.txt", getfield(GroundStats, k))
                end
                writedlm(namebase * "mc_and_final_eloc.txt", [energy; final_eloc])
                if save_figures
                    plot_statistics(GroundStats, namebase, n)
                end
            else
                writedlm(namebase * "mc_eloc.txt", energy)
            end
            for k in fieldnames(NETPARAMS)
                writedlm(namebase * "netparams-$k.txt", getfield(NetParams, k))
            end
        end
    end
    return energy
end

function init_print(NetSettings)
    @unpack n, modelname, pbc, γ_init, γ_decay, regulator, mag0, α, repetitions,
        mc_trials, therm_steps, stat_samples, mfree, iterative_inverse,
        writetofile =  NetSettings
    println("## Initializing $n spin system
## model = $modelname
## pbc = $pbc
## learning rate = $γ_init
## lrate-decay = $γ_decay
## S-matrix diag shift = $regulator
## workers = $(nworkers())
## ∑s_z=0-sector sampling = $mag0
## α = $α
## iterations = $repetitions
## samples per worker = $mc_trials
## thermalization steps = $therm_steps
## final samples per worker = $stat_samples
## matrixfree = $mfree
## conjugate gradient inversion = $iterative_inverse
## saving results = $writetofile")
end


using YaoBase

export ChainBlock, chain

"""
    ChainBlock{D} <: CompositeBlock{D}

`ChainBlock` is a basic construct tool to create
user defined blocks horizontically. It is a `Vector`
like composite type.
"""
struct ChainBlock{D} <: CompositeBlock{D}
    n::Int
    blocks::Vector{AbstractBlock{D}}
    function ChainBlock(n::Int, blocks::Vector{<:AbstractBlock{D}}) where {D}
        _check_block_sizes(blocks, n)
        return new{D}(n, blocks)
    end
end

ChainBlock(blocks::Vector{<:AbstractBlock{D}}) where {D} = ChainBlock(_check_block_sizes(blocks), blocks)
ChainBlock(blocks::AbstractBlock{D}...) where {D} =
    ChainBlock(collect(AbstractBlock{D}, blocks))

nqudits(c::ChainBlock) = c.n

"""
    chain(blocks...)

Return a [`ChainBlock`](@ref) which chains a list of blocks with same
[`nqudits`](@ref). If there is lazy evaluated
block in `blocks`, chain can infer the number of qudits and create an
instance itself.
"""
chain(blocks::AbstractBlock{D}...) where {D} = ChainBlock(blocks...)
function chain(blocks::Union{AbstractBlock{D},Function}...) where {D}
    blocks = filter(x->x isa AbstractBlock, blocks)
    return chain(map(x -> parse_block(_check_block_sizes(blocks...), x), blocks)...)
end

function chain(list::Vector{<:AbstractVector{D}}) where D
    return ChainBlock(list)
end

# if not all matrix block, try to put the number of qudits.
chain(n::Int, blocks...) = chain(map(x -> parse_block(n, x), blocks)...)
chain(n::Int, itr) = isempty(itr) ? chain(n) : chain(map(x -> parse_block(n, x), itr)...)
# disambiguity
# NOTE: we use parse_block here to make sure the behaviour are the same
chain(n::Int, it::Pair) = chain(n, parse_block(n, it))
chain(n::Int, f::Function) = chain(n, parse_block(n, f))
function chain(n::Int, block::AbstractBlock{D}) where D
    if n != nqudits(block)
        throw(QubitMismatchError("number of qudits mismatch: expect $n, got $(nqudits(block))"))
    end
    return ChainBlock(n, AbstractBlock{D}[block])
end
chain(blocks::Function...) = @λ(n -> chain(n, blocks...))
chain(it) = chain(it...) # forward iterator to vargs, so we could dispatch based on types
chain(it::Pair) = error("got $it, do you mean put($it)?")
chain(blocks...) = @λ(n -> chain(n, blocks))

"""
    chain(n)

Return an empty [`ChainBlock`](@ref) which can be used like a list of blocks.
"""
chain(n::Int; nlevel=2) = ChainBlock(n::Int, AbstractBlock{nlevel}[])

"""
    chain()

Return an lambda `n->chain(n)`.
"""
chain() = @λ(n -> chain(n))

subblocks(c::ChainBlock) = c.blocks
occupied_locs(c::ChainBlock) =
    Tuple(unique(Iterators.flatten(occupied_locs(b) for b in subblocks(c))))

chsubblocks(pb::ChainBlock{D}, blocks::Vector{<:AbstractBlock{D}}) where {D} =
    length(blocks) == 0 ? ChainBlock(pb.n, AbstractBlock{D}[]) : ChainBlock(pb.n, blocks)
chsubblocks(pb::ChainBlock, it) = chain(it...)

function mat(::Type{T}, c::ChainBlock{D}) where {T,D}
    if isempty(c.blocks)
        return IMatrix{D^nqudits(c),T}()
    else
        return prod(x -> mat(T, x), Iterators.reverse(c.blocks))
    end
end

function _apply!(r::AbstractRegister, c::ChainBlock)
    for each in c.blocks
        _apply!(r, each)
    end
    return r
end

cache_key(c::ChainBlock) = Tuple(cache_key(each) for each in c.blocks)

function Base.:(==)(lhs::ChainBlock{D}, rhs::ChainBlock{D}) where {D}
    nqudits(lhs) == nqudits(rhs) &&
    (length(lhs.blocks) == length(rhs.blocks)) &&
    all(lhs.blocks .== rhs.blocks)
end

Base.copy(c::ChainBlock{D}) where {D} = ChainBlock(c.n, copy(c.blocks))
Base.similar(c::ChainBlock{D}) where {D} = ChainBlock(c.n, empty!(similar(c.blocks)))
Base.getindex(c::ChainBlock, index) = getindex(c.blocks, index)
Base.getindex(c::ChainBlock, index::Union{UnitRange,Vector}) =
    ChainBlock(c.n, getindex(c.blocks, index))
Base.setindex!(c::ChainBlock{D}, val::AbstractBlock{D}, index::Integer) where {D} =
    (_check_block_sizes(c, val); setindex!(c.blocks, val, index); c)
Base.insert!(c::ChainBlock{D}, index::Integer, val::AbstractBlock{D}) where {D} =
    (_check_block_sizes(c, val); insert!(c.blocks, index, val); c)
Base.adjoint(blk::ChainBlock{D}) where {D} =
    ChainBlock(blk.n, AbstractBlock{D}[adjoint(b) for b in reverse(subblocks(blk))])
Base.lastindex(c::ChainBlock) = lastindex(c.blocks)
## Iterate contained blocks
Base.iterate(c::ChainBlock, st = 1) = iterate(c.blocks, st)
Base.length(c::ChainBlock) = length(c.blocks)
Base.eltype(c::ChainBlock) = eltype(c.blocks)
Base.eachindex(c::ChainBlock) = eachindex(c.blocks)
Base.popfirst!(c::ChainBlock) = popfirst!(c.blocks)
Base.pop!(c::ChainBlock) = pop!(c.blocks)
Base.push!(c::ChainBlock{D}, m::AbstractBlock{D}) where {D} = (_check_block_sizes(c, m); push!(c.blocks, m); c)
Base.push!(c::ChainBlock{D}, f::Function) where {D} = (push!(c.blocks, f(c.n)); c)
Base.pushfirst!(c::ChainBlock{D}, m::AbstractBlock{D}) where {D} =
    (_check_block_sizes(c, m); pushfirst!(c.blocks, m); c)
Base.pushfirst!(c::ChainBlock{D}, f::Function) where {D} = (pushfirst!(c.blocks, f(c.n)); c)
Base.append!(c::ChainBlock{D}, list::Vector{<:AbstractBlock{D}}) where {D} =
    (_check_block_sizes(c, list...); append!(c.blocks, list); c)
Base.append!(c1::ChainBlock{D}, c2::ChainBlock{D}) where {D} =
    (_check_block_sizes(c1, c2); append!(c1.blocks, c2.blocks); c1)
Base.prepend!(c1::ChainBlock{D}, list::Vector{<:AbstractBlock{D}}) where {D} =
    (_check_block_sizes(c1, list...); prepend!(c1.blocks, list); c1)
Base.prepend!(c1::ChainBlock{D}, c2::ChainBlock{D}) where {D} =
    (_check_block_sizes(c1, c2); prepend!(c1.blocks, c2.blocks); c1)

YaoBase.isunitary(c::ChainBlock) = all(isunitary, c.blocks) || isunitary(mat(c))
YaoBase.isreflexive(c::ChainBlock) =
    (iscommute(c.blocks...) && all(isreflexive, c.blocks)) || isreflexive(mat(c))
YaoBase.ishermitian(c::ChainBlock) =
    (all(isreflexive, c.blocks) && iscommute(c.blocks...)) || isreflexive(mat(c))


module Random

export uniform32, SplitMix64, Xoshiro256pp, UniformRNG, next, uniformfromindex, choose

@inline function uniform32(x::UInt64) :: Float32
    fract = UInt32((x >> 41) & 0b0_0000_0000_1111_1111_1111_1111_1111_111)
    k = (fract | 0b0_0111_1111_0000_0000_0000_0000_0000_000)
    reinterpret(Float32, k) - 1.0f0
end

# SplitMix64 is based on the description in http://prng.di.unimi.it/splitmix64.c
# and validated against 200 test vectors produced by that code.

mutable struct SplitMix64
    x::UInt64
end

@inline function next(sm::SplitMix64) :: UInt64
    sm.x += 0x9e3779b97f4a7c15
    z = sm.x
    z = (xor(z, z >> 30)) * 0xbf58476d1ce4e5b9
    z = (xor(z, z >> 27)) * 0x94d049bb133111eb

    xor(z, z >> 31)
end

# Xoshiro256++ is based on the description in http://prng.di.unimi.it/xoshiro256plusplus.c
# and validated against 200 test vectors produced by that code.

mutable struct Xoshiro256pp
    s0::UInt64
    s1::UInt64
    s2::UInt64
    s3::UInt64
end

@inline function rotl(x::UInt64, k::Int) :: UInt64
    (x << k) | (x >> (64 - k))
end

@inline function next(r::Xoshiro256pp) :: UInt64
    result = rotl(r.s0 + r.s3, 23) + r.s0

    t = r.s1 << 17
    r.s2 = xor(r.s2, r.s0)
    r.s3 = xor(r.s3, r.s1)
    r.s1 = xor(r.s1, r.s2)
    r.s0 = xor(r.s0, r.s3)

    r.s2 = xor(r.s2, t)
    r.s3 = rotl(r.s3, 45)

    result
end

# A wrapper around Xoshiro256pp to generate Float32s in the range
# 0.0 <= r < 1.0

struct UniformRNG
    x::Xoshiro256pp
end

@inline function next(r::UniformRNG) :: Float32
    uniform32(next(r.x))
end

# Factory function to create a UniformRNG from an index, which is
# normally the CUDA thread index.
@inline function uniformfromindex(index) :: UniformRNG
    seedgen = SplitMix64(index)
    s0 = next(seedgen)
    s1 = next(seedgen)
    s2 = next(seedgen)
    s3 = next(seedgen)
    UniformRNG(Xoshiro256pp(s0, s1, s2, s3))
end

end


#figuring out matrix
using Random
a = [1,2,3,4,5;6,7,8,9]
for i = 1:length(a)
    print(i)
end


# inplace
function constrain_all(
    m::AbstractMetricParams{T},
    us::AbstractArray{T,2},
    vs::AbstractArray{T,2},
    μ
) where {T<:Number}
    us_cols = eachcol(us)
    vs_cols = eachcol(vs)
    curried(u, v) = constrain(m, u, v, μ = μ)
    @. vs[1, :] = curried(us_cols, vs_cols)
    vs
end

# SMatrix -- do we actually want to support this?
function constrain_all(
    m::AbstractMetricParams{T},
    us::SMatrix{4,M,T,L},
    vs::SMatrix{4,M,T,L},
    μ
) where {M,T<:Number,L}
    alloc_vs = constrain_all(m, us, similar(vs), μ)
    SMatrix{4,M}(alloc_vs)
end

# vectors of vectors
function constrain_all(m::AbstractMetricParams{T}, us, vs, μ) where {T<:Number}
    curried(u, v) = constrain_all(m, u, v, μ)
    curried.(us, vs)
end

function constrain_all(
    m::AbstractMetricParams{T},
    u::AbstractVector{T},
    v::AbstractVector{T},
    μ
) where {T<:Number}
    v[1] = constrain(m, u, v, μ = μ)
    v
end

function constrain_all(
    m::AbstractMetricParams{T},
    u::StaticVector{S,T},
    v::StaticVector{S,T},
    μ
) where {S,T<:Number}
    mut = MVector{S,T}(v)
    mut[1] = constrain(m, u, v, μ = μ)
    SVector{S,T}(mut)
end


@inline DensityInterface.DensityKind(::Distribution) = DensityInterface.HasDensity()

for (di_func, d_func) in ((:logdensityof, :logpdf), (:densityof, :pdf))
    @eval begin
        DensityInterface.$di_func(d::Distribution, x) = $d_func(d, x)

        function DensityInterface.$di_func(d::UnivariateDistribution, x::AbstractArray)
            throw(ArgumentError("$(DensityInterface.$di_func) doesn't support multiple samples as an argument"))
        end

        function DensityInterface.$di_func(d::MultivariateDistribution, x::AbstractMatrix)
            throw(ArgumentError("$(DensityInterface.$di_func) doesn't support multiple samples as an argument"))
        end

        function DensityInterface.$di_func(d::MatrixDistribution, x::AbstractArray{<:AbstractMatrix{<:Real}})
            throw(ArgumentError("$(DensityInterface.$di_func) doesn't support multiple samples as an argument"))
        end
    end
end


"""
    The colormap from EigTool. This provides clearly
    distinguishable colors for a line-contour plot.
"""
function etcgrad()
    cm = [
        0   1.00000000000000   0.10000000000000
        0   0.92500000000000   0.17500000000000
        0   0.85000000000000   0.25000000000000
        0   0.77500000000000   0.32500000000000
        0   0.70000000000000   0.40000000000000
        0   0.62500000000000   0.47500000000000
        0   0.55000000000000   0.55000000000000
        0   0.47500000000000   0.62500000000000
        0   0.40000000000000   0.70000000000000
        0   0.32500000000000   0.77500000000000
        0   0.25000000000000   0.85000000000000
        0   0.17500000000000   0.92500000000000
        0   0.10000000000000   1.00000000000000
        0.06153846153846   0.09230769230769   1.00000000000000
        0.12307692307692   0.08461538461538   1.00000000000000
        0.18461538461538   0.07692307692308   1.00000000000000
        0.24615384615385   0.06923076923077   1.00000000000000
        0.30769230769231   0.06153846153846   1.00000000000000
        0.36923076923077   0.05384615384615   1.00000000000000
        0.43076923076923   0.04615384615385   1.00000000000000
        0.49230769230769   0.03846153846154   1.00000000000000
        0.55384615384615   0.03076923076923   1.00000000000000
        0.61538461538462   0.02307692307692   1.00000000000000
        0.67692307692308   0.01538461538462   1.00000000000000
        0.73846153846154   0.00769230769231   1.00000000000000
        0.80000000000000                  0   1.00000000000000
        0.80769230769231                  0   0.92307692307692
        0.81538461538462                  0   0.84615384615385
        0.82307692307692                  0   0.76923076923077
        0.83076923076923                  0   0.69230769230769
        0.83846153846154                  0   0.61538461538462
        0.84615384615385                  0   0.53846153846154
        0.85384615384615                  0   0.46153846153846
        0.86153846153846                  0   0.38461538461538
        0.86923076923077                  0   0.30769230769231
        0.87692307692308                  0   0.23076923076923
        0.88461538461538                  0   0.15384615384615
        0.89230769230769                  0   0.07692307692308
        0.90000000000000                  0                  0
        0.86923076923077                  0                  0
        0.83846153846154                  0                  0
        0.80769230769231                  0                  0
        0.77692307692308                  0                  0
        0.74615384615385                  0                  0
        0.71538461538462                  0                  0
        0.68461538461538                  0                  0
        0.65384615384615                  0                  0
        0.62307692307692                  0                  0
        0.59230769230769                  0                  0
        0.56153846153846                  0                  0
        0.53076923076923                  0                  0
        0.50000000000000                  0                  0
        0.54166666666667   0.05000000000000                  0
        0.58333333333333   0.10000000000000                  0
        0.62500000000000   0.15000000000000                  0
        0.66666666666667   0.20000000000000                  0
        0.70833333333333   0.25000000000000                  0
        0.75000000000000   0.30000000000000                  0
        0.79166666666667   0.35000000000000                  0
        0.83333333333333   0.40000000000000                  0
        0.87500000000000   0.45000000000000                  0
        0.91666666666667   0.50000000000000                  0
        0.95833333333333   0.55000000000000                  0
        1.00000000000000   0.60000000000000                  0
    ]
    nc = size(cm,1)
    cc = Vector{RGBA{Float64}}()
    vv = zeros(nc)
    for i in 1:nc
        col = RGBA(RGB(cm[i,:]...))
        push!(cc,col)
        vv[i]=(i-1)/(nc-1)
    end
    if :ContinuousColorGradient in names(Plots.PlotUtils, all=true)
        return PlotUtils.ContinuousColorGradient(cc,vv)
    else
        return PlotUtils.ColorGradient(cc,vv)
    end
end

const eigtool_cgrad = etcgrad()


export compute_compton_photon_flux, compute_compton_photons_per_bin

function compute_compton_photons_per_bin(energy_range, parameters)
    photon_number_per_bin = donthcomp(energy_range, parameters)
    return photon_number_per_bin
end


include(normpath(joinpath(@__FILE__, "..", "..", "exps", "common.jl")))
using IterTools
using CSV

function runconfig(setup::ScreenSetup,
                   lib::Library,
                   num_runs::Int64,
                   output_dir::Compat.String,
                   suppress_graph::Bool)

    info("Running config")
    methods = [venn, auprc]
    measures = [:inc, :dec, :incdec]
    genetypes = [:sigmoidal, :linear, :all]

    crisprtype = lib.cas9_behavior

    test_method_wrapper = (bc_counts, genes) -> test_methods_snr(bc_counts, genes, methods, measures, genetypes)

    # a little abuse of lexical scoping
    gen_plots = (bc_counts, genes) -> begin
        info("Generating plots")
        new_filename = joinpath(output_dir, "counts.svg")
        nopseudo = bc_counts[(bc_counts[:counts_bin1] .> 0.5) .& (bc_counts[:counts_bin2] .> 0.5), :]
        draw(SVG(new_filename, 10cm, 10cm), plot(nopseudo,
        x=:counts_bin1, y=:counts_bin2, color=:class, Scale.x_log10, Scale.y_log10,
        Theme(highlight_width=0pt), Coord.cartesian(fixed=true),
        Guide.title("$(typeof(crisprtype)) $(typeof(setup))"),
        Guide.xlabel("log counts t0"), Guide.ylabel("log counts t1")))
        new_filename = joinpath(output_dir, "volcano.svg")
        draw(SVG(new_filename, 10cm, 10cm), plot(genes,
        x=:mean_bin2_div_bin1, y=:pvalue_bin2_div_bin1, color=:class, Theme(highlight_width=0pt),
        Guide.title("$(typeof(crisprtype)) $(typeof(setup))"),
        Guide.xlabel("mean log2 fold change"), Guide.ylabel("-log10 pvalue")))
        test_methods_snr(bc_counts, genes, methods, measures, genetypes)
    end

    if suppress_graph
        funcs = [test_method_wrapper for i in 1:num_runs]
    else
        eval(:(using Gadfly))
        funcs = [gen_plots, [test_method_wrapper for i in 1:(num_runs - 1)]...]
    end

    # setup up runs
    runs = [(setup, lib, funcs[i], i) for i in 1:num_runs]

    results = pmap(args -> run_exp(args[1], args[2], args[3]; run_idx=args[4]), runs)

    info("Analyzing results")

    results = DataFrame(permutedims(hcat(results...), [2, 1]))
    mean_snr = mean(results[:x19]./results[:x20])
    std_snr = std(results[:x19]./results[:x20])
    snr_result = "SNR score = $(round(mean_snr, 3)) +/- $(round(std_snr, 3))"
    delete!(results, [:x19, :x20])
    hierarchy = vcat([hcat(item...) for item in IterTools.product(map(Symbol, methods), measures, genetypes)]...)
    new_names = [[:method, :measure, :genetype, :score]...; fieldnames(setup)...; :run_idx]
    results = construct_hierarchical_label(hierarchy, results, new_names)

    venn_result, auprc_result = "", ""

    grouped_info = by(results, [:method, :measure, :genetype]) do grouped_df
        n = size(grouped_df, 1)
        mean_score = mean(grouped_df[:score])
        std_score = std(float.(grouped_df[:score]))
        conf_int = 2.58 * std_score./sqrt(n)

        if (grouped_df[1, :measure] == :incdec) &&
           (grouped_df[1, :genetype] == :all)

            disp_score = round(mean_score, 3)
            conf_max = round(mean_score + conf_int, 3)
            conf_min = round(mean_score - conf_int, 3)

            if (grouped_df[1, :method] == :venn)
                venn_result = "Venn score = $disp_score, 95% conf int ($conf_min, $conf_max)"
            else
                auprc_result = "AUPRC score = $disp_score, 95% conf int ($conf_min, $conf_max)"
            end
        end

        DataFrame(
            std_score = std_score,
            mean_score = mean_score,
            conf_max = mean_score + conf_int,
            conf_min = mean_score - conf_int,
            n = n
        )
    end
    info("Saving results in $output_dir")
    CSV.write(joinpath(output_dir, "results_table.csv"), grouped_info)
    println("\nQuick results:\n##############\n$venn_result\n$auprc_result\n$snr_result")
end


function report(
    reportdir::String, invprob::InvestmentProblem{R,G1,G2,G3,I,T,P}
) where {R,G1,G2,G3,I,T,P}

    status = termination_status(invprob.model)

    if status != MOI.OPTIMAL
        status == MOI.OPTIMIZE_NOT_CALLED && error("Model has not yet been solved")
        error("Model was not solved to optimality: status is $(status)")
    end

    isdir(reportdir) || mkpath(reportdir)

    annualmarkets = DataFrame(
         scenario=String[], market=String[], demand=Float64[], price=Float64[])

    hourlymarkets = DataFrame(
        scenario=String[], market=String[], region=String[],
        period=String[], timestep=Int[], demand=Float64[], price=Float64[])

    dispatches = DataFrame(
        scenario=String[], market=String[], region=String[], gen=String[],
        period=String[], timestep=Int[], value=Float64[])

    for (scenario, isleaf) in scenarios(invprob)

        println("\n", scenario.name)
        oneoffdiscount, recurringdiscount = discount(scenario)
        periodweights = scenario.periodweights

        # Dispatch
        for market in [:energydispatch, :raisereserve, :lowerreserve],
            (r, regionname) in enumerate(invprob.regionnames),
            gentype in [:thermal, :variable, :storage]

            gentech = getfield(invprob.technologies, gentype)
            gennames = gentech.name

            if gentype == :storage
                gencapacities = gentech.maxdischarge
                dispatch = if market == :energydispatch
                    scenario.operations.storage.energydischarge .-
                    scenario.operations.storage.energycharge
                else
                    getfield(scenario.operations.storage, market)
                end
            else
                gencapacities = gentech.maxgen
                dispatch = getfield(getfield(scenario.operations, gentype), market)
            end

            marketname = market == :energydispatch ? "energy" : string(market)

            for (g, genname) in enumerate(gennames),
                (p, periodname) in enumerate(invprob.periodnames),
                t in 1:T

                    push!(dispatches,
                          (scenario=scenario.name, market=marketname,
                           region=regionname, gen=genname,
                           period=periodname, timestep=t,
                           value=value(dispatch[r,g,t,p])))

            end

        end

        # Capacity Market
        capmkt = scenario.markets.capacity
        capmkt_demand = value(capmkt.totalcontribution)
        capmkt_price =
            if capmkt_demand < capmkt.maxprice_capacity
                capmkt.maxprice
            elseif capmkt_demand < capmkt.midprice_capacity
                value(capmkt.seg1price)
            elseif capmkt_demand < capmkt.zeroprice_capacity
                value(capmkt.seg2price)
            else
                0
            end
        push!(annualmarkets, (
                scenario=scenario.name, market="capacity",
                demand=capmkt_demand, price=capmkt_price))

        # REC market
        recmkt = scenario.markets.rec
        recmkt_demand = recmkt.demand
        recmkt_price = dual(recmkt.marketclearing)
        push!(annualmarkets, (
            scenario=scenario.name, market="rec",
            demand=recmkt_demand, price=recmkt_price/recurringdiscount))

        # Hourly, regional market demand + prices
        for market in [:energy, :raisereserve, :lowerreserve],
            (r, regionname) in enumerate(invprob.regionnames)

            demand = getfield(scenario.markets, market).demand
            price = dual.(getfield(scenario.markets, market).marketclearing)

            for (p, periodname) in enumerate(invprob.periodnames), t in 1:T

                push!(hourlymarkets, (
                    scenario=scenario.name, market=string(market),
                    region=regionname, period=periodname, timestep=t,
                    demand=demand[r,t,p],
                    price=price[r,t,p]/periodweights[p]/recurringdiscount))

            end

        end

        if isleaf

            # TODO: Quantile plot of energy prices for each leaf scenario?

            # Save series of buildout progressions, by leaf scenario

            scen = scenario
            buildouts = DataFrame[]

            while !isnothing(scen)
                println(scen.name)
                push!(buildouts, investmentstates(scen))
                scen = scen.parent
            end

            CSV.write(joinpath(reportdir, "buildouts_$(scenario.name).csv"),
                      vcat(reverse(buildouts)...))

        end

    end

    dispatches.value[-1e-4 .< dispatches.value .< 0] .= 0.
    CSV.write(joinpath(reportdir, "dispatch.csv"), dispatches)
    CSV.write(joinpath(reportdir, "markets_hourly.csv"), hourlymarkets)
    CSV.write(joinpath(reportdir, "markets_annual.csv"), annualmarkets)

end

function investmentstates(scen::Scenario)

    result = DataFrame(
        scenario=String[], tech=String[], state=String[],
        count=Int[], capacity=Float64[])

    for techtype in [:thermal, :variable, :storage]

        techs = getfield(scen.investmentproblem.technologies, techtype)
        investments = getfield(scen.investments, techtype)

        for state in [:vesting, :holding, :building, :dispatching, :retired]

            counts = round.(Int, value.(getfield(investments, state)))

            for i in 1:length(techs.name)

                techname = techs.name[i]
                techcapacity = techtype == :storage ?
                    techs.maxdischarge[i] : techs.maxgen[i]

                count = sum(counts[:, i])
                push!(result, (scenario=scen.name, tech=techname,
                               state=string(state),
                               count=count,
                               capacity=count*techcapacity))

            end

        end

    end

    return result

end


using CUDA

function matmul(a, b, c, d, conf;
                transform_global_to_shared_a = Transform.Elementwise(),
                transform_global_to_shared_b = Transform.Elementwise(),
                transform_global_to_shared_c = Transform.Elementwise(),
                transform_shared_to_global_d = Transform.Elementwise(),
                transform_shared_to_regs_a = Transform.Elementwise(),
                transform_shared_to_regs_b = Transform.Elementwise(),
                transform_shared_to_regs_c = Transform.Elementwise(),
                transform_regs_to_shared_d = Transform.Elementwise(),
                epilogue = Epilogue.Default(),
                kernel = Kernel.matmul_singlestage)

    args = [a, b, c, d,
            transform_global_to_shared_a, transform_global_to_shared_b, transform_global_to_shared_c, transform_shared_to_global_d,
            transform_shared_to_regs_a, transform_shared_to_regs_b, transform_shared_to_regs_c, transform_regs_to_shared_d,
            epilogue,
            conf]

    hostkernel = @cuda launch=false kernel(args...)
    attributes(hostkernel.fun)[CUDA.FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES] = 64 * 1024
    hostkernel(args...; conf.launch_args...)
end


using Printf

export run_sim

# ============================================================================ #
"""
run_sim(d::LiveDemo, amp::Real=1.0, dur::Real=4.0, nrep::Integer=10)
"""
function run_sim(d::LiveDemo, amp::Real=1.0, dur::Real=4.0, nrep::Integer=10)
    stim = SineWave(d, amp, 4.0)
    return run_sim(d.net, stim, dur, nrep)
end
"""
run_sim(d::LiveDemo, stim::Stimulus, dur::Real=4.0, nrep::Integer=10)
"""
function run_sim(d::LiveDemo, stim::Stimulus, dur::Real=4.0, nrep::Integer=10)
    return run_sim(d.net, stim, dur, nrep)
end
# ---------------------------------------------------------------------------- #
"""
run_sim(net::LIFNetwork, stim::Stimulus, dur::Real, nrep::Integer,
    tstart::Time=0.0, record_vm::Bool=false)
"""
function run_sim(net::LIFNetwork, stim::Stimulus, dur::Real,
    nrep::Integer, tstart::Time=0.0, record_vm::Bool=false)

    #dur should be in seconds
    if dur > 100.0
        error("Simulation duration is really long... check your units!")
    end

    ncell = length(net)

    evt_ts = tstart:dur:tstart + dur*(nrep-1)

    fstim(id::Integer, t::Float64) = getstim(stim, id, t)
    # begin
        # if id == 1
        #     return amp*((0.5*sin(2.0*pi*(t*1e-3)*4.0))+0.5)
        # else
        #     return 0.0
        # end
    # end

    ts = [Vector{Float64}() for x in 1:ncell]

    vm = [zeros(1) for x in 1:ncell]
    if record_vm
        rec = collect(1:ncell)
    else
        rec = Int[]
    end

    @inbounds for kr = 1:nrep

        reset!(stim)
        Networks.reset!(net)
        d = simulate(net, fstim, dur*1e3, rec)

        @inbounds for k = 1:ncell
            append!(ts[k], d[:ts][k] .+ evt_ts[kr])
            if record_vm
                append!(vm[k], d[:vm][:,k])
            end
        end
    end

    return ts, vm, evt_ts
end
# =========================================================================== #


"""
To use this abstract type and avaliable functions,
create a subtype and implement the necessary functions.

Example:
    mutable struct Regression <: LD_Score_Regression
        ...

        function Regression(...)
            new(...)
        end
    end

    function update_weights(reg::Regression, ...)
    end

See `Hsq.jl` for a case study example.
"""

abstract type LD_Score_Regression end

# TODO need to find jackknife library or implement it
#function _delete_vals_tot(reg::LD_Score_Regression, jknife, Nbar, M) end
#function _delete_vals_part(reg::LD_Score_Regression, jknife, Nbar, M) end
#function _coef(reg::LD_Score_Regression, jknife, Nbar) end
#function _cat(reg::LD_Score_Regression, jknife, M, Nbar, coef, coef_cov) end
#function _tot(reg::LD_Score_Regression, cat, cat_cov) end
#function _prop(reg::LD_Score_Regression, jknife, M, Nbar, cat, tot) end
#function _enrichment(reg::LD_Score_Regression, M, M_tot, cat, tot) end
#function _intercept(reg::LD_Score_Regression, jknife) end
#function _combine_twostep_jknives(
#    reg::LD_Score_Regression, step1_jknife, step2_jknife, M_tot, c; Nbar=1,
#) end
#function _delete_vals_tot(reg::LD_Score_Regression) end

function aggregate(reg::LD_Score_Regression, y, x, N, M; intercept=nothing)
    if intercept == nothing
        intercept = reg.__null_intercept__
    end

    num = M * (mean(y) - intercept)
    denom = mean(x .* N)
    return num / denom
end

function ld_score_regression(
    reg::LD_Score_Regression,
    y, x, w, N, M, n_blocks, intercept, slow, step1_ii, old_weights,
)
    n_snp, reg.n_annot = size(x) # TODO fix this

    M_tot = sum(M)
    # shape should be [n_snp, 1]
    x_tot = sum(x, dims = 2)
    reg.constrain_intercept = intercept != nothing
    reg.intercept = intercept
    reg.n_blocks = n_blocks

    tot_agg = aggregate(reg, y, x_tot, N, M_tot; intercept=intercept)
    initial_w = update_weights(
        reg, x_tot, w, N, M_tot, tot_agg, intercept,
    )
    Nbar = mean(N)
    x = N .* x / Nbar

    if !reg.constrain_intercept
        x, x_tot = append_intercept(x), append_intercept(x_tot)
        yp = y
    else
        yp = y - intercept
        intercept_se = "NA"
    end

    reg.twostep_filtered = nothing
    if step1_ii != nothing && reg.constrain_intercept
        throw(ErrorException(
            "twostep is not compatible with constrain_intercept."
        ))
    elseif !(step1_ii == nothing) && reg.n_annot > 1
        throw(ErrorException(
            "twostep not compatible with partitioned LD Score yet."
        ))
    elseif !(step1_ii == nothing)
        n1 = sum(step1_ii)
        reg.twostep_filtered = n_snp - n1
        s = dropdims(step1_ii; dims=2)
        x1 = zeros(0)
        for (i, j) in enumerate(s)
            if s[i] == 1
                append!(x1, x[i])
            end
        end
        yp1, w1, N1, initial_w1 = map(
            (x) -> reshape(x[step1_ii], (n1, 1)), (yp, w, N, initial_w)
        )
    end

    # TODO finish implementing this function

    return reg
end


# args: dictfile outfile < input

using HDF5,JLD # ,DataStructures
using KUparser: Sentence, Position, DepRel

# Store postag/deprel strings as positive integer values
postag_hash = Dict{UTF8String,DepRel}()
deprel_hash = Dict{UTF8String,DepRel}()

# First argument is a dictionary file, could be an existing jld file
dict = load(ARGS[1])
for i=1:length(dict["postag"]); postag_hash[dict["postag"][i]] = i; end
for i=1:length(dict["deprel"]); deprel_hash[dict["deprel"][i]] = i; end

# ROOT deprel is stored as the special value 0
const ROOT="ROOT"
deprel_hash[ROOT] = 0

# The input is in conll07 format followed by wvec+cvec and read from stdin:
# ID,FORM,LEMMA,CPOSTAG,POSTAG,FEATS,HEAD,DEPREL,PHEAD,PDEPREL

corpus  = Sentence[]
s = nothing
mydims = 0
for l in eachline(STDIN)
    if l == "\n"
        @assert s != nothing
        push!(corpus, s)
        s = nothing
        continue
    end
    (id, form, lemma, cpostag, postag, feats, head, deprel, phead, pdeprel, wvec) = split(chomp(l), '\t')
    form = convert(UTF8String, form)
    postag = postag_hash[postag]
    deprel = deprel_hash[deprel]
    head = convert(Position, int(head))
    wvec = map(float32, split(wvec, ' '))
    if s == nothing
        mydims = length(wvec)
        s = Sentence()
        s.form = [form]
        s.postag = [postag]
        s.head = [head]
        s.deprel = [deprel]
        s.wvec = wvec''
    else
        @assert mydims == length(wvec)
        push!(s.form, form)
        push!(s.postag, postag)
        push!(s.head, head)
        push!(s.deprel, deprel)
        s.wvec = hcat(s.wvec, wvec)
    end
end

@assert s == nothing
@assert length(corpus) > 0
save(ARGS[2], "corpus", corpus, "postag", dict["postag"], "deprel", dict["deprel"])


struct ColumnBlockComponentCostModel{Tv, α_Col, β_Col} <: AbstractConnectivityModel
    α_col::α_Col
    β_col::β_Col
end

function ColumnBlockComponentCostModel{Tv}(; α_col=false, β_col=false) where {Tv}
    ColumnBlockComponentCostModel{Tv}(α_col, β_col)
end

@deprecate ColumnBlockComponentCostModel{Tv}(w_max, α_col, β_col) where {Tv, Ti} ConstrainedCost(ColumnBlockComponentCostModel{Tv}(α_col, β_col), VertexCount(), w_max)
function ColumnBlockComponentCostModel{Tv}(α_col::α_Col, β_col::β_Col) where {Tv, α_Col, β_Col}
    return ColumnBlockComponentCostModel{Tv, α_Col, β_Col}(α_col, β_col)
end

@inline cost_type(::Type{<:ColumnBlockComponentCostModel{Tv}}) where {Tv} = Tv

(mdl::ColumnBlockComponentCostModel{Tv, α_Col, β_Col})(n_vertices, n_pins, n_nets) where {Tv, α_Col, β_Col} = block_component(mdl.α_col, n_vertices) + n_nets * block_component(mdl.β_col, n_vertices)

struct BlockComponentCostModel{Tv, R, α_Row, α_Col, β_Row<:Tuple{Vararg{Any, R}}, β_Col<:Tuple{Vararg{Any, R}}}
    α_row::α_Row
    α_col::α_Col
    β_row::β_Row
    β_col::β_Col
end

function BlockComponentCostModel{Tv}(; α_row=false, α_col=false, β_row = (), β_col = ()) where {Tv}
    BlockComponentCostModel{Tv}(α_row, α_col, β_row, β_col)
end

function Base.permutedims(cst::BlockComponentCostModel{Tv}) where {Tv}
    return BlockComponentCostModel{Tv}(cst.α_col, cst.α_row, cst.β_col, cst.β_row)
end

@deprecate BlockComponentCostModel{Tv}(w_max, U, α_row, α_col, β_row, β_col) where {Tv, Ti} ConstrainedCost(BlockComponentCostModel{Tv}(α_row, α_col, β_row, β_col), VertexCount(), w_max)
function BlockComponentCostModel{Tv}(α_row::α_Row, α_col::α_Col, β_row::β_Row, β_col::β_Col) where {Tv, R, α_Row, α_Col, β_Row<:Tuple{Vararg{Any, R}}, β_Col<:Tuple{Vararg{Any, R}}}
    return BlockComponentCostModel{Tv, R, α_Row, α_Col, β_Row, β_Col}(α_row, α_col, β_row, β_col)
end

@inline cost_type(::Type{<:BlockComponentCostModel{Tv}}) where {Tv} = Tv

@propagate_inbounds block_component(f::F, w) where {F} = f(w)
@propagate_inbounds block_component(f::Number, w) = f
@propagate_inbounds block_component(f::Tuple, w) = f[w]
@propagate_inbounds block_component(f::AbstractArray, w) = f[w]

mutable struct BlockComponentCostStepOracle{Tv, Ti, Tc, R, Mdl<:BlockComponentCostModel{Tc, R}} <: AbstractOracleCost{Mdl}
    A::SparseMatrixCSC{Tv, Ti}
    Π_asg::MapPartition{Ti}
    Π_spl::SplitPartition{Ti}
    mdl::Mdl
    hst::Vector{Ti}
    Δ::Matrix{Tc}
    d::Vector{Tc}
    j::Ti
    j′::Ti
end

function oracle_stripe(hint::AbstractHint, mdl::BlockComponentCostModel{Tc, R}, A::SparseMatrixCSC{Tv, Ti}, Π; kwargs...) where {Tv, Ti, Tc, R}
    @inbounds begin
        m, n = size(A)
        K = length(Π)
        return BlockComponentCostStepOracle(A, convert(MapPartition, Π), convert(SplitPartition, Π), mdl, ones(Ti, K), undefs(Tc, R, n + 1), zeros(Tc, R), Ti(1), Ti(1))
    end
end

oracle_model(ocl::BlockComponentCostStepOracle) = ocl.mdl

@inline function (ocl::BlockComponentCostStepOracle{Tv, Ti, Tc, R})(j::Ti, j′::Ti, k...) where {Tv, Ti, Tc, R}
    @inbounds begin
        ocl_j = ocl.j
        ocl_j′ = ocl.j′
        A = ocl.A
        Π_asg = ocl.Π_asg.asg
        Π_spl = ocl.Π_spl.spl
        f = ocl.mdl
        d = ocl.d
        Δ = ocl.Δ
        pos = A.colptr
        idx = A.rowval
        hst = ocl.hst

        if j′ < ocl_j′
            ocl_j = Ti(1)
            ocl_j′ = Ti(1)
            zero!(d)
            one!(ocl.hst)
        end
        while ocl_j′ < j′
            q = pos[ocl_j′]
            q′ = pos[ocl_j′ + 1]

            for r = 1:R
                Δ[r, ocl_j′ + 1] = zero(Tc)
            end

            for _q = q:q′ - 1
                i = idx[_q]
                k = Π_asg[i]
                j₀ = hst[k] - 1
                u = Π_spl[k + 1] - Π_spl[k]
                if j₀  < ocl_j′
                    for r = 1:R
                        Δ[r, j₀ + 1] -= block_component(f.β_row[r], u)
                    end
                    for r = 1:R
                        Δ[r, ocl_j′ + 1] += block_component(f.β_row[r], u)
                    end
                end
                if j₀ < ocl_j
                    for r = 1:R
                        d[r] += block_component(f.β_row[r], u)
                    end
                end
                hst[k] = ocl_j′ + 1
            end

            ocl_j′ += 1
        end
        while j < ocl_j
            ocl_j -= 1
            for r = 1:R
                d[r] += Δ[r, ocl_j + 1]
            end
        end
        while j > ocl_j
            for r = 1:R
                d[r] -= Δ[r, ocl_j + 1]
            end
            ocl_j += 1
        end

        w = j′ - j
        c = block_component(f.α_col, w)
        for r = 1:R
            c += d[r] * block_component(f.β_col[r], w)
        end

        ocl.j = ocl_j
        ocl.j′ = ocl_j′
        return c
    end
end

function row_component_value(Π, mdl::BlockComponentCostModel)
    Π = convert(SplitPartition, Π)
    c_α = zero(cost_type(mdl))
    for k = 1:Π.K
        u = Π.spl[k + 1] - Π.spl[k]
        c_α += block_component(mdl.α_row, u)
    end
    return c_α
end

#=
function pack_stripe(A::SparseMatrixCSC{Tv, Ti}, method::DynamicTotalChunker{F}, args...; kwargs...) where {F<:BlockComponentCostModel, Tv, Ti}
    return pack_stripe(A, DynamicTotalChunker(ConstrainedCost(method.f, FeasibleCost(), Feasible())), args..., kwargs...)
end

function pack_stripe(A::SparseMatrixCSC{Tv, Ti}, method::DynamicTotalChunker{<:ConstrainedCost{F}}, Π; kwargs...) where {Tc, R, F<:BlockComponentCostModel{Tc, R}, Tv, Ti}
    @inbounds begin
        m, n = size(A)
        A_pos = A.colptr
        A_idx = A.rowval

        f = method.f.f
        w = oracle_stripe(StepHint(), method.f.w, A, Π) #TODO reverse outer loop so step oracles can have stationary j′
        w_max = method.f.w_max
        Π_asg = convert(MapPartition, Π).asg
        Π_spl = convert(DomainPartition, Π).spl
        K = length(Π)
        hst = fill(n + 1, K)
        Δ = zeros(Tc, R, n + 1)
        d = zeros(Tc, R)
        cst = Vector{cost_type(f)}(undef, n + 1) # cst[j] is the best cost of a partition from j to n
        spl = Vector{Int}(undef, n + 1)

        for r = 1:R
            Δ[r, n + 1] = zero(Tc)
        end
        cst[n + 1] = zero(Tc)
        for j = n:-1:1
            for q = A_pos[j] : A_pos[j + 1] - 1
                i = A_idx[q]
                k = Π_asg[i]
                u = Π_spl[k + 1] - Π_spl[k]
                if hst[k] > j
                    for r = 1:R
                        Δ[r, hst[k]] -= block_component(f.β_row[r], u)
                    end
                    for r = 1:R
                        Δ[r, j] += block_component(f.β_row[r], u)
                    end
                end
                hst[k] = j
            end
            zero!(d)
            for r = 1:R
                d[r] += Δ[r, j]
            end
            best_c = cst[j + 1]
            for r = 1:R
                best_c += d[r] * block_component(f.β_col[r], 1)
            end
            best_c += block_component(f.α_col, 1)
            best_j′ = j + 1
            for j′ = j + 2: n + 1
                if w(j, j′) > w_max
                    break
                end
                for r = 1:R
                    d[r] += Δ[r, j′ - 1]
                end
                c = cst[j′]
                for r = 1:R
                    c += d[r] * block_component(f.β_col[r], j′ - j)
                end
                c += block_component(f.α_col, j′ - j)
                if c < best_c
                    best_c = c
                    best_j′ = j′
                end
            end
            cst[j] = best_c
            spl[j] = best_j′
        end

        K = 0
        j = 1
        while j != n + 1
            j′ = spl[j]
            K += 1
            spl[K] = j
            j = j′
        end
        spl[K + 1] = j
        resize!(spl, K + 1)
        return SplitPartition{Ti}(K, spl)
    end
end
=#