include("keystate.jl")

mutable struct HostNode <: Ranger.AbstractNode
    base::Nodes.NodeData
    transform::Nodes.TransformProperties{Float64}

    # Collection for AnchorNode
    children::Array{Ranger.AbstractNode,1}

    # Colors
    color::Rendering.Palette

    # An icon so we can see the host
    polygon::Geometry.Polygon

    # Total AABB
    aabb::Geometry.AABB{Float64}

    # Hit detection for both projectiles and dragging
    detection::Nodes.Detection
    inside::Bool

    keystate::KeyState
    drag::DragState
    node_point::Geometry.Point{Float64}

    particle_system::Particles.ParticleSystem

    function HostNode(world::Ranger.World, name::String, parent::Ranger.AbstractNode)
        o = new()

        o.base = Nodes.NodeData(name, parent, world)
        o.transform = Nodes.TransformProperties{Float64}()
        o.children = Array{Ranger.AbstractNode,1}[]
        o.color = Rendering.White()

        o.aabb = Geometry.AABB{Float64}()

        o.polygon = Geometry.Polygon{Float64}()
        o.detection = Nodes.Detection(Rendering.Yellow(), Rendering.Red())
        o.inside = false
        o.keystate = KeyState(false, false, Int8(0))
        o.drag = DragState()
        o.node_point = Geometry.Point{Float64}()

        # Setup particle activator
        activator = Particles.Activator360()
        activator.max_life = 3.0

        # Create and activate particle system.
        o.particle_system = Particles.ParticleSystem(activator)
        o.particle_system.active = true
        o.particle_system.auto_trigger = true

        # Add a few particles to the system.
        for i in 1:Particles.MAX_PARTICLES
            p = Particles.Particle{Float64}()
            Particles.add_particle!(o.particle_system, p)

            # Add a particle node to the parent (i.e. a visual).
            pnode = BasicParticleNode(world, "ParticleNode:" * string(i), parent)
            p.visual = pnode
            pnode.color = Rendering.Red()
            pnode.base.id = i # assign an id for lookups.
            Nodes.set_scale!(pnode, Particles.MAX_PARTICE_SIZE) # an initial size

            # Important! We are adding the node to the parent and NOT this host node.
            # We want the particle to be independent of the emitter.
            push!(Nodes.get_children(parent), pnode)
        end

        build(o, world)
        o
    end
end

function build(node::HostNode, world::Ranger.World)
    Geometry.add_vertex!(node.polygon, 0.05, -0.05)
    Geometry.add_vertex!(node.polygon, 0.05, -0.15)
    Geometry.add_vertex!(node.polygon, 0.0, -0.25)
    Geometry.add_vertex!(node.polygon, -0.05, -0.25)
    Geometry.add_vertex!(node.polygon, -0.05, -0.15)
    Geometry.add_vertex!(node.polygon, -0.05, -0.05)

    Geometry.add_vertex!(node.polygon, -0.15, -0.05)
    Geometry.add_vertex!(node.polygon, -0.25, 0.0)
    Geometry.add_vertex!(node.polygon, -0.25, 0.05)
    Geometry.add_vertex!(node.polygon, -0.15, 0.05)
    Geometry.add_vertex!(node.polygon, -0.05, 0.05)

    Geometry.add_vertex!(node.polygon, -0.05, 0.15)
    Geometry.add_vertex!(node.polygon, 0.0, 0.25)
    Geometry.add_vertex!(node.polygon, 0.05, 0.25)
    Geometry.add_vertex!(node.polygon, 0.05, 0.15)
    Geometry.add_vertex!(node.polygon, 0.05, 0.05)

    Geometry.add_vertex!(node.polygon, 0.15, 0.05)
    Geometry.add_vertex!(node.polygon, 0.25, 0.0)
    Geometry.add_vertex!(node.polygon, 0.25, -0.05)
    Geometry.add_vertex!(node.polygon, 0.15, -0.05)
    Geometry.add_vertex!(node.polygon, 0.05, -0.05)

    Geometry.build!(node.polygon)

    Nodes.set_dirty!(node, true);
end

# --------------------------------------------------------
# Timing
# --------------------------------------------------------
function Nodes.update(node::HostNode, dt::Float64)
    Nodes.update!(node.detection, node)
    Particles.update!(node.particle_system, dt)

    pos = node.transform.position
    Particles.set_position!(node.particle_system, pos.x, pos.y)
end

# Called during visit/rendering which is AFTER update()
function Nodes.interpolate(node::HostNode, interpolation::Float64)
end

# --------------------------------------------------------
# Rendering
# --------------------------------------------------------
function Nodes.draw(node::HostNode, context::Rendering.RenderContext)
    if Nodes.is_dirty(node)
        Rendering.transform!(context, node.polygon)

        Nodes.set_dirty!(node, false)
    end

    Rendering.set_draw_color(context, node.color)
    Rendering.render_outlined_polygon(context, node.polygon, Rendering.CLOSED);

    aabb_color = Nodes.highlight_color(node.detection, node.inside)
    Rendering.set_draw_color(context, aabb_color)

    Geometry.set!(node.aabb, Nodes.get_bucket(node))
    Rendering.render_aabb_rectangle(context, node.aabb)

end

# --------------------------------------------------------
# Life cycle events
# --------------------------------------------------------
function Nodes.enter_node(node::HostNode, man::Nodes.NodeManager)
    println("enter ", node);
    Nodes.register_event_target(man, node);

    node.inside = Nodes.check!(node.detection, node, node.base.world)
end

function Nodes.exit_node(node::HostNode, man::Nodes.NodeManager)
    println("exit ", node);
    Nodes.unregister_event_target(man, node);
end

# --------------------------------------------------------
# Events
# --------------------------------------------------------
# Here we elect to receive keyboard events.
function Nodes.io_event(node::HostNode, event::Events.KeyboardEvent)
    # Events.print(event)
    if (event.keycode == Events.KEY_F && event.state == Events.KEY_PRESSED)
        pos = node.transform.position
        Particles.set_position!(node.particle_system, pos.x, pos.y)
        Particles.trigger_oneshot!(node.particle_system)
    end

    if (event.keycode == Events.KEY_E && event.state == Events.KEY_PRESSED)
        pos = node.transform.position
        Particles.set_position!(node.particle_system, pos.x, pos.y)
        Particles.explode!(node.particle_system)
    end

    if (event.keycode == Events.KEY_W && event.state == Events.KEY_PRESSED)
        node.particle_system.auto_trigger = !node.particle_system.auto_trigger
        if node.particle_system.auto_trigger
            println("Autotrigger ON")
        else
            println("Autotrigger OFF")
        end
    end

    set_state!(node.keystate, event)
end

function Nodes.io_event(node::HostNode, event::Events.MouseMotionEvent)
    set_motion_state!(node.drag, event.x, event.y, node)

    dragging = is_dragging(node.drag)
    # println(dragging, ", ", node.inside)
    if dragging && node.inside
        pos = node.transform.position
        Nodes.set_position!(node, pos.x + node.drag.delta.x, pos.y + node.drag.delta.y)
    end

    Nodes.set_device_point!(node.detection, Float64(event.x), Float64(event.y))
    node.inside = Nodes.check!(node.detection, node, node.base.world)
end

function Nodes.io_event(node::HostNode, event::Events.MouseButtonEvent)
    set_button_state!(node.drag, event.x, event.y, event.button, event.state, node)
end

# --------------------------------------------------------
# Grouping
# --------------------------------------------------------
function Nodes.get_children(node::HostNode)
    node.children
end

# --------------------------------------------------------
# Host specific functionality
# --------------------------------------------------------


using ADI
using ProgressLogging
using PSFModels: Gaussian
using PyCall
using SiriusB

fits = pyimport("astropy.io.fits")
pro = pyimport("proplot")

pro.use_style("ggplot")
pro.rc["image.cmap"] = "inferno"
pro.rc["image.origin"] = "lower"
pro.rc["grid"] = false

_label(a::Type{<:PCA}) = "PCA"
_label(a::Type{<:NMF}) = "NMF"
_label(a::Type{<:GreeDS}) = "GreeDS"

_epochs = Dict(
    "2020feb04" => "2020-02-04",
    "2020nov21" => "2020-11-21",
    "2020nov28" => "2020-11-28"
)

function plot_mosaic(res_cubes, angles; epoch, label, ncomps)
    py"""
    import proplot as pro
    fig, axs = pro.subplots(nrows=2, ncols=5, share=3, width="7.5in", space=0)
    """
    i = 0
    flats = collapse.(res_cubes, Ref(angles))
    vmax = mapreduce(maximum, max, flats)
    vmin = mapreduce(minimum, min, flats)

    for (n, flat) in zip(ncomps, flats)
        lab = "ncomp=$n"

        py"""
        m = axs[$i].imshow($flat, vmin=$vmin, vmax=$vmax)
        axs[$i].text(6, 6, $lab, color="w")
        """
        i += 1
    end

    ticks = range(40, 159, length=3)
    tick_labs = @. string(round(SiriusB.pxscale * (ticks - 99.5), digits=1))

    figname = make_filename_friendly(figuredir("reports", "$(epoch)_$(label)_mosaic.pdf"))
    py"""
    axs.format(
        yticks=$ticks,
        yticklabels=$tick_labs,
        ylabel="y [arcsec]",
        xticks=$ticks,
        xticklabels=$tick_labs,
        xlabel="x [arcsec]",
    )
    fig.colorbar(m, loc="r")
    fig.savefig($figname)
    """
end

