using LinearAlgebra: dot
using Plots
function init_position(c::Potentials.Configuration, dim)
    #=
    Initialise the positions for a series
    of particles.

    Parameters
    ----------
    N: int
        Number of particles in the
        system
    dim: int
        Number of dimensions

    Returns
    -------
    ndarray of floats
        Initial velocities for a series of
        particles (eVs/Åamu)
    =#
    N = c.num_atoms
    if dim == 2
        N = N-1
        r0 = 2 .*[cos.(2*pi/N .* (1:N))+1e-2*randn(N) sin.(2*pi/N.*(1:N))+1e-2*randn(N) 0.0 .*(1:N)]
        r0 = [r0; [0 0 0]]
    elseif dim == 3
        if N <= 7
            r0 = [-0.9523699364        0.0159052548       -0.0840802250;
            -0.2528949024        0.8598484003       -0.3332183372;
            -0.3357005052       -0.8500181306        0.2812584200;
            0.7960710440        0.5155096551       -0.1218613860;
            0.7448954577       -0.5412454414        0.2579017563;
            -0.0442130155        0.1953980103        0.5377653956;
            0.0442118576       -0.1953977484       -0.5377656238]
        elseif N == 38
            r0 = [0.163946667699999993 0.319538517000000022 1.71226674359999986;
            1.18570812759999988 -1.12890812809999996 -0.616849616600000039;
            1.47183871269999988 -0.0530362826999999995 0.944308599800000037;
            -0.156601128000000006 1.44806844500000009 -0.969234048100000023;
            1.43644861639999988 0.381943172100000006 -0.92280851779999995;
            -0.659303829199999969 -1.58099572070000005 -0.355825472500000017;
            -1.43644861639999988 -0.381943172100000006 0.92280851779999995;
            0.659303829199999969 1.58099572070000005 0.355825472600000026;
            -1.18570812749999988 1.12890812809999996 0.616849616600000039;
            0.156601128099999987 -1.44806844500000009 0.969234048100000023;
            -1.47183871259999988 0.0530362826999999995 -0.944308599700000029;
            -0.692522445699999989 0.622174532900000021 -1.48126712230000002;
            -0.903125893300000038 0.439946115799999982 1.43235882759999988;
            -0.943262934499999983 -0.888676767300000003 -1.17530822099999988;
            -0.163946667699999993 -0.319538517000000022 -1.71226674359999986;
            0.402594671899999978 -1.70081948270000005 -0.0772747840999999935;
            -1.71903085060000005 0.307018840199999976 0.107299306900000002;
            0.905297373100000025 1.32824468299999987 -0.690683359700000055;
            -0.402594671899999978 1.70081948270000005 0.0772747840999999935;
            1.71903085060000005 -0.307018840199999976 -0.107299306799999994;
            -0.905297373000000016 -1.32824468299999987 0.690683359700000055;
            0.943262934599999991 0.888676767300000003 1.17530822109999988;
            0.692522445799999997 -0.622174532900000021 1.48126712230000002;
            0.903125893400000046 -0.439946115799999982 -1.43235882749999988;
            -1.18439743100000006 -0.634064735699999993 -0.125267553799999987;
            0.120331581600000001 -1.00573849309999996 -0.891368452299999947;
            0.934263338099999952 -0.873132691499999947 0.430486517999999985;
            1.18439743110000006 0.634064735699999993 0.125267553799999987;
            -0.934263338000000054 0.873132691499999947 -0.430486517899999976;
            -0.370465674600000017 -0.501458934099999976 1.19658741639999988;
            0.370465674700000025 0.501458934099999976 -1.19658741639999988;
            -0.120331581600000001 1.00573849309999996 0.891368452299999947;
            0.649793600800000037 -0.185104513500000012 -0.381540884400000002;
            -0.405362064899999985 -0.0660416075000000019 -0.658322833699999999;
            0.124574169400000004 0.750628853799999995 -0.15200806280000001;
            -0.649793600800000037 0.185104513500000012 0.381540884400000002;
            0.405362064999999994 0.0660416075000000019 0.658322833699999999;
            -0.124574169299999996 -0.750628853799999995 0.15200806280000001]
        else
            r0 = rand(N, d)
        end

        r0 += 1e-2*randn(size(r0))
        com = mean(r0, dims = 1)
        r0 = r0 .- com
    else
        return false
    end

    for j = 1:c.num_atoms
        c.Positions[j] = Potentials.Position( r0[j, 1], r0[j, 2], r0[j, 3], c.Positions[j].type)
    end

    return c
end

function init_velocity(c::Potentials.Configuration, dim::Int, T::Float64)
    #=
    Initialise the velocities for a series
    of particles.

    Parameters
    ----------
    T: float
        Temperature of the system at
        initialisation (K)
    number_of_particles: int
        Number of particles in the
        system

    Returns
    -------
    ndarray of floats
        Initial velocities for a series of
        particles (eVs/Åamu)
    =#
    N = c.num_atoms
    r0 = randn(N, 3) .- 0.5
    if dim == 2
        r0[:, 3] .= 0.0
    end
    r0 = r0 * T
    for j = 1:c.num_atoms
        c.Velocities[j] = Potentials.Position( r0[j, 1], r0[j, 2], r0[j, 3], c.Positions[j].type)
    end
    return c
end

function update_velo(v::Vector, a::Vector, a1::Vector, dt::Float64)
    #=
    Update the particle velocities.

    Parameters
    ----------
    v: ndarray of floats
        The velocities of the particles in a
        single dimension (eVs/Åamu)
    a: ndarray of floats
        The accelerations of the particles in a
        single dimension at the previous
        timestep (eV/Åamu)
    a1: ndarray of floats
        The accelerations of the particles in a
        single dimension at the current
        timestep (eV/Åamu)
    dt: float
        The timestep length (s)

    Returns
    -------
    ndarray of floats:
        New velocities of the particles in a
        single dimension (eVs/Åamu)
    =#
    return v + 0.5 * (a + a1) * dt
end

function verlet(p::Potentials.ArbitraryPotential, c::Potentials.Configuration, dt::Float64, Temp:: Float64; γ = 1.0, k_b = 1.0/119.8)
   #=
    Update the particle positions and velocities by solving.
    dx/dt = dv/dt
    dv/dt = F(r)/m - γ * v + R/m = a

    Parameters
    ----------
    x : ndarray of floats
        The positions of the particles (Å)
    v: ndarray of floats
        The velocities of the particles
        single dimension (eVs/Åamu)
    dt: float
        The timestep length (s)

    Returns
    -------
    ndarray of floats:
        New velocities and positions of the particles in a
        single dimension (Å, eVs/Åamu)
    =#
    pos = vec(c.Positions)
    vels = vec(c.Velocities)
    R = sqrt(2*Temp*k_b*dt*γ)
    a = Potentials.force(c, p)
    a = a - γ .* vels
    noise = 0.5*R*randn(c.num_atoms, 3)
    if pos[1][3] == 0.0
        for j = 1:length(a)
            a[j][3] = 0.0
            noise[j, 3] = 0.0
        end
    end
    v_star = update_velo(vels, a, 0.0*a, dt)
    noise = noise .- mean( noise, dims = 1 )
    noise = [[noise[i, j] for j = 1:3] for i = 1:c.num_atoms]
    v_star += 0.5*noise
    # Second half-time step
    x_new = c.Positions + v_star*dt
    # println("Old Positions ")
    # show(stdout, "text/plain", c.Positions)
    # println(" ")
    # println("New positions ")
    # show(stdout, "text/plain", x_new)
    # println(" ")
    for (i, xi) in enumerate(x_new)
        c.Positions[i] = xi
    end
    a_new = Potentials.force(c, p)
    a_new = a_new - γ .* vels
    noise = 0.5*R*randn(c.num_atoms, 3)
    if pos[1][3] == 0.0
        for j = 1:length(a)
            a_new[j][3] = 0.0
            noise[j, 3] = 0.0
        end
    end
    noise = noise .- mean( noise, dims = 1 )
    noise = [[noise[i, j] for j = 1:3] for i = 1:c.num_atoms]
    v_new = update_velo(v_star, a, a_new, dt)
    v_new += 0.5*noise

    # Save
    for (i, (xi, vi)) in enumerate(zip(x_new, v_new))
        c.Positions[i] = xi
        c.Velocities[i] = Potentials.Position(vi[1], vi[2], vi[3], c.Positions[i].type)
    end
    return c
