text:
__precompile__()
#module SeisXcorrelation

using SeisIO, SeisNoise, Dates, FFTW, JLD2, Distributed, Sockets, ORCA

# read and write using jld2
include("io.jl")
# generate, sort, and classify station pairs
include("pairing.jl")
# miscellaneous utilities such as distance computation, windowing, and normalization
include("utils.jl")
# splitting utilities for high-order cross-correlation coda extraction
include("partition.jl")
# FFT wrapper for high-order cross-correlation
include("fft.jl")
# cross-correlation wrapper for high-order cross-correlation
include("correlate.jl")

export seisxcorrelation, seisxcorrelation_highorder

"""
    seisxcorrelation(data::Dict, tstamp::String, InputDict::Dict)

Compute cross-correlation save data in jld2 file with CorrData format.

# Arguments
- `data`::Dict    : Dictionary containing input data in the form: tstamp/stn => SeisChannel
- `tstamp::String,`    : Time stamp to read from JLD2 file.
- `InputDict::Dict`    : Dictionary containing IO, FFT, xcorr, and stacking parameters

# Output
- `tserrorList`::Array{String,1}    : Array containing the name of failed cross-correlation timestamp/stationpairs
- `basefoname.tstamp.jld2`    : contains CorrData structure with a hierarchical structure (CC function, metadata)

"""
function seisxcorrelation(data::Dict, tstamp::String, InputDict::Dict;verbose=false)
    # IO parameters
    basefoname = InputDict["basefoname"]
    time_unit  = InputDict["timeunit"]
    # FFT parameters
    freqmin    = Float64(InputDict["freqmin"])
    freqmax    = Float64(InputDict["freqmax"])
    fs         = Float64(InputDict["fs"])
    cc_len     = InputDict["cc_len"]
    cc_step    = InputDict["cc_step"]
    to_whiten  = InputDict["to_whiten"]
    time_norm  = InputDict["time_norm"]
    half_win   = InputDict["half_win"]
    water_level   = InputDict["water_level"]
    # correlation parameters
    corrtype   = InputDict["corrtype"]   # "xcorr," "acorr," or "xchancorr" (or any combination thereof)
    corrmethod = InputDict["corrmethod"] # "cross-correlation", "deconv", or "coherence"
    maxtimelag = InputDict["maxtimelag"]
    # stacking parameters
    stack      = InputDict["allstack"]
    stacktype  = InputDict["stacktype"]
    # dictionary to cache FFTs
    FFTDict = Dict{String, FFTData}()
    # list of stations that failed for this timestep
    tserrorList = []

    # split dataset names (keys of data) by "/" to get station list
    # assume form "$tstamp/$station"
    dsk = collect(keys(data))
    stlist = sort([string(split.(dsk[i], "/")[2]) for i=1:length(dsk)])

    # create output file for each time stamp, fill relevant info
    # outFile = jldopen("$basefoname.$tstamp.jld2", "a+")
    outFile = jldopen("$basefoname.$tstamp.jld2", true, true, true, IOStream)
    outFile["info/stationlist"] = stlist
    outFile["info/timeunit"] = time_unit

    println("$tstamp: Computing cross-correlations")
    stniter = 0 # counter to prevent computing duplicate xcorrs with reversed order
    # iterate over station list
    for stn1 in stlist
        stniter+=1

        # don't attempt FFT if this failed already
        if stn1 in tserrorList continue end

        # read station SeisChannels into SeisData before FFT
        S1 = try
            SeisData(data["$tstamp/$stn1"])
        catch;
            push!(tserrorList, "$tstamp/$stn1")
            filter!(a->a≠stn1, stlist)
            println("$tstamp: $stn1 encountered an error on read. Skipping.")
            continue
        end

        if S1.n > 1 @warn "SeisData contains multiple channels. Operating only on the first." end
        delete!(S1[1].misc, "kurtosis")
        delete!(S1[1].misc, "eqtimewindow")

        # make sure the data is the proper length to avoid dimension mismatch
        npts1 = Int(time_unit * S1[1].fs)
        if (length(S1[1].x) > npts1)
            S1[1].x=S1[1].x[1:npts1];
            idxtemp=findall(x->x<=npts1,S1[1].t[:,1])
            S1[1].t=S1[1].t[idxtemp,:];
        end
        # do not attempt fft if data was not available
        try
            if S1[1].misc["dlerror"] == 1
                push!(tserrorList, "$tstamp/$stn1")
                filter!(a->a≠stn1, stlist)
                println("$tstamp: $stn1 encountered an error: dlerror==1. Skipping.")
                continue
            elseif  sum(iszero.(S1[1].x)) > 0.5*length(S1[1].x) # skip if > 50% zeros
                push!(tserrorList, "$tstamp/$stn1")
                filter!(a->a≠stn1, stlist)
                println("$tstamp: $stn1 has >50% zeros. Skipping.")
                continue
            elseif  size(S1[1].t,1)>2
                push!(tserrorList, "$tstamp/$stn2")
                filter!(a->a≠stn2, stlist)
                println("$tstamp: $stn2 has $(size(S1[1].t,1)-2) gaps. Skipping.")
                continue
            end
        catch;
            # assume key "dlerror" does not exist (not downloaded via SeisDownload)
        end

        FFT1 = try
            # try to read FFT from cached FFTs
            if stn1 in keys(FFTDict)
                FFTDict[stn1]
            else
                if verbose; println("fft1 -ing");end
                #the following are new steps to do FFT
                process_raw!(S1,fs)
                R1=RawData(S1,cc_len,cc_step)
                R1.freqmin=freqmin;
                R1.freqmax=freqmax;
                R1.whitened=to_whiten;
                R1.time_norm=time_norm;

                detrend!(R1)
                taper!(R1)

                if time_norm == "one-bit" || time_norm == "phase"
                    if time_norm == "one-bit"
                        onebit!(R1);
                    end
                end
                FFT1 = compute_fft(R1)
                if to_whiten
                    whiten!(FFT1,freqmin,freqmax)
                end
                FFTDict[stn1] = FFT1
                FFT1
            end
        catch y
            push!(tserrorList, "$tstamp/$stn1")
            filter!(a->a≠stn1, stlist)
            println(y)
            exit()
            println("$tstamp: $stn1 encountered an error on FFT. Skipping.")
            continue
        end

        # iterate over station list again
        for stn2 in stlist[stniter:end]
            if stn2 in tserrorList continue end # don't attempt FFT if this failed already

            # see if this is an auto-, cross-, or xchan-correlation
            ct = get_corrtype([stn1, stn2])

            # autocorrelation
            if (ct=="acorr") && ("acorr" in corrtype)
                # set the stn2 FFT to the already computed FFT for stn1
                FFT2 = FFT1

            # cross- or cross-channel correlation
            elseif ct in corrtype
                # read station SeisChannels into SeisData before FFT
                S2 = try
                    SeisData(data["$tstamp/$stn2"])
                catch;
                    push!(tserrorList, "$tstamp/$stn2")
                    filter!(a->a≠stn2, stlist)
                    println("$tstamp: $stn2 encountered an error on read. Skipping.")
                    continue
                end

                if S1.n > 1 @warn "SeisData contains multiple channels. Operating only on the first." end
                delete!(S2[1].misc, "kurtosis")
                delete!(S2[1].misc, "eqtimewindow")
                # make sure the data is the proper length to avoid dimension mismatch
                npts2 = Int(time_unit * S2[1].fs)
                if (length(S2[1].x) > npts2)
                    S2[1].x=S2[1].x[1:npts2];
                    idxtemp=findall(x->x<=npts2,S2[1].t[:,1])
                    S2[1].t=S2[1].t[idxtemp,:];
                end
                # do not attempt fft if data was not available, >50% zeros, or has gaps.
                try
                    if S2[1].misc["dlerror"] == 1
                        push!(tserrorList, "$tstamp/$stn2")
                        filter!(a->a≠stn2, stlist)
                        println("$tstamp: $stn2 encountered an error: dlerror==1. Skipping.")
                        continue
                    elseif  sum(iszero.(S2[1].x)) > 0.5*length(S2[1].x) # skip if > 50% zeros
                        push!(tserrorList, "$tstamp/$stn2")
                        filter!(a->a≠stn2, stlist)
                        println("$tstamp: $stn2 has >50% zeros. Skipping.")
                        continue
                    elseif  size(S2[1].t,1)>2
                        push!(tserrorList, "$tstamp/$stn2")
                        filter!(a->a≠stn2, stlist)
                        println("$tstamp: $stn2 has $(size(S2[1].t,1)-2) gaps. Skipping.")
                        continue
                    end
                catch;
                    # assume that the key dlerror does not exist (data not downloaded via SeisDownload)
                end
                #
                # # make sure the data is the proper length to avoid dimension mismatch
                # npts2 = Int(time_unit * S2[1].fs)
                # if (length(S2[1].x) > npts2) S2[1].x=S2[1].x[1:npts2]; S2[1].t[end,1]=npts2 end

                # try to read FFT from cached FFTs, compute if not found
                FFT2 = try
                    if stn2 in keys(FFTDict)
                        FFTDict[stn2]
                    else
                        if verbose; println("fft2 -ing");end
                        #the following are new steps to do FFT
                        process_raw!(S2,fs)
                        R2=RawData(S2,cc_len,cc_step)
                        R2.freqmin=freqmin;
                        R2.freqmax=freqmax;
                        R2.whitened=to_whiten;
                        R2.time_norm=time_norm;

                        detrend!(R2)
                        taper!(R2)

                        if time_norm == "one-bit" || time_norm == "phase"
                            if time_norm == "one-bit"
                                onebit!(R2);
                            end
                        end
                        FFT2 = compute_fft(R2)
                        if to_whiten
                            whiten!(FFT2,freqmin,freqmax)
                        end
                        # FFT2 = compute_fft(S2, freqmin, freqmax, fs, Int(cc_step), Int(cc_len), to_whiten=to_whiten, time_norm=time_norm)

                        # store FFT2 in dict after smoothing
                        FFTDict[stn2] = FFT2
                        FFT2
                    end
                catch y
                    push!(tserrorList, "$tstamp/$stn2")
                    filter!(a->a≠stn2, stlist)
                    println("$tstamp: $stn2 encountered an error on FFT. Skipping.")
                    continue
                end
            else
                println("Skipping cross-correlation of $stn1 and $stn2.")
                continue
            end

            # TODO: find a way to store FFT2 for deconvolution without storing two sets of FFTs (smooth and unsmoothed)
            # deconvolution divides the cross spectrum by the squared power spectrum of the second signal
            # but each signal can be both source and receiver
            # smooth FFT2 if it hasn't been smoothed already
            if corrmethod == "coherence"
                if verbose; println("coherence -ing");end
                xcorr = compute_cc(coherence(FFT1, half_win,water_level), coherence(FFT2, half_win,water_level), maxtimelag, corr_type=corrmethod)
                # coherence!(FFT1, half_win,water_level)
                # coherence!(FFT2, half_win,water_level)
            elseif corrmethod == "deconv"
                if verbose; println("deconv -ing");end
                # deconvolution!(FFT2, half_win,water_level)
                xcorr = compute_cc(deconvolution(FFT1, half_win,water_level),FFT2,  maxtimelag, corr_type=corrmethod)
            else
                # compute correlation using SeisNoise.jl -- returns type CorrData
                xcorr = compute_cc(FFT1, FFT2, maxtimelag, corr_type=corrmethod)
            end

            varname = "$tstamp/$stn1.$stn2"
            try
                # compute distance between stations
                xcorr.misc["dist"] = dist(FFT1.loc, FFT2.loc)
                # save location of each station
                xcorr.misc["location"] = Dict(stn1=>FFT1.loc, stn2=>FFT2.loc)

                # stack over DL_time_unit
                if stack==true
                    if stacktype=="mean"
                        stack!(xcorr, allstack=true) #default is mean stacking/linear stacking.
                    else
                        stack!(xcorr, allstack=true,stacktype=stacktype)
                    end
                end

                outFile[varname] = xcorr
            catch
                println("$stn1 and $stn2 have no overlap at $tstamp.")
                push!(tserrorList, "$tstamp/$stn1")
            end
        end

        # release memory held by the FFT and time series for this station
        delete!(FFTDict, stn1)
        delete!(data, "$tstamp/$stn1")
    end
    outFile["info/errors"] = tserrorList
    close(outFile)
end

