module CovidAgent
using Agents
mutable struct GenericAgent <: AbstractAgent
    id::Int
    pos::Tuple{Int,Int} #region
    type::Int # 1 - nsec+1 HH, nsec + 2 - 2*(nsec+1) Firm, man / service / Food/public
    age:: Int # 1 young, 2 old
    cor::Int #infection status: 1 susceptible, 2 infected, 3 recovered, 4 infected with mutant
    vaccinated::Bool # true, if agent has been vaccinated
    cortime:: Int # time of infection
    infect:: Int # number of people infected by this agent
    emp:: Int #employer ID, 0 is unemployed
    shorttime:: Bool # on short time
    shorttimecounter:: Int # Number of weeks on short time
    homeoffice:: Bool # homeoffice
    worker:: Dict{Int,CovidAgent.GenericAgent} # list of firm employees
    workershorttime:: Dict{Int,CovidAgent.GenericAgent} # list of firm employees on short time
    income:: Float64 # average income
    saving::Float64 # savings of HH and firms
    consbud::Array{Float64,1} # consumption budget for HH for three cons types
    prevconsbud::Array{Float64,1} # consumption budget from previous period
    shopday::Array{Int,1} # array of shopping days in the week
    shoprepeat:: Array{Bool,1} # indicator repetition of shopping trip possible
    product::Float64 # productiviy of firm
    stock::Float64 # stock of product of firm
    demand::Float64 #firm's demand in previous week
    demandexp::Float64 # firm’s demand estimate
    rev::Float64 # firm revenues
    pastrev::Array{Float64,1} # array of last 4 weekly revenes
    markup::Float64 # firm markup
    fixedcosts::Float64 # firms fixed costs
end

end

# function for roulette wheel selection based on fitness vector vec
function conselect(vect::Array{Float64,1})
    vecn = size(vect)[1]
    vecsum = sum(vect)
    vecr = rand()*vecsum
    for i=1:vecn
        if sum(vect[1:i]) >= vecr
            return i
            break
        end
    end
end

function getagent(id::Int)
    return covidmodel[id]
end