end

function solve(p::Potentials.ArbitraryPotential, c0::Potentials.Configuration, dt::Float64, Nt::Int, Temp ::Array{Float64}; save_dt = 1 :: Int)
    N = Int(Nt / save_dt)
    r = Vector{Potentials.Configuration}(undef, N)
    r = fill!(r, c0)

    t = zeros( (N, 1));

    τ = 0.0
    count = 1
    c_temp = deepcopy(c0)
    for j = 2:(Nt)
        τ +=  dt
        # println("Time ", τ)
        c_temp = verlet(p, c_temp, dt, Temp[j-1])
        if (j-1) % save_dt == 0
            count += 1
            r[count] = deepcopy(c_temp)
            t[count] = τ
        end
    end

    return r, t
end

function get_positions(r::Vector{PotentialUQ.Potentials.Configuration})
    N = length(r)
    num_atoms = r[1].num_atoms

    pos = Vector{Vector{Float64}}(undef, N)
    for (i, ri) in enumerate(r)
        pos[i] = PotentialUQ.Potentials.norm.(ri.Positions)
    end
    return pos
end

function get_velocities(r::Vector{Potentials.Configuration})
    N = length(r)
    num_atoms = r[1].num_atoms

    pos = Vector{Vector{Float64}}(undef, N)
    for (i, ri) in enumerate(r)
        pos[i] = Potentials.norm.(ri.Velocities)
    end
    return pos
end

function calculate_rdf(r::Vector{Potentials.Configuration}; L = 2.5, maxbin = 1000)

    bins = vcat(range(0, L; length = maxbin+1)...)
    rdf = zeros(maxbin)
    bin_centers = 0.5 .* (bins[1:end-1] + bins[2:end])
    n = length(r)
    num_atoms = r[1].num_atoms
    dr = bin_centers[2] - bin_centers[1]
    hist = zeros(maxbin)
    for c ∈ r
        distances = Potentials.get_interparticle_distance(c)
        for d in distances
            bin = ceil(Int, d / dr)
            if (bin > 1) && (bin < maxbin)
                hist[bin] += 1
            end
        end
    end

    volume = 4 / 3 * π * L^3
    density = num_atoms / volume
    rdf = zeros(maxbin)
    for bin = 1:maxbin
        rlower = bins[bin]
        rupper = bins[bin+1]
        volume_of_shell = 4/3 * π * (rupper^3 - rlower^3)
        rdf[bin] = (hist[bin] / n) / volume_of_shell / density
    end

    return (bin_centers, rdf)
end

function calculate_mean_squared_displacement(r::Vector{Potentials.Configuration})
    n = length(r)
    num_atoms = r[1].num_atoms
    p0 = vec(r[1].Positions)
    disp = zeros(n)
    for i = 1:n
        for j = 1:num_atoms
            p = vec(r[i].Positions)
            dr = p - p0
            disp[i] += dot(dr, dr) / num_atoms
        end
    end
    return disp
end

function calculate_mass_diffusivity(r::Vector{Potentials.Configuration}, t)
    disp = calculate_mean_squared_displacement(r)
    D = sum( disp[end-100:end] ./ (6 .* t[end-100:end] ) ) / 100
    return D
end

function animate_atoms(r::Vector{Potentials.Configuration}, t; dim = 2, fps = 60, dT = 20, title = "lennard_jones_cluster.gif")
    n = length(r)
    num_atoms = r[1].num_atoms

    if dim == 3

        anim = @animate for i = 1:dT:n
            pos = hcat(vec(r[i].Positions)...)
            scatter(pos[1, :], pos[2, :], pos[3, :], xaxis = ("x", (-2, 2), -2:0.5:2),
                            yaxis = ("y", (-2, 2), -2:0.5:2),
                            zaxis = ("z", (-2, 2), -2:0.5:2),
                            markersize = 20,
                            c=colormap("Blues",num_atoms),
                            legend = false,
                            title = string("τ = ", floor(t[i]) ))
        end

    elseif dim==2
        anim = @animate for i = 1:dT:n
            pos = hcat(vec(r[i].Positions)...)
            scatter(pos[1, :], pos[2, :], xaxis = ("x", (-4, 4), -4:0.5:4),
                            yaxis = ("y", (-2, 2), -4:0.5:4),
                            c=colormap("Blues",num_atoms),
                            legend = false,
                            title = string("τ = ", floor(t[i]) ) )

        end
    end

    gif(anim, title, fps = fps)
end


# This is the file that executes on tasks, aka where the magic happens

function run_task(task)
    organization = fetch_organization(task)
    schedule = fetch_schedule(task)

    if organization["plan"] == "per-seat-v1"
        workers_raw = fetch_workers(task)
        api_availability = fetch_availability(task)
        workers = fill_missing_availability(schedule, api_availability, workers_raw, organization, task)
        workers = downgrade_availability(schedule, workers, organization, task)
    else  # Unassigned shifts mode
        workers = Dict()
        api_availability = Dict()
    end

    shifts = run_calculation(schedule, workers, organization)
    set_shifts(shifts, task, organization)
end

function run_calculation(schedule, workers_input, organization)
    # Do some merging into StaffJoy.jl format and run calculation

    env = deepcopy(schedule)

    # build coverage
    coverage = Array[]
    for day in days(organization)
        push!(coverage, schedule["demand"][day])
    end
    env["coverage"] = coverage

    env["shift_time_min"] = organization["min_shift_length"]
    env["shift_time_max"] = organization["max_shift_length"]
    env["intershift"] = organization["hours_between_shifts"]

    # time between coverage is 0 in manager by definition.
    # Windowing happens at scheduler level.
    env["time_between_coverage"] = 0

    workers = Dict()
    for worker in workers_input
        w = Dict()
        w["hours_min"] = int(worker["min_hours_per_week"])
        w["hours_max"] = int(worker["max_hours_per_week"])
        w["shift_time_min"] = int(worker["min_shift_length"])
        w["shift_time_max"] = int(worker["max_shift_length"])
        w["shift_count_min"] = int(worker["min_shifts_per_week"])
        w["shift_count_max"] = int(worker["max_shifts_per_week"])
        # Offset from absolute time to the 4AM thing
        if "no_shifts_after" in keys(worker)
            w["no_shifts_after"] = int(worker["no_shifts_after"]) - TIME_OFFSET
        end

        # Convert availability to array of arrays - no day of week
        w["availability"] = Array[]
        for day in days(organization)
            push!(w["availability"], worker["availability"][day])
        end

        workers[string(worker["id"])] = w
    end
    ok, shifts = StaffJoy.schedule(workers, env)
    if !ok
        Logging.err("Schedule failed")
        error()
    end
    return shifts