function plot_results(res_cubes, angles, contrast_curves; epoch, label, fwhm)
    py"""
    import proplot as pro
    fig, axs = pro.subplots([[1, 3, 3], [2, 3, 3]], share=0, width="7.5in")
    cycle = pro.Cycle("viridis_r", $(length(res_cubes)))
    lines = []
    """
    # plot the average STIM map and SLIM probability map
    ticks = range(20, 179, length=5)
    tick_labs = @. string(round(SiriusB.pxscale * (ticks - 99.5), digits=1))

    stimmap, slimmask = slimmap(res_cubes, angles; N=ceil(Int, π/4 * fwhm^2))
    slimprob = stimmap .* slimmask
    py"""
    axs[0].imshow($stimmap, vmin=0, vmax=1)
    axs[0].text(6, 6, "mean STIM", color="w")
    m = axs[1].imshow($slimprob, vmin=0, vmax=1)
    axs[1].text(6, 6, "SLIM map", color="w")
    fig.colorbar(m, loc="l", label="STIM probability")

    axs[0].format(
        ylabel="y [arcsec]",
        yticks=$ticks,
        yticklabels=$tick_labs,
        xticks="null"
    )
    axs[1].format(
        ylabel="y [arcsec]",
        xlabel="x [arcsec]",
        yticks=$ticks,
        yticklabels=$tick_labs,
        xticks=$ticks,
        xticklabels=$tick_labs
    )
    """
    # now plot contrast curves
    for curve in contrast_curves
        dist = curve.distance .* SiriusB.auscale
        py"""
        l = axs[2].plot($dist, $(curve.contrast), cycle=cycle)
        lines.extend(l)
        """
    end
    figname = make_filename_friendly(figuredir("reports", "$(epoch)_$(label)_contrast_curves.pdf"))
    py"""
    axs[2].text(0.25, 5e-4, $(_epochs[epoch]), fontsize=14)
    fig.colorbar(lines, loc="r", label="ncomp", values=range(1, 11))
    axs[2].dualx(lambda x: x * $(SiriusB.parallax), label="separation [arcsec]")
    ax2 = axs[2].alty(
        yticks=[0.09, 0.31, 0.54, 0.77, 1],
        yticklabels=["8", "6", "4", "2", "0"],
        label="Δ mag"
    )

    axs[2].format(
        ylim=(10**(-3.5), 1),
        yscale="log",
        yformatter="log",
        xlabel="projected separation [AU]",
        ylabel="5σ contrast",
        grid=True
    )
    fig.savefig($figname)
    """
end

function make_filename_friendly(filename)
    tmp = lowercase(filename)
    tmp = replace(tmp, " " => "_")
    tmp = replace(tmp, "(" => "-")
    tmp = replace(tmp,  ")" => "")
end

@progress name="epoch" for epoch in ["2020feb04", "2020nov21", "2020nov28"]
    cube = fits.getdata(datadir("epoch_$epoch", "processed", "$(epoch)_sirius-b_cube_calib_registered_crop.fits"))
    angles = fits.getdata(datadir("epoch_$epoch", "processed", "$(epoch)_sirius-b_pa.fits"))

    raw_median = collapse(cube)
    res = gaussian_fit(raw_median, (100.5, 100.5))
    fwhm = res[3]
    psf_model = Gaussian{eltype(cube)}(res[1:2], fwhm)
    @info "FWHM for epoch $epoch is $fwhm px"
    starphot = Metrics.estimate_starphot(cube, fwhm)
    psfphot = Metrics.estimate_starphot(psf_model, fwhm) * exp(res[4])
    @info "star photometry is ≈ $starphot (PSF photometry was $psfphot)"

    # naturally mask out inner FWHM
    av_cube = AnnulusView(cube; inner=fwhm)

    ncomps = 1:10
    @progress name="algorithm" for alg in [PCA, NMF, GreeDS]
        label = _label(alg)
        @withprogress name="ncomp" begin
            i = 1
            N = length(ncomps)
            outputs = map(ncomps) do n
                _alg = alg(n)
                fname = make_filename_friendly(datadir("epoch_$epoch", "residuals", "$(epoch)_sirius-b_residual_$label-$n.fits"))
                res_cube = load_or_produce(fname) do
                    subtract(_alg, av_cube; angles, fwhm)
                end
                fname_cc = make_filename_friendly(datadir("epoch_$epoch", "residuals", "$(epoch)_sirius-b_contrast-curve_$label-$n.csv"))
                cc = load_or_produce(fname_cc) do
                    contrast_curve(_alg, av_cube, angles, psf_model; fwhm, angles, starphot, nbranch=6)
                end
                @logprogress i/N
                i += 1
                res_cube, cc
            end
            res_cubes = map(first, outputs)
            contrast_curves = map(last, outputs)
            plot_mosaic(res_cubes, angles; epoch, label, ncomps)
            plot_results(res_cubes, angles, contrast_curves; epoch, label, fwhm)
        end
    end
end


module Potato

using JuMP
using XLSX
using DataFrames
using Distances
using UUIDs
using Statistics

include("struct.jl")
include("settings.jl")
include("base.jl")
include("core.jl")
include("search.jl")

#struct
export vtx,veh,col,node,dv,β,bound,S

#settings
export set_optimizer!
export get_optimizer
export reset_optimizer
export set_slack_coeff!
export set_surp_coeff!
export su_C
export sl_C
export silence!
export silent

#base
export extract!
export V,dist,K,T,d
export root

#core
export passes
export callSub,buildSub!
export master,getDuals
export sub,getCols
export colGen
export origin
export Q,s,f
export prioritize_vertex!,i_priority
export prioritize_vehicle!,k_priority
export fractional_column
export find_branch
export separate
export find_separator

#search
export leaf,traverse
export createBranch

end


# Clayton, N. S., Yu, K. S. and Dickinson, A. (2001), Scrub jays (Aphelocoma
# coerulescens) form integrated memories of the multiple features of caching
# episodes. Journal of Experimental Psychology: Animal Behavior Processes, 1:17-29
#
# Experiments 1 & 2 are merged because they are very similar and involve the
# same birds. Interleaved training trials have set index 0, test phase of
# experiment 1 set index 5 and test phase of experiment 2 has set index 6.

const CLAYTON0103_EXPERIMENTS = (:Clayton01_exp1, :Clayton01_exp2, :Clayton01_exp3,
                                 :Clayton01_exp4, :Clayton03_exp1, :Clayton03_exp2)
function summarize(::Experiment{:Clayton0103}, results)
    [summarize(EXPERIMENTS[key], res)
     for (key, res) in zip(CLAYTON0103_EXPERIMENTS, results)]
end
function _score(exp::Experiment{:Clayton0103}, result, metric)
    vcat([_score(EXPERIMENTS[key], res, metric)
          for (key, res) in zip(CLAYTON0103_EXPERIMENTS, result)]...)
end
function statistical_tests(exp::Experiment{:Clayton0103}, data)
    [statistical_tests(EXPERIMENTS[key], d)
     for (key, d) in zip(CLAYTON0103_EXPERIMENTS, data)]
end
function run!(::Experiment{:Clayton0103}, models)
    results = [run!(EXPERIMENTS[key], models) for key in CLAYTON0103_EXPERIMENTS]
    results[5] = vcat(results[1], results[5]) # data from Clayton01_exp1 was reused in Clayton03_exp2
    results
end

function cricketproportions(results)
    pcinspect = @where(results, :action .== "inspect", (|).(:set .== 4, :set .== 7),
                      :trial .== "pc", :RI .> 4)
    sort!(pcinspect, [:id, :foodtype, :RI, :set])
    pcinspectcricket = @where(pcinspect, :foodtype .== "other")
    pcinspectcricket.counts = float.(pcinspectcricket.counts)
    pcinspectcricket.counts ./= (pcinspectcricket.counts .+
                                @where(pcinspect, :foodtype .== "peanut").counts .+
                                1e-16) # I prefer 0 if no site is ever inspected over NaN because NaN propagates and may lead to NaN scores.
    pcinspectcricket
end
function _summarize(::Experiment{:Clayton01_exp1}, results)
    res = combine(groupby(@where(results, :action .== "inspect", :set .!= 0,
                                 (|).(:foodtype .!= "other", :set .== 5)),
                          [:group, :foodtype, :action, :RI, :trial, :set]),
             df -> DataFrame(μ = mean(df.counts), sem = sem(df.counts),
                             firstinspection = mean(df.firstinspection),
                             n = length(df.counts)))
    for g in groupby(@where(results, :action .== "cache", :set .== 5), :group)
        push!(res, [g.group[1], missing, "cache", missing, missing, 5,
                    mean(g.counts), sem(g.counts), 0, length(g.counts)])
    end
    res
end
firstinspect_dist(df::DataFrame; n = 4) = firstinspect_dist(df.firstinspection; n = n)
function firstinspect_dist(firstinspections; n = 4)
    patterncounts = zeros(Int, 2^n)
    for i in 1:n:length(firstinspections)
        p = 1
        for j in 0:n - 1
            firstinspections[i + j] == 0 || (p += 2^j)
        end
        patterncounts[p] += 1
    end
    patterncounts
