#=
Here I show the exact correspondence
my spiking implementation and code by
Litwin-Kumar / Auguste Schulz
Compare rates DURING STIMULATION
with a PSTH !
=#

# include inbuilt modules
using Distributions
using Dates
using LinearAlgebra
using Random
using Distributed
using Plots ; theme(:dark)
using Random ; Random.seed!(0)

# --------------------- include functions  ------------------------------------------------
# include runs the respective julia code, i.e. defined functions are then in the workspace

#function gennextsequencepretraining(stim::Array{Float64,2}, firstimg, novelimg; Nimg = 4, Ntrain = 20, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
function genstimparadigmpretraining(stim::Array{Float64,2}; Nimg = 4, Nass = 20, Ntrain = 20, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
	# generate a stimulus sequence
	# where each assembly is
	# Nimg Number of images per sequence
	# Nreps Number of seqeunce repititions in one block
	# Nseq Number of sequences
	# Nblocks Number of block repititions

	if Ntrain != 0
		Nstim = Nass * Ntrain
		tempstim = zeros(Nstim, 4)

		#tempstim[1,1] = firstimg
		# repeat sequence several times, i.e. assemblie numbers 123412341234...
		#tempstim[:,1] = repeat(firstimg:(firstimg+Nimg-1), outer = Nreps)
		# ensure that images are shuffled but never the same image after the other
		images = 1:Nass
		assemblysequence = shuffle(copy(images))
		for rep = 2:Ntrain
			shuffleimg = shuffle(images)
			#println("begin $(shuffleimg[1])")
			while shuffleimg[1] == assemblysequence[end]
				println("shuffle again")
				shuffleimg = shuffle(images)
				#println("new begin $(shuffleimg[1])")
			end
		assemblysequence = vcat(assemblysequence, shuffleimg)
		end
		tempstim[:,1] .= assemblysequence
		if stim[end,3] == 0
			tempstim[1,2] = stimstart
		else
			tempstim[1,2] = stim[end,3]+lenpause
		end
		tempstim[1,3] = tempstim[1,2] + lenstim
		tempstim[1,4] = strength

		for i = 2:Nstim#Nimg*Nreps #number in image times
			# # lenstim ms stimulation and lenpause ms wait time
			tempstim[i,2] = tempstim[i-1,2] + lenstim + lenpause
			tempstim[i,3] = tempstim[i,2]+lenstim
			tempstim[i,4] = strength
		end
		if stim[end,3] == 0
			stim = tempstim
		else
			stim = vcat(stim, tempstim)
		end
		#tempstim = nothing
	end
	return stim
end # function genstim

function gennextsequenceNonovel(stim::Array{Float64,2}, firstimg, novelimg; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
	# generate a stimulus sequence
	# where each assembly is
	# Nimg Number of images per sequence
	# Nreps Number of seqeunce repititions in one block
	# Nseq Number of sequences
	# Nblocks Number of block repititions
	Nstim = Nimg*Nreps
	tempstim = zeros(Nstim, 4)
    #tempstim[1,1] = firstimg
	tempstim[:,1] = repeat(firstimg:(firstimg+Nimg-1), outer = Nreps)

	if stim[end,3] == 0
		tempstim[1,2] = stimstart
	else
    	tempstim[1,2] = stim[end,3]+lenpause
	end
    tempstim[1,3] = tempstim[1,2] + lenstim
    tempstim[1,4] = strength

    for i = 2:Nimg*Nreps #number in image times
        # lenstim ms stimulation and lenpause ms wait time
        tempstim[i,2] = tempstim[i-1,2] + lenstim + lenpause
        tempstim[i,3] = tempstim[i,2]+lenstim
        tempstim[i,4] = strength
    end
	#tempstim[end-Nimg,1] = novelimg
	if stim[end,3] == 0
		stim = tempstim
	else
		stim = vcat(stim, tempstim)
	end
	#tempstim = nothing
    return stim
end # function genstim

function genstimparadigmNonovel(stimulus; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)

	firstimages = collect(1:Nimg:Nimg*Nseq)
	novelimg = firstimages .+ Nimg*Nseq
	novelrep = collect(Nimg*Nseq:Nimg*Nseq+Nimg).+1

	# for img in firstimages
	# 	stimulus = gennextsequence!(stimulus, img)
	# end
	blockonset = []
	storefirstimages = copy(firstimages)
	for b = 1:Nblocks
		if b == 1
			for img in firstimages
				append!(blockonset, size(stimulus,1) + 1)
				stimulus = gennextsequenceNonovel(stimulus, img, img + Nseq * Nimg, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
			end
		else
			shuffleimg = shuffle(firstimages)
			while shuffleimg[1] == storefirstimages[end]
				shuffleimg = shuffle(firstimages)
			end
			storefirstimages = vcat(storefirstimages, shuffleimg)
			for img in shuffleimg # since we shuffle ensure that it is still the correct novel image
				append!(blockonset, size(stimulus,1) + 1)
				stimulus = gennextsequenceNonovel(stimulus, img, img + Nseq * Nimg,Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
			end
		end
	end
	# ensure mapping from novel img to right value despite shuffling
	for i = 1:length(novelimg)
		stimulus[stimulus[:,1] .== novelimg[i],1] .= novelrep[i]
	end

	# ensure that blockonset 1 when we start with an emptz array
	if blockonset[1] == 2
		blockonset[1] = 1
	end
	return stimulus, convert(Array{Int64,1}, blockonset)

end

function genstimparadigmnovelcont(stimulus; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
  """	generate the stimulation paradigm ensure that for each block there is a different novel image and
  that two consecutive blocks are never the same

  generate seq. violation stimulation paradigm
  ABCABCABC....ABN1ABC
  DEFDEFDEF....DEN2DEF
  XYZXYZXYZ...XYN3XYZ
  DEFDEFDEF....DEN4DEF
  ABCABCABC....ABN1ABC
  """

    firstimages = collect(1:Nimg:Nimg*Nseq)
    novelimg = Nimg*Nseq + 1 # start counting from the first assembly after the core assemblies
    novelrep = collect(Nimg*Nseq:Nimg*Nseq+Nimg).+1

    storefirstimages = copy(firstimages)
    blockonset = []

    for b = 1:Nblocks
      #println("block $(b) -------------------")
      if b == 1
        for img in firstimages
          append!(blockonset, size(stimulus,1) + 1)
          stimulus = gennextsequencenovel(stimulus, img, novelimg, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
          novelimg += 1
        end
      else
        # ensure that the first images of two consecutive blocks differ
        shuffleimg = shuffle(firstimages)
        #println("begin $(shuffleimg[1])")
        while shuffleimg[1] == storefirstimages[end] && Nseq > 1 # added Nseq larger 1 cause otgerwise always the same
          println("shuffle again")
          shuffleimg = shuffle(firstimages)
          #println("new begin $(shuffleimg[1])")
        end
        storefirstimages = vcat(storefirstimages, shuffleimg)
        #println(storefirstimages)
        for img in shuffleimg  # since we shuffle ensure that it is still the correct novel image
          append!(blockonset, size(stimulus,1) + 1)
          stimulus = gennextsequencenovel(stimulus, img, novelimg ,Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
          novelimg += 1
        end
      end
    end

    # ensure that blockonset 1 when we start with an emptz array
    if blockonset[1] == 2
      blockonset[1] = 1
    end
    return stimulus, convert(Array{Int64,1}, blockonset)

  end

function gennextsequencenovel(stim::Array{Float64,2}, firstimg, novelimg; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
	"""generate a stimulus sequence with novel stim
	Nimg Number of images per sequence
	Nreps Number of seqeunce repititions in one block
	Nseq Number of sequences
	Nblocks Number of block repititions"""
	Nstim = Nimg*Nreps
	tempstim = zeros(Nstim, 4)
    #tempstim[1,1] = firstimg
	# repeat sequence several times, i.e. assemblie numbers 123412341234...
	tempstim[:,1] = repeat(firstimg:(firstimg+Nimg-1), outer = Nreps)
	if stim[end,3] == 0
		tempstim[1,2] = stimstart
	else
    	tempstim[1,2] = stim[end,3]+lenpause
	end
    tempstim[1,3] = tempstim[1,2] + lenstim
    tempstim[1,4] = strength

    for i = 2:Nimg*Nreps #number in image times
        # # lenstim ms stimulation and lenpause ms wait time
        tempstim[i,2] = tempstim[i-1,2] + lenstim + lenpause
        tempstim[i,3] = tempstim[i,2]+lenstim
        tempstim[i,4] = strength
    end
	tempstim[end-Nimg,1] = novelimg
	if stim[end,3] == 0
		stim = tempstim
	else
		stim = vcat(stim, tempstim)
	end
	#tempstim = nothing
    return stim
end # function genstim

function genstimparadigmssa(stimulus; Nimg = 4, Nreps = 20, Nseq = 5, Nblocks = 1, stimstart = 1000, lenstim = 1000, lenpause = 3000, strength = 8)
    """
    generate SSA stimulation paradigm
    A A A A A A A B A A A ... Block 1
    B B B B B 				 Block 2 ...
    Nreps 20 times
    Blocks
    """
    Nimg = 2
    Nreps = 5
    Nseq = 5
    Nblocks = 10
    stimulus, blockonset = genstimparadigmnovelcont(stimulus, Nimg = 2, Nreps = 5, Nseq = 5, Nblocks = 10, stimstart = stimstart, lenstim = 100, lenpause = 300, strength = strength )

    #last_onset = blockonset[1]
    freq = 1
    dev = 2
    for bb =1:length(blockonset)
      #idxfreq = collect(blockonset[bb]:(blockonset[bb]+Nimg*Nreps))
      freq_start = blockonset[bb]
      freq_end = blockonset[bb]+Nimg*Nreps
      dev_end = blockonset[bb]+Nimg*Nreps
      dev_start = blockonset[bb]+Nreps
      if bb == length(blockonset)
        stimulus[freq_start:end,1] .= freq
        #println(idxfreq)
        #idxdev = collect(blockonset[bb]+Nreps:Nreps:(blockonset[bb]+Nimg*Nreps))
        stimulus[dev_start:Nreps:end,1] .= dev
        #printl
      else
        stimulus[freq_start:freq_end,1] .= freq
        #println(idxfreq)
        #idxdev = collect(blockonset[bb]+Nreps:Nreps:(blockonset[bb]+Nimg*Nreps))
        stimulus[dev_start:Nreps:dev_end,1] .= dev
        #println(idxdev)
        # for the nect block switch deviant and frequent
      end
      temp =freq
      freq = dev
      dev = temp
    end

    return stimulus, blockonset
end

function initweights((Ne, Ni, p, Jee0, Jei0, Jie, Jii))
	""" initialise the connectivity weight matrix based on initial
	E-to-E Jee0 initial weight e to e plastic
	I-to-E Jei0 initial weight i to e plastic
	E-to-I Jie0 constant weight e to i not plastic
	I-to-I Jii constant weight i to i not plastic"""
	Ncells = Ne+Ni

	# initialise weight matrix
	# w[i,j] is the weight from pre- i to postsynaptic j neuron
	weights = zeros(Float64,Ncells,Ncells)
	weights[1:Ne,1:Ne] .= Jee0
	weights[1:Ne,(1+Ne):Ncells] .= Jie
	weights[(1+Ne):Ncells,1:Ne] .= Jei0
	weights[(1+Ne):Ncells,(1+Ne):Ncells] .= Jii
	# set diagonal elements to 0
	# for cc = 1:Ncells
	# 	weights[cc,cc] = 0
	# end
	weights[diagind(weights)] .= 0.0
	# ensure that the connection probability is only p
	weights = weights.*(rand(Ncells,Ncells) .< p)
	return weights
end

function weightpars(;Ne = 4000, Ni = 1000, p = 0.2 )
	"""Ne, Ni number of excitatory, inhibitory neurons
	p initial connection probability"""
	Jee0 = 2.86 #initial weight e to e plastic
	Jei0 = 48.7 #initial weight i to e plastic
	Jie = 1.27 #constant weight e to i not plastic
	Jii = 16.2 #constant weight i to i not plastic
	return Ne,Ni,p,Jee0,Jei0,Jie,Jii
end

function initassemblymembers(;Nassemblies = 20, pmember = .05, Nmembersmax = 300, Ne = 4000)
	"""Nassemblies number of assemblies
	pmember probability of belonging to any assembly
	Nmembersmax maximum number of neurons in a population (to set size of matrix)

	set up excitatory assemblies"""
	#seed = 1
	#Random.seed!(seed) # ensure same Assemblies when repeating
	assemblymembers = ones(Int,Nassemblies,Nmembersmax)*(-1)
	for pop = 1:Nassemblies
		members = findall(rand(Ne) .< pmember)
		assemblymembers[pop,1:length(members)] = members
	end
	#println(assemblymembers)
	return assemblymembers
end

function initinhibassemblymembers(;Nassemblies = 20, pmember = .15, Nmembersmax = 200, Ne = 4000, Ni = 1000)
	"""Nassemblies number of assemblies
	pmember probability of belonging to any assembly
	Nmembersmax maximum number of neurons in a population (to set size of matrix)

	set up inhibitory assemblies
	higher connection probability"""
	#seed = 1
	#Random.seed!(seed) # ensure same Assemblies when repeating
	inhibmembers = ones(Int,Nassemblies,Nmembersmax)*(-1)
	for pop = 1:Nassemblies
		members = findall(rand(Ni) .< pmember) .+ Ne
		inhibmembers[pop,1:length(members)] = members
	end
	#println(assemblymembers)
	return inhibmembers
end

#include("../simulation/runsimulation_inhibtuning.jl")
#include("../simulation/sequencefunctions.jl")
#include("../simulation/helperfunctions.jl")
#include("../evaluation/evaluationfunctions.jl")

# --------------------- initialise simulation --------------------------------------------

# Define number of excitatory and inhibitory neurons
const Ne = 4000
const Ni = 1000

Ncells = Ne + Ni
# Set integration timestep
dt 	= 0.1 #integration timestep in ms

ifiSTDP = false # true 		#  include inhibitory plasticity
ifwadapt = false	#  consider AdEx or Ex IF neurons

# --------------------- generate the stimulus --------------------------------------------

# In case this file is not run from the command line specify the ARGS list of strings
# Nimg 3, Nreps 20, Nseq 10, Nblocks 1, lenstim 300, strength 12, Ntrain 5, adjustfactor 1.0, adaptive neurons false, inhibfactor 0.1
ARGS = ["3", "20", "10", "1","300", "12", "5", "10", "0", "10"] # unique sequences
# ARGS = ["3", "20", "5", "10","300", "12", "5", "10", "0", "10"] # repeated sequences

# stimulus parameters
Nimg = parse(Int64, ARGS[1]) # number of stimuli per sequence
Nreps = parse(Int64, ARGS[2]) # number of repetitions per sequence cycle
Nseq = parse(Int64, ARGS[3]) # number of sequences
Nblocks = parse(Int64, ARGS[4]) # number of sequence block repetitions
stimstart = 1_000  # 4000 # start time of stimulation in ms
lenstimpre = parse(Int64, ARGS[5]) # duration of assembly stimulation per image in ms
lenpausepre = 0 # duration of no stimulation between two images in ms
strengthpre = parse(Int64, ARGS[6]) # strength of the stimulation in kHz added to baseline input
lenstim = lenstimpre # duration of assembly stimulation per image in ms
lenpause = lenpausepre # duration of no stimulation between two images in ms
strength = strengthpre # strength of the stimulation in kHz added to baseline input
Ntrain = parse(Int64, ARGS[7]) # number of pretraining iterations
Nass = (Nimg) * Nseq +  Nseq * Nblocks # total number of assemblies Nimg * Nseq + all novelty assemblies

# adjustment of learning rates
adjustfactor = parse(Int64, ARGS[8])/10
adjustfactorinhib = adjustfactor # ensure they are equal for inhibitory and exc. plasticity

# adaptive currents
if parse(Int64, ARGS[9]) == 1
	ifwadapt = true
end

# strength of the inhibitory tuning
inhibfactor = parse(Int64, ARGS[10])/100 # 100 more refined steering possibility to just switch it off

stimparams = [Nimg, Nreps, Nseq, Nblocks, stimstart, lenstim, lenpause, strength] # store stimulus param array in hdf5 file
stimparams_prestim = [Nimg, Nreps, Nseq, Nblocks, stimstart, lenstim, lenpause, strength, lenstimpre, lenpausepre, strengthpre, Ntrain] # store stimulus param array in hdf5 file

# initialise stimulus array
stimulus = zeros(1, 4)	# initialisation of the stimulus
# nohup julia initsim_sequence_violation_ARGS.jl 3 20 5 10 300 12 5 10 0 10 &> ../tmp/Pre5_Final_Run_Block_nonadapt_10_inhibtuning_nohetero_tunefactor_10_shortpretrain.txt &

# specify a tag which is added at the end of the filename to account for different stimulation paradigms
#tag = "repseq.h5"

# ---------------------- generate the stimulus --------------------------------

# specify aspects of the stimulation paradigm
withnovelty 	= true 	# novelty response
pretrainig 		= true	# include pretraining phase

# ----------- standard setting when the following three false -----------------
shuffledstim 	= false	# shuffle the stimulus randomly
lastshuffled 	= false	# shuffle the last two stimuli
reducednovelty 	= false	# reduce the stimulation strength of the novel stimulus

if withnovelty
	if pretrainig
		# --------- main stim generation function ---------------
		stimulus = genstimparadigmpretraining(stimulus, Nass = Nass, Ntrain = Ntrain, stimstart = stimstart, lenstim = lenstimpre, lenpause = lenpausepre, strength = strengthpre)
		lenpretrain = size(stimulus,1)
	end

	if shuffledstim
		if pretrainig
			stimulus, blockonset = genstimparadigmnovelshuffledpretrain(stimulus, lenpretrain, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength)
			tag = "repeatedsequencesshuffled.h5"

		else
			stimulus, blockonset = genstimparadigmnovelshuffled(stimulus, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength )
			tag = "repeatedsequencesshuffled.h5"
		end
	else # if shuffled stim

		# --------- main stim generation function ---------------
		stimulus, blockonset = genstimparadigmnovelcont(stimulus, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength )
		tag = "repeatedsequences.h5"

	end
else # if with novelty
	# even if no novelty assemblies are shown still train with all novelty assemblies
	if pretrainig
		#println(stimulus)
		stimulus = genstimparadigmpretraining(stimulus, Nass = Nass, Ntrain = Ntrain, stimstart = stimstart, lenstim = lenstimpre, lenpause = lenpausepre, strength = strengthpre)
		#println(stimulus)
		lenpretrain = size(stimulus,1)
		tag = "repeatedsequences_imprinting.h5"

	end
	tag = "repeatedsequencesnoNovelty.h5"
	stimulus, blockonset = genstimparadigmNonovel(stimulus, Nimg = Nimg, Nreps = Nreps, Nseq = Nseq, Nblocks = Nblocks, stimstart = stimstart, lenstim = lenstim, lenpause = lenpause, strength = strength )
end

# set length of pretraining to 0 if no pretraining
if Ntrain == 0
	lenpretrain = 0
end

# get sequence order
blockidx = collect(lenpretrain + 1:Nreps*Nimg:size(stimulus,1)) # get all stimulus indices when a new sequence starts
seqnumber = stimulus[blockidx,1]
for n = 2:Nseq
    seqnumber[seqnumber .== (1 + (n-1)*Nimg)] .= n
end

# Switch last two stimuli instead of novel stimulus
if lastshuffled
	if pretrainig

		idx = findall(stimulus[lenpretrain + 1:end,1] .> Nseq*Nimg)
		vals = stimulus[idx.+(lenpretrain),1]
		stimulus[idx.+(lenpretrain),1] .= stimulus[idx.+(lenpretrain-1),1]
		stimulus[idx.+(lenpretrain-1),1] .= stimulus[idx.+(lenpretrain-1),1] .+ 1
		tag = "pretrainshufflenovelty" * tag
	else
		idx = findall(stimulus[:,1] .> Nseq*Nimg)
		vals = stimulus[idx,1]
		stimulus[idx,1] = stimulus[idx.-1,1]
		stimulus[idx.-1,1] = stimulus[idx,1] .+ 1
		tag = "shufflenovelty" * tag

	end # if pretraining

end # if lastshuffled

# reduce novelty input by a certain factor a different one for each sequence
# to infer the relevance of the novelty stimulus strength
lastimages = collect(Nimg:Nimg:Nimg*Nseq)
reducefactor =  collect(0:(1/(Nseq-1)):1)
secondtolastimage = lastimages.-1

if reducednovelty
	if pretrainig

		idx = findall(stimulus[lenpretrain + 1:end,1] .> Nseq*Nimg)
		vals = stimulus[idx.+(lenpretrain),1]
		for i in idx
			for im = 1:length(secondtolastimage)
				if stimulus[i+(lenpretrain-1),1] == secondtolastimage[im]# if previous image is certain second to last image of sequence reduce by corresponding factor
					stimulus[i+(lenpretrain),4] = stimulus[i+(lenpretrain),4]*reducefactor[im]
				end
			end
		end
		tag = "reducednovelty$(reducefactor)" * tag
	else
		idx = findall(stimulus[:,1] .> Nseq*Nimg)
		vals = stimulus[idx,1]
		for i in idx
			for im = 1:length(secondtolastimage)
				if stimulus[i-1,1] == secondtolastimage[im]# if previous image is certain second to last image of sequence reduce by corresponding factor
					stimulus[i,4] = stimulus[i,4]*reducefactor[im]
				end
			end
		end
		tag = "reducednovelty$(reducefactor)" * tag
	end

end

# simulation run time
T = stimulus[end,3]+lenpause # last stimulation time + pause duration
println(T)

println("Simulation run time: $T")

# --------------------- initialise savefile ---------------------------------------------------------

# initialise savefile and avoid overwriting when identical parameters are used
datetime = Dates.format(Dates.now(), "yyyymmdd-HHMMSS")

filesavename = "seq_violation_$(inhibfactor)_dur$(T)msNblocks$(Nblocks)Ntrain$(Ntrain)lenstim$(lenstim)lenpause$(lenpause)Nreps$(Nreps)strength$(strength)wadapt$(ifwadapt)iSTDP$(ifiSTDP)RateAdjust$(adjustfactor)Time"

savefile = "../data/"*filesavename * datetime * tag
println(savefile)

# --------------------- initialise weights and assemblymembers -----------------------------------------

weights = initweights(weightpars(Ne = Ne, Ni = Ni))
# excitatory tuning
assemblymembers = initassemblymembers(Nassemblies = Nass,Ne = Ne)
# inhibitory tuning
inhibassemblies = initinhibassemblymembers(Nassemblies = Nass, Ne = Ne, Ni = Ni)

const assembly_vectors =
    map(eachrow(hcat(assemblymembers,inhibassemblies))) do r
    sort(filter(>(0),r))
end

# --------------------- store relevant initialisation parameters -------------------------------------------
#=
h5write(savefile, "initial/stimulus", stimulus)
h5write(savefile, "initial/lengthpretrain", lenpretrain)#
h5write(savefile, "initial/stimparams", stimparams)
h5write(savefile, "initial/stimparams_prestim", stimparams_prestim)
h5write(savefile, "initial/seqnumber", seqnumber)
h5write(savefile, "initial/idxblockonset", blockonset)
h5write(savefile, "initial/weights", weights)
h5write(savefile, "initial/assemblymembers", assemblymembers)
h5write(savefile, "initial/inhibassemblies", inhibassemblies)

h5write(savefile, "params/T", T)
h5write(savefile, "params/Ne", Ne)
h5write(savefile, "params/Ni", Ni)

=#

# Define storage decisions
Ndec = 1
storagedec = zeros(Bool,Ndec)
storagetimes = ones(Int,Ndec)*1000

# Stroage decisions
storagedec[1] = true  # store times when spikes occured

# --------------------- initialise spiketime matrix ID - t  -----------------------------------------

Tstore = 1000 # average duration of one spiketime matrix
avgrate = 10 # Hz
Nspikesmax = Ncells*Tstore*avgrate/1000 # Nr. neurons x run time in seconds x avg. firing rate
spiket = zeros(Int32,Int(Nspikesmax),2)

if Ntrain == 0 # switch back to 1 as it is used in run simulation julia indexing starts at 1
	lenpretrain = 1
end
# --------------------- run simulation ---------------------------------------------------------------
##

##############################################################################
#
# This code is part of the publication:
# https://www.biorxiv.org/content/10.1101/2020.11.30.403840v1
#
# The generation of cortical novelty responses through inhibitory plasticity
# Auguste Schulz*, Christoph Miehl*, Michael J. Berry II, Julijana Gjorgjieva
#
# * equal contribution
#
##############################################################################

# Additional information and instruction in README

function runsimulation_inhibtuning(ifiSTDP,ifwadapt,stimparams,stimulus::Array{Float64,2},
	weights::Array{Float64,2}, assemblymembers::Array{Int64,2},
	spiketimes::Array{Int32,2},storagedec::Array{Bool,1},
	storagetimes::Array{Int64,1}, savefile::String, lenpretrain,
	inhibassemblies::Array{Int64,2}; dt = 0.1, T = 2000,
	adjustfactor = 1, adjustfactorinhib=1, inhibtuing = true, inhibfactor = 0.1,
	bwfactor=100, tauw_adapt=150,v_start::Union{Nothing,Vector{Float64}}=nothing,
  psth_my_idxs=collect(1:100))

"""	Runs a new simulation of a plastic E-I spiking neural network model where both E and I
	neurons receive tuned input

	Inputs:
		ifiSTDP boolean: if inhibitory spike timing plasticity is active
		ifwadapt boolean: if an intrisic adaptive current is included for E neurons
		stimparams array: with stimulus parameters
		stimulus array: stimulus constisting of [stimulated assembly, start time, stop time, rate increase]
		weights array: initial weight matrix
		assemblymembers array: members of an excitatory assembly (group of E neurons tuned to the same stimulus)
		spiketimes array: spiketimes of all neuron [spike time, neuron id] (in julia speed is increased by prior initialisation)
		storagedec array of booleans: storage decisions what should be stored
		storagetimes array: how often each of the items should be stored
		savefile string: name of the file to store to
		lenpretrain int: duration of the pretraining phase prior to starting the real simulation in stimulus indices
		inhibassemblies array: members of an inhibitory "assembly" (group of I neurons tuned to the same stimulus)
		dt float: integration timestep (optional)
		T float: total duration of the run (optional)
		adjustfactor float: adjust factor of the E-to-E plasticity after pretraining to allow for switching off plasticity
		adjustfactorinhib float: adjust factor of the I-to-E plasticity after pretraining to allow for switching off plasticity
		inhibtuing boolean: if inhibitory neurons are tuned to stimuli as well
		inhibfactor float: the scaling factor by how much the inhibitory neurons are driven by an external stimulus compared to the excitatory tuned neurons
		bwfactor float: scaling factor of the adaptive currents (only required for adaptive experiments)
		tauw_adapt float: timescale of the adaptive currents (only required for adaptive experiments)

	Output:

		totalspikes array: set of last spikes
		several run parameters are stored in hdf5 files including all spiketimes of all neurons across the entire simulation
		"""

    # MODIFICATION TO SAVE E AND I PSTH

		# Naming convention
		# e corresponds to excitatory
		# i corresponds to inhibitory
		# x corresponds to external

		#membrane dynamics
		taue = 20 #e membrane time constant ms
		taui = 20 #i membrane time constant ms
		vreste = -70 #e resting potential mV
		vresti = -62 #i resting potential mV
		vpeak = 20 #cutoff for voltage.  when crossed, record a spike and reset mV
		eifslope = 2 #eif slope parameter mV
		C = 300 #capacitance pF
		erev = 0 #e synapse reversal potential mV
		irev = -75 #i synapse reversal potntial mV
		vth0 = -52 #initial spike voltage threshold mV
		thrchange = false # can be switched off to have vth constant at vth0
		ath = 10 #increase in threshold post spike mV
		tauth = 30 #threshold decay timescale ms
		vreset = -60 #reset potential mV
		taurefrac = 1 #absolute refractory period ms
		aw_adapt = 4 #adaptation parameter a nS conductance
		bw_adapt = bwfactor*0.805 #adaptation parameter b pA current
    #=
		if T > 1 # this should not be saved for the precompile run
			h5write(savefile, "params/MembraneDynamics", [taue,taui, vreste,vresti , vpeak , eifslope, C, erev, irev, vth0, thrchange, ath, tauth, vreset, taurefrac, aw_adapt,bw_adapt, tauw_adapt ])
			h5write(savefile, "params/STDPwadapt", [Int(ifiSTDP),Int(ifwadapt)])
			h5write(savefile, "params/bwfactor", bwfactor)
			h5write(savefile, "params/adjustfactor", adjustfactor)
			h5write(savefile, "params/adjustfactorinhib", adjustfactorinhib)

		end
    =#
		# total number of neurons
		Ncells = Ne+Ni

		# synaptic kernel
		tauerise = 1 #e synapse rise time
		tauedecay = 6 #e synapse decay time
		tauirise = .5 #i synapse rise time
		tauidecay = 2 #i synapse decay time

		# external input
		rex = 4.5 #external input rate to e (khz) since the timestep is one ms an input of 4.5 corresp to 4.5kHz
		rix = 2.5#2.25 #external input rate to i (khz) # on average this is reduced to 2.25
		println("Larger inhibitory input")
		# Ensure the overall inhibitory input remains the same when inhibiory tuning is included

		# ---------------- inhibitory tuning ------------------
		# reduce the overall inhibitory input by the amount added during each stimulus presentation
		memsinhib1 = inhibassemblies[1,inhibassemblies[1,:] .!= -1]
		# reduce the total inhibition based on the total added inhibition to stimulated inhibitory neurons
		added_inhib = length(memsinhib1)*stimulus[end,4]*inhibfactor
		println("initial rix ", rix)
		rix -= added_inhib/Ni
		println("reduced rix ", rix)

		# initial connectivity
		Jeemin = 1.78 #minimum ee weight  pF
		Jeemax = 21.4 #maximum ee weight pF

		Jeimin = 48.7 #minimum ei weight pF
		Jeimax = 243 #maximum ei weight pF

		Jex = 1.78 #external to e weight pF
		Jix = 1.27 #external to i weight pF
    #=
		if T > 1
			h5write(savefile, "params/Connectivity", [tauerise, tauedecay, tauirise, tauidecay, rex, rix, Jeemin, Jeemax, Jeimin, Jeimax, Jex, Jix])
		end
    =#
		#voltage based stdp (for alternative testing not used here)
		altd = .0008 #ltd strength pA/mV pairwise STDP LTD
		altp = .0014 #ltp strength pA/mV^2 triplet STDP LTP
		thetaltd = -70 #ltd voltage threshold mV
		thetaltp = -49 #ltp voltage threshold mV
		tauu = 10 #timescale for u variable ms
		tauv = 7 #timescale for v variable ms
		taux = 15 #timescale for x variable ms
    #=
		if T > 1
			h5write(savefile, "params/voltageSTDP", [altd, altp, thetaltd, thetaltp, tauu, tauv, taux])
		end
      =#

		#inhibitory stdp
		tauy = 20 #width of istdp curve ms
		eta = 1 #istdp learning rate pA
		r0 = .003 #target rate (khz)
    #=
		if T > 1
			h5write(savefile, "params/iSTDP", [tauy, eta, r0])
		end
    =#

		# triplet parameters
		tripletrule = false # true
		o1 = zeros(Float64,Ne);
		o2 = zeros(Float64,Ne);
		r1 = zeros(Float64,Ne);
		r2 = zeros(Float64,Ne);
    #=
		if T > 1
			h5write(savefile, "params/Triplet", Int(tripletrule))
		end
    =#
		tau_p = 16.8;        # in ms
		tau_m = 33.7;        # in s
		tau_x = 101.0;        # in s
		tau_y = 125.0;        # in s
		# init LTP and LTD variables
		A_2p = 7.5*10^(-10); # pairwise LTP disabled
		A_2m = 7.0*10^(-3);
		A_2m_eff = A_2m;       #effective A_2m, includes the sliding threshold
		A_3p = 9.3*10^(-3)
		A_3m = 2.3*10^(-4); # triplet LTP disabled
    #=
		if T > 1
			h5write(savefile, "params/TripletTausAs", [tau_p,tau_m, tau_x,tau_y , A_2p , A_2m, A_3p, A_3m])
		end
    =#
		# simulation parameters
		taurefrac = 1 #ms refractory preriod clamped for 1 ms
		dtnormalize = 20 #how often to normalize rows of ee weights ms heterosynaptic plasticity
		stdpdelay = 1000 #time before stdp is activated, to allow transients to die out ms
		dtsaveweights = 2000  # save weights  every 2000 ms
		# minimum and maximum of storing the weight matrices
		minwstore = 80
		modwstore = 10
    #=
		if T > 1
			h5write(savefile, "params/Normdt", dtnormalize)
			h5write(savefile, "params/dtsaveweights", dtsaveweights)
			h5write(savefile, "params/minwstore", minwstore)
			h5write(savefile, "params/modwstore", modwstore)
		end
    =#
		Nassemblies = size(assemblymembers,1) #number of assemblies
		Nmembersmax = size(assemblymembers,2) #maximum number of neurons in a population

		ttstimbegin = round(Integer,stimulus[lenpretrain,2]/dt) # is set to zero

		# stimulus parameters
		Nimg, Nreps, Nseq, Nblocks, stimstart, lenstim, lenpause, strength = stimparams

		# Initialisation of spike arrays ------------------------------

		# spike count and spike times

		if storagedec[1]
			totalspikes = zeros(Int,Ncells)
		else
			spiketimes = nothing
			totalspikes = nothing
			totsp = nothing
		end
		totsp::Int64 = 0; # total numberof spikes
		spmax::Int64 = size(spiketimes,1); # maximum recorded spikes

		# further arrays for storing inputs and synaptic integration
		forwardInputsE = zeros(Float64,Ncells) #sum of all incoming weights from excitatpry inputs both external and EE and IE
		forwardInputsI = zeros(Float64,Ncells) #sum of all incoming weights from inhibitory inputs both II and EI
		forwardInputsEPrev = zeros(Float64,Ncells) #as above, for previous timestep
		forwardInputsIPrev = zeros(Float64,Ncells)

		xerise = zeros(Float64,Ncells) #auxiliary variables for E/I currents (difference of exponentials)
		xedecay = zeros(Float64,Ncells)
		xirise = zeros(Float64,Ncells)
		xidecay = zeros(Float64,Ncells)

		expdist = Exponential()

		v = zeros(Float64,Ncells) #membrane voltage
		nextx = zeros(Float64,Ncells) #time of next external excitatory input
		sumwee0 = zeros(Float64,Ne) #initial summed e weight, for normalization
		Nee = zeros(Int,Ne) #number of e->e inputs, for normalization
		rx = zeros(Float64,Ncells) #rate of external input

		# initialisation of membrane potentials and poisson inputs
		for cc = 1:Ncells
      _v_init_rand = vreset + (vth0-vreset)*rand()
			v[cc] = something(v_start[cc],_v_init_rand)
			if cc <= Ne # excitatory neurons
				rx[cc] = rex
				nextx[cc] = rand(expdist)/rx[cc]
				for dd = 1:Ne
					sumwee0[cc] += weights[dd,cc]
					if weights[dd,cc] > 0
						Nee[cc] += 1
					end
				end
			else # inhibtory neurons
				rx[cc] = rix
				nextx[cc] = rand(expdist)/rx[cc]
			end
		end

		vth = vth0*ones(Float64,Ncells) #adaptive threshold
		wadapt = aw_adapt*(vreset-vreste)*ones(Float64,Ne) #adaptation current
		lastSpike = -100*ones(Ncells) #last time the neuron spiked
		trace_istdp = zeros(Float64,Ncells) #low-pass filtered spike train for istdp
		u_vstdp = vreset*zeros(Float64,Ne) # for voltage rule (not used here)
		v_vstdp = vreset*zeros(Float64,Ne) # for voltage rule (not used here)
		x_vstdp = zeros(Float64,Ne) # for voltage rule (not used here)

		# ---------------------------------- set up storing avg weights -----------------

		idxnovelty = zeros(Int32, 200)

		Nass = Nassemblies # number of distinct stimuli
		nm = zeros(Int32,Nass) # Number of memebers in this assembly
		inhibnm = zeros(Int32,Nass) # Number of memebers in this inhibitory "assembly"
		for mm = 1:Nass # loop over all assemblies
			#check when first -1 comes to determine number of neurons
			nm[mm] = sum(assemblymembers[mm,:].!=-1)
			inhibnm[mm] = sum(inhibassemblies[mm,:].!=-1)
		end

		# initialise avergae cross and inhbitiory assembly weights
		# Calculate them here instead of storing whole matrices (too large files)
		avgXassembly = zeros(Float64,Nass,Nass) # average cross assembly weights
		avgInhibXassembly = zeros(Float64,Nass,Nass) # average inhibitory to excitatory cross assembly weights
		avgItoassembly = zeros(Float64,Nass) # avergate I to E assembly weights
		avgnonmemstoassembly = zeros(Float64,Nass) # avergate weights from neurons not stimulus driven to exc. assemblies
		avgassemblytononmems = zeros(Float64,Nass) # avergate weights from exc. assemblies to neurons not stimulus driven
		avgassemblytonovelty = zeros(Float64,Nass) # avergate weights from exc. assemblies to novelty assemblies
		avgnoveltytoassembly = zeros(Float64,Nass) # avergate weights from exc. novelty assemblies to exc. assemblies

		# determine neurons not part of any excitatory assembly
		nonmems = collect(1:Ne)
		members = sort(unique(assemblymembers[assemblymembers .> 0]))
		deleteat!(nonmems, members)

		# total number of simulation steps, steps when to normalise, and save
		Nsteps = round(Int,T/dt)
		inormalize = round(Int,dtnormalize/dt)
		saveweights = round(Int,dtsaveweights/dt)

		# true time
		t::Float64 = 0.0
		tprev::Float64 = 0.0

		# counters how often a variable was saved
		weightstore::Integer = 0
		spikestore::Integer = 0
		novelstore::Integer = 0

		# bool counter if a neuron has just had a spike
		spiked = zeros(Bool,Ncells)

  # MODIFICATION TO SAVE E AND I PSTH
  psth_e = Vector{Int64}(undef,Nsteps)
  psth_i = Vector{Int64}(undef,Nsteps)
  psth_myidx = Vector{Int64}(undef,Nsteps)
  fill!(psth_e,0) ; fill!(psth_i,0) ; fill!(psth_myidx,0)

	#   ------------------------------------------------------------------------
	#
	#   				begin actual simulation
	#
	#   ------------------------------------------------------------------------

	for tt = 1:Nsteps

				if mod(tt,Nsteps/100) == 1  #print percent complete
					print("\r",round(Int,100*tt/Nsteps))
				end

				forwardInputsE[:] .= 0.
				forwardInputsI[:] .= 0.
				t = dt*tt
				tprev = dt*(tt-1)

				# iterate over all stimuli in the passed stimulus array
				for ss = 1:size(stimulus)[1]

					if (tprev<stimulus[ss,2]) && (t>=stimulus[ss,2])  #just entered stimulation period
						ass = round(Int,stimulus[ss,1]) # TODO: refactor somewhat unfortunate naming of assembly

						# if assembly is a novelty assembly
						if ass > Nimg*Nseq && (t>stimulus[lenpretrain,2])
							# name idxnovelty still stems from time before defining novelty assembly directly just novelty assembly
							idxnovelty = assemblymembers[ass,assemblymembers[ass,:] .!= -1]
							novelstore += 1
							println("Ass :$(ass) time $(t) > $(stimulus[lenpretrain,2])")
							println("novelstore :$(novelstore)")

							# h5write(savefile, "novelty/indices$(novelstore)", assemblymembers[ass,assemblymembers[ass,:] .!= -1])
							#h5write(savefile, "novelty/assembly$(novelstore)", ass)
						end

						# stimulus tuning: increase the external stimulus to rx of the corresponding E assembly memebers
						mems = assemblymembers[ass,assemblymembers[ass,:] .!= -1]
						rx[mems] .+= stimulus[ss,4]

						# ---------------------- inhibtuning --------------------
						if inhibtuing
							memsinhib = inhibassemblies[ass,inhibassemblies[ass,:] .!= -1]
							rx[memsinhib] .+= stimulus[ss,4]*inhibfactor
						end

					end # if just entered stim period

					if (tprev<stimulus[ss,3]) && (t>=stimulus[ss,3]) #just exited stimulation period
						ass = round(Int,stimulus[ss,1])

						mems = assemblymembers[ass,assemblymembers[ass,:] .!= -1]
						rx[mems] .-= stimulus[ss,4]

						if inhibtuing
							memsinhib = inhibassemblies[ass,inhibassemblies[ass,:] .!= -1]
							rx[memsinhib] .-= stimulus[ss,4]*inhibfactor
						end

					end # if just left stim period

				end #end loop over stimuli

				if mod(tt,inormalize) == 0 #excitatory synaptic normalization
					for cc = 1:Ne
						sumwee = 0.
						for dd = 1:Ne
							sumwee += weights[dd,cc]
						end # dd

						for dd = 1:Ne
							if weights[dd,cc] > 0.
								weights[dd,cc] -= (sumwee-sumwee0[cc])/Nee[cc]
								if weights[dd,cc] < Jeemin
									weights[dd,cc] = Jeemin
								elseif weights[dd,cc] > Jeemax
									weights[dd,cc] = Jeemax
								end # if
							end # if
						end # dd for
					end # cc for
					#println("Normalised...")
				end #end normalization

        #=
				if mod(tt,saveweights) == 0 #&& (weightstore < 20 || mod(weightstore,100) == 0)#excitatory synaptic normalization
					weightstore += 1
					#if (weightstore < minwstore || mod(weightstore,modwstore) == 0)
						#h5write(savefile, "dursim/weights$(weightstore)_$(tt)", weights)
						for pre = 1:Nass # use maximum to be sure to capture it both for sequencelength and variablerepetitions
							# for inhib to assembly get avg. weights here pre is actually post
							avgnonmemstoassembly[pre] = getXassemblyweight(nonmems, assemblymembers[Int(pre),1:nm[Int(pre)]], weights)
							avgassemblytononmems[pre] = getXassemblyweight(assemblymembers[Int(pre),1:nm[Int(pre)]],nonmems, weights)
							if sum(idxnovelty) == 0 # before first novelty arose still zero
								avgnoveltytoassembly[pre] = 0
								avgassemblytonovelty[pre] = 0
							else
								avgnoveltytoassembly[pre] = getXassemblyweight(idxnovelty[idxnovelty.>0], assemblymembers[Int(pre),1:nm[Int(pre)]], weights)
								avgassemblytonovelty[pre] = getXassemblyweight(assemblymembers[Int(pre),1:nm[Int(pre)]],idxnovelty[idxnovelty.>0], weights)
							end
							avgItoassembly[pre] = getXassemblyweight(collect(Ne+1:Ncells), assemblymembers[Int(pre),1:nm[Int(pre)]], weights)
							for post = 1:Nass
								avgXassembly[pre,post] = getXassemblyweight(assemblymembers[Int(pre),1:nm[Int(pre)]], assemblymembers[Int(post),1:nm[Int(post)]], weights)
								# determine the average inhibitory assembly to excitatory assembly weight  define inhibnm
								avgInhibXassembly[pre,post] = getXassemblyweight(inhibassemblies[Int(pre),1:inhibnm[Int(pre)]], assemblymembers[Int(post),1:nm[Int(post)]], weights)

							end # post assemblies
						end # pre assemblies

            #=
						h5write(savefile, "dursimavg/avgXassembly$(weightstore)_$(tt)", avgXassembly)
						h5write(savefile, "dursimavg/avgInhibXassembly$(weightstore)_$(tt)", avgInhibXassembly)

						h5write(savefile, "dursimavg/avgItoassembly$(weightstore)_$(tt)", avgItoassembly)
						h5write(savefile, "dursimavg/avgnonmemstoassembly$(weightstore)_$(tt)", avgnonmemstoassembly)
						h5write(savefile, "dursimavg/avgassemblytononmems$(weightstore)_$(tt)", avgassemblytononmems)
						h5write(savefile, "dursimavg/avgassemblytonovelty$(weightstore)_$(tt)", avgassemblytonovelty)
						h5write(savefile, "dursimavg/avgnoveltytoassembly$(weightstore)_$(tt)", avgnoveltytoassembly)

						h5write(savefile, "dursimavg/Itoneuron1$(weightstore)_$(tt)", weights[Ne+1:Ncells,1][weights[Ne+1:Ncells,1].>0])# store only non zero inhibitory to neuron 1/2 weights
						h5write(savefile, "dursimavg/Itoneuron2$(weightstore)_$(tt)", weights[Ne+1:Ncells,2][weights[Ne+1:Ncells,2].>0]) # store only non zero inhibitory to neuron 1/2 weights
            =#

				end # mod(tt,saveweights) == 0
        =#
				fill!(spiked,zero(Bool)) # reset spike bool without new memory allocation

				for cc = 1:Ncells
					trace_istdp[cc] -= dt*trace_istdp[cc]/tauy

					while(t > nextx[cc]) #external input
						nextx[cc] += rand(expdist)/rx[cc]
						if cc <= Ne
							forwardInputsEPrev[cc] += Jex
						else
							forwardInputsEPrev[cc] += Jix
						end
					end

					xerise[cc] += -dt*xerise[cc]/tauerise + forwardInputsEPrev[cc]
					xedecay[cc] += -dt*xedecay[cc]/tauedecay + forwardInputsEPrev[cc]
					xirise[cc] += -dt*xirise[cc]/tauirise + forwardInputsIPrev[cc]
					xidecay[cc] += -dt*xidecay[cc]/tauidecay + forwardInputsIPrev[cc]

					if cc <= Ne # excitatory
						if thrchange
						vth[cc] += dt*(vth0 - vth[cc])/tauth;
						end
						wadapt[cc] += dt*(aw_adapt*(v[cc]-vreste) - wadapt[cc])/tauw_adapt;
						u_vstdp[cc] += dt*(v[cc] - u_vstdp[cc])/tauu;
						v_vstdp[cc] += dt*(v[cc] - v_vstdp[cc])/tauv;
						x_vstdp[cc] -= dt*x_vstdp[cc]/taux;

						# triplet accumulators
						r1[cc] += -dt*r1[cc]/tau_p # exponential decay of all
			            r2[cc] += -dt*r2[cc]/tau_x
						o1[cc] += -dt*o1[cc]/tau_m
			            o2[cc] += -dt*o2[cc]/tau_y
					end

					if t > (lastSpike[cc] + taurefrac) #not in refractory period
						# update membrane voltage

						ge = (xedecay[cc] - xerise[cc])/(tauedecay - tauerise);
						gi = (xidecay[cc] - xirise[cc])/(tauidecay - tauirise);

						if cc <= Ne #excitatory neuron (eif), has adaptation
							if ifwadapt
								dv = (vreste - v[cc] + eifslope*exp((v[cc]-vth[cc])/eifslope))/taue + ge*(erev-v[cc])/C + gi*(irev-v[cc])/C - wadapt[cc]/C;
							else
								dv = (vreste - v[cc] + eifslope*exp((v[cc]-vth[cc])/eifslope))/taue + ge*(erev-v[cc])/C + gi*(irev-v[cc])/C;
							end
							v[cc] += dt*dv;
							if v[cc] > vpeak
								spiked[cc] = true
								wadapt[cc] += bw_adapt
							end
						else
							dv = (vresti - v[cc])/taui + ge*(erev-v[cc])/C + gi*(irev-v[cc])/C;
							v[cc] += dt*dv;
							if v[cc] > vth0
								spiked[cc] = true
							end
						end

						if spiked[cc] #spike occurred
							spiked[cc] = true;
							v[cc] = vreset;
							lastSpike[cc] = t;
							totalspikes[cc] += 1;
              # SAVE PSTH
              if cc in psth_my_idxs
                psth_myidx[tt] += 1
              end
              if cc <= Ne
                psth_e[tt] += 1
              else
                psth_i[tt] += 1
              end
							totsp += 1;
							if totsp < spmax
								spiketimes[totsp,1] = tt; # time index as a sparse way to save spiketimes
								spiketimes[totsp,2] = cc; # cell id
							elseif totsp == spmax
								spiketimes[totsp,1] = tt; # time index
								spiketimes[totsp,2] = cc; # cell id

								totsp = 0 # reset counter total number of spikes
								# store spiketimes
								spikestore += 1
								#h5write(savefile, "dursimspikes/spiketimes$(spikestore)", spiketimes)
							end

							trace_istdp[cc] += 1.;
							if cc<=Ne
								x_vstdp[cc] += 1. / taux;
							end

							if cc <= Ne && thrchange # only change for excitatory cells and when thrchange == true
								vth[cc] = vth0 + ath;
							end

							#loop over synaptic projections
							for dd = 1:Ncells # postsynaptic cells dd  - cc presynaptic cells
								if cc <= Ne #excitatory synapse
									forwardInputsE[dd] += weights[cc,dd];
								else #inhibitory synapse
									forwardInputsI[dd] += weights[cc,dd];
								end
							end

						end #end if(spiked)
					end #end if(not refractory)

					if ifiSTDP # select if iSTDP
						#istdp
						if spiked[cc] && (t > stdpdelay)
							if cc <= Ne #excitatory neuron fired, potentiate i inputs
								for dd = (Ne+1):Ncells
									if weights[dd,cc] == 0.
										continue
									end
									weights[dd,cc] += eta*trace_istdp[dd]

									if weights[dd,cc] > Jeimax
										weights[dd,cc] = Jeimax
									end
								end
							else #inhibitory neuron fired, modify outputs to e neurons
								for dd = 1:Ne
									if weights[cc,dd] == 0.
										continue
									end

									weights[cc,dd] += eta*(trace_istdp[dd] - 2*r0*tauy)
									if weights[cc,dd] > Jeimax
										weights[cc,dd] = Jeimax
									elseif weights[cc,dd] < Jeimin
										weights[cc,dd] = Jeimin
									end
								end
							end
						end #end istdp
					end # ifiSTDP

					if tripletrule

					#triplet, ltd component
					if spiked[cc] && (t > stdpdelay) && (cc <= Ne)
						r1[cc] = r1[cc] + 1 # incrememt r1 before weight update
						for dd = 1:Ne #depress weights from cc to all its postsyn cells
							# cc = pre dd = post
							if weights[cc,dd] == 0. # ignore connections that were not establishe in the beginning
								continue
							end

			                weights[cc,dd] -= o1[dd]*(A_2m + A_3m*r2[cc])

							if weights[cc,dd] < Jeemin
								weights[cc,dd] = Jeemin
							end

						end # for loop over Ne
						 r2[cc] = r2[cc] + 1 # increment after weight update
					end # ltd

					#triplet, ltp component
					if spiked[cc] && (t > stdpdelay) && (cc <= Ne)
						o1[cc] = o1[cc] + 1 # incrememt r1 before weight update
						# cc = post dd = pre
						for dd = 1:Ne #increase weights from cc to all its presyn cells dd
							if weights[dd,cc] == 0.
								continue
							end

							weights[dd,cc] += r1[dd]*(A_2p + A_3p*o2[cc]) #A_2p = 0

							if weights[dd,cc] > Jeemax
								weights[dd,cc] = Jeemax
							end

						end # loop over cells presynaptic
						o2[cc] = o2[cc] + 1 # increment after weight update

					end #ltp
        #=
				else # REMOVED not triplet but voltage rule

				#vstdp, ltd component
					if spiked[cc] && (t > stdpdelay) && (cc < Ne)
						for dd = 1:Ne #depress weights from cc to cj
							if weights[cc,dd] == 0.
								continue
							end

							if u_vstdp[dd] > thetaltd
								weights[cc,dd] -= altd*(u_vstdp[dd]-thetaltd)

								if weights[cc,dd] < Jeemin
									weights[cc,dd] = Jeemin

								end
							end
						end
					end #end ltd

					#vstdp, ltp component
					if (t > stdpdelay) && (cc <= Ne) && (v[cc] > thetaltp) && (v_vstdp[cc] > thetaltd)
						for dd = 1:Ne
							if weights[dd,cc] == 0.
								continue
							end

							weights[dd,cc] += dt*altp*x_vstdp[dd]*(v[cc] - thetaltp)*(v_vstdp[cc] - thetaltd)
							if weights[dd,cc] > Jeemax
								weights[dd,cc] = Jeemax
							end
						end
					end #end ltp
        =#
				end # if triplet rule
			end #end loop over cells
			forwardInputsEPrev = copy(forwardInputsE)
			forwardInputsIPrev = copy(forwardInputsI)

			# once the actual stimulation begins, eventually readjust the learning rate
			if tt == ttstimbegin

						tau_p = 16.8;        # in ms
						tau_m = 33.7;        # in s
						tau_x = 101.0;        # in s
						tau_y = 125.0;        # in s
						# init LTP and LTD variables
						A_2p = adjustfactor*7.5*10^(-10); # pairwise LTP disabled
						A_2m = adjustfactor*7.0*10^(-3);  #small learning rate
						A_2m_eff = A_2m;       #effective A_2m, includes the sliding threshold
						A_3p = adjustfactor*9.3*10^(-3)
						A_3m = adjustfactor*2.3*10^(-4); # triplet LTP disabled

						eta_old = eta
						eta = adjustfactorinhib*eta_old
						println("new eta iSTDP",eta)
            #=
						if T > 1
							h5write(savefile, "params/TripletTausAs_stim", [tau_p,tau_m, tau_x,tau_y , A_2p , A_2m, A_3p, A_3m])
							h5write(savefile, "params/eta_stim", eta)

							h5write(savefile, "params/adjustfactor2", adjustfactor)
							h5write(savefile, "params/adjustfactorinhib2", adjustfactorinhib)

							h5write(savefile, "dursim/weights", weights)

						end # if T > 1
            =#
			end # if tt == ttstimbegin
		end # tt loop over time
	print("\r")

	println("simulation finished")

	return totalspikes,psth_e,psth_i,psth_myidx
end # function simulation

using StatsBase
function do_psth_au(vals::Vector{<:Real},times_psth::AbstractVector{<:Real};dt=0.1E-3)
  times = dt/2:dt:Ttot-dt/2
  ret = fill(0.0,length(times_psth)-1)
  for (v,tt) in zip(vals,times)
    i = searchsortedfirst(times_psth,tt)-1
    ret[i] += v
  end
  dt_psth = times_psth[2] - times_psth[1]
  return ret ./ dt_psth
end

##

weights_start = permutedims(weights)

const Ttot = 10.0
const Ttot_ms = Int64(Ttot*1E3)

v_start = let vreset = -60.0 , vth0=-52.0
  doone =()->vreset + (vth0-vreset)*rand()
  [ doone() for i in 1:Ncells]
end

idxs_as1e = filter(<=(Ne),assembly_vectors[1])
N_as1e = length(idxs_as1e)

@time totalspikes,psth_e_au,psth_i_au,psth_myidx = runsimulation_inhibtuning(ifiSTDP,ifwadapt,stimparams,stimulus, weights,
  assemblymembers, spiket, storagedec, storagetimes,savefile, lenpretrain, inhibassemblies,
  adjustfactor = adjustfactor, adjustfactorinhib = adjustfactorinhib,
  inhibfactor = inhibfactor, T = Ttot_ms,v_start=v_start,psth_my_idxs=idxs_as1e)

totalrates_au = totalspikes ./ Ttot

weights_end = permutedims(weights)

# error("simulation done! STOPPP!")

const dt_psth = 50E-3

times_psth = let dtstep = dt_psth
  range(0,Ttot;step=dtstep)
end
psth_e_au_less =  do_psth_au(psth_e_au,times_psth) ./ Ne
psth_i_au_less =  do_psth_au(psth_i_au,times_psth) ./ Ni
psth_myidx_au_less =  do_psth_au(psth_myidx,times_psth) ./ N_as1e

#bar(midpoints(times_psth),psth_e_au_less)
#bar(midpoints(times_psth),psth_i_au_less)
# ---------------------- store final params ---------------------------------------------------------
#=
h5write(savefile, "postsim/weights", weights)
h5write(savefile, "postsim/spiketimes", spiket)
h5write(savefile, "postsim/totalspikes", totalspikes)
=#

# start evaluation
# include("initevalimmediate.jl")
##

# ---------------------- DO THINGS ---------------------------------------------------------

function _reldiff(a,b)
  if a+b == 0.0
    return 0.0
  else
    return (a-b)/(0.5(a+b))
  end
end

_ = let ws = weights_start, wn=weights_end
  di = wn .- ws
  direl  = map(ww->_reldiff(ww...),zip(wn,ws))
  # extrema(we[1:Ne,1:Ne] .- ws[1:Ne,1:Ne])
  #@show extrema(di[Ne+1:end,1:Ne])
  # heatmap(di[Ne+1:end,1:Ne])
  @show extrema(di)
  heatmap(di)
end

#bar(totalspikes)
#histogram(totalspikes)

## ---------------------- BUILD EQUIVALENT NETWORK -----------------------------------------

		# taue = 20 #e membrane time constant ms
		# taui = 20 #i membrane time constant ms
		# vreste = -70 #e resting potential mV
		# vresti = -62 #i resting potential mV
		# vpeak = 20 #cutoff for voltage.  when crossed, record a spike and reset mV
		# eifslope = 2 #eif slope parameter mV
		# C = 300 #capacitance pF
		# erev = 0 #e synapse reversal potential mV
		# irev = -75 #i synapse reversal potntial mV
		# vth0 = -52 #initial spike voltage threshold mV
		# thrchange = false # can be switched off to have vth constant at vth0
		# ath = 10 #increase in threshold post spike mV
		# tauth = 30 #threshold decay timescale ms
		# vreset = -60 #reset potential mV
		# taurefrac = 1 #absolute refractory period ms

		# synaptic kernel
		# tauerise = 1 #e synapse rise time
		# tauedecay = 6 #e synapse decay time
		# tauirise = .5 #i synapse rise time
		# tauidecay = 2 #i synapse decay time

		# external input
		# rex = 4.5 #external input rate to e (khz) since the timestep is one ms an input of 4.5 corresp to 4.5kHz
		# rix = 2.5#2.25 #external input rate to i (khz) # on average this is reduced to 2.25
    # initial rix 2.5
    # reduced rix 2.3044
		# Jex = 1.78 #external to e weight pF
		# Jix = 1.27 #external to i weight pF

# push!(LOAD_PATH, abspath(@__DIR__,".."))
using ProgressMeter
using SpikingRNNs; const global S = SpikingRNNs
using SparseArrays

dt = 0.1E-3
myτe = 20E-3 # seconds
myτi = 20E-3 # seconds
τrefr= 1E-3 # refractoriness
vth_e = 20.   # mV
vthexp = -52.0 # actual threshold for spike-generation
vth_i = vthexp
eifslope = 2.0
Cap = 300.0 #capacitance mF
v_rest_e = -70.0
v_rest_i = -62.0
v_rev_e = 0.0
v_rev_i = -75.0
v_reset_e = -60.0
v_reset_i = -60.0

# synaptic kernel
tauerise = 1E-3 #e synapse rise time
tauedecay = 6E-3 #e synapse decay time
taueplus,taueminus  = tauedecay, tauerise
tauirise = 0.5E-3 #i synapse rise time
tauidecay = 2E-3 #i synapse decay time
tauiplus,tauiminus = tauidecay,tauirise

# input parameters

in_rate_e = 4.5E3
in_rate_i = 2.3044E3 # 2.5 with correction
in_assembly_on = 12E3
Jin_e = 1.78
Jin_i = 1.27

##
nt_e = let sker = S.SKExpDiff(taueplus,taueminus)
  sgen = S.SpikeGenEIF(vthexp,eifslope)
  S.NTLIFConductance(sker,sgen,myτe,Cap,
   vth_e,v_reset_e,v_rest_e,τrefr,v_rev_e)
end
ps_e = S.PSLIFConductance(nt_e,Ne)

nt_i = let sker = S.SKExpDiff(tauiplus,tauiminus)
  sgen = S.SpikeGenNone()
  S.NTLIFConductance(sker,sgen,myτi,Cap,
   vth_i,v_reset_i,v_rest_i,τrefr,v_rev_i)
end
ps_i = S.PSLIFConductance(nt_i,Ni)

## Input : super easy! Just use the matrix `stimulus` to define a function!

function _is_in_assembly_old(neu::Integer,as::Integer)
  asrow = view(assemblymembers,as,:)
  return neu in asrow
end
function _is_in_assembly(neu_idx::Integer,as_vec::Vector{<:Integer})
  if as_vec[1] <= neu_idx <= as_vec[end]
    idx = searchsortedfirst(as_vec,neu_idx)
    return as_vec[idx] == neu_idx
  else
    return false
  end
end

const stimulus_sec = copy(stimulus)
stimulus_sec[:,[2,3]] ./= 1E3
stimulus_sec[:,4] .*= 1E3

function inputfun(t::R,i::Integer,low::R,
    stimulus_mat::Matrix{Float64},
    assembly_v::Vector{Vector{Int64}};scale_factor::Float64=1.0) where R<:Real
  # catch the first or last
  (!(stimulus_mat[1,2] <= t <= stimulus_mat[end,3])) && (return low)
  stimstarts = view(stimulus_mat,:,2)
  # find row number
  therow = searchsortedfirst(stimstarts,t) - 1
  # account for stim duration
  (t >= stimulus_mat[therow,3]) && (return low)
  # ok, the assembly is active!
  as_v = assembly_v[Int64(stimulus_mat[therow,1])]
  if _is_in_assembly(i,as_v)
    return low + scale_factor*stimulus_mat[therow,4] # add to stimulus
  else
    return low
  end
  return low # in case the compiler is dumb
end

function inputfun_upper(t::R,i::Integer,low::R,stimmat::Matrix{R},asv::Vector;
   scale_factor::Float64=1.0) where R<:Real
  return max(inputfun(t,i,low,stimmat,asv;scale_factor=scale_factor),
              inputfun(t+2dt,i,low,stimmat,asv;scale_factor=scale_factor))
end

nt_in_e = let low = in_rate_e,
  ratefun = (t,i) -> inputfun(t,i,low,stimulus_sec,assembly_vectors)
  ratefun_upper = (t,i) -> inputfun_upper(t,i,low,stimulus_sec,assembly_vectors)
  sker = S.SKExpDiff(taueplus,taueminus)
  sgen = S.SGPoissonFExact( ratefun,ratefun_upper )
  S.NTInputConductance(sgen,sker,v_rev_e)
end
ps_in_e = S.PSInputPoissonConductanceExact(nt_in_e,Jin_e,Ne)

nt_in_i = let inh_fact=inhibfactor,
  low = in_rate_i,
  ratefun = (t,i) -> inputfun(t,Ne+i,low,
    stimulus_sec,assembly_vectors;scale_factor=inh_fact) # add Ne to access inh indexes
  ratefun_upper = (t,i) -> inputfun_upper(t,Ne+i,low,stimulus_sec,assembly_vectors;
            scale_factor=inh_fact)
  sker = S.SKExpDiff(taueplus,taueminus)
  sgen = S.SGPoissonFExact(ratefun,ratefun_upper)
  S.NTInputConductance(sgen,sker,v_rev_e)
end
ps_in_i = S.PSInputPoissonConductanceExact(nt_in_i,Jin_i,Ni)

## plot input functions, out of curiosity

as_ordering = let Ntot=Ncells
  S.order_by_pattern_idxs(assembly_vectors,Ntot)
end

_ = let low_e = in_rate_e, low_i = in_rate_i,
  _scal = inhibfactor
  ratefun = function (t,i)
    if i<=Ne
      inputfun(t,i,low_e,stimulus_sec,assembly_vectors)
    else
      inputfun(t,i,low_i,stimulus_sec,assembly_vectors,scale_factor=_scal)
    end
  end
  times = range(0,60;length=200)
  nshow = Ncells
  neushow = as_ordering[1:nshow]
  vals = hcat([ ratefun.(t,neushow) for t in times]...)
  # vals_toplot = vals[as_ordering,:][1:nshow,:]
  heatmap(times,1:nshow,vals)
end

## connections
conn_ee = let w_ee = weights_start[1:Ne,1:Ne]
  S.ConnGeneralIF2(sparse(w_ee))
end
conn_ii = let w_ii = weights_start[Ne+1:end,Ne+1:end]
S.ConnGeneralIF2(sparse(w_ii))
end
conn_ei = let w_ei = weights_start[1:Ne,Ne+1:end]
  S.ConnGeneralIF2(sparse(w_ei))
end
conn_ie = let w_ie = weights_start[Ne+1:end,1:Ne]
  S.ConnGeneralIF2(sparse(w_ie))
end

## Populations

pop_e = S.Population(ps_e,(conn_ee,ps_e),(conn_ei,ps_i),
  (S.FakeConnection(),ps_in_e))

pop_i = S.Population(ps_i,(conn_ii,ps_i),(conn_ie,ps_e),
  (S.FakeConnection(),ps_in_i))
##
# that's it, let's make the network
myntw = S.RecurrentNetwork(dt,pop_e,pop_i)

# record spiketimes and internal potential
krec = 1
#n_e_rec = 1000
#n_i_rec = 1000
t_wup = 0.0
expected_rate = 15.0
#rec_state_e = S.RecStateNow(ps_e,krec,dt,Ttot;idx_save=collect(1:n_e_rec),t_warmup=t_wup)
#rec_state_i = S.RecStateNow(ps_i,krec,dt,Ttot;idx_save=collect(1:n_i_rec),t_warmup=t_wup)
rec_spikes_e = S.RecSpikes(ps_e,expected_rate,Ttot;idx_save=collect(1:Ne),t_warmup=t_wup)
rec_spikes_i = S.RecSpikes(ps_i,expected_rate,Ttot;idx_save=collect(1:Ni),t_warmup=t_wup)

## Run

times = (0:myntw.dt:Ttot)
nt = length(times)
# clean up
#S.reset!.([rec_state_e,rec_state_i])
S.reset!.([rec_spikes_e,rec_spikes_i])
S.reset!.([ps_e,ps_i])
S.reset!(conn_ei)
# initial conditions
ps_e.state_now .= v_start[1:Ne]
ps_i.state_now .= v_start[Ne+1:end]

@time begin
  @showprogress 5.0 "network simulation " for (k,t) in enumerate(times)
    #rec_state_e(t,k,myntw)
    #rec_state_i(t,k,myntw)
    rec_spikes_e(t,k,myntw)
    rec_spikes_i(t,k,myntw)
    S.dynamics_step!(t,myntw)
  end
end

# S.add_fake_spikes!(1.0vth_e,rec_state_e,rec_spikes_e)
# S.add_fake_spikes!(0.0,rec_state_i,rec_spikes_i)
##

rates_e = let rdic=S.get_mean_rates(rec_spikes_e,dt,Ttot)
  ret = fill(0.0,Ne)
  for (k,v) in pairs(rdic)
    ret[k] = v
  end
  ret
end

##
##

_ = let plt = plot()
  ts,psth_e = S.get_psth(collect(1:Ne),dt_psth,rec_spikes_e;Nneurons=Ne,Ttot=Ttot)
  bar!(plt,ts,psth_e;opacity=0.5,label="Dylan's")
  bar!(plt,ts,psth_e_au_less;opacity=0.5,title="excitation",label="Auguste's")
end

##
_ = let plt = plot()
  ts,psth_i = S.get_psth(collect(1:Ni),dt_psth,rec_spikes_i;Nneurons=Ni,Ttot=Ttot)
  bar!(plt,ts,psth_i;opacity=0.3,label="Dyl")
  bar!(plt,ts,psth_i_au_less;opacity=0.3,title="inhibition",lab="Aug")
end

##
_ = let plt = plot()
  ts,psth_myidx = S.get_psth(idxs_as1e,dt_psth,rec_spikes_e;Nneurons=Ni,Ttot=Ttot)
  bar!(plt,ts,psth_myidx;opacity=0.3,label="Dyl")
  bar!(plt,ts,psth_myidx_au_less;opacity=0.3,title="Assembly 1",lab="Aug")
end

##


all = [
    "eggs",
    "peanuts",
    "shellfish",
    "strawberries",
    "tomatoes",
    "chocolate",
    "pollen",
    "cats",
]

function allergic_to(score, allergen)
    return (score & (1 << (findfirst(x -> x==allergen, all) - 1)) ) != 0
end

function allergy_list(score)
    return filter(allergen -> allergic_to(score, allergen), all)
end


#=

Slice sampling for EGeoHiSSE

Ignacio Quintero Mächler

t(-_-t)

November 20 2017

=#

"""
    slice_sampler(tip_val     ::Dict{Int64,Array{Float64,1}},
                  ed          ::Array{Int64,2},
                  el          ::Array{Float64,1},
                  x           ::Array{Float64,1},
                  y           ::Array{Float64},
                  cov_mod     ::String,
                  out_file    ::String,
                  h           ::Int64;
                  constraints ::NTuple{N,String}  = (" ",),
                  niter       ::Int64             = 10_000,
                  nthin       ::Int64             = 10,
                  λpriors     ::Float64           = .1,
                  μpriors     ::Float64           = .1,
                  gpriors     ::Float64           = .1,
                  lpriors     ::Float64           = .1,
                  qpriors     ::Float64           = .1,
                  βpriors     ::NTuple{2,Float64} = (0.0, 5.0),
                  optimal_w   ::Float64           = 0.8,
                  screen_print::Int64             = 5) where {N}

Run slice-sampling Markov Chain given posterior function.
"""
function slice_sampler(lhf         ::Function,
                       p           ::Array{Float64,1},
                       fp          ::Array{Float64,1},
                       nnps        ::Array{Int64,1},
                       nps         ::Array{Int64,1},
                       phid        ::Array{Int64,1},
                       mvps        ::Array{Array{Int64,1},1},
                       nngps       ::Array{Array{Bool,1},1},
                       mvhfs       ::Array{Array{Int64,1},1},
                       npars       ::Int64,
                       niter       ::Int64,
                       nthin       ::Int64,
                       nburn       ::Int64,
                       ntakew      ::Int64,
                       winit       ::Float64,
                       optimal_w   ::Float64,
                       screen_print::Int64)

  # estimate optimal w
  p, fp, w =
    w_sampler(lhf, p, fp, nnps, nps, phid, mvps, nngps, mvhfs,
      npars, optimal_w, screen_print, nburn, ntakew, winit)

  # slice-sampler
  its, hlog, ps =
    loop_slice_sampler(lhf, p, fp, nnps, nps, phid, mvps, nngps, mvhfs,
      w, npars, niter, nthin, screen_print)

  # save samples
  R = hcat(its, hlog, ps)

  return R
end

# """
#     slice_sampler(tip_val    ::Dict{Int64,Array{Float64,1}},
#                   edges      ::Array{Int64,2},
#                   edlen      ::Array{Float64,1},
#                   out_file   ::String;
#                   constraints::String  = NaN,
#                   niter      ::Int64   = 10_000,
#                   nthin      ::Int64   = 10,
#                   model      ::String  = "musse",
#                   λpriors    ::Float64 = .1,
#                   μpriors    ::Float64 = .1,
#                   qpriors    ::Float64 = .1)

# Run slice-sampling Markov Chain for MuSSE model.
# """
# function slice_sampler(tip_val    ::Dict{Int64,Array{Float64,1}},
#                        edges      ::Array{Int64,2},
#                        edlen      ::Array{Float64,1},
#                        out_file   ::String;
#                        constraints::NTuple{N,String} = (" ",),
#                        niter      ::Int64            = 10_000,
#                        nthin      ::Int64            = 10,
#                        λpriors    ::Float64          = .1,
#                        μpriors    ::Float64          = .1,
#                        qpriors    ::Float64          = .1) where {N}

#   # k areas
#   k = length(tip_val[1])

#   # make ode and define parameters
#   mod_ode = make_musse(k)

#   npars = k + k*k
#   p = fill(0.2,npars)
#   δ = length(tip_val)/sum(edlen)
#   p[1:k]      .= δ + rand()    # set λs
#   p[(k+1):2k] .= p[1] - δ      # set μs

#   pardic = build_par_names(k)

#   # parameter update
#   pupd = Base.OneTo(npars)

#   #constraints
#   conp = set_constraints(constraints, pardic)

#   # remove contraints for being updated
#   pupd = setdiff(pupd, keys(conp))

#   # create likelihood, prior and posterior functions
#   llf = make_llf(tip_val, edges, edlen, mod_ode, p, sbrlen = sum(edlen))
#   lpf = make_lpf(λpriors, μpriors, qpriors, k)
#   lhf = make_lhf(llf, lpf, conp)

#   # set up slice-sampling
#   nlogs = fld(niter,nthin)
#   its   = zeros(Float64,nlogs)
#   h     = zeros(Float64,nlogs)
#   ps    = zeros(Float64,nlogs,npars)

#   lthin = 0::Int64
#   lit   = 0::Int64

#   # cat
#   printstyled("running MuSSE model \n", color=:green)

#   # estimate w
#   p, w = w_sampler(lhf, p, pupd, npars)

#   # start iterations
#   prog = Progress(niter, 5, "running slice-sampler....", 20)

#   hc = lhf(p)::Float64

#   for it in Base.OneTo(niter)

#     for j in pupd
#       S     = hc - Random.randexp()
#       L, R  = find_nonneg_int(p, j, S, lhf, w[j])
#       p, hc = sample_int(p, j, L, R, S, lhf)
#     end

#     # log samples
#     lthin += 1
#     if lthin == nthin
#       @inbounds begin
#         lit += 1
#         setindex!(its, it, lit)
#         setindex!(h,   hc, lit)
#         setindex!(ps,   p, lit, :)
#       end
#       lthin = 0
#     end

#     next!(prog)
#   end

#   # save samples
#   R = hcat(its, h, ps)

#   # column names
#   col_nam = ["Iteration", "Posterior"]

#   # add λ names
#   for i in 0:(k-1)
#     push!(col_nam, "lamdba$i")
#   end

#   # add μ names
#   for i in 0:(k-1)
#     push!(col_nam, "mu$i")
#   end

#   # add q names
#   for j in 0:(k-1), i in 0:(k-1)
#     if i == j
#       continue
#     end
#     push!(col_nam, "q$i$j")
#   end

#   R = vcat(reshape(col_nam, 1, lastindex(col_nam)), R)

#   writedlm(out_file*".log", R)

#   return R
# end


using HDF5
using Base.Test

@testset "readremote" begin

    # check that we can read the official HDF5 example files
    urlbase = "https://support.hdfgroup.org/ftp/HDF5/examples/files/exbyapi/"

    fcmp = [0 1 2 3 4 5 6;
        2 1.66667 2.4 3.28571 4.22222 5.18182 6.15385;
        4 2.33333 2.8 3.57143 4.44444 5.36364 6.30769;
        6 3 3.2 3.85714 4.66667 5.54545 6.46154]'
    icmp = [0 -1 -2 -3 -4 -5 -6;
        0 0 0 0 0 0 0;
        0 1 2 3 4 5 6;
        0 2 4 6 8 10 12]'
    const SOLID, LIQUID, GAS, PLASMA = 0, 1, 2, 3
    ecmp = [SOLID SOLID SOLID SOLID SOLID SOLID SOLID;
            SOLID LIQUID GAS PLASMA SOLID LIQUID GAS;
            SOLID GAS SOLID GAS SOLID GAS SOLID;
            SOLID PLASMA GAS LIQUID SOLID PLASMA GAS]'
    scmp = ["Parting", "is such", "sweet", "sorrow."]
    vicmp = Array{Int32}[[3,2,1],[1,1,2,3,5,8,13,21,34,55,89,144]]
    opq = Array{UInt8}[[0x4f, 0x50, 0x41, 0x51, 0x55, 0x45, 0x30],
                       [0x4f, 0x50, 0x41, 0x51, 0x55, 0x45, 0x31],
                       [0x4f, 0x50, 0x41, 0x51, 0x55, 0x45, 0x32],
                       [0x4f, 0x50, 0x41, 0x51, 0x55, 0x45, 0x33]]
    # For H5T_ARRAY
    AA = Array{Int,2}[
    [ 0   0   0;
     0  -1  -2;
     0  -2  -4;
     0  -3  -6;
     0  -4  -8],
    [ 0  1   2;
     1  1   1;
     2  1   0;
     3  1  -1;
     4  1  -2],
    [ 0  2  4;
     2  3  4;
     4  4  4;
     6  5  4;
     8  6  4],
    [ 0   3   6;
      3   5   7;
      6   7   8;
      9   9   9;
     12  11  10]]

    const savedir = joinpath(tempdir(), "h5")
    if !isdir(savedir)
        mkdir(savedir)
    end
    function getfile(name)
        file = joinpath(savedir, name)
        if !isfile(file)
            file = download(urlbase*name, file)
        end
        file
    end

    file = getfile("h5ex_t_floatatt.h5")
    fid = h5open(file, "r")
    dset = fid["DS1"]
    a = a_read(dset, "A1")
    @test norm(a - fcmp) < 1e-5
    close(fid)

    file = getfile("h5ex_t_float.h5")
    fid = h5open(file, "r")
    d = read(fid, "DS1")
    @test norm(d - fcmp) < 1e-5
    close(fid)

    file = getfile("h5ex_t_intatt.h5")
    fid = h5open(file, "r")
    dset = fid["DS1"]
    a = a_read(dset, "A1")
    @test a == icmp
    close(fid)

    file = getfile("h5ex_t_int.h5")
    fid = h5open(file, "r")
    d = read(fid, "DS1")
    @test d == icmp
    close(fid)

    if HDF5.h5_get_libversion() >= v"1.8.11"
      file = getfile("h5ex_t_enumatt.h5")
      fid = h5open(file, "r")
      dset = fid["DS1"]
      a = a_read(dset, "A1")
      @test a == ecmp
      close(fid)

      file = getfile("h5ex_t_enum.h5")
      fid = h5open(file, "r")
      d = read(fid, "DS1")
      @test d == ecmp
      close(fid)
    end

    file = getfile("h5ex_t_objrefatt.h5")
    fid = h5open(file, "r")
    dset = fid["DS1"]
    a = a_read(dset, "A1")
    g = fid[a[1]]
    @test isa(g, HDF5Group)
    ds2 = fid[a[2]]
    ds2v = read(ds2)
    @test isa(ds2v, Array{Int32})
    @test isempty(ds2v)
    close(fid)

    file = getfile("h5ex_t_objref.h5")
    fid = h5open(file, "r")
    d = read(fid, "DS1")
    g = fid[d[1]]
    @test isa(g, HDF5Group)
    ds2 = fid[d[2]]
    ds2v = read(ds2)
    @test isa(ds2v, Array{Int32})
    @test isempty(ds2v)
    close(fid)

    file = getfile("h5ex_t_stringatt.h5")
    fid = h5open(file, "r")
    dset = fid["DS1"]
    a = a_read(dset, "A1")
    @test a == scmp
    close(fid)

    file = getfile("h5ex_t_string.h5")
    fid = h5open(file, "r")
    d = read(fid, "DS1")
    @test d == scmp
    close(fid)

    file = getfile("h5ex_t_vlenatt.h5")
    fid = h5open(file, "r")
    dset = fid["DS1"]
    a = a_read(dset, "A1")
    @test a == vicmp
    close(fid)

    file = getfile("h5ex_t_vlen.h5")
    fid = h5open(file, "r")
    d = read(fid, "DS1")
    @test d == vicmp
    close(fid)

    file = getfile("h5ex_t_vlstringatt.h5")
    fid = h5open(file, "r")
    dset = fid["DS1"]
    a = a_read(dset, "A1")
    @test a == scmp
    close(fid)

    file = getfile("h5ex_t_vlstring.h5")
    fid = h5open(file, "r")
    d = read(fid, "DS1")
    @test d == scmp
    close(fid)

    file = getfile("h5ex_t_opaqueatt.h5")
    fid = h5open(file, "r")
    dset = fid["DS1"]
    a = a_read(dset, "A1")
    @test a.tag == "Character array"
    @test a.data == opq
    close(fid)

    file = getfile("h5ex_t_opaque.h5")
    fid = h5open(file, "r")
    d = read(fid, "DS1")
    @test d.tag == "Character array"
    @test d.data == opq
    close(fid)

    file = getfile("h5ex_t_array.h5")
    fid = h5open(file, "r")
    A = read(fid, "DS1")
    @test A == AA
    close(fid)

end


module CustomCollections

export
    ConstVector,
    ConstDict,
    NullDict,
    CacheElement,
    AbstractIndexDict,
    IndexDict,
    CacheIndexDict,
    SegmentedVector,
    SegmentedBlockDiagonalMatrix,
    UnorderedPair

export
    foreach_with_extra_args,
    isdirty,
    segments,
    ranges

using TypeSortedCollections
using DocStringExtensions
using LinearAlgebra
using UnsafeArrays

using Base: @propagate_inbounds

include("foreach_with_extra_args.jl")
include("ConstVector.jl")
include("ConstDict.jl")
include("NullDict.jl")
include("CacheElement.jl")
include("IndexDict.jl")
include("SegmentedVector.jl")
include("SegmentedBlockDiagonalMatrix.jl")
include("UnorderedPair.jl")

end # module


# Is this usable data?
function isconcrete{T<:Number}(x::T)
    !isna(x) && isfinite(x)
end

function isconcrete(x)
    !isna(x)
end

# How many concrete elements in an iterable
function concrete_length(xs)
    n = 0
    for x in xs
        if isconcrete(x)
            n += 1
        end
    end
    n
end

function nonzero_length(xs)
    n = 0
    for x in xs
        if x != 0
            n += 1
        end
    end
    n
end

# Create a new object of type T from a with missing values (i.e., those set to
# nothing) inherited from b.
function inherit{T}(a::T, b::T)
    c = copy(a)
    inherit!(c, b)
    c
end

function inherit!{T}(a::T, b::T)
    for field in T.names
        val = getfield(a, field)
        # TODO: this is a hack to let non-default labelers overide the defaults
        if val === nothing || val === string || val === fmt_float
            setfield(a, field, getfield(b, field))
        end
    end
    nothing
end

isnothing(u) = is(u, nothing)
issomething(u) = !isnothing(u)

negate(f) = x -> !f(x)

function has{T,N}(xs::AbstractArray{T,N}, y::T)
    for x in xs
        if x == y
            return true
        end
    end
    return false
end

function push!{T}(xs::Vector{T}, ys::T...)
    for y in ys
        push!(xs, y)
    end
end

Maybe(T) = Union(T, Nothing)

# Float64 -> String, trimming trailing zeros when appropriate.
# This is largely taken from cairo's function _cairo_dtostr.
function fmt_float(x::Float64)
    if x < 0.1
        a = @sprintf("%0.4f", x)
    else
        a = @sprintf("%f", x)
    end

    n = length(a)
    while a[n] == '0'
        n -= 1
    end

    if a[n] == '.'
        n -= 1
    end

    a[1:n]
end

function lerp(x::Float64, a::Float64, b::Float64)
    a + (b - a) * max(min(x, 1.0), 0.0)
end

# Generate a unique id, primarily for assigning IDs to SVG elements.
let next_unique_svg_id_num = 0
    global unique_svg_id
    function unique_svg_id()
        uid = @sprintf("id%d", next_unique_svg_id_num)
        next_unique_svg_id_num += 1
        uid
    end
end

# Remove any markup or whitespace from a string.
function escape_id(s::String)
    s = replace(s, r"<[^>]*>", "")
    s = replace(s, r" ", "_")
    s
end


"""
    FullyObservablePOMDP(mdp)

Turn `MDP` `mdp` into a `POMDP` where the observations are the states of the MDP.
"""
struct FullyObservablePOMDP{M,S,A} <: POMDP{S,A,S}
    mdp::M
end

function FullyObservablePOMDP(m::MDP)
    return FullyObservablePOMDP{typeof(m), statetype(m), actiontype(m)}(m)
end

mdptype(::Type{FullyObservablePOMDP{M,S,A}}) where {M,S,A} = M

function POMDPs.DDNStructure(::Type{M}) where M <: FullyObservablePOMDP
    MM = mdptype(M)
    add_obsnode(DDNStructure(MM))
end

add_obsnode(ddn) = add_node(ddn, :o, FunctionDDNNode((m,sp)->sp), (:sp,)) # for ::DDNStructure, but this is not declared yet POMDPs in v0.7.3

POMDPs.observations(pomdp::FullyObservablePOMDP) = states(pomdp.mdp)
POMDPs.obsindex(pomdp::FullyObservablePOMDP{S, A}, o::S) where {S, A} = stateindex(pomdp.mdp, o)

POMDPs.convert_o(T::Type{V}, o, pomdp::FullyObservablePOMDP) where {V<:AbstractArray} = convert_s(T, s, pomdp.mdp)
POMDPs.convert_o(T::Type{S}, vec::V, pomdp::FullyObservablePOMDP) where {S,V<:AbstractArray} = convert_s(T, vec, pomdp.mdp)

POMDPs.gen(::DDNNode{:o}, m::FullyObservablePOMDP, sp, rng) = sp

function POMDPs.observation(pomdp::FullyObservablePOMDP, a, sp)
    return Deterministic(sp)
end

function POMDPs.observation(pomdp::FullyObservablePOMDP, s, a, sp)
    return Deterministic(sp)
end

# inherit other function from the MDP type

POMDPs.states(pomdp::FullyObservablePOMDP) = states(pomdp.mdp)
POMDPs.actions(pomdp::FullyObservablePOMDP) = actions(pomdp.mdp)
POMDPs.transition(pomdp::FullyObservablePOMDP, s, a) = transition(pomdp.mdp, s, a)
POMDPs.initialstate_distribution(pomdp::FullyObservablePOMDP) = initialstate_distribution(pomdp.mdp)
POMDPs.initialstate(pomdp::FullyObservablePOMDP, rng::AbstractRNG) = initialstate(pomdp.mdp, rng)
POMDPs.isterminal(pomdp::FullyObservablePOMDP, s) = isterminal(pomdp.mdp, s)
POMDPs.discount(pomdp::FullyObservablePOMDP) = discount(pomdp.mdp)
POMDPs.stateindex(pomdp::FullyObservablePOMDP, s) = stateindex(pomdp.mdp, s)
POMDPs.actionindex(pomdp::FullyObservablePOMDP, a) = actionindex(pomdp.mdp, a)
POMDPs.convert_s(T::Type{V}, s, pomdp::FullyObservablePOMDP) where V<:AbstractArray = convert_s(T, s, pomdp.mdp)
POMDPs.convert_s(T::Type{S}, vec::V, pomdp::FullyObservablePOMDP) where {S,V<:AbstractArray} = convert_s(T, vec, pomdp.mdp)
POMDPs.convert_a(T::Type{V}, a, pomdp::FullyObservablePOMDP) where V<:AbstractArray = convert_a(T, a, pomdp.mdp)
POMDPs.convert_a(T::Type{A}, vec::V, pomdp::FullyObservablePOMDP) where {A,V<:AbstractArray} = convert_a(T, vec, pomdp.mdp)

POMDPs.gen(d::DDNNode, m::FullyObservablePOMDP, args...) = gen(d, m.mdp, args...)
POMDPs.gen(m::FullyObservablePOMDP, s, a, rng) = gen(m.mdp, s, a, rng)
POMDPs.reward(pomdp::FullyObservablePOMDP, s, a) = reward(pomdp.mdp, s, a)

# deprecated in POMDPs v0.8
add_obsnode(ddn::POMDPs.DDNStructureV7{(:s,:a,:sp,:r)}) = POMDPs.DDNStructureV7{(:s,:a,:sp,:o,:r)}()
add_obsnode(ddn::POMDPs.DDNStructureV7) = error("FullyObservablePOMDP only supports MDPs with the standard DDN Structure (DDNStructureV7{(:s,:a,:sp,:r)}) with POMDPs v0.7.")

POMDPs.generate_s(pomdp::FullyObservablePOMDP, s, a, rng::AbstractRNG) = generate_s(pomdp.mdp, s, a, rng)
POMDPs.generate_sr(pomdp::FullyObservablePOMDP, s, a, rng::AbstractRNG) = generate_sr(pomdp.mdp, s, a, rng)
POMDPs.n_actions(pomdp::FullyObservablePOMDP) = n_actions(pomdp.mdp)
POMDPs.n_states(pomdp::FullyObservablePOMDP) = n_states(pomdp.mdp)
function POMDPs.generate_o(pomdp::FullyObservablePOMDP, s, rng::AbstractRNG)
    return s
end


#
#  - split off the linear transformation
#  - remove the abstract oneparticle basis crap and incorporate it into
#    the Product1pBasis -> renamed as OnepBasis
#  - spec simply specifies in what order the wrapped basis generates the
#    basis functions, basically don't allow doing it lazy anymore
#  - P_qa(r) is actually P_q(r, a) so the only issue with this is
#    if we want different basis length for different species, but even then
#    we could just fill in zeros
#  - Linear1pTransform: just a linear mapping
#

# MORE NOTES
#  - consider leaving FVAL in here for now, but move it out later...
#  - changed getval(fval, X) to evaluate(fval, X)
#  - can we get away with only proviging gradients w.r.t. the
#    symbols used in this component, or do we need the whole thing?

import NamedTupleTools
using NamedTupleTools: namedtuple

@doc raw"""
`struct B1pComponent:` Wraps a one-particle basis component.

`basis` is a structure that can compute a basis, i.e. a vector of scalars,
real or complex. each element of this vector is "specified" by a NamedTuple,
stored in `spec`. E.g., for an `Rn` type basis the spec is would be
```julia
   spec = [ (n = 1,), (n=2, ), ... ]
```
while for a Ylm type basis it would be
```julia
   spec = [ (l=0, m=0), (l=1, m=-1), (l=1, m=0), ... ]
```

There is no convenience constructor, the `B1pComponent` must be constructed
"by hand", for which several wrapper functions are written.
"""
struct B1pComponent{ISYMS, TT, TB, FVAL}
   basis::TB
   fval::FVAL
   spec::Vector{NamedTuple{ISYMS, TT}}
   degrees::Vector{Int}
   label::String
   meta::Dict{String,Any}
   # ------------ derived fields
   invspec::Dict{NamedTuple{ISYMS, TT}, Int}
   # todo - fields for temporary arrays ...
end

function B1pComponent(basis, fval, spec::AbstractVector{<: NamedTuple},
                      degrees::Vector{Int},
                      label::AbstractString,
                      meta = Dict{String, Any}())
   spec1 = collect(spec)
   invspec = Dict{eltype(spec1), Int}()
   for (i, b) in enumerate(spec)
      invspec[b] = i
   end
   return B1pComponent(basis, fval, spec, degrees, label, meta, invspec)
end

# -------------- management of the basis specification, in particular the
#                interaction with Product1pBasis

getlabel(basis::B1pComponent) = basis.label

Base.length(basis::B1pComponent) = length(basis.spec)

_idxsyms(basis::B1pComponent{ISYMS}) where {ISYMS} = ISYMS

get_spec(basis::B1pComponent) = copy(basis.spec)

get_spec(basis::B1pComponent, i::Integer) = basis.spec[i]

# TODO - I don't think this is needed anymore, but the sparsification will
#        have to be moved in here and something be done about that...
# function set_spec!(basis::XScal1pBasis, spec)
#    basis.spec = identity.(spec)
#    basis.coeffs = zeros(eltype(basis.coeffs), length(spec), length(basis.P))
#    return basis
# end

function isadmissible(b::NamedTuple{BSYMS}, basis::B1pComponent) where {BSYMS}
   ISYMS = _idxsyms(basis)
   # this is an assert since it should ALWAYS be true, if not there is a bug
   @assert all(sym in BSYMS for sym in ISYMS)
   # project to the ISYMS and check it is in the b1pcomponent
   return haskey(basis.invspec, b[ISYMS])
end

# TODO - LATER
# function sparsify!(basis::XScal1pBasis, spec)
#    # spec is the part of the basis we keep
#    inds = Vector{Int}(undef, length(spec))
#    for (ib, b) in enumerate(spec)
#       iold = findall(isequal(b), basis.spec)
#       @assert length(iold) == 1
#       inds[ib] = iold[1]
#    end
#    # keep the original order of the basis since we assume it is ordered
#    # by some sensible notion of degree.
#    p = sortperm(inds)
#    basis.spec = spec[p]
#    basis.coeffs = basis.coeffs[p, :]
#    return basis
# end

symbols(basis::B1pComponent) = [ _idxsyms(basis)... ]

function indexrange(basis::B1pComponent)
   ISYMS = _idxsyms(basis)
   minidx = Dict([sym => minimum(b[sym] for b in basis.spec) for sym in ISYMS]...)
   maxidx = Dict([sym => maximum(b[sym] for b in basis.spec) for sym in ISYMS]...)
   return Dict([sym => minidx[sym]:maxidx[sym] for sym in ISYMS]...)
end

# this is needed to generate the product 1p basis - it returns the
# index of the 1p basis function component used in the 1p basis function
# specified by b
function get_index(basis::B1pComponent, b::NamedTuple)
   ISYMS = _idxsyms(basis)
   b1 = b[ISYMS]
   if !haskey(basis.invspec, b1)
      error("B1pComponent ($(basis.label) : can't find $(b1) in spec")
   end
   return basis.invspec[b1]
end

# -------------- degree calculations

# this should be revisited, doesn't feel clean yet

function degree(b::NamedTuple, basis::B1pComponent)
   ISYMS = _idxsyms(basis)
   idx = basis.invspec[ b[ISYMS] ]
   return basis.degrees[idx]
end

function degree(b::NamedTuple, basis::B1pComponent, weight::Dict)
   ISYMS = _idxsyms(basis)
   return weight[ISYMS[1]] * degree(b, basis)
end

# --------------- IO operations

==(P1::B1pComponent, P2::B1pComponent) = (
      (P1.basis == P2.basis) && (P1.spec == P2.spec) &&
      (P1.label == P2.label) && (P1.fval == P2.fval) &&
      (P1.degrees == P2.degrees) )

function write_dict(basis::B1pComponent)
   ISYMS = _idxsyms(basis)
   return Dict("__id__" => "ACE_B1pComponent",
                 "syms" => [ string.(ISYMS) ...],
                "basis" => write_dict(basis.basis),
                 "fval" => write_dict(basis.fval),
                 "spec" => convert.(Dict, basis.spec),
              "degrees" => basis.degrees,
                "label" => basis.label )
end

function read_dict(::Val{:ACE_B1pComponent}, D::Dict)
   basis = read_dict(D["basis"])
   ISYMS = tuple(Symbol.(D["syms"])...)
   spec = NamedTuple{ISYMS}.(namedtuple.(D["spec"]))
   fval = read_dict(D["fval"])
   return B1pComponent(basis, fval, spec, Int.(D["degrees"]), D["label"])
end

function show(io::IO, basis::B1pComponent)
   vsyms = Transforms.get_symbols(basis.fval)
   strvsyms = filter(!isequal(':'), "$vsyms")
   if length(vsyms) == 1
      # this doesn't work with unicode!?!?!?
      # strvsyms = strvsyms[1:end-2] * ")"
      strvsyms = prod(split(strvsyms, ','))
   end
   print(io, "B1pComponent( $(basis.label)$(strvsyms)")
   for (key, rg) in indexrange(basis)
      print(io, ", $key = $(rg)")
   end
   print(io, ")")
end

# ------------------- preparation for evaluation and managing temporaries

valtype(basis::B1pComponent) = valtype(basis.basis)

valtype(basis::B1pComponent, X::AbstractState) =
      valtype(basis.basis, evaluate(basis.fval, X))

function gradtype(basis::B1pComponent, X::AbstractState)
   x = evaluate(basis.fval, X)
   # gradient type of the inner basis which knows nothing about states
   TDB = gradtype(basis.basis, x)
   # now we need to incorporate the grad type of fval itself
   # dx = evaluate_d(basis.fval, X)
   # the gradient will be a product of a TDB times a TDVAL
   return Transforms.grad_type_dP(TDB, basis.fval, X)
      # dstate_type(valtype(basis, X), X)
end

acquire_B!(basis::B1pComponent, args...) =
         Vector{valtype(basis, args...)}(undef, length(basis))

release_B!(basis::B1pComponent, args...) = nothing

acquire_dB!(basis::B1pComponent, X::AbstractState) =
            Vector{gradtype(basis, X)}(undef, length(basis))

release_dB!(basis::B1pComponent, args...) = nothing

# ------------------------ Evaluation code
#                          this is basically an interface for the inner basis

evaluate(basis::B1pComponent, X::AbstractState) =
         evaluate(basis.basis, evaluate(basis.fval, X))

# evaluate!(B, basis::B1pComponent, X::AbstractState) =
#       evaluate!(B, basis.basis, )

function evaluate_d(basis::B1pComponent, X::AbstractState)
   B, dB = evaluate_ed(basis, X)
   release!(B)
   return dB
end

function evaluate_ed(basis::B1pComponent, X::AbstractState)
   x = evaluate(basis.fval, X)
   B, dP = evaluate_ed(basis.basis, x)
   dB = rrule_evaluate(dP, basis.fval, X)
   release!(dP)
   return B, dB
end

# this one we probably only need for training so ...
#   - we can relax the efficiency a bit
#   - we actually never need the evaluate_dd but only the
#     associated backpropagation operation which may be a bit simpler
#     to implement.
# so we will likely remove this entirely.
# evaluate_dd(basis::B1pComponent, X::AbstractState)

# -------------- AD codes

#=   *** TODO
import ChainRules: rrule, ZeroTangent, NoTangent

function _rrule_evaluate(basis::Scal1pBasis, X::AbstractState,
                         w::AbstractVector{<: Number})
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   a = _rrule_evaluate(basis.P, x, real.(w))
   TDX = ACE.dstate_type(a, X)
   return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
end

rrule(::typeof(evaluate), basis::Scal1pBasis, X::AbstractState) =
                  evaluate(basis, X),
                  w -> (NoTangent(), NoTangent(), _rrule_evaluate(basis, X, w))

function _rrule_evaluate_d(basis::Scal1pBasis, X::AbstractState,
                           w::AbstractVector)
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   w1 = [ _val(w, basis) for w in w ]
   a = _rrule_evaluate_d(basis.P, x, w1)
   TDX = ACE.dstate_type(a, X)
   return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
end

function rrule(::typeof(evaluate_d), basis::Scal1pBasis, X::AbstractState)
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   dB_ = evaluate_d(basis.P, x)
   TDX = dstate_type(valtype(basis, X), X)
   dB = [ TDX( NamedTuple{(_varsym(basis),)}( (dx,) ) )  for dx in dB_ ]
   return dB,
          w -> (NoTangent(), NoTangent(), _rrule_evaluate_d(basis, X, w))
end

=#


function accumulate!(
    ocn :: Ocean,
)

    for f in fieldnames(AccumulativeVariables)
        v = getfield(ocn.acc_vars, f)
        v .+= getfield(ocn, f)
    end

end

function avg_accumulate!(
    ocn :: Ocean;
    count :: Real,
)
    for f in fieldnames(AccumulativeVariables)

        ocn_field = getfield(ocn, f)
        acc_field = getfield(ocn.acc_vars, f)

        for i in eachindex(ocn_field)
            ocn_field[i] = acc_field[i] / count
        end

        # clear variable by default
        acc_field .= 0.0

    end

end


module EMLequil
	using DanaTypes
	using DotPlusInheritance
	using Reexport
	@reexport using ...reactors.EMLtank_basic
	import EMLtypes.length
	include("equil/equil_vap.jl")
	include("equil/equil_liq.jl")
end


module SpatQSims

using FisherySim
using Distributions
using NeutralLandscapes
using Plots
using Random
using HDF5
using CSV
using StructArrays
using Arrow

import FisherySim: simulate, domain
import Plots: plot

include("realizations.jl")
export
    generate_realizations,
    save_realizations,
    load_realization,
    load_movement,
    # save_realizations
    load_spatial_eq,
    load_popstate,
    # save_realizations,
    load_cathcability

include("prep_sims.jl")
export
    prep_sims

include("scale_devs.jl")
export
    scale_devs,
    transform_log_qdevs

include("simprep.jl")
export
    get_realization,
    SpatQSimPrep,
    movement,
    init_pop,
    log_qdevs,
    realization,
    prep_file

include("simspec.jl")
export
    SpatQSimSpec,
    # realization,
    # prep_file,
    domain,
    pop_dynamics,
    survey_targeting,
    comm_targeting,
    base_catchability,
    survey_catchability,
    comm_catchability,
    tweedie_shape,
    tweedie_dispersion,
    n_years,
    survey_vessel,
    comm_vessel,
    fleet

include("qdevscaling.jl")
export
    QDevScalingSpec,
    sim_value,
    comm_qdev_scale

include("sharedq.jl")
export
    SharedQSpec,
    sim_value,
    survey_catchability

include("prefintensity.jl")
export
    PrefIntensitySpec,
    sim_value,
    comm_targeting

include("densdepq.jl")
export
    DensityDependentQSpec,
    sim_value,
    comm_catchability

include("habq.jl")
export
    HabQSpec,
    HabQPrep,
    edge_hab,
    general_hab_distr,
    rocky_hab_distr,
    generate_habitat,
    general_hab_pref,
    rocky_hab_pref_gen,
    generate_habitat,
    move_rate,
    make_moveop,
    sim_value,
    survey_catchability,
    comm_catchability,
    prepare,
    setup,
    habitat,
    domain,
    movement,
    init_pop,
    prepare_habitat,
    get_habitat

include("sim_pars.jl")
export
    sim_values,
    sim_value_idx

include("files.jl")
export
    file_paths,
    simstudy_dir,
    simstudy_prefix,
    prep_path

include("simsetup.jl")
export
    SpatQSimSetup,
    simulate

include("run_sims.jl")
export
    run_sims

include("simresult.jl")
export
    SpatQSimResult,
    # realization,
    file_paths,
    save,
    save_pop_hdf5,
    save_pop_csv,
    save_catch_csv

include("plots.jl")
export
    catch_by_year,
    plot_gif,
    plot

# function run_simulation(scenario::Symbol,
#                         realization::Integer,
#                         prep_fn = "prep.h5")
#     # Don't need habitat for simulation
#     _, movement, init_pop, catchability_devs =
#         get_realization(realization, prep_fn)

#     # 100×100 gridded domain for all simulations
#     Ω = GriddedFisheryDomain()

#     #-Population dynamics-------------------------------------------------------
#     r = 0.06
#     K = 100.0

#     # K defined above to get spatial equilibrium
#     schaefer = Schaefer(r, K)

#     #-Targeting-----------------------------------------------------------------
#     # Commercial targeting depends on population, so isn't constructed prior to
#     # simulation.
#     survey_stations = vec(LinearIndices(Ω.n)[2:4:98, 5:10:95])
#     survey_targeting = FixedTargeting(survey_stations)
#     # Commercial preference is based on spatial distribution of biomass
#     comm_targeting = DynamicPreferentialTargeting(init_pop.P, identity)

#     #-Catchability--------------------------------------------------------------
#     survey_q_base = 0.2
#     comm_q_base = 0.2
#     if scenario == :pref
#         survey_catchability = Catchability(survey_q_base)
#         comm_catchability = Catchability(comm_q_base)
#     elseif scenario == :spat
#         survey_catchability = Catchability(survey_q_base)
#         comm_catchability = Catchability(comm_q_base .* catchability_devs)
#         comm_targeting = RandomTargeting()
#     elseif scenario == :combo
#         survey_catchability = Catchability(survey_q_base)
#         comm_catchability = Catchability(comm_q_base .* catchability_devs)
#     else
#         @error "scenario must be one of :pref, :spat, or :combo."
#     end

#     #-Vessels-------------------------------------------------------------------
#     # Tweedie power parameter
#     ξ = 1.84
#     # Tweedie variance parameter
#     ϕ = 1.2
#     survey_vessel = Vessel(survey_targeting,
#                            survey_catchability,
#                            ξ, ϕ)
#     comm_vessel = Vessel(comm_targeting,
#                          comm_catchability,
#                          ξ, ϕ)

#     #-Fleet---------------------------------------------------------------------
#     fleet = Fleet([survey_vessel, comm_vessel],
#                   [length(survey_vessel.target.locations),
#                    2_500])

#     #-Simulate------------------------------------------------------------------
#     simulate(init_pop, fleet, movement, schaefer, Ω, 25)
# end

# function run_scenarios(repl::Integer, flnm = "prep.h5")
#     Pdict = Dict{Symbol, Vector{PopState}}()
#     Cdict = Dict{Symbol, Vector{Catch}}()
#     for sc in [:pref, :spat, :combo]
#         P, C = run_simulation(sc, repl, flnm)
#         push!(Pdict, sc => P)
#         push!(Cdict, sc => C)
#     end
#     Pdict, Cdict
# end

# function save_scenarios(repl::Integer, Pdict, Cdict)
#     mkpath("repl_" * string(repl, pad = 2))
#     cd("repl_" * string(repl, pad = 2)) do
#         # Save population states to HDF5 file
#         flnm = "pop_" * string(repl, pad = 2) * ".h5"
#         dom_size = size(Pdict[:pref][1])
#         time_size = length(Pdict[:pref])
#         h5open(flnm, "w") do fid
#             for sc in keys(Pdict)
#                 sc_grp = g_create(fid, string(sc))
#                 sc_pop = create_dataset(sc_grp, "popstate", datatype(Float64),
#                                   dataspace(dom_size..., time_size))
#                 for yr in eachindex(Pdict[sc])
#                     sc_pop[:, :, yr] = Pdict[sc][yr].P
#                 end
#             end
#         end

#         for sc in keys(Pdict)
#             popstate = (pop = sum.(Pdict[sc]),)
#             CSV.write("popstate_" * string(repl, pad = 2) *
#                       "_" * string(sc) * ".csv",
#                       StructArray(popstate))
#         end

#         # Save catch records to CSV file
#         for sc in keys(Cdict)
#             flnm = "catch_" * string(repl, pad = 2) *
#                   "_" * string(sc) * ".csv"
#             CSV.write(flnm, StructArray(Cdict[sc]))
#         end
#     end
#     nothing
# end

# function get_coordref(Ω)
#     [(loc_idx = i, s1 = s1, s2 = s2)
#         for (i, (s1, s2)) in enumerate(vec(Ω.locs))]
# end

# function write_coordref(Ω)
#     coordref = get_coordref(Ω)
#     CSV.write("coordref.csv", coordref)
# end

end # module


"""
    build_crc3()

Build's required objects to model the CRC-3.

# Returns

- `crc3`: `AircraftSystems.Aircraft` object
- `default_parameters`: `<: AircraftSystems.Parameters` struct
- `actions`: list of functions required to solve the aerodynamics of the CRC-3
- `alphas`: a vector containing one element to hold the angle of attack
- `stepsymbol`: required input to solve the aerodynamis using `AircraftSystems`

"""
function build_crc3(; nwings = 2)

    #####
    ##### Define wings
    #####

    wing_lexs = fill([0.0, 0.0], nwings)
    wing_leys = fill([0.0, crc3_wing_b/2], nwings)
    wing_lezs = [[crc3_wing_spacing, crc3_wing_spacing] .* i for i in 0:nwings-1]
    wing_chords = fill([crc3_wing_c, crc3_wing_c], nwings)
    wing_thetas = fill([0.0, 0.0], nwings)
    wing_phis = fill([0.0, 0.0], nwings)
    wing_nspanwisepanels = fill(40, nwings)
    wing_nchordwisepanels = fill(1, nwings)

    #####
    ##### Define rotors
    #####

    alphas = [0.0]
    omegas = ones(nwings) .* 2000.0 * pi/30 # rad/s
    nblades = [2]
    rtip = [crc3_rotor_R]
    radii = fill(crc3_rotor_radii, length(nblades))
    rhub = fill(crc3_rotor_rhub, length(nblades))
    chords = fill(crc3_rotor_chord, length(nblades))

    # add collective
    twists = fill(crc3_rotor_twist, length(nblades)) # radians
    twist_extra = 0.0 * pi/180
    for (i, twist) in enumerate(twists)
        twist .+= twist_extra
    end

    # extra geometry
    index = fill(1, nwings)
    rotor_positions = [[-crc3_rotor_x, crc3_rotor_y, crc3_wing_spacing * i] for i in 0:nwings-1]
    rotor_orientations = fill([-1.0 * cos(crc3_rotor_cant), -sin(crc3_rotor_cant), 0.0], nwings)
    spin_directions = fill(true, nwings)

    #####
    ##### define airfoils
    #####
    Res = [5e3]
    Machs = [0.1]

    Res_lists = fill(Res,length(nblades))
    Ms_lists = fill(Machs,length(nblades))

    # contours
    contourdirectory = joinpath(topdirectory, "data", "airfoil", "contour")
    contour_paths = fill(joinpath(contourdirectory, "e212.dat"), length(nblades))
    # contour_paths = fill(fill(joinpath(contourdirectory, "e212.dat"), length(Res), length(Res)),length(nblades))

    # polars
    polar_directory = joinpath(topdirectory, "data", "airfoil", "polar")
    uncorrected_polar_paths = fill([joinpath(polar_directory, "crc3_rR06_Re5000_M0.1.txt")], length(nblades))
    corrected_polar_paths = fill([joinpath(polar_directory, "crc3_rR06_Re5000_M0.1_ext_rot.txt")], length(nblades))
    # uncorrected_polar_paths = fill([joinpath(polar_directory, "contour66.33_polar_5000.0_0.62.csv")], length(nblades))
    # corrected_polar_paths = fill([joinpath(polar_directory, "contour66.33_polar_5000.0_0.62_ext_rot.csv")], length(nblades))
    Res = [5000]
    Ms = [0.1]
    viterna_extrapolation = false
    rotation_correction = false

    plotstepi = 1:length(alphas)

    # Res_lists = fill(fill(Res, length(radii[1])), length(nblades))
    # Ms_lists = fill(fill(Machs, length(radii[1])), length(nblades))

    vinfs = [1.0] # dummy v

    #####
    ##### build system
    #####

    crc3, _, _, _, _, _, _, alphas, stepsymbol = AS.vlm_bem_template(vinfs, plotstepi, alphas,
        wing_lexs, wing_leys, wing_lezs, wing_chords, wing_thetas, wing_phis, wing_nspanwisepanels, wing_nchordwisepanels,
        omegas, nblades, rhub, rtip, radii, chords, twists,
        contour_paths, uncorrected_polar_paths, corrected_polar_paths, index,
        rotor_positions, rotor_orientations, spin_directions, Res_lists, Ms_lists;
            Vref = 1.0, symmetric = true, iref = 1, static_margin = 0.10, liftingline_x_over_c = 0.25,
            cambers = [fill((xc) -> 0, length(lex)) for lex in wing_lexs],
            spanwise_spacings = fill(AS.VL.Uniform(), length(wing_lexs)),
            chordwise_spacings = fill(AS.VL.Uniform(), length(wing_lexs)),
            wake_developement_factor = 1.0, # fully developed by default
            swirl_recovery_factor = 0.5, # as described in Veldhuis' paper
            surfacenames = ["default wing"],
            rotor_names = ["rotor 1"],
            mass = crc3_mass,
            inertia_x = 0.0,
            inertia_y = crc3_JInertia,
            inertia_z = 0.0,
            plot_directory = joinpath(topdirectory, "data","plots",TODAY),
            plot_base_name = "default",
            plot_extension = ".pdf",
            step_symbol = LS.L"\alpha "
    )

    #####
    ##### parameters
    #####

    solve_rotor_nondimensional_params = AS.solve_rotor_nondimensional(crc3, alphas)
    solve_rotor_wake_params = AS.solve_rotor_wake(crc3, alphas)
    solve_rotor_wake_params = solve_rotor_wake_params[[1,4,5,6,7,8]]
    solve_wing_CF_CM_params = AS.solve_wing_CF_CM(crc3, alphas)[2:end]
    lift_moment_distribution_params = AS.lift_moment_distribution(crc3, alphas)

    parameters = AS.ParamsSolveVLMBEM(solve_rotor_nondimensional_params..., solve_rotor_wake_params..., solve_wing_CF_CM_params..., lift_moment_distribution_params...)

    #####
    ##### Environment
    #####

    environment = AS.Environment()

    #####
    ##### Stall c_l
    #####

    cl_stall = 1.3

    #####
    ##### state/control
    #####

    nx = 7
    nu = 2

    return crc3, parameters, environment, alphas, stepsymbol, cl_stall, nx, nu
end