end


# Scaling for numerical stability - may not be needed
scale(x) = x / 10000

get_reference(S::TensorMeta, node::XNode) = S.reference_map[node]
graph(S::TensorMeta) = S.graph
Profiler.users(S::TensorMeta) = S.users

# Accessor methods
rb(I::ILPHolder) = I.read_bandwidth
wb(I::ILPHolder) = I.write_bandwidth
rba(I::ILPHolder) = I.read_bandwidth_async
wba(I::ILPHolder) = I.write_bandwidth_async

# Common Methods

# Length check because ngraph compilation is not 100 % consistent and can sometimes have
# a few more nodes than it began with ...
limit(F::Frame, args...) = limit(F.modeltype, args...)
limit(S::ILPHolder, i) = i > length(S.dram_limits) ? maxlimit(S) : S.dram_limits[i]
maxlimit(S::ILPHolder) = maximum(S.dram_limits)

predict(F::Frame) = objective_value(F.model)
descriptor(F::Frame, tensor::XTensor) = F.modeltype.descriptors[tensor]

#####
##### Entry Point
#####

# For procedurally building JuMP expressions for the ILP model
const AffExprType = typeof(AffExpr())

function create_model(modeltype::ILPHolder, profile_data::FunctionData)
    preprocess!(modeltype, profile_data)

    # Start with an empty model that we will progressively build.
    model = Model(AutoTM.getoptimizer())
    frame = Frame(modeltype, model, profile_data)

    # Going deep into JuMP here - the idea is to build the objective as a bunch of aff exprs
    # and eventually combine all of them together.
    model[:node_times] = Dict{String, AffExprType}()
    model[:tensor_async] = Dict{String, AffExprType}()
    model[:tensor_sync] = AffExprType()

    add_tensors!(frame)
    add_nodes!(frame)
    add_constraints!(frame)

    # Default objective is to just sum all of the node times
    objective_expr = model[:tensor_sync]
    for (node_name, node_times) in model[:node_times]
        # Check to see if there are overlapping async transfers.
        #
        # If so, take the max of the sum of the overlapping transfers and the node time.
        _async = get(model[:tensor_async], node_name, nothing)
        if isnothing(_async)
            add_to_expression!(objective_expr, node_times)
        else
            var = @variable(model, lower_bound = 0.0)
            @constraint(model, var >= node_times)
            @constraint(model, var >= _async)
            add_to_expression!(objective_expr, var)
        end
    end
    # Quick optimization to remove zero terms
    drop_zeros!(objective_expr)
    @objective(frame.model, Min, objective_expr)

    return frame
end

#####
##### Adding Tensors
#####

isfixed(frame::Frame, tensor) = descriptor(frame, tensor).isfixed
#isfixed(frame::Frame, tensor) = false
free_tensors(frame::Frame) = [t for t in tensors(frame.profile_data) if !isfixed(frame, t)]
fixed_tensors(frame::Frame) = [t for t in tensors(frame.profile_data) if isfixed(frame, t)]

function add_tensors!(frame::Frame)
    data = frame.profile_data
    modeltype = frame.modeltype

    # Create variables for the tensors and add flow constraints to the to the tensor graphs
    @variable(frame.model,
        tensor_graphs[
            tensor = free_tensors(frame),
            e = edges(graph(descriptor(frame, tensor)))
        ],
        Bin
    )

    @showprogress 1 "Creating Flow Formulation " for tensor in free_tensors(frame)
        g = graph(descriptor(frame, tensor))
        # Iterate through nodes in the graph - generating constraints based on the type
        # of node.
        for v in vertices(g)
            # Set flow coming out of the source node
            if getmeta(g, v).location == LOC_SOURCE
                @constraint(frame.model,
                    sum(tensor_graphs[tensor, e] for e in outedges(g, v)) == 1
                )

            # Set flow going into the sink node
            elseif getmeta(g, v).location == LOC_SINK
                @constraint(frame.model,
                    sum(tensor_graphs[tensor, e] for e in inedges(g, v)) == 1
                )

            # All other ops must conserve flow
            else
                oe = collect(outedges(g, v))
                ie = collect(inedges(g, v))
                @constraint(frame.model,
                    sum(tensor_graphs[tensor, e] for e in oe) -
                    sum(tensor_graphs[tensor, e] for e in ie) == 0
                )
            end
        end
    end

    #####
    ##### Add objective penalty for moving data
    #####

    add_movement_formulations!(frame)

    #####
    ##### Create variables to determine if a tensor is in DRAM.
    #####

    @variable(frame.model,
        tensor_in_dram[
            tensor = free_tensors(frame),
            user = nGraph.name.(users(descriptor(frame, tensor)))
        ],
        Bin
    )

    @variable(frame.model,
        tensor_in_dram_post[
            tensor = free_tensors(frame),
            user = nGraph.name.(users(descriptor(frame, tensor)))
        ],
        Bin
    )

    # A tensor in DRAM is live if any of its incoming edges are used.
    for tensor in free_tensors(frame)
        desc = descriptor(frame, tensor)
        g = graph(desc)

        # Create a container for the critical edge from LOC_SOURCE to LOC_SINK if this is
        # an argument tensor
        skip_edge = []
        if isarg(tensor)
            # Find the source to sink edge.
            # Will error if this edge doesn't exist, so serves as a form of error checking
            edge = find_edge(
                (g, e) -> getmeta(g, src(e)).location == LOC_SOURCE &&
                          getmeta(g, dst(e)).location == LOC_SINK,
                g
            )
            push!(skip_edge, edge)
        end

        for user in users(desc)
            # Get the DRAM for this op.
            verts = filter(
                v -> isdram(getmeta(g,v).location) && getmeta(g,v).op == user,
                vertices(g)
            )

            # Map `inedges` to `vertex_iter` and iterats over all those edges
            _iter = vflatten(Iterators.flatten(inedges.(Ref(g), verts)), skip_edge)
            for e in _iter
                @constraint(
                    frame.model,
                    tensor_in_dram[tensor, nGraph.name(user)] >= tensor_graphs[tensor, e]
                )
            end

            # If all incoming edges are not taken, tensor MUST not be in DRAM.
            @constraint(frame.model,
                sum(tensor_graphs[tensor, e] for e in _iter) >=
                    tensor_in_dram[tensor, nGraph.name(user)]
            )

            # Similary, set the post DRAM constraints
            _edges = filter(
                e -> (isdram(getmeta(g, src(e)).location) &&
                    # Need to check "LOC_SINK" for the static case
                    (
                        isdram(getmeta(g, dst(e)).location) ||
                        getmeta(g, dst(e)).location == LOC_SINK
                    ) && getmeta(g, src(e)).op == user),
                collect(edges(g))
            )
            _iter = vflatten(_edges, skip_edge)
            for e in _iter
                @constraint(
                    frame.model,
                    tensor_in_dram_post[tensor, nGraph.name(user)] >= tensor_graphs[tensor, e]
                )
            end

            @constraint(frame.model,
                sum(tensor_graphs[tensor, e] for e in _iter) >=
                    tensor_in_dram_post[tensor, nGraph.name(user)]
            )
        end
    end
    return nothing
