#
# Important directories
simdir          = "/remote/ceph2/group/gerda/data/mpik/gerda-simulations/gerda-gems-db/pss-dev"
simdir_old      = "/remote/ceph2/group/gerda/data/simulation/gerda-mage-sim"
mapping_file    = simdir_old*"/UTILS/det-data/ged-mapping.json"
parameters_file = simdir_old*"/UTILS/det-data/ged-parameters.json"

#########################################
# Create mapping dictionary ############
sim_to_channel = open(mapping_file) do f
    mapping  = JSON.parse(f)["mapping"]
    channels = Dict()
    for ged in keys(mapping)
        channels[mapping[ged]["sim"]] = [mapping[ged]["channel"], mapping[ged]["name"]]
    end
    return channels
end;

#########################################
# Create mapping dictionary ############
channel_to_name = open(mapping_file) do f
    mapping  = JSON.parse(f)["mapping"]
    channels = Dict()
    for ged in keys(mapping)
        channels[mapping[ged]["channel"]] = mapping[ged]["name"]
    end
    return channels
end;

#########################################
# Create mapping dictionary ############
parameters = open(parameters_file) do f
    temp = JSON.parse(f)
    temp["GTF45"] = temp["GTF45_2"]
    return temp
end;

#########################################
# Electronics Model #####################
function applyElectronics(pulse; Ts = 10e-9, GBP = 2750e6, tau = 180e-6, Kv = 150e3, Cd = 50e-12, Cf = 0.35e-12, Rf = 500e6)

    wop = GBP / (2 * pi * Kv)
    Cmod = Cf + Cd
    wmod = 1.0 / (Rf * Cmod)
    alfa = Cmod / (Cf * GBP)

    b0 = 1.0 / alfa
    a2 = 1.0
    a1 = 1.0 / alfa + wop + wmod
    a0 = 1.0 / (tau * alfa) + wmod*wop

    # then the transfer function in the *Laplace* s-domain looks like this:
    #                       b0
    #   T(s) = ----------------------------
    #             a2 * s^2 + a1 * s + a0

    # PolynomialRatio needs z-transform paramters: s- and z-domains can be connected by
    # the bilinear transform:
    #        z - 1
    # s = K -------- , K = Ts/2  , Ts - sampling period
    #        z + 1
    #
    # we can then convert T(s) to T(z):
    #              bz2 * z^2 + bz1 * z + bz0
    #   T(z) = -------------------------------
    #              az2 * z^2 + az1 * z + az0
    #

    K = 2/Ts

    az2 = 1.0   # normalized
    az1 = (2*a0 - 2*K^2)/(K^2 + a1*K + a0)
    az0 = (K^2 - a1*K + a0)/(K^2 + a1*K + a0)

    bz2 = b0/(K^2 + a1*K + a0)
    bz1 = 2*b0/(K^2 + a1*K + a0)
    bz0 = b0/(K^2 + a1*K + a0)

    myfilter = PolynomialRatio([bz2, bz1, bz0], [az2, az1, az0])

    filtered = filt(myfilter, vcat([0], diff(pulse)))

end

#########################################
# Extend simulated waveforms ############
function add_baseline_and_extend_tail(wv::RadiationDetectorSignals.RDWaveform{T,U,TV,UV}, n_baseline_samples::Int, total_waveform_length::Int) where {T,U,TV,UV}
    new_signal::Vector{eltype(UV)} = Vector{eltype(UV)}(undef, total_waveform_length)
    new_signal[1:n_baseline_samples] .= zero(eltype(UV))
    if length(wv.value) <= total_waveform_length - n_baseline_samples
        new_signal[n_baseline_samples+1:n_baseline_samples+length(wv.value)] = wv.value
        new_signal[n_baseline_samples+length(wv.value)+1:end] .= wv.value[end]
    else
        new_signal[n_baseline_samples+1:end] = wv.value[1:total_waveform_length - n_baseline_samples]
    end
    new_times = if TV <: AbstractRange
        range( zero(first(wv.time)), step = step(wv.time), length = total_waveform_length )
    else
        error("Not yet definted for timestamps of type `$(TV)`")
    end
    return RDWaveform( new_times, new_signal )
end

#########################################
# Read / Create h5 sim file #############
function readcreateh5(filepath::String, output_path::String; read::Bool=true, overwrite::Bool=false,
    max_refine::Int=4)
    config_filename = filepath
    ssd_config_file = filepath
    #output_path *= split(basename(filepath), ".")[1]*"/"
    if !isdir(output_path)
        mkdir(output_path)
    end
    h5_sim_file = joinpath(output_path, split(basename(config_filename), ".config")[1]*".h5")

    T = Float32
    is_file = isfile(h5_sim_file)

    if overwrite
        is_file = false
    end

    s = if is_file
        if read
                @info("File already exist. Read in starts..")

                s = HDF5.h5open(h5_sim_file, "r") do h5f
                    Simulation(LegendHDF5IO.readdata(h5f, "SSD_Simulation"));
            end;
            SSD.set_charge_drift_model!(s, ADLChargeDriftModel(T = SSD.get_precision_type(s.detector)));
            return s
        else
            @info("File already exist. Read = false")
        end
    else
        @info("File does not exist. Simulation starts..")
        s = Simulation{T}(ssd_config_file)

        calculate_electric_potential!(  s,
                                        convergence_limit = 1e-7,
                                        use_nthreads = Base.Threads.nthreads(),
                                        max_refinements = max_refine,
                                        max_n_iterations = 500000,
                                        init_grid_spacing = (T(1e-3), T(2e-3), T(2e-3)),
                                        refinement_limits = (T(1e-5), T(1e-1), T(1e-5)),
                                        min_grid_spacing  = (T(2e-6), T(1e-1), T(2e-6)),
                                        depletion_handling = true )
        SSD.calculate_electric_field!(s, n_points_in_φ = 72)
        SSD.set_charge_drift_model!(s, ADLChargeDriftModel(T = T))
        SSD.apply_charge_drift_model!(s)
        for contact in s.detector.contacts
            calculate_weighting_potential!( s, contact.id,
                                            convergence_limit = 1e-7,
                                            use_nthreads = Base.Threads.nthreads(),
                                            max_refinements = 4,
                                            max_n_iterations = 500000,
                                            init_grid_spacing = (T(1e-3), T(2e-3), T(2e-3)),
                                            refinement_limits = (T(1e-4), T(1e-1), T(1e-4)),
                                            min_grid_spacing  = (T(2e-6), T(1e-1), T(2e-6)),
                                            n_points_in_φ = 2,
                                            verbose = true )
        end
        # IO
        HDF5.h5open(h5_sim_file, "w") do h5f
            LegendHDF5IO.writedata( h5f, "SSD_Simulation", NamedTuple(s)  )
        end

        return s
    end
end

function read_AE_params(filename::String)
    fit_params = JSON.parsefile(filename)
    tt = Table(detector=[], set=[], energy=[], scale=[], sig=[], mu=[], C=[], l=[], e=[], t=[], scale_err=[], sig_err=[], mu_err=[], C_err=[], l_err=[], e_err=[], t_err=[])
    for det in keys(fit_params) # Detector
        for set in keys(fit_params[det])
            for E in keys(fit_params[det][set])
                if "bat_params" in keys(fit_params[det][set][E])
                    append!(tt, Table(
                            detector=[det], set=[set], energy=[parse(Int64, E)],
                            scale=[fit_params[det][set][E]["bat_params"][1][1]],
                            sig=[fit_params[det][set][E]["bat_params"][1][2]],
                            mu=[fit_params[det][set][E]["bat_params"][1][3]],
                            C=[fit_params[det][set][E]["bat_params"][1][4]],
                            l=[fit_params[det][set][E]["bat_params"][1][5]],
                            e=[fit_params[det][set][E]["bat_params"][1][6]],
                            t=[fit_params[det][set][E]["bat_params"][1][7]],
                            scale_err=[fit_params[det][set][E]["bat_uncert"][1][1]],
                            sig_err=[fit_params[det][set][E]["bat_uncert"][1][2]],
                            mu_err=[fit_params[det][set][E]["bat_uncert"][1][3]],
                            C_err=[fit_params[det][set][E]["bat_uncert"][1][4]],
                            l_err=[fit_params[det][set][E]["bat_uncert"][1][5]],
                            e_err=[fit_params[det][set][E]["bat_uncert"][1][6]],
                            t_err=[fit_params[det][set][E]["bat_uncert"][1][7]]
                        )
                    );
                end
            end
        end
    end
    return tt
end

function get_cut_values(fit_params, ch, set)
    fit_cut_values = []
    fit_survival =[]
    fit_survival_err =[]
    for cut in keys(fit_params[channel_to_name[ch]]["DEP_cuts"][set])
        if !(cut in ["cut_sum", "cut_fit"])
            push!(fit_cut_values, parse(Float64, cut))
            push!(fit_survival, fit_params[channel_to_name[ch]]["DEP_cuts"][set][cut]["survival"])
            push!(fit_survival_err, fit_params[channel_to_name[ch]]["DEP_cuts"][set][cut]["survival_err"])
        end
    end
    sum_cut_values = fit_params[channel_to_name[ch]]["DEP_cuts"][set * "_sum"]["cut_values"]
    sum_survival = fit_params[channel_to_name[ch]]["DEP_cuts"][set * "_sum"]["survival"]
    sum_survival_err = fit_params[channel_to_name[ch]]["DEP_cuts"][set * "_sum"]["survival_err"]
    return fit_cut_values, fit_survival, fit_survival_err, sum_cut_values, sum_survival, sum_survival_err
end
function get_cut(fit_params, ch, set)
    return fit_params[channel_to_name[ch]]["DEP_cuts"][set]["cut_fit"], fit_params[channel_to_name[ch]]["DEP_cuts"][set]["cut_sum"]
end

function find_intersect(samples, threshold::Real, noisefilter = 1)
    if noisefilter < 1 noisefilter = 1 end
    if length(samples) < 1 error("Empty array") end
    i = 1
    x = samples[i]
    counter = 0
    intersect = i
    if x != threshold
        findHigher = x < threshold
        while i < length(samples)
            i += 1
            x = samples[i]
            if (findHigher && (x >= threshold)) || (!findHigher && (x <= threshold))
                if counter == 0
                    intersect = i
                end
                if counter >= noisefilter-1
                    return intersect
                else
                    counter += 1
                end
            else
                counter = 0
            end
        end
        #println("No intersect found")
        return NaN
    end
    return NaN
end

simdir_old      = "/remote/ceph2/group/gerda/data/simulation/gerda-mage-sim"
mapping_file    = "$simdir_old/UTILS/det-data/ged-mapping.json"
parameters_file = "$simdir_old/UTILS/det-data/ged-parameters.json"

#>------- Moving Window Average -----------------------<#
#>-----------------------------------------------------<#

function movingaverage(X::Vector,numofele::Int)
    BackDelta = div(numofele,2)
    ForwardDelta = isodd(numofele) ? div(numofele,2) : div(numofele,2) - 1
    len = length(X)
    Y = similar(X)
    for n = 1:len
        lo = max(1,n - BackDelta)
        hi = min(len,n + ForwardDelta)
        Y[n] = mean(X[lo:hi])
    end
    return Y
end

function movingaverage(X::Vector,numofele::Int,iterations::Int)
    i = 1
    while i <= iterations
        X = movingaverage(X,numofele)
        i += 1
    end
    return X
end

function movingaverage(X::Vector,numofele::Int,BackDelta::Int64,ForwardDelta::Int64)
    len = length(X)
    Y = similar(X)
    for n = 1:len
        lo = max(1,n - BackDelta)
        hi = min(len,n + ForwardDelta)
        Y[n] = mean(X[lo:hi])
    end
    return Y
end

function movingaverage(X::Vector,numofele::Int,BackDelta::Int64,ForwardDelta::Int64,iterations::Int)
    i = 1
    while i <= iterations
        X = movingaverage(X,numofele,BackDelta,ForwardDelta)
        i += 1
    end
    return X
end

function mwa(samples::Vector, window::Integer)
#window length is in number of samples and has to be an odd number

    if !isodd(window)
        error("For this method the window size has to be an odd integer")
    end

    newsamples = zeros(eltype(samples), length(samples))
    p = convert(eltype(window), (window - 1) / 2)

	#padding before and after the original vector
    samples = vcat(ones(eltype(samples), p) .* samples[1],
		samples, ones(eltype(samples), p) .* samples[length(samples)])

    newsamples[1] = sum(samples[1:window])
    for i in 2:length(newsamples)
        newsamples[i] = newsamples[i-1] + samples[i+2*p] - samples[i-1]
    end

    newsamples./window #return after normalization

end

function multi_mwa(samples::Vector, window::Integer, reps::Integer)
    for i in 1:1:reps
        samples = mwa(samples, window)
    end
    return samples
end

function mwa(samples::Vector, window::Integer, reps::Integer)
    for i in 1:1:reps
        samples = mwa(samples, window)
    end
    return samples
end

#>------- Get average maximum -------------------------<#
#>-----------------------------------------------------<#
function get_avg_maximum(pulse, intervall)
    index = findall(x -> x == maximum(pulse), pulse)[1]
    intervall = (index - Int(floor(intervall/2))+1):1:(index + Int(ceil(intervall/2)))
    return mean(pulse[intervall])
end

#>------- Create mapping dictionary -------------------<#
#>-----------------------------------------------------<#
sim_to_channel = open(mapping_file) do f
    mapping  = JSON.parse(f)["mapping"]
    channels = Dict()
    for ged in keys(mapping)
        channels[mapping[ged]["sim"]] = [mapping[ged]["channel"], mapping[ged]["name"]]
    end
    return channels
end;

#>------- Create mapping dictionary -------------------<#
#>-----------------------------------------------------<#
channel_to_name = open(mapping_file) do f
    mapping  = JSON.parse(f)["mapping"]
    channels = Dict()
    for ged in keys(mapping)
        channels[mapping[ged]["channel"]] = mapping[ged]["name"]
    end
    return channels
end;
ctn  = deepcopy(channel_to_name)

#>------- Create mapping dictionary -------------------<#
#>-----------------------------------------------------<#
parameters = open(parameters_file) do f
    temp = JSON.parse(f)
    temp["GTF45"] = temp["GTF45_2"]
    return temp
end;

#>------- Create Coax dictionary ----------------------<#
#>-----------------------------------------------------<#
channel_to_coax = open(mapping_file) do f
    mapping  = JSON.parse(f)["mapping"]
    channels = Dict()
    for ged in keys(mapping)
        if mapping[ged]["channel"] in [8,9,10,27,28,29,36,37,38,39]
            channels[mapping[ged]["channel"]] = true
        else
            channels[mapping[ged]["channel"]] = false
        end
    end
    return channels
end;

#>------- Create BEGe dictionary ----------------------<#
#>-----------------------------------------------------<#
channel_to_bege = open(mapping_file) do f
    mapping  = JSON.parse(f)["mapping"]
    channels = Dict()
    for ged in keys(mapping)
        if mapping[ged]["channel"] in [8,9,10,27,28,29,36,37,38,39]
            channels[mapping[ged]["channel"]] = false
        else
            channels[mapping[ged]["channel"]] = true
        end
    end
    return channels
end;
ctb = deepcopy(channel_to_bege)