# one step in the dynamics
function model_step!(covidmodel)
    global hh, firms, lochh, unemplist, shorttimelist, cas, daycounter, noinf, noinfmut, unemp, consumption, weeklyconsumption, shoppers, nsec
    global regcons, regemp, actshoplist, pubacc, empcount, tau,trigger, proftax, totaldiv, totalfixedcosts
    global virus, unempcount, shorttimecount, oldcount, R0count,gdp, genhomeoffice, noinftraj
    global bankruptcypossible, inactivefirms, sumbailouts, totalaccounts, divperhh, totalsav

    global contact_count_traj, contact_work_traj, contact_social_traj, contact_shop_traj
    global contact_count, contact_work, contact_social, contact_shop, contact_num, nosuscep

    global updatefixedcosts, bankruptcount, curinfcount, curinfregcount

    # reset consumption counter
    consumption = zeros(nsec, k1, k2)
    #reset total savings
    totalsav = 0

    #determine mortaliy
    noinf = 0
    noinfmut = 0
    aginf = 0
    for i in union(hh[1],hh[2])
        if (getagent(i).cor ==2) || (getagent(i).cor ==4)
            noinf = noinf +1
        end
        if (getagent(i).cor ==4)
            noinfmut = noinfmut +1
        end
        if getagent(i).cor > 1
            aginf +=1
        end
    end
    push!(totinftraj,aginf)  # number of infected and recoverd
    mort = (min(hcap,icufrac*noinf) /(icufrac*noinf))*mortl  + (1 - min(hcap,icufrac*noinf) / (icufrac*noinf))*morth

    # determine first which of the young hh is travelling
    #lochhytemp = deepcopy(lochh[1,:,:]) # temp list of young at locations
    postemp = [(0,0) for i=1:nhh] # temp list of young positions

    lochhytemp = [Dict{Int,CovidAgent.GenericAgent}() for a = 1:2, i=1:k1, j=1:k2]
    for a=1:2
        for i=1:k1
            for j=1:k2
                for k in keys(lochh[a,i,j])
                    lochhytemp[a,i,j][k] = lochh[a,i,j][k]
                end
            end
        end
    end

    for i in hh[1]  # only young households travel
        agent = getagent(i)
        if rand() < ptravel
            posh = agent.pos
            posn = [rand(1:k1),rand(1:k2)]
            delete!(lochhytemp[posh[1],posh[2]],i)
            lochhytemp[posn[1],posn[2]][i] = agent
            postemp[i] = (posn[1],posn[2]) # save new position
        else
            postemp[i] =agent.pos # stick to default position
        end
    end

    for i in hh[2] # old agents do not travel
        postemp[i] =getagent(i).pos # stick to default position
    end

    # firm decisions once a week
    if daycounter == 0
        totaldiv = 0. # sums dividends of all firms
        totalfixedcosts = 0. # sums fixed costs of all firms
        proftax = 0. # total profit tax
        gdp = 0.
        inactivefirms = 0
        sumbailouts = 0.0
        badbankacc = 0.0
        actshoplist = [[] for tt=1:nsec, i=1:k1, j2 = 1:k2]

        totprice = zeros(nsec)
        pricecount = zeros(nsec)
        permlist = shuffle(firms)
        for id in permlist
            firm = getagent(nhh+id)
            if firm.type-(nsec+1) <= nsec # not a public sector office
                numworker = length(firm.worker)

                # firms collect profits, pay dividends, wages, taxes, update demand estimate
                profit = firm.rev - wage[firm.type-(nsec+1)]*numworker - firm.fixedcosts
                totalfixedcosts += firm.fixedcosts

                # determine dividend
                for revt = 1:3
                    firm.pastrev[revt] = firm.pastrev[revt+1]
                end
                firm.pastrev[4] = firm.rev
                divthr = savtar[firm.type-(nsec+1)]*mean(firm.pastrev) # dividend threhshold
                if firm.saving + (1-tau)*max(0,profit) > divthr
                    div = (1-tau)*max(0,profit)
                else
                    div = divrate*(1-tau)*max(0,profit)
                end
                totaldiv += div  # add to total dividends
                proftax += tau*max(0,profit)
                firm.saving += profit-tau*max(0,profit) - div

                # pricing
                if weeklyconsumption[firm.type-(nsec+1), firm.pos[1], firm.pos[2]] > 0
                    marketshare = firm.demand / weeklyconsumption[firm.type-(nsec+1), firm.pos[1], firm.pos[2]]
                    firm.markup = etmin[firm.type-(nsec+1)] + marketshare*(etmax[firm.type-(nsec+1)]-etmin[firm.type-(nsec+1)])
                else
                    firm.markup = etmin[firm.type-(nsec+1)]
                end

                firm.demandexp = (1 - rhodem)*firm.demandexp + rhodem*firm.demand # update demand expectation
                firm.demand = 0. # reset weekly demand
                firm.rev = 0. # reset weekly revenue

                # production, hiring/firing
                if firm.stock > 0
                    # mormal plan
                    prodplan = max(0,(firm.demandexp*(1 + buff[firm.type-(nsec+1)]) - (1-de[firm.type-(nsec+1)])*firm.stock))
                else
                    # stronger expansion if stock is empty
                    prodplan = max(0,(firm.demandexp*(1 + zerostockboost*buff[firm.type-(nsec+1)]) - (1-de[firm.type-(nsec+1)])*firm.stock))
                end
                labdem = ceil(prodplan / firm.product) - numworker

                if firm.saving < 0 && bailoutprogram
                    sumbailouts += - firm.saving
                    firm.saving = 0
                end

                if firm.saving < 0 && bankruptcypossible
                    prodplan = 0
                    labdem = -numworker
                end

                for (id2, stworker) in firm.workershorttime
                    stworker.shorttimecounter += 1

                    # move worker from short term to unemployed if:
                    # - short time program has ended
                    # - worker excedded maximum time in short time
                    # - firm is bankrupt
                    if (!shorttimeprogram || stworker.shorttimecounter > shorttimeperiod || (firm.saving < 0 && bankruptcypossible))
                        delete!(firm.workershorttime,id2)
                        delete!(shorttimelist[firm.type-(nsec+1),firm.pos[1],firm.pos[2]],id2)
                        unemplist[firm.type-(nsec+1),stworker.pos[1],stworker.pos[2]][id2] = stworker
                        stworker.shorttime = false
                        stworker.shorttimecounter = 0
                        stworker.emp = 0
                    end
                end

                if labdem < 0 # random firing
                    firelist = shuffle(collect(keys(firm.worker)))[1:Int(-labdem)]
                    for id2 in firelist
                        fireagent = getagent(id2)
                        if shorttimeprogram && rand() < shorttimeprob && !(firm.saving < 0 && bankruptcypossible)
                            delete!(firm.worker,id2)
                            firm.workershorttime[id2] = fireagent
                            shorttimelist[firm.type-(nsec+1),fireagent.pos[1],fireagent.pos[2]][id2] = fireagent # put on regional short time list
                            fireagent.shorttime = true
                            fireagent.shorttimecounter = 0
                        else
                            delete!(firm.worker,id2)
                            unemplist[firm.type-(nsec+1),fireagent.pos[1],fireagent.pos[2]][id2] = fireagent # put on regional unemploymentlist
                            fireagent.emp = 0
                        end
                    end
                    numworker += labdem
                elseif labdem > 0 # hiring
                    numhire = min(labdem)
                    numhirefromshorttime = min(numhire, length(firm.workershorttime))
                    numhirefrommarket = min(numhire - numhirefromshorttime, length(unemplist[firm.type-(nsec+1),firm.pos[1],firm.pos[2]])) # check if firm is rationed on labor market

                    if numhirefromshorttime > 0
                        hirelist = shuffle(collect(keys(firm.workershorttime)))[1:Int(numhirefromshorttime)] # search among unemployd from firms short time list
                        for id2 in hirelist
                            hireagent = getagent(id2)
                            firm.worker[id2] = hireagent
                            delete!(firm.workershorttime,id2) # remove from firms short time list
                            delete!(shorttimelist[firm.type-(nsec+1),firm.pos[1],firm.pos[2]],id2) # remove from regional unemploymentlist
                            hireagent.shorttime = false
                            hireagent.shorttimecounter = 0
                        end
                        numworker += numhirefromshorttime
                    end

                    if numhirefrommarket > 0
                        hirelist = shuffle(collect(keys(unemplist[firm.type-(nsec+1),firm.pos[1],firm.pos[2]])))[1:Int(numhirefrommarket)] # search among unemployd in region
                        for id2 in hirelist
                            hireagent = getagent(id2)
                            firm.worker[id2] = hireagent
                            delete!(unemplist[firm.type-(nsec+1),firm.pos[1],firm.pos[2]],id2) # remove from regional unemploymentlist
                            hireagent.emp = id
                        end
                        numworker += numhirefrommarket
                    end
                end

                firm.stock = (1-de[firm.type-(nsec+1)])*firm.stock + numworker*firm.product # update product stock with production
                gdp += (numworker*wage[firm.type-(nsec+1)]+firm.fixedcosts)*(1+firm.markup)/(1 - buff[firm.type-(nsec+1)]*de[firm.type-(nsec+1)])
                if firm.stock > 0
                    push!(actshoplist[firm.type-(nsec+1),firm.pos[1],firm.pos[2]],id) # if sotcvk positive add to actve shops
                end
            else # public sector office
                gdp += length(firm.worker)*wage[nsec+1]*(1+fcvcratio[nsec+1])
            end
            if updatefixedcosts  # update fixed costs
                firm.fixedcosts = fcvcratio[firm.type-(nsec+1)] * wage[firm.type-(nsec+1)]*length(firm.worker) # fixed costs here prop to vc
            else
                if firm.fixedcosts < fcvcratio[firm.type-(nsec+1)] * wage[firm.type-(nsec+1)]*length(firm.worker) # adjust fc only if firm expands
                    firm.fixedcosts = fcvcratio[firm.type-(nsec+1)] * wage[firm.type-(nsec+1)]*length(firm.worker)
                end
            end
            if firm.type < 2*(nsec+1) # temp delete
                totprice[firm.type-(nsec+1)] += (wage[firm.type-(nsec+1)] + firm.fixedcosts / (length(firm.worker)+10^(-5)))*(1+firm.markup)/ (firm.product*(1 - buff[firm.type-(nsec+1)]*de[firm.type-(nsec+1)]))
                pricecount[firm.type-(nsec+1)] += 1
            end
            if firm.saving < 0 && bankruptcypossible
                badbankacc -= firm.saving
                setdiff!(firms,(firm.id-nhh))
                setdiff!(actshoplist[firm.type-(nsec+1),firm.pos[1],firm.pos[2]],(firm.id-nhh))
                setdiff!(locf[firm.type-(nsec+1),firm.pos[1],firm.pos[2]],(firm.id-nhh))
                bankruptcount[firm.type-(nsec+1)] += 1
                kill_agent!(firm,covidmodel)
            end
        end # end firm loop
        #println(totprice ./ pricecount)

        weeklyconsumption = zeros(nsec, k1, k2)

        # government pays unemployment benefits, psensions and public sector, adjusts tax rate
        pubsur = proftax + tau*sum(empcount.*wage) + tau*(sum(empcount)+sum(unempcount)+sum(shorttimecount)+oldcount)*divperhh- (1-tau)*(sum(unempcount.*unempbenefit) + sum(shorttimecount.*shorttimewage) + oldcount*pension) - empcount[nsec+1]*wage[nsec+1]
        pubacc += pubsur

        # Cost of bailouts to the public account
        pubacc -= (sumbailouts + badbankacc)
        # --- End of tax period ---

        tauhat = max(0,(sum(unempcount.*unempbenefit) + sum(shorttimecount.*shorttimewage) + oldcount*pension + empcount[nsec+1]*wage[nsec+1] - debtfrac*pubacc) /
        ((proftax / tau) + sum(empcount.*wage) + sum(unempcount.*unempbenefit) + sum(shorttimecount.*shorttimewage) + oldcount*pension +
        (sum(empcount)+sum(unempcount)+sum(shorttimecount)+oldcount)*divperhh)) #taxrate needed to balance budget

        tau = (1-tauadj)*tau + tauadj*tauhat

        # reset employment, unemployment and shorttime counter
        empcount = zeros(nsec+1)
        unempcount = zeros(nsec)
        shorttimecount = zeros(nsec)
        oldcount = 0

        # calculate div per hh
        divperhh = (totaldiv+totalfixedcosts) / size(union(hh[1],hh[2]))[1]  # equasl share of total dividends for each hh
    end # end daycounter == 0

    #reset consumption counter
    consumption = zeros(nsec,k1,k2)

    #resent local shoppers list
    shoppers = [Dict{Int,CovidAgent.GenericAgent}() for tt=1:nsec, i=1:k1, j=1:k2]

    curinfcount = zeros(2)
    curinfregcount = zeros(2,k1,k2)
    # Household loop
    permlist = shuffle(union(hh[1],hh[2]))
    for ii in permlist
        agent = getagent(ii)
        # potential death of agent

        if (agent.cor == 2) || (agent.cor == 4) # if infected
            if rand() < mort[agent.age]
                setdiff!(hh[agent.age],agent.id)
                delete!(lochh[agent.age,agent.pos[1],agent.pos[2]],agent.id)
                delete!(lochh[agent.age,agent.pos[1],agent.pos[2]],agent.id)
                cas[agent.pos[1],agent.pos[2]][agent.age] += 1
                push!(R0count,agent.infect)
                if agent.age == 1
                    delete!(lochhytemp[postemp[agent.id][1],postemp[agent.id][2]],agent.id)
                    if agent.emp ==0
                        delete!(unemplist[agent.type,agent.pos[1],agent.pos[2]],agent.id)
                    else
                        employer = getagent(nhh+agent.emp)
                        if agent.shorttime
                            delete!(employer.workershorttime,agent.id)
                            delete!(shorttimelist[agent.type,agent.pos[1],agent.pos[2]],agent.id)
                        else
                            delete!(employer.worker,agent.id)
                        end
                    end
                end
                # transfer savings of HH to random other HH
                permlisttemp = shuffle(union(hh[1],hh[2]))
                heiragent = getagent(permlisttemp[1])
                heiragent.saving += agent.saving
                kill_agent!(agent,covidmodel)
                continue
            end
            agent.cortime = agent.cortime-1 # if not dying reduce recovery time by one
            if agent.cortime ==0
                agent.cor = 3
                push!(R0count,agent.infect)
            end
        end

    # hh determine consumption budgets and count unemp once a week
        if daycounter ==0
            agent.saving += sum(agent.consbud)   # non-spent consumption budget is added to savings
            actualincome = 0
            if agent.emp > 0 # employed

                if agent.shorttime
                    shorttimecount[agent.type] += 1
                    actualincome = (1-tau)*(shorttimewage[agent.type] + divperhh)
                    agent.income = (1-rhoinc)*agent.income +  rhoinc*(1-tau)*(shorttimewage[agent.type] + divperhh) # update income expectations
                else
                    empcount[agent.type] += 1
                    actualincome = (1-tau)*(wage[agent.type] + divperhh)
                    agent.income = (1-rhoinc)*agent.income +  rhoinc*(1-tau)*(wage[agent.type] + divperhh) # update income expectations

                end
            elseif agent.age==1 # unemployed young
                unempcount[agent.type] += 1
                agent.income = (1-rhoinc)*agent.income +  rhoinc*(1-tau)*(unempbenefit[agent.type] + divperhh)
                actualincome = (1-tau)*(unempbenefit[agent.type] + divperhh)
            else # old agent
                oldcount += 1
                agent.income = (1-rhoinc)*agent.income +  rhoinc*(1-tau)*(pension +  divperhh)
                actualincome = (1-tau)*(pension + divperhh)
            end

            totcons = max(0,agent.income + kap*(agent.saving - phi*agent.income))
            agent.consbud = fraccons*totcons

            if essentialsec > 0
                if agent.consbud[essentialsec] < (1-essentialadj) * agent.prevconsbud[essentialsec]
                    if totcons < (1-essentialadj) * agent.prevconsbud[essentialsec]
                        for i = 1:size(agent.consbud, 1)
                            agent.consbud[i] = 0
                        end
                        agent.consbud[essentialsec] = totcons
                    else
                        remainingcons = totcons - (1-essentialadj) * agent.prevconsbud[essentialsec]
                        tempfraccons = deepcopy(fraccons)
                        tempfraccons[essentialsec] = 0
                        sumfrac = sum(tempfraccons)
                        for i = 1:size(tempfraccons, 1)
                            tempfraccons[i] = tempfraccons[i] / sumfrac
                        end
                        agent.consbud = tempfraccons*remainingcons
                        agent.consbud[essentialsec] = (1-essentialadj) * agent.prevconsbud[essentialsec]
                    end
                end
            end

            agent.prevconsbud = deepcopy(agent.consbud)
            agent.saving += actualincome - totcons
            agent.shopday = rand(0:6,nsec) # determine shopping days for three sectors
            agent.shoprepeat = [true, true, true] # consumer repeats shopping once if rationed
            # calculate total savings
            totalsav += agent.saving
        end # end daycounter = 0

        for tt = 1:nsec
            if agent.shopday[tt] == daycounter && rand() < pshop[agent.age][tt] # agent goes shopping
                numshoph = min(size(actshoplist[tt,postemp[agent.id][1],postemp[agent.id][2]])[1],numshops[tt])
                shoppers[tt,postemp[agent.id][1],postemp[agent.id][2]][agent.id] = agent # if agent shops in region where she currently is
                if numshoph > 0 # only if there are potential suppliers
                    storelist = shuffle(actshoplist[tt,postemp[agent.id][1],postemp[agent.id][2]])[1:numshoph] # list of supplier agent observes
                    pricelist = Float64[]
                    for storeid in storelist
                        store = getagent(nhh + storeid)
                        push!(pricelist, exp( - gac*log((wage[store.type-(nsec+1)] + store.fixedcosts / (length(store.worker)+10^(-5)))*(1+store.markup)/ (store.product*(1 - buff[store.type-(nsec+1)]*de[store.type-(nsec+1)]) ))))
                    end
                    store = getagent(nhh + storelist[conselect(pricelist)]) # selected supplier
                    price = (wage[store.type-(nsec+1)] + store.fixedcosts / (length(store.worker)+10^(-5)))*(1+store.markup)/ (store.product*(1 - buff[store.type-(nsec+1)]*de[store.type-(nsec+1)]))
                    quant = min(agent.consbud[tt] / price, store.stock) # minimum of demand and stock
                    agent.consbud[tt] -= quant*price # update agent's consumer budget
                    store.stock -= quant # update stock of supplier
                    if store.stock < 0.0000001
                        setdiff!(actshoplist[store.type-(nsec+1),store.pos[1],store.pos[2]],(store.id-nhh)) # if stock is empty make store inactive
                    end
                    consumption[tt,postemp[agent.id][1],postemp[agent.id][2]] += quant # update total consumption
                    weeklyconsumption[tt,postemp[agent.id][1],postemp[agent.id][2]] += quant
                    store.rev += quant*price # update firm revenues
                    store.demand += quant # update demand
                end
                if agent.consbud[tt] > 0.000001 && agent.shoprepeat[tt]
                    agent.shopday[tt] += 1 # try again next day if rationed
                    agent.shoprepeat[tt] = false # only one repetition possible
                end
            end
        end
    end # end HH loop

    totalaccounts = pubacc
    for i in firms
        totalaccounts += getagent(nhh+i).saving
    end
    for i in union(hh[1], hh[2])
        totalaccounts += getagent(i).saving
    end
    # calculate unemployment
    unemp = 1 - (sum(empcount)+sum(shorttimecount))/size(hh[1])[1]

    #update total consumption
    push!(regcons, consumption)

    # start counting contatcs
    contact_count = 0
    contact_work = 0
    contact_social = 0
    contact_shop = 0
    contact_num = 0
    nosuscep = 0 # count susceptible