end

# Filter on edge type, sort by parent index to get the edges in execution order.
_find_edges(g, edgetype) = sort(
    filter(e -> getmeta(g, e).edgetype == edgetype, collect(edges(g))),
    by = src
)

# Formulation specific move node stuff
add_movement_formulations!(frame::Frame{ILPHolder{Static}}) = nothing
function add_movement_formulations!(frame::Frame)
    # Unpack variables
    data = frame.profile_data
    modeltype = frame.modeltype

    tensor_sync_expr = frame.model[:tensor_sync]
    tensor_async_dict = frame.model[:tensor_async]
    tensor_graphs = frame.model[:tensor_graphs]

    # A tensor is written to dram if:
    # - It was not created into PMEM
    # - Any edge from DRAM to PMEM is taken
    #
    # NOTE: We only pay the write cost once.
    @variable(frame.model, tensor_write[tensor = free_tensors(frame)], Bin)

    # Add objective terms for all read ops
    for tensor in free_tensors(frame)
        # Some unpacking
        g = graph(descriptor(frame, tensor))
        bytes = sizeof(tensor)

        # Take the ceiling of all these to ensure there's always a cost to moving.
        read_cost        = scale(bytes / rb(modeltype))
        write_cost       = scale(bytes / wb(modeltype))
        read_cost_async  = scale(bytes / rba(modeltype))
        write_cost_async = scale(bytes / wba(modeltype))

        # Collect Edges according to type.
        sync_reads = _find_edges(g, EDGE_SYNC_READ)
        sync_writes = _find_edges(g, EDGE_SYNC_WRITE)
        async_reads = _find_edges(g, EDGE_ASYNC_READ)
        async_writes = _find_edges(g, EDGE_ASYNC_WRITE)

        #####
        ##### Constraints for sync write variable
        #####

        # No Sync write if any async write
        for e in sync_writes
            @constraint(
                frame.model,
                tensor_write[tensor] >= tensor_graphs[tensor, e]
            )
        end

        # No sync write if no edges taken
        @constraint(
            frame.model,
            tensor_write[tensor] <= sum(tensor_graphs[tensor, e] for e in sync_writes)
        )

        #####
        ##### Constraints on async write variables
        #####

        # Assign each edge to the kernel it overlaps with.
        #
        # Just go overkill and grab all the edges, even though we end up only using a subset.
        kernels = Dict(e => getmeta(g,src(e)).op for e in edges(g))

        # Create read variables expressions
        for e in async_reads
            _expr = get!(tensor_async_dict, nGraph.name(kernels[e]), AffExprType())
            move_var = tensor_graphs[tensor, e]
            add_to_expression!(_expr, read_cost_async, move_var)
            dict_push!(frame.modeltype.async_move_vars, kernels[e], move_var)
        end

        # Create write variables.
        for e in async_writes
            _expr = get!(tensor_async_dict, nGraph.name(kernels[e]), AffExprType())
            move_var = tensor_graphs[tensor, e]
            add_to_expression!(_expr, write_cost_async, move_var)
            dict_push!(frame.modeltype.async_move_vars, kernels[e], move_var)
        end

        #####
        ##### Finally, add all the synchonous move costs.
        #####
        for e in sync_reads
            add_to_expression!(tensor_sync_expr, read_cost, tensor_graphs[tensor, e])
        end
        add_to_expression!(tensor_sync_expr, write_cost, tensor_write[tensor])
    end
    return nothing
end

# There's an issue when trying to reference whether or not a tensor is in DRAM.
#
# If we're on an op where the tensor is used, we have to look at the inputs to a
# graph verted with LOC_DRAM or LOC_PREAD to see if the tensor was fetched or already
# lived in dram.
#
# If we're on an op where a tensor is LIVE but not READ, we need to check the outgoing
# edge of the correct DRAM -> DRAM node to see if the tensor just lives around in DRAM.
function get_tensor_in_dram(F::Frame, tensor::XTensor, node::XNode)
    # First - check if the node is fixed. If it is fixed, then we can returna constant
    # depending on its location.
    if isfixed(F, tensor)
        location = collect(locations(tensor))
        if length(location) != 1
            @show location
            @show tensor
            @show node
            error()
        end
        if first(location) == DRAM
            return 1
        else
            return 0
        end
    end

    # Otherwise, look through the generated model to find the variable that will indicate
    # if this tensor is in DRAM or not.
    desc = descriptor(F, tensor)
    if in(node, users(desc))
        return F.model[:tensor_in_dram][tensor, nGraph.name(node)]
    else
        return F.model[:tensor_in_dram_post][tensor, nGraph.name(get_reference(desc, node))]
    end
end

function add_nodes!(F::Frame)
    data = F.profile_data

    # Create decision variables for all nodes that have a choice of backend algorithm.
    select_nodes = filter(x -> hasprofile(x) && can_select_algo(x), nodes(data))
    if !isempty(select_nodes)
        @info "Creating Algorithms Variables"
        @variable(
            F.model,
            algo_var[
                node = select_nodes,
                enum = enums(gettime(node))
            ],
            Bin
        )

        # Constrain so only one algorithm may be selected.
        for node in select_nodes
            @constraint(
                F.model,
                sum(algo_var[node, e] for e in enums(gettime(node))) == 1
            )
        end
    end

    for node in nodes(data)
        # We don't profile all ops, so perform a quick check to see if this is an op
        # the we have profile information for. If not, there's nothing to do as far as the
        # ILP model is concerned.
        hasprofile(node) || continue

        # The GPU path of this code will just return an all DRAM config - which will be
        # useful for generating the constraint that all kernel IO for the GPU case must
        # reside in GPU DRAM.
        #
        # The CPU path will yield a bunch of DRAM/PMEM combinations
        configs = configs_for(node)

        # Create a variable for each config.
        vars = @variable(F.model, [config = configs], Bin)

        for config in configs
            # Create an expression for the input and output locations
            expr = AffExprType()
            iter = Iterators.flatten((
                zip(config.inputs, inputs(node)),
                zip(config.outputs, outputs(node))
            ))

            for (location, tensor) in iter
                # use `jump_tensor` because it's really a JuMP variable that is returned
                # by this call.
                jump_tensor = get_tensor_in_dram(F, tensor, node)
                if location == DRAM
                    # This branch handles the jump_tensor being either an Int or a
                    # VariableRef just fine.
                    add_to_expression!(expr, jump_tensor)
                    @constraint(F.model, vars[config] <= jump_tensor)
                else
                    # The return value from `get_tensor_in_dram` can be a literal value,
                    # so make sure we handle this case.
                    if isa(jump_tensor, Int )
                        add_to_expression!(expr, jump_tensor)
                    else
                        add_to_expression!(expr, 1)
                        add_to_expression!(expr, -1, jump_tensor)
                    end
                    @constraint(F.model, vars[config] <= 1 - jump_tensor)
                end
            end

            @constraint(
                F.model,
                vars[config] + length(config.inputs) + length(config.outputs) >= 1 + expr
            )
        end

        # Add a valid contraint to help the solver
        @constraint(F.model, sum(vars[config] for config in configs) == 1)

        # Create an expression for this node's expected running time.
        node_times = AffExprType()
        for config in configs
            # If we can select the algorithm for this node, we need to generate some more
            # variables to "AND" the config with the algorithm selection to ensure that
            # we only get a single algorithm out at the end.
            #
            # If there are not multiple algorithms, then we don't have to worry about it.
            if can_select_algo(node, config)
                v = @variable(F.model, [enum = enums(gettime(node, config))], Bin)
                for enum in enums(gettime(node, config))
                    @constraint(F.model, v[enum] <= algo_var[node, enum])
                    @constraint(F.model, v[enum] <= vars[config])
                    @constraint(F.model, v[enum] + 1 >= vars[config] + algo_var[node, enum])

                    coeff = scale(Profiler.timeat(gettime(node, config), enum))
                    add_to_expression!(node_times, coeff, v[enum])
                end
            else
                coeff = scale(gettime(node, config))
                add_to_expression!(node_times, coeff, vars[config])
            end
        end
        F.model[:node_times][nGraph.name(node)] = node_times
    end
    return nothing
