text:
#=
  Package: FishEBM
  File: mortality.jl
  Devin Rose
  Functions for simulating all types of mortality
  Created: June 2016
=#

"""
  Description: Generates a harvest event based on the number of adult fish in
    the current environment. Operates on the 3 basins of Lake Huron, all divided
    into their specific zones.

  Returns: Operates directly on agent_db

  Last Update: October 2016
"""
function harvest!(effort::Float64, current_week::Int64, agent_db::Vector{EnviroAgent},
  enviro_a::EnvironmentAssumptions, adult_a::AdultAssumptions, agent_a::AgentAssumptions,
  hdf::DataFrame, zoneData::DataFrame; zonesToHarvest::Vector{Int64}=[0])

  classLength = length((agent_db[1]).weekNum)
  totalHarvested = fill(0, size(adult_a.catchability))
  zoneHarvest = fill(0, 18)

  # These three lines find the seasonal effort based on known harvesting data
  numYears = ceil(current_week/52)
  yearWeek = current_week - (52 * (numYears - 1))
  seasonalEffort = (-0.4*cos((1/4.138)*yearWeek) + 0.6) * effort

  # Find which zones should be harvested based on the time of year
  if zonesToHarvest[1] == 0
    if (yearWeek < 20 || yearWeek > 40)
      zonesToHarvest = [1,2,3,4,5,6]
    else
      zonesToHarvest = [7,8,9,10,11,12,13,14,15,16,17,18]
    end
  end

  for i = 1:length(enviro_a.harvestHash)
    # Check if agent is empty
    agentNum = enviro_a.harvestHash[i]
    if (isEmpty(agent_db[agentNum]) == false)
      for j = 1:classLength
        #Check if given cohort is an adult population
        if (findCurrentStage(current_week, agent_db[agentNum].weekNum[j], agent_a.growth)) == 4
          age = getAge(current_week, agent_db[agentNum].weekNum[j])
          if in(enviro_a.harvestZoneNum[i], zonesToHarvest)
            numHarvest = rand(Binomial(agent_db[agentNum].alive[j], adult_a.catchability[age - 1]*seasonalEffort))
          else
            numHarvest = rand(Binomial(agent_db[agentNum].alive[j], adult_a.catchability[age - 1]*(seasonalEffort*0.5)))
          end
          #Add/subtract harvest numbers to appropriate vectors
          agent_db[agentNum].harvest += numHarvest
          totalHarvested[age - 1] += numHarvest
          zoneHarvest[enviro_a.harvestZoneNum[i]] += numHarvest
          agent_db[agentNum].alive[j] -= numHarvest
        end #if findCurrentStage
      end #for j
    end #if isEmpty
  end #for i

  push!(hdf, (vcat(current_week, totalHarvested..., sum(totalHarvested))))
  push!(zoneData, (vcat(current_week, zoneHarvest..., sum(zoneHarvest))))
end

"""
  Description:  This function generates a mortality based on the stage of the
    fish and its corresponding natural mortality and its location within the
    habitat as described in EnvironmentAssumptions.

  Returns: Operates directly on agent_db and kdf

  Last update: September 2016
"""
function kill!(agent_db::Vector{EnviroAgent}, e_a::EnvironmentAssumptions, a_a::AgentAssumptions, current_week::Int64, kdf::DataFrame)
  classLength = length((agent_db[1]).weekNum)
  totalNatural = 0
  totalExtra = 0

  for i = 1:length(agent_db)
    #Check if class is empty. If not empty, continue with kill function. Otherwise skip to next agent
    if (isEmpty(agent_db[i]) == false)
      for j = 1:classLength

        #current_age = current_week - agent_db[i].weekNum[j]
        stage = findCurrentStage(current_week, agent_db[i].weekNum[j], a_a.growth)
        if agent_db[i].alive[j] > 0
          habitat = e_a.habitat[agent_db[i].locationID]

          #Number of fish killed follows binomial distribution with arguments of number of fish alive
          #and natural mortality in the form of a probability
          killedNatural = rand(Binomial(agent_db[i].alive[j], a_a.naturalmortality[habitat, stage]))
          agent_db[i].killedNatural[stage] += killedNatural
          totalNatural += killedNatural
          agent_db[i].alive[j] -= killedNatural
          if agent_db[i].alive[j] > 0
            if in(agent_db[i].locationID, e_a.risk) #Check if this particular locationID is in risk zone
              killedExtra = rand(Binomial(agent_db[i].alive[j], a_a.extramortality[stage]))
              agent_db[i].killedExtra[stage] += killedExtra
              totalExtra += killedExtra
              agent_db[i].alive[j] -= killedExtra
            end #if risk
          end #if agent_db[i].alive (inner)
        end #if agent_db[i].alive (outer)
      end #for j=1:classLength
    end #if isEmpty
  end #for i=1:length(agent_db)

  kdf[size(kdf)[1], 2] = totalNatural
  kdf[size(kdf)[1], 3] = totalExtra

  return agent_db
end

"""
  Description: This function applies transition probabilities to the adult fish
    population to regulate the adult age distribtion.

  Returns: Operates directly on agent_db

  Last update: October 2016
"""
function killAgeSpecific!(agent_db::Vector{EnviroAgent}, adult_a::AdultAssumptions, age_specific_pop::Vector{Int64},
  year_specific_cc::Float64, current_week::Int64, kdf::DataFrame)

  totalKilled = 0

  stockSize = fill(0, size(adult_a.naturalmortality))
  ageVector = fill(0, length(agent_db[1].weekNum))
  mortalityRate = fill(0., length(age_specific_pop))

  if isnan(adult_a.mortalitycompensation)
    compensation_factor = 1
  else
    compFactor = 2*(cdf(Normal(year_specific_cc, year_specific_cc/adult_a.mortalitycompensation), sum(age_specific_pop)))
  end

  @assert(0.0 <= compFactor <= 2.0, "Population regulation failed in killAgeSpecific!, respecify simulation parameters. Expected:0.0-2.0, Received:$compFactor")

  #find dynamic stock sizes
  for i = 1:(length(adult_a.naturalmortality))
    stockSize[i] = rand(Binomial(age_specific_pop[i], 1-adult_a.naturalmortality[i]*compFactor))
  end

  #Find desired mortality rates for dynamic stock sizes
  for j = 1:length(age_specific_pop)
    if age_specific_pop[j] != 0
      mortalityRate[j] = 1-(stockSize[j]/age_specific_pop[j])
    end
  end

  #Find the age (in years) of each adult cohort
  for i = 1:length(ageVector)
    ageVector[i] = getAge(current_week, agent_db[1].weekNum[i])
  end

  for j = 1:length(agent_db)
    if (isEmpty(agent_db[j]) == false)
      k = 1
      while ageVector[k] > 1 && k < length(ageVector)
        estimatedMortality = adult_a.naturalmortality[ageVector[k]-1]
        # Comment out the extra mortality for age8Plus adults
        #=if ageVector[k] == 8
          overAge = floor((current_week-agent_db[1].weekNum[k])/52)
          ageEffect = overAge-7
          estimatedMortality = ageEffect*estimatedMortality

          if estimatedMortality > 1.0
            estimatedMortality = 1.0
          end # if est.mort > 1
        end=#
        killedAdult = rand(Binomial(agent_db[j].alive[k], estimatedMortality))

        agent_db[j].alive[k] -= killedAdult
        totalKilled += killedAdult

        k+=1
      end # while adult cohort
    end # if empty
  end # for agent_db

  kdf[size(kdf)[1], 4] = totalKilled

end


text:
abstract type RatesCacheTrait end

struct HasRates <: RatesCacheTrait end

# populate the rates vector with each jump's propensity
@inline function update_jump_rates!(::HasRates, rates, total_rate, state, model)
  total_rate = zero(total_rate)

  for j in eachindex(rates)
      rate_j = rate(model, state, j)
      @inbounds rates[j] = rate_j
      @fastmath total_rate += rate_j
  end

  return total_rate
end

# populate the rates vector with each reaction channel's propensity evaluated at x,
# given that reaction R_k occured
@inline function update_jump_rates!(::HasRates, rates, total_rate, state, model, k)
  for j in dependents(model.dep_graph, k)
      rate_j = rate(model, state, j)
      @inbounds @fastmath total_rate += rate_j - rates[j]
      @inbounds rates[j] = rate_j
  end

  return total_rate
end

@inline function update_jump_rates!(::HasRates, rates, total_rate, state, model::InteractingParticleSystem, unused)
  xdiff = state.xdiff
  reactions = model.reactions

  for s in eachindex(xdiff)
    if xdiff[s]
      for j in eachindex(rates)
        rxn = reactions[j]

        if s == rxn.class
          rate_j = rate(model, state, j)
          @inbounds @fastmath total_rate += rate_j - rates[j]
          @inbounds rates[j] = rate_j
        end
      end
    end
  end

  return total_rate
end

# linear search
@inline function search_jump_rates(::HasRates, rates, total_rate)
  u = rand() * total_rate
  c = zero(total_rate)
  j = 0

  for rate in rates
    c += rate
    j += 1
    c >= u && break
  end

  return j
end

struct HasSums <: RatesCacheTrait end

# populate the rates vector with partial sums of the reaction channels' propensities evaluated at x
@inline function update_jump_rates!(::HasSums, rates, total_rate, state, model)
    rates[1] = rate(model, state, 1)

    for j in 2:length(rates)
        @fastmath rates[j] = rates[j - 1] + rate(model, state, j)
    end

    return rates[end]
end

# binary search on jump rates
@inline function search_jump_rates(::HasSums, rates, total_rate)
    return searchsortedfirst(rates, rand() * total_rate)
end

# search on jump times
@inline function search_jump_times(::HasRates, rates)
    n = length(rates)

    j = 1
    @inbounds s = randexp() / rates[1]

    for k in 2:n
        @inbounds w = randexp() / rates[k]
        if w < s
            j = k
            s = w
        end
    end

    return (j, s)
end

ext: jl


text:
"""
	EconomicDispatch(hours,
					 nodes::Nodes,
					 powerplants::PowerPlants,
					 renewables::Renewables,
					 )

Creates an economic dispatch model from power plants and renewables.
Solves using JuMP. Returns ED_mod, P_opt, P_R_opt, P_R_opt_dist, price.
"""
function EconomicDispatch(hours,
						  nodes::Nodes,
			   		      powerplants::PowerPlants,
						  renewables::Renewables,
			   			  )

	T = hours
	G = powerplants.unit
	R = renewables.unit
	J = nodes.id

	# Initialise JuMP model: Economic dispatch
	ED_mod = JuMP.Model(with_optimizer(Gurobi.Optimizer))

	# Variables
	@variable(ED_mod, P[T, G] >= 0)
	@variable(ED_mod, P_R[T, R] >= 0)

	# Constraints
	# Market clearing/power balance
	@constraint(ED_mod, MarketClearing[t = T],
	    sum(P[t, g] for g in G) +
			sum(P_R[t, r] for r in R) == nodes.systemload[t] -
			sum(nodes.exchange[j][t] for j in J));

	# Upper generation limit
	@constraint(ED_mod, GenerationLimitUp[g = G, t = T],
	    P[t, g] <= powerplants.pmax[g][t]);

	# Lower generation limit
	@constraint(ED_mod, GenerationLimitDown[g = G, t = T],
	    P[t, g] >= powerplants.pmin[g][t]);

	# Upper generation limit for renewables
	@constraint(ED_mod, ResGenerationLimitUp[r = R, t = T],
	    P_R[t, r] <= renewables.infeed[r][t]);

	# # Ramp-up limit
	# @constraint(ED_mod, RampUp[g = G, t = T; t > T[1]],
	# 	P[t, g] - P[t-1, g] <= powerplants.rup[g])
	#
	# # Ramp-down limit
	# @constraint(ED_mod, RampDown[g = G, t = T; t > T[1]],
	# 	-(P[t, g] - P[t-1, g]) <= powerplants.rdn[g])

	# Objective function
	@objective(ED_mod, Min,
		sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T));

	# @objective(ED_mod, Min,
	#     sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T) +
	# 		sum(0 * P_R[t, r] for r in R, t in T));

	# Initiate optimisation process
	JuMP.optimize!(ED_mod)

	# Export results
	P_opt = JuMP.value.(P)
	P_R_opt = JuMP.value.(P_R)
	price = JuMP.dual.(MarketClearing)

	# Correct RES dispatch
	# Total summed RES generation from dispatch for each t
	ResTotalGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)

	for t in T
		ResTotalGeneration[t] = sum(P_R_opt[t, :])
	end

	# Maximum output possible
	ResMaxGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)
	for t in T
		ResMaxGeneration[t] = sum(renewables.infeed[r][t] for r in R)
	end

	# Redistribution among RES
	ResUnitGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T), length(R)), T, R)

	for t in T
		for r in R
			ResUnitGeneration[t, r] = (renewables.infeed[r][t] / ResMaxGeneration[t]) *
			ResTotalGeneration[t]
		end
	end

	# Distributed, corrected P_R_opt
	P_R_opt_dist = ResUnitGeneration

	return(ED_mod,
	       P_opt,
		   P_R_opt,
		   P_R_opt_dist,
		   price
	       )
end