### infection loop from perspective of infected
    if virus
        for ii in union(hh[1],hh[2])
            agent = getagent(ii)
            if (agent.cor == 2) || (agent.cor == 4)
                global curinfcount[agent.age] += 1
                global curinfregcount[agent.age,agent.pos[1], agent.pos[2]] += 1
            end
            # possible infection
            if ((agent.cor==2)|| (agent.cor == 4))  && (agent.cortime < trec - corlatent) && (agent.cortime > trec - corlatent- corinf)
                contact_num += 1
                # first infection at the workplace
                if agent.emp > 0 && !agent.shorttime
                    if !agent.homeoffice || !genhomeoffice || rand() > phomeoffice # no home-office for agent
                        employer = getagent(nhh+agent.emp)
                        idlist = employer.worker
                        nummeet = rand(0:workmeet[employer.type-(nsec+1)])
                        for i in shuffle(collect(keys(idlist)))[1:min(nummeet,length(idlist))]
                            partner = getagent(i)
                            contact_work = contact_work + 1 # add contatcs
                            if !partner.vaccinated && (rand() < pinf[Int(agent.cor/2)]) && partner.cor == 1 && !(partner.homeoffice && genhomeoffice && rand() < phomeoffice) && !partner.shorttime
                                partner.cor=agent.cor
                                partner.cortime = trec
                                agent.infect +=1
                            end
                        end
                    end
                end

                # second infection through consumption
                for tt=1:nsec
                    if length(shoppers[tt,postemp[agent.id][1],postemp[agent.id][2]]) > 2
                        if sum(collect(keys(shoppers[tt,postemp[agent.id][1],postemp[agent.id][2]])) .== agent.id) > 0 # agent shopped in sector tt at that day
                            nummeet = rand(0:shopmeet[tt])
                            shopmeeth = Int(ceil(nummeet * length(shoppers[tt,postemp[agent.id][1],postemp[agent.id][2]]) * 7 / (length(lochh[1,postemp[agent.id][1],postemp[agent.id][2]]) + length(lochh[2,postemp[agent.id][1],postemp[agent.id][2]])))) # number of contacts when shopping, equal to nummeet if 1/7 of local pop is shopping
                            shopcontact = rand(collect(keys(delete!(shoppers[tt,postemp[agent.id][1],postemp[agent.id][2]],agent.id))),shopmeeth) #random selection from list of shoppers in same region
                            for i in shopcontact
                                partner = getagent(i)
                                contact_shop = contact_shop + 1 # add contacts
                                if !partner.vaccinated && (rand() < pinf[Int(agent.cor/2)]) &&  (partner.cor ==1)
                                    partner.cor=agent.cor
                                    partner.cortime = trec
                                    agent.infect +=1
                                end
                            end
                            shoppers[tt,postemp[agent.id][1],postemp[agent.id][2]][agent.id] = agent
                        end
                    end
                end

                # third infection through other social contacts
                if agent.age ==1 # young agent
                        # meet young
                        helplist = collect(keys(delete!(lochhytemp[postemp[agent.id][1],postemp[agent.id][2]],agent.id))) # list of potential young contacts in the region
                        if socialmaxyy >= 1
                            numcontact = rand(0:socialmaxyy) # number of contacts
                        else
                            if rand() < mean([0,socialmaxyy])
                                numcontact = 1
                            else
                                numcontact = 0
                            end
                        end
                        if numcontact > 0
                            helpcontact = shuffle(helplist)[1:min(numcontact,size(helplist)[1])] #list of contacts
                            for i in helpcontact
                                partner = getagent(i)
                                contact_social = contact_social + 1 # add contacts
                                if !partner.vaccinated && (rand() < pinf[Int(agent.cor/2)]) &&  (partner.cor ==1)
                                    partner.cor=agent.cor
                                    partner.cortime = trec
                                    agent.infect +=1
                                end
                            end
                        end
                        lochhytemp[postemp[agent.id][1],postemp[agent.id][2]][agent.id] = agent

                        #meet old
                        helplist = collect(keys(lochh[2,postemp[agent.id][1],postemp[agent.id][2]])) # list of potential old in the region
                        if socialmaxyo >= 1
                            numcontact = rand(0:socialmaxyo) # number of contacts
                        else
                            if rand() < mean([0,socialmaxyo])
                                numcontact = 1
                            else
                                numcontact = 0
                            end
                        end
                        if numcontact > 0
                            helpcontact = shuffle(helplist)[1:min(numcontact,size(helplist)[1])] #list of contacts
                            for i in helpcontact
                                partner = getagent(i)
                                contact_social = contact_social + 1 # add contacts
                                if !partner.vaccinated && (rand() < pinf[Int(agent.cor/2)]) &&  (partner.cor ==1)
                                    partner.cor=agent.cor
                                    partner.cortime = trec
                                    agent.infect +=1
                                end
                            end
                        end
                else # old agent
                    # meet young
                    helplist = collect(keys(lochhytemp[agent.pos[1],agent.pos[2]])) # list of potential young contacts in the region
                    if socialmaxoy >= 1
                        numcontact = rand(0:socialmaxoy) # number of contacts
                    else
                        if rand() < mean([0,socialmaxoy])
                            numcontact = 1
                        else
                            numcontact = 0
                        end
                    end
                    if numcontact > 0
                        helpcontact = shuffle(helplist)[1:min(numcontact,size(helplist)[1])] #list of contacts
                        for i in helpcontact
                            partner = getagent(i)
                            contact_social = contact_social + 1 # add contacts
                            if !partner.vaccinated && (rand() < pinf[Int(agent.cor/2)]) &&  (partner.cor ==1)
                                partner.cor=agent.cor
                                partner.cortime = trec
                                agent.infect +=1
                            end
                        end
                    end
                    #meet old
                    helplist = collect(keys(delete!(lochh[2,agent.pos[1],agent.pos[2]],agent.id))) # list of potential old in the region
                    if socialmaxoo >= 1
                        numcontact = rand(0:socialmaxoo) # number of contacts
                    else
                        if rand() < mean([0,socialmaxoo])
                            numcontact = 1
                        else
                            numcontact = 0
                        end
                    end
                    if numcontact > 0
                        helpcontact = shuffle(helplist)[1:min(numcontact,size(helplist)[1])] #list of contacts
                        for i in helpcontact
                            partner = getagent(i)
                            contact_social = contact_social + 1 # add contacts
                            if !partner.vaccinated && (rand() < pinf[Int(agent.cor/2)]) &&  (partner.cor ==1)
                                partner.cor=agent.cor
                                partner.cortime = trec
                                agent.infect +=1
                            end
                        end
                    end
                    lochh[2,agent.pos[1],agent.pos[2]][agent.id] = agent
                end
            end # if cor = 1

        end # summing over all young hh

        # collect contact data
        contact_count = contact_work + contact_social + contact_shop
        push!(contact_count_traj,contact_count/(contact_num + 10^(-5)))
        push!(contact_work_traj,contact_work/(contact_num + 10^(-5)))
        push!(contact_social_traj,contact_social/(contact_num + 10^(-5)))
        push!(contact_shop_traj,contact_shop/(contact_num + 10^(-5)))

    end # virus from infected perspective

    if daycounter == 0
        daycounter = 7
    end
    # reduce daycounter
    daycounter -=1
    return
end


using Gadfly
p = plot(y = [1, 2, 3, 4, 5])
img = SVG("plot3.svg", 8cm, 8cm)
draw(img, p)
println("Julia 👩🦰 sagt: I did it, Babe!")


"""
#### Read the particle-data
- select variables
- limit to a maximum range
- print the name of each data-file before reading it
- toggle verbose mode

```julia
getparticles(       dataobject::InfoType;
                    lmax::Real=dataobject.levelmax,
                    vars::Array{Symbol,1}=[:all],
                    stars::Bool=true,
                    xrange::Array{<:Any,1}=[missing, missing],
                    yrange::Array{<:Any,1}=[missing, missing],
                    zrange::Array{<:Any,1}=[missing, missing],
                    center::Array{<:Any,1}=[0., 0., 0.],
                    range_unit::Symbol=:standard,
                    print_filenames::Bool=false,
                    verbose::Bool=verbose_mode)
```
#### Returns an object of type PartDataType, containing the particle-data table, the selected and the simulation ScaleType and summary of the InfoType
```julia
return PartDataType()

# get an overview of the returned fields:
# e.g.:
julia> info = getinfo(100)
julia> particles  = getparticles(info)
julia> viewfields(particles)
#or:
julia> fieldnames(particles)
```

#### Arguments
##### Required:
- **`dataobject`:** needs to be of type: "InfoType", created by the function *getinfo*
##### Predefined/Optional Keywords:
- **`lmax`:** not defined
- **`stars`:** not defined
- **`var(s)`:** the selected particle variables in arbitrary order: :all (default), :cpu, :mass, :vx, :vy, :vz, :birth :metals, ...
- **`xrange`:** the range between [xmin, xmax] in units given by argument `range_unit` and relative to the given `center`; zero length for xmin=xmax=0. is converted to maximum possible length
- **`yrange`:** the range between [ymin, ymax] in units given by argument `range_unit` and relative to the given `center`; zero length for ymin=ymax=0. is converted to maximum possible length
- **`zrange`:** the range between [zmin, zmax] in units given by argument `range_unit` and relative to the given `center`; zero length for zmin=zmax=0. is converted to maximum possible length
- **`range_unit`:** the units of the given ranges: :standard (code units), :Mpc, :kpc, :pc, :mpc, :ly, :au , :km, :cm (of typye Symbol) ..etc. ; see for defined length-scales viewfields(info.scale)
- **`center`:** in units given by argument `range_unit`; by default [0., 0., 0.]; the box-center can be selected by e.g. [:bc], [:boxcenter], [value, :bc, :bc], etc..
- **`print_filenames`:** print on screen the current processed particle file of each CPU
- **`verbose`:** print timestamp, selected vars and ranges on screen; default: set by the variable `verbose_mode`

### Defined Methods - function defined for different arguments
- getparticles( dataobject::InfoType; ...) # no given variables -> all variables loaded
- getparticles( dataobject::InfoType, var::Symbol; ...) # one given variable -> no array needed
- getparticles( dataobject::InfoType, vars::Array{Symbol,1}; ...)  # several given variables -> array needed

#### Examples
```julia
# read simulation information
julia> info = getinfo(420)

# Example 1:
# read particle data of all variables, full-box, all levels
julia> particles = getparticles(info)

# Example 2:
# read particle data of all variables
# data range 20x20x4 kpc; ranges are given in kpc relative to the box (here: 48 kpc) center at 24 kpc
julia> particles = getparticles( info,
                                  xrange=[-10., 10.],
                                  yrange=[-10., 10.],
                                  zrange=[-2., 2.],
                                  center=[24., 24., 24.],
                                  range_unit=:kpc )

# Example 3:
# give the center of the box by simply passing: center = [:bc] or center = [:boxcenter]
# this is equivalent to center=[24.,24.,24.] in Example 2
# the following combination is also possible: e.g. center=[:bc, 12., 34.], etc.
julia> particles = getparticles(    info,
                                    xrange=[-10.,10.],
                                    yrange=[-10.,10.],
                                    zrange=[-2.,2.],
                                    center=[33., bc:, 10.],
                                    range_unit=:kpc )

# Example 4:
# read particle data of the variables mass and the birth-time, full-box, all levels
julia> particles = getparticles( info, [:mass, :birth] ) # use array for the variables

# Example 5:
# read particle data of the single variable mass, full-box, all levels
julia> particles = getparticles( info, :mass ) # no array for a single variable needed
...
```

"""
function getparticles( dataobject::InfoType, var::Symbol;
                    lmax::Real=dataobject.levelmax,
                    stars::Bool=true,
                    xrange::Array{<:Any,1}=[missing, missing],
                    yrange::Array{<:Any,1}=[missing, missing],
                    zrange::Array{<:Any,1}=[missing, missing],
                    center::Array{<:Any,1}=[0., 0., 0.],
                    range_unit::Symbol=:standard,
                    print_filenames::Bool=false,
                    verbose::Bool=verbose_mode)

    return  getparticles( dataobject, vars=[var],
                        lmax=lmax,
                        stars=stars,
                        xrange=xrange,
                        yrange=yrange,
                        zrange=zrange,
                        center=center,
                        range_unit=range_unit,
                        print_filenames=print_filenames,
                        verbose=verbose)