end

# Allocations in ngraph happen on 4096 bytes boundaries. For better accuracty, round
# up to the nearest multiple of 4096 before figuring out the number of bytes.
#
# Take the floor to introduce more zeros into the ILP formulation. This shouldn't really
# make much of a difference.
tensor_size(t::XTensor) = tensor_size(sizeof(t))
tensor_size(sz) = ceil(Int, ceil(Int, sz / 4096) * 4096 / 1E6)

function add_constraints!(F::Frame)
    # Unpack some variables
    data = F.profile_data

    for (index, tensors) in enumerate(live_tensors(data))
        node = nodes(data)[index]
        hasprofile(node) || continue
        F.modeltype.node_to_limit_index[nGraph.name(node)] = index

        # Add DRAM constraint for the workspace
        if can_select_algo(node)
            v = F.model[:algo_var]
            algo_expr = @expression(
                F.model,
                sum(
                    tensor_size(Profiler.bytesat(gettime(node), e)) *
                    v[node, e] for e in enums(gettime(node))
                )
            )
        else
            # If we can't select the algorithm, just create an empty expression that will
            # be hapily optimized away when we create the size constraint.
            algo_expr = AffExprType()
        end

        if !isempty(tensors)
            @constraint(F.model,
                algo_expr + sum(tensor_size(t) * get_tensor_in_dram(F, t, node)
                    for t in tensors
                    if !iszero(tensor_size(t))) <= limit(F, index)
            )
        end
    end

    return nothing
end

#####
##### Queries
#####

function find_shortcut_edge(frame, tensor)
    g = descriptor(frame, tensor).graph
    edge = find_edge(
        (g, e) -> getmeta(g, src(e)).location == LOC_SOURCE && getmeta(g, dst(e)).location == LOC_SINK,
        g
    )
    return edge
end

# Check if a given tensor is:
# 1. A function argument (nGraph function input or output)
# 2. Assigned to DRAM for the lifetime of the function
function islocalarg(frame, tensor::XTensor)
    # Determine if it's an argument
    isarg(tensor) || return false

    # Next - figure out if it was fixed in a memory pool.
    if isfixed(frame, tensor)
        location = first(locations(tensor))
        return location == DRAM
    end

    # Find the critical edge from LOC_SOURCE to LOC_SINK
    # This tensor is assigned to the local memory if this edge is taken (has a solved value
    # of one)
    edge = find_shortcut_edge(frame, tensor)

    return approx_one(frame.model[:tensor_graphs][tensor, edge])
end


using PyCall,DataFrames,RCall
unshift!(PyVector(pyimport("sys")["path"]), "")
@pyimport ncdata
len = length

function sorteqn(x,delim)
    r,p = split(string(x),delim)
    r = join(sort(split(r,'+')),'+')
    p = join(sort(split(p,'+')),'+')
    return join([r,p],delim)
end

organics = readtable("./mcm_3_3_1_reactions.csv", separator = ',', header = true)

reactions = map((x) -> vcat([join(matchall(r"[^\h]*",organics[x,4]),""),organics[x,2]]), 1:len(organics[1]))

organics[:reaction]=[sorteqn(i[1],"=") for i in reactions]

groups  =  [match(r"RO2NO\d?|RO2HO2|RO2|NO\d?|HO2|PAN|J",i[2]) for i in reactions]
allocations = [(r"\b(OH)\b","OH"),(r"\b(CH3O|HOCH2CH2O)\b","ALKOXYL"),(r"\b(HCOCO)\b","GLYOXL"),(r"\b(PPN)\b","PAN")]#[[regex,group]]

for i in 1:length(groups)
    groups[i] == nothing ? nothing : continue;
    rxn = split(reactions[i][1],'=')[1]
    groups[i] = match(r"\b(NO\d?|HO2)\b",rxn)

    if groups[i] == nothing

        for aloc in allocations
            if (ismatch( aloc[1] , rxn))
                groups[i] = aloc[2]
                @goto skip
            end
        end

        for j in split(rxn,"+")
            try
                if (ismatch(r"(\[O\-\]\[O\+\]\=C|\[C\]O\[O\]|\[O\]O\[C\])",mcm3[mcm3[:name].==j,:][:smiles][1]) )
                    groups[i] = "CRI"
                    @goto skip
                end
            catch
                groups[i] = "INORGANIC"
                @goto skip
            end
        end
        #print(reactions[i][1],groups[i],"\n")
        groups[i]="DISSOCIATION"
        @label skip
    end
end

function rmrx(i)
    try
        return i.match
    end
    return i
end

organics[Symbol("groups")]= [rmrx(i) for i in groups] #remove RegexMatch objects

allgroups = Set(organics[:groups])

###################
###################

eqnmatch = map((x)->ismatch(r"RO2|RO2NO2|RO2NO|RO2HO2|HO2",x),organics[:groups])

organics[Symbol("keep")]=eqnmatch

selection = organics[organics[:keep].==true,:]

matchingspecies = Set(matchall(r"\b\w+\b",join(selection[:reaction]," ")))

filename = "./edwards11.nc"

data = ncdata.get(filename)
specs = names!(DataFrame(data["spec"]),[Symbol(i)for i in data["sc"]])
rates = names!(DataFrame(data["rate"])[9:len(data["rc"])],[Symbol(i)for i in data["rc"][9:len(data["rc"])]])
#rates = rates[:,[Symbol(i) for i in filter(x->!ismatch(r"EMISS|DUMMY",x),data["rc"])]]
specs = specs[:,[Symbol(i) for i in filter(x->!ismatch(r"EMISS|DUMMY",x),data["sc"])]]
rates = rates[:,[Symbol(i) for i in filter(x->sum(rates[x])>0,names(rates))]]
joint = "\n"*(join(names(rates),"\n"))*"\n"

selects = intersect(Set(names(specs)),Set([Symbol(x) for x in matchingspecies]))

rtype = map((x)-> organics[organics[:reaction].== replace(sorteqn(string(x),"-->"),"-->","="),:groups],names(rates))