"""
	EconomicDispatch(hours,
					 nodes::Nodes,
					 powerplants::PowerPlants,
					 renewables::Renewables,
					 storages::Storages,
					 )

Creates an economic dispatch model from power plants, renewables, and storages.
Initial storage level is fixed to zero. Solves using JuMP. Returns ED_mod,
P_opt, P_R_opt, P_R_opt_dist, P_S_opt, D_S_opt, L_S_opt, price.
"""
### Economic dispatch with renewable energy sources and storages
function EconomicDispatch(hours,
						  nodes::Nodes,
			   		      powerplants::PowerPlants,
						  renewables::Renewables,
						  storages::Storages,
			   			  )

	T = hours
	G = powerplants.unit
	R = renewables.unit
	S = storages.unit
	J = nodes.id

	# Initialise JuMP model: Economic dispatch
	ED_mod = JuMP.Model(with_optimizer(Gurobi.Optimizer))

	# Variables
	@variable(ED_mod, P[T, G] >= 0)
	@variable(ED_mod, P_R[T, R] >= 0)

	# Storage voriables
	@variable(ED_mod, P_S[T, S] >= 0)	# Power generation from storage unit
	@variable(ED_mod, D_S[T, S] >= 0)	# Electricity demand from storage unit
	@variable(ED_mod, L_S[T, S] >= 0)	# Storage level

	# Fix storage level in period 0 to zero
	for s in S
		fix(L_S[T[1], s], 0; force = true)
	end

	# Constraints
	# Market clearing/power balance
	@constraint(ED_mod, MarketClearing[t = T],
	    sum(P[t, g] for g in G) +
		sum(P_R[t, r] for r in R) +
		sum(P_S[t, s] for s in S) -
		sum(D_S[t, s] for s in S)
			== nodes.systemload[t] -
			sum(nodes.exchange[j][t] for j in J));

	# Upper generation limit
	@constraint(ED_mod, GenerationLimitUp[g = G, t = T],
	    P[t, g] <= powerplants.pmax[g][t]);

	# Lower generation limit
	@constraint(ED_mod, GenerationLimitDown[g = G, t = T],
	    P[t, g] >= powerplants.pmin[g][t]);

	# Upper generation limit for renewables
	@constraint(ED_mod, ResGenerationLimitUp[r = R, t = T],
	    P_R[t, r] <= renewables.infeed[r][t]);

	# # Ramp-up limit
	# @constraint(ED_mod, RampUp[g = G, t = T; t > T[1]],
	# 	P[t, g] - P[t-1, g] <= powerplants.rup[g])
	#
	# # Ramp-down limit
	# @constraint(ED_mod, RampDown[g = G, t = T; t > T[1]],
	# 	-(P[t, g] - P[t-1, g]) <= powerplants.rdn[g])

	### Storage constraints
	# Maximum storage power output
	@constraint(ED_mod, StoragePowerOutput[s = S, t = T],
	    P_S[t, s] <= storages.power[s]);

	# Maximum storage power input
	@constraint(ED_mod, StoragePowerInput[s = S, t = T],
	    D_S[t, s] <= storages.power[s]);

	# Generation limit determined by remaining storage level
	@constraint(ED_mod, StorageLevelGen[s = S, t = T],
	    P_S[t, s] <= L_S[t, s]);

	# Storace capacity limit
	@constraint(ED_mod, StorageLevelCap[s = S, t = T],
	    L_S[t, s] <= storages.storage[s]);

	# Storage
	@constraint(ED_mod, Storage[s = S, t = T; t > T[1]],
		L_S[t-1, s] - P_S[t-1, s] + storages.efficiency[s] * D_S[t-1, s] == L_S[t, s]);

	# Objective function
	@objective(ED_mod, Min,
	    sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T));

	# @objective(ED_mod, Min,
	#     sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T) +
	# 		sum(0 * P_R[t, r] for r in R, t in T));

	# Initiate optimisation process
	JuMP.optimize!(ED_mod)

	# Export results
	P_opt = JuMP.value.(P)
	P_R_opt = JuMP.value.(P_R)
	P_S_opt = JuMP.value.(P_S)
	D_S_opt = JuMP.value.(D_S)
	L_S_opt = JuMP.value.(L_S)
	price = JuMP.dual.(MarketClearing)

	# Correct RES dispatch
	# Total summed RES generation from dispatch for each t
	ResTotalGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)

	for t in T
		ResTotalGeneration[t] = sum(P_R_opt[t, :])
	end

	# Maximum output possible
	ResMaxGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)
	for t in T
		ResMaxGeneration[t] = sum(renewables.infeed[r][t] for r in R)
	end

	# Redistribution among RES
	ResUnitGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T), length(R)), T, R)

	for t in T
		for r in R
			ResUnitGeneration[t, r] = (renewables.infeed[r][t] / ResMaxGeneration[t]) *
			ResTotalGeneration[t]
		end
	end

	# Distributed, corrected P_R_opt
	P_R_opt_dist = ResUnitGeneration

	return(ED_mod,
	       P_opt,
		   P_R_opt,
		   P_R_opt_dist,
		   P_S_opt,
		   D_S_opt,
		   L_S_opt,
		   price
	       )
end

"""
	EconomicDispatch(hours,
					 nodes::Nodes,
					 powerplants::PowerPlants,
					 renewables::Renewables,
					 storages::Storages,
					 P_S_opt_prev,
					 D_S_opt_prev,
					 L_S_opt_prev,
					 )

Creates an economic dispatch model from power plants, renewables, and storages.
Initial storage level, generation, and demand is initialised  with parameters
from the previous model run. Solves using JuMP. Returns ED_mod, P_opt, P_R_opt,
P_R_opt_dist, P_S_opt, D_S_opt, L_S_opt, price.
"""
function EconomicDispatch(hours,
						  nodes::Nodes,
			   		      powerplants::PowerPlants,
						  renewables::Renewables,
						  storages::Storages,
						  P_S_opt_prev,
						  D_S_opt_prev,
						  L_S_opt_prev,
			   			  )

	T = hours
	G = powerplants.unit
	R = renewables.unit
	S = storages.unit
	J = nodes.id

	# Initialise JuMP model: Economic dispatch
	ED_mod = JuMP.Model(with_optimizer(Gurobi.Optimizer))

	# Variables
	@variable(ED_mod, P[T, G] >= 0)
	@variable(ED_mod, P_R[T, R] >= 0)

	# Storage voriables
	@variable(ED_mod, P_S[T, S] >= 0)	# Power generation from storage unit
	@variable(ED_mod, D_S[T, S] >= 0)	# Electricity demand from storage unit
	@variable(ED_mod, L_S[T, S] >= 0)	# Storage level

	# Fix storage parameters in first hour to the last hour value of the
	# previous day slice.
	for s in S
		fix(P_S[T[1], s], P_S_opt_prev[(T[1]), s]; force = true)
		fix(D_S[T[1], s], D_S_opt_prev[(T[1]), s]; force = true)
		fix(L_S[T[1], s], L_S_opt_prev[(T[1]), s]; force = true)
	end

	# Constraints
	# Market clearing/power balance
	@constraint(ED_mod, MarketClearing[t = T],
	    sum(P[t, g] for g in G) +
		sum(P_R[t, r] for r in R) +
		sum(P_S[t, s] for s in S) -
		sum(D_S[t, s] for s in S)
			== nodes.systemload[t] -
			sum(nodes.exchange[j][t] for j in J));

	# Upper generation limit
	@constraint(ED_mod, GenerationLimitUp[g = G, t = T],
	    P[t, g] <= powerplants.pmax[g][t]);

	# Lower generation limit
	@constraint(ED_mod, GenerationLimitDown[g = G, t = T],
	    P[t, g] >= powerplants.pmin[g][t]);

	# Upper generation limit for renewables
	@constraint(ED_mod, ResGenerationLimitUp[r = R, t = T],
	    P_R[t, r] <= renewables.infeed[r][t]);

	# # Ramp-up limit
	# @constraint(ED_mod, RampUp[g = G, t = T; t > T[1]],
	# 	P[t, g] - P[t-1, g] <= powerplants.rup[g])
	#
	# # Ramp-down limit
	# @constraint(ED_mod, RampDown[g = G, t = T; t > T[1]],
	# 	-(P[t, g] - P[t-1, g]) <= powerplants.rdn[g])

	### Storage constraints
	# Maximum storage power output
	@constraint(ED_mod, StoragePowerOutput[s = S, t = T],
	    P_S[t, s] <= storages.power[s]);

	# Maximum storage power input
	@constraint(ED_mod, StoragePowerInput[s = S, t = T],
	    D_S[t, s] <= storages.power[s]);

	# Generation limit determined by remaining storage level
	@constraint(ED_mod, StorageLevelGen[s = S, t = T],
	    P_S[t, s] <= L_S[t, s]);

	# Storace capacity limit
	@constraint(ED_mod, StorageLevelCap[s = S, t = T],
	    L_S[t, s] <= storages.storage[s]);

	# Storage
	@constraint(ED_mod, Storage[s = S, t = T; t > T[1]],
		L_S[t-1, s] - P_S[t-1, s] + storages.efficiency[s] * D_S[t-1, s] == L_S[t, s]);

	# Objective function
	@objective(ED_mod, Min,
	    sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T));

	# @objective(ED_mod, Min,
	#     sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T) +
	# 		sum(0 * P_R[t, r] for r in R, t in T));

	# Initiate optimisation process
	JuMP.optimize!(ED_mod)

	# Export results
	P_opt = JuMP.value.(P)
	P_R_opt = JuMP.value.(P_R)
	P_S_opt = JuMP.value.(P_S)
	D_S_opt = JuMP.value.(D_S)
	L_S_opt = JuMP.value.(L_S)
	price = JuMP.dual.(MarketClearing)

	# Correct RES dispatch
	# Total summed RES generation from dispatch for each t
	ResTotalGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)

	for t in T
		ResTotalGeneration[t] = sum(P_R_opt[t, :])
	end

	# Maximum output possible
	ResMaxGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)
	for t in T
		ResMaxGeneration[t] = sum(renewables.infeed[r][t] for r in R)
	end

	# Redistribution among RES
	ResUnitGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T), length(R)), T, R)

	for t in T
		for r in R
			ResUnitGeneration[t, r] = (renewables.infeed[r][t] / ResMaxGeneration[t]) *
			ResTotalGeneration[t]
		end
	end

	# Distributed, corrected P_R_opt
	P_R_opt_dist = ResUnitGeneration

	return(ED_mod,
	       P_opt,
		   P_R_opt,
		   P_R_opt_dist,
		   P_S_opt,
		   D_S_opt,
		   L_S_opt,
		   price
	       )
end

#### 80%
### Economic dispatch with renewable energy sources and storages
function EconomicDispatch80(hours,
						  nodes::Nodes,
			   		      powerplants::PowerPlants,
						  renewables::Renewables,
						  storages::Storages,
			   			  )

	T = hours
	G = powerplants.unit
	R = renewables.unit
	S = storages.unit
	J = nodes.id

	# Initialise JuMP model: Economic dispatch
	ED_mod = JuMP.Model(with_optimizer(Gurobi.Optimizer))

	# Variables
	@variable(ED_mod, P[T, G] >= 0)
	@variable(ED_mod, P_R[T, R] >= 0)

	# Storage voriables
	@variable(ED_mod, P_S[T, S] >= 0)	# Power generation from storage unit
	@variable(ED_mod, D_S[T, S] >= 0)	# Electricity demand from storage unit
	@variable(ED_mod, L_S[T, S] >= 0)	# Storage level

	# Fix storage level in period 0 to zero
	for s in S
		fix(L_S[T[1], s], 0; force = true)
	end

	# Constraints
	# Market clearing/power balance
	@constraint(ED_mod, MarketClearing[t = T],
	    sum(P[t, g] for g in G) +
		sum(P_R[t, r] for r in R) +
		sum(P_S[t, s] for s in S) -
		sum(D_S[t, s] for s in S)
			== nodes.systemload[t] -
			sum(nodes.exchange[j][t] for j in J));

	# Upper generation limit
	@constraint(ED_mod, GenerationLimitUp[g = G, t = T],
	    P[t, g] <= powerplants.pmax[g][t]);

	# Lower generation limit
	@constraint(ED_mod, GenerationLimitDown[g = G, t = T],
	    P[t, g] >= powerplants.pmin[g][t]);

	# Upper generation limit for renewables
	@constraint(ED_mod, ResGenerationLimitUp[r = R, t = T],
	    P_R[t, r] <= renewables.infeed[r][t]);

	# # Ramp-up limit
	# @constraint(ED_mod, RampUp[g = G, t = T; t > T[1]],
	# 	P[t, g] - P[t-1, g] <= powerplants.rup[g])
	#
	# # Ramp-down limit
	# @constraint(ED_mod, RampDown[g = G, t = T; t > T[1]],
	# 	-(P[t, g] - P[t-1, g]) <= powerplants.rdn[g])

	### Storage constraints
	# Maximum storage power output
	@constraint(ED_mod, StoragePowerOutput[s = S, t = T],
	    P_S[t, s] <= storages.power[s]);

	# Maximum storage power input
	@constraint(ED_mod, StoragePowerInput[s = S, t = T],
	    D_S[t, s] <= storages.power[s]);

	# Generation limit determined by remaining storage level
	@constraint(ED_mod, StorageLevelGen[s = S, t = T],
	    P_S[t, s] <= L_S[t, s]);

	# Storace capacity limit
	@constraint(ED_mod, StorageLevelCap[s = S, t = T],
	    L_S[t, s] <= 0.8*storages.storage[s]);

	# Storage
	@constraint(ED_mod, Storage[s = S, t = T; t > T[1]],
		L_S[t-1, s] - P_S[t-1, s] + storages.efficiency[s] * D_S[t-1, s] == L_S[t, s]);

	# Objective function
	@objective(ED_mod, Min,
	    sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T));

	# @objective(ED_mod, Min,
	#     sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T) +
	# 		sum(0 * P_R[t, r] for r in R, t in T));

	# Initiate optimisation process
	JuMP.optimize!(ED_mod)

	# Export results
	P_opt = JuMP.value.(P)
	P_R_opt = JuMP.value.(P_R)
	P_S_opt = JuMP.value.(P_S)
	D_S_opt = JuMP.value.(D_S)
	L_S_opt = JuMP.value.(L_S)
	price = JuMP.dual.(MarketClearing)

	# Correct RES dispatch
	# Total summed RES generation from dispatch for each t
	ResTotalGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)

	for t in T
		ResTotalGeneration[t] = sum(P_R_opt[t, :])
	end

	# Maximum output possible
	ResMaxGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)
	for t in T
		ResMaxGeneration[t] = sum(renewables.infeed[r][t] for r in R)
	end

	# Redistribution among RES
	ResUnitGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T), length(R)), T, R)

	for t in T
		for r in R
			ResUnitGeneration[t, r] = (renewables.infeed[r][t] / ResMaxGeneration[t]) *
			ResTotalGeneration[t]
		end
	end

	# Distributed, corrected P_R_opt
	P_R_opt_dist = ResUnitGeneration

	return(ED_mod,
	       P_opt,
		   P_R_opt,
		   P_R_opt_dist,
		   P_S_opt,
		   D_S_opt,
		   L_S_opt,
		   price
	       )