end

function getparticles( dataobject::InfoType, vars::Array{Symbol,1};
                    lmax::Real=dataobject.levelmax,
                    stars::Bool=true,
                    xrange::Array{<:Any,1}=[missing, missing],
                    yrange::Array{<:Any,1}=[missing, missing],
                    zrange::Array{<:Any,1}=[missing, missing],
                    center::Array{<:Any,1}=[0., 0., 0.],
                    range_unit::Symbol=:standard,
                    print_filenames::Bool=false,
                    verbose::Bool=verbose_mode)

    return  getparticles( dataobject, vars=vars,
                                        lmax=lmax,
                                        stars=stars,
                                        xrange=xrange,
                                        yrange=yrange,
                                        zrange=zrange,
                                        center=center,
                                        range_unit=range_unit,
                                        print_filenames=print_filenames,
                                        verbose=verbose)
end

function getparticles( dataobject::InfoType;
                    lmax::Real=dataobject.levelmax,
                    vars::Array{Symbol,1}=[:all],
                    stars::Bool=true,
                    xrange::Array{<:Any,1}=[missing, missing],
                    yrange::Array{<:Any,1}=[missing, missing],
                    zrange::Array{<:Any,1}=[missing, missing],
                    center::Array{<:Any,1}=[0., 0., 0.],
                    range_unit::Symbol=:standard,
                    print_filenames::Bool=false,
                    verbose::Bool=verbose_mode)

    printtime("Get particle data: ", verbose)
    checkfortype(dataobject, :particles)

    #Todo: limit to a given lmax
    lmax=dataobject.levelmax # overwrite given lmax
    #checklevelmax(dataobject, lmax)
    isamr = checkuniformgrid(dataobject, lmax)
    #time = dataobject.time

    # create variabe-list and vector-mask (nvarh_corr) for getparticledata-function
    # print selected variables on screen
    nvarp_list, nvarp_i_list, nvarp_corr, read_cpu, used_descriptors = prepvariablelist(dataobject, :particles, vars, lmax, verbose)

    # convert given ranges and print overview on screen
    ranges = prepranges(dataobject, range_unit, verbose, xrange, yrange, zrange, center)

    # read particle-data of the selected variables
    if read_cpu
        if dataobject.descriptor.pversion == 0
            pos_1D, vars_1D, cpus_1D, identity_1D, levels_1D = getparticledata(  dataobject, length(nvarp_list), nvarp_corr, stars, lmax, ranges,
                                         print_filenames, verbose, read_cpu)
        elseif dataobject.descriptor.pversion > 0
            pos_1D, vars_1D, cpus_1D, identity_1D, family_1D, tag_1D, levels_1D = getparticledata(  dataobject, length(nvarp_list), nvarp_corr, stars, lmax, ranges,
                                         print_filenames, verbose, read_cpu)
        end
    else
        if dataobject.descriptor.pversion == 0
            pos_1D, vars_1D, identity_1D, levels_1D = getparticledata(  dataobject, length(nvarp_list), nvarp_corr, stars, lmax, ranges,
                                         print_filenames, verbose, read_cpu)
        elseif dataobject.descriptor.pversion > 0
            pos_1D, vars_1D, identity_1D, family_1D, tag_1D, levels_1D = getparticledata(  dataobject, length(nvarp_list), nvarp_corr, stars, lmax, ranges,
                                         print_filenames, verbose, read_cpu)
        end
    end

    if verbose
        @printf "Found %e particles\n" size(pos_1D)[2]
    end

    # prepare column names for the data table
    names_constr = preptablenames_particles(dataobject, dataobject.nvarp, nvarp_list, used_descriptors, read_cpu, lmax, dataobject.levelmin)

    if lmax != dataobject.levelmin # if AMR
        if dataobject.descriptor.pversion == 0
            Nkeys = [:level, :x, :y, :z, :id]
        elseif dataobject.descriptor.pversion > 0
            Nkeys = [:level, :x, :y, :z, :id, :family, :tag]
        end
    else # if uniform grid
        if dataobject.descriptor.pversion == 0
            Nkeys = [:x, :y, :z, :id]
        elseif dataobject.descriptor.pversion > 0
            Nkeys = [:x, :y, :z, :id, :family, :tag]
        end
    end

    # create data table
    # decouple pos_1D/vars_1D from ElasticArray with ElasticArray.data
    if read_cpu # read also cpu number related to particle
        if isamr
            if dataobject.descriptor.pversion == 0
                data = table( levels_1D[:],
                    pos_1D[1,:].data, pos_1D[2,:].data, pos_1D[3,:].data, identity_1D[:], cpus_1D[:],
                    [vars_1D[ nvarp_corr[i],: ].data for i in nvarp_i_list]...,
                    names=collect(names_constr), pkey=collect(Nkeys), presorted = false )
            elseif dataobject.descriptor.pversion > 0
                filter!(x->x≠6,nvarp_i_list)
                filter!(x->x≠5,nvarp_i_list)
                data = table( levels_1D[:],
                    pos_1D[1,:].data, pos_1D[2,:].data, pos_1D[3,:].data, identity_1D[:], family_1D[:], tag_1D[:], cpus_1D[:],
                    [vars_1D[ nvarp_corr[i],: ].data for i in nvarp_i_list]...,
                    names=collect(names_constr), pkey=collect(Nkeys), presorted = false )
            end

        else # if uniform grid
            if dataobject.descriptor.pversion == 0
                data = table(pos_1D[1,:].data, pos_1D[2,:].data, pos_1D[3,:].data, identity_1D[:], cpus_1D[:],
                    [vars_1D[ nvarp_corr[i],: ].data for i in nvarp_i_list]...,
                    names=collect(names_constr), pkey=collect(Nkeys), presorted = false )
            elseif dataobject.descriptor.pversion > 0
                filter!(x->x≠6,nvarp_i_list)
                filter!(x->x≠5,nvarp_i_list)
                data = table(pos_1D[1,:].data, pos_1D[2,:].data, pos_1D[3,:].data, identity_1D[:], family_1D[:], tag_1D[:], cpus_1D[:],
                    [vars_1D[ nvarp_corr[i],: ].data for i in nvarp_i_list]...,
                    names=collect(names_constr), pkey=collect(Nkeys), presorted = false )
            end
        end
    else
        if isamr
            if dataobject.descriptor.pversion == 0
                data = table( levels_1D[:],
                    pos_1D[1,:].data, pos_1D[2,:].data, pos_1D[3,:].data, identity_1D[:],
                    [vars_1D[ nvarp_corr[i],: ].data for i in nvarp_i_list]...,
                    names=collect(names_constr), pkey=collect(Nkeys), presorted = false )
            elseif dataobject.descriptor.pversion > 0
                filter!(x->x≠6,nvarp_i_list)
                filter!(x->x≠5,nvarp_i_list)
                data = table( levels_1D[:],
                    pos_1D[1,:].data, pos_1D[2,:].data, pos_1D[3,:].data, identity_1D[:], family_1D[:], tag_1D[:],
                    [vars_1D[ nvarp_corr[i],: ].data for i in nvarp_i_list]...,
                    names=collect(names_constr), pkey=collect(Nkeys), presorted = false )
            end
        else # if uniform grid
            if dataobject.descriptor.pversion == 0
                data = table(pos_1D[1,:].data, pos_1D[2,:].data, pos_1D[3,:].data, identity_1D[:],
                    [vars_1D[ nvarp_corr[i],: ].data for i in nvarp_i_list]...,
                    names=collect(names_constr), pkey=collect(Nkeys), presorted = false )
            elseif dataobject.descriptor.pversion > 0
                filter!(x->x≠6,nvarp_i_list)
                filter!(x->x≠5,nvarp_i_list)
                data = table(pos_1D[1,:].data, pos_1D[2,:].data, pos_1D[3,:].data, identity_1D[:], family_1D[:], tag_1D[:],
                    [vars_1D[ nvarp_corr[i],: ].data for i in nvarp_i_list]...,
                    names=collect(names_constr), pkey=collect(Nkeys), presorted = false )
            end

        end
    end

    printtablememory(data, verbose)

    partdata = PartDataType()
    partdata.data = data
    partdata.info = dataobject
    partdata.lmin = dataobject.levelmin
    partdata.lmax = lmax
    partdata.boxlen = dataobject.boxlen
    partdata.ranges = ranges
    partdata.selected_partvars = names_constr
    partdata.used_descriptors = used_descriptors
    partdata.scale = dataobject.scale
    return partdata

end