products= [split(i[4:len(i)],"+") for i in matchall(r"-->([A-z0-9+]*)",joint)]
reactants = [split(i[2:len(i)-1],"+") for i in matchall(r"\n([A-z0-9+]{1,60})([-->]{0,1})",joint)]

node_set = Set([ i[2:len(i)-1] for i in matchall(r"[\n+>-]{1}[A-z0-9]+[\n+>-]{1}",joint) ])
nodes = [i for i in node_set]
dict = Dict( [[i,x] for (x,i) in enumerate(nodes)])

fluxtype = []
flux = []
edges = []
for i in 1:length(reactants)
  rcol=[]
  #println(i)
  for j in reactants[i]
    coeff = match( r"([\d\s\.]*)(\D[\d\D]*)", j)
    dummy = specs[Symbol(coeff[2])]
    try
      push!(rcol,parse(Float32,coeff[1]))
    catch
      push!(rcol,1.0f0)
    end

    prod = 1
    for k in rcol
      prod*=k
    end

    push!(fluxtype, rtype[i])
    push!(flux, prod*rates[i])
    # make edges array
    for l in products[i]
        push!(edges,[match( r"([\d\s\.]*)(\D[\d\D]*)", j)[2],match( r"([\d\s\.]*)(\D[\d\D]*)", l)[2],i, fluxtype[i]])
    end

  end
end

  function smrxn(x,t,spcsel)
    loss = [flux[i[3]][t] for i in x if i[1]==spcsel]
    prod = [flux[i[3]][t] for i in x if i[2]==spcsel]

    #println(prod,loss)

    if len(loss)!=0
      loss = sum(loss)
    else
      loss = 0
    end

    if len(prod)!=0
      prod = sum(prod)
    else
      prod = 0
    end
    return(abs(prod-loss))
  end

  spcsel = "CH3COCH2O2"#"CH3O2"

newedges =  filter((x)-> ((Symbol(x[1]) in selects)&(  Symbol(x[2]) in selects)) , edges)
newedges =  filter((x)-> (((x[1]) == spcsel)|(  (x[2]) == spcsel)) , newedges)

RO2HO2 = filter((x)-> x[4]==AbstractString["RO2HO2"] , newedges)
RO2NO2 = filter((x)-> x[4]==AbstractString["RO2NO2"] , newedges)
RO2NO = filter((x)-> x[4]==AbstractString["RO2NO"] , newedges)
RO2 = filter((x)-> x[4]==AbstractString["RO2"] , newedges)
HO2 = filter((x)-> x[4]==AbstractString["RO2HO2"] , newedges)
NO2 =  filter((x)-> x[4]==AbstractString["NO2"] , newedges)
OH= filter((x)-> x[4]==AbstractString["OH"] , newedges)
NO3= filter((x)-> x[4]==AbstractString["NO3"] , newedges)
J= filter((x)-> x[4]==AbstractString["J"] , newedges)
NO = filter((x)-> x[4]==AbstractString["NO"] , newedges)
RO2NO3= filter((x)-> x[4]==AbstractString["RO2NO3"] , newedges)

file = open("results.csv", "w")
var = Float64(89)
typevar = string(typeof(var), "\n")

write(file,"t,ts,RO2HO2,RO2NO,RO2NO2,RO2,HO2,OH,NO3,J,NO,RO2NO3,concs,\n")
#names = "    ,ts,RO2HO2,RO2NO,RO2NO2,RO2,HO2,OH,NO3,J,NO,RO2NO3,conc        "
for t in 1:144#linspace(2,len(specs[2]),100)
         t = trunc(Int,t)
         write(file, string(t)*",");write(file, string(t)*",")
       newarr = [];push!(newarr,t);push!(newarr,t)
       for i in [RO2HO2,RO2NO,RO2NO2,RO2,HO2,OH,NO3,J,NO,RO2NO3,]
         val=smrxn(i,t,spcsel)
         write(file, string(smrxn(i,t,spcsel))*",")
         println("s2wwww ",t,spcsel)
         #push!(newarr,val)
       end
         #push!(newarr,specs[t,Symbol(specsel)])
          write(file, string(specs[t,Symbol(spcsel)])*",")
         write(file,"\n")
end
#writedlm("results.csv", dmy, ",")
close(file)

print("fini")


using POMDPs
using ContinuousPOMDPTreeSearchExperiments
using POMDPToolbox
using POMCP
using POMCPOW
using LightDarkPOMDPs
using ProgressMeter
using PmapProgressMeter
using ParticleFilters
using JLD
using CPUTime

N = 20

@everywhere begin
    using LightDarkPOMDPs
    using POMDPs
    using ContinuousPOMDPTreeSearchExperiments
    using POMDPToolbox
    using POMCP
    using POMCPOW
    using ParticleFilters

    pomdp = LightDark2DTarget(term_radius=0.1,
                              init_dist=SymmetricNormal2([2.0, 2.0], 5.0),
                              discount=0.95)
    p_rng = MersenneTwister(4)
    solvers = Dict{String, Union{Solver,Policy}}(
        #=
        "heuristic" => begin
            homing_policy = SimpleFeedback(gain=1.0, max_radius=10.0)
            exploring_up = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(1000), 0.05, copy(p_rng))
            homing_up = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(1000), 0.0, copy(p_rng))
            InfoGatherHeur(5.0, 0.03, homing_policy, exploring_up, homing_up)
        end,

        "pomcpdpw" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            rollout_policy = RadiusRandom(radius=10.0, rng=rng3)
            tree_queries = 1_000_000
            solver = POMCPDPWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    c=5.0,
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=OneStepValue(),
                                    default_action=Vec2(0.1, 0.1),
                                    rng=rng3
                                   )
        end,

        "bt_100_osv_100k" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            rollout_policy = SimpleFeedback(max_radius=10.0)
            tree_queries = 100_000
            node_updater = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(100), 0.05, rng3)
            solver = POMCPDPWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    c=5.0,
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=OneStepValue(),
                                    node_belief_updater=node_updater,
                                    default_action=Vec2(0.1, 0.1),
                                    rng=rng3
                                   )
        end,
        =#

        "pomcpow" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            tree_queries = 1_000_000
            solver = POMCPOWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    criterion=MaxUCB(5.0),
                                    final_criterion=MaxTries(),
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=OneStepValue(),
                                    rng=rng3
                                   )
        end,

        "pomcpow_10k" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            tree_queries = 10_000
            solver = POMCPOWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    criterion=MaxUCB(5.0),
                                    final_criterion=MaxTries(),
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=OneStepValue(),
                                    rng=rng3
                                   )
        end,

        "bt_100_ro_100k" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            rollout_policy = SimpleFeedback(max_radius=10.0)
            tree_queries = 100_000
            node_updater = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(100), 0.05, rng3)
            solver = POMCPDPWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    c=5.0,
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=PORolloutEstimator(rollout_policy, node_updater),
                                    node_belief_updater=node_updater,
                                    default_action=Vec2(0.1, 0.1),
                                    rng=rng3
                                   )
        end,

        "bt_100_osv_100k" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            rollout_policy = SimpleFeedback(max_radius=10.0)
            tree_queries = 100_000
            node_updater = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(100), 0.05, rng3)
            solver = POMCPDPWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    c=5.0,
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=OneStepValue(),
                                    node_belief_updater=node_updater,
                                    default_action=Vec2(0.1, 0.1),
                                    rng=rng3
                                   )
        end,

        "greedy" => SimpleFeedback(gain=1.0, max_radius=10.0)

    )

    updaters = Dict{String, Any}(
        "heuristic" => :heuristic
    )