function check_parameters(tt, td, GBP, tau, Cd, Cf)
    cut_events = Table( waveform    = [],
                        evtno       = [],
                        detno       = [],
                        det_ch      = [],
                        edep        = []);
    for i in 1:1:size(tt,1)
        pulse = tt[i].waveform.value
        filtered_pulse = applyElectronics(pulse; Ts = 10e-9, GBP = GBP, tau = tau, Kv = 150e3, Cd = Cd, Cf = Cf, Rf = 500e6)
        inter = find_intersect(filtered_pulse, maximum(filtered_pulse)/2, 5)
        if inter + 399 < length(pulse) && inter - 400 > 0
            cut  = filtered_pulse[(inter-400):1:(inter+400-1)]
            time = (0:10:7990)u"ns"
            cut_pulse = RDWaveform(time, cut)

            append!(cut_events, Table(  waveform    = [cut_pulse],
                                        evtno       = [tt[i].evtno],
                                        detno       = [tt[i].detno],
                                        det_ch      = [tt[i].det_ch],
                                        edep        = [tt[i].edep]))
        end
    end
    if size(cut_events,1) > 0
        sp = zeros(length(cut_events[1].waveform.value))
        for i in 1:1:size(cut_events,1)
            sp .+= cut_events[i].waveform.value
        end
        sp ./= size(cut_events,1)
        sp ./= maximum(sp)
        data_sp = td[1].superpulse.value ./ maximum(td[1].superpulse.value)
        #chi2_a = sum((sp[1:400] .- data_sp[1:400]).^2 ./ data_sp[1:400] )
        #chi2_b = sum((sp[401:end] .- data_sp[401:end]).^2 ./ data_sp[401:end] )
        chi2 = sum(abs.((sp .- data_sp).^2 ./ data_sp) )
        #return abs(chi2_a) + abs(chi2_b), [GBP, tau, Cd, Cf]
        return chi2, [GBP, tau, Cd, Cf]
    else
        return false
    end
end

function check_parameters_AE(tt, td, GBP, tau, Cd, Cf, ch; std_AE=2)
    cut_events = Table( waveform    = [],
                        evtno       = [],
                        detno       = [],
                        det_ch      = [],
                        edep        = []);
    for i in 1:1:size(tt,1)
        pulse = tt[i].waveform.value
        filtered_pulse = applyElectronics(pulse; Ts = 10e-9, GBP = GBP, tau = tau, Kv = 150e3, Cd = Cd, Cf = Cf, Rf = 500e6)
        inter = find_intersect(filtered_pulse, maximum(filtered_pulse)/2, 5)
        if inter + 399 < length(pulse) && inter - 400 > 0
            cut  = filtered_pulse[(inter-399):1:(inter+400)]
            time = (0:10:7990)u"ns"
            cut_pulse = RDWaveform(time, cut)

            append!(cut_events, Table(  waveform    = [cut_pulse],
                                        evtno       = [tt[i].evtno],
                                        detno       = [tt[i].detno],
                                        det_ch      = [tt[i].det_ch],
                                        edep        = [tt[i].edep]))
        end
    end
    if size(cut_events,1) > 0
        sp = zeros(length(cut_events[1].waveform.value))
        if channel_to_bege[ch]
            #> With A/E filter
            A = []
            for i in 1:1:size(cut_events, 1)
                diff_pulse = diff(cut_events.waveform[i].value ./sum(cut_events.edep[i]).val) # normalized to energy
                a = maximum(diff_pulse)
                push!(A, a)
            end
            µ = mean(A)
            sig = std(A)
            good = []
            for i in 1:1:size(cut_events,1)
                if A[i] >= µ - std_AE*sig && A[i] <= µ + std_AE*sig
                    sp  .+= cut_events[i].waveform.value
                    push!(good, i)
                end
            end
            sp ./= length(good)
        else
            #> Normal superposition of all pulses
            for i in 1:1:size(cut_events,1)
                sp .+= cut_events[i].waveform.value
            end
            sp ./= size(cut_events,1)
        end
        sp ./= maximum(sp)
        data_sp = td[1].superpulse.value ./ maximum(td[1].superpulse.value)
        diff_data = diff(data_sp)
        diff_sim  = diff(sp)
        #chi2_a = sum((sp[1:400] .- data_sp[1:400]).^2 ./ data_sp[1:400] )
        #chi2_b = sum((sp[401:end] .- data_sp[401:end]).^2 ./ data_sp[401:end] )
        non_zero = findall(x->x != 0, diff_data)
        chi2 = sum( abs.((diff_sim[non_zero] .- diff_data[non_zero]).^2 ./ diff_data[non_zero]) )
        return chi2, [GBP, tau, Cd, Cf]
    else
        return false
    end
end

function plot_param_set(tt, td, params, ch; std_AE=2)
    cut_events = Table( waveform    = [],
                        evtno       = [],
                        detno       = [],
                        det_ch      = [],
                        edep        = []);
    for i in 1:1:size(tt,1)
        pulse = tt[i].waveform.value
        filtered_pulse = applyElectronics(pulse; Ts = 10e-9, GBP = params[1], tau = params[2], Kv = 150e3, Cd = params[3], Cf = params[4], Rf = 500e6)
        inter = find_intersect(filtered_pulse, maximum(filtered_pulse)/2, 5)
        cut  = filtered_pulse[(inter-399):1:(inter+400)]
        time = (0:10:7990)u"ns"
        cut_pulse = RDWaveform(time, cut)

        append!(cut_events, Table(  waveform    = [cut_pulse],
                                    evtno       = [tt[i].evtno],
                                    detno       = [tt[i].detno],
                                    det_ch      = [tt[i].det_ch],
                                    edep        = [tt[i].edep]))
    end

    sp = zeros(length(cut_events[1].waveform.value))
    if channel_to_bege[ch]
        #> With A/E filter
        A = []
        for i in 1:1:size(cut_events, 1)
            diff_pulse = diff(cut_events.waveform[i].value ./sum(cut_events.edep[i]).val) # normalized to energy
            a = maximum(diff_pulse)
            push!(A, a)
        end
        µ = mean(A)
        sig = std(A)
        good = []
        for i in 1:1:size(cut_events,1)
            if A[i] >= µ - std_AE*sig && A[i] <= µ + std_AE*sig
                sp  .+= cut_events[i].waveform.value
                push!(good, i)
            end
        end
        sp ./= length(good)
    else
        #> Normal superposition of all pulses
        for i in 1:1:size(cut_events,1)
            sp .+= cut_events[i].waveform.value
        end
        sp ./= size(cut_events,1)
    end
    sp ./= maximum(sp)
    data_sp = td[1].superpulse.value ./ maximum(td[1].superpulse.value);
    diff_data = diff(data_sp);
    diff_sim  = diff(sp);

    chi2 = sum( abs.((diff_sim .- diff_data).^2 ./ diff_data) )
    p1 = plot(sp, label="Simulation: chi² of diff = "*string( round(chi2, digits=2) ));
    p1 = plot!(data_sp, color=:black, label="Superpulse (data)");
    p1 = plot!(title=channel_to_name[ch]*" | get electronics");
    p2 = plot(diff_data, label="Diff(data)");
    p2 = plot!(diff_sim, label="Diff(sim)");
    display(plot(p1,p2, layout=(2,1), xlim=[300, 500]))

end

function applyElectronics(pulse; Ts = 10e-9, GBP = 2750e6, tau = 180e-6, Kv = 150e3, Cd = 50e-12, Cf = 0.35e-12, Rf = 500e6)

    wop = GBP / (2 * pi * Kv)
    Cmod = Cf + Cd
    wmod = 1.0 / (Rf * Cmod)
    alfa = Cmod / (Cf * GBP)

    b0 = 1.0 / alfa
    a2 = 1.0
    a1 = 1.0 / alfa + wop + wmod
    a0 = 1.0 / (tau * alfa) + wmod*wop

    # then the transfer function in the *Laplace* s-domain looks like this:
    #                       b0
    #   T(s) = ----------------------------
    #             a2 * s^2 + a1 * s + a0

    # PolynomialRatio needs z-transform paramters: s- and z-domains can be connected by
    # the bilinear transform:
    #        z - 1
    # s = K -------- , K = Ts/2  , Ts - sampling period
    #        z + 1
    #
    # we can then convert T(s) to T(z):
    #              bz2 * z^2 + bz1 * z + bz0
    #   T(z) = -------------------------------
    #              az2 * z^2 + az1 * z + az0
    #

    K = 2/Ts

    az2 = 1.0   # normalized
    az1 = (2*a0 - 2*K^2)/(K^2 + a1*K + a0)
    az0 = (K^2 - a1*K + a0)/(K^2 + a1*K + a0)

    bz2 = b0/(K^2 + a1*K + a0)
    bz1 = 2*b0/(K^2 + a1*K + a0)
    bz0 = b0/(K^2 + a1*K + a0)

    myfilter = PolynomialRatio([bz2, bz1, bz0], [az2, az1, az0])

    filtered = filt(myfilter, vcat([0], diff(pulse)))

end

function add_baseline_and_extend_tail(wv::RadiationDetectorSignals.RDWaveform{T,U,TV,UV}, n_baseline_samples::Int, total_waveform_length::Int) where {T,U,TV,UV}
    new_signal::Vector{eltype(UV)} = Vector{eltype(UV)}(undef, total_waveform_length)
    new_signal[1:n_baseline_samples] .= zero(eltype(UV))
    if length(wv.value) <= total_waveform_length - n_baseline_samples
        new_signal[n_baseline_samples+1:n_baseline_samples+length(wv.value)] = wv.value
        new_signal[n_baseline_samples+length(wv.value)+1:end] .= wv.value[end]
    else
        new_signal[n_baseline_samples+1:end] = wv.value[1:total_waveform_length - n_baseline_samples]
    end
    new_times = if TV <: AbstractRange
        range( zero(first(wv.time)), step = step(wv.time), length = total_waveform_length )
    else
        error("Not yet definted for timestamps of type `$(TV)`")
    end
    return RDWaveform( new_times, new_signal )
end


module ActifMPCCmod

import Relaxation
import ParamSetmod
import NLPModels

"""
Type MPCC_actif : problème MPCC pénalisé avec slack et ensemble des contraintes actives

liste des constructeurs :
MPCC_actif(nlp::NLPModels.AbstractNLPModel,r::Float64,s::Float64,t::Float64,nb_comp::Int64,paramset::ParamSetmod.ParamSet,direction::Function,linesearch::Function)
MPCC_actif(nlp::NLPModels.AbstractNLPModel,r::Float64,s::Float64,t::Float64,w::Any,paramset::ParamSetmod.ParamSet,direction::Function,linesearch::Function)

liste des méthodes :
updatew(ma::MPCC_actif)
setf(ma::MPCC_actif, f::Function, xj::Any)
setw(ma::MPCC_actif, w::Any)
setbeta(ma::MPCC_actif,b::Float64)
sethess(ma::MPCC_actif,Hess::Array{Float64,2})

liste des accesseurs :
evalx(ma::MPCC_actif,x::Vector)
evald(ma::MPCC_actif,d::Vector)
redd(ma::MPCC_actif,d::Vector)
obj(ma::MPCC_actif,x::Vector)
ExtddDirection(ma::ActifMPCCmod.MPCC_actif,dr::Vector,xp::Vector,step::Float64)
grad(ma::MPCC_actif,x::Vector)
grad(ma::MPCC_actif,x::Vector,gradf::Vector)
hess(ma::MPCC_actif,x::Vector)
hess(ma::MPCC_actif,x::Vector,H::Array{Float64,2})

liste des fonctions :
LSQComputationMultiplier(ma::MPCC_actif,gradpen::Vector,xj::Vector)
RelaxationRule(ma::ActifMPCCmod.MPCC_actif,xj::Vector,lg::Vector,lh::Vector,lphi::Vector,wmax::Any)
PasMaxComp(ma::MPCC_actif,x::Vector,d::Vector)
AlphaChoix(alpha::Float64,alpha1::Float64,alpha2::Float64)
AlphaChoix(alpha::Float64,alpha1::Float64,alpha2::Float64,alpha3::Float64,alpha4::Float64)
PasMaxBound(ma::MPCC_actif,x::Vector,d::Vector) #TO DO
PasMax(ma::MPCC_actif,x::Vector,d::Vector)
"""

#TO DO List :
#Major :
#- Changer les w en set ou en liste
#- intégrer les contraintes de bornes sur x dans PasMax
#Minor :
#- copie de code dans PasMax
#- Plutôt que de passer le point courant -> stocker en point initial de mpcc ?

type MPCC_actif
 nlp::NLPModels.AbstractNLPModel # en fait on demande juste une fonction objectif, point initial, contraintes de bornes
 r::Float64
 s::Float64
 t::Float64

 w::Array{Bool,2} #matrice à 2 colonnes et de longueur 2nb_comp -- sparse

 n::Int64 #dans le fond est optionnel si on a nb_comp
 nb_comp::Int64

 #ensembles d'indices :
 w1::Array{Int64,1} # ensemble des indices (entre 0 et nb_comp) où la contrainte yG>=-r est active
 w2::Array{Int64,1} # ensemble des indices (entre 0 et nb_comp) où la contrainte yH>=-r est active
 w3::Array{Int64,1} # ensemble des indices (entre 0 et nb_comp) où la contrainte yG<=s+t*theta(yH,r) est active
 w4::Array{Int64,1} # ensemble des indices (entre 0 et nb_comp) où la contrainte yH<=s+t*theta(yG,r) est active
 wcomp::Array{Int64,1} #ensemble des indices (entre 0 et nb_comp) où la contrainte Phi<=0 est active
 w13c::Array{Int64,1} #ensemble des indices où les variables yG sont libres
 w24c::Array{Int64,1} #ensemble des indices où les variables yH sont libres
 wc::Array{Int64,1} #ensemble des indices des contraintes où yG et yH sont libres
 wcc::Array{Int64,1} #ensemble des indices des contraintes où yG et yH sont fixés

 wnew::Array{Bool,2} #dernières contraintes ajoutés

 #paramètres pour le calcul de la direction de descente
 crho::Float64 #constant such that : ||c(x)||_2 \approx crho*rho
 beta::Float64 #paramètre pour gradient conjugué
 Hess::Array{Float64,2} #inverse matrice hessienne approximée
 #Hd::Vector #produit inverse matrice hessienne et gradient (au lieu de la hessienne entière)

 paramset::ParamSetmod.ParamSet
 direction::Function #fonction qui calcul la direction de descente
 linesearch::Function #fonction qui calcul la recherche linéaire
end

"""
Constructeur recommandé pour MPCC_actif
"""
function MPCC_actif(nlp::NLPModels.AbstractNLPModel,r::Float64,s::Float64,t::Float64,nb_comp::Int64,paramset::ParamSetmod.ParamSet,direction::Function,linesearch::Function)

  n=length(nlp.meta.x0)-2*nb_comp
  xk=nlp.meta.x0[1:n]
  ygk=nlp.meta.x0[n+1:n+nb_comp]
  yhk=nlp.meta.x0[n+nb_comp+1:n+2*nb_comp]

  w=zeros(Bool,2*nb_comp,2)

  for l=1:nb_comp
   if ygk[l]==nlp.meta.lvar[n+l]
    w[l,1]=true;
   elseif ygk[l]==Relaxation.psi(yhk[l],r,s,t)
    w[l+nb_comp,1]=true;
   end
   if yhk[l]==nlp.meta.lvar[n+l+nb_comp]
    w[l,2]=true;
   elseif yhk[l]==Relaxation.psi(ygk[l],r,s,t)
    w[l+nb_comp,2]=true;
   end
  end

 return MPCC_actif(nlp,r,s,t,w,paramset,direction,linesearch)
end

