using Roots
using Statistics
using FileIO
using Parameters
include("random_matrices.jl")
include("network_sims_gpu.jl")


resultsdir = "results2"

etas = collect(-1.0f0:0.1f0:1.0f0)
ind = parse(Int64, ARGS[1])

N = 3000
eta = etas[ind]
s = 0.0f0

J_seeds = collect(5000:5019)
h0_seed = 6000

h0 = generate_nv(N, Float32, seed=h0_seed)

wait = 200.0f0
t_bound = 2000.0f0
dt = 0.01f0

startstep = round(Int, wait / dt) + 1


@unpack target_avg_var = load("vanilla_target_avg_var.jld2")

# Difference between the average variance of the network ensemble and target
function f(g)
    network_avg_vars = Vector{Float32}(undef, length(J_seeds))
    for (i, J_seed) in enumerate(J_seeds)
        J = generate_J(N, eta, seed=J_seed)
        A = sim_activity_gpu(J, g, h0; s=s, dt=dt, t_bound=t_bound,
                minimum_norm=1e-90)
        if size(A, 2) < startstep
            avg_var = 0
        else 
            avg_var = mean(A[:,startstep:end].^2)
        end
       network_avg_vars[i] = avg_var
    end
    return mean(network_avg_vars) - target_avg_var
end

if eta == -1.0f0
    lo_bound = 10f0
else
    lo_bound = 1 / (1 + eta)
end 
hi_bound = 20f0
@time g = find_zero(f, (lo_bound, hi_bound), Roots.A42(), verbose=true)

out = Dict{String,Any}()
@pack! out = N, g, eta, s, J_seeds, h0_seed, wait, t_bound, dt, target_avg_var
save(joinpath(resultsdir, "Rootfinding_E$(ind)_N=$N,eta=$eta,s=$s.jld2"), out)