end

function EconomicDispatch80(hours,
						  nodes::Nodes,
			   		      powerplants::PowerPlants,
						  renewables::Renewables,
						  storages::Storages,
						  P_S_opt_prev,
						  D_S_opt_prev,
						  L_S_opt_prev,
			   			  )

	T = hours
	G = powerplants.unit
	R = renewables.unit
	S = storages.unit
	J = nodes.id

	# Initialise JuMP model: Economic dispatch
	ED_mod = JuMP.Model(with_optimizer(Gurobi.Optimizer))

	# Variables
	@variable(ED_mod, P[T, G] >= 0)
	@variable(ED_mod, P_R[T, R] >= 0)

	# Storage voriables
	@variable(ED_mod, P_S[T, S] >= 0)	# Power generation from storage unit
	@variable(ED_mod, D_S[T, S] >= 0)	# Electricity demand from storage unit
	@variable(ED_mod, L_S[T, S] >= 0)	# Storage level

	# Fix storage parameters in first hour to the last hour value of the
	# previous day slice.
	for s in S
		fix(P_S[T[1], s], P_S_opt_prev[(T[1]), s]; force = true)
		fix(D_S[T[1], s], D_S_opt_prev[(T[1]), s]; force = true)
		fix(L_S[T[1], s], L_S_opt_prev[(T[1]), s]; force = true)
	end

	# Constraints
	# Market clearing/power balance
	@constraint(ED_mod, MarketClearing[t = T],
	    sum(P[t, g] for g in G) +
		sum(P_R[t, r] for r in R) +
		sum(P_S[t, s] for s in S) -
		sum(D_S[t, s] for s in S)
			== nodes.systemload[t] -
			sum(nodes.exchange[j][t] for j in J));

	# Upper generation limit
	@constraint(ED_mod, GenerationLimitUp[g = G, t = T],
	    P[t, g] <= powerplants.pmax[g][t]);

	# Lower generation limit
	@constraint(ED_mod, GenerationLimitDown[g = G, t = T],
	    P[t, g] >= powerplants.pmin[g][t]);

	# Upper generation limit for renewables
	@constraint(ED_mod, ResGenerationLimitUp[r = R, t = T],
	    P_R[t, r] <= renewables.infeed[r][t]);

	# # Ramp-up limit
	# @constraint(ED_mod, RampUp[g = G, t = T; t > T[1]],
	# 	P[t, g] - P[t-1, g] <= powerplants.rup[g])
	#
	# # Ramp-down limit
	# @constraint(ED_mod, RampDown[g = G, t = T; t > T[1]],
	# 	-(P[t, g] - P[t-1, g]) <= powerplants.rdn[g])

	### Storage constraints
	# Maximum storage power output
	@constraint(ED_mod, StoragePowerOutput[s = S, t = T],
	    P_S[t, s] <= storages.power[s]);

	# Maximum storage power input
	@constraint(ED_mod, StoragePowerInput[s = S, t = T],
	    D_S[t, s] <= storages.power[s]);

	# Generation limit determined by remaining storage level
	@constraint(ED_mod, StorageLevelGen[s = S, t = T],
	    P_S[t, s] <= L_S[t, s]);

	# Storace capacity limit
	@constraint(ED_mod, StorageLevelCap[s = S, t = T],
	    L_S[t, s] <= 0.8*storages.storage[s]);

	# Storage
	@constraint(ED_mod, Storage[s = S, t = T; t > T[1]],
		L_S[t-1, s] - P_S[t-1, s] + storages.efficiency[s] * D_S[t-1, s] == L_S[t, s]);

	# Objective function
	@objective(ED_mod, Min,
	    sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T));

	# @objective(ED_mod, Min,
	#     sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T) +
	# 		sum(0 * P_R[t, r] for r in R, t in T));

	# Initiate optimisation process
	JuMP.optimize!(ED_mod)

	# Export results
	P_opt = JuMP.value.(P)
	P_R_opt = JuMP.value.(P_R)
	P_S_opt = JuMP.value.(P_S)
	D_S_opt = JuMP.value.(D_S)
	L_S_opt = JuMP.value.(L_S)
	price = JuMP.dual.(MarketClearing)

	# Correct RES dispatch
	# Total summed RES generation from dispatch for each t
	ResTotalGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)

	for t in T
		ResTotalGeneration[t] = sum(P_R_opt[t, :])
	end

	# Maximum output possible
	ResMaxGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T)), T)
	for t in T
		ResMaxGeneration[t] = sum(renewables.infeed[r][t] for r in R)
	end

	# Redistribution among RES
	ResUnitGeneration = JuMP.Containers.DenseAxisArray(zeros(length(T), length(R)), T, R)

	for t in T
		for r in R
			ResUnitGeneration[t, r] = (renewables.infeed[r][t] / ResMaxGeneration[t]) *
			ResTotalGeneration[t]
		end
	end

	# Distributed, corrected P_R_opt
	P_R_opt_dist = ResUnitGeneration

	return(ED_mod,
	       P_opt,
		   P_R_opt,
		   P_R_opt_dist,
		   P_S_opt,
		   D_S_opt,
		   L_S_opt,
		   price
	       )
end

max_issues_repo_name: bobbyxiong/redispatch-ptg

max_forks_repo_name: bobbyxiong/redispatch-ptg


text:
module FluxOptTools

using LinearAlgebra, Optim, Flux, Zygote, RecipesBase

export veclength, optfuns

veclength(grads::Zygote.Grads) = sum(length(grads[p]) for p in grads.params)
veclength(params::Flux.Params) = sum(length, params.params)
veclength(x) = length(x)
Base.zeros(grads::Zygote.Grads) = zeros(veclength(grads))
Base.zeros(pars::Flux.Params) = zeros(veclength(pars))

function optfuns(loss, pars::Union{Flux.Params, Zygote.Params})
    grads = Zygote.gradient(loss, pars)
    p0 = zeros(pars)
    copy!(p0, pars)
    gradfun = function (g,w)
        copy!(pars, w)
        grads = Zygote.gradient(loss, pars)
        copy!(g, grads)
    end
    lossfun = function (w)
        copy!(pars, w)
        loss()
    end
    fg! = function (F,G,w)
        copy!(pars, w)
        if G != nothing
            l, back = Zygote.pullback(loss, pars)
            grads = back(1)
            copy!(G, grads)
            return l
        end
        if F != nothing
            return loss()
        end
    end
    lossfun, gradfun, fg!, p0
end

@recipe function lossplot(loss::Function, pars::Flux.Params; l=0.1, npoints=30)
    p       = zeros(pars)
    copy!(p,pars)
    pcopy   = deepcopy(p)
    n0      = norm(p)
    dx,dy   = randn(length(p)),randn(length(p))
    dx     *= n0*l/norm(dx)
    dy     *= n0*l/norm(dy)
    pertvec = LinRange(-1,1,npoints)
    losses = map(Iterators.product(pertvec,pertvec)) do (lx,ly)
        pi = p + lx*dx + ly*dy
        copy!(pars, pi)
        loss()
    end
    copy!(pars, pcopy)
    seriestype --> :contour
    pertvec,pertvec,losses
end

include("SLBFGS.jl")

end # module


text:
"""
    resonance_eigenproblem(sim, k, ka=0, kb=0) -> A,B,σ
"""
function resonance_eigenproblem(sim::Simulation, k::Number, ka::Number=0, kb::Number=0)
    N,args = sim.dis.N,(sim.dis.coordinate_system,sim.bnd.∂Ω,sim.bnd.bc,sim.bnd.bl)
    (∇²,fs,s),defs = laplacian(N,args...,k,ka,kb,0)
    A = spzeros(ComplexF64,prod(N),prod(N))
    for i ∈ eachindex(∇²)
        A += ∇²[i]*fs[i](k,ka,kb)
    end
    A = (A + SparseMatrixCSC(transpose(A)))/2
    B = s*spdiagm(0=>-sim.sys.ε[:])
    return A::SparseMatrixCSC{ComplexF64,Int}, B, k^2
end

"""
    resonance_eigenproblem(sim, k, ka=0, kb=0; η=0) -> A,B,σ
"""
function cf_eigenproblem(sim::Simulation, k::Number, ka::Number=0, kb::Number=0; η::Number=0)
    N,args = sim.dis.N,(sim.dis.coordinate_system,sim.bnd.∂Ω,sim.bnd.bc,sim.bnd.bl)
    (∇²,fs,s),defs = laplacian(N,args...,k,ka,kb,0)
    A = spzeros(eltype(∇²[1]),prod(N),prod(N))
    for i ∈ eachindex(∇²)
        A += ∇²[i]*fs[i](k,ka,kb)
    end
    A = (A + SparseMatrixCSC(transpose(A)))/2 - s*spdiagm(0=>-sim.sys.ε[:]*k^2)
    B = s*spdiagm(0=>-sim.sys.F[:]*k^2)
    return A, B, η
end


text:
import Statistics: mean, cov
import Random.rand
import LinearAlgebra.logdet

"""
    Gaussian{(:μ,:Σ)}
    Gaussian{(:F,:Γ)}

Mitosis provides the measure `Gaussian` based on MeasureTheory.jl,
with a mean `μ` and covariance `Σ` parametrization,
or parametrised by natural parameters `F = Γ μ`, `Γ = Σ⁻¹`.

# Usage:

    Gaussian(μ=m, Σ=C)
    p = Gaussian{(:μ,:Σ)}(m, C)
    Gaussian(F=C\\m, Γ=inv(C))

    convert(Gaussian{(:F,:Γ)}, p)

    rand(rng, p)
"""
struct Gaussian{P,T} <: AbstractMeasure
    par::NamedTuple{P,T}
end
Gaussian{P}(nt::NamedTuple{P,T}) where {P,T} = Gaussian{P,T}(nt)
Gaussian{P}(args...) where {P} = Gaussian(NamedTuple{P}(args))
Gaussian(;args...) = Gaussian(args.data)
Gaussian{P}(;args...) where {P} = Gaussian{P}(args.data)

# the following propagates uncertainty if `μ` is `Gaussian`
Gaussian(par::NamedTuple{(:μ,:Σ),Tuple{T,S}}) where {T<:Gaussian,S} = Gaussian((;μ = mean(par.μ),Σ=par.Σ +cov(par.μ)))
Gaussian{P}(par::NamedTuple{(:μ,:Σ),Tuple{T,S}}) where {P,T<:Gaussian,S} = Gaussian{P}((;μ = mean(par.μ),Σ=par.Σ +cov(par.μ)))

Base.getproperty(p::Gaussian, s::Symbol) = getproperty(getfield(p, :par), s)

const GaussianOrNdTuple{P} = Union{Gaussian{P},NamedTuple{P}}

Base.keys(p::Gaussian{P}) where {P} = P
params(p::Gaussian) = getfield(p, :par)
## Basics

Base.:(==)(p1::Gaussian, p2::Gaussian) = mean(p1) == mean(p2) && cov(p1) == cov(p2)
Base.isapprox(p1::Gaussian, p2::Gaussian; kwargs...) =
    isapprox(mean(p1), mean(p2); kwargs...) && isapprox(cov(p1), cov(p2); kwargs...)

mean(p::Gaussian{(:μ, :Σ)}) = p.μ
mean(p::Gaussian{(:Σ,)}) = Zero()
cov(p::Gaussian{(:μ, :Σ)})  = p.Σ
meancov(p) = mean(p), cov(p)

precision(p::Gaussian{(:μ, :Σ)}) = inv(p.Σ)

mean(p::Gaussian{(:F, :Γ)}) = p.Γ\p.F
cov(p::Gaussian{(:F, :Γ)}) = inv(p.Γ)
precision(p::Gaussian{(:F, :Γ)}) = p.Γ
norm_sqr(x) = dot(x,x)
dim(p::Gaussian{(:F, :Γ)}) = length(p.F)
dim(p::Gaussian) = length(mean(p))
dim(p::Gaussian{(:Σ,)}) = size(p.Σ, 1)
whiten(p::Gaussian{(:μ, :Σ)}, x) = lchol(p.Σ)\(x - p.μ)
unwhiten(p::Gaussian{(:μ, :Σ)}, z) = lchol(p.Σ)*z + p.μ
function whiten(p::Gaussian{(:F, :Γ)}, x)
    L= lchol(p.Γ)
    L*x - L'\p.F