function MPCC_actif(nlp::NLPModels.AbstractNLPModel,r::Float64,s::Float64,t::Float64,w::Array{Bool,2},paramset::ParamSetmod.ParamSet,direction::Function,linesearch::Function)

 nb_comp=Int(size(w,1)/2)
 n=length(nlp.meta.x0)-2*nb_comp
 w1=find(w[1:nb_comp,1])
 w2=find(w[1:nb_comp,2])
 w3=find(w[nb_comp+1:2*nb_comp,1])
 w4=find(w[nb_comp+1:2*nb_comp,2])
 wcomp=find(w[nb_comp+1:2*nb_comp,1] .| w[nb_comp+1:2*nb_comp,2])
 w13c=find(.!w[1:nb_comp,1] .& .!w[nb_comp+1:2*nb_comp,1])
 w24c=find(.!w[1:nb_comp,2] .& .!w[nb_comp+1:2*nb_comp,2])
 wc=find(.!w[1:nb_comp,1] .& .!w[1:nb_comp,2] .& .!w[nb_comp+1:2*nb_comp,1] .& .!w[nb_comp+1:2*nb_comp,2])
 wcc=find((w[1:nb_comp,1] .| w[nb_comp+1:2*nb_comp,1]) .& (w[1:nb_comp,2] .| w[nb_comp+1:2*nb_comp,2]))

 wnew=zeros(Bool,0,0)

 crho=1.0
 beta=0.0
 Hess=eye(n+2*nb_comp)

 return MPCC_actif(nlp,r,s,t,w,n,nb_comp,w1,w2,w3,w4,wcomp,w13c,w24c,wc,wcc,wnew,crho,beta,Hess,paramset,direction,linesearch)
end

"""
Methodes pour le MPCC_actif
"""
#Mise à jour des composantes liés à w
function updatew(ma::MPCC_actif)

 ma.w1=find(ma.w[1:ma.nb_comp,1])
 ma.w2=find(ma.w[1:ma.nb_comp,2])
 ma.w3=find(ma.w[ma.nb_comp+1:2*ma.nb_comp,1])
 ma.w4=find(ma.w[ma.nb_comp+1:2*ma.nb_comp,2])
 ma.wcomp=find(ma.w[ma.nb_comp+1:2*ma.nb_comp,1] .| ma.w[ma.nb_comp+1:2*ma.nb_comp,2])
 ma.w13c=find(.!ma.w[1:ma.nb_comp,1] .& .!ma.w[ma.nb_comp+1:2*ma.nb_comp,1])
 ma.w24c=find(.!ma.w[1:ma.nb_comp,2] .& .!ma.w[ma.nb_comp+1:2*ma.nb_comp,2])
 ma.wc=find(.!ma.w[1:ma.nb_comp,1] .& .!ma.w[1:ma.nb_comp,2] .& .!ma.w[ma.nb_comp+1:2*ma.nb_comp,1] .& .!ma.w[ma.nb_comp+1:2*ma.nb_comp,2])
 ma.wcc=find((ma.w[1:ma.nb_comp,1] .| ma.w[ma.nb_comp+1:2*ma.nb_comp,1]) .& (ma.w[1:ma.nb_comp,2] .| ma.w[ma.nb_comp+1:2*ma.nb_comp,2]))
 return ma
end

#Mise à jour de w
function setw(ma::MPCC_actif, w::Array{Bool,2})

 ma.wnew=w .& .!ma.w
 ma.w=w
 return updatew(ma)
end

function setbeta(ma::MPCC_actif,b::Float64)
 ma.beta=b
 return ma
end

function setcrho(ma::MPCC_actif,crho::Float64)
 ma.crho=crho
 return ma
end

function sethess(ma::MPCC_actif,Hess::Array{Float64,2})
 ma.Hess=Hess
 return ma
end

"""
Renvoie le vecteur x=[x,yg,yh] au complet
"""
function evalx(ma::MPCC_actif,x::Vector)
 #construction du vecteur de taille n+2nb_comp que l'on évalue :
 xf=ma.s*ones(ma.n+2*ma.nb_comp)
 xf[1:ma.n]=x[1:ma.n]
 xf[ma.w13c+ma.n]=x[ma.n+1:ma.n+length(ma.w13c)]
 xf[ma.w24c+ma.n+ma.nb_comp]=x[ma.n+length(ma.w13c)+1:ma.n+length(ma.w13c)+length(ma.w24c)]

 #on regarde les variables yG fixées :
 xf[ma.w1+ma.n]=ma.nlp.meta.lvar[ma.w1+ma.n]
 xf[ma.w3+ma.n]=Relaxation.psi(xf[ma.w3+ma.n+ma.nb_comp],ma.r,ma.s,ma.t)
 #on regarde les variables yH fixées :
 xf[ma.w2+ma.n+ma.nb_comp]=ma.nlp.meta.lvar[ma.w2+ma.n+ma.nb_comp]
 xf[ma.w4+ma.n+ma.nb_comp]=Relaxation.psi(xf[ma.w4+ma.n],ma.r,ma.s,ma.t)

 return xf
end

"""
Renvoie la direction d au complet (avec des 0 aux actifs)
"""
function evald(ma::MPCC_actif,d::Vector)
 df=zeros(ma.n+2*ma.nb_comp)
 df[1:ma.n]=d[1:ma.n]
 df[ma.w13c+ma.n]=d[ma.n+1:ma.n+length(ma.w13c)]
 df[ma.w24c+ma.nb_comp+ma.n]=d[ma.n+length(ma.w13c)+1:ma.n+length(ma.w13c)+length(ma.w24c)]
 return df
end

"""
Renvoie la direction d réduite
"""
function redd(ma::MPCC_actif,d::Vector)

  df=zeros(ma.n+length(ma.w13c)+length(ma.w24c))
  df[1:ma.n]=d[1:ma.n]
  df[ma.n+1:ma.n+length(ma.w13c)]=d[ma.w13c+ma.n]
  df[ma.n+length(ma.w13c)+1:ma.n+length(ma.w13c)+length(ma.w24c)]=d[ma.w24c+ma.nb_comp+ma.n]

 return df
end

function redd(ma::MPCC_actif,d::Vector,w::Array{Int64,1})

  df=zeros(length(w))
  df[1:length(w)]=d[w]

 return df
end

"""
Evalue la fonction objectif d'un MPCC actif : x
"""
function obj(ma::MPCC_actif,x::Vector)

 if length(x)==ma.n+2*ma.nb_comp
  return NLPModels.obj(ma.nlp,x)
 else
  return NLPModels.obj(ma.nlp,evalx(ma,x))
 end

end

"""
Calcul la direction étendue à partir de la direction du sous-domaine
complète une version étendue de la direction de descente :
dr : la direction dans le sous-espace actif
xp : le nouvel itéré
step : le pas utilisé pour calculé xp

utilise la formule suivante :
yG+=yG+alpha*dyG --> dyG=(yG+-yG)/alpha
"""
function ExtddDirection(ma::ActifMPCCmod.MPCC_actif,dr::Vector,xp::Vector,step::Float64)
 d=evald(ma,dr) #evald rempli les trous par des 0
 x=evalx(ma,xp)

 d[1:ma.n]=dr[1:ma.n]

 psip=Relaxation.psi(x[ma.n+1:ma.n+2*ma.nb_comp],ma.r,ma.s,ma.t)
 psi=Relaxation.psi(x[ma.n+1:ma.n+2*ma.nb_comp]-step*d[ma.n+1:ma.n+2*ma.nb_comp],ma.r,ma.s,ma.t)
 #d[ma.w1]=0 #yG fixé
 #d[ma.w2]=0 #yH fixé

 d[ma.n+ma.w3]=(psip[ma.w3+ma.nb_comp]-psi[ma.w3+ma.nb_comp])/step #yG fixé
 d[ma.n+ma.nb_comp+ma.w4]=(psip[ma.w4]-psi[ma.w4])/step #yH fixé
 #d[ma.n+ma.w13c]=dr[ma.n+ma.w13c] #yG est libre
 #d[ma.n+ma.nb_comp+ma.w24c]=dr[ma.n+ma.nb_comp+ma.w24c] #yG est libre

 return d
end

"""
Evalue le gradient de la fonction objectif d'un MPCC actif
x est le vecteur réduit
"""
function grad(ma::MPCC_actif,x::Vector)

 #on calcul xf le vecteur complet
 xf=evalx(ma,x)
 #construction du vecteur gradient de taille n+2nb_comp
 gradf=NLPModels.grad(ma.nlp,xf)

 return length(x)==ma.n+2*ma.nb_comp?gradf:grad(ma,x,gradf)
end

function grad(ma::MPCC_actif,x::Vector,gradf::Vector)

 #on calcul xf le vecteur complet
 xf=evalx(ma,x)
 #construction du vecteur gradient de taille n+2nb_comp
 #gradf=NLPModels.grad(ma.nlp,xf)

 gradg=Array{Float64}
 # Conditionnelles pour gérer le cas où w1 et w3 est vide
 #if isempty(ma.w1) && isempty(ma.w3)
 if isempty(ma.w4)
  gradg=zeros(length(ma.w13c))
 #elseif !isempty(ma.w13c) #certaines variables sont fixés
 elseif !isempty(ma.w4) #certaines variables sont fixés
  tmp=zeros(ma.nb_comp)
  #tmp[ma.w3]=Relaxation.dpsi(xf[ma.w3+ma.n+ma.nb_comp],ma.r,ma.s,ma.t).*gradf[ma.w3+ma.n]
  tmp[ma.w4]=Relaxation.dpsi(xf[ma.w4+ma.n],ma.r,ma.s,ma.t).*gradf[ma.w4+ma.nb_comp+ma.n]
  gradg=redd(ma,tmp,ma.w13c)
 else #ma.w13c est vide
  gradg=Float64[]
 end

 gradh=Array{Float64,1}
 #if isempty(ma.w2) && isempty(ma.w4)
 if isempty(ma.w3)
  gradh=zeros(length(ma.w24c))
 #elseif !isempty(ma.w24c)
 elseif !isempty(ma.w3)
  tmp=zeros(ma.nb_comp)
  #tmp[ma.w4]=Relaxation.dpsi(xf[ma.w4+ma.n],ma.r,ma.s,ma.t).*gradf[ma.w4+ma.nb_comp+ma.n]
  tmp[ma.w3]=Relaxation.dpsi(xf[ma.w3+ma.n+ma.nb_comp],ma.r,ma.s,ma.t).*gradf[ma.w3+ma.n]
  gradh=redd(ma,tmp,ma.w24c)
 else
  gradh=Float64[]
 end

 return vcat(gradf[1:ma.n],gradf[ma.w13c+ma.n]+gradg,gradf[ma.w24c+ma.nb_comp+ma.n]+gradh)
end

"""
Evalue la matrice hessienne de la fonction objectif d'un MPCC actif
x est le vecteur réduit
"""
function hess(ma::MPCC_actif,x::Vector)

 #on calcul xf le vecteur complet
 xf=evalx(ma,x)

 #construction de la hessienne de taille (n+2nb_comp)^2

 #H=NLPModels.hess(ma.nlp,xf) #renvoi la triangulaire inférieure tril(H,-1)'
 #H=H+tril(H,-1)'

 H=ma.nlp.H(xf)
 H=H+tril(H,-1)'

 if ma.nb_comp>0
  return hess(ma,x,H)
 else
  return H
 end
end

"""
A partir de la hessienne complète (ou une approximation) : calcul la hessienne dans le sous-espace actif
"""
function hess(ma::MPCC_actif,x::Vector,H::Array{Float64,2})

 #on calcul xf le vecteur complet
 xf=evalx(ma,x)
 nred=length(x)
 nnb=ma.n+ma.nb_comp;nnbt=ma.n+2*ma.nb_comp;
 #construction du vecteur gradient de taille n+2nb_comp
 gradf=NLPModels.grad(ma.nlp,xf)

#la hessienne des variables du sous-espace (nredxnred)
 Hred=vcat(hcat(H[1:ma.n,1:ma.n],H[1:ma.n,ma.n+ma.w13c],H[1:ma.n,nnb+ma.w24c]),
             hcat(H[ma.n+ma.w13c,1:ma.n],H[ma.n+ma.w13c,ma.n+ma.w13c],H[ma.n+ma.w13c,nnb+ma.w24c]),
             hcat(H[nnb+ma.w24c,1:ma.n],H[nnb+ma.w24c,ma.n+ma.w13c],H[nnb+ma.w24c,ma.n+ma.nb_comp+ma.w24c]))

 if isempty(ma.w4)
  hessg=sparse(zeros(length(ma.w13c),nred))
 else
  hessg=sparse(zeros(length(ma.w13c),nred))

  w4r=zeros(Int64,length(ma.w4))
  for i=1:length(ma.w4)
   w4r[i]=findfirst(x->x==ma.w4[i],ma.w13c)
  end
  hessg=diagm(Relaxation.ddpsi(xf[ma.w4+ma.n],ma.r,ma.s,ma.t).*gradf[ma.w4+ma.nb_comp+ma.n])
  hessg+=diagm(Relaxation.dpsi(xf[ma.w4+ma.n],ma.r,ma.s,ma.t))*H[ma.w4+ma.nb_comp+ma.n,ma.w4+ma.nb_comp+ma.n]

  Hred[ma.n+w4r,ma.n+w4r]+=hessg
 end

 if isempty(ma.w3)
  hessh=sparse(zeros(length(ma.w24c),nred))
 else
  hessh=sparse(zeros(length(ma.w24c),nred))

  w3r=zeros(Int64,length(ma.w3))
  for i=1:length(ma.w3)
   w3r[i]=findfirst(x->x==ma.w3[i],ma.w24c)
  end
  hessh=diagm(Relaxation.ddpsi(xf[ma.w3+ma.nb_comp+ma.n],ma.r,ma.s,ma.t).*gradf[ma.w3+ma.n])
  hessh+=diagm(Relaxation.dpsi(xf[ma.w3+ma.nb_comp+ma.n],ma.r,ma.s,ma.t))*H[ma.w3+ma.n,ma.w3+ma.n]

  Hred[ma.n+length(ma.w13c)+w3r,ma.n+length(ma.w13c)+w3r]+=hessh
 end

 return Hred
end