end
function statistical_tests(exp::Experiment{:Clayton01_exp1}, data)
    tests = Test[]
    for g in groupby(@where(data, :action .== "inspect", :set .== 4), :RI)
        push!(tests, Test("first inspect set 4 RI $(g.RI[1])",
                          chisqtest(firstinspect_dist(@where(g, :group .== "degrade")),
                                    firstinspect_dist(@where(g, :group .== "replenish"))),
                          exp.tests))
    end
    datatest = @where(data, :action .== "inspect", :set .== 5)
    for g in groupby(datatest, :RI)
        RI = g.RI[1]
        id = "inspect test RI $RI"
        push!(tests, Test(id,
                          splitplotanova(g,
                                         exp.tests[id].locals,
                                         betweenfactors = :group,
                                         withinfactors = :foodtype,
                                         id = :id, data = :counts),
                          exp.tests))
        push!(tests, Test("first inspect test RI $RI",
                          chisqtest(firstinspect_dist(@where(g, :group .== "degrade")),
                                    firstinspect_dist(@where(g, :group .== "replenish"))),
                          exp.tests))
        if RI == 100
            for subg in groupby(g, :group)
                id = "inspect test RI 100 $(subg.group[1])"
                push!(tests, Test(id,
                                  anova(subg, @formula(counts ~ foodtype),
                                        exp.tests[id].locals),
                                  exp.tests))
            end
        end
    end
    datatestcache = @where(data, :action .== "cache", :set .== 5)
    for g in groupby(datatest, :RI)
        RI = g.RI[1]
        id = "cache test RI $RI"
        push!(tests, Test(id,
                          splitplotanova(g,
                                         exp.tests[id].locals,
                                         betweenfactors = :group,
                                         withinfactors = :foodtype,
                                         id = :id, data = :counts),
                          exp.tests))
    end
    TestCollection(tests)
end
value(x::Quantity) = x.val
value(x) = div(x, 60)
@inline function Clayton01_trial!(m, results, group, id, set,
                        foodtype1, foodtype2, wait;
                        degrade, pilfer, bothtrayspresent = false,
                        waitbetween = 0u"hr", specificothername = false)
    foodtype1string = lowercase(string(foodtype1))
    foodtype2string = foodtype1 == Peanut && !specificothername ? "other" : lowercase(string(foodtype2))
    trial = "$(foodtype1string[1])$(lowercase(string(foodtype2)[1]))"
    remove!(m, Any)
    wait!(m, 4u"hr")
    tray1 = Tray(hack(1)) # one half of tray
    tray2 = Tray(hack(2)) # second half of tray
    if bothtrayspresent
        cover!(tray2)
        add!(m, tray2)
    end
    add!(m, tray1)
    add!(m, foodtype1, 50)
    wait!(m, 15u"minute")
    remove!(m, Any)
    push!(results, [group, foodtype1string, id, set, trial, value(wait), "cache", countcache(tray1), missing])
    if waitbetween > 0u"hr"
        add!(m, MaintenanceDiet)
        wait!(m, waitbetween - 4u"hr")
        remove!(m, Any)
        wait!(m, 4u"hr")
    end
    if bothtrayspresent
        uncover!(tray2)
        cover!(tray1)
        add!(m, tray1)
    end
    add!(m, tray2)
    add!(m, foodtype2, 50)
    wait!(m, 15u"minute")
    remove!(m, Any)
    uncover!(tray1)
    push!(results, [group, foodtype2string, id, set, trial, value(wait), "cache",
                    countcache(tray2), missing])
    if wait > 4u"hr"
        add!(m, MaintenanceDiet)
        wait!(m, wait - 4u"hr")
        remove!(m, MaintenanceDiet)
    end
    if degrade
        degrade!(tray2)
    end
    if pilfer
        pilfer!(tray2)
        pilfer!(tray1)
    end
    wait!(m, 4u"hr")
    add!(m, tray1)
    add!(m, tray2)
    inspectionobserver = InspectionObserver()
    add!(m, inspectionobserver)
    wait!(m, 5u"minute")
    remove!(m, Any)
    push!(results, [group, foodtype1string, id, set, trial, value(wait), "inspect",
                    countinspections(inspectionobserver, tray1),
                    firstinspection(inspectionobserver, tray1)])
    push!(results, [group, foodtype2string, id, set, trial, value(wait), "inspect",
                    countinspections(inspectionobserver, tray2),
                    firstinspection(inspectionobserver, tray2)])
    add!(m, MaintenanceDiet)
    wait!(m, 25u"minute" + 15u"hr")
end
function isdegrade(group, foodtype, wait)
    if group === "degrade"
        foodtype == Mealworm && wait > 4u"hr" ||
        foodtype == Cricket && wait > 28u"hr" ||
        foodtype == Waxworm && wait > 4u"hr"
    elseif group == "ripen"
        foodtype == Waxworm && wait <= 4u"hr"
    else
        false
    end
end
function run!(::Experiment{:Clayton01_exp1}, models)
    results = DataFrame(group = Union{String,Missing}[],
                        foodtype = Union{String,Missing}[],
                        id = Int[], set = Int[], trial = Union{String,Missing}[],
                        RI = Union{Int,Missing}[],
                        action = String[], counts = Int[],
                        firstinspection = Union{Missing,Int}[])
    id = 0
    for group in ["degrade", "replenish"]
        for i in 1:8
            id += 1
            m = models[id]
            # training (May 1999)
            for set in 1:4
                ris = [4u"hr", 28u"hr", 100u"hr"]
                isodd(set) && reverse!(ris)
                for wait in ris
                    for foodtype in [Mealworm, Cricket]
                        Clayton01_trial!(m, results, group, id, set, Peanut, foodtype, wait,
                               degrade = isdegrade(group, foodtype, wait), pilfer = false)
                    end
                end
            end
            # test
            ris = [28u"hr", 100u"hr", 4u"hr"]
            for wait in ris
                for foodtype in [Mealworm, Cricket]
                    Clayton01_trial!(m, results, group, id, 5, Peanut, foodtype, wait,
                           degrade = false, pilfer = true)
                    Clayton01_trial!(m, results, group, id, 0, Peanut, foodtype, 4u"hr",
                           degrade = false, pilfer = false)
                end
            end
        end
    end
    results
end

function _summarize(::Experiment{:Clayton01_exp2}, results)
    res = combine(groupby(@where(results, :action .== "inspect", :set .!= 0),
                          [:group, :foodtype, :action, :RI, :trial, :set]),
             df -> DataFrame(μ = mean(df.counts), sem = sem(df.counts),
                             firstinspection = mean(df.firstinspection),
                             n = length(df.counts)))
    for g in groupby(@where(results, :action .== "cache", :set .>= 6),
                     [:foodtype, :set])
        push!(res, [missing, g.foodtype[1], "cache", missing, missing, g.set[1],
                    mean(g.counts), sem(g.counts), 0, length(g.counts)])
    end
    res
end
function statistical_tests(exp::Experiment{:Clayton01_exp2}, data)
    tests = Test[]
    for g in groupby(@where(data, :set .== 6), :action)
        action  = g.action[1]
        id = "$action test exp 2"
        push!(tests, Test(id,
                          splitplotanova(g,
                                         exp.tests[id].locals,
                                         betweenfactors = :group,
                                         withinfactors = [:foodtype, :RI],
                                         id = :id, data = :counts),
                          exp.tests))
        if action == "inspect"
            for subg in groupby(g, :RI)
                id = "$action test exp 2 RI $(subg.RI[1])"
                push!(tests, Test(id,
                                  splitplotanova(subg,
                                                 exp.tests[id].locals,
                                                 betweenfactors = :group,
                                                 withinfactors = :foodtype,
                                                 id = :id, data = :counts),
                                  exp.tests))
            end
            for subg in groupby(g, :group)
                group = subg.group[1]
                id = "$action test exp 2 group $group"
                push!(tests, Test(id,
                                  anova(subg, @formula(counts ~ foodtype * RI), exp.tests[id].locals),
                                  exp.tests))
                if group == "degrade"
                    for subsubg in groupby(subg, :RI)
                        id = "$action test exp 2 group $group RI $(subsubg.RI[1])"
                        push!(tests, Test(id,
                                          anova(subsubg, @formula(counts ~ foodtype), exp.tests[id].locals),
                                          exp.tests))
                    end
                end
            end
        end
    end
    datatest2inspect = sort(@where(data, :set .== 6, :action .== "inspect"), [:id, :RI, :foodtype])
    push!(tests, Test("first inspect test exp 2",
                      chisqtest(firstinspect_dist(@where(datatest2inspect, :group .== "degrade")),
                                firstinspect_dist(@where(datatest2inspect, :group .== "replenish"))),
                      exp.tests))
    TestCollection(tests)
end
function run!(::Experiment{:Clayton01_exp2}, models)
    results = DataFrame(group = Union{String,Missing}[],
                        foodtype = Union{String,Missing}[],
                        id = Int[], set = Int[], trial = Union{String,Missing}[],
                        RI = Union{Int,Missing}[],
                        action = String[], counts = Int[],
                        firstinspection = Union{Missing,Int}[])
    id = 0
    for group in ["degrade", "replenish"]
        for i in 1:8
            id += 1
            m = models[id]
            # training before experiment 2 (July 1999)
            for foodtype in [Mealworm, Cricket]
                Clayton01_trial!(m, results, group, id, 0, Peanut, foodtype, 28u"hr",
                       degrade = isdegrade(group, foodtype, 28u"hr"), pilfer = false)
            end
            # test experiment 2
            ris = [4u"hr", 28u"hr"]
            isodd(id) && reverse!(ris)
            for wait in ris
                Clayton01_trial!(m, results, group, id, 6, Mealworm, Cricket, wait,
                       degrade = false, pilfer = true)
                # interleaved training trials
                for foodtype in [Mealworm, Cricket]
                    Clayton01_trial!(m, results, group, id, 0, Peanut, foodtype, 28u"hr",
                           degrade = isdegrade(group, foodtype, 28u"hr"), pilfer = false)
                end
            end
        end
    end
    results
end