function preptablenames_particles(dataobject::InfoType, nvarp::Int, nvarp_list::Array{Int, 1}, used_descriptors::Dict{Any,Any}, read_cpu::Bool, lmax::Real, levelmin::Real)

    if read_cpu
        if lmax != levelmin # if AMR
            if dataobject.descriptor.pversion == 0
                names_constr = [:level, :x, :y, :z, :id, :cpu]
            elseif dataobject.descriptor.pversion > 0
                names_constr = [:level, :x, :y, :z, :id, :family, :tag, :cpu]
            end
        else # if uniform grid
            if dataobject.descriptor.pversion == 0
                names_constr = [:x, :y, :z, :id, :cpu]
            elseif dataobject.descriptor.pversion > 0
                names_constr = [:x, :y, :z, :id, :family, :tag, :cpu]
            end
        end
                    #, Symbol("x"), Symbol("y"), Symbol("z")
    else
        if lmax != levelmin # if AMR
            if dataobject.descriptor.pversion == 0
                names_constr = [:level, :x, :y, :z, :id]
            elseif dataobject.descriptor.pversion > 0
                names_constr = [:level, :x, :y, :z, :id, :family, :tag]
            end
        else # if uniform grid
            if dataobject.descriptor.pversion == 0
                names_constr = [:x, :y, :z, :id]
            elseif dataobject.descriptor.pversion > 0
                names_constr = [:x, :y, :z, :id, :family, :tag]
            end
        end
    end

    for i=1:nvarp
        if in(i, nvarp_list)
            #if  length(used_descriptors) == 0 || !haskey(used_descriptors, i)
            if dataobject.descriptor.pversion == 0
                if i == 1
                    append!(names_constr, [Symbol("vx")] )
                elseif i == 2
                    append!(names_constr, [Symbol("vy")] )
                elseif i == 3
                    append!(names_constr, [Symbol("vz")] )
                elseif i == 4
                    append!(names_constr, [Symbol("mass")] )
                elseif i == 5
                    append!(names_constr, [Symbol("birth")] )
                elseif i > 5
                    append!(names_constr, [Symbol("var$i")] )
                end
            elseif dataobject.descriptor.pversion > 0
                if i == 1
                    append!(names_constr, [Symbol("vx")] )
                elseif i == 2
                    append!(names_constr, [Symbol("vy")] )
                elseif i == 3
                    append!(names_constr, [Symbol("vz")] )
                elseif i == 4
                    append!(names_constr, [Symbol("mass")] )
                elseif i == 7
                    append!(names_constr, [Symbol("birth")] )
                elseif i == 8
                    append!(names_constr, [Symbol("metals")] )
                elseif i > 8
                    append!(names_constr, [Symbol("var$i")] )
                end
            end
            #else append!(names_constr, [used_descriptors[i]] )

            #end
        end
    end

    return names_constr
end


module IntrinsicPersistence

using Distances
using LightGraphs
using NearestNeighbors
using RecipesBase
using SimpleWeightedGraphs
using StaticArrays

using Random
using SparseArrays

include("geodesiccomplex.jl")
include("equilateraliterator.jl")
include("persistence.jl")
include("plotting.jl")

export persistence, GeodesicComplex, diameter, perimeter, points

end


abstract type BitMask{T<:Unsigned} end

Base.broadcastable(x::BitMask) = Ref(x)

function generate_bitmask_flags(type, decl)
    identifier, value = decl.args
    :(const $identifier = $type($value))
end

"""
    @bitmask_flag BitFlags::UInt32 begin
        FLAG_A = 1
        FLAG_B = 2
        FLAG_C = 4
    end

Enumeration of bitmask flags that can be combined with `&`, `|` and `xor`, forbidding the combination of flags from different bitmasks.
"""
macro bitmask_flag(typedecl, expr)
    if Base.is_expr(typedecl, :(::), 2)
        type, eltype = typedecl.args
    else
        error("First argument to @bitmask_flag must be of the form 'type::eltype'")
    end
    decls = filter(x -> typeof(x) ≠ LineNumberNode, expr.args)

    exs = [
        :(
            Base.@__doc__ struct $type <: BitMask{$eltype}
                val::$eltype
            end
        );
        generate_bitmask_flags.(type, decls)
    ]

    Expr(:block, esc.(exs)...)
end

(&)(a::BitMask, b::BitMask) = error("Bitwise operation not allowed between incompatible bitmasks '$(typeof(a))', '$(typeof(b))'")
(|)(a::BitMask, b::BitMask) = error("Bitwise operation not allowed between incompatible bitmasks '$(typeof(a))', '$(typeof(b))'")
xor(a::BitMask, b::BitMask) = error("Bitwise operation not allowed between incompatible bitmasks '$(typeof(a))', '$(typeof(b))'")
isless(a::BitMask, b::BitMask) = error("Bitwise operation not allowed between incompatible bitmasks '$(typeof(a))', '$(typeof(b))'")
(==)(a::BitMask, b::BitMask) = error("Operation not allowed between incompatible bitmasks '$(typeof(a))', '$(typeof(b))'")
in(a::BitMask, b::BitMask) = error("Operation not allowed between incompatible bitmasks '$(typeof(a))', '$(typeof(b))'")

(&)(a::T, b::T) where {T <: BitMask} = T(a.val & b.val)
(|)(a::T, b::T) where {T <: BitMask} = T(a.val | b.val)
xor(a::T, b::T) where {T <: BitMask} = T(xor(a.val, b.val))
isless(a::T, b::T) where {T <: BitMask} = isless(a.val, b.val)
(==)(a::T, b::T) where {T <: BitMask} = a.val == b.val
in(a::T, b::T) where {T <: BitMask} = a & b == a

(&)(a::T, b::Integer) where {T <: BitMask} = T(a.val & b)
(|)(a::T, b::Integer) where {T <: BitMask} = T(a.val | b)
xor(a::T, b::Integer) where {T <: BitMask} = T(xor(a.val, b))
isless(a::T, b::Integer) where {T <: BitMask} = isless(a.val, b)
in(a::T, b::Integer) where {T <: BitMask} = a & b == a

(&)(a::Integer, b::T) where {T <: BitMask} = b & a
(|)(a::Integer, b::T) where {T <: BitMask} = b | a
xor(a::Integer, b::T) where {T <: BitMask} = xor(b, a)
isless(a::Integer, b::T) where {T <: BitMask} = isless(a, b.val) # need b.val to prevent stackoverflow
in(a::Integer, b::T) where {T <: BitMask} = a | b == b

(::Type{T})(bm::BitMask) where {T <: Integer} = T(bm.val)

convert(T::Type{<:Integer}, bm::BitMask) = T(bm.val)
convert(T::Type{<:BitMask}, val::Integer) = T(val)

typemax(T::Type{<:BitMask{_T}}) where {_T} = T(typemax(_T))


using Arrhenius
using PyCall

```
1. config Conda.jl
2. add PyCall
3. build PyCall
```
ct = pyimport("cantera")

gas = ct.Solution("gri30")

gas.X


# unit test for the options factory
using COSMO, Test, LinearAlgebra, SparseArrays, Pardiso

@testset "Options Factory" begin

P = spzeros(1, 1)
q = [1]
c1 = COSMO.Constraint([1.], [-1.], COSMO.Nonnegatives)

iparms = Dict{Int64, Int64}()
iparms[13] = 100
num_threads = 2

settings = COSMO.Settings(kkt_solver = with_options(COSMO.MKLPardisoKKTSolver, iparm = iparms, num_threads = num_threads, msg_level_on = false))
model = COSMO.Model()
assemble!(model, P, q, c1, settings = settings)
model.ρvec = settings.rho * ones(size(model.p.A, 1))
COSMO._make_kkt_solver!(model)

  @test typeof(model.kkt_solver) <: COSMO.MKLPardisoKKTSolver
  # @test Pardiso.get_nprocs(model.kkt_solver.ps) == num_threads
  @test Pardiso.get_iparm(model.kkt_solver.ps, 13) == iparms[13]

  ## check that this works via the MathOptInterface wrapper as well
  optimizer = COSMO.Optimizer(time_limit = 10., kkt_solver = with_options(COSMO.MKLPardisoKKTSolver, iparm = iparms, num_threads = num_threads, msg_level_on = false))
  @test model.settings.kkt_solver == optimizer.inner.settings.kkt_solver

  # assemble model when only KKT solver type is provided
  settings = COSMO.Settings(kkt_solver = COSMO.QdldlKKTSolver)
  model = COSMO.Model()
  assemble!(model, P, q, c1, settings = settings)
  model.ρvec = settings.rho * ones(size(model.p.A, 1))
  COSMO._make_kkt_solver!(model)
  @test typeof(model.kkt_solver) <: COSMO.QdldlKKTSolver

  # try to assemble with wrong type
  struct MockAccelerator end
  @test_throws MethodError COSMO.Settings(kkt_solver = with_options(MockAccelerator, mem = 5))

end


export PointMassInSphere

# Simple double integrator
mutable struct PointMassInSphere{T<:AbstractFloat} <: Robot
  mass::T
  r::T
  hard_limit_vel::T
  hard_limit_accel::T
  btCollisionObject
end
function PointMassInSphere{T}() where T
  # Default values
  mass = 1.0
  r = 0.1
  hard_limit_vel = 1.0
  hard_limit_accel = 1.0

  btCollisionObject = BulletCollision.convex_hull_cylinder(SVector{3}(zeros(3)), SVector{3}(0.,0.,1.), r)

  return PointMassInSphere{T}(mass, r, hard_limit_vel, hard_limit_accel, btCollisionObject)
end
PointMassInSphere(::Type{T} = Float64; kwargs...) where {T} = PointMassInSphere{T}(; kwargs...)


abstract type AbstractProgressiveHedging end

StochasticPrograms.num_subproblems(ph::AbstractProgressiveHedging) = StochasticPrograms.num_subproblems(ph.structure)

# Initialization #
# ======================================================================== #
function initialize!(ph::AbstractProgressiveHedging, penaltyterm::AbstractPenaltyterm)
    # Initialize progress meter
    ph.progress.thresh = sqrt(ph.parameters.ϵ₁ ^ 2 + ph.parameters.ϵ₂ ^ 2)
    # Initialize subproblems
    initialize_subproblems!(ph, scenarioproblems(ph.structure), penaltyterm)
    # Initialize penalty parameter (if applicable)
    ph.data.δ₁ = 1.0
    initialize_penalty!(ph)
    # Finish initialization
    finish_initilization!(ph, penalty(ph))
end
# ======================================================================== #

# Functions #
# ======================================================================== #
function decision(ph::AbstractProgressiveHedging)
    return ph.ξ
end

function decision(ph::AbstractProgressiveHedging, index::MOI.VariableIndex)
    i = something(findfirst(i -> i == index, ph.decisions.undecided), 0)
    if iszero(i)
        throw(MOI.InvalidIndex(index))
    end
    return decision(ph)[i]
end

function current_objective_value(ph::AbstractProgressiveHedging, Qs::AbstractVector)
    return sum(Qs)
end
current_objective_value(ph) = current_objective_value(ph, ph.subobjectives)

function get_objective_value(ph::AbstractProgressiveHedging)
    if !isempty(ph.Q_history)
        return ph.Q_history[end]
    else
        return calculate_objective_value(ph)
    end
end

function objective_value(ph::AbstractProgressiveHedging)
    return ph.data.Q
end