"""

    seisxcorrelation_highorder(tstamp::String, corrstationlist::Array{String,2}, data::JLD2.JLDFile, InputDict::Dict)

Compute C2 or C3 cross-correlation and save data in jld2 file with CorrData format.

# Arguments
- `tstamp::String,`    : Time stamp to read from JLD2 file.
- `corrstationlist::Array{String,1},`    : List of cross-correlation names, e.g. BP.CCRB..BP1.BP.EADB..BP1
- `data::JLD2.JLDFile`    : JLD2 file object to read data from.
- `InputDict::Dict`    : Dictionary containing IO, FFT, xcorr, and stacking parameters

# Output
- `basefoname.tstamp.jld2`    : contains SeisData structure with a hierarchical structure (CC function, metadata)

"""
function seisxcorrelation_highorder(data::Dict, tstamp::String, corrstationlist::Array{String,2}, allowable_pairs::Array{String,1}, allowable_vs::Array{String,1}, InputDict::Dict)
    # IO parameters
    basefoname = InputDict["basefoname"]
    # Partitioning parameters
    start_lag  = InputDict["start_lag"]
    win_len    = InputDict["window_len"]
    # FFT parameters
    freqmin    = InputDict["freqmin"]
    freqmax    = InputDict["freqmax"]
    fs         = InputDict["fs"]
    cc_len     = InputDict["cc_len"]
    cc_step    = InputDict["cc_step"]
    # correlation parameters
    corrtype   = InputDict["corrtype"]
    maxtimelag = InputDict["maxtimelag"]
    # stacking parameters
    stack      = InputDict["allstack"]

    pairiter = 1 # counter to prevent computing duplicate xcorrs with reversed order

    # dictionary to cache FFTs
    FFTDict = Dict{String, Array{FFTData,1}}()

    # list of stations that failed for this timestep
    tserrorList = []

    xcorrlist = deepcopy(corrstationlist)

    # convert samples given in [s] to the number of samples
    if typeof(start_lag) == Float64
        start_lag = convert(Int64, start_lag*fs)
    end
    if typeof(win_len) == Float64
        win_len  = convert(Int64, win_len*fs)
    end

    outFile = jldopen("$basefoname.$tstamp.jld2", "a+")
    outFile["info/corrstationlist"] = xcorrlist

    println("$tstamp: Computing cross-correlations")
    for p1=1:length(xcorrlist[1,:])
        # first station pair name and its reverse
        pair1 = xcorrlist[:, p1]
        p1name = pair1[1]*"."*pair1[2]
        p1namerev = pair1[2]*"."*pair1[1]

        for p2=pairiter:length(xcorrlist[1,:])
            # second station pair name
            pair2 = xcorrlist[:, p2]
            p2name = pair2[1]*"."*pair2[2]

            # read data only if we have 3 unique stations
            virt_src = intersect(pair1, pair2)

            if length(virt_src) != 1 continue end

            # set stations used as receivers for c3
            stn1 = setdiff(pair1, virt_src)[1]
            stn2 = setdiff(pair2, virt_src)[1]

            # keep string instead of array for ease of use
            virt_src = virt_src[1]

            # read data only if we have 3 unique station pairs
            if get_corrtype([stn1, stn2]) != "xcorr" continue end
            # read data only if the virtual source is approved
            if virt_src ∉ allowable_vs continue end

            # read data only if the station pair is approved
            if ("$stn1.$stn2" ∉ allowable_pairs) && ("$stn2.$stn1" ∉ allowable_pairs) continue end

            # load FFT1 from dict if it exists
            if "$(virt_src)/$stn1" in keys(FFTDict)
                FFT1_neg, FFT1_pos = FFTDict["$(virt_src)/$stn1"]
            else
                # load data and (optionally) reverse it to get virt_src-stn config
                # copy struct to prevent in-place reversal from messing up future xcorrs
                xcorr1 = copy(data["$tstamp/$p1name"])

                if size(xcorr1.corr, 2)>1 stack!(xcorr1) end

                # reverse xcorr if we loaded stn1-virt_src to get virt_src-stn1
                if virt_src==pair1[2]
                    xcorr1.corr=reverse(xcorr1.corr, dims=1)
                    vs_loc = xcorr1.misc["location"][pair1[2]]
                    stn1_loc = xcorr1.misc["location"][pair1[1]]
                else
                    vs_loc = xcorr1.misc["location"][pair1[1]]
                    stn1_loc = xcorr1.misc["location"][pair1[2]]
                end

                # partition xcorr into positive and negative windows
                xcorr1.corr = partition(xcorr1.corr, start_lag, win_len)

                # compute FFT
                FFT1_neg, FFT1_pos = compute_fft_c3(xcorr1, freqmin, freqmax, fs, cc_step, cc_len)

                FFTDict["$virt_src/$stn1"] = [FFT1_neg, FFT1_pos]
            end

            # load FFT2 from dict if it exists
            if "$virt_src/$stn2" in keys(FFTDict)
                FFT2_neg, FFT2_pos = FFTDict["$virt_src/$stn2"]
            else
                # load data and (optionally) reverse it to get virt_src-stn config
                # copy struct to prevent in-place reversal from messing up future xcorrs
                xcorr2 = copy(data["$tstamp/$p2name"])

                if size(xcorr2.corr, 2)>1 stack!(xcorr2) end

                # reverse xcorr if we loaded stn1-virt_src to get virt_src-stn1
                if virt_src==pair2[2]
                    xcorr2.corr=reverse(xcorr2.corr, dims=1)
                    stn2_loc = xcorr2.misc["location"][pair2[1]]
                else
                    stn2_loc = xcorr2.misc["location"][pair2[2]]
                end

                # partition xcorr into positive and negative windows
                xcorr2.corr = partition(xcorr2.corr, start_lag, win_len)

                # compute FFT
                FFT2_neg, FFT2_pos = compute_fft_c3(xcorr2, freqmin, freqmax, fs, cc_step, cc_len)

                FFTDict["$virt_src/$stn2"] = [FFT2_neg, FFT2_pos]
            end

            #compute xcorr - 1st half of dim=2 is neg lag, 2nd half is pos lag
            xcorr_c3_neg = compute_cc_c3(FFT1_neg, FFT2_neg, maxtimelag)
            xcorr_c3_neg.name = "$stn1.$stn2.$(virt_src)_neg"
            xcorr_c3_pos = compute_cc_c3(FFT1_pos, FFT2_pos, maxtimelag)
            xcorr_c3_pos.name = "$stn1.$stn2.$(virt_src)_pos"

            # stacking
            if stack
                stack!(xcorr_c3_neg, allstack=true)
                stack!(xcorr_c3_pos, allstack=true)
            end

            # save cross-correlation
            varname = "$tstamp/$stn1.$stn2/$virt_src"
            try
                outFile[varname] = [xcorr_c3_neg, xcorr_c3_pos]
            catch
                println("$varname encountered an error on saving to JLD2.")
                push!(tserrorList, varname)
                continue
            end
        end

        # release memory held by the FFT for this station pair in FFTDict
        delete!(FFTDict, p1name)
        delete!(FFTDict, p1namerev)
        # release memory held by the input data
        delete!(data, "$tstamp/$p1name")
        delete!(data, "$tstamp/$p1namerev")

        pairiter += 1
    end
    outFile["erorrs"] = tserrorList
    close(outFile)
end

#end # module


text:
module MuseScore
using MIDI
using DefaultApplication

const MUSESCORE = @static Sys.iswindows() ? "MuseScore" : "mscore"
const MUSESCORE_EXISTS = [false]

function test_musescore()
    if !MUSESCORE_EXISTS[1]
        r = try
            r = run(`$(MUSESCORE) -v`)
        catch
            false
        end
        if r == false || ((typeof(r) == Base.Process) && r.exitcode != 0)
            throw(SystemError(
            """
            The command `$(MUSESCORE) -v` did not run, which probably means that
            MuseScore is not accessible from the command line. Please first install MuseScore
            on your computer and then add it to your PATH."""
            ))
        end
    end
    global MUSESCORE_EXISTS[1] = true
end

"""
    musescore(file, notes | midi; display = true)
Use the open source software "MuseScore" to create a score and save the
output to `file`. By default it will also display the created `file`,
which can be either a `.pdf` or a `.png`.

If given a `.png` the actual file name will end with `-1`, `-2` etc.
for each page of the score. Notice that MuseScore must be accessible from the
command line for this function to work.
"""
function musescore(file, notes; display = true)

    @assert file[end-3:end] ∈ (".png", ".pdf")
    MUSESCORE_EXISTS[1] || test_musescore()

    tdir = tempdir()
    midi = writeMIDIfile(tdir*"/tempmid.mid", notes)
	midipath = joinpath(tdir, "tempmid.mid")

    if file[end-3:end] == ".png"
        cmd = `$MUSESCORE -n -T 10 -o $(file) $(midipath)`
        muspng = file[1:end-4]*"-1.png"
    else
        cmd = `$MUSESCORE -n -o $(file) $(midipath)`
        muspng = file
    end
	run(cmd)
    rm(tdir*"/tempmid.mid")
    display && DefaultApplication.open(muspng)
end

end#module

ext: jl


text:
""" The category of finite sets and relations, and its skeleton.
"""
module FinRelations
export BoolRig, FinRel, FinRelation, FinRelationCallable, FinRelationMatrix,
  force

import Base: +, *
using AutoHashEquals
import FunctionWrappers: FunctionWrapper

using ...GAT
using ...Theories: DistributiveBicategoryRelations
import ...Theories: dom, codom, id, compose, ⋅, ∘, dagger, dunit, dcounit,
  otimes, ⊗, munit, braid, oplus, ⊕, mzero, swap,
  mcopy, Δ, mmerge, ∇, delete, ◊, create, □, plus, zero, coplus, cozero,
  pair, copair, proj1, proj2, coproj1, coproj2, meet, join, top, bottom
import ..FinCats: force
using ..Matrices
using ..Matrices: zero_matrix

# Rig of booleans
#################

""" The rig of booleans.

This struct is needed because in base Julia, the product of booleans is another
boolean, but the sum of booleans is coerced to an integer: `true + true == 2`.
"""
@auto_hash_equals struct BoolRig <: Number
  value::Bool
end

+(x::BoolRig, y::BoolRig) = BoolRig(x.value || y.value)
*(x::BoolRig, y::BoolRig) = BoolRig(x.value && y.value)

Base.zero(::Type{BoolRig}) = BoolRig(false)
Base.one(::Type{BoolRig}) = BoolRig(true)
Base.float(x::BoolRig) = float(x.value)
Base.show(io::IO, x::BoolRig) = print(io, x.value)

# Category of finite relations
##############################

""" Object in the category of finite sets and relations.

See also: [`FinSet`](@ref).
"""
@auto_hash_equals struct FinRel{S,T}
  set::S
end

FinRel(i::Int) = FinRel{Int,Int}(i)
FinRel(set::S) where {T,S<:AbstractSet{T}} = FinRel{S,T}(set)

Base.eltype(::Type{FinRel{S,T}}) where {S,T} = T
Base.iterate(s::FinRel, args...) = iterate(iterable(s), args...)
Base.length(s::FinRel) = length(iterable(s))
Base.in(s::FinRel, elem) = in(s, iterable(s))
iterable(s::FinRel{Int}) = 1:s.set
iterable(s::FinRel{<:AbstractSet}) = s.set

""" Binary relation between finite sets.

A morphism in the category of finite sets and relations. The relation can be
represented implicitly by an arbitrary Julia function mapping pairs of elements
to booleans or explicitly by a matrix (dense or sparse) taking values in the rig
of booleans ([`BoolRig`](@ref)).
"""
abstract type FinRelation{S,S′,T,T′} end

FinRelation(R::Function, args...) = FinRelationCallable(R, args...)
FinRelation(R::AbstractMatrix, args...) = FinRelationMatrix(R, args...)

""" Relation in FinRel defined by a callable Julia object.
"""
@auto_hash_equals struct FinRelationCallable{S,S′,T,T′} <: FinRelation{S,S′,T,T′}
  # Field `rel` is usually a `Function` but can be any Julia callable.
  rel::FunctionWrapper{Bool,Tuple{T,T′}}
  dom::FinRel{S,T}
  codom::FinRel{S′,T′}

  FinRelationCallable(R, dom::FinRel{S,T}, codom::FinRel{S′,T′}) where {S,S′,T,T′} =
    new{S,S′,T,T′}(FunctionWrapper{Bool,Tuple{T,T′}}(R), dom, codom)
end

FinRelationCallable(R, dom::Int, codom::Int) =
  FinRelationCallable(R, FinRel(dom), FinRel(codom))

(R::FinRelationCallable{S,T})(x::T, y::T) where {S,T} = R.rel(x, y)

""" Relation in FinRel represented by a boolean matrix.

Boolean matrices are also known as logical matrices or relation matrices.
"""
@auto_hash_equals struct FinRelationMatrix{M<:AbstractMatrix{BoolRig}} <: FinRelation{Int,Int,Int,Int}
  rel::M
end

function FinRelationMatrix(R::AbstractMatrix{BoolRig}, dom::Int, codom::Int)
  @assert size(R) == (dom, codom)
  FinRelationMatrix(R)
end
function FinRelationMatrix(R::AbstractMatrix{BoolRig},
                           dom::FinRel{Int}, codom::FinRel{Int})
  FinRelationMatrix(R, length(dom), length(codom))
end

(R::FinRelationMatrix)(x, y) = R.rel[y, x].value

function force(::Type{T}, R::FinRelation) where T <: AbstractMatrix{BoolRig}
  m, n = length(dom(R)), length(codom(R))
  M = zero_matrix(T, n, m)
  for i in 1:m, j in 1:n
    if R(i,j)
      M[j,i] = BoolRig(true)
    end
  end
  FinRelationMatrix(M)
end
force(::Type{T}, R::FinRelationMatrix{T}) where T <: AbstractMatrix{BoolRig} = R
force(R::FinRelation) = force(Matrix{BoolRig}, R)