function _summarize(::Experiment{:Clayton01_exp3}, results)
    res = [combine(groupby(@where(results, :action .== "inspect", :set .== 1),
                           [:group, :foodtype, :action, :trial]),
             df -> DataFrame(μ = mean(df.counts), sem = sem(df.counts),
                             firstinspection = mean(df.firstinspection),
                             n = length(df.counts)));
           combine(groupby(@where(results, :action .== "cache", :set .== 1), :foodtype),
              df -> DataFrame(μ = mean(df.counts), sem = sem(df.counts),
                              group = missing, action = "cache", trial = missing,
                              firstinspection = missing, n = length(df.counts)))]
end
function statistical_tests(exp::Experiment{:Clayton01_exp3}, data)
    tests = Test[]
    testdata = @where(data, :set .== 1)
    sort!(testdata, [:id, :foodtype, :trial])
    for g in groupby(testdata, :action)
        action = g.action[1]
        push!(tests, Test(action,
                          splitplotanova(g,
                                         exp.tests[action].locals,
                                         betweenfactors = :group,
                                         withinfactors = [:foodtype, :trial],
                                         id = :id, data = :counts),
                          exp.tests))
        if action == "inspect"
            for subg in groupby(g, :trial)
                id = "$action $(subg.trial[1])"
                push!(tests, Test(id,
                                  splitplotanova(subg,
                                                 exp.tests[id].locals,
                                                 betweenfactors = :group,
                                                 withinfactors = :foodtype,
                                                 id = :id, data = :counts),
                                  exp.tests))
            end
            for subg in groupby(g, :group)
                group = subg.group[1]
                id = "$action $group"
                push!(tests, Test(id,
                                  splitplotanova(subg,
                                                 exp.tests[id].locals,
                                                 betweenfactors = [],
                                                 withinfactors = [:foodtype, :trial],
                                                 id = :id, data = :counts),
                                  exp.tests))
                if group == "degrade"
                    for subsubg in groupby(subg, :trial)
                        trial = subsubg.trial[1]
                        condbigger = trial == "mc" ? "cricket" : "mealworm"
                        condsmaller = trial == "mc" ? "mealworm" : "cricket"
                        push!(tests, Test("$action $group $trial",
                                          is_significantly_bigger(subsubg,
                                                                  conditionkey = :foodtype,
                                                                  condbigger = condbigger,
                                                                  condsmaller = condsmaller,
                                                                  valuekey = :counts),
                                          exp.tests))
                    end
                end
            end
            push!(tests, Test("first inspectes",
                              chisqtest(firstinspect_dist(@where(g, :group .== "degrade")),
                                        firstinspect_dist(@where(g, :group .== "replenish"))),
                              exp.tests))
        end
    end
    TestCollection(tests)
end
function run!(::Experiment{:Clayton01_exp3}, models)
    results = DataFrame(group = Union{String,Missing}[],
                        foodtype = Union{String,Missing}[],
                        id = Int[], set = Int[], trial = Union{String,Missing}[],
                        RI = Union{Int,Missing}[],
                        action = String[], counts = Int[],
                        firstinspection = Union{Missing,Int}[])
    id = 0
    for group in ["degrade", "replenish"]
        for i in 1:8
            id += 1
            m = models[id]
            # training (November 1999)
            for wait in [28u"hr", 4u"hr"]
                for foodtype in [Mealworm, Cricket]
                    Clayton01_trial!(m, results, group, id, 0, Peanut, foodtype, wait,
                           degrade = isdegrade(group, foodtype, wait), pilfer = false)
                end
            end
            # testing
            trials = [(Cricket, Mealworm), (Mealworm, Cricket)]
            i >= 4 && reverse!(trials)
            for (food1, food2) in trials
                Clayton01_trial!(m, results, group, id, 1, food1, food2, 4u"hr",
                       degrade = false, pilfer = true, bothtrayspresent = true,
                       waitbetween = 24u"hr")
                # interleaved training
                for foodtype in [Mealworm, Cricket]
                    Clayton01_trial!(m, results, group, id, 0, Peanut, foodtype, 4u"hr",
                           degrade = false, pilfer = false)
                end
            end
        end
    end
    results
end

function _summarize(::Experiment{:Clayton01_exp4}, results)
    res = [combine(groupby(@where(results, :action .== "inspect", :set .== 1),
                           [:group, :foodtype, :action, :RI]),
             df -> DataFrame(μ = mean(df.counts), sem = sem(df.counts),
                             firstinspection = mean(df.firstinspection),
                             n = length(df.counts)));
           combine(groupby(@where(results, :action .== "cache", :set .== 1),
                           [:group, :foodtype]),
              df -> DataFrame(μ = mean(df.counts), sem = sem(df.counts),
                              action = "cache", RI = missing,
                              firstinspection = missing, n = length(df.counts)))]
end
function statistical_tests(exp::Experiment{:Clayton01_exp4}, data)
    tests = Test[]
    datacache = @where(data, :action .== "cache", :set .== 1)
    id = "cache"
    push!(tests, Test(id,
                      splitplotanova(datacache,
                                     exp.tests[id].locals,
                                     betweenfactors = :group,
                                     withinfactors = [:foodtype, :RI],
                                     id = :id, data = :counts),
                      exp.tests))
    for g in groupby(datacache, :foodtype)
        id = "cache $(g.foodtype[1])"
        push!(tests, Test(id,
                          splitplotanova(g,
                                         exp.tests[id].locals,
                                         betweenfactors = :group,
                                         withinfactors = :RI,
                                         id = :id, data = :counts),
                          exp.tests))
    end
    datacache = @where(data, :action .== "inspect", :set .== 1)
    sort!(datacache, [:id, :foodtype, :RI])
    id = "inspect"
    push!(tests, Test(id,
                      splitplotanova(datacache,
                                     exp.tests[id].locals,
                                     betweenfactors = :group,
                                     withinfactors = [:foodtype, :RI],
                                     id = :id, data = :counts),
                      exp.tests))
    for g in groupby(datacache, :RI)
        id = "inspect $(g.RI[1])"
        push!(tests, Test(id,
                          splitplotanova(g,
                                         exp.tests[id].locals,
                                         betweenfactors = :group,
                                         withinfactors = :foodtype,
                                         id = :id, data = :counts),
                          exp.tests))
    end
    datacacheegrade = @where(datacache, :group .== "degrade")
    id = "inspect degrade"
    push!(tests, Test(id,
                      splitplotanova(datacacheegrade,
                                     exp.tests[id].locals,
                                     betweenfactors = [],
                                     withinfactors = [:foodtype, :RI],
                                     id = :id, data = :counts),
                      exp.tests))
    for g in groupby(datacacheegrade, :RI)
        id = "inspect degrade $(g.RI[1])"
        push!(tests, Test(id,
                          splitplotanova(g,
                                         exp.tests[id].locals,
                                         betweenfactors = [],
                                         withinfactors = :foodtype,
                                         id = :id, data = :counts),
                          exp.tests))
    end
    push!(tests, Test("first inspectes",
                      chisqtest(firstinspect_dist(@where(datacache, :group .== "degrade")),
                                firstinspect_dist(@where(datacache, :group .== "replenish"))),
                      exp.tests))
    TestCollection(tests)
end
function run!(::Experiment{:Clayton01_exp4}, models)
    results = DataFrame(group = Union{String,Missing}[],
                        foodtype = Union{String,Missing}[],
                        id = Int[], set = Int[], trial = Union{String,Missing}[],
                        RI = Union{Int,Missing}[],
                        action = String[], counts = Int[],
                        firstinspection = Union{Missing,Int}[])
    id = 0
    for group in ["degrade", "replenish"]
        for i in 1:8
            id += 1
            m = models[id]
            # training (December 1999)
            Clayton01_trial!(m, results, group, id, 0, Peanut, Mealworm, 4u"hr",
                   degrade = false, pilfer = false)
            remove!(m, Any)
            wait!(m, 4u"hr")
            tray1 = Tray(hack(1)) # one half of tray
            tray2 = Tray(hack(2)) # second half of tray
            add!(m, tray1)
            add!(m, Peanut, 50)
            wait!(m, 15u"minute")
            remove!(m, Any)
            push!(results, [group, "peanut", id, 1, "pm", 28, "cache", countcache(tray1), missing])
            add!(m, tray2)
            add!(m, Mealworm, 50)
            wait!(m, 15u"minute")
            remove!(m, Any)
            push!(results, [group, "mealworm", id, 1, "pm", 28, "cache", countcache(tray2), missing])
            add!(m, MaintenanceDiet)
            wait!(m, 30u"minute" + 23u"hr")
            tray3 = Tray(hack(1)) # one half of tray
            tray4 = Tray(hack(2)) # second half of tray
            add!(m, tray3)
            add!(m, Peanut, 50)
            wait!(m, 15u"minute")
            remove!(m, Any)
            push!(results, [group, "peanut", id, 1, "pm", 4, "cache", countcache(tray3), missing])
            add!(m, tray4)
            add!(m, Mealworm, 50)
            wait!(m, 15u"minute")
            remove!(m, Any)
            push!(results, [group, "mealworm", id, 1, "pm", 4, "cache", countcache(tray4), missing])
            wait!(m, 4u"hr")
            pilfer!(tray1); pilfer!(tray2); pilfer!(tray3); pilfer!(tray4)
            add!(m, tray1); add!(m, tray2)
            inspectionobserver = InspectionObserver()
            add!(m, inspectionobserver)
            wait!(m, 5u"minute")
            remove!(m, Any)
            push!(results, [group, "peanut", id, 1, "pm", 28, "inspect",
                            countinspections(inspectionobserver, tray1),
                            firstinspection(inspectionobserver, tray1)])
            push!(results, [group, "mealworm", id, 1, "pm", 28, "inspect",
                            countinspections(inspectionobserver, tray2),
                            firstinspection(inspectionobserver, tray2)])
            add!(m, tray3); add!(m, tray4)
            inspectionobserver = InspectionObserver()
            add!(m, inspectionobserver)
            wait!(m, 5u"minute")
            remove!(m, Any)
            push!(results, [group, "peanut", id, 1, "pm", 4, "inspect",
                            countinspections(inspectionobserver, tray3),
                            firstinspection(inspectionobserver, tray3)])
            push!(results, [group, "mealworm", id, 1, "pm", 4, "inspect",
                            countinspections(inspectionobserver, tray4),
                            firstinspection(inspectionobserver, tray4)])
        end
    end
    results