end

solver_keys = keys(solvers)
rewards = Dict{String, AbstractVector{Float64}}()
# state_hists = Dict{String, AbstractVector{AbstractVector{state_type(pomdp)}}}()
counts = Dict{String, AbstractVector{Int}}()
times = Dict{String, AbstractVector{Float64}}()
steps = Dict{String, AbstractVector{Int}}()

for (j, sk) in enumerate(solver_keys)
    s_rewards = SharedArray(Float64, N)
    s_counts = SharedArray(Int, N)
    # s_hists = SharedArray(AbstractVector{state_type(pomdp)}, N)
    s_times = SharedArray(Float64, N)
    s_steps = SharedArray(Float64, N)
    prog = Progress(N, desc="$sk ($j of $(length(solver_keys)))...")
    pmap(prog, 1:N) do i
        sim_rng = MersenneTwister(i)
        up_rng = MersenneTwister(i+100_000)
        if isa(solvers[sk], Policy)
            policy = deepcopy(solvers[sk])
        else
            policy = solve(solvers[sk], pomdp)
        end
        if get(updaters, sk, :standard) == :standard
            up = ObsAdaptiveParticleFilter(deepcopy(pomdp), LowVarianceResampler(100_000), 0.05, up_rng)
        elseif updaters[sk] == :heuristic
            up = updater(policy)
        end
        sim = HistoryRecorder(max_steps=40, rng=sim_rng)
        pomdp.count = 0
        s_times[i] = @CPUelapsed hist = simulate(sim, deepcopy(pomdp), policy, up)
        s_rewards[i] = discounted_reward(hist)
        s_counts[i] = pomdp.count
        s_steps[i] = n_steps(hist)
    end
    rewards[sk] = sdata(s_rewards)
    @show mean(rewards[sk])
    # state_hists[sk] = sdata(s_hists)
    counts[sk] = sdata(s_counts)
    times[sk] = sdata(s_times)
    steps[sk] = sdata(s_steps)
    @show sum(counts[sk])/sum(steps[sk])
end

for k in solver_keys
    println("$k mean: $(mean(rewards[k])) sem: $(std(rewards[k])/sqrt(N))")
    println("$k time: $(mean(times[k])) sim counts: $(mean(counts[k])) steps: $(mean(steps[k]))")
end

filename = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "data", "compare_$(Dates.format(now(), "E_d_u_HH_MM")).jld")
println("saving to $filename...")
@save(filename, pomdp, solver_keys, solvers, rewards, updaters, counts, steps, times)
println("done.")


include("prepareDB.jl")
##
fdata = filter(r ->
    r.Area != "PIR" &&
    r.Area != "ENT" &&
    r.Area != "OT" &&
    r.Areac!=
    # r.Mouse != "FM14",
    r.Valid_frs,
    fulldata)
    # lvA = union(fdata.Area)
    # fdata.Area = categorical(fdata.Area; levels = lvA)
##
function prov_define_bold(df0, signal; av_pres = false, method = :Peak, peak_to = :all, nmax = 4, window = 4:10, nanmath = true)
    fmean, fsem = nanmath ? (safemean, safesem) : (mean, sem)
    if  method == :Max
        #average response over all presentation
        df1 = average_presentation(df0,signal; nanmath = nanmath)
        gd1 = groupby(df1,[:Genotype, :Mouse, :Stim, :Area, :Concentration])
        signal = :Bold
        #find the first n max values in the bold signal
        df2 = combine(gd1, signal => (s -> FMRIOpto.max_bold(s; max = nmax, nanmath = nanmath)) => :Bold)
    elseif method == :Window
        #average all bold values per mouse and concentration in a time window
        df1 = combine(groupby(df0, [:Time,:Concentration,:Area,:Stim,:Mouse,:Genotype]), signal => (s -> mean(skipmissing(s))) => :Bold)
        filter!(r -> r.Time in window, df1)
        gd1 = groupby(df1,[:Genotype, :Mouse, :Stim, :Area, :Concentration])
        df2 = combine(gd1, :Bold => (s -> mean(skipmissing(s))) => :Bold)
    elseif method == :None
        df2 = filter(r -> r.Time in window, df0)
        df2.Bold = df2[:,signal]
    else
        error("unrecognised bold processing method:/n options are :Peak, :Max, :Timewindow")
    end
    return df2
end

function prov_plot_bold(df0; nanmath = true, lim = (-0.4,0.7))
    fmean = nanmath ? safemean : mean
    fsem = nanmath ? safesem : sem
    #calculate mean repsonse over concentrations per area
    df = combine(groupby(df0,[:Area, :Stim, :Concentration, :Genotype]),  :Bold .=> [fmean, fsem] .=> [:mean, :sem])
    df1 = filter(r-> r.Genotype == "sert",df)
    # identify regions with missing concentration info
    res = combine(groupby(df1,[:Area]), :mean => (m -> .!(any(isnan.(m)))) => :to_keep)
    to_keep = filter(r -> r.to_keep, res)[:,:Area]
    # prepare a grid plot with the results for each regions
    df1.Color = [x ? :red : :black for x in df1.Stim]
    # posdict = Dict(v=>p for (v,p) in zip(union(df1.Concentration), collect(1:length(union(df1.Concentration)))))
    # df1.Pos = [get(posdict, x,0) for x in df1.Concentration]
    # sort!(df1,[:Stim,:Pos])
    sort!(df1,[:Stim,:Concentration])
    # fconc && filter!(r -> r.Concentration != 0.0001 ,df1)
    plt = []
        combine(groupby(df1, :Area)) do dd
            pp = @df dd scatter(:Concentration, :mean, yerror = :sem, group = :Stim, color = :Color,
                title = :Area[1], legend = false, ylims = lim)
            @df dd plot!(:Concentration, :mean, color = :Color, group = :Stim)
            Plots.abline!(0,0,color = :black, linestyle = :dash)
            push!(plt,pp)
        end
    return plt, df, to_keep
end
function prov_odour_response(df0; farea = :none)
    # use only control data to determine odour concentration dose dependency
    df1 = filter(r -> !r.Stim, df0)
    #remove lowest concentration
    # fconc && filter!(r -> r.Concentration != 0.0001, df1)
    # filters area with insufficient response, identified from plot_bold function
    farea != :none && filter!(r -> r.Area in farea, df1)
    # identify if the model requires to incorporate time
    if iscolumn(df0,:Time)
        f0 = @formula(Bold ~ 1 + Time + Concentration + (1+Time|Mouse))
        f1 = @formula(Bold ~ 1 + Time + Concentration*Area + (1+Time|Mouse))
    else
        f0 = @formula(Bold ~ 1 + Concentration + (1|Mouse))
        f1 = @formula(Bold ~ 1 + Concentration*Area + (1|Mouse))
    end
    # #replace NaN values used for plotting to missing for model
    # df1.Bold = [isnan(x) ? missing : x for x in df1.Bold]
    # model fit and test
    m0 = fit(MixedModel, f0, df1)
    m1 = fit(MixedModel, f1, df1)
    mp = MixedModels.likelihoodratiotest(m0,m1)
    # identifies area that significantly respond to odour concentration
    res = DataFrame(Area = coefnames(m1), P = m1.pvalues, Z = m1.beta ./ m1.stderror)
    res2 = filter!(r -> occursin("Concentration & Area: ", r.Area), res)
    res2.Area = replace.(res2.Area, "Concentration & Area: " => "")
    to_keep = filter(r -> r.P <= 0.049, res2)[:,:Area]

    return m1, mp, to_keep