""" FinRel as a distributive bicategory of relations.

FIXME: Many methods only work for `FinRel{Int}`. The default methods should
assume `FinRel{<:AbstractSet}` and the case `FinRel{Int}` should be handled
specially.
"""
@instance DistributiveBicategoryRelations{FinRel, FinRelation} begin
  dom(R::FinRelation) = R.dom
  codom(R::FinRelation) = R.codom

  id(A::FinRel) = FinRelation((x,y) -> x == y, A, A)

  function compose(R::FinRelation, S::FinRelation)
    @assert codom(R) == dom(S)
    FinRelation(dom(R), codom(S)) do x,z
      any(R(x,y) && S(y,z) for y in codom(R))
    end
  end

  # Multiplicative monoidal structure

  otimes(A::FinRel, B::FinRel) = FinRel(length(A) * length(B))
  munit(::Type{FinRel}) = FinRel(1)

  function otimes(R::FinRelation, S::FinRelation)
    # Indexing is consistent with that of Kronecker products.
    dom_proj = CartesianIndices((length(dom(S)), length(dom(R))))
    cod_proj = CartesianIndices((length(codom(S)), length(codom(R))))
    FinRelation(dom(R)⊗dom(S), codom(R)⊗codom(S)) do x, y
      R(dom_proj[x][2], cod_proj[y][2]) && S(dom_proj[x][1], cod_proj[y][1])
    end
  end

  function braid(A::FinRel, B::FinRel)
    dom_proj = CartesianIndices((length(B), length(A)))
    cod_proj = CartesianIndices((length(A), length(B)))
    FinRelation(A⊗B, B⊗A) do x, y
      dom_proj[x][1] == cod_proj[y][2] && dom_proj[x][2] == cod_proj[y][1]
    end
  end

  function mcopy(A::FinRel)
    n = length(A)
    proj = CartesianIndices((n, n))
    FinRelation((x,y) -> x == proj[y][1] && x == proj[y][2], A, A⊗A)
  end
  function mmerge(A::FinRel)
    n = length(A)
    proj = CartesianIndices((n, n))
    FinRelation((x,y) -> proj[x][1] == y && proj[x][2] == y, A⊗A, A)
  end
  delete(A::FinRel) = FinRelation((x,y) -> true, A, munit(FinRel))
  create(A::FinRel) = FinRelation((x,y) -> true, munit(FinRel), A)

  dagger(R::FinRelation) = FinRelation((x,y) -> R(y,x), codom(R), dom(R))
  dunit(A::FinRel) = compose(create(A), mcopy(A))
  dcounit(A::FinRel) = compose(mmerge(A), delete(A))

  # Additive monoidal structure

  oplus(A::FinRel, B::FinRel) = FinRel(length(A) + length(B))
  mzero(::Type{FinRel}) = FinRel(0)

  function oplus(R::FinRelation, S::FinRelation)
    m, n = length(dom(R)), length(codom(R))
    FinRelation(dom(R)⊕dom(S), codom(R)⊕codom(S)) do x, y
      (x <= m && y <= n && R(x, y)) || (x > m && y > n && S(x-m, y-n))
    end
  end

  function swap(A::FinRel, B::FinRel)
    m, n = length(A), length(B)
    FinRelation((x,y) -> x == y-n || x-m == y, A⊕B, B⊕A)
  end

  function plus(A::FinRel)
    n = length(A)
    FinRelation((x,y) -> x == y || x-n == y, A⊕A, A)
  end
  function coplus(A::FinRel)
    n = length(A)
    FinRelation((x,y) -> x == y || x == y-n, A, A⊕A)
  end
  zero(A::FinRel) = FinRelation((x,y) -> false, mzero(FinRel), A)
  cozero(A::FinRel) = FinRelation((x,y) -> false, A, mzero(FinRel))

  pair(R::FinRelation, S::FinRelation) = coplus(dom(R))⋅(R⊕S)
  copair(R::FinRelation, S::FinRelation) = (R⊕S)⋅plus(codom(R))
  proj1(A::FinRel, B::FinRel) = id(A)⊕cozero(B)
  proj2(A::FinRel, B::FinRel) = cozero(A)⊕id(B)
  coproj1(A::FinRel, B::FinRel) = id(A)⊕zero(B)
  coproj2(A::FinRel, B::FinRel) = zero(A)⊕id(B)

  # Lattice structure

  function meet(R::FinRelation, S::FinRelation)
    @assert dom(R) == dom(S) && codom(R) == codom(S)
    FinRelation((x,y) -> R(x,y) && S(x,y), dom(R), codom(R))
  end
  function join(R::FinRelation, S::FinRelation)
    @assert dom(R) == dom(S) && codom(R) == codom(S)
    FinRelation((x,y) -> R(x,y) || S(x,y), dom(R), codom(R))
  end
  top(A::FinRel, B::FinRel) = FinRelation((x,y) -> true, A, B)
  bottom(A::FinRel, B::FinRel) = FinRelation((x,y) -> false, A, B)
end

# For relation matrices, delegate to the category of matrices.
const FinRelMat = FinRelationMatrix

dom(R::FinRelMat) = FinRel(size(R.rel, 2))
codom(R::FinRelMat) = FinRel(size(R.rel, 1))

compose(R::FinRelMat, S::FinRelMat) = FinRelMat(compose(R.rel, S.rel))
otimes(R::FinRelMat, S::FinRelMat) = FinRelMat(otimes(R.rel, S.rel))
oplus(R::FinRelMat, S::FinRelMat) = FinRelMat(oplus(R.rel, S.rel))
dagger(R::FinRelMat) = FinRelMat(transpose(R.rel))

meet(R::FinRelMat, S::FinRelMat) = FinRelMat(R.rel .* S.rel)
join(R::FinRelMat, S::FinRelMat) = FinRelMat(R.rel .+ S.rel)
pair(R::FinRelMat, S::FinRelMat) = FinRelMat(pair(R.rel, S.rel))
copair(R::FinRelMat, S::FinRelMat) = FinRelMat(copair(R.rel, S.rel))

end


text:
using Revise
using Test
using NNFEM
using PyCall
using PyPlot
using JLD2
using ADCME
using MAT
using LinearAlgebra

include("nnutil.jl")
nntype = "piecewise"