end


module Exp
include("./observable.jl")

experiments = Array{Dict{String,Array{Float64,1}},1}(undef, length(observables))
error_bars = Array{Dict{String,Array{Float64,1}},1}(undef, length(observables))

experiments[observables_index("Phosphorylated_MEKc")] = Dict(
    "EGF" => [0.000,0.773,0.439,0.252,0.130,0.087,0.080,0.066],
    "HRG" => [0.000,0.865,1.000,0.837,0.884,0.920,0.875,0.789],
)
error_bars[observables_index("Phosphorylated_MEKc")] = Dict(
    "EGF" => [0.000,0.030,0.048,0.009,0.009,0.017,0.012,0.008] ./ sqrt(3),
    "HRG" => [0.000,0.041,0.000,0.051,0.058,0.097,0.157,0.136] ./ sqrt(3),
)

experiments[observables_index("Phosphorylated_ERKc")] = Dict(
    "EGF" => [0.000,0.867,0.799,0.494,0.313,0.266,0.200,0.194],
    "HRG" => [0.000,0.848,1.000,0.971,0.950,0.812,0.747,0.595],
)
error_bars[observables_index("Phosphorylated_ERKc")] = Dict(
    "EGF" => [0.000,0.137,0.188,0.126,0.096,0.087,0.056,0.012] ./ sqrt(3),
    "HRG" => [0.000,0.120,0.000,0.037,0.088,0.019,0.093,0.075] ./ sqrt(3),
)

experiments[observables_index("Phosphorylated_RSKw")] = Dict(
    "EGF" => [0,0.814,0.812,0.450,0.151,0.059,0.038,0.030],
    "HRG" => [0,0.953,1.000,0.844,0.935,0.868,0.779,0.558],
)
error_bars[observables_index("Phosphorylated_RSKw")] = Dict(
    "EGF" => [0,0.064,0.194,0.030,0.027,0.031,0.043,0.051] ./ sqrt(3),
    "HRG" => [0,0.230,0.118,0.058,0.041,0.076,0.090,0.077] ./ sqrt(3),
)

experiments[observables_index("Phosphorylated_cFos")] = Dict(
    "EGF" => [0,0.060,0.109,0.083,0.068,0.049,0.027,0.017],
    "HRG" => [0,0.145,0.177,0.158,0.598,1.000,0.852,0.431],
)
error_bars[observables_index("Phosphorylated_cFos")] = Dict(
    "EGF" => [0,0.003,0.021,0.013,0.016,0.007,0.003,0.002] ./ sqrt(3),
    "HRG" => [0,0.010,0.013,0.001,0.014,0.000,0.077,0.047] ./ sqrt(3),
)

experiments[observables_index("Phosphorylated_CREBw")] = Dict(
    "EGF" => [0,0.446,0.030,0.000,0.000],
    "HRG" => [0,1.000,0.668,0.460,0.340],
)
error_bars[observables_index("Phosphorylated_CREBw")] = Dict(
    "EGF" => [0,0.0,0.0,0.0,0.0] ./ sqrt(3),
    "HRG" => [0,0.0,0.0,0.0,0.0] ./ sqrt(3),
)

experiments[observables_index("cfos_mRNA")] = Dict(
    "EGF" => [0,0.181,0.476,0.518,0.174,0.026,0.000],
    "HRG" => [0,0.353,0.861,1.000,0.637,0.300,0.059],
)
error_bars[observables_index("cfos_mRNA")] = Dict(
    "EGF" => [0.017,0.004,0.044,0.004,0.023,0.007,0.008] ./ sqrt(3),
    "HRG" => [0.017,0.006,0.065,0.044,0.087,0.023,0.001] ./ sqrt(3),
)

experiments[observables_index("cFos_Protein")] = Dict(
    "EGF" => [0,0.078,0.216,0.240,0.320,0.235],
    "HRG" => [0,0.089,0.552,0.861,1.000,0.698],
)
error_bars[observables_index("cFos_Protein")] = Dict(
    "EGF" => [0,0.036,0.028,0.056,0.071,0.048] ./ sqrt(3),
    "HRG" => [0,0.021,0.042,0.063,0.000,0.047] ./ sqrt(3),
)

experiments[observables_index("dusp_mRNA")] = Dict(
    "EGF" => [0.000,0.177,0.331,0.214,0.177,0.231],
    "HRG" => [0.000,0.221,0.750,1.000,0.960,0.934],
)
error_bars[observables_index("dusp_mRNA")] = Dict(
    "EGF" => [0.033,0.060,0.061,0.032,0.068,0.050] ./ sqrt(3),
    "HRG" => [0.027,0.059,0.094,0.124,0.113,0.108] ./ sqrt(3),
)

function get_timepoint(obs_name::String)::Vector{Float64}
    if obs_name in ["Phosphorylated_MEKc", "Phosphorylated_ERKc",
                    "Phosphorylated_RSKw", "Phosphorylated_cFos"]
        return [0., 300., 600., 900., 1800., 2700., 3600., 5400.]

    elseif obs_name == "Phosphorylated_CREBw"
        return [0., 600., 1800., 3600., 5400.]

    elseif obs_name == "cfos_mRNA"
        return [0.,600.,1200.,1800.,2700.,3600.,5400.]

    elseif obs_name in ["cFos_Protein", "dusp_mRNA"]
        return [0.,900.,1800.,2700.,3600.,5400.]
    end
end
end # module


using Test, TimeSpans, Dates

using TimeSpans: contains, nanoseconds_per_sample

@testset "basic TimeSpan code paths" begin
    t = TimeSpan(Nanosecond(rand(UInt32)))
    @test t == TimeSpan(t)
    @test t == TimeSpan(start(t), stop(t))
    @test t == TimeSpan(start(t), start(t))
    @test t == TimeSpan(start(t), start(t) + Nanosecond(1))
    @test contains(t, t)
    @test overlaps(t, t)
    @test start(t) ∈ t
    @test !(stop(t) ∈ t)
    @test stop(t) + Nanosecond(1) ∉ t
    @test shortest_timespan_containing([t]) == t
    @test shortest_timespan_containing((t,t,t)) == t
    @test duration(TimeSpan(start(t), stop(t) + Nanosecond(100))) == Nanosecond(101)
    @test duration(start(t)) == Nanosecond(1)
    @test_throws ArgumentError TimeSpan(4, 2)
    @test istimespan(t)
    @test istimespan(start(t))
    @test !istimespan(1)
    @test !istimespan(1:10)
    by = Second(rand(1:10))
    @test translate(t, by) === TimeSpan(start(t) + Nanosecond(by), stop(t) + Nanosecond(by))
    @test translate(t, -by) === TimeSpan(start(t) - Nanosecond(by), stop(t) - Nanosecond(by))
    @test repr(TimeSpan(6149872364198, 123412345678910)) == "TimeSpan(01:42:29.872364198, 34:16:52.345678910)"
end

@testset "format_duration" begin
    @test TimeSpans.format_duration(3723004005006) == "01:02:03.004005006"
    @test TimeSpans.format_duration(-3723004005006) == "-01:02:03.004005006"
end

@testset "contains(::TimeSpan...)" begin
    @test contains(TimeSpan(10, 20), TimeSpan(10, 20))
    @test contains(TimeSpan(10, 20), TimeSpan(11, 19))
    @test contains(TimeSpan(11, 20), TimeSpan(11, 19))
    @test contains(TimeSpan(10, 19), TimeSpan(11, 19))
    @test !contains(TimeSpan(10, 20), TimeSpan(11, 21))
    @test !contains(TimeSpan(11, 20), TimeSpan(10, 19))
    @test !contains(TimeSpan(10, 19), TimeSpan(10, 21))
    @test !contains(TimeSpan(11, 19), TimeSpan(10, 20))
    @test contains(TimeSpan(1, 10), Nanosecond(4))
end

@testset "overlaps(::TimeSpan...)" begin
    @test overlaps(TimeSpan(10, 20), TimeSpan(10, 20))
    @test overlaps(TimeSpan(10, 20), TimeSpan(11, 19))
    @test overlaps(TimeSpan(11, 20), TimeSpan(11, 19))
    @test overlaps(TimeSpan(10, 19), TimeSpan(11, 19))
    @test overlaps(TimeSpan(10, 20), TimeSpan(11, 21))
    @test overlaps(TimeSpan(11, 20), TimeSpan(10, 19))
    @test overlaps(TimeSpan(10, 19), TimeSpan(10, 21))
    @test overlaps(TimeSpan(11, 19), TimeSpan(10, 20))
    @test !overlaps(TimeSpan(20, 30), TimeSpan(10, 20))
    @test !overlaps(TimeSpan(10, 20), TimeSpan(20, 30))
    @test !overlaps(TimeSpan(10, 20), TimeSpan(21, 30))
    @test !overlaps(TimeSpan(21, 30), TimeSpan(10, 20))
end