function log!(ph::AbstractProgressiveHedging; optimal = false, status = nothing)
    @unpack Q, δ₁, δ₂, iterations = ph.data
    @unpack ϵ₁, ϵ₂, keep, offset, indent = ph.parameters
    # Early termination log
    if status != nothing && ph.parameters.log
        ph.progress.thresh = Inf
        ph.progress.printed = true
        val = if status == MOI.INFEASIBLE
            Inf
        elseif status == MOI.DUAL_INFEASIBLE
            -Inf
        else
            0.0
        end
        ProgressMeter.update!(ph.progress, val,
                              showvalues = [
                                  ("$(indentstr(indent))Objective", val),
                                  ("$(indentstr(indent))Early termination", status),
                                  ("$(indentstr(indent))Iterations", iterations)
                              ], keep = keep, offset = offset)
        return nothing
    end
    push!(ph.Q_history, Q)
    push!(ph.primal_gaps, δ₁)
    push!(ph.dual_gaps, δ₂)
    ph.data.iterations += 1
    if ph.parameters.log
        δ = sqrt(δ₁ ^ 2 + δ₂ ^ 2)
        if δ <= ph.progress.thresh && !(δ₁ <= ϵ₁ && δ₂ <= ϵ₂)
            δ = 1.01 * ph.progress.thresh
        end
        δ = optimal ? 0.0 : δ
        ProgressMeter.update!(ph.progress, δ,
                              showvalues = [
                                  ("$(indentstr(indent))Objective",Q),
                                  ("$(indentstr(indent))Primal gap", δ₁),
                                  ("$(indentstr(indent))Dual gap", δ₂),
                                  ("$(indentstr(indent))Iterations", iterations)
                              ], keep = keep, offset = offset)
    end
end

function indentstr(n::Integer)
    return repeat(" ", n)
end

function check_optimality(ph::AbstractProgressiveHedging)
    @unpack ϵ₁, ϵ₂ = ph.parameters
    @unpack δ₁, δ₂ = ph.data
    return δ₁ <= ϵ₁ && δ₂ <= ϵ₂
end
# ======================================================================== #
function show(io::IO, ph::AbstractProgressiveHedging)
    println(io, typeof(ph).name.name)
    println(io, "State:")
    show(io, ph.data)
    println(io, "Parameters:")
    show(io, ph.parameters)
end

function show(io::IO, ::MIME"text/plain", ph::AbstractProgressiveHedging)
    show(io, ph)
end


#=

Want to process `@constraint f(f(x)) ∈ [0.3, 0.4]`
where `f(x) = 4x * (1-x)`

Given code `f(f(x))`, we need `f_forward` and `f_backward`.

Each "copy" of `f` uses the same actual forward and back functions,
`f.forward` and `f.backward`.

```
@function f(x) = 4x * (1-x)
```
should generate these forward and backward functions, and register the function
`f`.

"""
=#

type FunctionObject
    input_args::Vector{Symbol}
    output_args::Vector{Symbol}
    local_vars::Vector{Symbol}
    code::Expr
end

doc"""
A `ConstraintFunction` contains the created forward and backward
code
"""
type ConstraintFunction{F <: Function, G <: Function}
    input::Vector{Symbol}  # input arguments for forward function
    output::Vector{Symbol} # output arguments for forward function
    forward::F
    backward::G
end

type FunctionArguments
    # input::Vector{Symbol}  # input arguments for forward function
    # output::Vector{Symbol}  # output arguments
    # generated::Vector{Symbol} # local variables generated

    input
    generated
    return_arguments
end

const registered_functions = Dict{Symbol, FunctionArguments}()

@doc """
`@function` registers a function to be used in forwards and backwards mode.

Example: `@function f(x, y) = x^2 + y^2`
"""  # this docstring does not work!

@eval macro ($(:function))(ex)   # workaround to define macro @function

    # (f, args, code) = @match ex begin
    #     ( f_(args__) = code_ ) => (f, args, code)
    # end

    (f, args, code) = match_function(ex)

    @show f, args, code

    #root, all_vars, generated, code2 = IntervalConstraintProgramming.flatten!(code)
    return_arguments, flatAST = IntervalConstraintProgramming.flatten!(code)

    @show return_arguments

    #@show root, all_vars, generated, code2

    # rearrange so actual return arguments come first:

    @show flatAST
    @show return_arguments

    # make into an array:
    if !(isa(return_arguments, Array))
        return_arguments = [return_arguments]
    end

    flatAST.intermediate = setdiff(flatAST.intermediate, return_arguments)

    println("HERE")
    flatAST.intermediate = [return_arguments; flatAST.intermediate]

    forward_code = forward_pass(flatAST) #root, all_vars, generated, code2)
    backward_code = backward_pass(flatAST) #root, all_vars, generated, code2)

    @show forward_code, backward_code

    registered_functions[f] = FunctionArguments(flatAST.variables, flatAST.intermediate, return_arguments)

    return quote
        #$(esc(Meta.quot(f))) = ConstraintFunction($(all_vars), $(generated), $(forward_code), $(backward_code))
        $(esc(f)) = ConstraintFunction($(flatAST.variables), $(flatAST.intermediate), $(forward_code), $(backward_code))
        #registered_functions[$(Meta.quot(f))] =  ConstraintFunction($(all_vars), $(generated), $(forward_code), $(backward_code))
        #$(Meta.quot(f)) =  ConstraintFunction($(all_vars), $(generated), $(forward_code), $(backward_code))
    end
end

function match_function(ex)

    try
        (f, args, body) =
            @match ex begin
             ( (f_(args__) = body_) |
              (function f_(args__) body_ end) ) => (f, args, body)
           end

         return (f, args, rmlines(body))  # rmlines is from MacroTools package

    catch
        throw(ArgumentError("$ex does not have the form of a function"))
    end
end

doc"""
Generate code for an anonymous function with given
input arguments, output arguments, and code block.
"""
function make_function(input_args, output_args, code)

    input = Expr(:tuple, input_args...)  # make a tuple of the variables
    output = Expr(:tuple, output_args...)  # make a tuple of the variables

    new_code = copy(code)
    push!(new_code.args, :(return $output))

    complete_code = :( $input -> $new_code )

    return GeneratedFunction(input_args, output_args, complete_code)
end


"""
    VerticalStructure

Vertical memory structure. Decomposes stochastic program by stages.