end
function unwhiten(p::Gaussian{(:F, :Γ)}, z)
    L = lchol(p.Γ)
    L\(z + L'\p.F)
end
whiten(p::Gaussian{(:Σ,)}, x) = lchol(p.Σ)\x
unwhiten(p::Gaussian{(:Σ,)}, z) = lchol(p.Σ)*z
sqmahal(p::Gaussian, x) = norm_sqr(whiten(p, x))

rand(p::Gaussian) = rand(Random.GLOBAL_RNG, p)
randwn(rng::AbstractRNG, x::Vector) = randn!(rng, zero(x))
function randwn(rng::AbstractRNG, x)
    map(xi -> randn(rng, typeof(xi)), x)
end
rand(rng::AbstractRNG, p::Gaussian) = unwhiten(p, randwn(rng, mean(p)))

_logdet(p::Gaussian{(:μ,:Σ)}) = _logdet(p.Σ, dim(p))
_logdet(p::Gaussian{(:Σ,)}) = logdet(p.Σ)
MeasureTheory.logdensity(p::Gaussian, x) = -(sqmahal(p,x) + _logdet(p) + dim(p)*log(2pi))/2
MeasureTheory.density(p::Gaussian, x) = exp(logdensity(p, x))
function MeasureTheory.logdensity(p::Gaussian{(:F,:Γ)}, x)
    C = cholesky(sym(p.Γ))
    -x'*p.Γ*x/2 + x'*p.F - p.F'*(C\p.F)/2  + logdet(C)/2 - dim(p)*log(2pi)/2
end

function logdensity0(p::Gaussian{(:F,:Γ)})
    C = cholesky(sym(p.Γ))
     - p.F'*(C\p.F)/2  + logdet(C)/2 - dim(p)*log(2pi)/2
end

function Base.convert(::Type{Gaussian{(:F,:Γ)}}, p::Gaussian{(:μ,:Σ)})
    Γ = inv(p.Σ)
    return Gaussian{(:F,:Γ)}(Γ*p.μ, Γ)
end
function Base.convert(::Type{Gaussian{(:μ,:Σ)}}, p::Gaussian{(:F,:Γ)})
    Σ = inv(p.Γ)
    return Gaussian{(:μ,:Σ)}(Σ*p.F, Σ)
end

## Algebra

Base.:+(p::Gaussian{P}, x) where {P} = Gaussian{P}(mean(p) + x, p.par[2])
Base.:+(x, p::Gaussian) = p + x

Base.:-(p::Gaussian, x) = p + (-x)
Base.:*(M, p::Gaussian{P}) where {P} = Gaussian{P}(M * mean(p), Σ = M * cov(p) * M')

⊕(p1::Gaussian{(:μ, :Σ)}, p2::Gaussian{(:μ, :Σ)}) = Gaussian{(:μ, :Σ)}(p1.μ + p2.μ, p1.Σ + p2.Σ)
⊕(x, p::Gaussian) = x + p
⊕(p::Gaussian, x) = p + x

## Conditionals and filtering

"""
    conditional(p::Gaussian, A, B, xB)

Conditional distribution of `X[i for i in A]` given
`X[i for i in B] == xB` if ``X ~ P``.
"""
function conditional(p::Gaussian{(:μ, :Σ)}, A, B, xB)
    Z = p.Σ[A,B]*inv(p.Σ[B,B])
    Gaussian{(:μ, :Σ)}(p.μ[A] + Z*(xB - p.μ[B]), p.Σ[A,A] - Z*p.Σ[B,A])
end


text:
function CN(ALPHA,BETA)
      A=[0    0.0180    0.0380    0.0560    0.0640    0.0740    0.0790
         0    0.0190    0.0420    0.0570    0.0770    0.0860    0.0900
         0    0.0180    0.0420    0.0590    0.0760    0.0930    0.1060
         0    0.0190    0.0420    0.0580    0.0740    0.0890    0.1060
         0    0.0190    0.0430    0.0580    0.0730    0.0800    0.0960
         0    0.0180    0.0390    0.0530    0.0570    0.0620    0.0800
         0    0.0130    0.0300    0.0320    0.0290    0.0490    0.0680
         0    0.0070    0.0170    0.0120    0.0070    0.0220    0.0300
         0    0.0040    0.0040    0.0020    0.0120    0.0280    0.0640
         0   -0.0140   -0.0350   -0.0460   -0.0340   -0.0120    0.0150
         0   -0.0170   -0.0470   -0.0710   -0.0650   -0.0020    0.0110
         0   -0.0330   -0.0570   -0.0730   -0.0410   -0.0130   -0.0010]
      #
      S= 0.2 * ALPHA
      K= floor(S)+3
      if(K <= 1)
          K= 2
      end
      if(K >=  12)
          K=  11
      end
      DA= S - K +3
      L = K + fix( 1.1*sign(DA) )
      #
      S= .2* abs(BETA)
      M= floor(S) + 1
      if(M <=  1)
          M= 2
      end
      if(M >=  7)
          M= 6
      end
      DB= S - M +1
      N= M + fix( 1.1*sign(DB) )
      #
      K=convert(Int64, K)
      L=convert(Int64, L)
      M=convert(Int64, M)
      N=convert(Int64, N)
      T= A[K,M]
      U= A[K,N]
      V= T + abs(DA) * (A[L,M] - T)
      W= U + abs(DA) * (A[L,N] - U)
      DUM= V + (W-V)  * abs(DB)
      return cnv = DUM * sign(BETA)
end


text:
using Base.Test
using CUDArt
using CUDNN
@show CUDNN_VERSION # this is set at runtime in CUDNN.jl, not fixed in types.jl

# Uncomment this if you want lots of messages:
# Base.Test.default_handler(r::Base.Test.Success) = info("$(r.expr)")
Base.Test.default_handler(r::Base.Test.Failure) = info("FAIL: $(r.expr)")

# See which operations support which dimensions:
function testdims()
    for n=3:8
        @show n
        try
            dims = [(n+1):-1:2]
            x = CudaArray(rand(dims...))
            y = CudaArray(rand(dims...))
            dx = CudaArray(rand(dims...))
            dy = CudaArray(rand(dims...))
            CUDNN.cudnnActivationBackward(CUDNN.CUDNN_ACTIVATION_RELU, y, dy, x, dx)
            CUDNN.cudnnActivationForward(CUDNN.CUDNN_ACTIVATION_RELU, x)
            CUDNN.cudnnScaleTensor(x, pi)
            CUDNN.cudnnSetTensor(x, pi)
            CUDNN.cudnnTransformTensor(2.0, x, 3.0, y)
            bdim = ones(dims); bdim[end-1]=dims[end-1]
            bias = CudaArray(rand(bdim...))
            CUDNN.cudnnAddTensor(bias, x)
        catch y
            println(y)
        end
    end
end

# Or rtfm:
# The interface of cuDNN has been generalized so that data sets with other than two
# spatial dimensions (e.g., 1D or 3D data) can be supported in future releases.
# o Note: while the interface now allows arbitrary N-dimensional tensors, most
# routines in this release remain limited to two spatial dimensions. This may
# be relaxed in future releases based on community feedback.
# o As a BETA preview in this release, the convolution forward, convolution
# weight and data gradient, and cudnnSetTensor/cudnnScaleTensor routines
# now support 3D datasets through the “Nd” interface. Note, however, that
# these 3D routines have not yet been tuned for optimal performance. This will
# be improved in future releases.

using CUDNN: cudnnGetVersion, cudnnGetErrorString, CUDNN_STATUS_SUCCESS
@show cudnnGetVersion()
@show bytestring(cudnnGetErrorString(CUDNN_STATUS_SUCCESS))

# TODO: handle type conflict between CUDArt and CUDNN when handling streams
# using CUDNN: cudnnCreate, cudnnDestroy, cudnnHandle_t, cudnnSetStream, cudnnGetStream, cudaStream_t
# hptr = cudnnHandle_t[0]
# cudnnCreate(hptr)
# sptr = cudaStream_t[0]
# CUDArt.rt.cudaStreamCreate(sptr) # library already checks == CUDArt.rt.cudaSuccess
# cudnnSetStream(hptr[1], sptr[1]) # library already checks == CUDNN_STATUS_SUCCESS in cudnnCheck
# tptr = cudaStream_t[0]
# cudnnGetStream(hptr[1], tptr)
# @test sptr[1] == tptr[1]
# CUDArt.rt.cudaStreamDestroy(sptr[1])
# cudnnDestroy(hptr[1])

using CUDNN: cudnnTensorDescriptor_t, cudnnCreateTensorDescriptor, cudnnSetTensor4dDescriptor, CUDNN_TENSOR_NCHW, CUDNN_DATA_DOUBLE, cudnnDataType_t, cudnnGetTensor4dDescriptor
d = cudnnTensorDescriptor_t[0]
cudnnCreateTensorDescriptor(pointer(d))
cudnnSetTensor4dDescriptor(d[1], CUDNN_TENSOR_NCHW, CUDNN_DATA_DOUBLE, 2, 3, 4, 5)
dt = cudnnDataType_t[0]
for n in (:sn, :sc, :sh, :sw, :tn, :tc, :th, :tw); @eval $n = Cint[0]; end
cudnnGetTensor4dDescriptor(d[1],dt,sn,sc,sh,sw,tn,tc,th,tw)
@test (dt[1],sn[1],sc[1],sh[1],sw[1],tn[1],tc[1],th[1],tw[1]) == (CUDNN_DATA_DOUBLE, 2, 3, 4, 5, 60, 20, 5, 1)

using CUDNN: cudnnGetTensorNdDescriptor
nd=4; nbDims=Cint[0]; dimA=Array(Cint,nd); strideA=Array(Cint,nd)
cudnnGetTensorNdDescriptor(d[1],nd,dt,nbDims,dimA,strideA)
@test (nbDims, dimA, strideA) == (Cint[4], Cint[2,3,4,5], Cint[60,20,5,1])

using CUDNN: TD
x = ones(5,4,3,2)
tx = CudaArray(x)
@test to_host(tx) == x
@test cudnnGetTensorNdDescriptor(TD(tx)) == (CUDNN_DATA_DOUBLE, 4, (reverse(size(x))...), (reverse(strides(x))...))

using CUDNN: cudnnTransformTensor
y = ones(5,4,3,2)
ty = CudaArray(y)
@test to_host(cudnnTransformTensor(2, tx, 3, ty)) == 2x+3y

using CUDNN: cudnnAddTensor
b = rand(1,1,3,1)
tb = CudaArray(b)
x = rand(5,4,3,2)
tx = CudaArray(x)
@test to_host(cudnnAddTensor(tb,tx)) == x .+ b
x = rand(5,4,3,2)
tx = CudaArray(x)
@test to_host(cudnnAddTensor(tb,tx;alpha=2,beta=3)) == 2b .+ 3x

using CUDNN: cudnnSetTensor
@test to_host(cudnnSetTensor(tx, pi)) == fill!(ones(size(tx)), pi)

using CUDNN: cudnnScaleTensor
x = rand(5,4,3,2)
tx = CudaArray(x)
@test to_host(cudnnScaleTensor(tx, pi)) == x .* pi

using CUDNN: cudnnActivationForward, CUDNN_ACTIVATION_RELU
myrelu(x,y)=(copy!(y,x);for i=1:length(y); (y[i]<zero(y[i]))&&(y[i]=zero(y[i])); end; y)
x = rand(5,4,3,2) - 0.5; tx = CudaArray(x)
y = zeros(5,4,3,2); ty = CudaArray(y)
@test to_host(cudnnActivationForward(tx, ty, mode=CUDNN_ACTIVATION_RELU)) == myrelu(x, y)

using CUDNN: cudnnActivationBackward
dy = (rand(5,4,3,2) - 0.5); tdy = CudaArray(dy)
dx = zeros(5,4,3,2); tdx = CudaArray(dx)
myrelu(y,dy,dx)=(copy!(dx,dy);for i=1:length(y); (y[i]==zero(y[i]))&&(dx[i]=zero(dx[i])); end; dx)
@test to_host(cudnnActivationBackward(ty, tdy, tx, tdx, mode=CUDNN_ACTIVATION_RELU)) == myrelu(y,dy,dx)

using CUDNN: CUDNN_ACTIVATION_SIGMOID
mysigm(x,y)=(for i=1:length(y); y[i]=(1.0/(1.0+exp(-x[i]))); end; y)
epseq(x,y)=(maximum(abs(x-y)) < 1e-7)
x = rand(5,4,3,2) - 0.5; tx = CudaArray(x)
y = zeros(5,4,3,2); ty = CudaArray(y)
@test epseq(to_host(cudnnActivationForward(tx, ty, mode=CUDNN_ACTIVATION_SIGMOID)), mysigm(x, y))

dy = (rand(5,4,3,2) - 0.5); tdy = CudaArray(dy)
dx = zeros(5,4,3,2); tdx = CudaArray(dx)
mysigm(y,dy,dx)=(for i=1:length(dx); dx[i]=dy[i]*y[i]*(1.0-y[i]); end; dx)
@test epseq(to_host(cudnnActivationBackward(ty, tdy, tx, tdx, mode=CUDNN_ACTIVATION_SIGMOID)), mysigm(y,dy,dx))

using CUDNN: CUDNN_ACTIVATION_TANH
mytanh(x,y)=(for i=1:length(y); y[i]=tanh(x[i]); end; y)
x = rand(5,4,3,2) - 0.5; tx = CudaArray(x)
y = zeros(5,4,3,2); ty = CudaArray(y)
@test epseq(to_host(cudnnActivationForward(tx, ty, mode=CUDNN_ACTIVATION_TANH)), mytanh(x, y))

dy = (rand(5,4,3,2) - 0.5); tdy = CudaArray(dy)
dx = zeros(5,4,3,2); tdx = CudaArray(dx)
mytanh(y,dy,dx)=(for i=1:length(dx); dx[i]=dy[i]*(1.0+y[i])*(1.0-y[i]); end; dx)
@test epseq(to_host(cudnnActivationBackward(ty, tdy, tx, tdx, mode=CUDNN_ACTIVATION_TANH)), mytanh(y,dy,dx))

using CUDNN: cudnnSoftmaxForward, CUDNN_SOFTMAX_MODE_CHANNEL, CUDNN_SOFTMAX_MODE_INSTANCE
x = (rand(5,4,3,2) - 0.5)
tx = CudaArray(x)
Base.zeros(a::AbstractCudaArray)=cudnnSetTensor(similar(a), zero(eltype(a)))
ty = zeros(tx)
@test epseq(to_host(cudnnSoftmaxForward(tx, ty; mode=CUDNN_SOFTMAX_MODE_INSTANCE)), exp(x)./sum(exp(x), (1,2,3)))
@test epseq(to_host(cudnnSoftmaxForward(tx, ty; mode=CUDNN_SOFTMAX_MODE_CHANNEL)), exp(x)./sum(exp(x), 3))

# Do we feed the gold probabilities p or the output gradients (1-p/q) to cudnnSoftmaxBackward?
# It turns out we feed 1-p/q, which is numerically unstable.
# x: unnormalized logp
# q: normalized estimates  qi=(exp xi)/(Σ exp xj)  q=softmaxForward(x)
# p: gold answers
# J = -Σ p log q
# dJ/dq = 1-p/q
# dJ/dx = q-p

using CUDNN: cudnnSoftmaxBackward
# Model probabilities:
x = (rand(5,4) - 0.5); tx = CudaArray(reshape(x,(1,1,5,4))); tq = zeros(tx)
cudnnSoftmaxForward(tx,tq); q = squeeze(to_host(tq),(1,2))
# Gold probabilities:
z = (rand(5,4) - 0.5); tz = CudaArray(reshape(z,(1,1,5,4))); tp = zeros(tz)
cudnnSoftmaxForward(tz,tp); p = squeeze(to_host(tp),(1,2))
# Gradient wrt softmax output:
dq = 1-p./q; tdq = CudaArray(reshape(dq,(1,1,5,4)))
# Gradient wrt softmax input:
tdx = similar(tdq)
cudnnSoftmaxBackward(tq, tdq, tdx)
dx = squeeze(to_host(tdx),(1,2))
@test epseq(dx,q-p)

# Discovering what pooling does:
# using CUDNN: cudnnPoolingDescriptor_t, cudnnCreatePoolingDescriptor, cudnnSetPooling2dDescriptor
# # x = rand(5,4,1,1)
# x = reshape(Float64[1:20;], 5, 4, 1, 1)
# tx = CudaArray(x)
# pdptr = Array(cudnnPoolingDescriptor_t, 1)
# cudnnCreatePoolingDescriptor(pdptr)
# pd = pdptr[1]
# using CUDNN: CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING, CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING, CUDNN_POOLING_MAX
# #cudnnSetPooling2dDescriptor(pd, CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING, 3, 3, 0, 0, 1, 1)
# cudnnSetPooling2dDescriptor(pd, CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING, 3, 3, 0, 0, 1, 1)
# #cudnnSetPooling2dDescriptor(pd, CUDNN_POOLING_MAX, 3, 3, 0, 0, 1, 1)
# ty = CudaArray(ones(10,9,1,1))
# using CUDNN: cudnnHandle, ptr, cudnnPoolingForward
# cudnnPoolingForward(cudnnHandle, pd, ptr(1,tx), tx.desc, tx.data.ptr, ptr(0,ty), ty.desc, ty.data.ptr)
# y = to_host(ty)
# dump(x)
# dump(y)

using CUDNN: CUDNN_POOLING_MAX, CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING, CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING, CUDNN_NOT_PROPAGATE_NAN
using CUDNN: PD, cudnnGetPoolingNdDescriptor
pd = PD(2, 3, 2, 1, CUDNN_POOLING_MAX)
if CUDNN_VERSION >= 4000
    @test cudnnGetPoolingNdDescriptor(pd) == (CUDNN_POOLING_MAX, CUDNN_NOT_PROPAGATE_NAN, 2, (3,3), (2,2), (1,1))
else
    @test cudnnGetPoolingNdDescriptor(pd) == (CUDNN_POOLING_MAX, 2, (3,3), (2,2), (1,1))
end
# free(pd)

using CUDNN: cudnnPoolingForward, cudnnPoolingBackward, cudnnGetPoolingNdForwardOutputDim, cudnnGetPoolingNdForwardOutputDim_buggy
x = reshape(Float64[1:20;], 5, 4, 1, 1); tx = CudaArray(x)

# 1.0   6.0  11.0  16.0
# 2.0   7.0  12.0  17.0
# 3.0   8.0  13.0  18.0
# 4.0   9.0  14.0  19.0
# 5.0  10.0  15.0  20.0

# 3 size, 0 pad, 1 stride
ty1 = CudaArray(zeros(3, 2, 1, 1))
pd1 = PD(2, 3, 0, 1, CUDNN_POOLING_MAX)
@test cudnnGetPoolingNdForwardOutputDim(pd1, tx) == (3,2,1,1)
@test cudnnGetPoolingNdForwardOutputDim_buggy(pd1, tx) == (3,2,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty1; pd=pd1)),(3,4)) == [13 18; 14 19; 15 20.]
ty2 = CudaArray(zeros(3, 2, 1, 1))
pd2 = PD(2, 3, 0, 1, CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING)
@test cudnnGetPoolingNdForwardOutputDim(pd2, tx) == (3,2,1,1)
@test cudnnGetPoolingNdForwardOutputDim_buggy(pd2, tx) == (3,2,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty2; pd=pd2)),(3,4)) == [7 12; 8 13; 9 14.]
ty3 = CudaArray(zeros(3, 2, 1, 1))
pd3 = PD(2, 3, 0, 1, CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING)
@test cudnnGetPoolingNdForwardOutputDim(pd3, tx) == (3,2,1,1)
@test cudnnGetPoolingNdForwardOutputDim_buggy(pd3, tx) == (3,2,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty3, pd=pd3)),(3,4)) == [7 12; 8 13; 9 14.]