@testset "shortest_timespan_containing(spans)" begin
    @test shortest_timespan_containing([TimeSpan(1, 2),
                                        TimeSpan(5, 10),
                                        TimeSpan(2, 3)]) == TimeSpan(1, 10)
    @test shortest_timespan_containing([TimeSpan(3, 7),
                                        TimeSpan(1, 10),
                                        TimeSpan(2, 5)]) == TimeSpan(1, 10)
end

@testset "time <--> index conversion" begin
    @test_throws ArgumentError time_from_index(200, 0)
    @test time_from_index(100, 1) == Nanosecond(0)
    @test time_from_index(100, 301:600) == TimeSpan(Second(3), Second(6))
    @test time_from_index(100, 101:101) == TimeSpan(Second(1))
    @test_throws ArgumentError index_from_time(200, Nanosecond(-1))
    @test index_from_time(100, Nanosecond(0)) == 1
    @test index_from_time(100, TimeSpan(Second(3), Second(6))) == 301:600
    @test index_from_time(100, TimeSpan(Second(1))) == 101:101
    # test non-integer sample rates
    rate = 100.66
    ns_per_sample = nanoseconds_per_sample(rate)
    for i in 1:1000
        t = Nanosecond(ceil(Int, (i - 1) * ns_per_sample))
        @test index_from_time(rate, t) == i
        @test time_from_index(rate, i) == t
    end
end

@testset "`in` and `findall`" begin
    @test findall(in(TimeSpan(1, 10)), Nanosecond.(5:15)) == 1:5
    @test findall(in(TimeSpan(1, 10)), map(Nanosecond, (9,10,11))) == 1:1
    @test in(TimeSpan(1,2))(Nanosecond(1))
    @test !in(TimeSpan(1,2))(Nanosecond(2))
end

@testset "index_from_time" begin
    @testset "docstring" begin
        @test index_from_time(1, Second(0)) == 1
        @test index_from_time(1, Second(1)) == 2
        @test index_from_time(100, Millisecond(999)) == 100
        @test index_from_time(100, Millisecond(1000)) == 101
    end

    @testset "floating-point precision" begin
        ns = Nanosecond((2 * 60 + 30) * 1e9)
        @test index_from_time(200, ns) == 30001
        @test index_from_time(200e0, ns) == 30001
        @test index_from_time(200f0, ns) == 30001
    end
end


# Nucleotides
# ===========
#
# DNA and RNA nucleotide types.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/Bio.jl/blob/master/LICENSE.md

# Valid nucleotides are represented in bytes using just the four low-order bits
abstract Nucleotide
bitstype 8 DNANucleotide <: Nucleotide
bitstype 8 RNANucleotide <: Nucleotide

# Conversion from/to integers
# ---------------------------

Base.convert(::Type{DNANucleotide}, nt::UInt8) = reinterpret(DNANucleotide, nt)
Base.convert(::Type{RNANucleotide}, nt::UInt8) = reinterpret(RNANucleotide, nt)
Base.convert(::Type{UInt8}, nt::DNANucleotide) = reinterpret(UInt8, nt)
Base.convert(::Type{UInt8}, nt::RNANucleotide) = reinterpret(UInt8, nt)
Base.convert{T<:Number,S<:Nucleotide}(::Type{T}, nt::S) = convert(T, UInt8(nt))
Base.convert{T<:Number,S<:Nucleotide}(::Type{S}, nt::T) = convert(S, UInt8(nt))

# Arithmetic and Order
# --------------------

# These methods are necessary when deriving some algorithims
# like iteration, sort, comparison, and so on.
@compat begin
    Base.:-{N<:Nucleotide}(x::N, y::N) = Int(x) - Int(y)
    Base.:-{N<:Nucleotide}(x::N, y::Integer) = x + (-y)
    Base.:+{N<:Nucleotide}(x::N, y::Integer) = reinterpret(N, (UInt8(x) + y % UInt8) & 0b1111)
    Base.:+{N<:Nucleotide}(x::Integer, y::N) = y + x
end
Base.isless{N<:Nucleotide}(x::N, y::N) = isless(UInt8(x), UInt8(y))

# Nucleotide encoding definition
# ------------------------------

# DNA Nucleotides

"DNA Invalid Nucleotide"
const DNA_INVALID = convert(DNANucleotide, 0b10000) # Indicates invalid DNA when converting string

# lookup table for characters
const char_to_dna = [DNA_INVALID for _ in 0x00:0x7f]
const dna_to_char = Vector{Char}(16)

# compatibility bits
const compatbits_nuc = zeros(UInt8, 16)

# derived from "The DDBJ/ENA/GenBank Feature Table Definition"
# §7.4.1 Nucleotide base code (IUPAC)
# http://www.insdc.org/documents/feature_table.html#7.4.1
for (code, (nt, doc, compat)) in enumerate([
        ('A', "DNA Adenine",                               0b0001),
        ('C', "DNA Cytosine",                              0b0010),
        ('G', "DNA Guanine",                               0b0100),
        ('T', "DNA Thymine",                               0b1000),
        ('M', "DNA Adenine or Cytosine",                   0b0011),
        ('R', "DNA Adenine or Guanine",                    0b0101),
        ('W', "DNA Adenine or Thymine",                    0b1001),
        ('S', "DNA Cytosine or Guanine",                   0b0110),
        ('Y', "DNA Cytosine or Thymine",                   0b1010),
        ('K', "DNA Guanine or Thymine",                    0b1100),
        ('V', "DNA Adenine, Cytosine or Guanine",          0b0111),
        ('H', "DNA Adenine, Cytosine or Thymine",          0b1011),
        ('D', "DNA Adenine, Guanine or Thymine",           0b1101),
        ('B', "DNA Cytosine, Guanine or Thymine",          0b1110),
        ('N', "DNA Adenine, Cytosine, Guanine or Thymine", 0b1111)])
    var = Symbol("DNA_", nt)
    @eval begin
        @doc $doc const $var = convert(DNANucleotide, $(UInt8(code - 1)))
        char_to_dna[$(Int(nt + 1))] = char_to_dna[$(Int(lowercase(nt) + 1))] = $var
        dna_to_char[$(code)] = $nt
        compatbits_nuc[$(code)] = $compat
    end
end

"DNA Gap"
const DNA_Gap = convert(DNANucleotide, 0b1111)
char_to_dna[Int('-') + 1] = DNA_Gap
dna_to_char[0b1111 + 1] = '-'
compatbits_nuc[0b1111 + 1] = 0b0000

"Returns Any DNA Nucleotide (DNA_N)"
nnucleotide(::Type{DNANucleotide}) = DNA_N
invalid_nucleotide(::Type{DNANucleotide}) = DNA_INVALID

Base.colon(start::DNANucleotide, stop::DNANucleotide) = SymbolRange(start, stop)

Base.isvalid(::Type{DNANucleotide}, x::Integer) = 0 ≤ x < 16
Base.isvalid(nt::DNANucleotide) = nt ≤ DNA_Gap
isambiguous(nt::DNANucleotide) = nt > DNA_T
alphabet(::Type{DNANucleotide}) = DNA_A:DNA_Gap
gap(::Type{DNANucleotide}) = DNA_Gap

# RNA Nucleotides

"Invalid RNA Nucleotide"
const RNA_INVALID = convert(RNANucleotide, 0b10000) # Indicates invalid RNA when converting string

# lookup table for characters
const char_to_rna = [RNA_INVALID for _ in 0x00:0x7f]
const rna_to_char = Vector{Char}(16)

for (code, (nt, doc)) in enumerate([
        ('A', "RNA Adenine"                             ),
        ('C', "RNA Cytosine"                            ),
        ('G', "RNA Guanine"                             ),
        ('U', "RNA Uracil"                              ),
        ('M', "RNA Adenine or Cytosine"                 ),
        ('R', "RNA Adenine or Guanine"                  ),
        ('W', "RNA Adenine or Uracil"                   ),
        ('S', "RNA Cytosine or Guanine"                 ),
        ('Y', "RNA Cytosine or Uracil"                  ),
        ('K', "RNA Guanine or Uracil"                   ),
        ('V', "RNA Adenine, Cytosine or Guanine"        ),
        ('H', "RNA Adenine, Cytosine or Uracil"         ),
        ('D', "RNA Adenine, Guanine or Uracil"          ),
        ('B', "RNA Cytosine, Guanine or Uracil"         ),
        ('N', "RNA Adenine, Cytosine, Guanine or Uracil")])
    var = Symbol("RNA_", nt)
    @eval begin
        @doc $doc const $var = convert(RNANucleotide, $(UInt8(code - 1)))
        char_to_rna[$(Int(nt + 1))] = char_to_rna[$(Int(lowercase(nt) + 1))] = $var
        rna_to_char[$(code)] = $nt
    end
end

"RNA Gap"
const RNA_Gap = convert(RNANucleotide, 0b1111)
char_to_rna[Int('-') + 1] = RNA_Gap
rna_to_char[0b1111 + 1] = '-'

"Returns Any RNA Nucleotide (RNA_N)"
nnucleotide(::Type{RNANucleotide}) = RNA_N
invalid_nucleotide(::Type{RNANucleotide}) = RNA_INVALID

Base.colon(start::RNANucleotide, stop::RNANucleotide) = SymbolRange(start, stop)

Base.isvalid(::Type{RNANucleotide}, x::Integer) = 0 ≤ x < 16
Base.isvalid(nt::RNANucleotide) = nt ≤ RNA_Gap
isambiguous(nt::RNANucleotide) = nt > RNA_U
alphabet(::Type{RNANucleotide}) = RNA_A:RNA_Gap
gap(::Type{RNANucleotide}) = RNA_Gap

# Compatibility
# -------------