"""
struct VerticalStructure{N, M, SP <: NTuple{M, AbstractScenarioProblems}} <: AbstractBlockStructure{N}
    decisions::NTuple{N, Decisions}
    first_stage::JuMP.Model
    scenarioproblems::SP

    function VerticalStructure(decisions::NTuple{N, Decisions}, scenarioproblems::NTuple{M,AbstractScenarioProblems}) where {N, M}
        M == N - 1 || error("Inconsistent number of stages $N and number of scenario types $M")
        SP = typeof(scenarioproblems)
        return new{N,M,SP}(decisions, Model(), scenarioproblems)
    end
end

function StochasticStructure(decisions::NTuple{N, Decisions}, scenario_types::ScenarioTypes{M}, instantiation::Union{Vertical, DistributedVertical}) where {N, M}
    scenarioproblems = ntuple(Val(M)) do i
        ScenarioProblems(scenario_types[i], instantiation)
    end
    return VerticalStructure(decisions, scenarioproblems)
end

function StochasticStructure(decisions::NTuple{N, Decisions}, scenarios::NTuple{M, Vector{<:AbstractScenario}}, instantiation::Union{Vertical, DistributedVertical}) where {N, M}
    scenarioproblems = ntuple(Val(M)) do i
        ScenarioProblems(scenarios[i], instantiation)
    end
    return VerticalStructure(decisions, scenarioproblems)
end

# Base overloads #
# ========================== #
function Base.print(io::IO, structure::VerticalStructure{N}) where N
    print(io, "Stage 1\n")
    print(io, "============== \n")
    print(io, structure.first_stage)
    for stage in 2:N
        print(io, "\nStage $stage\n")
        print(io, "============== \n")
        for (id, subproblem) in enumerate(subproblems(structure, stage))
            @printf(io, "Subproblem %d (p = %.2f):\n", id, probability(structure, stage, id))
            print(io, subproblem)
            print(io, "\n")
        end
    end
end
function Base.print(io::IO, structure::VerticalStructure{2})
    print(io, "First-stage \n")
    print(io, "============== \n")
    print(io, structure.first_stage)
    print(io, "\nSecond-stage \n")
    print(io, "============== \n")
    for (id, subproblem) in enumerate(subproblems(structure, 2))
        @printf(io, "Subproblem %d (p = %.2f):\n", id, probability(structure, 2, id))
        print(io, subproblem)
        print(io, "\n")
    end
end
# ========================== #

# MOI #
# ========================== #
function MOI.get(structure::VerticalStructure, attr::MOI.AbstractModelAttribute)
    return MOI.get(backend(structure.first_stage), attr)
end
function MOI.get(structure::VerticalStructure, attr::MOI.AbstractVariableAttribute, index::MOI.VariableIndex)
    return MOI.get(backend(structure.first_stage), attr, index)
end
function MOI.get(structure::VerticalStructure, attr::Type{MOI.VariableIndex}, name::String)
    return MOI.get(backend(structure.first_stage), attr, name)
end
function MOI.get(structure::VerticalStructure, attr::MOI.AbstractConstraintAttribute, cindex::MOI.ConstraintIndex)
    return MOI.get(backend(structure.first_stage), attr, cindex)
end
function MOI.get(structure::VerticalStructure, attr::ScenarioDependentModelAttribute)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    return MOI.get(scenarioproblems(structure, attr.stage), attr)
end
function MOI.get(structure::VerticalStructure, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    return MOI.get(scenarioproblems(structure, attr.stage), attr, index)
end
function MOI.get(structure::VerticalStructure, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    return MOI.get(scenarioproblems(structure, attr.stage), attr, ci)
end

function MOI.set(structure::VerticalStructure, attr::MOI.AbstractModelAttribute, value)
    MOI.set(backend(structure.first_stage), attr, value)
    return nothing
end
function MOI.set(structure::VerticalStructure, attr::MOI.Silent, flag)
    # Silence master
    MOI.set(backend(structure.first_stage), attr, flag)
    # Silence subproblems
    MOI.set(scenarioproblems(structure), attr, flag)
    return nothing
end
function MOI.set(structure::VerticalStructure, attr::MOI.AbstractVariableAttribute,
                 index::MOI.VariableIndex, value)
    MOI.set(backend(structure.first_stage), attr, index, value)
    return nothing
end
function MOI.set(structure::VerticalStructure, attr::MOI.AbstractConstraintAttribute,
                 ci::MOI.ConstraintIndex, value)
    MOI.set(backend(structure.first_stage), attr, ci, value)
    return nothing
end
function MOI.set(structure::VerticalStructure, attr::ScenarioDependentModelAttribute, value)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    MOI.set(scenarioproblems(subproblem, attr.stage), attr, value)
    return nothing
end
function MOI.set(structure::VerticalStructure, attr::ScenarioDependentVariableAttribute,
                 index::MOI.VariableIndex, value)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    MOI.set(scenarioproblems(structure, attr.stage), attr, index, value)
    return nothing
end
function MOI.set(structure::VerticalStructure, attr::ScenarioDependentConstraintAttribute,
                 ci::MOI.ConstraintIndex, value)
    n = num_scenarios(structure, attr.stage)
    1 <= attr.scenario_index <= n || error("Scenario index $attr.scenario_index not in range 1 to $n.")
    MOI.set(scenarioproblems(structure, attr.stage), attr, ci, value)
    return nothing
end

function MOI.is_valid(structure::VerticalStructure, index::MOI.VariableIndex, stage::Integer)
    stage == 1 || error("No scenario index specified.")
    return MOI.is_valid(backend(structure.first_stage), index)
end
function MOI.is_valid(structure::VerticalStructure, ci::MOI.ConstraintIndex, stage::Integer)
    stage == 1 || error("No scenario index specified.")
    return MOI.is_valid(backend(structure.first_stage), ci)
end
function MOI.is_valid(structure::VerticalStructure{N}, ci::MOI.ConstraintIndex, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    return MOI.is_valid(scenarioproblems(structure, stage), ci, scenario_index)
end

function MOI.add_constraint(structure::VerticalStructure, f::SingleDecision, s::MOI.AbstractSet)
    return MOI.add_constraint(backend(structure.first_stage), f, s)
end
function MOI.add_constraint(structure::VerticalStructure, f::SingleDecision, s::MOI.AbstractSet, stage::Integer, scenario_index::Integer)
    stage == 1 && error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    return MOI.add_constraint(scenarioproblems(structure, stage), f, s, scenario_index)
end

function MOI.delete(structure::VerticalStructure{N}, index::MOI.VariableIndex, stage::Integer) where N
    stage == 1 || error("No scenario index specified.")
    JuMP.delete(structure.first_stage, DecisionRef(structure.first_stage, index))
    # Remove known decisions
    for s in 2:N
        JuMP.delete(scenarioproblems(structure, s), index)
    end
    return nothing
end
function MOI.delete(structure::VerticalStructure{N}, indices::Vector{MOI.VariableIndex}, stage::Integer) where N
    stage == 1 || error("No scenario index specified.")
    JuMP.delete(structure.first_stage, DecisionRef.(structure.first_stage, indices))
    # Remove known decisions
    for s in 2:N
        JuMP.delete(scenarioproblems(structure, s), indices)
    end
    return nothing
end
function MOI.delete(structure::VerticalStructure, ci::MOI.ConstraintIndex, stage::Integer)
    stage == 1 || error("No scenario index specified.")
    MOI.delete(backend(structure.first_stage), ci)
    return nothing
end
function MOI.delete(structure::VerticalStructure, cis::Vector{<:MOI.ConstraintIndex}, stage::Integer)
    stage == 1 || error("No scenario index specified.")
    MOI.delete(backend(structure.first_stage), cis)
    return nothing
end
function MOI.delete(structure::VerticalStructure{N}, ci::MOI.ConstraintIndex, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    MOI.delete(scenarioproblems(structure, stage), ci, scenario_index)
    return nothing
end
function MOI.delete(structure::VerticalStructure{N}, cis::Vector{<:MOI.ConstraintIndex}, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    MOI.delete(scenarioproblems(structure, stage), cis, scenario_index)
    return nothing
end

# JuMP #
# ========================== #
function JuMP.fix(structure::VerticalStructure{N}, index::MOI.VariableIndex, stage::Integer, val::Number) where N
    dref = DecisionRef(structure.first_stage, index)
    fix(dref, val)
    update_decision_state!(scenarioproblems(structure), index, Known)
    return nothing
end
function JuMP.fix(structure::VerticalStructure{N}, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer, val::Number) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    fix(scenarioproblems(structure, stage), index, scenario_index, val)
    return nothing
end
function JuMP.unfix(structure::VerticalStructure{N}, index::MOI.VariableIndex, stage::Integer) where N
    dref = DecisionRef(structure.first_stage, index)
    unfix(dref)
    return nothing
end
function JuMP.unfix(structure::VerticalStructure{N}, index::MOI.VariableIndex, stage::Integer, scenario_index::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    unfix(scenarioproblems(structure, stage), index, scenario_index)
    return nothing
end
function JuMP.objective_function_type(structure::VerticalStructure)
    return jump_function_type(structure.first_stage,
                              MOI.get(structure, MOI.ObjectiveFunctionType()))
end

function JuMP.objective_function(structure::VerticalStructure, FunType::Type{<:AbstractJuMPScalar})
    MOIFunType = moi_function_type(FunType)
    func = MOI.get(structure.first_stage,
                   MOI.ObjectiveFunction{MOIFunType}())::MOIFunType
    return jump_function(structure.first_stage, func)
end

function JuMP.objective_function(structure::AbstractBlockStructure, proxy::JuMP.Model, stage::Integer, FunType::Type{<:AbstractJuMPScalar})
    if stage == 1
        return objective_function(structure, FunType)
    else
        return objective_function(proxy, FunType)
    end
end

function JuMP._moi_optimizer_index(structure::VerticalStructure, index::VI)
    return decision_index(backend(structure.first_stage), index)
end
function JuMP._moi_optimizer_index(structure::VerticalStructure, ci::CI)
    return decision_index(backend(structure.first_stage), ci)
end

function JuMP.set_objective_coefficient(structure::VerticalStructure{N}, index::VI, var_stage::Integer, stage::Integer, coeff::Real) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    var_stage <= stage || error("Can only modify coefficient in current stage of decision or subsequent stages from where decision is taken.")
    if var_stage == 1
        if stage == 1
            dref = DecisionRef(structure.first_stage, index)
            set_objective_coefficient(structure.first_stage, dref, coeff)
        else
            set_objective_coefficient(scenarioproblems(structure, stage), index, coeff)
        end
    else
        for scenario_index in 1:num_scenarios(structure, stage)
            set_objective_coefficient(scenarioproblems(structure, stage), index, scenario_index, coeff)
        end
    end
    return nothing
end
function JuMP.set_objective_coefficient(structure::VerticalStructure{N}, index::VI, var_stage::Integer, stage::Integer, scenario_index::Integer, coeff::Real) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    if var_stage == 1
        set_objective_coefficient(scenarioproblems(structure, stage), index, scenario_index, coeff)
    else
        set_objective_coefficient(scenarioproblems(structure, stage), index, scenario_index, coeff)
    end
    return nothing
end

function JuMP.set_normalized_coefficient(structure::VerticalStructure,
                                         ci::CI{F,S},
                                         index::VI,
                                         value) where {T, F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}}, S}
    MOI.modify(backend(structure.first_stage), ci,
               DecisionCoefficientChange(index, convert(T, value)))
    return nothing
end
function JuMP.set_normalized_coefficient(structure::VerticalStructure{N},
                                         ci::CI{F,S},
                                         index::VI,
                                         stage::Integer,
                                         scenario_index::Integer,
                                         value) where {N, T, F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}}, S}
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    set_normalized_coefficient(scenarioproblems(structure, stage), ci, index, scenario_index, value)
    return nothing
end

function JuMP.normalized_coefficient(structure::VerticalStructure{N},
                                     ci::CI{F,S},
                                     index::VI,
                                     stage::Integer,
                                     scenario_index::Integer) where {N, T, F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}}, S}
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    return normalized_coefficient(scenarioproblems(structure, stage), ci, index, scenario_index)
end

function JuMP.set_normalized_rhs(structure::VerticalStructure,
                                 ci::CI{F,S},
                                 value) where {T,
                                               F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}},
                                               S <: MOIU.ScalarLinearSet{T}}
    MOI.set(backend(structure.first_stage), MOI.ConstraintSet(), ci,
            S(convert(T, value)))
    return nothing
end
function JuMP.set_normalized_rhs(structure::VerticalStructure{N},
                                 ci::CI{F,S},
                                 stage::Integer,
                                 scenario_index::Integer,
                                 value) where {N,
                                               T,
                                               F <: Union{AffineDecisionFunction{T}, QuadraticDecisionFunction{T}},
                                               S <: MOIU.ScalarLinearSet{T}}
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage > 1 || error("There are no scenarios in the first stage.")
    n = num_scenarios(structure, stage)
    1 <= scenario_index <= n || error("Scenario index $scenario_index not in range 1 to $n.")
    set_normalized_rhs(scenarioproblems(structure, stage), ci, scenario_index, value)
    return nothing
end

function DecisionRef(proxy::JuMP.Model, structure::VerticalStructure, index::VI)
    return DecisionRef(structure.first_stage, index)
end

function JuMP.jump_function(structure::VerticalStructure{N},
                            proxy::JuMP.Model,
                            stage::Integer,
                            f::MOI.AbstractFunction) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    if stage == 1
        return JuMP.jump_function(structure.first_stage, f)
    else
        return JuMP.jump_function(proxy, f)
    end
end

# Getters #
# ========================== #
function structure_name(structure::VerticalStructure)
    return "Vertical"
end

deferred_first_stage(structure::VerticalStructure, ::Val{1}) = num_variables(first_stage(structure)) == 0

function decision(structure::VerticalStructure{N}, index::MOI.VariableIndex, stage::Integer) where N
    stage == 1 || error("No scenario index specified.")
    return decision(structure.decisions[stage], index)
end
# ========================== #

# Setters
# ========================== #
function update_known_decisions!(structure::VerticalStructure)
    update_known_decisions!(structure.first_stage)
    return nothing
end


"""
    setup_problem(problem, parameters::Parameters)

Set up a fixed-dimensional optimization problem.

There are several `setup_problem()` method that accept different type of
`problem` argument:
    * `OptimizationProblem`
    * function (`:NumDimensions` has to be specified in `parameters`)
    * `FunctionBasedProblemFamily` (`:NumDimensions` has to be specified in `parameters`)
"""
function setup_problem(problem::OptimizationProblem, parameters::Parameters)
    return problem, parameters
end

function setup_problem(family::FunctionBasedProblemFamily, parameters::Parameters)
    # If an anydim problem was given the dimension param must have been specified.
    if parameters[:NumDimensions] == :NotSpecified
        throw(ArgumentError("You MUST specify NumDimensions= when a problem family is given"))
    end
    problem = instantiate(family, parameters[:NumDimensions])

    return problem, parameters
end

function setup_problem(func, parameters::Parameters)
    ss = check_and_create_search_space(parameters)

    # Now create an optimization problem with the given information. We currently reuse the type
    # from our pre-defined problems so some of the data for the constructor is dummy.
    problem = FunctionBasedProblem(func, "<unknown>", parameters[:FitnessScheme], ss,
                                   parameters[:TargetFitness])

    # validate fitness: create a random solution from the search space and ensure that fitness(problem) returns fitness_type(problem).
    ind = rand_individual(search_space(problem))
    res = fitness(ind, problem)
    fitnessT = fitness_type(problem)
    if !isa(res, fitnessT)
      throw(ArgumentError("The supplied fitness function does NOT return the expected fitness type $(fitnessT)"*
                          "when called with a potential solution "*
                          "(when called with $(ind) it returned $(res) of type $(typeof(res)) so we cannot optimize it!"))
    end

    return problem, parameters
end

"""
    bboptimize(problem[; parameters::Associative, kwargs...])

Solve given optimization `problem`.

See `setup_problem()` for the types of `problem` supported.
In addition, the `problem` could be `OptController` containing the
results of the previous optimization runs.