"""
LSQComputationMultiplier(ma::MPCC_actif,x::Vector) :
ma MPCC_Actif
xj in n+2nb_comp
gradpen in n+2nb_comp

calcul la valeur des multiplicateurs de Lagrange pour la contrainte de complémentarité en utilisant moindre carré
"""
function LSQComputationMultiplier(ma::MPCC_actif,gradpen::Vector,xj::Vector)

 dg=Relaxation.dphi(xj[ma.n+1:ma.n+ma.nb_comp],xj[ma.n+ma.nb_comp+1:ma.n+2*ma.nb_comp],ma.r,ma.s,ma.t)
 gx=dg[1:ma.nb_comp];gy=dg[ma.nb_comp+1:2*ma.nb_comp]

 #matrices des contraintes actives : (lg,lh,lphi)'*A=b
 nw1=length(ma.w1)
 nw2=length(ma.w2)
 nwcomp=length(ma.wcomp)
 Dlg=-diagm(ones(nw1))
 Dlh=-diagm(ones(nw2))
 Dlphig=diagm(collect(gx)[ma.wcomp])
 Dlphih=diagm(collect(gy)[ma.wcomp])
 A=[hcat(Dlg,zeros(nw1,nw2+2*nwcomp));hcat(zeros(nw2,nw1),Dlh,zeros(nw2,2*nwcomp));hcat(zeros(nwcomp,nw1+nw2),Dlphig,Dlphih)] # nw1+nw2+nwcomp x nw1+nw2+2nwcomp
 #second membre
 b=-[gradpen[ma.w1+ma.n];gradpen[ma.wcomp+ma.n];gradpen[ma.w2+ma.n+ma.nb_comp];gradpen[ma.wcomp+ma.n+ma.nb_comp]] #nw1+nw2+2nwcomp
 #on calcule la solution par pseudo-inverse :
 l=pinv(A')*b

 lk=zeros(3*ma.nb_comp)
 lk[ma.w1]=l[1:nw1]
 lk[ma.nb_comp+ma.w2]=l[nw1+1:nw1+nw2]
 lk[2*ma.nb_comp+ma.wcomp]=l[nw1+nw2+1:nw1+nw2+nwcomp]

 return lk[1:ma.nb_comp],lk[ma.nb_comp+1:2*ma.nb_comp],lk[2*ma.nb_comp+1:3*ma.nb_comp]
end

"""
Définit la règle de relaxation de l'ensemble des contraintes
RelaxationRule(ma::ActifMPCCmod.MPCC_actif,xj::Vector,lg,lh,lphi,wmax)
xj          : de taille n+2nb_comp
(lg,lh,phi) : valeur multiplicateurs
wmax        : ensemble des contraintes qui viennent d'être ajouté

output : MPCC_actif (avec les ensembles de contraintes actives mis à jour
"""
function RelaxationRule(ma::ActifMPCCmod.MPCC_actif,xj::Vector,lg::Vector,lh::Vector,lphi::Vector,wmax::Array{Bool,2})

  copy_wmax=copy(wmax)

  # Relaxation de l'ensemble d'activation : désactive toutes les contraintes négatives
  ma.w[find(x -> x<0,[lg;lphi;lh;lphi])]=zeros(Bool,length(find(x -> x<0,[lg;lphi;lh;lphi])))
  # Règle d'anti-cyclage : on enlève pas une contrainte qui vient d'être ajouté.
  ma.w[find(x->x==1.0,copy_wmax)]=ones(Bool,length(find(x->x==1.0,copy_wmax)))

 return ActifMPCCmod.updatew(ma)
end

"""
Calcul le pas maximum que l'on peut prendre dans une direction d (par rapport à la contrainte de complémentarité relaxé)
d : direction réduit
xj : itéré réduit

output :
alpha : le pas maximum
w_save : l'ensemble des contraintes qui vont devenir actives si choisit alphamax
"""
function PasMaxComp(ma::MPCC_actif,x::Vector,d::Vector)

 #initialisation
 alpha=Inf #pas maximum que l'on peut prendre
 w_save=copy(ma.w) #double tableau des indices avec les contraintes activent en x+alpha*d

 #les indices où la première composante est libre
 for i in ma.w13c
  wr13=findfirst(x->x==i,ma.w13c) #l'indice relatif dans les variables libre
  iw13c=ma.n+wr13

  bloque=(i in ma.w2) && (i in ma.w4)
  if !(i in ma.w24c) && !bloque && d[iw13c]<0
   #on prend le plus petit entre x+alpha*dx>=-r et s+tTheta(x+alpha*dx-s)>=-r
   alpha11=(ma.nlp.meta.lvar[iw13c]-x[iw13c])/d[iw13c]
   alpha12=(Relaxation.invpsi(ma.nlp.meta.lvar[iw13c],ma.r,ma.s,ma.t)-x[iw13c])/d[iw13c]

   alphag=AlphaChoix(alpha,alpha11,alpha12)
   if alphag<=alpha
    alpha=alphag
    w_save=copy(ma.w)
   end

   #update of the active set
   if alpha11==alpha
    w_save[i,1]=true
   end
   if alpha12==alpha
    w_save[i+ma.nb_comp,2]=true
    w_save[i,2]=true
   end

  elseif bloque
   alpha=0.0
  end
 end

 #c'est un copie coller d'au dessus => exporter dans une fonction
 #les indices où la deuxième composante est libre
 for i in ma.w24c
  wr24=findfirst(x->x==i,ma.w24c) #l'indice relatif dans les variables libre
  iw24c=ma.n+length(ma.w13c)+wr24

  bloque=(i in ma.w1) && (i in ma.w3)
  if !(i in ma.w13c) && !bloque && d[iw24c]<0
   #on prend le plus petit entre y+alpha*dy>=-r et s+tTheta(y+alpha*dy-s)>=-r
   alpha21=(ma.nlp.meta.lvar[iw24c]-x[iw24c])/d[iw24c]
   alpha22=(Relaxation.invpsi(ma.nlp.meta.lvar[iw24c],ma.r,ma.s,ma.t)-x[iw24c])/d[iw24c]

   alphah=AlphaChoix(alpha,alpha21,alpha22)
   if alphah<=alpha
    alpha=alphah
    w_save=copy(ma.w)
   end

   #on met à jour les contraintes
   if alpha21==alpha
    w_save[i,2]=true
   end
   if alpha22==alpha
    w_save[i+ma.nb_comp,1]=true
    w_save[i,1]=true
   end
  elseif bloque
   alpha=0.0
  end
 end

 #enfin les indices où les deux sont libres
 for i in ma.wc
  #yG-psi(yH)=0 ou yH-psi(yG)=0
  wr1=findfirst(x->x==i,ma.w13c) #l'indice relatif dans les variables libre
  wr2=findfirst(x->x==i,ma.w24c) #l'indice relatif dans les variables libre
  iwr1=wr1+ma.n;iwr2=wr2+length(ma.w13c)+ma.n;

  #alphac=Relaxation.AlphaThetaMax(x[i+ma.n],d[i+ma.n],x[i+length(ma.w13c)+ma.n],d[i+length(ma.w13c)+ma.n],ma.r,ma.s,ma.t)
  alphac=Relaxation.AlphaThetaMax(x[iwr1],d[iwr1],x[iwr2],d[iwr2],ma.r,ma.s,ma.t)
  #yG-tb=0
  #alphac11=d[i+ma.n]<0 ? (ma.nlp.meta.lvar[ma.n+i]-x[i+ma.n])/d[i+ma.n] : Inf
  alphac11=d[iwr1]<0 ? (ma.nlp.meta.lvar[iwr1]-x[iwr1])/d[iwr1] : Inf
  #yH-tb=0
  #alphac21=d[i+length(ma.w13c)+ma.n]<0 ? (ma.nlp.meta.lvar[ma.n+length(ma.w13c)+i]-x[i+length(ma.w13c)+ma.n])/d[i+length(ma.w13c)+ma.n] : Inf
  alphac21=d[iwr2]<0 ? (ma.nlp.meta.lvar[iwr2]-x[iwr2])/d[iwr2] : Inf

  alphagh=AlphaChoix(alpha,alphac[1],alphac[2],alphac11,alphac21)

  if alphagh<=alpha
   alpha=alphagh
   w_save=copy(ma.w)
  end

   if alphac[1]==alpha
    w_save[i+ma.nb_comp,1]=true
   end
   if alphac[2]==alpha
    w_save[i+ma.nb_comp,2]=true
   end
   if alphac11==alpha
    w_save[i,1]=true
   end
   if alphac21==alpha
    w_save[i,2]=true
   end

 end #fin boucle for ma.wc

 return alpha,w_save,Array(w_save .& .!ma.w)
end

"""
Met à jour alpha si :
1) il est plus petit
2) il est non-nul
"""
function AlphaChoix(alpha::Float64,alpha1::Float64,alpha2::Float64)
 return AlphaChoix(alpha,alpha1,alpha2,0.0,0.0)
end

function AlphaChoix(alpha::Float64,alpha1::Float64,alpha2::Float64,alpha3::Float64,alpha4::Float64)
 prec=eps(Float64)
 a=alpha1,alpha2,alpha3,alpha4
 a=a[find(x->x>=prec,collect(a))]
 if isempty(a)
  a=max(alpha1,alpha2,alpha3,alpha4)
 end

 return min(minimum(a),alpha)
end

"""
Calcul le pas maximum que l'on peut prendre dans une direction d par rapport aux contraintes de bornes sur x.
d : direction réduit
xj : itéré réduit

output :
alpha : le pas maximum
w_save : l'ensemble des contraintes qui vont devenir actives si choisit alphamax
"""
function PasMaxBound(ma::MPCC_actif,x::Vector,d::Vector)
 if max(l-x)>0 || max(x-u)>0
  println("Error PasMaxBound: infeasible x")
 end
 #ma.nlp.meta.lvar[1:ma.n]
 #ma.nlp.meta.uvar[1:ma.n]
 return #TO DO : nécessite d'ajouter les x dans l'ensemble des contraintes actives
end

"""
Calcul le pas maximum que l'on peut prendre dans une direction d
d : direction réduit
xj : itéré réduit

output :
alpha : le pas maximum
w_save : l'ensemble des contraintes qui vont devenir actives si choisit alphamax
"""
function PasMax(ma::MPCC_actif,x::Vector,d::Vector)

 if ma.nb_comp>0
  #on récupère les infos sur la contrainte de complémentarité
  alpha,w_save,w_new=PasMaxComp(ma,x,d)
 else
  alpha=Inf;w_save=zeros(Bool,0,0);w_new=ma.w;
 end
 #alpha,w_save,w_new=PasMaxBound(ma,x,d)

 if alpha<0.0
  println("PasMax error: pas maximum négatif.")
  return
 end

 return alpha,w_save,w_new
end

#end of module
end


du[1]=(-4).*(2+P1.^2+(-2).*2.^(1/2).*P2+P2.^2+2.*P1.*((-1).*2.^(1/2)+P2) +(Q1+(-1).*Q2).^2).^(-2).*(2+P1.^2+2.*2.^(1/2).*P2+P2.^2+2.*P1.*(   2.^(1/2)+P2)+(Q1+(-1).*Q2).^2).^(-2).*(P2.^2+Q2.^2).*(2+P1.^2+2.*   P1.*(2.^(1/2)+(-1).*P2)+(-2).*2.^(1/2).*P2+P2.^2+(Q1+Q2).^2).^(-2)   .*(2+P1.^2+2.*2.^(1/2).*P2+P2.^2+(-2).*P1.*(2.^(1/2)+P2)+(Q1+Q2) .^2).^(-2).*((P1.^2+Q1.^2).*(2+(-2).*2.^(1/2).*(P1+P2)+(P1+P2).^2+  (Q1+(-1).*Q2).^2).^(-1).*(2+2.*2.^(1/2).*(P1+P2)+(P1+P2).^2+(Q1+(   -1).*Q2).^2).^(-1).*(P2.^2+Q2.^2).*(2+2.*2.^(1/2).*(P1+(-1).*P2)+(   P1+(-1).*P2).^2+(Q1+Q2).^2).^(-1).*(2+(P1+(-1).*P2).^2+2.*2.^(1/2)  .*((-1).*P1+P2)+(Q1+Q2).^2).^(-1)).^(-1/2).*(3.*P1.^9+24.*P1.^6.*   P2.*Q1.*Q2+4.*P1.^7.*((-4)+(-2).*P2.^2+3.*Q1.^2+2.*Q2.^2)+8.*   P1.^4.*P2.*Q1.*Q2.*(4+(-2).*P2.^2+7.*Q1.^2+2.*Q2.^2)+4.*P1.^3.*   Q1.^2.*(12+3.*P2.^4+3.*Q1.^4+(-4).*Q2.^2+3.*Q2.^4+2.*Q1.^2.*((-2)+   Q2.^2)+(-2).*P2.^2.*((-2)+Q1.^2+Q2.^2))+8.*P2.*Q1.^3.*Q2.*(P2.^4+(   Q1.^2+(-1).*Q2.^2).^2+2.*P2.^2.*(2+Q1.^2+Q2.^2)+(-4).*(3+Q1.^2+   Q2.^2))+2.*P1.^5.*(12+3.*P2.^4+9.*Q1.^4+(-4).*Q2.^2+3.*Q2.^4+8.*   Q1.^2.*((-2)+Q2.^2)+(-2).*P2.^2.*((-2)+4.*Q1.^2+Q2.^2))+(-8).*   P1.^2.*P2.*Q1.*Q2.*(P2.^4+(-5).*Q1.^4+Q2.^4+(-8).*Q1.^2.*(2+Q2.^2)   +(-4).*(3+Q2.^2)+2.*P2.^2.*(2+4.*Q1.^2+Q2.^2))+P1.*((-1).*P2.^8+   3.*Q1.^8+8.*Q1.^2.*((-2)+Q2.^2).*(2+Q2.^2).^2+(-1).*(2+Q2.^2).^4+(   -4).*P2.^6.*((-2)+2.*Q1.^2+Q2.^2)+(-2).*Q1.^4.*(20+36.*Q2.^2+5.*   Q2.^4)+(-2).*P2.^4.*(12+5.*Q1.^4+(-4).*Q2.^2+3.*Q2.^4+4.*Q1.^2.*((   -2)+Q2.^2))+4.*P2.^2.*((-1).*((-2)+Q2.^2).*(2+Q2.^2).^2+Q1.^4.*(   18+23.*Q2.^2)+2.*Q1.^2.*(4+20.*Q2.^2+Q2.^4))));
du[2]=4.*(P1.^2+Q1.^2).*(2+P1.^2+(-2).*2.^(1/2).*P2+P2.^2+2.*P1.*((-1).*  2.^(1/2)+P2)+(Q1+(-1).*Q2).^2).^(-2).*(2+P1.^2+2.*2.^(1/2).*P2+  P2.^2+2.*P1.*(2.^(1/2)+P2)+(Q1+(-1).*Q2).^2).^(-2).*(2+P1.^2+2.*   P1.*(2.^(1/2)+(-1).*P2)+(-2).*2.^(1/2).*P2+P2.^2+(Q1+Q2).^2).^(-2)   .*(2+P1.^2+2.*2.^(1/2).*P2+P2.^2+(-2).*P1.*(2.^(1/2)+P2)+(Q1+Q2)   .^2).^(-2).*((P1.^2+Q1.^2).*(2+(-2).*2.^(1/2).*(P1+P2)+(P1+P2).^2+   (Q1+(-1).*Q2).^2).^(-1).*(2+2.*2.^(1/2).*(P1+P2)+(P1+P2).^2+(Q1+(   -1).*Q2).^2).^(-1).*(P2.^2+Q2.^2).*(2+2.*2.^(1/2).*(P1+(-1).*P2)+(   P1+(-1).*P2).^2+(Q1+Q2).^2).^(-1).*(2+(P1+(-1).*P2).^2+2.*2.^(1/2)   .*((-1).*P1+P2)+(Q1+Q2).^2).^(-1)).^(-1/2).*((-3).*P2.^9+(-24).*   P1.*P2.^6.*Q1.*Q2+8.*P1.*P2.^4.*Q1.*Q2.*((-4)+2.*P1.^2+(-2).*   Q1.^2+(-7).*Q2.^2)+4.*P2.^7.*(4+2.*P1.^2+(-2).*Q1.^2+(-3).*Q2.^2)+  P2.*((P1.^4+2.*P1.^2.*((-2)+Q1.^2)+(2+Q1.^2).^2).^2+8.*(P1.^6+  P1.^4.*((-2)+Q1.^2)+(-1).*((-2)+Q1.^2).*(2+Q1.^2).^2+(-1).*P1.^2.*  (4+20.*Q1.^2+Q1.^4)).*Q2.^2+2.*(20+5.*P1.^4+36.*Q1.^2+5.*Q1.^4+(-2).*P1.^2.*(18+23.*Q1.^2)).*Q2.^4+(-3).*Q2.^8)+(-4).*P2.^3.*  Q2.^2.*(12+3.*P1.^4+(-4).*Q1.^2+3.*Q1.^4+2.*((-2)+Q1.^2).*Q2.^2+   3.*Q2.^4+(-2).*P1.^2.*((-2)+Q1.^2+Q2.^2))+(-8).*P1.*Q1.*Q2.^3.*(   P1.^4+(Q1.^2+(-1).*Q2.^2).^2+2.*P1.^2.*(2+Q1.^2+Q2.^2)+(-4).*(3+   Q1.^2+Q2.^2))+(-2).*P2.^5.*(12+3.*P1.^4+(-4).*Q1.^2+3.*Q1.^4+8.*((  -2)+Q1.^2).*Q2.^2+9.*Q2.^4+(-2).*P1.^2.*((-2)+Q1.^2+4.*Q2.^2))+8.*  P1.*P2.^2.*Q1.*Q2.*((-12)+P1.^4+(-4).*Q1.^2+Q1.^4+(-8).*(2+Q1.^2)   .*Q2.^2+(-5).*Q2.^4+2.*P1.^2.*(2+Q1.^2+4.*Q2.^2)));
du[3]=4.*(2+P1.^2+(-2).*2.^(1/2).*P2+P2.^2+2.*P1.*((-1).*2.^(1/2)+P2)+(Q1+(-1).*Q2).^2).^(-2).*(2+P1.^2+2.*2.^(1/2).*P2+P2.^2+2.*P1.*( 2.^(1/2)+P2)+(Q1+(-1).*Q2).^2).^(-2).*(P2.^2+Q2.^2).*(2+P1.^2+2.* P1.*(2.^(1/2)+(-1).*P2)+(-2).*2.^(1/2).*P2+P2.^2+(Q1+Q2).^2).^(-2) .*(2+P1.^2+2.*2.^(1/2).*P2+P2.^2+(-2).*P1.*(2.^(1/2)+P2)+(Q1+Q2) .^2).^(-2).*((P1.^2+Q1.^2).*(2+(-2).*2.^(1/2).*(P1+P2)+(P1+P2).^2+(Q1+(-1).*Q2).^2).^(-1).*(2+2.*2.^(1/2).*(P1+P2)+(P1+P2).^2+(Q1+( -1).*Q2).^2).^(-1).*(P2.^2+Q2.^2).*(2+2.*2.^(1/2).*(P1+(-1).*P2)+(P1+(-1).*P2).^2+(Q1+Q2).^2).^(-1).*(2+(P1+(-1).*P2).^2+2.*2.^(1/2).*((-1).*P1+P2)+(Q1+Q2).^2).^(-1)).^(-1/2).*(3.*Q1.^9+24.*P1.*P2.*Q1.^6.*Q2+4.*Q1.^7.*(4+3.*P1.^2+2.*P2.^2+(-2).*Q2.^2)+8.*P1.*P2.* Q1.^4.*Q2.*((-4)+7.*P1.^2+2.*P2.^2+(-2).*Q2.^2)+4.*P1.^2.*Q1.^3.*(12+3.*P1.^4+4.*P2.^2+3.*P2.^4+(-2).*(2+P2.^2).*Q2.^2+3.*Q2.^4+2.*P1.^2.*(2+P2.^2+(-1).*Q2.^2))+2.*Q1.^5.*(12+9.*P1.^4+4.*P2.^2+3.*P2.^4+(-2).*(2+P2.^2).*Q2.^2+3.*Q2.^4+8.*P1.^2.*(2+P2.^2+(-1).* Q2.^2))+8.*P1.*P2.*Q1.^2.*Q2.*(5.*P1.^4+8.*P1.^2.*((-2)+P2.^2+(-1) .*Q2.^2)+4.*(3+(-1).*P2.^2+Q2.^2)+(-1).*(P2.^2+Q2.^2).^2)+8.*P1.^3.*P2.*Q2.*(P1.^4+4.*((-3)+P2.^2+(-1).*Q2.^2)+2.*P1.^2.*(2+( -1).*P2.^2+Q2.^2)+(P2.^2+Q2.^2).^2)+Q1.*(3.*P1.^8+(-1).*(P2.^4+2.* P2.^2.*((-2)+Q2.^2)+(2+Q2.^2).^2).^2+(-2).*P1.^4.*(20+5.*P2.^4+ 36.*Q2.^2+5.*Q2.^4+(-2).*P2.^2.*(18+23.*Q2.^2))+8.*P1.^2.*(P2.^6+ P2.^4.*((-2)+Q2.^2)+(-1).*((-2)+Q2.^2).*(2+Q2.^2).^2+(-1).*P2.^2.*(4+20.*Q2.^2+Q2.^4))));
du[4]=(-4).*(P1.^2+Q1.^2).*(2+P1.^2+(-2).*2.^(1/2).*P2+P2.^2+2.*P1.*((-1).*2.^(1/2)+P2)+(Q1+(-1).*Q2).^2).^(-2).*(2+P1.^2+2.*2.^(1/2).*P2+P2.^2+2.*P1.*(2.^(1/2)+P2)+(Q1+(-1).*Q2).^2).^(-2).*((-8).*P1.*P2.^3.*Q1.*(P1.^4+4.*((-3)+P2.^2+(-1).*Q1.^2)+2.*P1.^2.*(2+(-1).*P2.^2+Q1.^2)+(P2.^2+Q1.^2).^2)+(P1.^8+(-3).*P2.^8+8.*P2.^2.*((-2)+Q1.^2).*(2+Q1.^2).^2+(2+Q1.^2).^4+4.*P1.^6.*((-2)+(-2).*P2.^2+Q1.^2)+2.*P1.^4.*(12+8.*P2.^2+5.*P2.^4+(-4).*(1+P2.^2).*Q1.^2+3.*Q1.^4)+2.*P2.^4.*(20+36.*Q1.^2+5.*Q1.^4)+4.*P1.^2.*(((-2)+Q1.^2).*(2+Q1.^2).^2+(-1).*P2.^4.*(18+23.*Q1.^2)+2.*P2.^2.*(4+20.*Q1.^2+Q1.^4))).*Q2+8.*P1.*P2.*Q1.*(P1.^4+(-5).*P2.^4+Q1.^4+8.*P2.^2.*(2+Q1.^2)+(-4).*(3+Q1.^2)+2.*P1.^2.*(2+(-4).*P2.^2+Q1.^2)).*Q2.^2+( -4).*P2.^2.*(12+3.*P1.^4+4.*P2.^2+3.*P2.^4+(-2).*(2+P2.^2).*Q1.^2+3.*Q1.^4+2.*P1.^2.*(2+P2.^2+(-1).*Q1.^2)).*Q2.^3+(-8).*P1.*P2.* Q1.*(2.*P1.^2+7.*P2.^2+(-2).*(2+Q1.^2)).*Q2.^4+(-2).*(12+3.*P1.^4+ 9.*P2.^4+(-4).*Q1.^2+3.*Q1.^4+P1.^2.*(4+8.*P2.^2+(-2).*Q1.^2)+(-8) .*P2.^2.*((-2)+Q1.^2)).*Q2.^5+(-24).*P1.*P2.*Q1.*Q2.^6+(-4).*(4+ 2.*P1.^2+3.*P2.^2+(-2).*Q1.^2).*Q2.^7+(-3).*Q2.^9).*(2+P1.^2+2.* P1.*(2.^(1/2)+(-1).*P2)+(-2).*2.^(1/2).*P2+P2.^2+(Q1+Q2).^2).^(-2) .*(2+P1.^2+2.*2.^(1/2).*P2+P2.^2+(-2).*P1.*(2.^(1/2)+P2)+(Q1+Q2) .^2).^(-2).*((P1.^2+Q1.^2).*(2+(-2).*2.^(1/2).*(P1+P2)+(P1+P2).^2+ (Q1+(-1).*Q2).^2).^(-1).*(2+2.*2.^(1/2).*(P1+P2)+(P1+P2).^2+(Q1+(-1).*Q2).^2).^(-1).*(P2.^2+Q2.^2).*(2+2.*2.^(1/2).*(P1+(-1).*P2)+( P1+(-1).*P2).^2+(Q1+Q2).^2).^(-1).*(2+(P1+(-1).*P2).^2+2.*2.^(1/2) .*((-1).*P1+P2)+(Q1+Q2).^2).^(-1)).^(-1/2);


##### Multi-physics types

using CLIMAParameters.Planet: T_freeze, cv_l

using CloudMicrophysics.Microphysics_0M
using CloudMicrophysics.Microphysics_1M
import CloudMicrophysics

const CM1M = CloudMicrophysics.Microphysics_1M

export RainSnow_1M
export WarmRain_1M
export Diffusion
export Subsidence
export RemovePrecipitation

struct Subsidence{FT} <: TendencyDef{Source}
    D::FT
end

prognostic_vars(::Subsidence) = (Mass(), Energy(), TotalMoisture())

# Subsidence includes tendencies in Mass, Energy and TotalMoisture equations:

subsidence_velocity(subsidence::Subsidence{FT}, z::FT) where {FT} =
    -subsidence.D * z

struct PressureGradient <: TendencyDef{Flux{FirstOrder}} end
struct Pressure <: TendencyDef{Flux{FirstOrder}} end
struct Advect <: TendencyDef{Flux{FirstOrder}} end
struct Diffusion <: TendencyDef{Flux{SecondOrder}} end
struct MoistureDiffusion <: TendencyDef{Flux{SecondOrder}} end

"""
    RemovePrecipitation <: TendencyDef{Source}

A sink to `q_tot` when cloud condensate is exceeding a threshold.
The threshold is defined either in terms of condensate or supersaturation.
The removal rate is implemented as a relaxation term
in the CloudMicrophysics.jl Microphysics_0M module.
The default thresholds and timescale are defined in CLIMAParameters.jl.
"""
struct RemovePrecipitation <: TendencyDef{Source}
    " Set to true if using qc based threshold"
    use_qc_thr::Bool
end

prognostic_vars(::RemovePrecipitation) = (Mass(), Energy(), TotalMoisture())

"""
    PrecipitationFlux <: TendencyDef{Flux{FirstOrder}}

Computes the precipitation flux as a sum of air velocity and terminal velocity
multiplied by the advected variable.
"""
struct PrecipitationFlux <: TendencyDef{Flux{FirstOrder}} end

function remove_precipitation_sources(s::RemovePrecipitation, atmos, args)
    @unpack state, aux = args
    @unpack ts = args.precomputed

    FT = eltype(state)

    q = PhasePartition(ts)
    λ::FT = liquid_fraction(ts)
    I_l::FT = internal_energy_liquid(ts)
    I_i::FT = internal_energy_ice(ts)
    Φ::FT = gravitational_potential(atmos.orientation, aux)

    S_qt::FT = 0
    param_set = parameter_set(atmos)
    if s.use_qc_thr
        S_qt = remove_precipitation(param_set, q)
    else
        q_vap_sat = q_vap_saturation(ts)
        S_qt = remove_precipitation(param_set, q, q_vap_sat)
    end

    S_e::FT = (λ * I_l + (1 - λ) * I_i + Φ) * S_qt

    return (; S_ρ_qt = state.ρ * S_qt, S_ρ_e = state.ρ * S_e)
end

"""
    WarmRain_1M <: TendencyDef{Source}

A collection of source/sink terms related to 1-moment warm rain microphysics.
The microphysics process rates are implemented
in the CloudMicrophysics.jl Microphysics_1M module.
"""
struct WarmRain_1M <: TendencyDef{Source} end

prognostic_vars(::WarmRain_1M) =
    (Mass(), Energy(), TotalMoisture(), LiquidMoisture(), Rain())

function warm_rain_sources(atmos, args, ts)
    @unpack state, aux = args

    FT = eltype(state)

    q = PhasePartition(ts)
    T::FT = air_temperature(ts)
    I_l::FT = internal_energy_liquid(ts)
    q_rai::FT = state.precipitation.ρq_rai / state.ρ
    Φ::FT = gravitational_potential(atmos.orientation, aux)
    param_set = parameter_set(atmos)

    # autoconversion
    src_q_rai_acnv = conv_q_liq_to_q_rai(param_set, q.liq)
    # accretion
    src_q_rai_accr = accretion(
        param_set,
        CM1M.LiquidType(),
        CM1M.RainType(),
        q.liq,
        q_rai,
        state.ρ,
    )
    # rain evaporation
    src_q_rai_evap = evaporation_sublimation(
        param_set,
        CM1M.RainType(),
        q,
        q_rai,
        state.ρ,
        T,
    )

    S_qr::FT = src_q_rai_acnv + src_q_rai_accr + src_q_rai_evap
    S_ql::FT = -src_q_rai_acnv - src_q_rai_accr
    S_qt::FT = -S_qr
    S_e::FT = S_qt * (I_l + Φ)

    return (;
        S_ρ_qt = state.ρ * S_qt,
        S_ρ_ql = state.ρ * S_ql,
        S_ρ_qr = state.ρ * S_qr,
        S_ρ_e = state.ρ * S_e,
    )
end

"""
    RainSnow_1M <: TendencyDef{Source}

A collection of source/sink terms related to 1-moment rain and snow microphysics
The microphysics process rates are implemented
in the CloudMicrophysics.jl Microphysics_1M module.
"""
struct RainSnow_1M <: TendencyDef{Source} end

prognostic_vars(::RainSnow_1M) = (
    Mass(),
    Energy(),
    TotalMoisture(),
    LiquidMoisture(),
    IceMoisture(),
    Rain(),
    Snow(),
)

function rain_snow_sources(atmos, args, ts)
    @unpack state, aux = args

    FT = eltype(state)
    param_set = parameter_set(atmos)

    q_rai::FT = state.precipitation.ρq_rai / state.ρ
    q_sno::FT = state.precipitation.ρq_sno / state.ρ
    q = PhasePartition(ts)
    T::FT = air_temperature(ts)
    I_d::FT = internal_energy_dry(ts)
    I_v::FT = internal_energy_vapor(ts)
    I_l::FT = internal_energy_liquid(ts)
    I_i::FT = internal_energy_ice(ts)
    _T_freeze::FT = T_freeze(param_set)
    _L_f::FT = latent_heat_fusion(ts)
    _cv_l::FT = cv_l(param_set)
    Φ::FT = gravitational_potential(atmos.orientation, aux)

    # temporary vars for summming different source terms
    S_qr::FT = FT(0)
    S_qs::FT = FT(0)
    S_ql::FT = FT(0)
    S_qi::FT = FT(0)
    S_qt::FT = FT(0)
    S_e::FT = FT(0)

    # source of rain via autoconversion
    tmp = conv_q_liq_to_q_rai(param_set, q.liq)
    S_qr += tmp
    S_ql -= tmp
    S_e -= tmp * (I_l + Φ)

    # source of snow via autoconversion
    tmp = conv_q_ice_to_q_sno(param_set, q, state.ρ, T)
    S_qs += tmp
    S_qi -= tmp
    S_e -= tmp * (I_i + Φ)

    # source of rain water via accretion cloud water - rain
    tmp = accretion(
        param_set,
        CM1M.LiquidType(),
        CM1M.RainType(),
        q.liq,
        q_rai,
        state.ρ,
    )
    S_qr += tmp
    S_ql -= tmp
    S_e -= tmp * (I_l + Φ)

    # source of snow via accretion cloud ice - snow
    tmp = accretion(
        param_set,
        CM1M.IceType(),
        CM1M.SnowType(),
        q.ice,
        q_sno,
        state.ρ,
    )
    S_qs += tmp
    S_qi -= tmp
    S_e -= tmp * (I_i + Φ)

    # sink of cloud water via accretion cloud water - snow
    tmp = accretion(
        param_set,
        CM1M.LiquidType(),
        CM1M.SnowType(),
        q.liq,
        q_sno,
        state.ρ,
    )
    if T < _T_freeze # cloud droplets freeze to become snow)
        S_qs += tmp
        S_ql -= tmp
        S_e -= tmp * (I_i + Φ)
    else # snow melts, both cloud water and snow become rain
        α::FT = _cv_l / _L_f * (T - _T_freeze)
        S_ql -= tmp
        S_qs -= tmp * α
        S_qr += tmp * (1 + α)
        S_e -= tmp * ((1 + α) * I_l - α * I_i + Φ)
    end

    # sink of cloud ice via accretion cloud ice - rain
    tmp1 = accretion(
        param_set,
        CM1M.IceType(),
        CM1M.RainType(),
        q.ice,
        q_rai,
        state.ρ,
    )
    # sink of rain via accretion cloud ice - rain
    tmp2 = accretion_rain_sink(param_set, q.ice, q_rai, state.ρ)
    S_qi -= tmp1
    S_e -= tmp1 * (I_i + Φ)
    S_qr -= tmp2
    S_e += tmp2 * _L_f
    S_qs += tmp1 + tmp2

    # accretion rain - snow
    if T < _T_freeze
        tmp = accretion_snow_rain(
            param_set,
            CM1M.SnowType(),
            CM1M.RainType(),
            q_sno,
            q_rai,
            state.ρ,
        )
        S_qs += tmp
        S_qr -= tmp
        S_e += tmp * _L_f
    else
        tmp = accretion_snow_rain(
            param_set,
            CM1M.RainType(),
            CM1M.SnowType(),
            q_rai,
            q_sno,
            state.ρ,
        )
        S_qs -= tmp
        S_qr += tmp
        S_e -= tmp * _L_f
    end

    # rain evaporation sink (it already has negative sign for evaporation)
    tmp = evaporation_sublimation(
        param_set,
        CM1M.RainType(),
        q,
        q_rai,
        state.ρ,
        T,
    )
    S_qr += tmp
    S_e -= tmp * (I_l + Φ)

    # snow sublimation/deposition source/sink
    tmp = evaporation_sublimation(
        param_set,
        CM1M.SnowType(),
        q,
        q_sno,
        state.ρ,
        T,
    )
    S_qs += tmp
    S_e -= tmp * (I_i + Φ)

    # snow melt
    tmp = snow_melt(param_set, q_sno, state.ρ, T)
    S_qs -= tmp
    S_qr += tmp
    S_e -= tmp * _L_f

    # total qt sink is the sum of precip sources
    S_qt = -S_qr - S_qs

    return (;
        S_ρ_qt = state.ρ * S_qt,
        S_ρ_ql = state.ρ * S_ql,
        S_ρ_qi = state.ρ * S_qi,
        S_ρ_qr = state.ρ * S_qr,
        S_ρ_qs = state.ρ * S_qs,
        S_ρ_e = state.ρ * S_e,
    )
end


# random functions that dispatch either to CURAND or GPUArrays' generic RNG

using Random

export rand_logn!, rand_poisson!

# native RNG

"""
    CUDA.RNG()

A random number generator using `rand()` in a device kernel.

See also: [CUDA.Philox2x32](@ref)
"""
mutable struct RNG <: AbstractRNG
    seed::UInt32
    counter::UInt32

    function RNG(seed::Integer)
        new(seed%UInt32, 0)
    end
end

RNG() = RNG(Random.rand(UInt32))

function Random.seed!(rng::RNG, seed::Integer)
    rng.seed = seed % UInt32
    rng.counter = 0
end

Random.seed!(rng::RNG) = seed!(rng, Random.rand(UInt32))

function Random.rand!(rng::RNG, A::AnyCuArray)
    function kernel(A::AbstractArray{T}, seed::UInt32, counter::UInt32) where {T}
        device_rng = Random.default_rng()

        # initialize the state
        @inbounds Random.seed!(device_rng, seed, counter)

        # grid-stride loop
        threadId = threadIdx().x
        window = (blockDim().x - 1) * gridDim().x
        offset = (blockIdx().x - 1) * blockDim().x
        while offset < length(A)
            i = threadId + offset
            if i <= length(A)
                @inbounds A[i] = Random.rand(device_rng, T)
            end

            offset += window
        end

        return
    end

    kernel = @cuda launch=false name="rand!" kernel(A, rng.seed, rng.counter)
    config = launch_configuration(kernel.fun; max_threads=64)
    threads = max(32, min(config.threads, length(A)))
    blocks = min(config.blocks, cld(length(A), threads))
    kernel(A, rng.seed, rng.counter; threads=threads, blocks=blocks)

    new_counter = Int64(rng.counter) + length(A)
    overflow, remainder = fldmod(new_counter, typemax(UInt32))
    rng.seed += overflow     # XXX: is this OK?
    rng.counter = remainder

    A
end

function Random.randn!(rng::RNG, A::AnyCuArray{<:T}) where {T<:Union{AbstractFloat,Complex{<:AbstractFloat}}}
    function kernel(A::AbstractArray{T}, seed::UInt32, counter::UInt32)
        device_rng = Random.default_rng()

        # initialize the state
        @inbounds Random.seed!(device_rng, seed, counter)

        # grid-stride loop
        threadId = threadIdx().x
        window = (blockDim().x - 1) * gridDim().x
        offset = (blockIdx().x - 1) * blockDim().x
        while offset < length(A)
            i = threadId + offset
            j = threadId + offset + window
            if i <= length(A)
                # Box–Muller transform
                U1 = Random.rand(device_rng, T)
                U2 = Random.rand(device_rng, T)
                Z0 = sqrt(T(-2.0)*log(U1))*cos(T(2pi)*U2)
                Z1 = sqrt(T(-2.0)*log(U1))*sin(T(2pi)*U2)
                @inbounds A[i] = Z0
                if j <= length(A)
                    @inbounds A[j] = Z1
                end
            end

            offset += 2*window
        end

        return
    end

    kernel = @cuda launch=false name="rand!" kernel(A, rng.seed, rng.counter)
    config = launch_configuration(kernel.fun; max_threads=64)
    threads = max(32, min(config.threads, length(A)÷2))
    blocks = min(config.blocks, cld(length(A)÷2, threads))
    kernel(A, rng.seed, rng.counter; threads=threads, blocks=blocks)

    new_counter = Int64(rng.counter) + length(A)
    overflow, remainder = fldmod(new_counter, typemax(UInt32))
    rng.seed += overflow     # XXX: is this OK?
    rng.counter = remainder

    A
end

function default_rng()
    cuda = CUDA.active_state()

    # every task maintains library state per device
    LibraryState = @NamedTuple{rng::RNG}
    states = get!(task_local_storage(), :RNG) do
        Dict{CuContext,LibraryState}()
    end::Dict{CuContext,LibraryState}

    # get library state
    @noinline function new_state(cuda)
        # CUDA RNG objects are cheap, so we don't need to cache them
        (; rng=RNG())
    end
    state = get!(states, cuda.context) do
        new_state(cuda)
    end

    return state.rng
end

# old native RNG

# we keep this for the GPUArrays.jl tests

const idle_gpuarray_rngs = HandleCache{CuContext,GPUArrays.RNG}()

function GPUArrays.default_rng(::Type{<:CuArray})
    cuda = CUDA.active_state()

    # every task maintains library state per device
    LibraryState = @NamedTuple{rng::GPUArrays.RNG}
    states = get!(task_local_storage(), :GPUArraysRNG) do
        Dict{CuContext,LibraryState}()
    end::Dict{CuContext,LibraryState}

    # get library state
    @noinline function new_state(cuda)
        new_rng = pop!(idle_gpuarray_rngs, cuda.context) do
            N = attribute(cuda.device, DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK)
            buf = CuArray{NTuple{4, UInt32}}(undef, N)
            GPUArrays.RNG(buf)
        end

        finalizer(current_task()) do task
            push!(idle_gpuarray_rngs, cuda.context, new_rng) do
                # no need to do anything, as the RNG is collected by its finalizer
            end
        end

        Random.seed!(new_rng)
        (; rng=new_rng)
    end
    state = get!(states, cuda.context) do
        new_state(cuda)
    end

    return state.rng
end

# generic functionality

function Random.rand!(rng::Union{RNG,CURAND.RNG,GPUArrays.RNG}, A::AbstractArray{T}) where {T}
    B = CuArray{T}(undef, size(A))
    Random.rand!(rng, B)
    copyto!(A, B)
end

function Random.rand(rng::Union{RNG,CURAND.RNG,GPUArrays.RNG}, T::Type)
    assertscalar("scalar rand")
    A = CuArray{T}(undef, 1)
    Random.rand!(rng, A)
    A[]
end

# RNG-less interface

cuda_rng() = default_rng()
curand_rng() = CURAND.default_rng()

function seed!(seed=Base.rand(UInt64))
    Random.seed!(cuda_rng(), seed)
    Random.seed!(curand_rng(), seed)
end

# CURAND in-place
Random.rand!(A::CURAND.UniformArray) = Random.rand!(curand_rng(), A)
Random.randn!(A::CURAND.NormalArray; kwargs...) = Random.randn!(curand_rng(), A; kwargs...)
rand_logn!(A::CURAND.LognormalArray; kwargs...) = CURAND.rand_logn!(curand_rng(), A; kwargs...)
rand_poisson!(A::CURAND.PoissonArray; kwargs...) = CURAND.rand_poisson!(curand_rng(), A; kwargs...)

# CURAND out-of-place
rand(T::CURAND.UniformType, dims::Dims) = Random.rand(curand_rng(), T, dims)
randn(T::CURAND.NormalType, dims::Dims; kwargs...) = Random.randn(curand_rng(), T, dims; kwargs...)
rand_logn(T::CURAND.LognormalType, dims::Dims; kwargs...) = CURAND.rand_logn(curand_rng(), T, dims; kwargs...)
rand_poisson(T::CURAND.PoissonType, dims::Dims; kwargs...) = CURAND.rand_poisson(curand_rng(), T, dims; kwargs...)

# support all dimension specifications
rand(T::CURAND.UniformType, dim1::Integer, dims::Integer...) =
    Random.rand(curand_rng(), T, Dims((dim1, dims...)))
randn(T::CURAND.NormalType, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(curand_rng(), T, Dims((dim1, dims...)); kwargs...)
rand_logn(T::CURAND.LognormalType, dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_logn(curand_rng(), T, Dims((dim1, dims...)); kwargs...)
rand_poisson(T::CURAND.PoissonType, dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_poisson(curand_rng(), T, Dims((dim1, dims...)); kwargs...)

# native in-place
Random.rand!(A::AnyCuArray) = Random.rand!(cuda_rng(), A)
Random.randn!(A::AnyCuArray) = Random.randn!(cuda_rng(), A)
rand_logn!(A::AnyCuArray; kwargs...) =
    error("CUDA.jl does not support generating lognormally-distributed random numbers of type $(eltype(A))")
rand_poisson!(A::AnyCuArray; kwargs...) =
    error("CUDA.jl does not support generating Poisson-distributed random numbers of type $(eltype(A))")

# native out-of-place
rand(T::Type, dims::Dims) = Random.rand!(CuArray{T}(undef, dims...))
randn(T::Type, dims::Dims; kwargs...) = Random.randn!(CuArray{T}(undef, dims...); kwargs...)
rand_logn(T::Type, dims::Dims; kwargs...) = rand_logn!(CuArray{T}(undef, dims...); kwargs...)
rand_poisson(T::Type, dims::Dims; kwargs...) = rand_poisson!(CuArray{T}(undef, dims...); kwargs...)

# support all dimension specifications
rand(T::Type, dim1::Integer, dims::Integer...) =
    Random.rand!(CuArray{T}(undef, dim1, dims...))
randn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn!(CuArray{T}(undef, dim1, dims...); kwargs...)
rand_logn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_logn!(CuArray{T}(undef, dim1, dims...); kwargs...)
rand_poisson(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_poisson!(CuArray{T}(undef, dim1, dims...); kwargs...)

# untyped out-of-place
rand(dim1::Integer, dims::Integer...) =
    Random.rand(curand_rng(), Dims((dim1, dims...)))
randn(dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(curand_rng(), Dims((dim1, dims...)); kwargs...)
rand_logn(dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_logn(curand_rng(), Dims((dim1, dims...)); kwargs...)
rand_poisson(dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_poisson(curand_rng(), Dims((dim1, dims...)); kwargs...)


using FixedPointNumbers

@testset "Parse" begin
    r8(x) = reinterpret(N0f8, x)

    # Color parsing
    # named-color
    redN0f8 = parse(Colorant, "red")
    @test colorant"red" === redN0f8
    @test redN0f8 === RGB{N0f8}(1,0,0)
    @test parse(RGB{Float64}, "red") === RGB{Float64}(1,0,0)
    @test isa(parse(HSV, "blue"), HSV)
    @test_throws ErrorException parse(Colorant, "p ink")
    @test parse(Colorant, "transparent") === RGBA{N0f8}(0,0,0,0)
    @test parse(Colorant, "\nSeaGreen ") === RGB{N0f8}(r8(0x2E),r8(0x8B),r8(0x57))
    @test parse(Colorant, "sea GREEN") === colorant"seagreen"

    # hex-color
    @test parse(Colorant, "#D0FF58") === RGB(r8(0xD0),r8(0xFF),r8(0x58))
    @test parse(Colorant, "0xd0ff58") === RGB(r8(0xD0),r8(0xFF),r8(0x58))
    @test parse(Colorant, "#FB0") === RGB(r8(0xFF),r8(0xBB),r8(0x00))
    @test parse(Colorant, "#FB0A") === RGBA(r8(0xFF),r8(0xBB),r8(0x00),r8(0xAA))
    @test parse(Colorant, "0xFB0A") === ARGB(r8(0xBB),r8(0x00),r8(0xAA),r8(0xFF))
    @test parse(Colorant, "#FFBB00AA") === RGBA(r8(0xFF),r8(0xBB),r8(0x00),r8(0xAA))
    @test parse(Colorant, "0xFFBB00AA") === ARGB(r8(0xBB),r8(0x00),r8(0xAA),r8(0xFF))
    @test_throws ErrorException parse(Colorant, "#BAD05")
    @test_throws ErrorException parse(Colorant, "#BAD0007")
    @test_throws ErrorException parse(Colorant, "#BAD000009")

    # rgb()
    @test parse(Colorant, "rgb(55,217,127)")      === RGB{N0f8}(r8(0x37),r8(0xd9),r8(0x7f))
    @test colorant"    rgb(    55, 217, 127 )   " === RGB{N0f8}(r8(0x37),r8(0xd9),r8(0x7f))
    @test parse(Colorant, "rgb(22%,85%,50%)")     === RGB{N0f8}(r8(0x38),r8(0xd9),r8(0x80))
    @test parse(Colorant, "rgba(55,217,127,0.5)") === RGBA{N0f8}(r8(0x37),r8(0xd9),r8(0x7f),0.5)
    @test parse(Colorant, "rgb( 55,217,127,50%)") === RGBA{N0f8}(r8(0x37),r8(0xd9),r8(0x7f),0.5) # CSS Color Module Level 4
    @test parse(Colorant, "rgb( 55 217 127 /.5)") === RGBA{N0f8}(r8(0x37),r8(0xd9),r8(0x7f),0.5) # CSS Color Module Level 4
    @test parse(Colorant, "rgb(55, 85%, 50%)")    === RGB{N0f8}(r8(0x37),r8(0xd9),r8(0x80)) # this is invalid according to CSS spec.
    @test_throws ErrorException parse(Colorant, "rgb(21.6%,85%,50%)") # this is valid but not supported

    # hsl()
    @test parse(Colorant, "hsl(120, 100%, 50%)")  === HSL{Float32}(120,1.0,.5)
    @test colorant"    hsl(    120, 100%, 50% ) " === HSL{Float32}(120,1.0,.5)
    @test parse(RGB{N0f8},"hsl(120, 100%, 50%)")  === convert(RGB{N0f8}, HSL{Float32}(120,1.0,.5))
    @test_throws ErrorException  parse(Colorant, "hsl(120, 100, 50)")
    @test_throws ErrorException  parse(Colorant, "hsl(120%,100%,50%)")
    @test parse(Colorant, "hsla(120,50%,7%, .6)") === HSLA{Float32}(120,.5,.07,.6)
    @test parse(Colorant, "hsl( 120,50%,7%,60%)") === HSLA{Float32}(120,.5,.07,.6) # CSS Color Module Level 4
    @test parse(Colorant, "hsl( 120 50% 7% / 1)") === HSLA{Float32}(120,.5,.07, 1) # CSS Color Module Level 4

    @test parse(Colorant, :red) === colorant"red"
    @test_deprecated parse(Colorant, colorant"red") === colorant"red"

end


using Fed
using Flux
using SHA

struct Server
    model::Chain
    dataset::DataLoader

    Server() = new(
        newmodel(),
        get_testdata(isflat=true)
    )
end

function start_server()
    server = Server()
    config = newconfig(9090, false)

    # config
    host = "127.0.0.1"
    port = 8080
    weights = flatten_model(server.model)
    strategy = Fed.Server.federated_averaging
    eval_hook = (weights::Vector{Float32}) -> evaluate(server, weights)

    central_node = Fed.Server.CentralNode(
        host,
        port,
        weights,
        strategy,
        eval_hook,
        config
    )

    # start the server
    @info "Server started on [http://$host:$port]"
    Fed.Server.start(central_node)
end


function write_nl_file(f::IO, m::AmplNLMathProgModel)
    write_nl_header(f, m)

    if m.ncon > 0
        write_nl_c_blocks(f, m)
    end

    if has_objective(m)
        write_nl_o_block(f, m)
    end

    if m.ncon > 0
        write_nl_d_block(f, m)
    end

    write_nl_x_block(f, m)

    if m.ncon > 0
        write_nl_r_block(f, m)
    end

    write_nl_b_block(f, m)

    if m.ncon > 0
        write_nl_k_block(f, m)
        write_nl_j_blocks(f, m)
    end

    if has_objective(m)
        write_nl_g_block(f, m)
    end
end

function has_objective(m::AmplNLMathProgModel)
    return !isempty(m.lin_obj) || (isa(m.obj, Expr))
end

function write_nl_header(f, m::AmplNLMathProgModel)
    # Line 1: Always the same
    println(f, "g3 1 1 0")
    # Line 2: vars, constraints, objectives, ranges, eqns, logical constraints
    n_ranges = sum(m.r_codes .== 0)
    n_eqns = sum(m.r_codes .== 4)
    nobj = has_objective(m) ? 1 : 0
    println(f, " $(m.nvar) $(m.ncon) $nobj $n_ranges $n_eqns 0")
    # Line 3: nonlinear constraints, objectives
    nlc = sum(m.conlinearities .== :Nonlin)
    nlo = m.objlinearity == :Nonlin ? 1 : 0
    println(f, " $nlc $nlo")
    # Line 4: network constraints: nonlinear, linear
    println(f, " 0 0")
    # Line 5: nonlinear vars in constraints, objectives, both
    nonlinear_obj = m.varlinearities_obj .== :Nonlin
    nonlinear_con = m.varlinearities_con .== :Nonlin
    nonlinear = (nonlinear_con + nonlinear_obj) .> 0
    nonlinear_both = (nonlinear_con + nonlinear_obj) .> 1
    nlvc = sum(nonlinear_con .> 0)
    nlvo = sum(nonlinear_obj .> 0)
    nlvb = sum(nonlinear_both)
    println(f, " $nlvc $nlvo $nlvb")
    # Line 6: linear network variables; functions; arith, flags
    println(f, " 0 0 0 1")  # flags set to 1 to get suffixes in .sol file
    # Line 7: discrete variables: binary, integer, nonlinear (b,c,o)
    binary = m.vartypes .== :Bin
    integer = m.vartypes .== :Int
    discrete = binary + integer .> 0
    # Julia 0.6 syntax
    nbv = sum(binary + .!nonlinear .> 1)
    niv = sum(integer + .!nonlinear .> 1)
    nlvbi = sum(nonlinear_both + discrete .> 1)
    nlvci = sum(nonlinear_con - nonlinear_obj + discrete .> 1)
    nlvoi = sum(nonlinear_obj - nonlinear_con + discrete .> 1)
    println(f, " $nbv $niv $nlvbi $nlvci $nlvoi")
    # Line 8: nonzeros in Jacobian, gradients
    nzc = sum(m.j_counts)
    nzo = length(m.lin_obj)
    println(f, " $nzc $nzo")
    # Line 9: max name lengths: constraints, variables
    println(f, " 0 0")
    # Line 10: common exprs: b,c,o,c1,o1
    println(f, " 0 0 0 0 0")
end

# Nonlinear constraint trees
function write_nl_c_blocks(f, m::AmplNLMathProgModel)
    for index in 0:(m.ncon - 1)
        i = m.c_index_map_rev[index]
        println(f, "C$index")
        write_nl_expr(f, m, m.constrs[i])
    end
end

# Nonlinear objective tree
function write_nl_o_block(f, m::AmplNLMathProgModel)
    println(f, string("O0 ", sense_to_nl[m.sense]))
    write_nl_expr(f, m, m.obj)
end

# Initial dual guesses - unused
function write_nl_d_block(f, m::AmplNLMathProgModel)
    println(f, "d$(m.ncon)")
    for index in 0:(m.ncon - 1)
        i = m.c_index_map_rev[index]
        println(f, "$index 0")
    end
end

# Initial primal guesses
function write_nl_x_block(f, m::AmplNLMathProgModel)
    println(f, "x$(m.nvar)")
    for index in 0:(m.nvar - 1)
        i = m.v_index_map_rev[index]
        println(f, "$index $(m.x_0[i])")
    end
end

# Constraint bounds
function write_nl_r_block(f, m::AmplNLMathProgModel)
    println(f, "r")
    for index in 0:(m.ncon - 1)
        i = m.c_index_map_rev[index]
        lower = m.g_l[i]
        upper = m.g_u[i]
        rel = m.r_codes[i]
        if rel == 0
            println(f, "$rel $lower $upper")
        elseif rel == 1
            println(f, "$rel $upper")
        elseif rel == 2
            println(f, "$rel $lower")
        elseif rel == 3
            println(f, "$rel")
        elseif rel == 4
            println(f, "$rel $lower")
        end
    end
end

# Variable bounds
function write_nl_b_block(f, m::AmplNLMathProgModel)
    println(f, "b")
    for index in 0:(m.nvar - 1)
        i = m.v_index_map_rev[index]
        lower = m.x_l[i]
        upper = m.x_u[i]
        if lower == -Inf
            if upper == Inf
                println(f, "3")
            else
                println(f, "1 $upper")
            end
        else
            if lower == upper
                println(f, "4 $lower")
            elseif upper == Inf
                println(f, "2 $lower")
            else
                println(f, "0 $lower $upper")
            end
        end
    end
end

# Jacobian counts
function write_nl_k_block(f, m::AmplNLMathProgModel)
    println(f, "k$(m.nvar - 1)")
    total = 0
    for index = 0:(m.nvar - 2)
        i = m.v_index_map_rev[index]
        total += m.j_counts[i]
        println(f, total)
    end
end

# Linear constraint expressions
function write_nl_j_blocks(f, m::AmplNLMathProgModel)
    for index in 0:(m.ncon - 1)
        i = m.c_index_map_rev[index]
        num_vars = length(m.lin_constrs[i])
        # Only print linear blocks with vars, .nl file is malformed otherwise
        if num_vars > 0
            println(f, "J$index $num_vars")

            # We need to output .nl index and constraint coeff, ordered by .nl
            # index. `lin_constrs` contains our variable index as the key and
            # constraint coeff as the value

            # Assemble tuples of (.nl index, constraint value)
            output = collect(zip(
                (m.v_index_map[j] for j in keys(m.lin_constrs[i])),
                values(m.lin_constrs[i]))
            )

            # Loop through output in .nl index order
            for (index2, value) in sort(output)
                println(f, "$index2 $value")
            end
        end
    end
end

# Linear objective expression
function write_nl_g_block(f, m::AmplNLMathProgModel)
    println(f, string("G0 ", length(m.lin_obj)))
    for index in 0:(m.nvar - 1)
        i = m.v_index_map_rev[index]
        if i in keys(m.lin_obj)
            println(f, "$index $(m.lin_obj[i])")
        end
    end
end

# Convert an expression tree (with .nl formulae only) to .nl format
write_nl_expr(f, m, c) = println(f, string(c))
# Handle numerical constants e.g. pi
write_nl_expr(f, m, c::Symbol) =  write_nl_expr(f, m, float(eval(c)))
function write_nl_expr(f, m, c::Real)
    println(f, nl_number(c == round(Integer, c) ? round(Integer, c) : c))
end
write_nl_expr(f, m, c::LinearityExpr) = write_nl_expr(f, m, c.c)
function write_nl_expr(f, m, c::Expr)
    if c.head == :ref
        # Output variable as `v$index`
        if c.args[1] == :x
            @assert isa(c.args[2], Int)
            println(f, nl_variable(m.v_index_map[c.args[2]]))
        else
            error("Unrecognized reference expression $c")
        end
    elseif c.head == :call
        # Output function as `o$opcode`
        println(f, nl_operator(c.args[1]))
        if c.args[1] in nary_functions
            # Output nargs on subsequent line if n-ary function
            println(f, (string(length(c.args) - 1)))
        end
        map(arg -> write_nl_expr(f, m, arg), c.args[2:end])

    elseif c.head == :comparison
        # .nl only handles binary comparison
        @assert length(c.args) == 3
        # Output comparison type first, followed by args
        println(f, nl_operator(c.args[2]))
        map(arg -> write_nl_expr(f, m, arg), c.args[1:2:end])

    elseif c.head in [:&&, :||]
        # Only support binary and/or for now
        @assert length(c.args) == 2
        println(f, nl_operator(c.head))
        map(arg -> write_nl_expr(f, m, arg), c.args)
    else
        error("Unrecognized expression $c")
    end
end

nl_variable(index::Integer) = "v$index"
nl_number(value::Real) = "n$value"

function nl_operator(operator::Symbol)
    if !haskey(func_to_nl, operator)
        error("translation of the function \"$operator\" to NL is not defined")
    end
    return "o$(func_to_nl[operator])"
end


using Test, MadNLP, MadNLPHSL, MadNLPTests

testset = [
    [
        "HSL-Ma27",
        ()->MadNLP.Optimizer(
            linear_solver=MadNLPMa27,
            print_level=MadNLP.ERROR),
        []
    ],
    [
        "HSL-Ma57",
        ()->MadNLP.Optimizer(
            linear_solver=MadNLPMa57,
            print_level=MadNLP.ERROR),
        []
    ],
    [
        "HSL-Ma77",
        ()->MadNLP.Optimizer(
            linear_solver=MadNLPMa77,
            print_level=MadNLP.ERROR),
        ["unbounded"]
    ],
    [
        "HSL-Ma86",
        ()->MadNLP.Optimizer(
            linear_solver=MadNLPMa86,
            print_level=MadNLP.ERROR),
        []
    ],
    [
        "HSL-Ma97",
        ()->MadNLP.Optimizer(
            linear_solver=MadNLPMa97,
            print_level=MadNLP.ERROR),
        []
    ],
]

@testset "MadNLPHSL test" begin
    for (name,optimizer_constructor,exclude) in testset
        test_madnlp(name,optimizer_constructor,exclude)
    end
end


# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# McCormick.jl
# A McCormick operator library in Julia
# See https://github.com/PSORLab/McCormick.jl
#############################################################################
# src/forward_operators/arithmetic.jl
# Contains definitions of +, -, /, *, promotions, conversion, one, zero.
#############################################################################

# Defines functions required for linear algebra packages
@inline nan(::Type{MC{N,T}}) where {N, T <: RelaxTag} = MC{N,T}(NaN, NaN, Interval{Float64}(NaN),
                                                                fill(NaN, SVector{N,Float64}),
																fill(NaN, SVector{N,Float64}), true)
@inline nan(x::MC{N,T}) where {N, T <: RelaxTag} = MC{N,T}(NaN, NaN, Interval{Float64}(NaN),
                                                           fill(NaN, SVector{N,Float64}),
														   fill(NaN, SVector{N,Float64}), true)

@inline one(::Type{MC{N,T}}) where {N, T <: RelaxTag} = MC{N,T}(1.0, 1.0, one(Interval{Float64}),
                                                                zero(SVector{N,Float64}),
																zero(SVector{N,Float64}), true)
@inline one(x::MC{N,T}) where {N, T <: RelaxTag} = MC{N,T}(1.0, 1.0, one(Interval{Float64}),
                                                           zero(SVector{N,Float64}),
														   zero(SVector{N,Float64}), true)

@inline zero(::Type{MC{N,T}}) where {N, T <: RelaxTag} = MC{N,T}(0.0, 0.0, zero(Interval{Float64}),
                                                                 zero(SVector{N,Float64}),
																 zero(SVector{N,Float64}), true)
@inline zero(x::MC{N,T}) where {N, T <: RelaxTag} = zero(MC{N,T})

@inline real(x::MC) = x
@inline dist(x1::MC, x2::MC) = max(abs(x1.cc - x2.cc), abs(x1.cv - x2.cv))
@inline eps(x::MC) = max(eps(x.cc), eps(x.cv))
@inline mid(x::MC) = mid(x.Intv)

# Unsafe addition
@inline function plus_kernel(x::MC{N,T}, y::MC{N,T}, z::Interval{Float64}) where {N, T <: Union{NS, MV, Diff}}
	MC{N,T}(x.cv + y.cv, x.cc + y.cc, z, x.cv_grad + y.cv_grad, x.cc_grad + y.cc_grad, (x.cnst && y.cnst))
end
@inline +(x::MC, y::MC) = plus_kernel(x, y, x.Intv + y.Intv)
@inline plus_kernel(x::MC, y::Interval{Float64}) = x
@inline +(x::MC) = x

@inline minus_kernel(x::MC{N,T}, z::Interval{Float64}) where {N, T <: RelaxTag} = MC{N,T}(-x.cc, -x.cv, z, -x.cc_grad, -x.cv_grad, x.cnst)
@inline -(x::MC) = minus_kernel(x, -x.Intv)
@inline -(x::MC{N,T}, y::MC{N,T}) where {N, T <: RelaxTag} = minus_kernel(x, y, x.Intv - y.Intv)

# Unsafe subtraction
@inline function minus_kernel(x::MC{N,T}, y::MC{N,T}, z::Interval{Float64}) where {N, T <: Union{NS, MV, Diff}}
	MC{N,T}(x.cv - y.cc, x.cc - y.cv, z, x.cv_grad - y.cc_grad, x.cc_grad - y.cv_grad, (x.cnst && y.cnst))
end

################## CONVERT THROUGH BINARY DEFINITIONS #########################
# Unsafe scalar addition
@inline function plus_kernel(x::MC{N,T}, y::Float64, z::Interval{Float64}) where {N, T <: Union{NS, MV, Diff}}
	MC{N,T}(x.cv + y, x.cc + y, z, x.cv_grad, x.cc_grad, x.cnst)
end
@inline +(x::MC, y::Float64) = plus_kernel(x, y, x.Intv + y)
@inline +(y::Float64, x::MC) = plus_kernel(x, y, x.Intv + y)
@inline +(x::MC{N,T}, y::Interval{Float64}) where {N, T<:RelaxTag} = x + MC{N,T}(y)
@inline +(y::Interval{Float64}, x::MC{N,T}) where {N, T<:RelaxTag} = x + MC{N,T}(y)

@inline plus_kernel(x::MC, y::C, z::Interval{Float64}) where {C <: NumberNotRelax} = plus_kernel(x, convert(Float64, y), z)
@inline plus_kernel(x::C, y::MC, z::Interval{Float64}) where {C <: NumberNotRelax} = plus_kernel(y, convert(Float64, x), z)
@inline +(x::MC, y::C) where {C <: NumberNotRelax} = x + convert(Float64, y)
@inline +(y::C, x::MC) where {C <: NumberNotRelax} = x + convert(Float64, y)

# Unsafe scalar subtraction
@inline function minus_kernel(x::MC{N,T}, c::Float64, z::Interval{Float64}) where {N, T <: Union{NS, MV, Diff}}
	MC{N,T}(x.cv - c, x.cc - c, z, x.cv_grad, x.cc_grad, x.cnst)
end
@inline function minus_kernel(c::Float64, x::MC{N,T}, z::Interval{Float64}) where {N, T <: Union{NS, MV, Diff}}
	MC{N,T}(c - x.cc, c - x.cv, z, -x.cc_grad, -x.cv_grad, x.cnst)
end
@inline -(x::MC, c::Float64) = minus_kernel(x, c, x.Intv - c)
@inline -(c::Float64, x::MC) = minus_kernel(c, x, c - x.Intv)
@inline -(x::MC{N,T}, y::Interval{Float64}) where {N, T<:RelaxTag} = x - MC{N,T}(y)
@inline -(y::Interval{Float64}, x::MC{N,T}) where {N, T<:RelaxTag} = MC{N,T}(y) - x

@inline minus_kernel(x::MC, y::C, z::Interval{Float64}) where {C <: NumberNotRelax} = minus_kernel(x, convert(Float64, y), z)
@inline minus_kernel(y::C, x::MC, z::Interval{Float64}) where {C <: NumberNotRelax} = minus_kernel(convert(Float64, y), x, z)
@inline -(x::MC, c::C) where {C <: NumberNotRelax} = x - convert(Float64,c)
@inline -(c::C, x::MC) where {C <: NumberNotRelax} = convert(Float64,c) - x

# Unsafe Scalar Multiplication
@inline function mult_kernel(x::MC{N,T}, c::Float64, z::Interval{Float64}) where {N, T <: Union{NS, MV, Diff}}
	if c >= 0.0
		zMC = MC{N,T}(c*x.cv, c*x.cc, z, c*x.cv_grad, c*x.cc_grad, x.cnst)
	else
		zMC = MC{N,T}(c*x.cc, c*x.cv, z, c*x.cc_grad, c*x.cv_grad, x.cnst)
	end
	return zMC
end
@inline *(x::MC, c::Float64) = mult_kernel(x, c, c*x.Intv)
@inline *(c::Float64, x::MC) = mult_kernel(x, c, c*x.Intv)
@inline *(x::MC{N,T}, y::Interval{Float64}) where {N, T<:RelaxTag} = x*MC{N,T}(y)
@inline *(y::Interval{Float64}, x::MC{N,T}) where {N, T<:RelaxTag} = MC{N,T}(y)*x

@inline mult_kernel(x::MC, c::C, z::Interval{Float64}) where {C <: NumberNotRelax} = mult_kernel(x, convert(Float64, c), z)
@inline mult_kernel(c::C, x::MC, z::Interval{Float64}) where {C <: NumberNotRelax} =  mult_kernel(x, convert(Float64, c), z)
@inline *(c::C, x::MC) where {C <: NumberNotRelax}  = x*Float64(c)
@inline *(x::MC, c::C) where {C <: NumberNotRelax}  = x*Float64(c)

# Unsafe scalar division
@inline div_kernel(x::MC, y::Float64, z::Interval{Float64}) = mult_kernel(x, inv(y), z)
@inline div_kernel(x::Float64, y::MC, z::Interval{Float64}) = mult_kernel(inv(y), x, z)
@inline div_kernel(x::MC, y::C, z::Interval{Float64}) where {C <: NumberNotRelax}  = mult_kernel(x, inv(y), z)
@inline div_kernel(x::C, y::MC, z::Interval{Float64}) where {C <: NumberNotRelax}  = mult_kernel(inv(y), x, z)
@inline /(x::MC, y::Float64) = x*inv(y)
@inline /(x::Float64, y::MC) = x*inv(y)
@inline /(x::MC, y::C) where {C <: NumberNotRelax} = x*inv(convert(Float64,y))
@inline /(x::C, y::MC) where {C <: NumberNotRelax} = convert(Float64,x)*inv(y)
@inline /(x::MC{N,T}, y::Interval{Float64}) where {N, T<:RelaxTag} = x/MC{N,T}(y)
@inline /(y::Interval{Float64}, x::MC{N,T}) where {N, T<:RelaxTag} = MC{N,T}(y)/x

# Maximization
@inline max_kernel(c::Float64, x::MC, z::Interval{Float64}) = max_kernel(x, c, z)
@inline max_kernel(x::MC, c::C, z::Interval{Float64}) where {C <: NumberNotRelax} = max_kernel(x, convert(Float64, c), z)
@inline max_kernel(c::C, x::MC, z::Interval{Float64}) where {C <: NumberNotRelax} = max_kernel(x, convert(Float64, c), z)

@inline max(c::Float64, x::MC) = max_kernel(x, c, max(x.Intv, c))
@inline max(x::MC, c::C) where {C <: NumberNotRelax} = max_kernel(x, convert(Float64, c), max(x.Intv, c))
@inline max(c::C, x::MC) where {C <: NumberNotRelax} = max_kernel(x, convert(Float64, c), max(x.Intv, c))
@inline max(x::MC{N,T}, y::Interval{Float64}) where {N, T<:RelaxTag} = max(x, MC{N,T}(y))
@inline max(y::Interval{Float64}, x::MC{N,T}) where {N, T<:RelaxTag} = max(MC{N,T}(y), x)

# Minimization
@inline min_kernel(x::MC, c::C, z::Interval{Float64}) where {C <: NumberNotRelax} = min_kernel(x, convert(Float64, c), z)
@inline min_kernel(c::C, x::MC, z::Interval{Float64}) where {C <: NumberNotRelax} = min_kernel(x, convert(Float64, c), z)

@inline min(c::Float64, x::MC) = min_kernel(x, c, min(x.Intv, c))
@inline min(x::MC, c::C) where {C <: NumberNotRelax} = min_kernel(x, convert(Float64, c), min(x.Intv, c))
@inline min(c::C, x::MC) where {C <: NumberNotRelax} = min_kernel(x, convert(Float64, c), min(x.Intv, c))
@inline min(x::MC{N,T}, y::Interval{Float64}) where {N, T<:RelaxTag} = min(x, MC{N,T}(y))
@inline min(y::Interval{Float64}, x::MC{N,T}) where {N, T<:RelaxTag} = min(MC{N,T}(y), x)

# Add fma function
@inline fma(x::MC, y::Float64, z::Float64) = x*y + z
@inline fma(x::MC, y::MC, z::Float64) = x*y + z
@inline fma(x::MC, y::Float64, z::MC) = x*y + z
@inline fma(x::MC, y::MC, z::MC) = x*y + z
@inline fma(x::Float64, y::MC, z::Float64) = x*y + z
@inline fma(x::Float64, y::MC, z::MC) = x*y + z
@inline fma(x::Float64, y::Float64, z::MC) = x*y + z

@inline fma(x::MC, y::Float64, z::Float64, q::Interval{Float64}) = x*y + z
@inline fma(x::MC, y::MC, z::Float64, q::Interval{Float64}) = x*y + z
@inline fma(x::MC, y::Float64, z::MC, q::Interval{Float64}) = x*y + z
@inline fma(x::MC, y::MC, z::MC, q::Interval{Float64}) = x*y + z
@inline fma(x::Float64, y::MC, z::Float64, q::Interval{Float64}) = x*y + z
@inline fma(x::Float64, y::MC, z::MC, q::Interval{Float64}) = x*y + z
@inline fma(x::Float64, y::Float64, z::MC, q::Interval{Float64}) = x*y + z

# Promote and Convert
promote_rule(::Type{MC{N,T}}, ::Type{S}) where {S<:NumberNotRelax, N, T <: RelaxTag} = MC{N,T}
promote_rule(::Type{MC{N,T}}, ::Type{S}) where {S<:Real, N, T <: RelaxTag} = MC{N,T}

convert(::Type{MC{N,T}}, x::S) where {S<:NumberNotRelax, N, T <: RelaxTag} = MC{N,T}(Interval{Float64}(x))
convert(::Type{MC{N,T}}, x::S) where {S<:AbstractInterval, N, T <: RelaxTag} = MC{N,T}(Interval{Float64}(x.lo, x.hi))
Interval(x::MC{N,T}) where {N,T<:RelaxTag} = x.Intv
Interval{Float64}(x::MC{N,T}) where {N,T<:RelaxTag} = x.Intv


using Base.Test
using Base.Threads
using Bio.Seq
using Bio.Intervals
using Libz
using DataFrames
using Gadfly

importall Bio.Intervals

include("pairs.jl")
include("energy.jl")
include("duplex.jl")
include("intervals.jl")
include("io.jl")
include("trie.jl")

function trie_profile( seq, range, num_steps; raw_output="" )
   stepsize = div(length(seq), num_steps)
   steps    = collect(50:stepsize:length(seq))
   build    = Vector{Float64}(length(steps))
   trav     = Vector{Float64}(length(steps))
   bulge    = Vector{Float64}(length(steps))
   mismatch = Vector{Float64}(length(steps))
   bulgemat = Vector{Float64}(length(steps))
   dep      = Vector{Int}(length(steps))
   cnt      = Vector{Int}(length(steps))
   int      = Vector{Int}(length(steps))

   df = DataFrame(length=[], buildtime=[], travtime=[], depth=[])

   if raw_output != ""
      const io = open( raw_output, "w" )
      const stream = ZlibDeflateOutputStream( io )
      const output = true
   else
      const output = false
   end

   for i in range
      for j in 1:length(steps)
         trie, build[j]   = @timed DuplexTrie{DNAAlphabet{2},UInt8}( seq[1:steps[j]], min(i, 5):i )
         cnt[j] = nodecount( trie.fwd )
         dep[j] = i
         res, trav[j]     = @timed traverse( trie, 1:10000 )
         int[j] = length(res)
         res, bulge[j]    = @timed traverse( trie, 1:10000, bulge_max=1 )
         res, mismatch[j] = @timed traverse( trie, 1:10000, mismatch_max=1 )
         res, bulgemat[j] = @timed traverse( trie, 1:10000, bulge_max=1, mismatch_max=1 )
      end
      if output
         for j in 1:length(steps)
            write( stream, string(i) * "\t" )
            write( stream, string(j) * "\t" )
            write( stream, string(steps[j]) * "\t" )
            write( stream, string(cnt[j]) * "\t" )
            write( stream, string(build[j]) * "\t" )
            write( stream, string(trav[j]) * "\t" )
            write( stream, string(bulge[j]) * "\t" )
            write( stream, string(mismatch[j]) * "\t" )
            write( stream, string(bulgemat[j]) * "\t" )
            write( stream, string(int[j]) * "\n" )
         end
      else
         df = join(df, DataFrame(length=steps, buildtime=build, travtime=trav, depth=dep), on=[:length, :buildtime, :travtime, :depth], kind = :outer)
      end
   end
   if !output
      buildplt = plot(df, x=:length, y=:buildtime, color=:depth, Geom.smooth)
      travplt  = plot(df, x=:length, y=:travtime,  color=:depth, Geom.smooth)

      return df,hstack(buildplt, travplt)
   else
      close(stream)
      close(io)
   end
end

function plot_profile( filename::String )
   df = readtable(filename, separator = '\t', header=false, names=[:depth, :step, :length, :nodecount, :buildtime,
                                                                   :travtime, :travbulgetime, :travmistime, :travbulgemistime, :duplexes])

   buildplt = plot(df, x=:length, y=:buildtime, color=:depth, Geom.smooth,
         Coord.cartesian(ymin=0),
         Guide.xlabel("Sequence Length"),
         Guide.ylabel("Trie Build Time (sec)"))
   nodeplt  = plot(df, x=:length, y=:nodecount,  color=:depth, Geom.line,
         Coord.cartesian(ymin=0),
         Guide.xlabel("Sequence Length"),
         Guide.ylabel("Node Count"))
   travplt  = plot(df, x=:length, y=:travtime,  color=:depth, Geom.smooth,
         Coord.cartesian(ymin=0),
         Guide.xlabel("Sequence Length"),
         Guide.ylabel("Traversal Time (sec)"))
  duplextravplt = plot(df, x=:duplexes, y=:travtime, color=:depth, Geom.line,
         Coord.cartesian(ymin=0),
         Guide.xlabel("Returned Duplexes"),
         Guide.ylabel("Traversal Time (sec)"))

   draw(PDF("profile.pdf", 9inch, 8inch), vstack(hstack(buildplt, nodeplt), hstack(travplt, duplextravplt)))

   res = stack(df, [:buildtime, :travtime, :travbulgetime, :travmistime, :travbulgemistime])
   timeplot = plot(res, x=:length, y=:value, color=:variable, Geom.smooth, Coord.cartesian(ymin=0, fixed=1.0),
                   Guide.xlabel("Sequence Length"),
                   Guide.ylabel("Traversal Time (sec)"))
   draw(PDF("profile_time.pdf", 4inch, 4inch), timeplot)
end


# Converts the weigths of a PyTorch model to a Flux model from Metalhead

# PyTorch need to be installed

# Tested on ResNet and VGG models

using Flux
import Metalhead
using DataStructures
using Statistics
using BSON
using PyCall
using Images
using Test

torchvision = pyimport("torchvision")
torch = pyimport("torch")

modellib = [
    ("vgg11", () -> Metalhead.VGG(11), torchvision.models.vgg11),
    ("vgg13", () -> Metalhead.VGG(13), torchvision.models.vgg13),
    ("vgg16", () -> Metalhead.VGG(16), torchvision.models.vgg16),
    ("vgg19", () -> Metalhead.VGG(19), torchvision.models.vgg19),
    ("resnet18", () -> Metalhead.ResNet(18), torchvision.models.resnet18),
    ("resnet34", () -> Metalhead.ResNet(34), torchvision.models.resnet34),
    ("resnet50", () -> Metalhead.ResNet(50), torchvision.models.resnet50),
    ("resnet101",() -> Metalhead.ResNet(101),torchvision.models.resnet101),
    ("resnet152",() -> Metalhead.ResNet(152),torchvision.models.resnet152),
]

function _list_state(node::Flux.BatchNorm,channel,prefix)
    # use the same order of parameters than PyTorch
    put!(channel, (prefix * ".γ", node.γ)) # weigth (learnable)
    put!(channel, (prefix * ".β", node.β)) # bias (learnable)
    put!(channel, (prefix * ".μ", node.μ))  # running mean
    put!(channel, (prefix * ".σ²", node.σ²)) # running variance
end

function _list_state(node::Union{Flux.Conv,Flux.Dense},channel,prefix)
    put!(channel, (prefix * ".weight", node.weight))

    if node.bias !== Flux.Zeros()
        put!(channel, (prefix * ".bias", node.bias))
    end
end

_list_state(node,channel,prefix) = nothing

function _list_state(node::Union{Flux.Chain,Flux.Parallel},channel,prefix)
    for (i,n) in enumerate(node.layers)
        _list_state(n,channel,prefix * ".layers[$i]")
    end
end

function list_state(node; prefix = "model")
    Channel() do channel
        _list_state(node,channel,prefix)
    end
end

for (modelname,jlmodel,pymodel) in modellib

    model = jlmodel()
    pytorchmodel = pymodel(pretrained=true)

    state = OrderedDict(list_state(model.layers))

    # pytorchmodel.state_dict() looses the order
    state_dict = OrderedDict(pycall(pytorchmodel.state_dict,PyObject).items())
    pytorch_pp = OrderedDict((k,v.numpy()) for (k,v) in state_dict if !occursin("num_batches_tracked",k))

    # loop over all parameters
    for ((flux_key,flux_param),(pytorch_key,pytorch_param)) in zip(state,pytorch_pp)
        if size(flux_param) == size(pytorch_param)
            # Dense weight and vectors
            flux_param .= pytorch_param
        elseif size(flux_param) == reverse(size(pytorch_param))
            tmp = pytorch_param
            tmp = permutedims(tmp,ndims(tmp):-1:1)

            if ndims(flux_param)  == 4
                # convolutional weights
                flux_param .= reverse(tmp,dims=(1,2))
            else
                flux_param .= tmp
            end
        else
            @debug begin
                @show size(flux_param), size(pytorch_param)
            end
            error("incompatible shape $flux_key $pytorch_key")
        end
    end

    @info "saving model $modelname"
    BSON.@save joinpath(@__DIR__,"weights","$(modelname).bson") model
end