end

function prov_stim_interaction(df0, farea)
    # filters area lacking concentration response identified by odour_response function
    df1 = filter(r -> r.Genotype == "sert" && r.Area in farea, df0)
    # fconc && filter!(r -> r.Concentration != 0.0001, df0)
    # identify if the model requires to incorporate time
    if iscolumn(df0,:Time)
        f1 = @formula(Bold ~ 1 + Time + Concentration*Area + Stim + zerocorr(1+Time+Concentration|Mouse) + zerocorr(1+Time+Concentration|Area))
        f2 = @formula(Bold ~ 1 + Time + Concentration + Area + Stim + Concentration & Area + Concentration & Area & Stim + zerocorr(1+Time+Concentration|Mouse) + zerocorr(1+Time+Concentration|Area))
        f3 = @formula(Bold ~ 1 + Time + Concentration*Area*Stim + zerocorr(1+Time+Concentration|Mouse) + zerocorr(1+Time+Concentration|Area))
    else
        f1 = @formula(Bold ~ 1 + Concentration*Area + Stim + zerocorr(1+Concentration|Mouse) + zerocorr(1+Concentration|Area))
        f2 = @formula(Bold ~ 1 + Concentration + Area + Stim + Concentration & Area + Concentration & Area & Stim + zerocorr(1+Concentration|Mouse) + zerocorr(1+Concentration|Area))
        f3 = @formula(Bold ~ 1 + Concentration*Area*Stim + zerocorr(1+Concentration|Mouse) + zerocorr(1+Concentration|Area))
    end
    #replace NaN values used for plotting to missing for model
    # df1.Bold = [isnan(x) ? missing : x for x in df1.Bold]
    m1 = fit(MixedModel, f1,df1)
    m2 = fit(MixedModel, f2,df1)
    m3 = fit(MixedModel, f3,df1)
    p_simple = MixedModels.likelihoodratiotest(m1,m2)
    p_complex = MixedModels.likelihoodratiotest(m2,m3)
    res = model_res(m2)
    filter!(r -> occursin("Stim", r.Area), res)
    sort!(res,:P)
    return df1, m2, m3, p_complex, res
end

function prov_analysis_bold(df0, signal; fconc = true, av_pres = true, method = :Peak, peak_to = :all, nmax = 6, window = 4:12, logarithmic = false, nanmath = true, lim = (-0.4,0.7))
    #adjust bold signal
    bold_df = prov_define_bold(df0,signal; method = method, av_pres = av_pres, peak_to = peak_to, nmax = nmax, window = window, nanmath = nanmath)
    fconc && filter!(r -> r.Concentration != 0.0001, bold_df)
    logarithmic && (bold_df.Concentration = log10.(bold_df.Concentration))
    if method == :None
        gd = groupby(bold_df,[:Genotype, :Mouse, :Stim, :Area, :Concentration])
        pltdf = combine(gd, :Bold => (b -> mean(skipmissing(b))) => :Bold)
    else
        pltdf = transform(bold_df, :Bold => ByRow(b -> (ismissing(b) ? NaN : b)) => :Bold)
    end
    plt, plt_df, to_keep = prov_plot_bold(pltdf; nanmath = true, lim = lim)
    display(plot(plt..., size = (900*1.5,1200*1.5)))#layout = (6,3))
    ## find areas that respond to concentration
    m1, mp, to_keep2 = prov_odour_response(bold_df; farea = to_keep)
    # assess stim concentration interaction
    stim_df, m2, m3, p_complex, res = prov_stim_interaction(bold_df, to_keep2)
    show(m2)
    show(m3)
    show(p_complex)
    return bold_df, plt, plt_df, to_keep, m1, mp, to_keep2, stim_df, m2, m3, p_complex, res
end
##
bold_df, plt, plt_df, to_keep, m1, mp, to_keep2, stim_df, m2, m3, p_complex, res =
    prov_analysis_bold(fdata, :BOLD_90; fconc = true, av_pres = true, method = :Window,
    peak_to = :all, nmax = 6, window = 4:10, logarithmic = true,
    nanmath = true, lim = (-0.4,0.7))
##
union(bold_df.Genotype)
test_df = filter(r -> r.Genotype == "sert" && r.Area in to_keep && r.Area in to_keep2, bold_df)
f1 = @formula(Bold ~ 1 + Concentration * Area + Stim + Stim & Concentration + zerocorr(1+Concentration|Mouse) + zerocorr(1+Concentration|Area))
f2 = @formula(Bold ~ 1 + Concentration * Area * Stim + zerocorr(1+Concentration|Mouse) + zerocorr(1+Concentration|Area))
m1 = fit(MixedModel,f1,test_df)
m2 = fit(MixedModel,f2,test_df)
MixedModels.likelihoodratiotest(m1,m2)
##
f_add = @formula(Bold ~ 1 + Concentration * Area + Stim + Stim & Area + Stim & Concentration + zerocorr(1+Concentration|Mouse) + zerocorr(1+Concentration|Area))
f_mul = @formula(Bold ~ 1 + Concentration * Area + Stim + Stim & Concentration + Stim & Concentration & Area + zerocorr(1+Concentration|Mouse) + zerocorr(1+Concentration|Area))
m_add = fit(MixedModel,f_add,test_df)
m_mul = fit(MixedModel,f_mul,test_df)
##
function AICc(model)
    aic(model) + ((2*(dof(model)^2) + 2*dof(model))/(nobs(model) - dof(model) - 1))
end
function AICc_test(candidate, simpler)
    exp((AICc(candidate) - AICc(simpler))/2)
end

AICc_test(m_add,m_mul)

MixedModels.likelihoodratiotest(m_add,m2)
MixedModels.likelihoodratiotest(m_mul,m2)
##

df00 = fdata
signal = :BOLD_90
fconc = true
av_pres = true
method = :Window
peak_to = :all
nmax = 6
window = 4:10
logarithmic = true

df0 = prov_define_bold(df00, signal; av_pres = av_pres, method = :Window, window = window, nanmath = true)
fconc && filter!(r -> r.Concentration != 0.0001, df0)
logarithmic && (df0.Concentration = log10.(df0.Concentration))
open_html_table(df0)
pltdf = transform(df0, :Bold => ByRow(b -> (ismissing(b) ? NaN : b)) => :Bold)
plt, plt_df, to_keep = prov_plot_bold(pltdf)
plot(plt...)
sum(isnan.(df00.BOLD_90))
sum(ismissing.(df00.BOLD_90))
m1, mp, to_keep2 = prov_odour_response(df0; farea = to_keep)
to_keep2


function model()
    T = Float32
    x = Var()
    y = Embedding(T,10000,100)(x)
end

function getscore(m, s::State)
    fs = features(s)
    score = 0.0

end