The optimization method parameters could be specified via `kwargs` or `parameters` arguments.

Returns `OptimizationResults` instance.

See also `bbsetup()` and [`BlackBoxOptim.OptRunController`](@ref) for a full list of supported parameters.
"""
function bboptimize(optctrl::OptController; kwargs...)
    if length(kwargs) > 0
        update_parameters!(optctrl, kwargs2dict(kwargs))
    end
    run!(optctrl)
end

function bboptimize(functionOrProblem, parameters::Parameters = EMPTY_PARAMS; kwargs...)
    optctrl = bbsetup(functionOrProblem, parameters; kwargs...)
    run!(optctrl)
end

"""
    bbsetup(problem[; parameters::Associative, kwargs...])

Set up optimization method for a given problem.

See `setup_problem()` for the types of `problem` supported.
The optimization method parameters could be specified via `kwargs` or `parameters` arguments.

Returns the initialized `OptController` instance. To actually run the method
call `bboptimize()` or `run!()`.

See also [BlackBoxOptim.OptRunController](@ref) for a full list of supported parameters.
"""
function bbsetup(functionOrProblem, parameters::Parameters = EMPTY_PARAMS; kwargs...)
    parameters = chain(convert(ParamsDict, parameters), kwargs2dict(kwargs))
    problem, params = setup_problem(functionOrProblem, chain(DefaultParameters, parameters))
    check_valid!(params)

    optimizer_func = chain(SingleObjectiveMethods, MultiObjectiveMethods)[params[:Method]]
    optimizer = optimizer_func(problem, params)

    # Now set up a controller for this problem. This will handle
    # application of optimizer, checking for termination conditions
    # as well as collecting the statistics about the number of function evals,
    # keep an archive and top list of candidates.
    ctrl = OptController(optimizer, problem, params)

    return ctrl
end


using ReinforcementLearningEnvironmentGym, Flux, ReinforcementLearning
# List all envs

listallenvs()

# CartPole example

env = GymEnv("CartPole-v0")
learner = DQN(Chain(Dense(4, 48, relu), Dense(48, 24, relu), Dense(24, 2)),
                  updateevery = 1, updatetargetevery = 100,
                  startlearningat = 50, minibatchsize = 32,
                  doubledqn = false, replaysize = 10^3,
                  opttype = x -> ADAM(x, .0005))
x = RLSetup(learner, env, ConstantNumberEpisodes(10),
            callbacks = [Progress(), EvaluationPerEpisode(TimeSteps()),
                         Visualize(wait = 0)])
@info("Before learning.")
run!(x)
pop!(x.callbacks)
x.stoppingcriterion = ConstantNumberEpisodes(400)
@time learn!(x)
x.stoppingcriterion = ConstantNumberEpisodes(10)
push!(x.callbacks, Visualize(wait = 0))
@info("After learning.")
run!(x)


# Copyright (c) 2019 Simon Castano
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.

@testset "Transaction" begin
    @testset "TxIn" begin
        txin_hex = "7b1eabe0209b1fe794124575ef807057c77ada2138ae4fa8d6c4de0398a14f3f00000000494830450221008949f0cb400094ad2b5eb399d59d01c14d73d8fe6e96df1a7150deb388ab8935022079656090d7f6bac4c9a94e0aad311a4268e082a725f8aeae0573fb12ff866a5f01ffffffff"
        raw = hex2bytes(txin_hex)
        txin = TxIn(IOBuffer(raw))
        @test typeof(txin.prevout) == Outpoint
        @test bytes2hex(txin.prevout.txid) == "3f4fa19803dec4d6a84fae3821da7ac7577080ef75451294e71f9b20e0ab1e7b"
        @test txin.prevout.index == 0
        want = hex2bytes("30450221008949f0cb400094ad2b5eb399d59d01c14d73d8fe6e96df1a7150deb388ab8935022079656090d7f6bac4c9a94e0aad311a4268e082a725f8aeae0573fb12ff866a5f01")
        @test txin.scriptsig.data == [want]
        @test txin.sequence == 0xffffffff
    end
    @testset "TxOut" begin
        txout_hex = "f0ca052a010000001976a914cbc20a7664f2f69e5355aa427045bc15e7c6c77288ac"
        raw = hex2bytes(txout_hex)
        txout = TxOut(IOBuffer(raw))
        @test txout.value == 4999990000
        @test txout.scriptpubkey.data == [[0x76], [0xa9], hex2bytes("cbc20a7664f2f69e5355aa427045bc15e7c6c772"), [0x88], [0xac]]
    end

    raw_tx = hex2bytes("0100000001813f79011acb80925dfe69b3def355fe914bd1d96a3f5f71bf8303c6a989c7d1000000006b483045022100ed81ff192e75a3fd2304004dcadb746fa5e24c5031ccfcf21320b0277457c98f02207a986d955c6e0cb35d446a89d3f56100f4d7f67801c31967743a9c8e10615bed01210349fc4e631e3624a545de3f89f5d8684c7b8138bd94bdd531d2e213bf016b278afeffffff02a135ef01000000001976a914bc3b654dca7e56b04dca18f2566cdaf02e8d9ada88ac99c39800000000001976a9141c4bc762dd5423e332166702cb75f40df79fea1288ac19430600")
    raw_bip141 = hex2bytes("0100000000010115e180dc28a2327e687facc33f10f2a20da717e5548406f7ae8b4c811072f85601000000000fffffff0100b4f505000000001976a9141d7cd6c75c2e86f4cbf98eaed221b30bd9a0b92888ac02483045022100df7b7e5cda14ddf91290e02ea10786e03eb11ee36ec02dd862fe9a326bbcb7fd02203f5b4496b667e6e281cc654a2da9e4f08660c620a1051337fa8965f727eb19190121038262a6c6cec93c2d3ecd6c6072efea86d02ff8e3328bbd0242b20af3425990ac00000000")
    @testset "Parsing" begin
        @testset "Legacy" begin
            tx = Tx(IOBuffer(raw_tx))
            @test tx.version == 1
            @test tx.locktime == 0x00064319
            @testset "Inputs" begin
                @test length(tx.inputs) == 1
                want = hex2bytes("d1c789a9c60383bf715f3f6ad9d14b91fe55f3deb369fe5d9280cb1a01793f81")
                @test tx.inputs[1].prevout.txid == want
                @test tx.inputs[1].prevout.index == 0
                want = hex2bytes("6b483045022100ed81ff192e75a3fd2304004dcadb746fa5e24c5031ccfcf21320b0277457c98f02207a986d955c6e0cb35d446a89d3f56100f4d7f67801c31967743a9c8e10615bed01210349fc4e631e3624a545de3f89f5d8684c7b8138bd94bdd531d2e213bf016b278a")
                @test serialize(tx.inputs[1].scriptsig) == want
                @test tx.inputs[1].sequence == 0xfffffffe
            end
            @testset "Outputs" begin
                @test length(tx.outputs) == 2
                want = 32454049
                @test tx.outputs[1].value == want
                want = hex2bytes("1976a914bc3b654dca7e56b04dca18f2566cdaf02e8d9ada88ac")
                @test serialize(tx.outputs[1].scriptpubkey) == want
                want = 10011545
                @test tx.outputs[2].value == want
                want = hex2bytes("1976a9141c4bc762dd5423e332166702cb75f40df79fea1288ac")
                @test serialize(tx.outputs[2].scriptpubkey) == want
            end
        end
        @testset "Segwit" begin
            tx = Tx(IOBuffer(raw_bip141))
            @test tx.version == 0x00000001
            @test tx.flag == 0x01
            hex1 = "3045022100df7b7e5cda14ddf91290e02ea10786e03eb11ee36ec02dd862fe9a326bbcb7fd02203f5b4496b667e6e281cc654a2da9e4f08660c620a1051337fa8965f727eb191901"
            hex2 = "038262a6c6cec93c2d3ecd6c6072efea86d02ff8e3328bbd0242b20af3425990ac"
            @test tx.witnesses[1].data[1] == hex2bytes(hex1)
            @test tx.witnesses[1].data[2] == hex2bytes(hex2)
        end
    end
    @testset "Serialize" begin
        tx = Tx(IOBuffer(raw_tx))
        @test serialize(tx) == raw_tx

        tx = Tx(IOBuffer(raw_bip141))
        @test serialize(tx) == raw_bip141
    end
    @testset "Is CoinbaseTx" begin
        raw_tx = hex2bytes("01000000010000000000000000000000000000000000000000000000000000000000000000ffffffff5e03d71b07254d696e656420627920416e74506f6f6c20626a31312f4542312f4144362f43205914293101fabe6d6d678e2c8c34afc36896e7d9402824ed38e856676ee94bfdb0c6c4bcd8b2e5666a0400000000000000c7270000a5e00e00ffffffff01faf20b58000000001976a914338c84849423992471bffb1a54a8d9b1d69dc28a88ac00000000")
        stream = IOBuffer(raw_tx)
        tx = Tx(stream)
        @test iscoinbase(tx)
    end
    @testset "coinbase_height" begin
        raw_tx = hex2bytes("01000000010000000000000000000000000000000000000000000000000000000000000000ffffffff5e03d71b07254d696e656420627920416e74506f6f6c20626a31312f4542312f4144362f43205914293101fabe6d6d678e2c8c34afc36896e7d9402824ed38e856676ee94bfdb0c6c4bcd8b2e5666a0400000000000000c7270000a5e00e00ffffffff01faf20b58000000001976a914338c84849423992471bffb1a54a8d9b1d69dc28a88ac00000000")
        stream = IOBuffer(raw_tx)
        tx = Tx(stream)
        @test coinbase_height(tx) == 465879
        raw_tx = hex2bytes("0100000001813f79011acb80925dfe69b3def355fe914bd1d96a3f5f71bf8303c6a989c7d1000000006b483045022100ed81ff192e75a3fd2304004dcadb746fa5e24c5031ccfcf21320b0277457c98f02207a986d955c6e0cb35d446a89d3f56100f4d7f67801c31967743a9c8e10615bed01210349fc4e631e3624a545de3f89f5d8684c7b8138bd94bdd531d2e213bf016b278afeffffff02a135ef01000000001976a914bc3b654dca7e56b04dca18f2566cdaf02e8d9ada88ac99c39800000000001976a9141c4bc762dd5423e332166702cb75f40df79fea1288ac19430600")
        stream = IOBuffer(raw_tx)
        tx = Tx(stream)
        @test_throws AssertionError coinbase_height(tx)
    end
end


#
#  UpdateSystem.jl
#  CollagenPolymer
#
#  Created by Christopher Revell on 30/03/2020.
#
#

module UpdateSystem

# Import Julia packages
using LinearAlgebra
using StaticArrays

# Import local program modules

@views function updateSystem!(pos,F,W,t,Δt,D,kT,nParticles)

    pos .+= F[:,1]*(Δt*D/kT) .+ W.*sqrt(2.0*D*Δt)
    fill!(F,@SVector zeros(3))

    return t += Δt

end

export updateSystem!

end