dy1 = reshape(Float64[1:6;], 3, 2, 1, 1);
tdy1 = CudaArray(dy1)
tdx1 = zeros(tx)
@test squeeze(to_host(cudnnPoolingBackward(ty1, tdy1, tx, tdx1; pd=pd1)),(3,4)) == [0 0 0 0;0 0 0 0;0 0 1 4;0 0 2 5;0 0 3 6.]
tdx2 = zeros(tx)
@test epseq(squeeze(to_host(cudnnPoolingBackward(ty2, tdy1, tx, tdx2; pd=pd2)),(3,4)), [1/9 5/9 5/9 4/9;3/9 12/9 12/9 9/9;6/9 21/9 21/9 15/9;5/9 16/9 16/9 11/9;3/9 9/9 9/9 6/9])
tdx3 = zeros(tx)
@test epseq(squeeze(to_host(cudnnPoolingBackward(ty3, tdy1, tx, tdx3; pd=pd3)),(3,4)), [1/9 5/9 5/9 4/9;3/9 12/9 12/9 9/9;6/9 21/9 21/9 15/9;5/9 16/9 16/9 11/9;3/9 9/9 9/9 6/9])

# 3 size, 1 pad, 1 stride
ty4 = CudaArray(zeros(5, 4, 1, 1))
pd4 = PD(2, 3, 1, 1, CUDNN_POOLING_MAX)
@test cudnnGetPoolingNdForwardOutputDim(pd4, tx) == (5,4,1,1)
@test cudnnGetPoolingNdForwardOutputDim_buggy(pd4, tx) == (5,4,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty4; pd=pd4)),(3,4)) == [7 12 17 17; 8 13 18 18; 9 14 19 19; 10 15 20 20; 10 15 20 20.]
ty5 = CudaArray(zeros(5, 4, 1, 1))
pd5 = PD(2, 3, 1, 1, CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING)
@test cudnnGetPoolingNdForwardOutputDim(pd5, tx) == (5,4,1,1)
@test cudnnGetPoolingNdForwardOutputDim_buggy(pd5, tx) == (5,4,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty5; pd=pd5)),(3,4)) == [16/9 39/9 69/9 56/9; 3 7 12 87/9; 33/9 8 13 93/9; 39/9 9 14 11; 28/9 57/9 87/9 68/9]
ty6 = CudaArray(zeros(5, 4, 1, 1))
pd6 = PD(2, 3, 1, 1, CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING)
@test cudnnGetPoolingNdForwardOutputDim(pd6, tx) == (5,4,1,1)
@test cudnnGetPoolingNdForwardOutputDim_buggy(pd6, tx) == (5,4,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty6; pd=pd6)),(3,4)) == [16/4 39/6 69/6 56/4; 27/6 7 12 87/6; 33/6 8 13 93/6; 39/6 9 14 99/6; 28/4 57/6 87/6 68/4]

dy4 = reshape(Float64[1:20;], 5, 4, 1, 1);
tdy4 = CudaArray(dy4)
tdx4 = zeros(tx)
@test squeeze(to_host(cudnnPoolingBackward(ty4, tdy4, tx, tdx4; pd=pd4)),(3,4)) == [0 0 0 0;0 1 6 11+16;0 2 7 12+17;0 3 8 13+18;0 4+5 9+10 14+15+19+20.]
tdx5 = zeros(tx)
@test epseq(squeeze(to_host(cudnnPoolingBackward(ty5, tdy4, tx, tdx5; pd=pd5)),(3,4)), [16/9 39/9 69/9 56/9; 3 7 12 87/9; 33/9 8 13 93/9; 39/9 9 14 11; 28/9 57/9 87/9 68/9])
tdx6 = zeros(tx)
@test epseq(squeeze(to_host(cudnnPoolingBackward(ty6, tdy4, tx, tdx6; pd=pd6)),(3,4)), [2.361111111111111 5.527777777777778 11.777777777777779 10.0;  3.75 8.36111111111111 17.11111111111111 14.444444444444445;  4.166666666666666 8.5 16.0 13.333333333333332;  5.972222222222222 11.472222222222221 20.22222222222222 16.666666666666668; 4.583333333333334 8.63888888888889 14.88888888888889 12.222222222222221])

# 3 size, 1 pad, 2 stride
ty7 = CudaArray(zeros(3, 3, 1, 1))
pd7 = PD(2, 3, 1, 2, CUDNN_POOLING_MAX)
@test cudnnGetPoolingNdForwardOutputDim(pd7, tx) == (3,3,1,1)
@show cudnnGetPoolingNdForwardOutputDim_buggy(pd7, tx), (3,3,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty7; pd=pd7)),(3,4)) == [7 17 17; 9 19 19; 10 20 20.]
# Note below that if the window falls outside the padding, the denom can be less than 9!
ty8 = CudaArray(zeros(3, 3, 1, 1))
pd8 = PD(2, 3, 1, 2, CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING)
@test cudnnGetPoolingNdForwardOutputDim(pd8, tx) == (3,3,1,1)
@show cudnnGetPoolingNdForwardOutputDim_buggy(pd8, tx), (3,3,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty8; pd=pd8)),(3,4)) == [16/9 69/9 33/6; 33/9 13 54/6; 28/9 87/9 39/6]
ty9 = CudaArray(zeros(3, 3, 1, 1))
pd9 = PD(2, 3, 1, 2, CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING)
@test cudnnGetPoolingNdForwardOutputDim(pd9, tx) == (3,3,1,1)
@show cudnnGetPoolingNdForwardOutputDim_buggy(pd9, tx), (3,3,1,1)
@test squeeze(to_host(cudnnPoolingForward(tx, ty9; pd=pd9)),(3,4)) == [16/4 69/6 33/2; 33/6 13 54/3; 28/4 87/6 39/2]

dy7 = reshape(Float64[1:9;], 3, 3, 1, 1);
tdy7 = CudaArray(dy7)
tdx7 = zeros(tx)
# TODO: check this answer
@test squeeze(to_host(cudnnPoolingBackward(ty7, tdy7, tx, tdx7; pd=pd7)),(3,4)) == [0 0 0 0;0 1 0 11;0 0 0 0;0 2 0 13;0 3 0 15.]
tdx8 = zeros(tx)
# TODO: check this answer
@test epseq(squeeze(to_host(cudnnPoolingBackward(ty8, tdy7, tx, tdx8; pd=pd8)),(3,4)), [1/9 5/9 4/9 29/18; 1/3 4/3 1 7/2; 2/9 7/9 5/9 17/9; 5/9 16/9 11/9 73/18; 1/3 1 2/3 39/18])
tdx9 = zeros(tx)
@test epseq(squeeze(to_host(cudnnPoolingBackward(ty9, tdy7, tx, tdx9; pd=pd9)),(3,4)), [0.25 0.9166666666666666 0.6666666666666666 4.166666666666667; 0.5833333333333333 1.8055555555555554 1.2222222222222223 7.388888888888889; 0.3333333333333333 0.8888888888888888 0.5555555555555556 3.2222222222222223; 1.0833333333333333 2.638888888888889 1.5555555555555556 8.722222222222221; 0.75 1.75 1.0 5.5])

# 3D pooling support added in v3
x10 = reshape(Float64[1:60;], 5, 4, 3, 1, 1); tx10 = CudaArray(x10)
ty10 = CudaArray(zeros(3, 2, 1, 1, 1))
pd10 = PD(3, 3, 0, 1, CUDNN_POOLING_MAX)
# TODO: add 3d pooling tests
# @show cudnnPoolingForward(tx10, ty10; pd=pd10)

# Filters are basically the same as tensors:

using CUDNN: cudnnFilterDescriptor_t, cudnnCreateFilterDescriptor, cudnnSetFilter4dDescriptor, CUDNN_TENSOR_NCHW, CUDNN_DATA_DOUBLE, cudnnDataType_t, cudnnGetFilter4dDescriptor
d = cudnnFilterDescriptor_t[0]
cudnnCreateFilterDescriptor(d)
cudnnSetFilter4dDescriptor(d[1], CUDNN_DATA_DOUBLE, 2, 3, 4, 5)
dt = cudnnDataType_t[0]
for n in (:sk, :sc, :sh, :sw); @eval $n = Cint[0]; end
cudnnGetFilter4dDescriptor(d[1],dt,sk,sc,sh,sw)
@test (dt[1],sn[1],sc[1],sh[1],sw[1],tn[1],tc[1],th[1],tw[1]) == (CUDNN_DATA_DOUBLE, 2, 3, 4, 5, 60, 20, 5, 1)