# * Auto-generated code by `ae_to_code`
# aedictae_scaled = matread("Data/learned_nn.mat"); # using MAT
aedictae_scaled = matread("Data/train_neural_network_from_fem.mat")
Wkey = "$(nntype)backslashfully_connectedbackslashweightscolon0"
Wkey = "$(nntype)backslashfully_connected_1backslashweightscolon0"
Wkey = "$(nntype)backslashfully_connected_2backslashweightscolon0"
Wkey = "$(nntype)backslashfully_connected_3backslashweightscolon0"
Wkey = "$(nntype)backslashfully_connected_4backslashweightscolon0"
function nnae_scaled(net)
        # @show size(net)
        # return reshape(reshape(net[1:3],1,3)*H0,1,3)/stress_scale
        W0 = aedictae_scaled["$(nntype)backslashfully_connectedbackslashweightscolon0"]; b0 = aedictae_scaled["$(nntype)backslashfully_connectedbackslashbiasescolon0"];
        isa(net, Array) ? (net = net * W0 .+ b0') : (net = net *W0 + b0)
        isa(net, Array) ? (net = tanh.(net)) : (net=tanh(net))
        W1 = aedictae_scaled["$(nntype)backslashfully_connected_1backslashweightscolon0"]; b1 = aedictae_scaled["$(nntype)backslashfully_connected_1backslashbiasescolon0"];
        isa(net, Array) ? (net = net * W1 .+ b1') : (net = net *W1 + b1)
        isa(net, Array) ? (net = tanh.(net)) : (net=tanh(net))
        W2 = aedictae_scaled["$(nntype)backslashfully_connected_2backslashweightscolon0"]; b2 = aedictae_scaled["$(nntype)backslashfully_connected_2backslashbiasescolon0"];
        isa(net, Array) ? (net = net * W2 .+ b2') : (net = net *W2 + b2)
        isa(net, Array) ? (net = tanh.(net)) : (net=tanh(net))
        W3 = aedictae_scaled["$(nntype)backslashfully_connected_3backslashweightscolon0"]; b3 = aedictae_scaled["$(nntype)backslashfully_connected_3backslashbiasescolon0"];
        isa(net, Array) ? (net = net * W3 .+ b3') : (net = net *W3 + b3)
        isa(net, Array) ? (net = tanh.(net)) : (net=tanh(net))
        W4 = aedictae_scaled["$(nntype)backslashfully_connected_4backslashweightscolon0"]; b4 = aedictae_scaled["$(nntype)backslashfully_connected_4backslashbiasescolon0"];
        isa(net, Array) ? (net = net * W4 .+ b4') : (net = net *W4 + b4)
        return net
end

# testtype = "PlaneStressPlasticity"
testtype = "NeuralNetwork2D"
prop = Dict("name"=> testtype, "rho"=> 8000.0, "E"=> 200e+9, "nu"=> 0.45,
"sigmaY"=>0.3e+9, "K"=>1/9*200e+9, "nn"=>post_nn)
include("NNPlatePull_Domain.jl")

@load "Data/domain.jld2" domain
stress_scale = 1.0e10
X, Y = prepare_strain_stress_data2D(domain)
x = constant(X)
y = nn(X[:,1:3], X[:,4:6], X[:,7:9])

sess = Session(); init(sess)
# ADCME.load(sess, "Data/learned_nn.mat")
ADCME.load(sess, "Data/train_neural_network_from_fem.mat")

O = run(sess, y)
using Random; Random.seed!(233)
VisualizeStress2D(Y, O, 20)


text:
Pkg.add("VennDiagrams")

using VennDiagrams

c1 = [0:5] # Can be any iterable
c2 = [3:10]

p = venn(c1, c2)
draw(PDF("beaut1.pdf", 8cm, 10cm), p)

c3 = [0:2:12]

p = venn(c1, c2, c3, proportional=false)
draw(PDF("beaut3.pdf", 8cm, 10cm), p)

max_stars_repo_name: Ellon-M/visualizations

max_issues_repo_name: Ellon-M/visualizations

max_forks_repo_name: Ellon-M/visualizations


text:
using Images

img = load("sample-images/cats-3061372_640.jpg")
save("cats-small-file-size.jpg", img) # save file in JPG format
save("cats-high-quality.png", img) # save file in PNG format


text:
using StaticArrays
using JuLIP: Atoms, neighbourlist, cutoff
using NeighbourLists: max_neigs, sites
using LinearAlgebra: norm

"""
Store a TB Hamiltonian in RI format (i.e. the bond integrals)
"""
mutable struct SKHamiltonianRI{T}
   skh::SKH
   model::SKModel
   Nat::Int
   # ----------- off-diagonal blocks
   Iat::Vector{Int}            # first atom index i
   Jat::Vector{Int}            # second atom index j
   R::Vector{SVector{3,T}}     # X[j] - X[i]
   V::Matrix{T}                # bond integral values
   # ----------- diagonal blocks (require they  are diagonal for now!)
   D::Matrix{T}                # diagonal entries Nat x norbitals
end

Base.length(H::SKHamiltonianRI) = length(H.Iat)

# -----
# Prototypes

function eval_bonds! end
function eval_diag! end
function filter_nhd end
function getskh end

# ----------------------------------------------------------------
#    Two-Centre Case

function assembleRI(at::Atoms, model::TwoCentreModel)
   nlist = neighbourlist(at, cutoff(model))
   skh = getskh(model)

   # copy / allocate space
   Nat = length(at)
   Iat = copy(nlist.i)
   Jat = copy(nlist.j)
   # R = copy(nlist.R)
   R = [ ijR[3] for ijR in pairs(nlist) ]
   V = zeros(length(R), nbonds(skh))

   # call the function barriers
   temp = MVector(zeros(nbonds(skh))...)
   _assembleRI!(V, R, model, temp)

   # and for the diagonal entries ...
   D = zeros(Nat, nbonds(skh))
   _assembleRI_diag!(D, model, nlist, temp)

   return SKHamiltonianRI(skh, model, Nat, Iat, Jat, R, V, D)
end

function _assembleRI!(V, R, model::TwoCentreModel, temp)
   for n = 1:length(R)
      eval_bonds!(temp, model, norm(R[n]))
      V[n, :] .= temp
   end
   return V
end

# TODO: at the moment, this is generic -> next allow SKmodel to have
#       general diagonal blocks
function _assembleRI_diag!(D, model::SKModel, nlist, temp)
   for (i, j, r, R) in sites(nlist)
      eval_diag!(temp, model, r, R)
      D[i, :] .= temp
   end
   return D
end

# ----------------------------------------------------------------
#    General Case

function assembleRI(at::Atoms, model::SKModel)
   nlist = neighbourlist(at, offsite_cutoff(model))
   skh = getskh(model)

   # copy / allocate space
   Nat = length(at)
   nb = nbonds(skh)
   Iat = zeros(Int, nb)
   Jat = zeros(Int, nb)
   R = zeros(JVecF, nb)
   V = zeros(length(R), nb)

   # call the function barrier
   temp = MVector(zeros(nbonds(skh))...)
   _assembleRI!(Iat, Jat, R, V, nlist, model, temp)

   # and for the diagonal entries ...
   D = zeros(Nat, nbonds(skh))
   _assembleRI_diag!(D, model, nlist, temp)

   return SKHamiltonianRI(skh, model, Nat, Iat, Jat, R, V, D)
end

function _assembleRI!(Iat, Jat, R, V, nlist, model, temp)
   idx = 0
   R_filtered = sizehint!(JVecF[], max_neigs(nlist))  # allocate without growing the array
   for (i, J, r, R) in sites(nlist)
      for (nj, j) in enumerate(J)
         Rji  = R[j]  # R[j] = X[j] - X[i]
         # filter -> Ks, Rs = [X[k]-X[i] for k in ...]
         for k in Jh
            Rki = R[k]
            if (k != j) && filter_nhd(Rji, Rki)
               push!(R_filtered, Rki)
            end
         end
         # evaluate the bonds
         fill!(temp, 0.0)
         eval_bonds!(temp, model, Rji, R_filtered)
         empty!(R_filtered)    # this keeps the memory allocated
         # write into V
         idx += 1
         V[idx, :] .= temp
      end
   end
   return V
end

# ----------------------------------------------------------------
#    Assemble the Hamiltonian Matrix in Cartesian format

"""
number of local orbitals
"""
get_norbloc(H::SKH) = sum( 2*get_l(o)+1 for o in H.orbitals )

"""
construct global Hamiltonian matrix indices from atom indices and
local orbital indices
"""
global_index(iat, iloc, nat, norbloc) = (iat-1) * norbloc .+ iloc

function local_indices(nb::Integer, skh::SKH)
   io1, io2 = skh.b2o[nb] # orbital indices
   I1, I2 = skh.locorbidx[io1], skh.locorbidx[io2]   # local block indices
   # convert to SVectors
   return SVector(I1...), SVector(I2...)
end

function Base.Matrix(HRI::SKHamiltonianRI{T}) where {T}
   # ------------------------------------------------------------------------
   # this "outer" matrix assembly is intentionally type unstable; it prepares
   # everything needed for a fast assembly in _assemble_full_inner!, which is the
   # function barrier beyond which type stability should be strictly enforced
   # ------------------------------------------------------------------------
   norbloc = get_norbloc(HRI.skh)
   norbglob = HRI.Nat * norbloc

   # allocate the global hamiltonian matrix
   H = zeros(norbglob, norbglob)

   for nb = 1:length(HRI.skh.bonds)
      # compute the local orbital indices corresponding to the current bond type
      # TODO make sure these indices are SVectors
      irowloc, jcolloc = local_indices(nb, HRI.skh)
      @assert ( (irowloc isa SVector{N, <: Integer} where {N}) &&
                (jcolloc isa SVector{N, <: Integer} where {N}) )
      # call the function barrier
      _assemble_full_inner!(H, HRI, nb, HRI.skh.bonds[nb], irowloc, jcolloc, norbloc)
   end

   _assemble_full_inner_diag!(H, HRI, SVector((1:norbloc)...))  # TODO: split into TwoCentre / General

   return H
end

# TODO: some ideas to speed this up
#    - irowloc, jcolloc could become type parameters
#    - the global indices could be incremented rather than recomputed
#    - compute only upper triangular part, copy the lower triangular part
function _assemble_full_inner!(H, HRI, nb, b, irowloc, jcolloc, norbloc)
   sig = sksign(b)
   sigt = sksignt(b)
   for (n, (iat, jat, R)) in enumerate( zip(HRI.Iat, HRI.Jat, HRI.R) )
      V = HRI.V[n, nb]
      U = R / norm(R)
      φ, θ = carttospher(U[1], U[2], U[3])
      E12 = CodeGeneration.sk_gen(b, φ, θ)
      irowglob = global_index(iat, irowloc, HRI.Nat, norbloc)
      jcolglob = global_index(jat, jcolloc, HRI.Nat, norbloc)
      H[irowglob, jcolglob] += sig * V * E12
      if irowloc != jcolloc # non-diagonal blocks must be copied
         irowglobt = global_index(iat, jcolloc, HRI.Nat, norbloc)
         jcolglobt = global_index(jat, irowloc, HRI.Nat, norbloc)
         H[irowglobt, jcolglobt] += sigt * V * E12'
      end
   end
   return H
end

function _assemble_full_inner_diag!(H, HRI, iloc)
   norbloc = length(iloc)
   for nat = 1:HRI.Nat
      iglob = global_index(nat, iloc, HRI.Nat, norbloc)
      for j = 1:norbloc
         H[iglob[j], iglob[j]] = HRI.D[nat, j]
      end
   end
   return H
end

# TODO:
# - k-dependence
# - overlaps


text:
# class constructors
# convenience type constructor:
# enable `R[:x,:y,:z,...]` as short for `MultivariatePolynomial{R,N,Id}`
function getindex(R::Type{<:Ring}, s::Symbol, t::Symbol...)
    vs = collect((s, t...))
    N = length(vs)
    Id = sintern(vs)
    new_class(MultivariatePolynomial{R,N,Id,Int,Tuple{N}}, vs)
end
function getindex(R::Type{<:Ring}, s::AbstractVector{Symbol}, t::AbstractVector{Symbol}...)
    blocks = collect((s, t...))
    construct(R, blocks)
end

function lextend(::Type{P}, s::Symbol, t::Symbol...) where {R,P<:MultivariatePolynomial{R}}
    blocks = [collect((s, t...)), varblocks(P)...]
    construct(R, blocks)
end

function construct(::Type{R}, blocks) where R<:Ring
    vs = vcat(blocks...)
    N = length(vs)
    Id = sintern(vs)
    n = length(blocks)
    T = n == 1 ? Int : NTuple{n,Int}
    B = Tuple{[length(x) for x in blocks]...}
    new_class(MultivariatePolynomial{R,N,Id,T,B}, vs)
end

import Base: copy, promote_rule
import Base: +, -, *, zero, one, ==, hash, isless, iszero, isone

# make new copy
copy(p::P) where P<:MultivariatePolynomial = P(copy(p.ind), copy(p.coeff))

# promotion and conversion
_promote_rule(::Type{<:MultivariatePolynomial{R,M,X}}, ::Type{<:Polynomial}) where {X,M,R} = Base.Bottom
_promote_rule(::Type{MultivariatePolynomial{R,N,X,T,B}}, ::Type{MultivariatePolynomial{S,N,X,T,B}}) where {X,N,R,S,T,B} = MultivariatePolynomial{promote_type(R,S),N,X,T,B}
_promote_rule(::Type{P}, ::Type{S}) where {R,N,X,B,T,S<:Ring,P<:MultivariatePolynomial{R,N,X,T,B}} = MultivariatePolynomial{promote_type(R,S),N,X,T,B}
promote_rule(::Type{P}, ::Type{S}) where {R,N,X,T,B,S<:Union{Integer,Rational},P<:MultivariatePolynomial{R,N,X,T,B}} = MultivariatePolynomial{promote_type(R,S),N,X,T,B}

function (P::Type{MultivariatePolynomial{R,N,X,T,B}})(a::MultivariatePolynomial{S,N,X,T,B}) where {R,N,X,T,B,S}
    P(a.ind, R.(a.coeff))
end

function (P::Type{<:MultivariatePolynomial{R,N,X,T}})(a::MultivariatePolynomial{S}) where {R,N,X,T,S}
    deg(a) <= 0 && return P(LC(a))
    vp = varnames(P)
    va = varnames(a)
    pos = positionsin(va, vp)
    n = length(a.ind)
    pind = Vector{T}(undef, n)
    for i = 1:n
        xa = index2expo(a, i)
        checkpositions(pos, xa, va, vp)
        xp = reindex2(xa, pos, N)
        pind[i] = expo2ordblock(P, xp)
    end
    perm = sortperm(pind)
    pind = pind[perm]
    pc = a.coeff[perm]
    P(pind, pc)
end
function (P::Type{<:MultivariatePolynomial{R,N,X,T}})(a::UnivariatePolynomial{S}) where {R,N,X,T,S}

    deg(a) <= 0 && return P(LC(a))
    vp = varnames(P)
    va = varnames(a)
    issubset(va, vp) || throw(ArgumentError("Variable :$(va[1]) not contained in $vp"))
    pos = positionsin(va, vp)
    ac = a.coeff
    n = length(ac)
    pind = Vector{T}(undef, n)
    perm = Vector{Int}(undef, n)
    j = 0
    for i = 1:n
        aci = ac[i]
        if !iszero(aci)
            j += 1
            xa = [i-1]
            perm[j] = i
            xp = reindex2(xa, pos, N)
            pind[j] = expo2ordblock(P, xp)
        end
    end
    resize!(perm, j)
    resize!(pind, j)
    P(pind, ac[perm])
end

function (P::Type{<:MultivariatePolynomial{S}})(a::S) where S
    iszero(a) ? zero(P) : P(one(P).ind, [a])
end
function (P::Type{<:MultivariatePolynomial{S}})(a::T) where {S,T}
    iszero(a) ? zero(P) : P(one(P).ind, [S(a)])
end

deg(p::MultivariatePolynomial) = isempty(p.ind) ? -1 : sum(multideg(p))
isunit(a::MultivariatePolynomial) = deg(a) == 0 && isunit(a.coeff[1])
ismonom(p::MultivariatePolynomial) = length(p.ind) <= 1

"""
    derive(p::MultivariatePolynomial, (d1,...dN))

Derive polynomial with N variables with respect to variables `1` ... `N` with the
respective degrees `d1` ... `dN`.
"""
function derive(p::P, d::NTuple{N,<:Integer}) where {S,N,P<:MultivariatePolynomial{S,N}}
    all(d .>= 0) || throw(ArgumentError("negative derivation degrees not allowed"))
    ind = similar(p.ind)
    coeff = similar(p.coeff)
    j = 0
    for i = 1:length(p.ind)
        a = index2expo(p, i)
        b = a .- d
        if all(b .>= 0)
            c = p.coeff[i]
            for k = 1:N
                for f = a[k]-d[k]+1:a[k]
                    c *= f
                end
            end
            if !iszero(c)
                j += 1
                coeff[j] = c
                ind[j] = expo2ordblock(P, b)
            end
        end
    end
    resize!(coeff, j)
    resize!(ind, j)
    P(ind, coeff)
end

"""
    monom(<:MultivariatePolynomial, expos::Vector{Int})

Return monic monomial with given exponents. (`[1,0,...]` corresponds to first variable).
"""
function monom(P::Type{<:MultivariatePolynomial{S,N}}, xv::Vector{<:Integer}) where {N,S}
    n = length(xv)
    n == 0 && return zero(P)
    length(xv) != N && throw(ArgumentError("multivariate monom needs exponents for all $N variables"))
    P([expo2ordblock(P, xv)], [1])
end

"""
    generators(P)

Return an array of all monomial generators (variables) of a polynomial type `P`.
"""
generators(P::Type{<:UnivariatePolynomial}) = [monom(P, 1)]
function generators(P::Type{<:MultivariatePolynomial{S,N}}) where {S,N}
    v = zeros(Int, N)
    res = Vector{P}(undef, N)
    for i = 1:N
        v[i] = 1
        res[i] = monom(P, v)
        v[i] = 0
    end
    res
end

function multideg(p::MultivariatePolynomial{S,N}) where {S,N}
    isempty(p.ind) ? Int[] : index2expo(p, length(p.ind))
end

function LM(p::P) where {S,N,P<:MultivariatePolynomial{S,N}}
    n = length(p.ind)
    n == 0 && return zero(P)
    P([p.ind[n]], S[1])
end

function LT(p::P) where {S,N,P<:MultivariatePolynomial{S,N}}
    n = length(p.ind)
    n == 0 && return zero(P)
    P([p.ind[n]], [p.coeff[n]])
end

"""
    index2expo(p::Polynomial, i::Integer)

Return the tuple of variable exponents stored at this index (of `p.ind`).
"""
function index2expo(p::MultivariatePolynomial{S,N,X,<:Integer}, i::Integer) where {S,N,X}
    ord2expo(p.ind[i], N)
end

function index2expo(p::MultivariatePolynomial{S,N,X,<:Tuple,B}, i::Integer) where {S,N,X,B<:Tuple}
    vp = tupcon(B)
    m = length(vp)
    vcat([ord2expo(p.ind[i][k], vp[k]) for k = 1:m]...)
end

# arithmetic
function zero(::Type{<:T}) where {S,T<:MultivariatePolynomial{S}}
    T(Int[], S[])
end
function one(::Type{<:T}) where {S,T<:MultivariatePolynomial{S}}
    T([zeroindex(T)], S[1])
end

-(p::T) where T<:MultivariatePolynomial = T(p.ind, -p.coeff)
-(a::T, b::T) where T<:MultivariatePolynomial = +(a, -b)
function *(p::T, a::Integer) where T<:MultivariatePolynomial
    iszero(a) && return zero(T)
    isone(a) && return p
    pc = p.coeff .* a
    pind = copy(p.ind)
    for i = length(pc):-1:1
        if iszero(pc[i])
            deleteat!(pind, i)
            deleteat!(pc, i)
        end
    end
    T(pind, pc)
end

*(a::Integer, p::T) where T<:MultivariatePolynomial = p * a
==(a::T, b::T) where T<:MultivariatePolynomial = a.ind == b.ind && a.coeff == a.coeff
function hash(a::MultivariatePolynomial, h::UInt)
    n = deg(a)
    n < 0 ? hash(0, h) : n == 0 ? hash(LC(a), h) : hash(a.ind, hash(a.coeff, h))
end
function isless(a::T, b::T) where T<:MultivariatePolynomial
    m = length(a.ind)
    n = length(b.ind)
    if m == 0
        n != 0
    elseif n == 0
        false
    else
        a.ind[m] < b.ind[n]
    end
end

function +(a::T...) where T<:MultivariatePolynomial
    n = length(a)
    n >  0 || throw(ArgumentError("+ requires at least one argument"))
    n == 1 && return a[1]
    c = similar(a[1].coeff)
    d = similar(a[1].ind)
    j = 0
    p = ones(Int, n)
    pm = [getindex(x, 1) for x in a]
    bound = maxindex(T)

    while true
        m, imin = findmin(pm)
        m == bound && break
        cj = a[imin].coeff[p[imin]]
        p[imin] += 1
        pm[imin] = getindex(a[imin], p[imin])
        for i = imin+1:n
            if pm[i] == m
                cj += a[i].coeff[p[i]]
                p[i] += 1
                pm[i] = getindex(a[i], p[i])
            end
        end
        if !iszero(cj)
            j += 1
            if j > length(d)
                resize!(d, 2*j)
                resize!(c, 2*j)
            end
            d[j] = m
            c[j] = cj
        end
    end
    resize!(c, j)
    resize!(d, j)
    T(d, c)
end

function *(a::T, b::T) where {N,S,T<:MultivariatePolynomial{S,N}}
    m = length(a.ind)
    n = length(b.ind)
    m >= n || return *(b, a)
    n == 0 && return zero(T)

    c = similar(a.coeff)
    d = similar(a.ind)
    j = 0
    p = ones(Int, n)
    pm = [exposum(a, 1, b, j) for j in 1:n]
    bound = maxindex(T)

    while true
        min, imin = findmin(pm)
        min == bound && break
        cj = a.coeff[p[imin]] * b.coeff[imin]
        p[imin] += 1
        pm[imin] = exposum(a, p[imin], b, imin)
        for i = imin+1:n
            if pm[i] == min
                cj += a.coeff[p[i]] * b.coeff[i]
                p[i] += 1
                pm[i] = exposum(a, p[i], b, i)
            end
        end
        if !iszero(cj)
            j += 1
            if j > length(d)
                resize!(d, 2*j)
                resize!(c, 2*j)
            end
            d[j] = min
            c[j] = cj
        end
    end
    resize!(c, j)
    resize!(d, j)
    T(d, c)
end

(p::MultivariatePolynomial)(a, b...) = evaluate(p, a, b...)
function evaluate(p::T, a::Union{Ring,Int,Rational}...) where {N,S,T<:MultivariatePolynomial{S,N}}
    length(a) != N && throw(ArgumentError("wrong number of arguments of polynomial with $N variables"))
    n = length(p.ind)
    R = promote_type(S, typeof.(a)...)
    deg(p) < 0 && return zero(R)
    deg(p) == 0 && return R(p.coeff[1])
    vdeg = maximum(hcat(index2expo.(p, 1:n)...), dims=2)
    xpot = [Vector{R}(undef, vdeg[i]) for i = 1:N]
    # precalculate all required monoms.
    for i = 1:N
        m = vdeg[i]
        if m > 0
            ai = bi = a[i]
            xpot[i][1] = bi
            for k = 2:m
                bi *= ai
                xpot[i][k] = bi
            end
        end
    end
    s = zero(R)
    for j = 1:n
        ex = index2expo(p, j)
        t = p.coeff[j]
        for i = 1:N
            if ex[i] > 0
                t *= xpot[i][ex[i]]
            end
        end
        s += t
    end
    s
end
"""
    expo2ordblock(::Type{<:Polynomial}, a::Vector{Int})

Return `Integer` or tuple of `Integer` representing the monomial ordering.
In the first case, or if the tuple has length one, that is the `:grevlex` order,
otherwise a block order, based on `:grevlex`, including `:lex`, if each variable
block consists of one variable.
"""
function expo2ordblock(::Type{P}, a::AbstractVector{<:Integer}) where {R,N,X,T,P<:MultivariatePolynomial{R,N,X,T,Tuple{N}}}
    expo2ord(a)
end

function expo2ordblock(::Type{P}, a::AbstractVector{<:Integer}) where {R,N,X,T,M,B,P<:MultivariatePolynomial{R,N,X,NTuple{M,T},B}}

    t = tupcon(B)
    res = Vector{T}(undef, M)
    j = 0
    for i = 1:M
        d = t[i]
        res[i] = expo2ord(a[j+1:j+d])
        j += d
    end
    tuple(res...)
end

# extract constants from Tuple{1,2,3...}
tupcon(::Type{Tuple{A}}) where A = (A,)
tupcon(::Type{Tuple{A,B}}) where {A,B} = (A, B)
tupcon(A::Type{<:Tuple}) = tuple(A.parameters...)

#= Tuple mappings. See for reference:
https://stackoverflow.com/questions/26932409/compact-storage-coefficients-of-a-multivariate-polynomial
https://en.wikipedia.org/wiki/Combinatorial_number_system

This constructs the degrevlex total ordering of monomials.
=#

"""
    expo2ord(a)

bijective mapping from tuples of non-negative integers to positive integers.

The induced order of tuples is `degrevlex`.
"""
function expo2ord(a::AbstractVector{<:Integer})
    c = similar(a)
    d = length(a)
    d == 0 && return c + 1
    ci = s = sum = a[1]
    for i = 2:d
        s += a[i]
        ci = s + i - 1
        sum += binomial(ci, i)
    end
    sum + 1
end

"""
    ord2expo(n, d)

bijective mapping from integers to `d`-tuples of integers.

The induced order of tuples is `degrevlex`.
"""
function ord2expo(n::T, d::Int) where T<:Integer
    c = Vector{T}(undef, d)
    n < 1 && throw(ArgumentError("index must be positive but is $n"))
    n -= 1
    for i = d:-1:1
        ci, b = cbin(n, i)
        c[i] = ci
        n -= b
    end
    ci = c[d]
    for i = d:-1:2
        cp = c[i-1]
        c[i] = ci - cp - 1
        ci = cp
    end
    c
end

"""
    cbin(n, d)

Calculate the greatest integer `c` such that `binomial(c, d) <= n`
"""
function cbin(n::T, d::Int) where T<:Integer
    d <= 0 && throw(ArgumentError("tuple size > 0 required, but is $d"))
    d == 1 && return n, n
    n == 0 && return d-1, n
    n == 1 && return d, n
    c = T(floor((n * sqrt(2pi*d))^(1/d) * d / ℯ + d/2))
    if c <= d
        c = d
        b = T(1)
    else
        b = binomial(c, d)
    end
    b == n && return c, b
    bp = b
    while b <= n
        bp = b
        c += 1
        b = b * c ÷ (c-d)
    end
    b >= n > bp && return c-1, bp
    while b > n
        bp = b
        b = b * (c-d) ÷ c
        c -= 1
    end
    return c, b
end

# the isless function for the `degrevlex` ordering of monomial exponents
function degrevlex(a::V, b::V) where V<:AbstractVector{<:Integer}
    length(a) != length(b) && throw(ArgumentError("lengths of vectors must match"))
    dega = sum(a)
    degb = sum(b)
    dega != degb && return dega < degb
    for j = length(a):-1:1
        aj = a[j]
        bj = b[j]
        if aj != bj
            return bj > aj # mind the reversal!
        end
    end
    false
end

# multiplication of monomials
function indexsum(x::T, y::T, d::Int) where T<:Integer
    x > 0 && y > 0 || return 0
    expo2ord(ord2expo(x, d) + ord2expo(y, d))
end

zeroindex(P::Type{<:MultivariatePolynomial}) = fillindex(one, P)
maxindex(P::Type{<:MultivariatePolynomial}) = fillindex(typemax, P)

function fillindex(f, ::Type{<:P}) where {R,N,X,T,P<:MultivariatePolynomial{R,N,X,T,Tuple{N}}}
    f(T)
end
function fillindex(f, ::Type{<:P}) where {R,N,X,T,M,P<:MultivariatePolynomial{R,N,X,NTuple{M,T}}}
    ft = f(T)
    ntuple(x->ft, M)
end

"""
    exposum(a::Polynomial, i, b::Polynomial, j)

Return the sum of variable exponents at `a.ind[i]` and `b.ind[j].
Realizes multiplication of monomials.
If one of the coefficients is undefined, return `maxindex` symbolizing zero.
"""
function exposum(pa::P, i::Integer, pb::P, j::Integer) where {R,N,X,T,P<:MultivariatePolynomial{R,N,X,T,Tuple{N}}}
    a = pa.ind
    b = pb.ind
    (isassigned(a, i) && isassigned(b, j)) || return maxindex(P)
    indexsum(a[i], b[j], N)
end

function exposum(pa::P, i::Integer, pb::P, j::Integer) where {R,N,X,T,B,P<:MultivariatePolynomial{R,N,X,T,B}}
    a = pa.ind
    b = pb.ind
    (isassigned(a, i) && isassigned(b, j)) || return maxindex(P)
    ai = a[i]
    bi = b[j]
    vp = tupcon(B)
    ntuple(k->indexsum(ai[k], bi[k], vp[k]), length(ai))
end

function getindex(pa::P, i::Integer) where P<:MultivariatePolynomial
    isassigned(pa.ind, i) ? pa.ind[i] : maxindex(P)
end

function divides(x::V, y::V) where V<:AbstractVector{<:Integer}
    all(x .<= y)
end

"""
    varnames(P)

Return array of variable names of polynomial or polynomial type `P`.
"""
varnames(::Type{T}) where {R,X,T<:UnivariatePolynomial{R,X}} = Symbol[X]
varnames(::Type{T}) where T<:MultivariatePolynomial = gettypevar(T).varnames
varnames(p::T) where T<:Polynomial = varnames(T)
varnames(::Type{T}, n::Integer) where T<:MultivariatePolynomial = varblocks(T)[n]

function varblocks(::Type{P}) where {R,N,X,T,B,P<:MultivariatePolynomial{R,N,X,T,B}}
    vars = varnames(P)
    t = tupcon(B)
    M = length(t)
    res = Vector{Vector{Symbol}}(undef, M)
    j = 0
    for i = 1:M
        d = t[i]
        res[i] = vars[j+1:j+d]
        j += d
    end
    res
end

function showvar(io::IO, var::MultivariatePolynomial{S,N}, n::Integer) where {N,S}
    ex = index2expo(var, n)
    vn = varnames(var)
    start = true
    for i = 1:N
        x = ex[i]
        x <= 0 && continue
        !start && print(io, '*')
        print(io, vn[i])
        x > 1 && print(io, '^', x)
        start = false
    end
end

function isconstterm(p::P, n::Integer) where P<:MultivariatePolynomial
    n <= 0 || p.ind[n] == zeroindex(P)
end

divrem(f::P, id::Ideal{P}) where P<:Polynomial = divrem(f, id.base)

function divrem(f::P, g::P) where P<:MultivariatePolynomial
    if 0 == deg(f) == deg(g)
        return P.(divrem(LC(f), LC(g)))
    end
    a, s, d = pdivrem(f, [g])
    isunit(d) ? (a[1], s) : (zero(P), f)
end
function divrem(f::P, g::AbstractVector{P}) where P<:MultivariatePolynomial
    a, s, d = pdivrem(f, g)
    isunit(d) ? (a, s) : (zeros(P, length(g)), f)
end

# division and Gröbner base calculation

function pdivrem(f::P, g::P) where {T,N,P<:MultivariatePolynomial{T,N}}

    lig = multideg(g)
    xif = 0
    lif = lig
    for i = length(f.ind):-1:1
        li = index2expo(f, i)
        if divides(lig, li)
            lif = li
            xif = i
            break
        end
    end
    xif == 0 && return zero(P), f, one(T)
    c = f.coeff[xif]
    d = LC(g)
    q = monom(P, lif .- lig)
    if isone(d)
        k = q * c
        k, f - g * k, one(T)
    elseif isunit(d)
        k = q * (c / d)
        k, f - g * k, one(T)
    else
        x = gcd(c, d)
        h = d / x
        k = q * (c / x)
        k, f * h - g * k, h
    end
end

function pdivrem(f::P, G::AbstractVector{P}) where {T,P<:MultivariatePolynomial{T}}
    f0 = f
    fp = zero(P)
    a = zeros(P, length(G))
    dd = one(T)
    while f !== fp && !iszero(f)
        fp = f
        for (i, g) in enumerate(G)
            ffp = f
            k, f, d = pdivrem(f, g)
            if f !== ffp
                if !isone(d)
                    a .*= d
                    dd *= d
                end
                a[i] += k
            elseif iszero(f)
                break
            end
        end
    end
    a, f, dd
end

"""
    SPOL(f, g)

Calculate the S-polynomial (`SPOL`) of the polynomials `f`, `g`.
"""
function SPOL(f::P, g::P) where P<:MultivariatePolynomial
    lcf = LC(f)
    lcg = LC(g)
    lif = multideg(f)
    lig = multideg(g)

    h = gcd(lcf, lcg)
    af = lcf / h
    ag = lcg / h
    bf = max.(lif - lig, 0)
    bg = max.(lig - lif, 0)
    monom(P, bg) * ag * f - monom(P, bf) * af * g
end

using Base.Iterators

# find all in one
function buchbergerall(f::AbstractVector{P}) where P<:MultivariatePolynomial
    n = length(f)
    C = [(i,j) for i=1:n for j = i+1:n]
    buchberger(f, C)
end

# incrementally adding vectors
function buchberger(A::AbstractVector{P}) where P<:MultivariatePolynomial
    n = length(A)
    n == 0 && return P[]
    a = A[1:1]
    for i = 2:n
        a = buchberger(a, A[i:i])
    end
    a
end

# assume f is already gröbner base
function buchberger(f::VP, h::VP) where {P<:MultivariatePolynomial,VP<:AbstractVector{P}}
    n = length(f)
    m = length(h)
    g = vcat(f, h)
    C = [(i,n+j) for i = 1:n for j = 1:m]
    buchberger(g, C)
end

# Buchberger's algorithm - C contains the remaining critical pairs (i,j)
function buchberger(f::AbstractVector{P}, C::Vector{Tuple{Int,Int}}) where P<:MultivariatePolynomial
    n = length(f)
    g = copy(f)
    while !isempty(C)
        k = select_critical_pair(C, g)
        i, j = C[k]
        p, q = g[i], g[j]
        deleteat!(C, k)
        criterion(g, C, i, j) && continue
        pq = SPOL(p, q)
        a, s, d = pdivrem(pq, g)
        if !iszero(s) && isone(d)
            push!(g, s)
            n += 1
            append!(C, [(i,n) for i = 1:n-1])
            cleanup!(C, g)
        end
    end
    g
end

function select_critical_pair(C::AbstractVector{Tuple{Int,Int}}, g::AbstractVector)
    n = length(C)
    kmin = 1
    return kmin
    degp(k) = sum(max.((multideg(g[C[k][1]])), multideg(g[C[k][2]])))
    dmin = degp(1)
    for k = 2:n
        dk = degp(k)
        if dk < dmin
            dmin = dk
            kmin = k
        end
    end
    kmin
end

function criterion(G::AbstractVector{<:Polynomial}, C::AbstractVector, i::Int, k::Int)
    f = G[i]
    g = G[k]
    fx = multideg(f)
    gx = multideg(g)
    sum(fx .* gx) == 0 && return true # product criterion - no powers in common
    for j = 1:length(G)
        (j == i || j == k) && continue
        hx = multideg(G[j])
        if all(hx .<= max.(fx, gx)) && !in(minmax(i, j), C) && !in(minmax(k,j), C)
            return true # chain criterion - already removed
        end
    end
    return false
end

function cleanup!(C::AbstractVector{Tuple{Int,Int}}, g)
    n = length(C)
    for k = 1:n
        i, j = C[k]
        if i > 0 && j > 0 && criterion(g, C, i, j)
            C[k] = (0, 0)
        end
    end
    for k = n:-1:1
        i, j = C[k]
        if i <= 0
            deleteat!(C, k)
        end
    end
end

# eliminiate generators with leading terms spanned by other leading terms
function minimize!(H::AbstractVector{P}) where P<:MultivariatePolynomial
    n = length(H)
    for i = 1:n
        f = H[i]
        if !iszero(f)
            lif = multideg(H[i])
            for g in H
                if !iszero(g) && f != g
                    if all(lif .>= multideg(g))
                        cf = LC(f)
                        cg = LC(g)
                        if iszero(rem(cf, cg))
                            H[i] = zero(P)
                        end
                    end
                end
            end
        end
    end
    j = 0
    for i = 1:n
        f = H[i]
        if !iszero(f)
            j += 1
            if i != j
                H[j] = f
            end
        end
    end
    resize!(H, j)
    for i = 1:j
        f = H[i]
        lcu = inv(lcunit(f))
        if !isone(lcu)
            f = f * lcu
        end
        H[i] = f
    end
    H
end

# reduced Gröbner base
function reduce!(H::AbstractVector{P}) where P<:Polynomial
    n = length(H)
    for i = 1:n
        f = H[i]
        a, g, c = pdivrem(f, [g for g in H if g != f && !iszero(g)])
        if g !== f && isone(c)
            H[i] = g
        end
    end
    sort_unique!(H, rev=true)
    j = findlast(iszero, H)
    if j !== nothing
        resize!(H, j-1)
    end
    H
end

"""
    groebnerbase(H::AbstractVector{<:Polynomial})

Calculate the reduced groebner base from a set of generators of an ideal `<H>`.

see for example:
https://en.wikipedia.org/wiki/Gr%C3%B6bner_basis
http://www.crypto.rub.de/imperia/md/content/may/12/ws1213/kryptanal12/13_buchberger.pdf
"""
function groebnerbase(H::AbstractVector{P}) where P<:Polynomial
    buchberger(H) |> minimize! |> reduce!
end

# utility functions
"""
    positionsin(va, vb)

Return integer array of the same size as va. Element `i` is the index of first `va[i]`
in `vb` or zero, if missing.
"""
function positionsin(va::T, vb::T) where T<:AbstractVector
    n = length(va)
    res = Vector{Int}(undef, n)
    for i = 1:n
        p = findfirst(isequal(va[i]), vb)
        res[i] = ifelse(p === nothing, 0, p)
    end
    res
end

"""
    reindex2(values, pos, n)

Return an array of integers of size `n` such that `result[pos] = val`.
Default values in `result` are zero.
The indices with `pos[i] == 0` are silently ignored.
"""
function reindex2(xa::AbstractVector, pos::AbstractVector{<:Integer}, n::Integer)
    res = zeros(Int, n)
    for i = 1:length(pos)
        posi = pos[i]
        if posi != 0
            res[posi] = xa[i]
        end
    end
    res
end

function checkpositions(pos::AbstractVector{<:Integer}, xa::AbstractVector, va, vp)
    findfirst(iszero, pos) == nothing && return pos
    i = findfirst(i->iszero(pos[i]) && !iszero(xa[i]), 1:length(pos))
    if i !== nothing
        throw(ArgumentError("Variable :$(va[i]) not contained in $vp."))
    end
end


text:
## gmms.jl  Some functions for a Gaussia Mixture Model
## (c) 2013--2014 David A. van Leeuwen

import Base.LinAlg.AbstractTriangular

## uninitialized constructor, defaults to Float64
"""
`GMM(n::Int, d::Int, kind::Symbol=:diag)` initializes a GMM with means 0 and Indentity covariances
"""
function GMM(n::Int, d::Int; kind::Symbol=:diag)
    w = ones(n)/n
    μ = zeros(n, d)
    if kind == :diag
        Σ = ones(n, d)
    elseif kind == :full
        Σ = UpperTriangular{Float64, Matrix{Float64}}[UpperTriangular(eye(d)) for i=1:n]
    else
        error("Unknown kind")
    end
    hist = [History(@sprintf "Initialization n=%d, d=%d, kind=%s" n d kind)]
    GMM(w, μ, Σ, hist, 0)
end

Base.eltype{T}(gmm::GMM{T}) = T

## switch between full covariance and inverse cholesky decomposition representations.
"""
`covar(GMM.Σ)` extracts the covariances Σ (which may be encoded as chol(inv(Σ))
"""
covar{T}(ci::AbstractTriangular{T}) = (c = inv(ci); c * c')
cholinv{T}(Σ::Matrix{T}) = chol(inv(cholfact(0.5(Σ+Σ'))))

"""
`kind(::GMM)` returns the kind of GMM, either `:diag` or `:full`
"""
function kind(g::GMM{T, CT}) where {T, CT <: DiagCov}
    return :diag
end
function kind(g::GMM{T, CT}) where {T, CT <: FullCov}
    return :full
end

## This may clash with STatsBase
"""
`weights(::GMM)` returns the weights `w`, or priors, of the Gaussians in the mixture
"""
weights(gmm::GMM) = gmm.w
"`means(::GMM)` returns the means `μ` of the Gaussians in the mixture"
means(gmm::GMM) = gmm.μ
"`covars(::GMM)` returns the covariance matrices Σ of the Gaussians in the mixture."
covars{T}(gmm::GMM{T,DiagCov{T}}) = gmm.Σ
covars{T}(gmm::GMM{T,FullCov{T}}) = [covar(ci) for ci in gmm.Σ]

"`nparams(::GMM)` returns the number of free parameters in the GMM"
function nparams(gmm::GMM)
    gmmkind = kind(gmm)
    if gmmkind ==:diag
        gmm.n * (1 + 2gmm.d) - 1
#        sum(map(length, (gmm.w, gmm.μ, gmm.Σ)))-1
    elseif gmmkind == :full
        gmm.n * (1 + gmm.d + div((gmm.d+1)gmm.d, 2)) - 1
#        sum(map(length, (gmm.w, gmm.μ))) - 1 + div(gmm.n * gmm.d * (gmm.d+1), 2)
    else
        error("Unknown kind")
    end
end

"`addhist!(::GMM, s)` adds a comment `s` to the GMMM"
function addhist!(gmm::GaussianMixture, s::AbstractString)
    info(s)
    push!(gmm.hist, History(s))
    gmm
end

## copy a GMM, deep-copying all internal information.
function Base.copy(gmm::GMM)
    w = copy(gmm.w)
    μ = copy(gmm.μ)
    Σ = deepcopy(gmm.Σ)
    hist = copy(gmm.hist)
    g = GMM(w, μ, Σ, hist, gmm.nx)
    addhist!(g, "Copy")
end

## create a full cov GMM from a diag cov GMM (for testing full covariance routines)
"`full(::GMM)` turns a diagonal covariance GMM into a full-covariance GMMM"
function Base.full{T}(gmm::GMM{T})
    if kind(gmm) == :full
        return gmm
    end
    Σ = convert(FullCov{T}, [UpperTriangular(diagm(vec(1./√gmm.Σ[i,:]))) for i=1:gmm.n])
    new = GMM(copy(gmm.w), copy(gmm.μ), Σ, copy(gmm.hist), gmm.nx)
    addhist!(new, "Converted to full covariance")
end

"""`diag(::GMM)` turns a full-covariance GMM into a diagonal-covariance GMM, by ignoring
off-diagonal elements"""
function Base.diag{T}(gmm::GMM{T})
    if kind(gmm) == :diag
        return gmm
    end
    Σ = Array{T}(gmm.n, gmm.d)
    for i=1:gmm.n
        Σ[i,:] = 1./abs2(diag(gmm.Σ[i]))
    end
    new = GMM(copy(gmm.w), copy(gmm.μ), Σ, copy(gmm.hist), gmm.nx)
    addhist!(new, "Converted to diag covariance")
end

function Base.show(io::IO, h::History)
    println(io, Libc.strftime(h.t), ": ", h.s)
end

history(gmm::GaussianMixture) = gmm.hist

function Base.show(io::IO, ::MIME"text/plain", hist::Vector{History})
    t0 = hist[1].t
    println(io, "GMM trained from ", Libc.strftime(t0), " to ", Libc.strftime(last(hist).t))
    for h in hist
        s = split(h.s, "\n")
        print(io, @sprintf("%6.3f\t%s\n", h.t-t0, s[1]))
        for i = 2:length(s)
            print(io, @sprintf("%6s\t%s\n", " ", s[i]))
        end
    end
end

# compute the ranges to be displayed, plus a total index comprising all ranges.
function compute_range(maxn, n)
    if maxn < n
        hn = div(maxn,2)
        r = (1:hn, n-hn+1:n)
    else
        r = (1:n,)
    end
    totr = vcat(map(collect, r)...)
    r, totr
end

## we could improve this a lot
function Base.show{T}(io::IO, mime::MIME"text/plain", gmm::GMM{T})
    println(io, @sprintf("GMM{%s} with %d components in %d dimensions and %s covariance", T, gmm.n, gmm.d, kind(gmm)))
    gmmkind = kind(gmm)
    if gmmkind == :diag
        maxngauss = clamp((displaysize(io)[1] - 1) ÷ 6gmm.n, 1, gmm.n)
        ranges, = compute_range(maxngauss, gmm.n)
        for (i, r) in enumerate(ranges)
            if i > 1
                println(io, "⋮")
            end
            for j in r
                println(io, @sprintf "Mix %d: weight %f" j gmm.w[j]);
                println(io, "  mean: ", gmm.μ[j,:])
                println(io, "  variance: ", gmm.Σ[j,:])
            end
        end
    elseif gmmkind == :full
        nlinesneeded = gmm.n * (4 + gmm.d) + 1
        if displaysize(io)[1] > nlinesneeded
            ranges = (1:gmm.n,)
        else
            maxngauss = clamp((displaysize(io)[1] - 1) ÷ (4 + gmm.d), min(gmm.n,3), gmm.n)
            ranges, = compute_range(maxngauss, gmm.n)
            println(maxngauss, " ", ranges)
        end
        for (i, r) in enumerate(ranges)
            if i > 1
                println(io, "⋮")
            end
            for j in r
                println(io, @sprintf "Mix %d: weight %f" j gmm.w[j]);
                println(io, " mean: ", gmm.μ[j,:])
                print(io, " covariance: ")
                show(io, mime, covar(gmm.Σ[j]))
                println(io)
            end
        end
    else
        printf("Unknown kind")
    end
end

## some routines for conversion between float types
#    @doc """`convert(GMM{::Type}, GMM)` convert the GMM to a different floating point type""" ->
function Base.convert{Td,Ts}(::Type{GMM{Td}}, gmm::GMM{Ts})
    Ts == Td && return gmm
    h = vcat(gmm.hist, History(string("Converted to ", Td)))
    w = map(Td, gmm.w)
    μ = map(Td, gmm.μ)
    gmmkind = kind(gmm)
    if gmmkind == :full
        Σ = map(eltype(FullCov{Td}),  gmm.Σ)
    elseif gmmkind == :diag
        Σ = map(Td, gmm.Σ)
    else
        error("Unknown kind")
    end
    GMM(w, μ, Σ, h, gmm.nx)
end
function Base.convert{Td,Ts}(::Type{VGMM{Td}}, vg::VGMM{Ts})
    Ts == Td && return vg
    h = vcat(vg.hist, History(string("Converted to ", Td)))
    W = map(eltype(FullCov{Td}), vg.W)
    π = convert(GMMprior{Td}, vg.π)
    VGMM(vg.n, vg.d, π, map(Td,vg.α), map(Td, vg.β), map(Td,vg.m),
    map(Td, vg.ν), W, h)
end
function Base.convert{Td,Ts}(::Type{GMMprior{Td}}, p::GMMprior{Ts})
    Ts == Td && return p
    GMMprior(map(Td, p.α₀), map(Td, p.β₀), map(Td, p.m₀), map(Td, p.ν₀), map(Td, p.W₀))
end


text:
#
# CoaXPress.jl -
#
# Julia interface to CoaXPress cameras via Phoenix (PHX) library.
#
#------------------------------------------------------------------------------
#
# This file is part of the `Phoenix.jl` package which is licensed under the MIT
# "Expat" License.
#
# Copyright (C) 2017-2019, Éric Thiébaut (https://github.com/emmt/Phoenix.jl).
# Copyright (C) 2016, Éric Thiébaut & Jonathan Léger.
#

export
    CXP_MAGIC,
    CXP_STANDARD,
    CXP_REVISION,
    CXP_XML_MANIFEST_SIZE,
    CXP_XML_MANIFEST_SELECTOR,
    CXP_XML_VERSION,
    CXP_XML_SCHEME_VERSION,
    CXP_XML_URL_ADDRESS,
    CXP_IIDC2_ADDRESS,
    CXP_DEVICE_VENDOR_NAME,
    CXP_DEVICE_MODEL_NAME,
    CXP_DEVICE_MANUFACTURER_INFO,
    CXP_DEVICE_VERSION,
    CXP_DEVICE_SERIAL_NUMBER,
    CXP_DEVICE_USER_ID,
    CXP_WIDTH_ADDRESS,
    CXP_HEIGHT_ADDRESS,
    CXP_ACQUISITION_MODE_ADDRESS,
    CXP_ACQUISITION_START_ADDRESS,
    CXP_ACQUISITION_STOP_ADDRESS,
    CXP_PIXEL_FORMAT_ADDRESS,
    CXP_DEVICE_TAP_GEOMETRY_ADDRESS,
    CXP_IMAGE1_STREAM_ID_ADDRESS,
    CXP_CONNECTION_RESET,
    CXP_DEVICE_CONNECTION_ID,
    CXP_MASTER_HOST_CONNECTION_ID,
    CXP_CONTROL_PACKET_SIZE_MAX,
    CXP_STREAM_PACKET_SIZE_MAX,
    CXP_CONNECTION_CONFIG,
    CXP_CONNECTION_CONFIG_DEFAULT,
    CXP_TEST_MODE,
    CXP_TEST_ERROR_COUNT_SELECTOR,
    CXP_TEST_ERROR_COUNT,
    CXP_TEST_PACKET_COUNT_TX,
    CXP_TEST_PACKET_COUNT_RX,
    CXP_HS_UP_CONNECTION,
    CXP_MANUFACTURER

#---------------------------------------------------------------------------
# TYPES

"""

Abstract type `Register` is used to store (among others) the address (as an
`UInt32` value) of a CoaXPress register.  All concrete sub-types must have a
first member `addr::UInt32`.

"""
abstract type Register end

## Manage to automatically convert a register to a pointer to its
## address member. See refpointer.jl
#function unsafe_convert{T<:Register}(::Type{Ptr{UInt32}}, rgst::T)
#    if isbits(T)
#        return convert(Ptr{UInt32}, data_pointer_from_objref(rgst))
#    else
#        return convert(Ptr{UInt32}, data_pointer_from_objref(rgst.addr))
#    end
#end

"""

`RegisterValue{T,A}(addr)` is a value of type `T` stored at address `addr` with
access mode `A`.

"""
struct RegisterValue{T,A} <: Register where {T,A<:AccessMode}
    addr::UInt32
end

Base.eltype(::RegisterValue{T,A}) where {T,A<:AccessMode} = T

const RegisterEnum{A} = RegisterValue{UInt32,A} where {A<:AccessMode}

"""

`RegisterAddress{T,A}(addr)` stores at address `addr` the address of another
register for a value of type `T` with access mode `A`.

"""
struct RegisterAddress{T,A} <: Register where {T,A<:AccessMode}
    addr::UInt32
end

"""

`RegisterCommand{T}(addr,val)` is to write a constant of type `T` and value
`val` at address `addr`.

"""
struct RegisterCommand{T} <: Register where {T}
    addr::UInt32
    value::T
end

"""

`RegisterString{N,A}(addr)` is a fixed length string of length `N` stored at
address `addr` with access mode `A`.

"""
struct RegisterString{N,A} <: Register where {N,A<:AccessMode}
    addr::UInt32
end

#---------------------------------------------------------------------------
# CONSTANTS

# Value returned by reading at CXP_STANDARD register.
const CXP_MAGIC = UInt32(0xC0A79AE5)

# Addresses of bootstrap CoaXPress registers common to all CoaXPress compliant
# devices.
const CXP_STANDARD                    = RegisterValue{UInt32,ReadOnly}(0x00000000)
const CXP_REVISION                    = RegisterValue{UInt32,ReadOnly}(0x00000004)
const CXP_XML_MANIFEST_SIZE           = RegisterValue{UInt32,ReadOnly}(0x00000008)
const CXP_XML_MANIFEST_SELECTOR       = RegisterValue{UInt32,ReadWrite}(0x0000000C)
const CXP_XML_VERSION                 = RegisterValue{UInt32,ReadOnly}(0x00000010)
const CXP_XML_SCHEME_VERSION          = RegisterValue{UInt32,ReadOnly}(0x00000014)
const CXP_XML_URL_ADDRESS             = RegisterValue{UInt32,ReadOnly}(0x00000018)
const CXP_IIDC2_ADDRESS               = RegisterValue{UInt32,ReadOnly}(0x0000001C)
const CXP_DEVICE_VENDOR_NAME          = RegisterString{32,ReadOnly}(0x00002000)
const CXP_DEVICE_MODEL_NAME           = RegisterString{32,ReadOnly}(0x00002020)
const CXP_DEVICE_MANUFACTURER_INFO    = RegisterString{48,ReadOnly}(0x00002040)
const CXP_DEVICE_VERSION              = RegisterString{32,ReadOnly}(0x00002070)
const CXP_DEVICE_SERIAL_NUMBER        = RegisterString{16,ReadOnly}(0x000020B0)
const CXP_DEVICE_USER_ID              = RegisterString{16,ReadWrite}(0x000020C0)
const CXP_WIDTH_ADDRESS               = RegisterAddress{UInt32,ReadWrite}(0x00003000)
const CXP_HEIGHT_ADDRESS              = RegisterAddress{UInt32,ReadWrite}(0x00003004)
const CXP_ACQUISITION_MODE_ADDRESS    = RegisterAddress{UInt32,ReadWrite}(0x00003008)
const CXP_ACQUISITION_START_ADDRESS   = RegisterAddress{UInt32,WriteOnly}(0x0000300C)
const CXP_ACQUISITION_STOP_ADDRESS    = RegisterAddress{UInt32,WriteOnly}(0x00003010)
const CXP_PIXEL_FORMAT_ADDRESS        = RegisterAddress{UInt32,ReadWrite}(0x00003014)
const CXP_DEVICE_TAP_GEOMETRY_ADDRESS = RegisterAddress{UInt32,ReadWrite}(0x00003018)
const CXP_IMAGE1_STREAM_ID_ADDRESS    = RegisterAddress{UInt32,ReadWrite}(0x0000301C)
const CXP_CONNECTION_RESET            = RegisterValue{UInt32,ReadWrite}(0x00004000) # read doc.
const CXP_DEVICE_CONNECTION_ID        = RegisterValue{UInt32,ReadOnly}(0x00004004)
const CXP_MASTER_HOST_CONNECTION_ID   = RegisterValue{UInt32,ReadWrite}(0x00004008)
const CXP_CONTROL_PACKET_SIZE_MAX     = RegisterValue{UInt32,ReadOnly}(0x0000400C)
const CXP_STREAM_PACKET_SIZE_MAX      = RegisterValue{UInt32,ReadWrite}(0x00004010)
const CXP_CONNECTION_CONFIG           = RegisterEnum{ReadWrite}(0x00004014)
const CXP_CONNECTION_CONFIG_DEFAULT   = RegisterValue{UInt32,ReadOnly}(0x00004018)
const CXP_TEST_MODE                   = RegisterValue{UInt32,ReadWrite}(0x0000401C)
const CXP_TEST_ERROR_COUNT_SELECTOR   = RegisterValue{UInt32,ReadWrite}(0x00004020)
const CXP_TEST_ERROR_COUNT            = RegisterValue{UInt32,ReadWrite}(0x00004024)
const CXP_TEST_PACKET_COUNT_TX        = RegisterValue{UInt64,ReadWrite}(0x00004028) # read doc.
const CXP_TEST_PACKET_COUNT_RX        = RegisterValue{UInt64,ReadWrite}(0x00004030) # read doc.
const CXP_HS_UP_CONNECTION            = RegisterValue{UInt32,ReadOnly}(0x0000403C)
const CXP_MANUFACTURER                = RegisterValue{Nothing,Unreachable}(0x00006000)

#---------------------------------------------------------------------------


text:
using DiffEqBase, DiffEqBiological, DiffEqJump, DiffEqProblemLibrary.JumpProblemLibrary, Plots, Statistics
gr()
fmt = :png
JumpProblemLibrary.importjumpproblems()

methods = (Direct(),DirectFW(),FRM(),FRMFW(),SortingDirect(),NRM(),DirectCR(),RSSA())
shortlabels = [string(leg)[12:end-2] for leg in methods]
jprob   = prob_jump_multistate
tf      = 10.0*jprob.tstop
prob    = DiscreteProblem(jprob.u0, (0.0,tf), jprob.rates)
rn      = jprob.network
varlegs = ["A_P", "A_bound_P", "A_unbound_P", "RLA_P"]
varsyms = [
    [:S7,:S8,:S9],
    [:S9],
    [:S7,:S8],
    [:S7]
]
varidxs = []
for vars in varsyms
    push!(varidxs, [findfirst(isequal(sym),rn.syms) for sym in vars])
end

p = []
for (i,method) in enumerate(methods)
    jump_prob = JumpProblem(prob, method, rn, save_positions=(false,false))
    sol = solve(jump_prob, SSAStepper(), saveat=tf/1000.)
    solv = zeros(1001,4)
    for (i,varidx) in enumerate(varidxs)
        solv[:,i] = sum(sol[varidx,:], dims=1)
    end
    if i < length(methods)
        push!(p, plot(sol.t,solv,title=shortlabels[i],legend=false,format=fmt))
    else
        push!(p, plot(sol.t,solv,title=shortlabels[i],legend=true,labels=varlegs,format=fmt))
    end
end
plot(p...,format=fmt)

function run_benchmark!(t, jump_prob, stepper)
    sol = solve(jump_prob, stepper)
    @inbounds for i in 1:length(t)
        t[i] = @elapsed (sol = solve(jump_prob, stepper))
    end
end

nsims = 100
benchmarks = Vector{Vector{Float64}}()
for method in methods
    jump_prob = JumpProblem(prob, method, rn, save_positions=(false,false))
    stepper = SSAStepper()
    t = Vector{Float64}(undef, nsims)
    run_benchmark!(t, jump_prob, stepper)
    push!(benchmarks, t)
end

medtimes = Vector{Float64}(undef,length(methods))
stdtimes = Vector{Float64}(undef,length(methods))
avgtimes = Vector{Float64}(undef,length(methods))
for i in 1:length(methods)
    medtimes[i] = median(benchmarks[i])
    avgtimes[i] = mean(benchmarks[i])
    stdtimes[i] = std(benchmarks[i])
end
using DataFrames

df = DataFrame(names=shortlabels,medtimes=medtimes,relmedtimes=(medtimes/medtimes[1]),avgtimes=avgtimes, std=stdtimes, cv=stdtimes./avgtimes)

sa = [text(string(round(mt,digits=3),"s"),:center,12) for mt in df.medtimes]
bar(df.names,df.relmedtimes,legend=:false, fmt=fmt)
scatter!(df.names, .05 .+ df.relmedtimes, markeralpha=0, series_annotations=sa, fmt=fmt)
ylabel!("median relative to Direct")
title!("Multistate Model")

using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


text:
module Requests

using MySQL

SQL() = DBInterface.connect(
        MySQL.Connection,
        ENV["MYSQL_HOST"],
        ENV["MYSQL_USER"],
        ENV["MYSQL_ROOT_PASSWORD"],
        db = ENV["MYSQL_DATABASE"]
    )

function insert_channel(conn::MySQL.Connection, values...)
    ret = true
    try
        DBInterface.execute(conn, """INSERT INTO channel (name,
        description, protected, password, owner) VALUES ('$(values[1])',
        'tmp', $(values[2]), '$(values[3])', $(values[4]))""")
    catch err
        ret = false
    end
    ret
end

function insert_account(conn::MySQL.Connection, name, password)
    ret = true
    try
        DBInterface.execute(conn, """INSERT INTO user (name, password)
        VALUES ('$name', '$password')""")
    catch err
        ret = false
    end
    ret
end

function get_channel(conn::MySQL.Connection, name)
    r = DBInterface.execute(conn, """SELECT id, description,
    name, protected, password, owner FROM channel WHERE name='$name'""")
    if (length(r) == 0)
        return ([])
    end
    first(r)
end

function get_account(conn::MySQL.Connection, name)
    r = DBInterface.execute(conn, """SELECT id, password, name
    FROM user WHERE name='$name'""")
    if (length(r) == 0)
        return ([])
    end
    first(r)
end

function get_lobby(conn::MySQL.Connection)
    r = DBInterface.execute(conn, """SELECT id, name, description,
    protected, password, owner FROM channel WHERE id=1""")
    if (length(r) == 0)
        return ([])
    end
    first(r)
end

export SQL

end # Requests


text:
using SpkCore

include("./GLMBasis.jl")
using .GLMBasis

# ============================================================================ #
abstract type AbstractBasis end
# ============================================================================ #
struct NullBasis <: AbstractBasis
    bin_size::Float64
end
NullBasis() = NullBasis(1.0)
nbasis(n::NullBasis) = 1
project(n::NullBasis, ts::Vector{Float64}, evt::Vector{Float64}) = evt
# ============================================================================ #
mutable struct DefaultBasis <: AbstractBasis
    length::Int
    offset::Int
    bin_size::Float64
    changed::Bool
end
function DefaultBasis(;length::Integer=0, offset::Integer=0, bin_size::Real=0.001)
    @assert(length > 0, "Required length argument must be > 0")
    return DefaultBasis(length, offset, bin_size, false)
end
# ---------------------------------------------------------------------------- #
nbasis(b::DefaultBasis) = b.length
# ---------------------------------------------------------------------------- #
function project(x::DefaultBasis, ts::Vector{Float64}, evt::Vector{Float64})
    k = x.offset
    return permutedims(psth(ts, evt, -(x.length + k - 1):-k, x.bin_size)[1], (2,1))
end
# ============================================================================ #
# CosineBasis(length=100, n=4, b=4, ortho=false, bin_size=0.001)
mutable struct CosineBasis <: AbstractBasis
    length::Int
    offset::Int
    nbasis::Int
    b::Float64
    ortho::Bool
    bin_size::Float64
    basis::Matrix{Float64}
    changed::Bool
end
function CosineBasis(;length::Int=0, offset::Integer=0, nbasis::Int=8, b::Real=4.0, ortho::Bool=false, bin_size::Real=0.001)
    @assert(length > 0, "Required length argument must be > 0")
    return CosineBasis(length, offset, nbasis, b, ortho, bin_size, Matrix{Float64}(undef, 0, 0), true)
end
# ---------------------------------------------------------------------------- #
nbasis(x::CosineBasis) = x.nbasis
# ---------------------------------------------------------------------------- #
function generate(x::CosineBasis)
    if x.changed
        xpeaks = [0, round(Int, x.length * (1 - 1.5 / x.nbasis))]
        basis, ortho_basis = make_stim_basis(0, x.nbasis, xpeaks, x.b, x.length)
        if size(basis) == size(x.basis)
            x.basis .= x.ortho ? ortho_basis : basis
        else
            x.basis = x.ortho ? ortho_basis : basis
        end
        if any(isnan, x.basis)
            error("NaN detected in basis!!!")
        end
        x.changed = false
    end
    return x.basis
end
# ---------------------------------------------------------------------------- #
function project(x::CosineBasis, ts::Vector{Float64}, evt::Vector{Float64})
    idx = ts2idx.(ts, x.bin_size)
    if ts === evt
        evt_idx = idx
    else
        evt_idx = ts2idx.(evt, x.bin_size)
    end
    tidx = evt_idx .- x.offset
    if tidx[1] < 1
        filter!(>(0), tidx)
    end
    return spikefilter(idx, reverse(generate(x), dims=1), (1, evt_idx[end]))[tidx, :]
end
# ============================================================================ #
function Base.setproperty!(b::AbstractBasis, k::Symbol, v)
    setfield!(b, k, v)
    setfield!(b, :changed, true)
    return b
end
# ============================================================================ #
function set_parameter!(x::T; args...) where T <: AbstractBasis
    for (k,v) in args
        if hasfield(T, k)
            setfield!(x, k, v)
            x.changed = true
        end
    end
end
# ============================================================================ #
function Base.:*(a::Vector{<:Real},  x::CosineBasis)
    return (x.basis' * x.basis) \ (x.basis' * a)
end
Base.:*(x::CosineBasis, a::Vector{<:Real}) = x.basis * a

@inline Base.:*(a::Vector{<:Real}, x::AbstractBasis) = a
@inline Base.:*(x::AbstractBasis, a::Vector{<:Real}) = a
# ============================================================================ #


text: Dict{String,Int32}(zip(keys, values))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
"""
    poisson_rng(p, n, shape=1; seed=nothing)

Generate a `shape` element array of random variables from a Poisson(`λ`) distribution. Optionally you can set a specific seed.

# Notes

The Poisson distribution is given by:

``f(x, λ) = \\frac{λ^x}{x!} e^λ \\quad x = 0,1,2,…``

# Examples

```julia-repl
julia> poisson_rng(3)
1×1 Matrix{Int64}:
 7

julia> poisson_rng(10, 5)
5×1 Matrix{Int64}:
 13
 11
 10
  8
 15

julia> poisson_rng(10, (5,5))
5×5×1 Array{Int64, 3}:
[:, :, 1] =
 11  15   9  11   9
  8  15  13  10   9
 11  12   4  10   6
  7   9  13  11   7
 13   7  10  10  14

```

# References

R. Larson, A. Odoni. Urban operations research. Prentice-Hall, New Jersey, 1981.

G. Last, M. Penrose. Lectures on the poisson process. Cambridge University Press, 2017.
"""
function poisson_rng(λ::Real, shape::Union{Int, Tuple{Vararg{Int}}}=1; seed::Union{Int, Nothing}=nothing)
    n = ceil(Int, λ*5)  # ensure n is integer
    U = expon_rng(λ, (shape..., n), seed=seed)
    X = sum(cumsum(U, dims=ndims(U)) .< 1, dims=ndims(U))
    return X
end


text:
# preprocess.jl
# Runs clustering algorithms to harmonize set of constraints over New York City
# Authored on 6/22/15 by Brandon Zeng

function outputPreprocessedConstraints(manhattan::Manhattan;radius::Int=40, numClusters::Int64=50, minRides::Int64=4, sampleSize::Int64=5000, overwrite::Bool=false, times::String = "1214")
	"""
	Checks if particular preprocessed file already exists. If so, does nothing. Otherwise, generates it.
	"""
	if overwrite || !isfile("../Travel_times/travel_times_r$(radius)_wd_$(times)_clust$(numClusters)_rides$(sampleSize)_minr$(minRides).csv")
		println("**** Selecting travel times ****")
		inputName = "../Travel_times/travel_times_r$(radius)_wd_$(times)"
		preprocess(inputName, numClusters, minRides, sampleSize, manhattan)
	else
		println("**** Selected travel times found ****")
	end
end

function outputPreprocessedConstraints(manhattan::Manhattan, trainOrTest::String; radius::Int=40, numClusters::Int64=50, minRides::Int64=4, sampleSize::Int64=5000, overwrite::Bool=false, times::String="1214")
	"""
	Checks if particular preprocessed file already exists. If so, does nothing. Otherwise, generates it.
	"""
	if overwrite || !isfile("../Travel_times/$(trainOrTest)_r$(radius)_wd_$(times)_clust$(numClusters)_rides$(sampleSize)_minr$(minRides).csv")
		println("**** Selecting travel times ****")
		inputName = "../Travel_times/$(trainOrTest)_r$(radius)_wd_$(times)"
		preprocess(inputName, numClusters, minRides, sampleSize, manhattan)
	else
		println("**** Selected travel times found ****")
	end
end

function preprocess(inputName::String, numClusters::Int64, minRides::Int64, sampleSize::Int64, manhattan::Manhattan)
	"""
	Clusters the nodes in Manhattan into numClusters clusters. Then reduces the number of constraints to sampleSize, so that the constraints cover as many cluster pairs as possible.
	"""
	println("---- Clustering ----")

	# Load list of nodes that are part of a highway (to exclude rides with these nodes)
	highwayNodes = load("Cities/Saved/highwayNodes.jld", "highwayNodes")

	# Load Manhattan and cluster nodes geographically
	man = manhattan
	coordinatesArray = Float64[]
	for c in man.positions
		push!(coordinatesArray, c.x)
		push!(coordinatesArray, c.y)
	end
	coordinateMatrix = reshape(coordinatesArray, 2, 6134)
	R = kmeans(coordinateMatrix, numClusters, maxiter = 1000)

	# Read in rides
	println("---- Reading CSV ----")
	rides = readcsv("$(inputName).csv")

	nodes = rides[:,1]
	average = rides[:,2]
	stddev = rides[:,3]
	stderr = rides[:,4]
	numRides = rides[:,5]

	function extractNodes(nodePair::SubString{ASCIIString})
		source = parse(Int, split(nodePair, ";")[1]) + 1
		dest = parse(Int, split(nodePair, ";")[2]) + 1
		return source, dest
	end

	function classify(source::Int, dest::Int, clusterAssignments::Vector{Int64})
		return (clusterAssignments[source], clusterAssignments[dest])
	end

	function flightDistance(source::Int, dest::Int)
		return norm([man.positions[source].x - man.positions[dest].x, man.positions[source].y - man.positions[dest].y])
	end

	givenRides = Dict{Tuple{Int64,Int64}, Vector{Int}}()
	selectedRides = Dict{Tuple{Int64,Int64}, Vector{Bool}}()

	println("---- Loading rides ----")
	# Load rides into memory
	for i = 1:length(nodes)
		if i % 10000 == 0
			@printf("                                \r")
			@printf("Progress: %.1f%% completed\r", 100*i/length(nodes))
		end
		# Check that there are enough rides and that the time is not ridiculously small
		if numRides[i] >= minRides && average[i] >= 60
			src, dest = extractNodes(nodes[i])
			# Check that we are not starting or ending on a highway (makes no sense)
			if !(src in highwayNodes) && !(dest in highwayNodes) && src != dest && flightDistance(src, dest) >= 300
				clusterPair = classify(src, dest, R.assignments)
				if haskey(givenRides, clusterPair)
					push!(givenRides[clusterPair], i)
					push!(selectedRides[clusterPair], false)
				else
					givenRides[clusterPair] = [i]
					selectedRides[clusterPair] = [false]
				end
			end
		end
	end

	# Print out some info about how many rides we have
	lengths = [length(givenRides[key]) for key in keys(givenRides)]
	println("Number of cluster pairs with data: ", length(givenRides))
	println("Average number of rides/cluster pair: ", length(nodes)/length(givenRides))
	println("Median number of rides/cluster pair: ", median(lengths))
	totalSelections = sampleSize
	clusterSelection = round(Int, totalSelections / length(givenRides)) + 1
	result = Int[]

	println("---- Selecting rides ----")
	# Select number of rides from each cluster ("cut out the fat")
	for key in keys(givenRides)
		if length(givenRides[key]) <= clusterSelection
			result = vcat(result, givenRides[key])
		else
			count = 0
			while count < clusterSelection
				index = rand(1:length(givenRides[key]))
				if !selectedRides[key][index]
					selectedRides[key][index] = true
					push!(result, givenRides[key][index])
					count += 1
				end
			end
		end
	end

	# Sort the ride ids thus obtained so that we can quickly extract them
	sort!(result)
	writecsv("$(inputName)_clust$(numClusters)_rides$(sampleSize)_minr$(minRides).csv", rides[result,:])
end

max_stars_repo_name: sebmart/NYC-timing

max_issues_repo_name: sebmart/NYC-timing

max_forks_repo_name: sebmart/NYC-timing