compatbits(nt::Nucleotide) = compatbits_nuc[reinterpret(UInt8, nt)+1]

"""
    iscompatible(x, y)

Return `true` if and only if `x` and `y` are compatible with each other (i.e.
`x` and `y` can be the same symbol).

`x` and `y` must be the same type (`DNANucleotide`, `RNANucleotide` or `AminoAcid`).

# Examples

```julia
julia> iscompatible(DNA_A, DNA_A)
true

julia> iscompatible(DNA_C, DNA_N)  # DNA_N can be DNA_C
true

julia> iscompatible(DNA_C, DNA_R)  # DNA_R (A or G) cannot be DNA_C
false

julia> iscompatible(AA_A, AA_X)    # AA_X can be AA_A
true

```
"""
iscompatible{T}(x::T, y::T) = compatbits(x) & compatbits(y) != 0

# Conversion from Char
# --------------------

function Base.convert(::Type{DNANucleotide}, c::Char)
    @inbounds return c <= '\x7f' ? char_to_dna[Int(c) + 1] : DNA_INVALID
end

function unsafe_ascii_byte_to_nucleotide(T::Type{DNANucleotide}, c::UInt8)
    @inbounds return char_to_dna[c + 1]
end

function Base.convert(::Type{RNANucleotide}, c::Char)
    @inbounds return c <= '\x7f' ? char_to_rna[Int(c) + 1] : RNA_INVALID
end

function unsafe_ascii_byte_to_nucleotide(T::Type{RNANucleotide}, c::UInt8)
    @inbounds return char_to_rna[c + 1]
end

# Conversion to Char
# ------------------

Base.convert(::Type{Char}, nt::DNANucleotide) = dna_to_char[convert(UInt8, nt) + 1]
Base.convert(::Type{Char}, nt::RNANucleotide) = rna_to_char[convert(UInt8, nt) + 1]

# Basic functions
# ---------------

function Base.show(io::IO, nt::DNANucleotide)
    if isvalid(nt)
        if nt == DNA_Gap
            write(io, "DNA_Gap")
        else
            write(io, "DNA_", Char(nt))
        end
    else
        write(io, "Invalid DNA Nucleotide")
    end
    return
end

function Base.show(io::IO, nt::RNANucleotide)
    if isvalid(nt)
        if nt == RNA_Gap
            write(io, "RNA_Gap")
        else
            write(io, "RNA_", Char(nt))
        end
    else
        write(io, "Invalid RNA Nucleotide")
    end
    return
end

function Base.print(io::IO, nt::Nucleotide)
    if !isvalid(nt)
        throw(ArgumentError("nucleotide is invalid"))
    end
    write(io, Char(nt))
    return
end


import ONNXNaiveNASflux: fluxlayers, sources, actfuns, invariantops, pseudotransparentops, optype, nodes
using ONNXNaiveNASflux.NaiveNASflux

# Logging to avoid CI timeouts
@info "  Test padding and sources"

@testset "Read padding" begin
    import ONNXNaiveNASflux: prev

    @test prev(2) == 2
    @test prev([1,2]) == [1,2]
    @test prev([1,2,3,4]) == [2,4,1,3]
    @test prev([1,2,3,4,5,6]) == [3,6,2,5,1,4]
end

@testset "Sources" for tc in
    (
    (name="test_constant", ninputs=0, noutputs=1),
    )

    model, gb, inputs, outputs = prepare_node_test(tc.name, tc.ninputs, tc.noutputs)

    @testset "$(tc.name) op $(optype(node))" for node in nodes(gb)
        @test haskey(sources, optype(node))
        res = sources[optype(node)](node.attribute, params(node)...)

        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]
    end

    @testset "$(tc.name) graph" begin
        cg = load(model)
        res = cg()
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]

        # Also test that it we get the same thing by serializing and then deserializing
        io = PipeBuffer()
        save(io, cg)
        cg = load(io)
        res = cg()
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]
    end
end

# For testing since ONNX states that recurrent layers take 3D input while flux uses
# an Array of 2D Arrays
function (l::Flux.Recur)(x::AbstractArray{T, 3}) where T
    # ONNX shape for RNNs inputs is [seq_length, batch_size, input_size]
    # ONNX.jl reverses this to [input_size, batch_size, seq_length]
    # Unstacking it to a sequence of [input_size, batch_size]
    inseq =Flux.unstack(x, 3)
    out = nothing
    for inpt in inseq
         out = l(inpt)
     end
    # Just to turn it back to ONNX shape.
    # In the testdata only the last output in the sequence is present in the reference
    return reshape(out, size(out)..., 1)
end

@info "  Test Flux layers"

@testset "Fluxlayer $(tc.name)" for tc in
    (
    (name="test_averagepool_1d_default", ninputs=1, noutputs=1),
    # (name="test_averagepool_2d_ceil", ninputs=1, noutputs=1), Not supported!
    (name="test_averagepool_2d_default", ninputs=1, noutputs=1),
    #(name="test_averagepool_2d_pads", ninputs=1, noutputs=1), Not supported!
    (name="test_averagepool_2d_strides", ninputs=1, noutputs=1),
    (name="test_averagepool_3d_default", ninputs=1, noutputs=1),
    (name="test_basic_conv_with_padding", ninputs=2, noutputs=1),
    (name="test_basic_conv_without_padding", ninputs=2, noutputs=1),
    (name="test_batchnorm_epsilon", ninputs=5, noutputs=1),
    (name="test_batchnorm_example", ninputs=5, noutputs=1),
    (name="test_conv_with_strides_and_asymmetric_padding", ninputs=2, noutputs=1),
    (name="test_conv_with_strides_no_padding", ninputs=2, noutputs=1),
    (name="test_conv_with_strides_padding", ninputs=2, noutputs=1),
    (name="test_dropout_default", ninputs=1, noutputs=1),
    (name="test_dropout_random", ninputs=1, noutputs=1),
    #(name="test_gemm_all_attributes", ninputs=3, noutputs=1), Not supported!
    (name="test_gemm_alpha", ninputs=3, noutputs=1),
    (name="test_gemm_beta", ninputs=3, noutputs=1),
    #(name="test_gemm_default_matrix_bias", ninputs=3, noutputs=1), Not supported!
    (name="test_gemm_default_no_bias", ninputs=2, noutputs=1),
    (name="test_gemm_default_scalar_bias", ninputs=3, noutputs=1),
    (name="test_gemm_default_single_elem_vector_bias", ninputs=3, noutputs=1),
    (name="test_gemm_default_vector_bias", ninputs=3, noutputs=1),
    (name="test_gemm_default_zero_bias", ninputs=3, noutputs=1),
    #(name="test_gemm_transposeA", ninputs=3, noutputs=1), Not supported!
    (name="test_gemm_transposeB", ninputs=3, noutputs=1),
    (name="test_lstm_defaults", ninputs=3, noutputs=1),
    (name="test_lstm_with_initial_bias", ninputs=4, noutputs=1),
    # (name="test_lstm_with_peepholes", ninputs=8, noutputs=1), Not supported!
    (name="test_maxpool_1d_default", ninputs=1, noutputs=1),
    #(name="test_maxpool_2d_ceil", ninputs=1, noutputs=1), Not supported!
    (name="test_maxpool_2d_default", ninputs=1, noutputs=1),
    #(name="test_maxpool_2d_dilations", ninputs=1, noutputs=1), Not supported!
    #(name="test_maxpool_2d_pads", ninputs=1, noutputs=1), Not supported!
    (name="test_maxpool_2d_strides", ninputs=1, noutputs=1),
    (name="test_maxpool_3d_default", ninputs=1, noutputs=1),
    (name="test_maxpool_3d_default", ninputs=1, noutputs=1),
    (name="test_rnn_seq_length", ninputs=4, noutputs=1),
    )

    model, gb, inputs, outputs = prepare_node_test(tc.name, tc.ninputs, tc.noutputs)

    @testset "$(tc.name) op $(optype(node))" for node in nodes(gb)
        @test haskey(fluxlayers, optype(node))
        op = fluxlayers[optype(node)](node.attribute, params(node)...)

        res = op(Float32.(inputs[1]))
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]
    end

    @testset "$(tc.name) graph" begin
        cg = load(model)
        res = cg(Float32.(inputs[1]))
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]

        # Also test that it we get the same thing by serializing and then deserializing
        io = PipeBuffer()
        save(io, cg)
        cg = load(io)
        res = cg(Float32.(inputs[1]))
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]
    end
end

@info "  Test Flux activation functions"

@testset "Activation functions $(tc.name)" for tc in
    (
    (name="test_elu", ninputs=1, noutputs=1),
    (name="test_elu_default", ninputs=1, noutputs=1),
    (name="test_elu_example", ninputs=1, noutputs=1),
    (name="test_relu", ninputs=1, noutputs=1),
    (name="test_selu", ninputs=1, noutputs=1),
    (name="test_selu_default", ninputs=1, noutputs=1),
    (name="test_selu_example", ninputs=1, noutputs=1),
    )

    model, gb, inputs, outputs = prepare_node_test(tc.name, tc.ninputs, tc.noutputs)

    @testset "$(tc.name) op $(optype(node))" for node in nodes(gb)
        @test haskey(actfuns, optype(node))
        op = actfuns[optype(node)](node.attribute, params(node)...)
        @test op.(inputs[1]) ≈ outputs[1]

        @test haskey(invariantops, optype(node))
        bcop = invariantops[optype(node)](node.attribute, params(node)...)
        @test bcop(inputs[1]) ≈ outputs[1]

    end
end

@info "  Test stateless ops"