using CUDNN: cudnnGetFilterNdDescriptor
nd=4; nbDims=Cint[0]; dimA=Array(Cint,nd)
cudnnGetFilterNdDescriptor(d[1],nd,dt,nbDims,dimA)
@test (nbDims, dimA) == (Cint[4], Cint[2,3,4,5])

using CUDNN: FD
x = rand(5,4,3,2)
tx = CudaArray(x)
@test to_host(tx) == x
@test cudnnGetFilterNdDescriptor(FD(tx)) == (CUDNN_DATA_DOUBLE, CUDNN_TENSOR_NCHW, 4, (reverse(size(x))...))

# Convolution

using CUDNN: CD, cudnnGetConvolutionNdDescriptor, CUDNN_CONVOLUTION
cd = CD(2, 0, 1, 1, CUDNN_CONVOLUTION, Float32)
if CUDNN_VERSION >= 3000
@test cudnnGetConvolutionNdDescriptor(cd) == (2, (0,0), (1,1), (1,1), CUDNN_CONVOLUTION, Float32)
else # if CUDNN_VERSION >= 3000
@test cudnnGetConvolutionNdDescriptor(cd) == (2, (0,0), (1,1), (1,1), CUDNN_CONVOLUTION)
end # if CUDNN_VERSION >= 3000
# Note: upscale other than (1,1) gives unsupported error.
# Note: not sure if we need to expose the ConvolutionDescriptor or just have options for convolution.
# Note: need to understand upscale.  how often do we need non-default padding and stride?
# Note: conv vs xcor

using CUDNN: cudnnGetConvolutionNdForwardOutputDim
@test cudnnGetConvolutionNdForwardOutputDim(CudaArray(ones(12,8,3,6)), CudaArray(ones(5,4,3,2))) == (8,5,2,6)
# Does not work for dimensions other than 4 yet:
# @show cudnnGetConvolutionNdForwardOutputDim(CudaArray(ones(13,12,11,3,10)), CudaArray(ones(6,5,4,3,2)))

using CUDNN: CUDNN_CONVOLUTION_FWD_NO_WORKSPACE, CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT
using CUDNN: CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM, CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM, CUDNN_CONVOLUTION_FWD_ALGO_GEMM, CUDNN_CONVOLUTION_FWD_ALGO_DIRECT, CUDNN_CONVOLUTION_FWD_ALGO_FFT, CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING

using CUDNN: cudnnGetConvolutionForwardAlgorithm
src = CudaArray(ones(102,101,3,100))
flt = CudaArray(ones(25,15,3,99))
dst = CudaArray(ones(cudnnGetConvolutionNdForwardOutputDim(src, flt)))
@test cudnnGetConvolutionForwardAlgorithm(src, flt, dst; preference=CUDNN_CONVOLUTION_FWD_NO_WORKSPACE) == CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM
@test cudnnGetConvolutionForwardAlgorithm(src, flt, dst; preference=CUDNN_CONVOLUTION_FWD_PREFER_FASTEST) == CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM
@test cudnnGetConvolutionForwardAlgorithm(src, flt, dst; preference=CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT) == CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM

# TODO: check if the 0 answers are accurate
using CUDNN: cudnnGetConvolutionForwardWorkspaceSize
@test cudnnGetConvolutionForwardWorkspaceSize(src, flt, dst; algorithm=CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM) == 0
@test cudnnGetConvolutionForwardWorkspaceSize(src, flt, dst; algorithm=CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM) == 4500 # size(flt,1)*size(flt,2)*size(flt,3)*sizeof(Cint)
@test cudnnGetConvolutionForwardWorkspaceSize(src, flt, dst; algorithm=CUDNN_CONVOLUTION_FWD_ALGO_GEMM) == 1812432704 # this is from v4, v3:6107400000
@test cudnnGetConvolutionForwardWorkspaceSize(src, flt, dst; algorithm=CUDNN_CONVOLUTION_FWD_ALGO_DIRECT) == 0
@test cudnnGetConvolutionForwardWorkspaceSize(src, flt, dst; algorithm=CUDNN_CONVOLUTION_FWD_ALGO_FFT) == 0
@test cudnnGetConvolutionForwardWorkspaceSize(src, flt, dst; algorithm=CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING) == 0

using CUDNN: cudnnConvolutionForward
x = reshape(Float64[1:20;], 5, 4, 1, 1); tx = CudaArray(x)
w = reshape(Float64[1:4;], 2, 2, 1, 1); tw = CudaArray(w)
@test squeeze(to_host(cudnnConvolutionForward(tx, tw)),(3,4)) == [29 79 129; 39 89 139; 49 99 149; 59 109 159.]

using CUDNN: CD, CUDNN_CROSS_CORRELATION
cdesc = CD(tx; mode=CUDNN_CROSS_CORRELATION)
@test squeeze(to_host(cudnnConvolutionForward(tx, tw; cd=cdesc)),(3,4)) == [51 101 151;61 111 161;71 121 171;81 131 181.]

using CUDNN: cudnnConvolutionBackwardBias, cudnnConvolutionBackwardFilter, cudnnConvolutionBackwardData
x = rand(5,4,3,2); tx = CudaArray(x)
w = rand(2,2,3,4); tw = CudaArray(w)
ty = cudnnConvolutionForward(tx, tw)
dy = rand(size(ty)); tdy = CudaArray(dy)
@test epseq(to_host(cudnnConvolutionBackwardBias(tdy)), sum(dy,(1,2,4)))

# TODO: put a more meaningful test here...
tdx = zeros(tx)
@test size(cudnnConvolutionBackwardData(tw, tdy, tdx)) == size(x)

tdw = zeros(tw)
@test size(cudnnConvolutionBackwardFilter(tx, tdy, tdw)) == size(w)

# Testing correspondence with conv2
x = rand(5,4); tx = CudaArray(reshape(x, (5,4,1,1)))
w = rand(3,3); tw = CudaArray(reshape(w, (3,3,1,1)))
padding=map(x->x-1,size(w))
@test epseq(squeeze(to_host(cudnnConvolutionForward(tx, tw; cd=CD(tx; padding=padding))),(3,4)), conv2(x,w))
@test epseq(squeeze(to_host(conv2(tx, tw)), (3,4)), conv2(x,w))

if CUDNN_VERSION >= 3000
    # LRN: TODO: add tests
    x = rand(5,4); tx = CudaArray(reshape(x, (5,4,1,1)))
    ty = similar(tx)
    cudnnLRNCrossChannelForward(tx,ty)
    y = squeeze(to_host(ty), (3,4))
    dy = rand(5,4); tdy = CudaArray(reshape(x, (5,4,1,1)))
    tdx = similar(tdy)
    cudnnLRNCrossChannelBackward(ty, tdy, tx, tdx)
    dx = squeeze(to_host(tdx), (3,4))

    # DivisiveNormalization
    x = rand(5,4); tx = CudaArray(reshape(x, (5,4,1,1)))
    ty = similar(tx)
    cudnnDivisiveNormalizationForward(tx,ty)
    y = squeeze(to_host(ty), (3,4))
    dy = rand(5,4); tdy = CudaArray(reshape(x, (5,4,1,1)))
    tdx = similar(tdy)
    cudnnDivisiveNormalizationBackward(ty, tdy, tdx)
    dx = squeeze(to_host(tdx), (3,4))

    # println(x)
    # println(y)
    # println(dy)
    # println(dx)
end # if CUDNN_VERSION >= 3000

:ok


text:
"""
TcpInstruments is an ongoing effort to provide a simple unified
interface to common lab equipment.

To connect to an instrument use:
```
handle = initialize({Instrument-Type}, "IP-ADDRESS:PORT")
```

To prevent hardcoding in scripts and/or to make it simplier for
your lab or organization to keep track of all the equipment and
ip addresses a configuration file can also be specified.

To use Orchard Ultrasound's latest lab config use:
```
create_config()
```
This will create a yaml file in your home directory: ~/.tcp_instruments.yml

This yaml file will be loaded everytime you use this package.

You can also create a project-specific config by creating
the config in your project root directory instead of your home
directory. You can do this with:
```
create_config(pwd())
```

Once you have created a config file you can change it with
```
edit_config()
```

If the instrument has its address in the config file you can now connect with:
```
handle = initialize({Instrument-Type})
```

To see the different types of devices you can interface with
use `help>Instrument`.
"""
module TcpInstruments

using Sockets
using Unitful
using Unitful: s, ms, μs, ns, ps

import InstrumentConfig: initialize, terminate

export Instrument
export Oscilloscope, MultiMeter, PowerSupply, WaveformGenerator, ImpedanceAnalyzer

export save, load

export initialize, terminate, reset
export remote_mode, local_mode
export query, write, info, connect!, close!

# Power Supply
export enable_output, disable_output, get_output
export set_current_limit, get_current_limit
export set_voltage, get_voltage
export set_voltage_limit, get_voltage_limit
export set_channel, get_channel

    # Scope
export get_data
export lpf_on, lpf_off, get_lpf_state
export set_impedance_1Mohm, set_impedance_50ohm, get_impedance
export get_coupling
export get_function, set_function
export get_frequency, set_frequency
export get_amplitude, set_amplitude
export get_voltage_offset, set_voltage_offset
export set_burst_mode_gated, set_burst_mode_triggered, get_burst_mode
export get_mode, set_mode_burst, set_mode_cw

# Prologix
export set_prologix_chan, get_prologix_chan, scan_prologix

# DMM
export get_tc_temperature, set_tc_type
export get_current
export get_resistance
export set_temp_unit_celsius, set_temp_unit_farenheit, set_temp_unit_kelvin
export get_temp_unit

# Impedance
export get_bandwidth, set_bandwidth
export get_frequency_limits, set_frequency_limits
export get_num_data_points, set_num_data_points
export get_volt_dc, set_volt_dc
export get_volt_ac, set_volt_ac
export get_volt_limit_dc, set_volt_limit_dc

# Signal Generator
export get_frequency, set_frequency
export get_amplitude, set_amplitude
export get_burst_num_cycles, set_burst_num_cycles
export get_time_offset, set_time_offset
export get_voltage_offset, set_voltage_offset
export get_burst_period, set_burst_period
export get_mode, set_mode_burst, set_mode_cw

# Devices
## Impedance Analyzer
export Agilent4294A, Agilent4395A
## Multimeter
export KeysightDMM34465A
## Scope
export AgilentDSOX4024A, AgilentDSOX4034A
## Power Supply
export AgilentE36312A, SRSPS310, VersatilePower
## Waveform Generator
export Keysight33612A

export scan_network

include("util.jl")
include("config.jl")

# common instrument containers
include("instrument.jl")
include("common_commands.jl")

# instruments
include("oscilloscope/all.jl")
include("power_supply/all.jl")
include("signal_generator/all.jl")
include("impedance_analyzer/all.jl")
include("multimeter/all.jl")

include("emulator/emulator.jl")
end #endmodule


text:
path = dirname(@__FILE__)
@testset "Check prior overloading" begin
    custom_settings = Dict{Symbol, Setting}(
        :date_forecast_start  => Setting(:date_forecast_start, quartertodate("2015-Q4")))
    m = AnSchorfheide(custom_settings = custom_settings, testing = true)
    dsgevar = DSGE.DSGEVAR(m, collect(keys(m.exogenous_shocks)))
    @test prior(m) == prior(m.parameters)
    @test prior(dsgevar) == prior(dsgevar.dsge) == prior(m.parameters)
end