@testset "Invariant op $(tc.name)" for tc in
    (
    (name="test_flatten_axis0", ninputs=1, noutputs=1, fd=pseudotransparentops),
    (name="test_flatten_axis1", ninputs=1, noutputs=1, fd=pseudotransparentops),
    (name="test_flatten_axis2", ninputs=1, noutputs=1, fd=pseudotransparentops),
    (name="test_flatten_axis3", ninputs=1, noutputs=1, fd=pseudotransparentops),
    (name="test_flatten_default_axis", ninputs=1, noutputs=1, fd=pseudotransparentops),
    (name="test_flatten_negative_axis1", ninputs=1, noutputs=1, fd=pseudotransparentops),
    (name="test_flatten_negative_axis2", ninputs=1, noutputs=1, fd=pseudotransparentops),
    (name="test_flatten_negative_axis3", ninputs=1, noutputs=1, fd=pseudotransparentops),
    (name="test_globalaveragepool", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_globalaveragepool_precomputed", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_globalmaxpool", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_globalmaxpool_precomputed", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reduce_mean_default_axes_keepdims_example", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reduce_mean_default_axes_keepdims_random", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reduce_mean_do_not_keepdims_example", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reduce_mean_do_not_keepdims_random", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reduce_mean_keepdims_example", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reduce_mean_keepdims_random", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reduce_mean_negative_axes_keepdims_example", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reduce_mean_negative_axes_keepdims_random", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_reshape_extended_dims", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_reshape_negative_dim", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_reshape_negative_extended_dims", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_reshape_one_dim", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_reshape_reduced_dims", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_reshape_reordered_all_dims", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_reshape_reordered_last_dims", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_reshape_zero_and_negative_dim", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_reshape_zero_dim", ninputs=2, noutputs=1, fd=pseudotransparentops),
    (name="test_softmax_axis_0", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_softmax_axis_1", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_softmax_axis_2", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_softmax_default_axis", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_softmax_example", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_softmax_large_number", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_softmax_negative_axis", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_squeeze", ninputs=1, noutputs=1, fd=invariantops),
    (name="test_squeeze_negative_axes", ninputs=1, noutputs=1, fd=invariantops)
    )

    model, gb, inputs, outputs = prepare_node_test(tc.name, tc.ninputs, tc.noutputs)

    @testset "$(tc.name) op $(optype(node))" for node in nodes(gb)
        @test haskey(tc.fd, optype(node))
        op = tc.fd[optype(node)](node.attribute, params(node)...)
        res = op(inputs[1])
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]
    end

    @testset "$(tc.name) graph" begin
        cg = load(model, size(inputs[1]))
        res = cg(inputs[1])
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]

        # Also test that it we get the same thing by serializing and then deserializing
        io = PipeBuffer()
        save(io, cg)
        cg = load(io, size(inputs[1]))
        res = cg(inputs[1])
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]
    end
end

@testset "Vertex $(tc.name)" for tc in
    (
    (name="test_add", ninputs=2, noutputs=1),
    #(name="test_add_bcast", ninputs=2, noutputs=1), # Op is supported, but we get the wrong idea about what type of inputvertex to create from 3D input
    (name="test_concat_1d_axis_0", ninputs=2, noutputs=1),
    (name="test_concat_1d_axis_negative_1", ninputs=2, noutputs=1),
    (name="test_concat_2d_axis_0", ninputs=2, noutputs=1),
    (name="test_concat_2d_axis_1", ninputs=2, noutputs=1),
    (name="test_concat_2d_axis_negative_1", ninputs=2, noutputs=1),
    (name="test_concat_2d_axis_negative_2", ninputs=2, noutputs=1),
    (name="test_concat_3d_axis_0", ninputs=2, noutputs=1),
    (name="test_concat_3d_axis_1", ninputs=2, noutputs=1),
    (name="test_concat_3d_axis_2", ninputs=2, noutputs=1),
    (name="test_concat_3d_axis_negative_1", ninputs=2, noutputs=1),
    (name="test_concat_3d_axis_negative_2", ninputs=2, noutputs=1),
    (name="test_concat_3d_axis_negative_3", ninputs=2, noutputs=1),
    (name="test_mul", ninputs=2, noutputs=1),
    )

    model, gb, inputs, outputs = prepare_node_test(tc.name, tc.ninputs, tc.noutputs)

    @testset "$(tc.name) graph" begin
        cg = load(model)
        res = cg(inputs[1:length(cg.inputs)]...)
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]

        # Also test that it we get the same thing by serializing and then deserializing
        io = PipeBuffer()
        save(io, cg)
        cg = load(io)
        res = cg(inputs[1:length(cg.inputs)]...)
        @test size(res) == size(outputs[1])
        @test res ≈ outputs[1]
    end
end

@testset "Deserialize with inputs" begin
    using NaiveNASflux: GenericFlux2D, layertype

    function sumgraph()
        ivs = denseinputvertex.(["in1", "in2"], 4)
        g_org = CompGraph(ivs, "out" >> ivs[1] + ivs[2])
        pb = PipeBuffer()
        save(pb, g_org, "in1" => missing, "in2" => missing)
        return pb
    end

    insize(t::Tuple) = ONNXNaiveNASflux.int_size(t[NaiveNASflux.actdim(length(t))])
    insize(p::Pair) = p |> last |> insize
    @testset "Input format $inshapes" for inshapes in (
        ((4,1), (4,1)),
        ("in1" => (4,1), "in2" => (4,1)),
        ((4,missing), (4, :B)),
        ((:I, 3), (:I, 4))
    )
        expsizes =  insize.(inshapes) |> collect
        g_new = if any(==(0), expsizes)
            @test_logs (:warn, r"No valid input sizes") load(sumgraph(), inshapes...)
        else
            load(sumgraph(), inshapes...)
        end
        @test nout.(g_new.inputs) == expsizes
        @test layertype.(g_new.inputs) == [GenericFlux2D(), GenericFlux2D()]
    end

    inshape(t::Tuple) = t |> length |> ONNXNaiveNASflux.guess_layertype
    @testset "Mixshape format $inshapes" for inshapes in (
        ((1,1,5,1), (5,1)),
        ((5,1), (1,1,5,1)),
    )
        g_new = load(sumgraph(), inshapes...)
        @test nout.(g_new.inputs) == [5, 5]
        @test layertype.(g_new.inputs) == inshape.(inshapes |> collect)
    end

    @testset "Malformed input $inshapes" for inshapes in (
        ((4,1), (4,1), (4,1)),
        ("in1" => (4,1), "in2" => (4,1), "in2" => (4,1)),
        ("in1" => (4,1), "notin2" => (4,1))
    )
        @test_throws AssertionError load(sumgraph(), inshapes...)
    end
end

@testset "Deserialize with merging" begin
    function remodel(f, args...)
        pb = PipeBuffer()
        save(pb, f, args...)
        return load(pb, args...)
    end

    @testset "Merge activation function" begin
        m = remodel(Dense(3,4, relu), (3, missing))
        @test nvertices(m) == 2
        @test layer(m.outputs[1]).σ == relu
    end

    @testset "Merge Reshape and $gp" for gp in (
        ONNXNaiveNASflux.globalmeanpool,
        ONNXNaiveNASflux.globalmaxpool
    )
        m = remodel(Chain(
            Conv((3,3), 3 => 3),
            x -> gp(x, xf -> reshape(xf, 3, :)),
        ), (4, 4, 3, missing))

        @test nvertices(m) == 3
    end

    @testset "Merge constant" begin
        m = remodel(x -> relu.(x) .+ (3))
        @test nvertices(m) == 3
        @test m([-1,1]) == [3, 4]
    end
end

@testset "Infer CompGraph shapes" begin
    function remodel(f, args...;kwargs...)
        pb = PipeBuffer()
        save(pb, f, args...)
        return load(pb,; kwargs...)
    end

    @testset "Simple Dense" begin
        v1 = denseinputvertex("v1", 3)
        v2 = fluxvertex("v2", Dense(nout(v1), 4), v1)
        g = remodel(CompGraph(v1, v2), missing)

        @test nout(inputs(g)[]) == nout(v1)
    end

    @testset "After invariant" begin
        v1 = denseinputvertex("v1", 3)
        v2 = invariantvertex("v2", identity, v1)
        v3 = fluxvertex("v3", Dense(nout(v2), 4), v2)
        g = remodel(CompGraph(v1, v3), missing)

        @test nout(inputs(g)[]) == nout(v1)
    end

    @testset "Can't infer after concat" begin
        # I suppose in this case we could infer it as we know nin(v3) = 2 * nout(v1)
        # Seems like too much of an edge case to be worth considering though
        v1 = denseinputvertex("v1", 3)
        v2 = concat("v2", v1, v1)
        v3 = fluxvertex("v3", Dense(nout(v2), 4), v2)
        g = @test_logs (:warn, r"No valid input sizes") remodel(CompGraph(v1, v3), (missing, :B))

        @test nout(inputs(g)[]) == 0
    end

    @testset "Flatten$label" for (label, layerfun, exputilsize) in
        (
        ("", identity, 1),
        ("with ActivationContribution", ActivationContribution, 60),
        )
        using ONNXNaiveNASflux: Flatten, create_vertex_default, defaultutility

        v1 = conv2dinputvertex("v1", 3)
        v2 = fluxvertex("v2", Conv((2,2), 3=>5), v1)
        v3 = absorbvertex("v3", Flatten(-1), v2)

        g = remodel(CompGraph(v1, v3), (5,4,3,:B); vfun=(args...) -> create_vertex_default(args...; layerfun))

        @test nout(inputs(g)[]) == nout(v1)
        @test nout(g[end]) == 60
        @test length(defaultutility(g[end])) == exputilsize
    end
end