@testset "Check DSGE likelihood and posterior calculations" begin
    custom_settings = Dict{Symbol, Setting}(
        :date_forecast_start  => Setting(:date_forecast_start, quartertodate("2015-Q4")))
    m = AnSchorfheide(custom_settings = custom_settings, testing = true)

    file = "$path/../reference/posterior.jld2"
    data = Matrix{Float64}(load(file, "data")')
    lh_expected = load(file, "lh_expected")
    post_expected = load(file, "post_expected")

    lh = likelihood(m, data)
    @test lh_expected ≈ lh

    post = DSGE.posterior(m, data)
    @test post_expected ≈ post

    x = map(α->α.value, m.parameters)
    post_at_start = DSGE.posterior!(m, x, data)
    @test post_expected ≈ post_at_start

    # Ensure if we are not evaluating at start vector, then we do not get the reference
    # posterior
    global y = x .+ 0.01
    post_not_at_start = DSGE.posterior!(m, y, data)
    ϵ = 1.0
    @test abs(post_at_start - post_not_at_start) > ϵ

    # Check keywords
    m.parameters[2].value = 3.
    @test !isinf(likelihood(m, data; sampler = false))
    @test isinf(likelihood(m, data; sampler = true))
end

@testset "Check DSGEVAR likelihood and posterior calculations" begin
    dsge = Model1002("ss10")
    obs_i = [dsge.observables[:obs_nominalrate], dsge.observables[:obs_gdp],
             dsge.observables[:obs_gdpdeflator]]
    m = DSGE.DSGEVAR(dsge, collect(keys(dsge.exogenous_shocks)))
    DSGE.update!(m; observables = [:obs_nominalrate, :obs_gdp, :π_t], lags = 4, λ = 0.5)
    dsge_data =
        df_to_matrix(dsge,
                     CSV.read(joinpath(path, "../reference/test_dsgevar_likelihood_dsge_data.csv")))[obs_i, :]
    lh = likelihood(m, dsge_data)
    @test lh ≈ DSGE.dsgevar_likelihood(m, dsge_data) ≈
        load(joinpath(path, "../reference/test_dsgevar_likelihood_dsge.jld2"), "llh")

    post = DSGE.posterior(m, dsge_data)
    @test post ≈ lh + prior(m)

    x = map(α -> α.value, m.dsge.parameters)
    post_at_start = DSGE.posterior!(m, x, dsge_data)
    @test post ≈ post_at_start

    # Ensure if we are not evaluating at start vector, then we do not get the reference
    # posterior
    global y = x .+ 0.01
    post_not_at_start = DSGE.posterior!(m, y, dsge_data)
    ϵ = 1.0
    @test abs(post_at_start - post_not_at_start) > ϵ

    # Check keywords
    m.dsge.parameters[1].value = 3.
    @test !isinf(likelihood(m, dsge_data; sampler = false))
    @test isinf(likelihood(m, dsge_data; sampler = true))
end


text:
using Images, ShapeFromShading
using Test

@testset "ShapeFromShading" begin
    include("shah.jl")
    include("pentland.jl")
    include("discreteshape.jl")
    include("discreteshapebound.jl")
    include("photometric.jl")
    include("deterministic_custom.jl")
    include("deterministic.jl")
end

@testset "Integration" begin
    include("frankot.jl")
    include("path.jl")
    include("splitpath.jl")
    include("horn.jl")
    include("durou.jl")
end

@testset "SyntheticSurface" begin
    include("synthsphere.jl")
    include("synthvase.jl")
    include("tent.jl")
    include("dem.jl")
    include("cake.jl")
    include("cake2.jl")
    include("ripple.jl")
    include("ripple2.jl")
    include("synthgaussian.jl")
    include("supergaussian.jl")
    include("estimatealbedo.jl")
end


text:
using Genie, Genie.Renderer.Html, Stipple, StipplePlotly

Genie.config.log_requests = false

xx = -π:(2π/250):π

xxs = -3.0:0.2:3.0

# Data:

pl1 = PlotData(
    x = xx, y = sin.(xx), plot = StipplePlotly.Charts.PLOT_TYPE_SCATTER,
    name = "sine", mode = "lines", xaxis = "x", yaxis = "y", line = PlotlyLine(color = "rgb(0,0,192)", dash="solid")
)

pl2 = PlotData(
    x = xx, y = sinh.(xx), plot = StipplePlotly.Charts.PLOT_TYPE_SCATTER,
    name = "sinh", mode = "lines", xaxis = "x2", yaxis = "y2", line = PlotlyLine(color = "rgb(0,192,0)", dash="dot")
)

pl3 = PlotData(
    x = xx, y = cos.(xx), plot = StipplePlotly.Charts.PLOT_TYPE_SCATTER,
    name = "cosine", mode = "lines", xaxis = "x3", yaxis = "y3", line = PlotlyLine(color = "rgb(192,0,0)", dash="dash")
)

pl4 = PlotData(
    x = xx, y = cosh.(xx), plot = StipplePlotly.Charts.PLOT_TYPE_SCATTER,
    name = "cosh", mode = "lines", xaxis = "x4", yaxis = "y4", line = PlotlyLine(color = "rgb(192,0,192)", dash="dashdot")
)

ps1 = PlotData(
    x = xxs, y = sin.(xxs) .+ rand(Float64, size(xxs)) .- 0.5, plot = StipplePlotly.Charts.PLOT_TYPE_SCATTER,
    name = "sine", mode = "markers", xaxis = "x", yaxis = "y", marker = PlotDataMarker(color="rgb(0,0,192)", symbol="circle", size=10, opacity=0.5)
)

ps2 = PlotData(
    x = xxs, y = sinh.(xxs) .+ 3.0 .* rand(Float64, size(xxs)) .- 1.5, plot = StipplePlotly.Charts.PLOT_TYPE_SCATTER,
    name = "sinh", mode = "markers", xaxis = "x2", yaxis = "y2", marker = PlotDataMarker(color = "rgb(0,192,0)", symbol="circle-open", size=14)
)

ps3 = PlotData(
    x = xxs, y = cos.(xxs) .+ rand(Float64, size(xxs)) .- 0.5, plot = StipplePlotly.Charts.PLOT_TYPE_SCATTER,
    name = "cosine", mode = "markers", xaxis = "x3", yaxis = "y3", marker = PlotDataMarker(color = "rgb(192,0,0)", symbol="diamond", size=10, opacity=0.5)
)

ps4 = PlotData(
    x = xxs, y = cosh.(xxs) .+ 3.0 .* rand(Float64, size(xxs)) .- 1.5, plot = StipplePlotly.Charts.PLOT_TYPE_SCATTER,
    name = "cosh", mode = "markers", xaxis = "x4", yaxis = "y4", marker = PlotDataMarker(color = "rgb(192,0,192)", symbol="diamond-open", size=3)
)

plotdata = [ps1, ps2, ps3, ps4, pl1, pl2, pl3, pl4];

# Layout

layout = PlotLayout(
    title = PlotLayoutTitle(text="Multiple Mixed Subplots", font=Font(24)),
    showlegend = false,
    grid = PlotLayoutGrid(rows = 2, columns = 2, pattern = "independent"),
    xaxis = [
        PlotLayoutAxis(xy = "x", index = 1, ticks = "outside", showline = true, zeroline = false),
        PlotLayoutAxis(xy = "x", index = 2, ticks = "outside", showline = true, zeroline = false),
        PlotLayoutAxis(xy = "x", index = 3, ticks = "outside", showline = true, zeroline = false, title="range (arb. units)"),
        PlotLayoutAxis(xy = "x", index = 4, ticks = "outside", showline = true, zeroline = false, title="range (arb. units)")
    ],
    yaxis = [
        PlotLayoutAxis(xy = "y", index = 1, ticks = "outside", showline = true, zeroline = false, title="response A"),
        PlotLayoutAxis(xy = "y", index = 2, ticks = "outside", showline = true, zeroline = false, title="response B"),
        PlotLayoutAxis(xy = "y", index = 3, ticks = "outside", showline = true, zeroline = false, title="response C"),
        PlotLayoutAxis(xy = "y", index = 4, ticks = "outside", showline = true, zeroline = false, title="response D")
    ],
)

Base.@kwdef mutable struct Model <: ReactiveModel
  data::R{Vector{PlotData}} = plotdata, READONLY
  layout::R{PlotLayout} = layout, READONLY
  config::R{PlotConfig} = PlotConfig(), READONLY
end

model = Stipple.init(Model())

function ui()
  page(
    vm(model), class="container", [
        plot(:data, layout = :layout, config = :config)
    ]
  ) |> html
end

route("/", ui)

up()


text:
abstract type Score end
# expectation fields, must be mutable!
# - prior::Prior

# make scores iterable
Base.iterate(score::Score, state = 0) = state > 0 ? nothing : (score, state + 1)
Base.length(score::Score) = 1

# reduce Base.show() to Base.string()
Base.show(io::IO, score::Score) = print(io, string(score))
Base.show(io::IO, ::MIME"application/prs.juno.inline", score::Score) = print(io, string(score))
Base.string(score::Score) = @sprintf "ToDo: implement Base.string for %s" typeof(score)

# make scores callable
(score::Score)(args...; kwargs...) = evaluate(score, args...; kwargs...)

# update the prior of a score
update!(score::TS, x::TI, n::TI) where {TS<:Score,TI<:Integer} = score.prior = update(score.prior, x, n)

function update(score::Score, prior::Prior)
    error("not implemented")
end

function evaluate(score::TS, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR, p::TR)::TR where {TI<:Integer,TR<:Real,TS<:Score}

    error(@sprintf "ToDo: implement bad.evaluate for %s" typeof(score))
end
function evaluate(score::TS, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR) where {TI<:Integer,TR<:Real,TS<:Score}
    expectation(
        p -> evaluate(score, x1, n1, x2, n2, c2, p),
        update(score.prior, x1 + x2, n1 + n2)
    )
end

function evaluate(score::TS, design::TD, x1::TI, x2::TI, p::TR) where {TI<:Integer,TR<:Real,TD<:AbstractDesign,TS<:Score}

    evaluate(score, x1, n1(design), x2, n2(design, x1), c2(design, x1), p)
end
function evaluate(score::TS, design::TD, x1::TI, x2::TI) where {TI<:Integer,TD<:AbstractDesign,TS<:Score}
    expectation(
        p -> evaluate(score, design, x1, x2, p),
        update(score.prior, x1 + x2, n1(design) + n2(design, x1))
    )
end

# conditional expectation given x1
function evaluate(score::TS, x1::TI, n1::TI, n2::TI, c2::TR, p::TR; partial_stage_two::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TR<:Real,TS<:Score}

    @assert 0 <= partial_stage_two[1] <= partial_stage_two[2] <= n2
    X2 = collect(partial_stage_two[1]:n2)
    return sum(
        evaluate.(score, x1, n1, X2, n2, c2, p) .*
        pmf.(X2, n2, p; partial = partial_stage_two)
    )
end
function evaluate(score::TS, x1::TI, n1::TI, n2::TI, c2::TR; partial_stage_two::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TR<:Real,TS<:Score}
    expectation(
        p -> evaluate(score, x1, n1, n2, c2, p; partial_stage_two = partial_stage_two),
        update(score.prior, x1 + partial_stage_two[1], n1 + partial_stage_two[2])
    )
end

function evaluate(score::TS, design::TD, x1::TI, p::TR; partial_stage_two::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TR<:Real,TD<:AbstractDesign,TS<:Score}

    return evaluate(score, x1, n1(design), n2(design, x1), c2(design, x1), p; partial_stage_two = partial_stage_two)
end
function evaluate(score::TS, design::TD, x1::TI; partial_stage_two::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TD<:AbstractDesign,TS<:Score}
    expectation(
        p -> evaluate(score, design, x1, p; partial_stage_two = partial_stage_two),
        update(score.prior, x1 + partial_stage_two[1], n1(design) + partial_stage_two[2])
    )
end

# marginalise X1, X2
function evaluate(score::TS, design::TD, p::TR; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TR<:Real,TD<:AbstractDesign,TS<:Score}

    @assert 0 <= partial_stage_one[1] <= partial_stage_one[2] <= n1(design)
    XX = sample_space(design)
    XX = XX[XX[:,1] .>= partial_stage_one[1], :]
    return sum(
        evaluate.(score, design, XX[:,1], XX[:,2], p) .*
        pmf.(XX[:,2], n2.(design, XX[:,1]), p) .*
        pmf.(XX[:,1], n1(design), p; partial = partial_stage_one)
    )
end
function evaluate(score::TS, design::TD; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TD<:AbstractDesign,TS<:Score}

    expectation(
        p -> evaluate(score, design, p; partial_stage_one = partial_stage_one),
        update(score.prior, partial_stage_one[1], partial_stage_one[1])
    )
end

# compute Pr[X1=x1]*E_prior[s(x1, n1, X2, n2(x1), c2(x1))] = score_integrand(x1, n1, n2, c2)
# this is the only thing that actually used during optmisation
# feel free to implement specific, more efficient versions for each score!
function integrand_x1(score::TS, x1::TI, n1::TI, n2::TI, c2::TR; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TS<:Score,TI<:Integer,TR<:Real}
    score(x1, n1, n2, c2) * pmf(x1, n1, score.prior; partial = partial_stage_one)
end

mutable struct SampleSize{TP<:Prior} <: Score
    prior::TP
end
Base.string(score::SampleSize) = @sprintf "SampleSize<%s>" string(score.prior)

function evaluate(score::SampleSize, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR, p::TR)::TR where {TI<:Integer,TR<:Real,TD<:AbstractDesign}
    n1 + n2
end
# no custom score_integrand necessary - fallback to default

update(score::SampleSize, prior::Prior) = SampleSize(prior)

mutable struct Power{TP<:Prior,TR<:Real} <: Score
    prior::TP
    bounds::Tuple{TR,TR}
    Power{TP,TR}(prior::TP, bounds::Tuple{TR,TR}) where{TP<:Prior,TR<:Real} = new(condition(prior; low = bounds[1], high = bounds[2]), bounds)
end
Power(prior::Prior; bounds::Tuple{Real,Real} = Tuple(bounds(prior))) = Power{typeof(prior),eltype(bounds)}(prior, bounds)
Base.string(score::Power) = @sprintf "Power<%s>" string(score.prior)

update(score::Power, prior::Prior) = Power(prior; bounds = score.bounds)

function evaluate(score::Power, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR, p::TR)::TR where {TI<:Integer,TR<:Real}
    (score.bounds[1] <= p <= score.bounds[2]) ? (x2 > c2) : 0.0
end

# since we integrate power, it is easier to do that in one wash, also
# closed form cdf is available, note that the score prior is already properly conditioned
function integrand_x1(score::Power, x1::TI, n1::TI, n2::TI, c2::TR; partial_stage_one::Tuple{TI,TI} = (0, 0) ) where {TS<:Score,TI<:Integer,TR<:Real}

    expectation( # conditional expected power
        p -> 1 - cdf(c2, n2, p),
        update(score.prior, x1, n1)
    ) * pmf( # conditional pmf given partial observations
        x1, n1, score.prior; partial = partial_stage_one
    )
end

mutable struct CompositeScore <: Score
    ω
    components
end
Base.string(score::CompositeScore) = "CompositeScore"

Base.convert(CompositeScore, score::Score) = CompositeScore([1.0], [score])
Base.convert(::Type{T}, score::T)  where {T<:Score} = score

*(ω::Real, score::Score) = CompositeScore([ω], [score])
*(ω::Real, score::CompositeScore) = CompositeScore(ω .* score.ω, score.components)
function +(score1::T1, score2::T2) where {T1,T2<:Score}

    score1, score2 = convert.(CompositeScore, (score1, score2))
    return CompositeScore(vcat(score1.ω, score2.ω), vcat(score1.components, score2.components))
end
function -(score1::T1, score2::T2) where {T1,T2<:Score}

    score1, score2 = convert.(CompositeScore, (score1, score2))
    return CompositeScore(vcat(score1.ω, -score2.ω), vcat(score1.components, score2.components))
end
function update!(score::CompositeScore, x::TI, n::TI) where {TS<:Score,TI<:Integer}
    map(component -> update!(component, x, n), score.components)
end

function update(score::CompositeScore, prior::Prior)
    score = deepcopy(score)
    score.components = map(component -> update(component, prior), score.components)
    return score
end

evaluate(score::CompositeScore, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR, p::TR) where {TI<:Integer,TR<:Real} = sum( score.ω .* evaluate.(score.components, x1, n1, x2, n2, c2, p) )
evaluate(score::CompositeScore, x1::TI, n1::TI, x2::TI, n2::TI, c2::TR) where {TI<:Integer,TR<:Real} = sum( score.ω .* evaluate.(score.components, x1, n1, x2, n2, c2) )
evaluate(score::CompositeScore, design::TD, x1::TI, x2::TI, p::TR) where {TI<:Integer,TR<:Real,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, x1, x2, p) )
evaluate(score::CompositeScore, design::TD, x1::TI, x2::TI) where {TI<:Integer,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, x1, x2, c2) )
evaluate(score::CompositeScore, x1::TI, n1::TI, n2::TI, c2::TR, p::TR) where {TI<:Integer,TR<:Real} = sum( score.ω .* evaluate.(score.components, x1, n1, n2, c2, p) )
evaluate(score::CompositeScore, x1::TI, n1::TI, n2::TI, c2::TR) where {TI<:Integer,TR<:Real} = sum( score.ω .* evaluate.(score.components, x1, n1, n2, c2) )
evaluate(score::CompositeScore, design::TD, x1::TI, p::TR) where {TI<:Integer,TR<:Real,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, x1, p) )
evaluate(score::CompositeScore, design::TD, x1::TI) where {TI<:Integer,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, x1) )
evaluate(score::CompositeScore, design::TD, p::TR; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TR<:Real,TI<:Integer,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design, p; partial_stage_one = partial_stage_one) )
evaluate(score::CompositeScore, design::TD; partial_stage_one::Tuple{TI,TI} = (0, 0)) where {TI<:Integer,TD<:AbstractDesign} = sum( score.ω .* evaluate.(score.components, design; partial_stage_one = partial_stage_one) )

function integrand_x1(score::CompositeScore, x1::TI, n1::TI, n2::TI, c2::TR; partial_stage_one::Tuple{TI,TI} = (0, 0) ) where {TS<:Score,TI<:Integer,TR<:Real}
    return sum( score.ω .* integrand_x1.(score.components, x1, n1, n2, c2; partial_stage_one = partial_stage_one) )
end


text:
using NiLang
using Test

@testset "auto bcast" begin
    a = AutoBcast([1.0, 2.0, 3.0])
    @instr NEG(a)
    @test a.x == [-1.0,-2.0,-3.0]
    a = AutoBcast([1.0, 2.0, 3.0])
    @instr INC(a)
    @test a.x == [2.0,3.0,4.0]
    @instr DEC(a)
    @test a.x == [1.0,2.0,3.0]

    a = AutoBcast([false, true, true])
    @instr FLIP(a)
    @test a.x == [true, false, false]

    a = AutoBcast([1.0, 2.0, 3.0])
    b = AutoBcast([1.0, 2.0, 4.0])
    @instr a += b
    @test a.x == [2,4,7.0]
    @test b.x == [1,2,4.0]
    @instr SWAP(a, b)
    @test b.x == [2,4,7.0]
    @test a.x == [1,2,4.0]

    a = AutoBcast([1.0, 2.0, 3.0])
    b = 2.0
    @instr a += b
    @test a.x == [3,4,5.0]
    @test b == 2.0

    a = AutoBcast([1.0, 2.0, 3.0])
    b = AutoBcast([1.0, 2.0, 4.0])
    c = AutoBcast([1.0, 2.0, 1.0])
    @instr a += b * c
    @test a.x == [2,6,7.0]
    @test b.x == [1,2,4.0]
    @test c.x == [1,2,1.0]

    a = AutoBcast([1.0, 2.0, 3.0])
    b = 2.0
    c = AutoBcast([1.0, 2.0, 1.0])
    @instr a += b * c
    @test a.x == [3,6,5.0]
    @test b == 2.0
    @test c.x == [1,2,1.0]

    a = AutoBcast([1.0, 2.0, 3.0])
    b = AutoBcast([1.0, 2.0, 4.0])
    c = 3.0
    @instr a += b * c
    @test a.x == [4,8,15.0]
    @test b.x == [1,2,4.0]
    @test c == 3.0

    a = AutoBcast([1.0, 2.0, 3.0])
    b = 2.0
    c = 3.0
    @instr a += b * c
    @test a.x == [7,8,9.0]
    @test b == 2.0
    @test c == 3.0

    @test zero(AutoBcast{Int,3}) == AutoBcast([0, 0, 0])
end


text:
using Xerxes
using Test

@testset "Xerxes.jl" begin
    prog = Xerxes.Program(
        [
            1.0 -1 1 0 0
            2.0 -1 0 1 0
            0    1 0 0 1
        ],
        [1.0, 3, 5],
        [4.0, 3, 0, 0, 0],
        [1, 2],
        [3, 4, 5],
        [0.0, 0, 1, 3, 5],
        [-4.0, -3, 0, 0, 0]
    )

    status = Xerxes.primal_iterate!(prog)
    @test status == Xerxes.Nonoptimal
    @test issetequal(prog.basic_vars, [1, 4, 5])
    @test issetequal(prog.nonbasic_vars, [3, 2])
    @test prog.A[:, sort(prog.basic_vars)] ==
        [ 1 0 0
          2 1 0
          0 0 1 ]
    @test prog.A[:, sort(prog.nonbasic_vars)] ==
        [ -1 1
          -1 0
           1 0 ]
    @test prog.x[sort(prog.basic_vars)] == [1, 1, 5]
    @test prog.z[sort(prog.nonbasic_vars)] == [-7, 4]

    status = Xerxes.primal_iterate!(prog)
    @test status == Xerxes.Nonoptimal
    @test issetequal(prog.basic_vars, [1, 2, 5])
    @test issetequal(prog.nonbasic_vars, [3, 4])
    @test prog.A[:, sort(prog.basic_vars)] ==
        [ 1 -1 0
          2 -1 0
          0  1 1 ]
    @test prog.A[:, sort(prog.nonbasic_vars)] ==
        [ 1 0
          0 1
          0 0 ]
    @test prog.x[sort(prog.basic_vars)] == [2, 1, 4]
    @test prog.z[sort(prog.nonbasic_vars)] == [-10, 7]

    status = Xerxes.primal_iterate!(prog)
    @test status == Xerxes.Nonoptimal
    @test issetequal(prog.basic_vars, [1, 2, 3])
    @test issetequal(prog.nonbasic_vars, [5, 4])
    @test prog.A[:, sort(prog.basic_vars)] ==
        [ 1 -1 1
          2 -1 0
          0  1 0 ]
    @test prog.A[:, sort(prog.nonbasic_vars)] ==
        [ 0 0
          1 0
          0 1 ]
    @test prog.x[sort(prog.basic_vars)] == [4, 5, 2]
    @test prog.z[sort(prog.nonbasic_vars)] == [2, 5]

    status = Xerxes.primal_iterate!(prog)
    @test status == Xerxes.Optimal
end

max_stars_repo_name: dlozeve/Xerxes

max_issues_repo_name: dlozeve/Xerxes

max_forks_repo_name: dlozeve/Xerxes


text:
module PlotDendrite
export plotDendrite
using Compose
const offset = 1 # how many synapse widths correspond to the connector width

function _draw_branch(root; fill_connectors=false, connector_style= fill_connectors ? [] : [stroke("black"), linewidth(1)], style=[stroke("black")], h_frac=0.5, font_style=[stroke("black"), fill("white"), fontsize(12), linewidth(0.1)], kwargs...)
    branch_drawings = []
    branch_root_loc = Tuple{Float64,Float64}[]
    branch_total_lengths = Int64[]
    branch_total_branches = Int64[]
    for branch ∈ root.branches
        branch_drawing, branch_total_length, branch_total_branches_, by0, by1 = _draw_branch(branch; fill_connectors=fill_connectors, connector_style=connector_style, style=style, h_frac=h_frac, font_style=font_style, kwargs...)
        push!(branch_drawings, branch_drawing)
        push!(branch_root_loc, (by0, by1))
        push!(branch_total_lengths, branch_total_length)
        push!(branch_total_branches, branch_total_branches_)
    end
    total_branches = isempty(branch_total_branches) ? 1 : sum(branch_total_branches)
    branch_length = length(root.SC) + offset
    branches_total_length = isempty(branch_total_lengths) ? 0 : maximum(branch_total_lengths)
    total_length = branches_total_length + branch_length
    width  = branch_length/total_length
    height  = (1/total_branches)*h_frac*root.DW
    x0 = 1-width
    y0 = (1-height)*0.5 + 0.5/total_length#0.5-0.5*height# + 0.75*1/(aspect)/total_length
    lw = 1/total_length

    branch_heights = branch_total_branches./total_branches
    branch_ys      = cumsum([0.0;branch_heights[1:end-1]])

    drawing = (
            context(0.01,0.01, 0.98, 0.98),
            # draw stem
            (context(x0,0,width,1), fill(root.DC),
                (context(offset/branch_length, y0, 1-offset/branch_length, height),
                    # draw rectangle representing the stem
                    rectangle(),
                    # write name
                    (context(), text(0.5, 0.5, root.name, hcenter, vcenter), font_style...)
                ),
                # draw connectors
                if fill_connectors # filled connectors in parent context
                    [[
                        polygon([(0,branch_ys[i]+by0*branch_heights[i]),(0,branch_ys[i]+by1*branch_heights[i]),(offset/branch_length,y0+height),(offset/branch_length,y0)]) for (i,(by0,by1)) ∈ enumerate(branch_root_loc)
                    ]; connector_style...]
                else # bezier connectors in their own context
                    [(context(), [
                        curve((0,branch_ys[i]+0.5*(by0+by1)*branch_heights[i]), (0.7*offset/branch_length,branch_ys[i]+0.5*(by0+by1)*branch_heights[i]), (0.3*offset/branch_length,y0+height*i/(length(root.branches)+1)), (offset/branch_length,y0+height*i/(length(root.branches)+1))) for (i,(by0,by1)) ∈ enumerate(branch_root_loc)
                    ]..., connector_style...)]
                end...,

                if isempty(root.branches)
                    # draw end-cap
                    [circle(offset/total_length, y0+0.5*height, 0.5h*height)]
                else
                    []
                end...,

            ),

            # draw sub-branches
            [
                # draw actual sub-branch
                (context(x0-branch_total_lengths[i]/total_length, branch_ys[i], branch_total_lengths[i]/total_length, branch_heights[i]), branch_drawings[i][2:end]...) for i ∈ eachindex(branch_drawings)
            ]...,

            # draw synapses
            [
                (context(), arc(x0+(offset+i-1+0.5)*lw, y0, 0.4*lw, π, 2π), fill(lc)) for (i,lc) ∈ enumerate(root.SC)
            ]...,

            style...
        )
    return drawing, total_length, total_branches, y0, y0+height
end

function plotDendrite(args...;canvas=nothing, caption="", caption_options=[stroke("white"), fill("black"), fontsize(8)], background="white", kwargs...)
    drawing, total_length, total_branches, _,_ = _draw_branch(args...; kwargs...)
    composed_drawing = compose(context(), (context(), drawing, (context(), text(1,1,caption,hright,vbottom), caption_options...), (context(), rectangle(), fill(background))))
    if canvas != nothing
        draw(canvas, composed_drawing)
    end
    return (drawing=composed_drawing, width=total_length, height=total_branches)
end

end


text:
#!/usr/local/bin/julia
#
# This script prepares a data set in current folder.
# For each of chromosome 1-23, there is a set gzipped Plink files:
# hapmap3.chr.1.bed.gz, hapmap3.chr.1.bim.gz, hapmap3.chr.1.fam.gz
# hapmap3.chr.2.bed.gz, hapmap3.chr.2.bim.gz, hapmap3.chr.2.fam.gz
# ...
# hapmap3.chr.23.bed.gz, hapmap3.chr.23.bim.gz, hapmap3.chr.23.fam.gz
# There is also a csv file "covariate.txt" that contains trait and covariates.
#

# install and load Julia packages
using Pkg
if haskey(Pkg.installed(), "SnpArrays")
    Pkg.update("SnpArrays")
else
    Pkg.add(PackageSpec(url="https://github.com/OpenMendel/SnpArrays.jl.git"))
end
if haskey(Pkg.installed(), "OrdinalMultinomialModels")
    Pkg.update("OrdinalMultinomialModels")
else
    Pkg.add(PackageSpec(url="https://github.com/OpenMendel/OrdinalMultinomialModels.jl.git"))
end
if haskey(Pkg.installed(), "OrdinalGWAS")
    Pkg.update("OrdinalGWAS")
else
    Pkg.add(PackageSpec(url="https://github.com/OpenMendel/OrdinalGWAS.jl.git"))
end
using OrdinalMultinomialModels, OrdinalGWAS, SnpArrays

# split hapmap3 data according to chromosome
datadir = normpath(joinpath(dirname(pathof(OrdinalGWAS)), "../data/"))
SnpArrays.split_plink(datadir * "hapmap3", :chromosome; prefix = "hapmap3.chr.")
# compresse Plink files for chromosome 1-23
for chr in 1:23
    plinkfile = "hapmap3.chr." * string(chr)
    SnpArrays.compress_plink(plinkfile)
end
# delete uncompressed chromosome Plink files
for chr in 1:26
    plinkfile = "hapmap3.chr." * string(chr)
    rm(plinkfile * ".bed", force=true)
    rm(plinkfile * ".bim", force=true)
    rm(plinkfile * ".fam", force=true)
end
# copy covariate.txt file
cp(datadir * "covariate.txt", joinpath(pwd(), "covariate.txt"))

max_stars_repo_name: Hua-Zhou/PolrGWAS

max_issues_repo_name: Hua-Zhou/PolrGWAS