text:
module solve_eqm
# Write your package code here.
using NLsolve

#Calculates the Sigmoid
function logit(x, fixedParam)
    return 1 ./ (1.0 .+ exp.(-fixedParam.lambda .* x))
end

#Solves for steady state distribution of workforce across age stages
function solveStages(var, fixedParam)
    NC_f    = var[1];
    YC_f    = var[2];
    PL_f    = var[3];
    D_f     = var[4];

    NC_m    = var[5];
    YC_m    = var[6];
    PL_m    = var[7];
    D_m     = var[8];

    # 4 equations for women
    eq1f = NC_f + YC_f + PL_f + D_f - 0.5;
    eq2f = fixedParam.chi*(NC_f+YC_f) - (fixedParam.etaf+fixedParam.gamma)*PL_f; # flows into and out of PL
    eq3f = fixedParam.etaf*PL_f - (fixedParam.chi+fixedParam.gamma)*YC_f; # flows into and out of YC
    eq4f = fixedParam.gamma*(NC_f+YC_f+PL_f) - fixedParam.phi*D_f;

    # 4 equations for men
    eq1m = NC_m + YC_m + PL_m + D_m - 0.5;
    eq2m = fixedParam.chi*(NC_m+YC_m) - (fixedParam.etam+fixedParam.gamma)*PL_m; # flows into and out of PL
    eq3m = fixedParam.etam*PL_m - (fixedParam.chi+fixedParam.gamma)*YC_m; # flows into and out of YC
    eq4m = fixedParam.gamma*(NC_m+YC_m+PL_m) - fixedParam.phi*D_m;

    return [eq1f, eq2f, eq3f, eq4f, eq1m, eq2m, eq3m, eq4m];
end

#Functions for updating male value functions
#Male No Child
function getPmNC(x, y,v0,fxp,par,Pi,gem,SmNC,UmNC,PmD_hat,PmPL_hat,PmNC_hat,epsm,alpha,output,tau)
    expectedPmNC = 0.0
    dmNC1 = SmNC[x, :] .- SmNC[x, y]
    #Calculate the value function if no hc increment in the next round
    val1 = fxp.deltam_NC * (Pi[y] + UmNC[x]) +
           fxp.lambdae * par.sigma * sum(v0 .* dmNC1 .* logit(dmNC1, fxp)) * (SmNC[x, y] > 0) +
           fxp.gamma * PmD_hat[x, y] + fxp.chi * PmPL_hat[x, y] +
           (1 - fxp.deltam_NC - fxp.gamma - fxp.chi) * PmNC_hat[x, y]
    #Check if it is possible for HC to grow
    if (x % 7 == 0)
        expectedPmNC = val1
    else
        dmNC2 = SmNC[x+1, :] .- SmNC[x+1, y]
        #Calculate the value function if hc increments in the next round
        val2 = fxp.deltam_NC * (Pi[y] + UmNC[x+1]) +
               fxp.lambdae * par.sigma * sum(v0 .* dmNC2 .* logit(dmNC2, fxp)) * (SmNC[x+1, y] > 0) +
               fxp.gamma * PmD_hat[x+1, y] + fxp.chi * PmPL_hat[x+1, y] +
               (1 - fxp.deltam_NC - fxp.gamma - fxp.chi) * PmNC_hat[x+1, y]
        #weight the expected value by probability of HC growth
        if (y % 7 == 0)
            expectedPmNC = gem[7] * val2 + (1 - gem[7]) * val1
        else
            expectedPmNC = gem[y%7] * val2 + (1 - gem[y%7]) * val1
        end
    end
    #return sum of value at current round and expected value for next perioid
    return (1 - tau[1]) * output[x, y] + epsm[x] * alpha[y] + fxp.beta * expectedPmNC
end

#Male Parental Leave
function getPmPL(x, y,Pi,fxp,UmPL,PmD_hat,PmYC_hat,PmPL_hat,epsm,alpha,output)
    val = fxp.deltam_YC * (Pi[y] + UmPL[x]) +
          fxp.gamma * PmD_hat[x, y] +
          fxp.etam * PmYC_hat[x, y] +
          (1 - fxp.deltam_YC - fxp.gamma - fxp.etam) * PmPL_hat[x, y]

    return fxp.Transfer * output[x, y] + epsm[x] * alpha[y] + fxp.beta * val
end

#Male Young Child
function getPmYC(x, y,v0,Pi,fxp,par,epsm,gem,tau,output,alpha,UmYC,SmYC,PmD_hat,PmPL_hat,PmYC_hat)
    expectedPmYC = 0.0
    dmYC1 = SmYC[x, :] .- SmYC[x, y]
    #Calculate the value function if no hc increment in the next round
    val1 = fxp.deltam_YC * (Pi[y] + UmYC[x]) +
           fxp.lambdae * par.sigma * sum(v0 .* dmYC1 .* logit(dmYC1, fxp)) * (SmYC[x, y] > 0) +
           fxp.gamma * PmD_hat[x, y] + fxp.chi * PmPL_hat[x, y] +
           (1 - fxp.deltam_YC - fxp.gamma - fxp.chi) * PmYC_hat[x, y]
    #Check if it is possible for HC to grow
    if (x % 7 == 0)
        expectedPmYC = val1
    else
        dmYC2 = SmYC[x+1, :] .- SmYC[x+1, y]
        #Calculate the value function if hc increments in the next round
        val2 = fxp.deltam_YC * (Pi[y] + UmYC[x+1]) +
               fxp.lambdae * par.sigma * sum(v0 .* dmYC2 .* logit(dmYC2, fxp)) * (SmYC[x+1, y] > 0) +
               fxp.gamma * PmD_hat[x+1, y] + fxp.chi * PmPL_hat[x+1, y] +
               (1 - fxp.deltam_YC - fxp.gamma - fxp.chi) * PmYC_hat[x+1, y]
        if (y % 7 == 0)
            expectedPmYC = gem[7] * val2 + (1 - gem[7]) * val1
        else
            expectedPmYC = gem[y%7] * val2 + (1 - gem[y%7]) * val1
        end
    end
    return (1 - tau[1]) * output[x, y] + epsm[x] * alpha[y] + fxp.beta * expectedPmYC
end

#Male Done with Children
function getPmD(x,y,v0,Pi,fxp,par,output,gem,tau,epsm,alpha,SmD,PmD_hat,UmD)
    expectedPmD = 0.0
    dmD1 = SmD[x, :] .- SmD[x, y]
    #Calculate the value function if no hc increment in the next round
    val1 = fxp.phi * Pi[y] + fxp.delta * (Pi[y] + UmD[x]) +
           (1 - fxp.phi - fxp.delta) * PmD_hat[x, y] +
           fxp.lambdae * par.sigma * sum(v0 .* dmD1 .* logit(dmD1, fxp)) * (SmD[x, y] > 0)
    #Check if it is possible for HC to grow
    if (x % 7 == 0)
        expectedPmD = val1
    else
        dmD2 = SmD[x+1, :] .- SmD[x+1, y]
        #Calculate the value function if hc increments in the next round
        val2 = fxp.phi * Pi[y] + fxp.delta * (Pi[y] + UmD[x+1]) +
               (1 - fxp.phi - fxp.delta) * PmD_hat[x+1, y] +
               fxp.lambdae * par.sigma * sum(v0 .* dmD2 .* logit(dmD2, fxp)) * (SmD[x+1, y] > 0)
        if (y % 7 == 0)
            expectedPmD = gem[7] * val2 + (1 - gem[7]) * val1
        else
            expectedPmD = gem[y%7] * val2 + (1 - gem[y%7]) * val1
        end
    end
    return (1 - tau[1]) * output[x, y] + epsm[x] * alpha[y] + fxp.beta * expectedPmD
end

#Functions for updating female value functions
#Female No Child
function getPfNC(x, y,v0,fxp,par,Pi,gef,SfNC,UfNC,PfD_hat,PfPL_hat,PfNC_hat,epsf,alpha,output,tau)
    expectedPfNC = 0.0
    dfNC1 = SfNC[x, :] .- SfNC[x, y]
    #Calculate the value function if no hc increment in the next round
    val1 = fxp.deltaf_NC * (Pi[y] + UfNC[x]) +
           fxp.lambdae * par.sigma * sum(v0 .* dfNC1 .* logit(dfNC1, fxp)) * (SfNC[x, y] > 0) +
           fxp.gamma * PfD_hat[x, y] + fxp.chi * PfPL_hat[x, y] +
           (1 - fxp.deltaf_NC - fxp.gamma - fxp.chi) * PfNC_hat[x, y]
    #Check if it is possible for HC to grow
    if (x % 7 == 0)
        expectedPfNC = val1
    else
        dfNC2 = SfNC[x+1, :] .- SfNC[x+1, y]
        #Calculate the value function if hc increments in the next round
        val2 = fxp.deltaf_NC * (Pi[y] + UfNC[x+1]) +
               fxp.lambdae * par.sigma * sum(v0 .* dfNC2 .* logit(dfNC2, fxp)) * (SfNC[x+1, y] > 0) +
               fxp.gamma * PfD_hat[x+1, y] + fxp.chi * PfPL_hat[x+1, y] +
               (1 - fxp.deltaf_NC - fxp.gamma - fxp.chi) * PfNC_hat[x+1, y]
        if (y % 7 == 0)
            expectedPfNC = gef[7] * val2 + (1 - gef[7]) * val1
        else
            expectedPfNC = gef[y%7] * val2 + (1 - gef[y%7]) * val1
        end
    end
    #return value at current round and expected value for next perioid
    return (1 - tau[1]) * output[x, y] + epsf[x] * alpha[y] + fxp.beta * expectedPfNC
end

#Female Parental Leave
function getPfPL(x, y,Pi,par,fxp,UfPL,PfD_hat,PfYC_hat,PfPL_hat,epsf,alpha,output)
    val = fxp.deltaf_YC * (Pi[y] + UfPL[x]) +
          fxp.gamma * PfD_hat[x, y] +
          fxp.etaf * PfYC_hat[x, y] +
          (1 - fxp.deltaf_YC - fxp.gamma - fxp.etaf) * PfPL_hat[x, y]

    return fxp.Transfer * output[x, y] + (par.M + epsf[x]) * alpha[y] + fxp.beta * val
end

#Female Young Child
function getPfYC(x, y,v0,Pi,fxp,par,epsf,gef,tau,output,alpha,UfYC,SfYC,PfD_hat,PfPL_hat,PfYC_hat)
    expectedPfYC = 0.0
    dfYC1 = SfYC[x, :] .- SfYC[x, y]
    #Calculate the value function if no hc increment in the next round
    val1 = fxp.deltaf_YC * (Pi[y] + UfYC[x]) +
           fxp.lambdae * par.sigma * sum(v0 .* dfYC1 .* logit(dfYC1, fxp)) * (SfYC[x, y] > 0) +
           fxp.gamma2 * PfD_hat[x, y] + fxp.chi * PfPL_hat[x, y] +
           (1 - fxp.deltaf_YC - fxp.gamma - fxp.chi) * PfYC_hat[x, y]
    #Check if it is possible for HC to grow
    if (x % 7 == 0)
        expectedPfYC = val1
    else
        dfYC2 = SfYC[x+1, :] .- SfYC[x+1, y]
        #Calculate the value function if hc increments in the next round
        val2 = fxp.deltaf_YC * (Pi[y] + UfYC[x+1]) +
               fxp.lambdae * par.sigma * sum(v0 .* dfYC2 .* logit(dfYC2, fxp)) * (SfYC[x+1, y] > 0) +
               fxp.gamma2 * PfD_hat[x+1, y] + fxp.chi * PfPL_hat[x+1, y] +
               (1 - fxp.deltaf_YC - fxp.gamma - fxp.chi) * PfYC_hat[x+1, y]
        if (y % 7 == 0)
            expectedPfYC = gef[7] * val2 + (1 - gef[7]) * val1
        else
            expectedPfYC = gef[y%7] * val2 + (1 - gef[y%7]) * val1
        end
    end
    return (1 - tau[1]) * output[x, y] + (par.M + epsf[x]) * alpha[y] + fxp.beta * expectedPfYC
end

#Female Done with Children
function getPfD(x,y,v0,Pi,fxp,par,output,gef,tau,epsf,alpha,SfD,PfD_hat,UfD)
    expectedPfD = 0.0
    dfD1 = SfD[x, :] .- SfD[x, y]
    #Calculate the value function if no hc increment in the next round
    val1 = fxp.phi * Pi[y] + fxp.delta * (Pi[y] + UfD[x]) +
           (1 - fxp.phi - fxp.delta) * PfD_hat[x, y] +
           fxp.lambdae * par.sigma * sum(v0 .* dfD1 .* logit(dfD1, fxp)) * (SfD[x, y] > 0)
    #Check if it is possible for HC to grow
    if (x % 7 == 0)
        expectedPfD = val1
    else
        dfD2 = SfD[x+1, :] .- SfD[x+1, y]
        #Calculate the value function if hc increments in the next round
        val2 = fxp.phi * Pi[y] + fxp.delta * (Pi[y] + UfD[x+1]) +
               (1 - fxp.phi - fxp.delta) * PfD_hat[x+1, y] +
               fxp.lambdae * par.sigma * sum(v0 .* dfD2 .* logit(dfD2, fxp)) * (SfD[x+1, y] > 0)
        if (y % 7 == 0)
            expectedPfD = gef[7] * val2 + (1 - gef[7]) * val1
        else
            expectedPfD = gef[y%7] * val2 + (1 - gef[y%7]) * val1
        end
    end
    return (1 - tau[1]) * output[x, y] + epsf[x] * alpha[y] + fxp.beta * expectedPfD
end

#Solve for equilibrium
function solve(par, set, fxp, init)
    kk = set.Ny * set.Na;
    ii = set.Nx * set.Ne;

    epsf = repeat([-2, 0, 2] .* fxp.sdf .+ par.muf * fxp.mum, inner = 7)
    epsm = repeat([-2, 0, 2] .* fxp.sdm .+ fxp.mum, inner = 7)
    alpha = repeat(fxp.alpha,inner=7);
    #humanCapital = repeat(fxp.hc, outer = 3);
    #jobProductivityType = repeat(fxp.p,outer=3);
    #jobAmenityType = repeat(fxp.alpha,inner=7);
    output = par.K .* [(par.a * hc^fxp.rho + (1 - par.a) * p^fxp.rho)^(1 / fxp.rho) for hc in repeat(fxp.hc, outer = 3), p in repeat(fxp.p, outer = 3)]
    outputb = repeat([par.b * hc for hc in fxp.hc], outer = 3)
    e = Dict("stages" => nlsolve(x -> solveStages(x, fxp), 12.50 .* ones(typeof(12.50), 8)).zero, "converge" => 1, "Pineg" => 0)
    NC_f = e["stages"][1]
    YC_f = e["stages"][2]
    PL_f = e["stages"][3]
    D_f = e["stages"][4]
    NC_m = e["stages"][5]
    YC_m = e["stages"][6]
    PL_m = e["stages"][7]
    D_m = e["stages"][8]
    gem = par.d1m .+ par.d2m .* fxp.p
    gef = par.d1f .+ par.d2f .* fxp.p

    #unemployed workers u(x), in different stages
    umNC_0 = vec(init.umNC)
    umPL_0 = vec(init.umYC / 2)
    umYC_0 = vec(init.umYC / 2)
    umD_0 = vec(init.umD)
    ufNC_0 = vec(init.ufNC)
    ufPL_0 = vec(init.ufYC / 2)
    ufYC_0 = vec(init.ufYC / 2)
    ufD_0 = vec(init.ufD)
    #employed workers h(x, y, fxp.alpha)
    hmNC_0 = init.hmNC
    hmYC_0 = init.hmYC
    hmPL_0 = init.hmPL
    hmD_0 = init.hmD
    hfNC_0 = init.hfNC
    hfYC_0 = init.hfYC
    hfPL_0 = init.hfPL
    hfD_0 = init.hfD
    #vacancies
    v0 = vec(init.v)
    #Guess Value Functions
    # value in unemployment
    UmNC_0 = vec(init.UmNC)
    UmPL_0 = vec(init.UmYC)
    UmYC_0 = vec(init.UmYC)
    UmD_0 = vec(init.UmD)
    UfNC_0 = vec(init.UfNC)
    UfPL_0 = vec(init.UfYC)
    UfYC_0 = vec(init.UfYC)
    UfD_0 = vec(init.UfD)
    #value in vacancy Π₀(y,α)
    Pi0 = vec(init.Pi)

    #value in surplus
    #Male NC
    SmNC = init.SmNC
    SmNCPos = copy(SmNC)
    SmNCPos[SmNCPos.<0] .= 0
    #Male PL
    SmPL = init.SmPL
    SmPLPos = copy(SmPL)
    SmPLPos[SmPLPos.<0] .= 0
    #Male YC
    SmYC = init.SmYC
    SmYCPos = copy(SmYC)
    SmYCPos[SmYCPos.<0] .= 0
    #Male D
    SmD = init.SmD
    SmDPos = copy(SmD)
    SmDPos[SmDPos.<0] .= 0

    #Female NC
    SfNC = init.SfNC
    SfNCPos = copy(SfNC)
    SfNCPos[SfNCPos.<0] .= 0
    #Female PL
    SfPL = init.SfPL
    SfPLPos = copy(SfPL)
    SfPLPos[SmPLPos.<0] .= 0
    #Female YC
    SfYC = init.SfYC
    SfYCPos = copy(SfYC)
    SfYCPos[SfYCPos.<0] .= 0
    #Female D
    SfD = init.SfD
    SfDPos = copy(SfD)
    SfDPos[SfDPos.<0] .= 0

    #Probably don't need this
    PmNC_0 = SmNC + [x + y for x in UmNC_0, y in Pi0]
    PmPL_0 = SmPL + [x + y for x in UmPL_0, y in Pi0]
    PmYC_0 = SmYC + [x + y for x in UmYC_0, y in Pi0]
    PmD_0 = SmD + [x + y for x in UmD_0, y in Pi0]

    PfNC_0 = SfNC + [x + y for x in UfNC_0, y in Pi0]
    PfPL_0 = SfPL + [x + y for x in UfPL_0, y in Pi0]
    PfYC_0 = SfYC + [x + y for x in UfYC_0, y in Pi0]
    PfD_0 = SfD + [x + y for x in UfD_0, y in Pi0]

    tau = fxp.Transfer * sum(sum(hmPL_0 + hfPL_0, dims = 1) .* output, dims = 1) / (sum(sum(hmNC_0 + hfNC_0 + hmYC_0 + hfYC_0 + hmD_0 + hfD_0, dims = 1) .* output, dims = 1))
    PmNC = zeros(Float64, 21, 21)
    PmPL = zeros(Float64, 21, 21)
    PmYC = zeros(Float64, 21, 21)
    PmD = zeros(Float64, 21, 21)

    PfNC = zeros(Float64, 21, 21)
    PfPL = zeros(Float64, 21, 21)
    PfYC = zeros(Float64, 21, 21)
    PfD = zeros(Float64, 21, 21)

    UmNC = zeros(Float64, 21)
    UmPL = zeros(Float64, 21)
    UmYC = zeros(Float64, 21)
    UmD = zeros(Float64, 21)
    UfNC = zeros(Float64, 21)
    UfPL = zeros(Float64, 21)
    UfYC = zeros(Float64, 21)
    UfD = zeros(Float64, 21)
    effU = sum(umNC_0 + ufNC_0 + fxp.s1 * (umYC_0 + umD_0 + ufYC_0 + ufD_0)) +
           fxp.s2 * sum(sum(hmNC_0 + hmYC_0 + hmD_0 + hfNC_0 + hfYC_0 + hfD_0, dims = 1))
    V = sum(v0)
    fxp = merge(fxp, (lambdau_NC = fxp.theta / (effU * V)^0.5, lambdau_YC = fxp.theta / (effU * V)^0.5 * fxp.s1, lambdae = fxp.theta / (effU * V)^0.5 * fxp.s2))
    maxit = 10000
    change_dist = zeros(maxit, 1)
    tol = 1e-9
    for it1 = 1:maxit
        #Value of Unemployment
        #Male NC
        UmNC = [x + fxp.beta * (fxp.gamma * y + fxp.chi * z +
                                fxp.lambdau_NC * par.sigma * sum(v0 .* SmNCPos[i, :]) +
                                (1 - fxp.gamma - fxp.chi) * m)
                for (x, y, z, m, i) in zip(outputb, UmD_0, UmPL_0, UmNC_0, 1:21)]

        #Male PL
        UmPL = [x + fxp.beta * (fxp.gamma * y + fxp.etam * z +
                                (1 - fxp.gamma - fxp.etam) * m)
                for (x, y, z, m) in zip(outputb, UmD_0, UmYC_0, UmPL_0)]

        #Male YC
        UmYC = [x + fxp.beta * (fxp.gamma * y + fxp.chi * z +
                                fxp.lambdau_YC * par.sigma * sum(v0 .* SmYCPos[i, :]) +
                                (1 - fxp.gamma - fxp.chi) * m)
                for (x, y, z, m, i) in zip(outputb, UmD_0, UmPL_0, UmYC_0, 1:21)]

        #Male D
        UmD = [x + fxp.beta * ((1 - fxp.phi) * y +
                               fxp.lambdau_YC * par.sigma * sum(v0 .* SmDPos[i, :]))
               for (x, y, i) in zip(outputb, UmD_0, 1:21)]

        #Female NC
        UfNC = [x + fxp.beta * (fxp.gamma * y + fxp.chi * z +
                                fxp.lambdau_NC * par.sigma * sum(v0 .* SfNCPos[i, :]) +
                                (1 - fxp.gamma - fxp.chi) * m)
                for (x, y, z, m, i) in zip(outputb, UfD_0, UfPL_0, UfNC_0, 1:21)]

        #Female PL
        UfPL = [x + fxp.beta * (fxp.gamma * y + fxp.etaf * z +
                                (1 - fxp.gamma - fxp.etaf) * m)
                for (x, y, z, m) in zip(outputb, UfD_0, UfYC_0, UfPL_0)]
        #Female YC
        UfYC = [x + fxp.beta * (fxp.gamma * y + fxp.chi * z +
                                fxp.lambdau_YC * par.sigma * sum(v0 .* SfYCPos[i, :]) +
                                (1 - fxp.gamma - fxp.chi) * m)
                for (x, y, z, m, i) in zip(outputb, UfD_0, UfPL_0, UfYC_0, 1:21)]
        #Female D
        UfD = [x + fxp.beta * ((1 - fxp.phi) * y +
                               fxp.lambdau_YC * par.sigma * sum(v0 .* SfDPos[i, :]))
               for (x, y, i) in zip(outputb, UfD_0, 1:21)]

        Pi = zeros(Float64, 21)
        for j = 1:kk
            dmNC = -1 .* SmNC .+ SmNC[:, j]
            dmYC = -1 .* SmYC .+ SmYC[:, j]
            dmD = -1 .* SmD .+ SmD[:, j]

            dfNC = -1 .* SfNC .+ SfNC[:, j]
            dfYC = -1 .* SfYC .+ SfYC[:, j]
            dfD = -1 .* SfD .+ SfD[:, j]

            #Vacancy Value Calculation
            Pi[j] = -fxp.c + fxp.beta * (
                fxp.lambdau_NC * (1 - par.sigma) * sum(umNC_0 .* SmNCPos[:, j]) +
                fxp.lambdau_YC * (1 - par.sigma) * (sum(umYC_0 .* SmYCPos[:, j]) + sum(umD_0 .* SmDPos[:, j])) +
                fxp.lambdau_NC * (1 - par.sigma) * sum(ufNC_0 .* SfNCPos[:, j]) +
                fxp.lambdau_YC * (1 - par.sigma) * (sum(ufYC_0 .* SfYCPos[:, j]) + sum(ufD_0 .* SfDPos[:, j])) +
                fxp.lambdae * (1 - par.sigma) * (
                    sum(hmNC_0 .* dmNC .* logit(dmNC, fxp)) +
                    sum(hmYC_0 .* dmYC .* logit(dmYC, fxp)) +
                    sum(hmD_0 .* dmD .* logit(dmD, fxp)) +
                    sum(hfNC_0 .* dfNC .* logit(dfNC, fxp)) +
                    sum(hfYC_0 .* dfYC .* logit(dfYC, fxp)) +
                    sum(hfD_0 .* dfD .* logit(dfD, fxp))
                )
                + Pi0[j]
            )
        end

        # joint value P
        PmNC_hat = PmNC_0 .* (SmNC .> 0) + [x + y for x in UmNC, y in Pi] .* (SmNC .< 0)
        PmPL_hat = PmPL_0 .* (SmPL .> 0) + [x + y for x in UmPL, y in Pi] .* (SmPL .< 0)
        PmYC_hat = PmYC_0 .* (SmYC .> 0) + [x + y for x in UmYC, y in Pi] .* (SmYC .< 0)
        PmD_hat = PmD_0 .* (SmD .> 0) + [x + y for x in UmD, y in Pi] .* (SmD .< 0)

        PfNC_hat = PfNC_0 .* (SfNC .> 0) + [x + y for x in UfNC, y in Pi] .* (SfNC .< 0)
        PfPL_hat = PfPL_0 .* (SfPL .> 0) + [x + y for x in UfPL, y in Pi] .* (SfPL .< 0)
        PfYC_hat = PfYC_0 .* (SfYC .> 0) + [x + y for x in UfYC, y in Pi] .* (SfYC .< 0)
        PfD_hat = PfD_0 .* (SfD .> 0) + [x + y for x in UfD, y in Pi] .* (SfD .< 0)
        for i in 1:ii
            for j in 1:kk
                PmNC[i, j] = getPmNC(i,j,v0,fxp,par,Pi,gem,SmNC,UmNC,PmD_hat,PmPL_hat,PmNC_hat,epsm,alpha,output,tau)
                PmPL[i, j] = getPmPL(i,j,Pi,fxp,UmPL,PmD_hat,PmYC_hat,PmPL_hat,epsm,alpha,output)
                PmYC[i, j] = getPmYC(i,j,v0,Pi,fxp,par,epsm,gem,tau,output,alpha,UmYC,SmYC,PmD_hat,PmPL_hat,PmYC_hat)
                PmD[i, j] = getPmD(i,j,v0,Pi,fxp,par,output,gem,tau,epsm,alpha,SmD,PmD_hat,UmD)

                PfNC[i, j] = getPfNC(i,j,v0,fxp,par,Pi,gef,SfNC,UfNC,PfD_hat,PfPL_hat,PfNC_hat,epsf,alpha,output,tau)
                PfPL[i, j] = getPfPL(i,j,Pi,par,fxp,UfPL,PfD_hat,PfYC_hat,PfPL_hat,epsf,alpha,output)
                PfYC[i, j] = getPfYC(i,j,v0,Pi,fxp,par,epsf,gef,tau,output,alpha,UfYC,SfYC,PfD_hat,PfPL_hat,PfYC_hat)
                PfD[i, j] = getPfD(i,j,v0,Pi,fxp,par,output,gef,tau,epsf,alpha,SfD,PfD_hat,UfD)
            end
        end

        SmNC_0 = PmNC - [x + y for x in UmNC, y in Pi]
        SmPL_0 = PmPL - [x + y for x in UmPL, y in Pi]
        SmYC_0 = PmYC - [x + y for x in UmYC, y in Pi]
        SmD_0 = PmD - [x + y for x in UmD, y in Pi]

        SfNC_0 = PfNC - [x + y for x in UfNC, y in Pi]
        SfPL_0 = PfPL - [x + y for x in UfPL, y in Pi]
        SfYC_0 = PfYC - [x + y for x in UfYC, y in Pi]
        SfD_0 = PfD - [x + y for x in UfD, y in Pi]

        change1 = sum((SmNC - SmNC_0) .^ 2 + (SmYC - SmYC_0) .^ 2 + (SmPL - SmPL_0) .^ 2 + (SmD - SmD_0) .^ 2 + (SfNC - SfNC_0) .^ 2 + (SfYC - SfYC_0) .^ 2 + (SfPL - SfPL_0) .^ 2 + (SfD - SfD_0) .^ 2)

        if change1 < tol
            println("converged in " * string(it1) * " iterations, with change:" * string(change1))
            break
        end
        UmNC_0 = UmNC
        UmPL_0 = UmPL
        UmYC_0 = UmYC
        UmD_0 = UmD
        UfNC_0 = UfNC
        UfPL_0 = UfPL
        UfYC_0 = UfYC
        UfD_0 = UfD

        Pi0 = Pi

        PmNC_0 = PmNC
        PmYC_0 = PmYC
        PmPL_0 = PmPL
        PmD_0 = PmD
        PfNC_0 = PfNC
        PfYC_0 = PfYC
        PfPL_0 = PfPL
        PfD_0 = PfD

        SmNC = SmNC_0
        SmNCPos = copy(SmNC_0)
        SmNCPos[SmNCPos.<0] .= 0

        SmYC = SmYC_0
        SmYCPos = copy(SmYC)
        SmYCPos[SmYCPos.<0] .= 0

        SmPL = SmPL_0
        SmPLPos = copy(SmPL)
        SmPLPos[SmPLPos.<0] .= 0

        SmD = SmD_0
        SmDPos = copy(SmD)
        SmDPos[SmDPos.<0] .= 0

        SfNC = SfNC_0
        SfNCPos = copy(SfNC)
        SfNCPos[SfNCPos.<0] .= 0

        SfYC = SfYC_0
        SfYCPos = copy(SfYC)
        SfYCPos[SfYCPos.<0] .= 0

        SfPL = SfPL_0
        SmPLPos = copy(SmPL)
        SmPLPos[SmPLPos.<0] .= 0

        SfD = SfD_0
        SfDPos = copy(SfD)
        SfDPos[SfDPos.<0] .= 0
    end
    return PmNC
end

export solve
end


text:
yearmonthday(x::NanoDate) = yearmonthday(days(x.datetime))
yearmonth(x::NanoDate) = yearmonth(days(x.datetime))
monthday(x::NanoDate) = monthday(days(x.datetime))

for P in (:Year, :Quarter, :Month, :Week, :Day, :Hour, :Minute, :Second, :Millisecond)
    p = Symbol(lowercase(String(P)))
    @eval begin
        Dates.$P(nd::NanoDate) = $P(nd.datetime)
        Dates.$p(nd::NanoDate) = $p(nd.datetime)
    end
end

@inline Dates.microsecond(nd::NanoDate) = div(value(nd.nanosecs), 1000)
Dates.Microsecond(nd::NanoDate) = Microsecond(microsecond(nd))

@inline Dates.nanosecond(nd::NanoDate) = rem(value(nd.nanosecs), 1000)
Dates.Nanosecond(nd::NanoDate) = Nanosecond(nanosecond(nd))

ext: jl


text:
export CUSOLVERError

struct CUSOLVERError <: Exception
    code::cusolverStatus_t
end

Base.convert(::Type{cusolverStatus_t}, err::CUSOLVERError) = err.code

Base.showerror(io::IO, err::CUSOLVERError) =
    print(io, "CUSOLVERError: ", description(err), " (code $(reinterpret(Int32, err.code)), $(name(err)))")

name(err::CUSOLVERError) = string(err.code)

## COV_EXCL_START
function description(err)
    if err.code == CUSOLVER_STATUS_SUCCESS
        "the operation completed successfully"
    elseif err.code == CUSOLVER_STATUS_NOT_INITIALIZED
        "the library was not initialized"
    elseif err.code == CUSOLVER_STATUS_ALLOC_FAILED
        "the resource allocation failed"
    elseif err.code == CUSOLVER_STATUS_INVALID_VALUE
        "an invalid value was used as an argument"
    elseif err.code == CUSOLVER_STATUS_ARCH_MISMATCH
        "an absent device architectural feature is required"
    elseif err.code == CUSOLVER_STATUS_EXECUTION_FAILED
        "the GPU program failed to execute"
    elseif err.code == CUSOLVER_STATUS_INTERNAL_ERROR
        "an internal operation failed"
    elseif err.code == CUSOLVER_STATUS_MATRIX_TYPE_NOT_SUPPORTED
        "the matrix type is not supported."
    else
        "no description for this error"
    end
end
## COV_EXCL_STOP

## API call wrapper

# outlined functionality to avoid GC frame allocation
@noinline function throw_api_error(res)
    throw(CUSOLVERError(res))
end

function initialize_api()
    CUDA.initialize_cuda_context()
end

macro check(ex)
    quote
        res = @retry_reclaim isequal(CUSOLVER_STATUS_ALLOC_FAILED) $(esc(ex))
        if res != CUSOLVER_STATUS_SUCCESS
            throw_api_error(res)
        end

        return
    end
end


text:
using DataValues
using Test
using Dates

@testset "Operations" begin

time_a = DataValue("14:00:00")
date_a = DataValue("01/01/2012")

@test DateTime(time_a,"HH:MM:SS") == DataValue(DateTime("14:00:00", "HH:MM:SS"))
@test DateTime(DataValue{String}(), "HH:MM:SS") == DataValue{DateTime}()

@test Date(date_a,"mm/dd/yyyy") == DataValue(Date("01/01/2012", "mm/dd/yyyy"))
@test Date(DataValue{String}(), "mm/dd/yyyy") == DataValue{Date}()

@test abs(DataValue(3.)) == DataValue(3.)
@test abs(DataValue{Float64}()) == DataValue{Float64}()

@test log(DataValue(3.)) == DataValue(log(3.))
@test log(DataValue{Float64}()) == DataValue{Float64}()

@test min(DataValue(3), DataValue(2)) == DataValue(2)
@test min(DataValue(3), DataValue{Int}()) == DataValue{Int}()
@test min(DataValue{Int}(), DataValue(3)) == DataValue{Int}()
@test min(DataValue{Int}(), DataValue{Int}()) == DataValue{Int}()

@test min(3, DataValue(2)) == DataValue(2)
@test min(3, DataValue{Int}()) == DataValue{Int}()
@test min(DataValue{Int}(), 3) == DataValue{Int}()
@test min(DataValue{Int}(2), 3) == DataValue{Int}(2)
@test min(DataValue{Int}(), DataValue{Int}()) == DataValue{Int}()

@test DataValue(3) + DataValue(2.) == DataValue(5.)
@test DataValue(3) + DataValue{Float64}() == DataValue{Float64}()
@test DataValue{Int}() + DataValue(2.) == DataValue{Float64}()
@test DataValue{Int}() + DataValue{Float64}() == DataValue{Float64}()
@test DataValue(3) + NA == DataValue{Int}()
@test NA + DataValue(2.) == DataValue{Float64}()
@test NA + NA == NA

@test 3 + DataValue(2.) == DataValue(5.)
@test 3 + DataValue{Float64}() == DataValue{Float64}()
@test 3 + NA == DataValue{Int}()

@test DataValue(3) + 2. == DataValue(5.)
@test DataValue{Int}() + 2. == DataValue{Float64}()
@test NA + 2. == DataValue{Float64}()

end


text:
# For rubberband, we draw the selection region on the front canvas, and repair
# by copying from the back. Note that the front canvas has
#     user coordinates = device coordinates,
# so device_to_user doesn't do anything. The back canvas has
#     user coordinates = image pixel coordinates,
# so is the correct reference for anything dealing with image pixels.
type RubberBand
    pos1::Vec2
    pos2::Vec2
    moved::Bool
end

function rbdraw(r::GraphicsContext, rb::RubberBand)
    rectangle(r, rb.pos1.x, rb.pos1.y, rb.pos2.x-rb.pos1.x, rb.pos2.y-rb.pos1.y)
    set_line_width(r, 1)
    set_dash(r, [3.0,3.0], 3.0)
    set_source_rgb(r, 1, 1, 1)
    stroke_preserve(r)
    set_dash(r, [3.0,3.0], 0.0)
    set_source_rgb(r, 0, 0, 0)
    stroke_preserve(r)
end

# callback_done is executed when the user finishes drawing the rubberband.
# Its syntax is callback_done(canvas, boundingbox), where the boundingbox is
# in user coordinates.
function rubberband_start(c::Canvas, x, y, callback_done::Function)
    # Copy the surface to another buffer, so we can repaint the areas obscured by the rubberband
    r = getgc(c)
    Graphics.save(r)
    reset_transform(r)
    ctxcopy = copy(r)
    rb = RubberBand(Vec2(x,y), Vec2(x,y), false)
    callbacks_old = (c.mouse.button1motion, c.mouse.button1release)
    c.mouse.button1motion = (c, x, y) -> rubberband_move(c, rb, x, y, ctxcopy)
    c.mouse.button1release = (c, x, y) -> rubberband_stop(c, rb, x, y, ctxcopy, callbacks_old, callback_done)
end

function rubberband_move(c::Canvas, rb::RubberBand, x, y, ctxcopy)
    r = getgc(c)
    if rb.moved
        # Erase the previous rubberband by copying from back surface to front
        set_source(r, ctxcopy)
        # Since the path was already created and preserved, we just modify its properties
        set_line_width(r, 2)
        set_dash(r, Float64[])
        stroke(r)
    end
    rb.moved = true
    # Draw the new rubberband
    rb.pos2 = Vec2(x, y)
    rbdraw(r, rb)
    reveal(c)
    Tk.update()
end

function rubberband_stop(c::Canvas, rb::RubberBand, x, y, ctxcopy, callbacks_old, callback_done)
    c.mouse.button1motion = callbacks_old[1]
    c.mouse.button1release = callbacks_old[2]
    if !rb.moved
        return
    end
    r = getgc(c)
    set_source(r, ctxcopy)
    set_line_width(r, 2)
    stroke(r)
    reveal(c)
    restore(r)
    Tk.update()
    x1, y1 = rb.pos1.x, rb.pos1.y
    if abs(x1-x) > 2 || abs(y1-y) > 2
        # It moved sufficiently, let's execute the callback
        xu, yu = device_to_user(r, x, y)
        x1u, y1u = device_to_user(r, x1, y1)
        zoombb = BoundingBox(min(x1u,xu), max(x1u,xu), min(y1u,yu), max(y1u,yu))
        callback_done(c, zoombb)
    end
end


text:
"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	thermal_commit(EP::Model, inputs::Dict, Reserves::Int)

This function defines the operating constraints for thermal power plants subject to unit commitment constraints on power plant start-ups and shut-down decision ($y \in UC$).

We model capacity investment decisions and commitment and cycling (start-up, shut-down) of thermal generators using the integer clustering technique developed in [Palmintier, 2011](https://pennstate.pure.elsevier.com/en/publications/impact-of-unit-commitment-constraints-on-generation-expansion-pla), [Palmintier, 2013](https://dspace.mit.edu/handle/1721.1/79147), and [Palmintier, 2014](https://ieeexplore.ieee.org/document/6684593). In a typical binary unit commitment formulation, each unit is either on or off. With the clustered unit commitment formulation, one or more cluster(s) of similar generators are clustered by type and zone (typically using heat rate and fixed O\&M cost to create clusters), and the integer commitment state variable for each cluster varies from zero to the number of units in the cluster, $\frac{\Delta^{total}_{y,z}}{\Omega^{size}_{y,z}}$. As discussed in \cite{Palmintier2014}, this approach replaces the large set of binary commitment decisions and associated constraints, which scale directly with the number of individual units, with a smaller set of integer commitment states and  constraints, one for each cluster $y$. The dimensionality of the problem thus scales with the number of units of a given type in each zone, rather than by the number of discrete units, significantly improving computational efficiency. However, this method entails the simplifying assumption that all clustered units have identical parameters (e.g., capacity size, ramp rates, heat rate) and that all committed units in a given time step $t$ are operating at the same power output per unit.

**Power balance expression**

This function adds the sum of power generation from thermal units subject to unit commitment ($\Theta_{y \in UC,t \in T,z \in Z}$) to the power balance expression.

**Startup and shutdown events (thermal plant cycling)**

*Capacitated limits on unit commitment decision variables*

Thermal resources subject to unit commitment ($y \in \mathcal{UC}$) adhere to the following constraints on commitment states, startup events, and shutdown events, which limit each decision to be no greater than the maximum number of discrete units installed (as per the following three constraints):

```math
\begin{aligned}
\nu_{y,z,t} \leq \frac{\Delta^{\text{total}}_{y,z}}{\Omega^{size}_{y,z}}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

```math
\begin{aligned}
\chi_{y,z,t} \leq \frac{\Delta^{\text{total}}_{y,z}}{\Omega^{size}_{y,z}}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

```math
\begin{aligned}
\zeta_{y,z,t} \leq \frac{\Delta^{\text{total}}_{y,z}}{\Omega^{size}_{y,z}}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```
(See Constraints 1-3 in the code)

where decision $\nu_{y,z,t}$ designates the commitment state of generator cluster $y$ in zone $z$ at time $t$, decision $\chi_{y,z,t}$ represents number of startup decisions, decision $\zeta_{y,z,t}$ represents number of shutdown decisions, $\Delta^{\text{total}}_{y,z}$ is the total installed capacity, and parameter $\Omega^{size}_{y,z}$ is the unit size.

*Commitment state constraint linking start-up and shut-down decisions*

Additionally, the following constarint maintains the commitment state variable across time, $\nu_{y,z,t}$, as the sum of the commitment state in the prior, $\nu_{y,z,t-1}$, period plus the number of units started in the current period, $\chi_{y,z,t}$, less the number of units shut down in the current period, $\zeta_{y,z,t}$:

```math
\begin{aligned}
&\nu_{y,z,t} =\nu_{y,z,t-1} + \chi_{y,z,t} - \zeta_{y,z,t}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}^{interior} \\
&\nu_{y,z,t} =\nu_{y,z,t +\tau^{period}-1} + \chi_{y,z,t} - \zeta_{y,z,t}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}^{start}
\end{aligned}
```
(See Constraint 4 in the code)

Like other time-coupling constraints, this constraint wraps around to link the commitment state in the first time step of the year (or each representative period), $t \in \mathcal{T}^{start}$, to the last time step of the year (or each representative period), $t+\tau^{period}-1$.

**Ramping constraints**

Thermal resources subject to unit commitment ($y \in UC$) adhere to the following ramping constraints on hourly changes in power output:

```math
\begin{aligned}
	\Theta_{y,z,t-1} - \Theta_{y,z,t} &\leq  \kappa^{down}_{y,z} \cdot \Omega^{size}_{y,z} \cdot (\nu_{y,z,t} - \chi_{y,z,t}) & \\[6pt]
	\qquad & - \: \rho^{min}_{y,z} \cdot \Omega^{size}_{y,z} \cdot \chi_{y,z,t} & \hspace{0.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}  \\[6pt]
	\qquad & + \: \text{min}( \rho^{max}_{y,z,t}, \text{max}( \rho^{min}_{y,z}, \kappa^{down}_{y,z} ) ) \cdot \Omega^{size}_{y,z} \cdot \zeta_{y,z,t} &
\end{aligned}
```

```math
\begin{aligned}
	\Theta_{y,z,t} - \Theta_{y,z,t-1} &\leq  \kappa^{up}_{y,z} \cdot \Omega^{size}_{y,z} \cdot (\nu_{y,z,t} - \chi_{y,z,t}) & \\[6pt]
	\qquad & + \: \text{min}( \rho^{max}_{y,z,t}, \text{max}( \rho^{min}_{y,z}, \kappa^{up}_{y,z} ) ) \cdot \Omega^{size}_{y,z} \cdot \chi_{y,z,t} & \hspace{0.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T} \\[6pt]
	\qquad & - \: \rho^{min}_{y,z} \cdot \Omega^{size}_{y,z} \cdot \zeta_{y,z,t} &
\end{aligned}
```
(See Constraints 5-6 in the code)

where decision $\Theta_{y,z,t}$ is the energy injected into the grid by technology $y$ in zone $z$ at time $t$, parameter $\kappa_{y,z,t}^{up|down}$ is the maximum ramp-up or ramp-down rate as a percentage of installed capacity, parameter $\rho_{y,z}^{min}$ is the minimum stable power output per unit of installed capacity, and parameter $\rho_{y,z,t}^{max}$ is the maximum available generation per unit of installed capacity. These constraints account for the ramping limits for committed (online) units as well as faster changes in power enabled by units starting or shutting down in the current time step.

**Minimum and maximum power output**

If not modeling regulation and spinning reserves, thermal resources subject to unit commitment adhere to the following constraints that ensure power output does not exceed minimum and maximum feasible levels:

```math
\begin{aligned}
	\Theta_{y,z,t} \geq \rho^{min}_{y,z} \times \Omega^{size}_{y,z} \times \nu_{y,z,t}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

```math
\begin{aligned}
	\Theta_{y,z,t} \leq \rho^{max}_{y,z} \times \Omega^{size}_{y,z} \times \nu_{y,z,t}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

(See Constraints 7-8 the code)

If modeling reserves and regulation, these constraints are replaced by those established in this ```thermal_commit_reserves()```.

**Minimum and maximum up and down time**

Thermal resources subject to unit commitment adhere to the following constraints on the minimum time steps after start-up before a unit can shutdown again (minimum up time) and the minimum time steps after shut-down before a unit can start-up again (minimum down time):

```math
\begin{aligned}
	\nu_{y,z,t} \geq \displaystyle \sum_{\hat{t} = t-\tau^{up}_{y,z}}^t \chi_{y,z,\hat{t}}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

```math
\begin{aligned}
	\frac{\overline{\Delta_{y,z}} + \Omega_{y,z} - \Delta_{y,z}}{\Omega^{size}_{y,z}} -  \nu_{y,z,t} \geq \displaystyle \sum_{\hat{t} = t-\tau^{down}_{y,z}}^t \zeta_{y,z,\hat{t}}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```
(See Constraints 9-10 in the code)

where $\tau_{y,z}^{up|down}$ is the minimum up or down time for units in generating cluster $y$ in zone $z$.

Like with the ramping constraints, the minimum up and down constraint time also wrap around from the start of each time period to the end of each period.
It is recommended that users of GenX must use longer subperiods than the longest min up/down time if modeling UC. Otherwise, the model will report error.
"""
function thermal_commit(EP::Model, inputs::Dict, Reserves::Int)

	println("Thermal (Unit Commitment) Resources Module")

	dfGen = inputs["dfGen"]

	T = inputs["T"]     # Number of time steps (hours)
	Z = inputs["Z"]     # Number of zones
	G = inputs["G"]     # Number of resources

	hours_per_subperiod = inputs["hours_per_subperiod"] #total number of hours per subperiod

	THERM_COMMIT = inputs["THERM_COMMIT"]
	START_SUBPERIODS = inputs["START_SUBPERIODS"]
	INTERIOR_SUBPERIODS = inputs["INTERIOR_SUBPERIODS"]

	### Expressions ###

	## Power Balance Expressions ##
	@expression(EP, ePowerBalanceThermCommit[t=1:T, z=1:Z],
		sum(EP[:vP][y,t] for y in intersect(THERM_COMMIT, dfGen[dfGen[!,:Zone].==z,:][!,:R_ID])))

	EP[:ePowerBalance] += ePowerBalanceThermCommit

	### Constraints ###

	### Capacitated limits on unit commitment decision variables (Constraints #1-3)
	@constraints(EP, begin
		[y in THERM_COMMIT, t=1:T], EP[:vCOMMIT][y,t] <= EP[:eTotalCap][y]/dfGen[!,:Cap_Size][y]
		[y in THERM_COMMIT, t=1:T], EP[:vSTART][y,t] <= EP[:eTotalCap][y]/dfGen[!,:Cap_Size][y]
		[y in THERM_COMMIT, t=1:T], EP[:vSHUT][y,t] <= EP[:eTotalCap][y]/dfGen[!,:Cap_Size][y]
	end)

	# Commitment state constraint linking startup and shutdown decisions (Constraint #4)
	@constraints(EP, begin
		# For Start Hours, links first time step with last time step in subperiod
		[y in THERM_COMMIT, t in START_SUBPERIODS], EP[:vCOMMIT][y,t] == EP[:vCOMMIT][y,(t+hours_per_subperiod-1)] + EP[:vSTART][y,t] - EP[:vSHUT][y,t]
		# For all other hours, links commitment state in hour t with commitment state in prior hour + sum of start up and shut down in current hour
		[y in THERM_COMMIT, t in INTERIOR_SUBPERIODS], EP[:vCOMMIT][y,t] == EP[:vCOMMIT][y,t-1] + EP[:vSTART][y,t] - EP[:vSHUT][y,t]
	end)

	### Maximum ramp up and down between consecutive hours (Constraints #5-6)

	## For Start Hours
	# Links last time step with first time step, ensuring position in hour 1 is within eligible ramp of final hour position
		# rampup constraints
	@constraint(EP,[y in THERM_COMMIT, t in START_SUBPERIODS],
		EP[:vP][y,t]-EP[:vP][y,(t+hours_per_subperiod-1)] <= dfGen[!,:Ramp_Up_Percentage][y]*dfGen[!,:Cap_Size][y]*(EP[:vCOMMIT][y,t]-EP[:vSTART][y,t])
			+ min(inputs["pP_Max"][y,t],max(dfGen[!,:Min_Power][y],dfGen[!,:Ramp_Up_Percentage][y]))*dfGen[!,:Cap_Size][y]*EP[:vSTART][y,t]
			- dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vSHUT][y,t])

		# rampdown constraints
	@constraint(EP,[y in THERM_COMMIT, t in START_SUBPERIODS],
		EP[:vP][y,(t+hours_per_subperiod-1)]-EP[:vP][y,t] <= dfGen[!,:Ramp_Dn_Percentage][y]*dfGen[!,:Cap_Size][y]*(EP[:vCOMMIT][y,t]-EP[:vSTART][y,t])
			- dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vSTART][y,t]
			+ min(inputs["pP_Max"][y,t],max(dfGen[!,:Min_Power][y],dfGen[!,:Ramp_Dn_Percentage][y]))*dfGen[!,:Cap_Size][y]*EP[:vSHUT][y,t])

	## For Interior Hours
		# rampup constraints
	@constraint(EP,[y in THERM_COMMIT, t in INTERIOR_SUBPERIODS],
		EP[:vP][y,t]-EP[:vP][y,t-1] <= dfGen[!,:Ramp_Up_Percentage][y]*dfGen[!,:Cap_Size][y]*(EP[:vCOMMIT][y,t]-EP[:vSTART][y,t])
			+ min(inputs["pP_Max"][y,t],max(dfGen[!,:Min_Power][y],dfGen[!,:Ramp_Up_Percentage][y]))*dfGen[!,:Cap_Size][y]*EP[:vSTART][y,t]
			-dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vSHUT][y,t])

		# rampdown constraints
	@constraint(EP,[y in THERM_COMMIT, t in INTERIOR_SUBPERIODS],
		EP[:vP][y,t-1]-EP[:vP][y,t] <= dfGen[!,:Ramp_Dn_Percentage][y]*dfGen[!,:Cap_Size][y]*(EP[:vCOMMIT][y,t]-EP[:vSTART][y,t])
			-dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vSTART][y,t]
			+min(inputs["pP_Max"][y,t],max(dfGen[!,:Min_Power][y],dfGen[!,:Ramp_Dn_Percentage][y]))*dfGen[!,:Cap_Size][y]*EP[:vSHUT][y,t])

	### Minimum and maximum power output constraints (Constraints #7-8)
	if Reserves == 1
		# If modeling with regulation and reserves, constraints are established by thermal_commit_reserves() function below
		EP = thermal_commit_reserves(EP, inputs)
	else
		@constraints(EP, begin
			# Minimum stable power generated per technology "y" at hour "t" > Min power
			[y in THERM_COMMIT, t=1:T], EP[:vP][y,t] >= dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]

			# Maximum power generated per technology "y" at hour "t" < Max power
			[y in THERM_COMMIT, t=1:T], EP[:vP][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]
		end)
	end

	### Minimum up and down times (Constraints #9-10)
	for y in THERM_COMMIT

		## up time
		Up_Time = Int(floor(dfGen[!,:Up_Time][y]))
		Up_Time_HOURS = [] # Set of hours in the summation term of the maximum up time constraint for the first subperiod of each representative period
		for s in START_SUBPERIODS
			Up_Time_HOURS = union(Up_Time_HOURS, (s+1):(s+Up_Time-1))
		end

		@constraints(EP, begin
			# cUpTimeInterior: Constraint looks back over last n hours, where n = dfGen[!,:Up_Time][y]
			[t in setdiff(INTERIOR_SUBPERIODS,Up_Time_HOURS)], EP[:vCOMMIT][y,t] >= sum(EP[:vSTART][y,e] for e=(t-dfGen[!,:Up_Time][y]):t)

			# cUpTimeWrap: If n is greater than the number of subperiods left in the period, constraint wraps around to first hour of time series
			# cUpTimeWrap constraint equivalant to: sum(EP[:vSTART][y,e] for e=(t-((t%hours_per_subperiod)-1):t))+sum(EP[:vSTART][y,e] for e=(hours_per_subperiod_max-(dfGen[!,:Up_Time][y]-(t%hours_per_subperiod))):hours_per_subperiod_max)
			[t in Up_Time_HOURS], EP[:vCOMMIT][y,t] >= sum(EP[:vSTART][y,e] for e=(t-((t%hours_per_subperiod)-1):t))+sum(EP[:vSTART][y,e] for e=((t+hours_per_subperiod-(t%hours_per_subperiod))-(dfGen[!,:Up_Time][y]-(t%hours_per_subperiod))):(t+hours_per_subperiod-(t%hours_per_subperiod)))

			# cUpTimeStart:
			# NOTE: Expression t+hours_per_subperiod-(t%hours_per_subperiod) is equivalant to "hours_per_subperiod_max"
			[t in START_SUBPERIODS], EP[:vCOMMIT][y,t] >= EP[:vSTART][y,t]+sum(EP[:vSTART][y,e] for e=((t+hours_per_subperiod-1)-(dfGen[!,:Up_Time][y]-1)):(t+hours_per_subperiod-1))
		end)

		## down time
		Down_Time = Int(floor(dfGen[!,:Down_Time][y]))
		Down_Time_HOURS = [] # Set of hours in the summation term of the maximum down time constraint for the first subperiod of each representative period
		for s in START_SUBPERIODS
			Down_Time_HOURS = union(Down_Time_HOURS, (s+1):(s+Down_Time-1))
		end

		# Constraint looks back over last n hours, where n = dfGen[!,:Down_Time][y]
		# TODO: Replace LHS of constraints in this block with eNumPlantsOffline[y,t]
		@constraints(EP, begin
			# cDownTimeInterior: Constraint looks back over last n hours, where n = inputs["pDMS_Time"][y]
			[t in setdiff(INTERIOR_SUBPERIODS,Down_Time_HOURS)], EP[:eTotalCap][y]/dfGen[!,:Cap_Size][y]-EP[:vCOMMIT][y,t] >= sum(EP[:vSHUT][y,e] for e=(t-dfGen[!,:Down_Time][y]):t)

			# cDownTimeWrap: If n is greater than the number of subperiods left in the period, constraint wraps around to first hour of time series
			# cDownTimeWrap constraint equivalant to: EP[:eTotalCap][y]/dfGen[!,:Cap_Size][y]-EP[:vCOMMIT][y,t] >= sum(EP[:vSHUT][y,e] for e=(t-((t%hours_per_subperiod)-1):t))+sum(EP[:vSHUT][y,e] for e=(hours_per_subperiod_max-(dfGen[!,:Down_Time][y]-(t%hours_per_subperiod))):hours_per_subperiod_max)
			[t in Down_Time_HOURS], EP[:eTotalCap][y]/dfGen[!,:Cap_Size][y]-EP[:vCOMMIT][y,t] >= sum(EP[:vSHUT][y,e] for e=(t-((t%hours_per_subperiod)-1):t))+sum(EP[:vSHUT][y,e] for e=((t+hours_per_subperiod-(t%hours_per_subperiod))-(dfGen[!,:Down_Time][y]-(t%hours_per_subperiod))):(t+hours_per_subperiod-(t%hours_per_subperiod)))

			# cDownTimeStart:
			# NOTE: Expression t+hours_per_subperiod-(t%hours_per_subperiod) is equivalant to "hours_per_subperiod_max"
			[t in START_SUBPERIODS], EP[:eTotalCap][y]/dfGen[!,:Cap_Size][y]-EP[:vCOMMIT][y,t]  >= EP[:vSHUT][y,t]+sum(EP[:vSHUT][y,e] for e=((t+hours_per_subperiod-1)-(dfGen[!,:Down_Time][y]-1)):(t+hours_per_subperiod-1))
		end)
	end

	## END Constraints for thermal units subject to integer (discrete) unit commitment decisions

	return EP
end

@doc raw"""
	thermal_commit_reserves(EP::Model, inputs::Dict)

This function is called by the ```thermal_commit()``` function when regulation and reserves constraints are active and defines reserve related constraints for thermal power plants subject to unit commitment constraints on power plant start-ups and shut-down decisions.

**Maximum contributions to frequency regulation and reserves**

When modeling frequency regulation and reserves contributions, thermal units subject to unit commitment adhere to the following constraints which limit the maximum contribution to regulation and reserves in each time step to a specified maximum fraction ($,\upsilon^{rsv}_{y,z}$) of the commitment capacity in that time step ($(\Omega^{size}_{y,z} \cdot \nu_{y,z,t})$):

```math
\begin{aligned}
	f_{y,z,t} \leq \upsilon^{reg}_{y,z} \times \rho^{max}_{y,z,t} (\Omega^{size}_{y,z} \times \nu_{y,z,t}) \hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

```math
\begin{aligned}
	r_{y,z,t} \leq \upsilon^{rsv}_{y,z} \times \rho^{max}_{y,z,t} (\Omega^{size}_{y,z} \times \nu_{y,z,t}) \hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

where $f_{y,z,t}$ is the frequency regulation contribution limited by the maximum regulation contribution $\upsilon^{reg}_{y,z}$, and $r_{y,z,t}$ is the reserves contribution limited by the maximum reserves contribution $\upsilon^{rsv}_{y,z}$. Limits on reserve contributions reflect the maximum ramp rate for the thermal resource in whatever time interval defines the requisite response time for the regulation or reserve products (e.g., 5 mins or 15 mins or 30 mins). These response times differ by system operator and reserve product, and so the user should define these parameters in a self-consistent way for whatever system context they are modeling.

**Minimum and maximum power output**

When modeling frequency regulation and spinning reserves contributions, thermal resources subject to unit commitment adhere to the following constraints that ensure the sum of power output and reserve and/or regulation contributions do not exceed minimum and maximum feasible power output:

```math
\begin{aligned}
	\Theta_{y,z,t} - f_{y,z,t} \geq \rho^{min}_{y,z} \times Omega^{size}_{y,z} \times \nu_{y,z,t}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

```math
\begin{aligned}
	\Theta_{y,z,t} + f_{y,z,t} + r_{y,z,t} \leq \rho^{max}_{y,z,t} \times \Omega^{size}_{y,z} \times \nu_{y,z,t}
	\hspace{1.5cm} \forall y \in \mathcal{UC}, \forall z \in \mathcal{Z}, \forall t \in \mathcal{T}
\end{aligned}
```

Note there are multiple versions of these constraints in the code in order to avoid creation of unecessary constraints and decision variables for thermal units unable to provide regulation and/or reserves contributions due to input parameters (e.g. ```Reg_Max=0``` and/or ```RSV_Max=0```).
"""
function thermal_commit_reserves(EP::Model, inputs::Dict)

	println("Thermal (Unit Commitment) Reserves Module")

	dfGen = inputs["dfGen"]

	T = inputs["T"]     # Number of time steps (hours)

	THERM_COMMIT = inputs["THERM_COMMIT"]

	THERM_COMMIT_REG_RSV = intersect(THERM_COMMIT, inputs["REG"], inputs["RSV"]) # Set of thermal resources with both regulation and spinning reserves

	THERM_COMMIT_REG = intersect(THERM_COMMIT, inputs["REG"]) # Set of thermal resources with regulation reserves
	THERM_COMMIT_RSV = intersect(THERM_COMMIT, inputs["RSV"]) # Set of thermal resources with spinning reserves

	THERM_COMMIT_NO_RES = setdiff(THERM_COMMIT, THERM_COMMIT_REG, THERM_COMMIT_RSV) # Set of thermal resources with no reserves

	THERM_COMMIT_REG_ONLY = setdiff(THERM_COMMIT_REG, THERM_COMMIT_RSV) # Set of thermal resources only with regulation reserves
	THERM_COMMIT_RSV_ONLY = setdiff(THERM_COMMIT_RSV, THERM_COMMIT_REG) # Set of thermal resources only with spinning reserves

	if !isempty(THERM_COMMIT_REG_RSV)
		@constraints(EP, begin
			# Maximum regulation and reserve contributions
			[y in THERM_COMMIT_REG_RSV, t=1:T], EP[:vREG][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Reg_Max][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]
			[y in THERM_COMMIT_REG_RSV, t=1:T], EP[:vRSV][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Rsv_Max][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]

			# Minimum stable power generated per technology "y" at hour "t" and contribution to regulation must be > min power
			[y in THERM_COMMIT_REG_RSV, t=1:T], EP[:vP][y,t]-EP[:vREG][y,t] >= dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]

			# Maximum power generated per technology "y" at hour "t"  and contribution to regulation and reserves up must be < max power
			[y in THERM_COMMIT_REG_RSV, t=1:T], EP[:vP][y,t]+EP[:vREG][y,t]+EP[:vRSV][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]
		end)
	end

	if !isempty(THERM_COMMIT_REG)
		@constraints(EP, begin
			# Maximum regulation and reserve contributions
			[y in THERM_COMMIT_REG, t=1:T], EP[:vREG][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Reg_Max][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]

			# Minimum stable power generated per technology "y" at hour "t" and contribution to regulation must be > min power
			[y in THERM_COMMIT_REG, t=1:T], EP[:vP][y,t]-EP[:vREG][y,t] >= dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]

			# Maximum power generated per technology "y" at hour "t"  and contribution to regulation must be < max power
			[y in THERM_COMMIT_REG, t=1:T], EP[:vP][y,t]+EP[:vREG][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]
		end)
	end

	if !isempty(THERM_COMMIT_RSV)
		@constraints(EP, begin
			# Maximum regulation and reserve contributions
			[y in THERM_COMMIT_RSV, t=1:T], EP[:vRSV][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Rsv_Max][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]

			# Minimum stable power generated per technology "y" at hour "t" must be > min power
			[y in THERM_COMMIT_RSV, t=1:T], EP[:vP][y,t] >= dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]

			# Maximum power generated per technology "y" at hour "t"  and contribution to reserves up must be < max power
			[y in THERM_COMMIT_RSV, t=1:T], EP[:vP][y,t]+EP[:vRSV][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]
		end)
	end

	if !isempty(THERM_COMMIT_NO_RES)
		@constraints(EP, begin
			# Minimum stable power generated per technology "y" at hour "t" > Min power
			[y in THERM_COMMIT_NO_RES, t=1:T], EP[:vP][y,t] >= dfGen[!,:Min_Power][y]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]

			# Maximum power generated per technology "y" at hour "t" < Max power
			[y in THERM_COMMIT_NO_RES, t=1:T], EP[:vP][y,t] <= inputs["pP_Max"][y,t]*dfGen[!,:Cap_Size][y]*EP[:vCOMMIT][y,t]
		end)
	end

	return EP
end

max_stars_repo_name: Betristor/LTESOM

max_issues_repo_name: Betristor/LTESOM

max_forks_repo_name: Betristor/LTESOM


text:
using LinearAlgebra

function linear_h(θ, x)
    # linear_h: The linear hypothesis regressor.
    #
    # :param θ: parameters for our linear regressor; shape (1, features)
    # :type θ: np.ndarray
    # :param x: input that model is predicting; shape (samples, features)
    # :type x: np.ndarray
    # :return: The predictions of our model on inputs X; shape (samples, 1)
    # :rtype: np.ndarray
    #
    θ .⋅ x #Computes the dot product
end

function linear_grad_h(θ, x)
    # linear_h: The gradient of the linear hypothesis regressor.
    #
    # :param θ: parameters for our linear regressor; shape (1, features)
    # :type θ: np.ndarray
    # :param x: input that model is predicting; shape (samples, features)
    # :type x: np.ndarray
    # :return: The gradient of our linear regressor; shape (samples, features)
    # :rtype: np.ndarray
    #
    x
end

function parabolic_h(θ, x)
    # parabolic_h: The parabolic hypothesis regressor.
    #
    # :param θ: parameters for our parabolic regressor; shape (1, features)
    # :type θ: np.ndarray
    # :param x: input that model is predicting; shape (samples, features)
    # :type x: np.ndarray
    # :return: The predictions of our model on inputs X; shape (samples, 1)
    # :rtype: np.ndarray
    #
    θ ⋅ x.^2
end

function parabolic_grad_h(θ, x)
    # parabolic_grad_h: The gradient of the parabolic hypothesis regressor.
    #
    # :param θ: parameters for our parabolic regressor; shape (1, features)
    # :type θ: np.ndarray
    # :param x: input that model is predicting; shape is (samples, features)
    # :type x: np.ndarray
    # :return: The gradient of our parabolic regressor; shape (samples, features)
    # :rtype: np.ndarray
    #
    x.^2
end

# Add your own hypotheses if you want

function loss_f1(h, θ, x, y)
    # loss_f1 returns the loss for special function f1.
    #
    # This function is for demonstration purposes, since it ignores
    # data points x and y.
    #
    # :param h: hypothesis function that is being used
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param θ: The parameters for our model, must be of shape (2,)
    # :type θ: np.ndarray of shape (-1, 2)
    # :param x: A matrix of samples and their respective features.
    # :type x: np.ndarray of shape (samples, features)
    # :param y: The expected targets our model is attempting to match
    # :type y: np.ndarray of shape (samples,)
    # :return: Return the function evaluation of θ, x, y
    # :rtype: int or np.ndarray of shape (θ.shape[1],)
    #
    θ = reshape(θ, :, 2)
    w1 = θ[:, 1]
    w2 = θ[:, 2]
    @. (
        -2 * exp(-((w1 - 1) * (w1 - 1) + w2 * w2) / 0.2)
        + -3 * exp(-((w1 + 1) * (w1 + 1) + y * y) / 0.2)
        + w1 * w1
        + w2 * w2
    )
end

#TODO: Does not work, Python implementation doesn't work either
function grad_loss_f1(h, grad_h, θ, x, y)
    # grad_loss_f1 returns the gradients for the loss of the f1 function.
    #
    # This function is for demonstration purposes, since it ignores
    # data points x and y.
    #
    # :param h: The hypothesis function that predicts our output given weights
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: The gradient function of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param θ: The parameters for our model.
    # :type θ: np.ndarray of shape (-1, 2)
    # :param x: A matrix of samples and their respective features.
    # :type x: np.ndarray of shape (samples, features)
    # :param y: The expected targets our model is attempting to match
    # :type y: np.ndarray of shape (samples,)
    # :return: gradients for the loss function along the two axes
    # :rtype: np.ndarray
    #
    θ = reshape(θ, :, 2)
    w1 = θ[:, 1]
    w2 = θ[:, 2]
    step = 1e-7
    grad_w1 = (loss_f1(h, θ, w1 .+ step, w2) .- loss_f1(h, θ, w1, w2)) ./ step
    grad_w2 = (loss_f1(h, θ, w1, w2 .+ step) .- loss_f1(h, θ, w1, y)) ./ step
    hcat(grad_w1, grad_w2)
end

function l2_loss(h, grad_h, θ, x, y)
    # l2_loss: standard l2 loss.
    #
    # The l2 loss is defined as (h(x) - y)^2. This is usually used for linear
    # regression in the sum of squares.
    #
    # :param h: hypothesis function that models our data (x) using θ
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: function for the gradient of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param θ: The parameters of our hypothesis function of shape (1, features)
    # :type θ: np.ndarray
    # :param x: Input matrix of shape (samples, features)
    # :type x: np.ndarray
    # :param y: The expected targets our model is attempting to match, of shape (samples, 1)
    # :type y: np.ndarray
    # :return: The l2 loss value
    # :rtype: float
    #
    sum((h(θ, x) .- y).^2)
end

function grad_l2_loss(h, grad_h, θ, x, y)
    # grad_l2_loss: The gradient of the standard l2 loss.
    #
    # The gradient of l2 loss is given by d/dx[(h(x) - y)^2] which is
    # evaluated to 2*(h(x) - y)*h'(x).
    #
    # :param h: hypothesis function that models our data (x) using θ
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: function for the gradient of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param θ: The parameters of our hypothesis fucntion
    # :type θ: np.ndarray
    # :param x: Input matrix of shape (samples, features)
    # :type x: np.ndarray
    # :param y: The expected targets our model is attempting to match, of shape (samples, 1)
    # :type y: np.ndarray
    # :return: The l2 loss gradient of shape (1, features)
    # :rtype: np.ndarray
    #
    [sum(2 .* (h(θ, x) .- y) .* grad_h(θ, x))]
end

# ============================================================================
# YOUR CODE GOES HERE:
# ============================================================================

function grad_descent(h, grad_h, loss_f, grad_loss_f, x, y, steps)
    # grad_descent: gradient descent algorithm on a hypothesis class.
    #
    # This does not use the matrix operations from numpy, this function
    # uses the brute force calculations
    #
    # :param h: hypothesis function that models our data (x) using θ
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: function for the gradient of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param loss_f: loss function that we will be optimizing on
    # :type loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param grad_loss_f: the gradient of the loss function we are optimizing
    # :type grad_loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param x: Input matrix of shape (samples, features)
    # :type x: np.ndarray
    # :param y: The expected targets our model is attempting to match, of shape (samples, 1)
    # :param y: np.ndarray
    # :param steps: number of steps to take in the gradient descent algorithm
    # :type steps: int
    # :return: Ideal weights of shape (1, features), and the list of weights through time
    # :rtype: tuple[np.ndarray, np.ndarray]
    #
    θ = rand(Float64, 1, size(x)[2])
    history = [θ]
    for t in 1:steps
        grad = grad_loss_f(h, grad_h, θ, x, y)
        for i in 1:size(θ)[2]
            θ[i] -= 0.001*grad[i]
        end
        push!(history, θ)
    end
    θ, history
end

function stochastic_grad_descent(h, grad_h, loss_f, grad_loss_f, x, y, steps)
    # grad_descent: gradient descent algorithm on a hypothesis class.
    #
    # This does not use the matrix operations from numpy, this function
    # uses the brute force calculations
    #
    # :param h: hypothesis function that models our data (x) using θ
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: function for the gradient of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param loss_f: loss function that we will be optimizing on
    # :type loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param grad_loss_f: the gradient of the loss function we are optimizing
    # :type grad_loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param x: Input matrix of shape (samples, features)
    # :type x: np.ndarray
    # :param y: The expected targets our model is attempting to match, of shape (samples, 1)
    # :param y: np.ndarray
    # :param steps: number of steps to take in the gradient descent algorithm
    # :type steps: int
    # :return: Ideal weights of shape (1, features), and the list of weights through time
    # :rtype: tuple[np.ndarray, np.ndarray]
    #
    θ = rand(Float64, 1, size(x)[2])
    history = [θ]
    for t in 1:steps
        ix = rand((1:length(x)))
        grad = grad_loss_f(h, grad_h, θ, x[ix], y[ix])
        for i in 1:size(θ)[2]
            θ[i] -= .01*grad[i]
        end
        push!(history, θ)
    end
    θ, history
end

function minibatch_grad_descent(h, grad_h, loss_f, grad_loss_f, x, y, steps, batch_size=10)
    # grad_descent: gradient descent algorithm on a hypothesis class.
    #
    # This does not use the matrix operations from numpy, this function
    # uses the brute force calculations
    #
    # :param h: hypothesis function that models our data (x) using θ
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: function for the gradient of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param loss_f: loss function that we will be optimizing on
    # :type loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param grad_loss_f: the gradient of the loss function we are optimizing
    # :type grad_loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param x: Input matrix of shape (samples, features)
    # :type x: np.ndarray
    # :param y: The expected targets our model is attempting to match, of shape (samples, 1)
    # :param y: np.ndarray
    # :param steps: number of steps to take in the gradient descent algorithm
    # :type steps: int
    # :return: Ideal weights of shape (1, features), and the list of weights through time
    # :rtype: tuple[np.ndarray, np.ndarray]
    #
    θ = rand(Float64, 1, size(x)[2])
    history = [θ]
    for t in 1:steps
        minibatch = rand((1,length(x)), batch_size)
        grad = grad_loss_f(h, grad_h, θ, x[minibatch], y[minibatch])
        for i in 1:size(θ)[2]
            θ[i] -= .001*grad[i]
        end
        push!(history, θ)
    end
    θ, history
end

function matrix_gd(h, grad_h, loss_f, grad_loss_f, x, y, steps)
    # grad_descent: gradient descent algorithm on a hypothesis class.
    #
    # This does not use the matrix operations from numpy, this function
    # uses the brute force calculations
    #
    # :param h: hypothesis function that models our data (x) using θ
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: function for the gradient of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param loss_f: loss function that we will be optimizing on
    # :type loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param grad_loss_f: the gradient of the loss function we are optimizing
    # :type grad_loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param x: Input matrix of shape (samples, features)
    # :type x: np.ndarray
    # :param y: The expected targets our model is attempting to match, of shape (samples, 1)
    # :param y: np.ndarray
    # :param steps: number of steps to take in the gradient descent algorithm
    # :type steps: int
    # :return: Ideal weights of shape (1, features), and the list of weights through time
    # :rtype: tuple[np.ndarray, np.ndarray]
    #
    θ = rand(Float64, 1, size(x)[2])
    history = [θ]
    for t in 1:steps
        grad = grad_loss_f(h, grad_h, θ, x, y)
        θ .-= 0.001.*grad
        push!(history, θ)
    end
    θ, history
end

function matrix_sgd(h, grad_h, loss_f, grad_loss_f, x, y, steps)
    # grad_descent: gradient descent algorithm on a hypothesis class.
    #
    # This does not use the matrix operations from numpy, this function
    # uses the brute force calculations
    #
    # :param h: hypothesis function that models our data (x) using θ
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: function for the gradient of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param loss_f: loss function that we will be optimizing on
    # :type loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param grad_loss_f: the gradient of the loss function we are optimizing
    # :type grad_loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param x: Input matrix of shape (samples, features)
    # :type x: np.ndarray
    # :param y: The expected targets our model is attempting to match, of shape (samples, 1)
    # :param y: np.ndarray
    # :param steps: number of steps to take in the gradient descent algorithm
    # :type steps: int
    # :return: Ideal weights of shape (1, features), and the list of weights through time
    # :rtype: tuple[np.ndarray, np.ndarray]
    #
    θ = rand(Float64, 1, size(x)[2])
    history = [θ]
    for t in 1:steps
        ix = rand((1:length(x)))
        grad = grad_loss_f(h, grad_h, θ, x[ix], y[ix])
        θ .-= .01*grad
        push!(history, θ)
    end
    θ, history
end

function matrix_minibatch_gd(h, grad_h, loss_f, grad_loss_f, x, y, steps, batch_size=10)
    # grad_descent: gradient descent algorithm on a hypothesis class.
    #
    # This does not use the matrix operations from numpy, this function
    # uses the brute force calculations
    #
    # :param h: hypothesis function that models our data (x) using θ
    # :type h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param grad_h: function for the gradient of our hypothesis function
    # :type grad_h: typing.Callable[[np.ndarray, np.ndarray], np.ndarray]
    # :param loss_f: loss function that we will be optimizing on
    # :type loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param grad_loss_f: the gradient of the loss function we are optimizing
    # :type grad_loss_f: typing.Callable[
    #     [
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     typing.Callable[[np.ndarray, np.ndarray], np.ndarray],
    #     np.ndarray,
    #     np.ndarray,
    #     np.ndarray
    #     ],
    #     np.ndarray]
    # :param x: Input matrix of shape (samples, features)
    # :type x: np.ndarray
    # :param y: The expected targets our model is attempting to match, of shape (samples, 1)
    # :param y: np.ndarray
    # :param steps: number of steps to take in the gradient descent algorithm
    # :type steps: int
    # :return: Ideal weights of shape (1, features), and the list of weights through time
    # :rtype: tuple[np.ndarray, np.ndarray]
    #
    θ = rand(Float64, 1, size(x)[2])
    history = [θ]
    for t in 1:steps
        minibatch = rand((1,length(x)), batch_size)
        grad = grad_loss_f(h, grad_h, θ, x[minibatch], y[minibatch])
        θ .-= .001*grad
        push!(history, θ)
    end
    θ, history
end

# ============================================================================
# TESTING FUNCTIONS:
# ============================================================================
function test_gd(grad_des_f)
    x = reshape([-3:0.1:4 ...], :, 1)
    y = 2 .* x
    x_support = [0, 4]
    y_support = [-0.1, 200]
    steps = 500

    θ, hist = grad_des_f(linear_h, linear_grad_h, l2_loss, grad_l2_loss, x, y, steps)
    1.99 .< θ .< 2.01
end

macro test(func)
    append!(results, test_gd(@eval $func))
end

results = []
@test grad_descent
@test stochastic_grad_descent
@test minibatch_grad_descent
@test matrix_gd
@test matrix_sgd
@test matrix_minibatch_gd

results

max_stars_repo_name: EliasLittle/IntSys-Education

max_issues_repo_name: EliasLittle/IntSys-Education

max_forks_repo_name: EliasLittle/IntSys-Education


text:
# ---------------------------------------------------------------------------- #
#
#   ioSTL.jl
#
#   STL i/o routines
#
#   intrico
#   Spring 2018
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

"""
    writeFacetSTLB( vert::SVector{3,SVector{3,Float64}}, fid::IOStream )

Writes one STL facet for binary file.
"""
function writeFacetSTLB( vert::SVector{3,SVector{3,Float64}}, nfac::Vector{Int64},
                         fid::IOStream, edgWrite::Vector{Bool} )

    vec1   = vert[2] - vert[1]
    vec2   = vert[3] - vert[1]
    vec1   = cross( vec1, vec2 )
    normal = vec1 / norm(vec1)

    # normals
    write( fid, Float32( normal[1] ) )
    write( fid, Float32( normal[2] ) )
    write( fid, Float32( normal[3] ) )

    # Vertex 1
    write( fid, Float32( vert[1][1] ) )
    write( fid, Float32( vert[1][2] ) )
    write( fid, Float32( vert[1][3] ) )

    # Vertex 2
    write( fid, Float32( vert[2][1] ) )
    write( fid, Float32( vert[2][2] ) )
    write( fid, Float32( vert[2][3] ) )

    # Vertex 3
    write( fid, Float32( vert[3][1] ) )
    write( fid, Float32( vert[3][2] ) )
    write( fid, Float32( vert[3][3] ) )

    # End facet
    write( fid, UInt16(0) )

    nfac[1] += 1

end

function writeFacetSTLB( vert::SVector{3,SVector{3,Float64}}, nfac::Vector{Int64},
                         fid::Vector{IOStream}, edgWrite::Vector{Bool} )

    for jj in 1:length( fid )

        if edgWrite[jj]
            writeFacetSTLB( vert, [0], fid[jj], edgWrite )
            nfac[jj] += 1
        end

    end

end

"""
    writeFacetSTLB( vert::SVector{3,SVector{3,Float64}}, fid::IOStream )

Writes one STL facet for binary file.
"""
function writeFacetSTLB( vert::SVector{3,SVector{3,Float64}}, fid::IOStream )

    vec1   = vert[2] - vert[1]
    vec2   = vert[3] - vert[1]
    vec1   = cross( vec1, vec2 )
    normal = vec1 / norm(vec1)

    # normals
    write( fid, Float32( normal[1] ) )
    write( fid, Float32( normal[2] ) )
    write( fid, Float32( normal[3] ) )

    # Vertex 1
    write( fid, Float32( vert[1][1] ) )
    write( fid, Float32( vert[1][2] ) )
    write( fid, Float32( vert[1][3] ) )

    # Vertex 2
    write( fid, Float32( vert[2][1] ) )
    write( fid, Float32( vert[2][2] ) )
    write( fid, Float32( vert[2][3] ) )

    # Vertex 3
    write( fid, Float32( vert[3][1] ) )
    write( fid, Float32( vert[3][2] ) )
    write( fid, Float32( vert[3][3] ) )

    # End facet
    write( fid, UInt16(0) )

end

"""
    writeFacetSTLB( vert::SVector{3,SVector{3,Float64}}, fid::IOStream )

Writes one STL facet for binary file.
"""
function writeFacetSTLB( vert::SVector{3,SVector{3,Float64}},
                         normal::SVector{3,Float64}, fid::IOStream )

    # normals
    write( fid, Float32( normal[1] ) )
    write( fid, Float32( normal[2] ) )
    write( fid, Float32( normal[3] ) )

    # Vertex 1
    write( fid, Float32( vert[1][1] ) )
    write( fid, Float32( vert[1][2] ) )
    write( fid, Float32( vert[1][3] ) )

    # Vertex 2
    write( fid, Float32( vert[2][1] ) )
    write( fid, Float32( vert[2][2] ) )
    write( fid, Float32( vert[2][3] ) )

    # Vertex 3
    write( fid, Float32( vert[3][1] ) )
    write( fid, Float32( vert[3][2] ) )
    write( fid, Float32( vert[3][3] ) )

    # End facet
    write( fid, UInt16(0) )

end


text:
using DataDeps

register(DataDep(
  "vote-for-clinton",
  "http://www.stat.ufl.edu/~winner/data/clinton1.txt",
  "http://www.stat.ufl.edu/~winner/data/clinton1.dat",
  "51ae0be6fbf96ce4f9b36684a036c6e6c144549bb673ccd4ff045696ed6d9382",
  post_fetch_method=(path -> begin
    # IDs have spaces in them which is the delimiter
    lines = readlines(path)
    map!(x -> lstrip(x[30:end]), lines, lines)
    write(path, join(lines, '\n'))

    UCIData.process_dataset(path,
        target_index=1,
        feature_indices=2:10,
        delim=' ',
        ignorerepeated=true,
    )
  end),
))

max_stars_repo_name: JackDunnNZ/uci-data

max_issues_repo_name: JackDunnNZ/uci-data

max_forks_repo_name: JackDunnNZ/uci-data


text:
using Oceananigans.Grids: Flat

#####
##### Base difference operators
#####

@inline δxᶜᵃᵃ(i, j, k, grid, u) = @inbounds u[i+1, j, k] - u[i,   j, k]
@inline δxᶠᵃᵃ(i, j, k, grid, c) = @inbounds c[i,   j, k] - c[i-1, j, k]

@inline δyᵃᶜᵃ(i, j, k, grid, v) = @inbounds v[i, j+1, k] - v[i, j,   k]
@inline δyᵃᶠᵃ(i, j, k, grid, c) = @inbounds c[i, j,   k] - c[i, j-1, k]

@inline δzᵃᵃᶜ(i, j, k, grid, w) = @inbounds w[i, j, k+1] - w[i, j,   k]
@inline δzᵃᵃᶠ(i, j, k, grid, c) = @inbounds c[i, j,   k] - c[i, j, k-1]

#####
##### Difference operators acting on functions
#####

@inline δxᶜᵃᵃ(i, j, k, grid, f::F, args...) where F<:Function = f(i+1, j, k, grid, args...) - f(i,   j, k, grid, args...)
@inline δxᶠᵃᵃ(i, j, k, grid, f::F, args...) where F<:Function = f(i,   j, k, grid, args...) - f(i-1, j, k, grid, args...)

@inline δyᵃᶜᵃ(i, j, k, grid, f::F, args...) where F<:Function = f(i, j+1, k, grid, args...) - f(i, j,   k, grid, args...)
@inline δyᵃᶠᵃ(i, j, k, grid, f::F, args...) where F<:Function = f(i, j,   k, grid, args...) - f(i, j-1, k, grid, args...)

@inline δzᵃᵃᶜ(i, j, k, grid, f::F, args...) where F<:Function = f(i, j, k+1, grid, args...) - f(i, j, k,   grid, args...)
@inline δzᵃᵃᶠ(i, j, k, grid, f::F, args...) where F<:Function = f(i, j, k,   grid, args...) - f(i, j, k-1, grid, args...)

#####
##### Support for Flat Earths
#####

@inline δxᶜᵃᵃ(i, j, k, grid::AG{FT, Flat, TY, TZ}, u) where {FT, TY, TZ} = zero(FT)
@inline δxᶠᵃᵃ(i, j, k, grid::AG{FT, Flat, TY, TZ}, c) where {FT, TY, TZ} = zero(FT)

@inline δyᵃᶜᵃ(i, j, k, grid::AG{FT, TX, Flat, TZ}, v) where {FT, TX, TZ} = zero(FT)
@inline δyᵃᶠᵃ(i, j, k, grid::AG{FT, TX, Flat, TZ}, c) where {FT, TX, TZ} = zero(FT)

@inline δzᵃᵃᶜ(i, j, k, grid::AG{FT, TX, TY, Flat}, w) where {FT, TX, TY} = zero(FT)
@inline δzᵃᵃᶠ(i, j, k, grid::AG{FT, TX, TY, Flat}, c) where {FT, TX, TY} = zero(FT)

@inline δxᶜᵃᵃ(i, j, k, grid::AG{FT, Flat, TY, TZ}, f::F, args...) where {FT, TY, TZ, F<:Function} = zero(FT)
@inline δxᶠᵃᵃ(i, j, k, grid::AG{FT, Flat, TY, TZ}, f::F, args...) where {FT, TY, TZ, F<:Function} = zero(FT)

@inline δyᵃᶜᵃ(i, j, k, grid::AG{FT, TX, Flat, TZ}, f::F, args...) where {FT, TX, TZ, F<:Function} = zero(FT)
@inline δyᵃᶠᵃ(i, j, k, grid::AG{FT, TX, Flat, TZ}, f::F, args...) where {FT, TX, TZ, F<:Function} = zero(FT)

@inline δzᵃᵃᶜ(i, j, k, grid::AG{FT, TX, TY, Flat}, f::F, args...) where {FT, TX, TY, F<:Function} = zero(FT)
@inline δzᵃᵃᶠ(i, j, k, grid::AG{FT, TX, TY, Flat}, f::F, args...) where {FT, TX, TY, F<:Function} = zero(FT)


text:
using Base.LinAlg.BLAS
using Distributions
import StatsBase.fit
import StatsBase.coef
import StatsBase: sample, sample!

typealias Mat{T} AbstractArray{T, 2}
typealias Vec{T} AbstractArray{T, 1}
typealias Gaussian Normal

"""
Distribution with a single possible value. Used e.g. during sampling
to provide stable result equal to provided means:

sample(Degenerate, means) = means
"""
type Degenerate <: Distribution{Distributions.Univariate,
                               Distributions.Discrete}
end

## types

@runonce begin

    # TODO: why these dosctrings are not attached to these types?
    """
    Base type for all RBMs. Takes type parameters:

     * T - type of RBM parameters (weights, biases, etc.), input and output.
       By default, Float64 is used
     * V - type of visible units
     * H - type of hidden units
    """
    abstract AbstractRBM{T,V,H}

    """
    Restricted Boltzmann Machine, parametrized by element type T, visible
    unit type V and hidden unit type H.
    """
    type RBM{T,V,H} <: AbstractRBM{T,V,H}
        W::Matrix{T}         # matrix of weights between vis and hid vars
        vbias::Vector{T}     # biases for visible variables
        hbias::Vector{T}     # biases for hidden variables
    end

end

"""
Construct RBM. Parameters:

 * T - type of RBM parameters (e.g. weights and biases; by default, Float64)
 * V - type of visible units
 * H - type of hidden units
 * n_vis - number of visible units
 * n_hid - number of hidden units

Optional parameters:

 * sigma - variance to use during parameter initialization

"""
function RBM(T::Type, V::Type, H::Type,
             n_vis::Int, n_hid::Int; sigma=0.01)
    RBM{T,V,H}(map(T, rand(Normal(0, sigma), (n_hid, n_vis))),
             zeros(n_vis), zeros(n_hid))
end

RBM(V::Type, H::Type, n_vis::Int, n_hid::Int; sigma=0.01) =
    RBM(Float64, V, H, n_vis, n_hid; sigma=sigma)

# some well-known RBM kinds

"""Same as RBM{Float64,Degenerate,Bernoulli}"""
BernoulliRBM(n_vis::Int, n_hid::Int; sigma=0.01) =
    RBM(Float64, Degenerate, Bernoulli, n_vis, n_hid; sigma=sigma)

"""Same as RBM{Float64,Gaussian,Bernoulli}"""
GRBM(n_vis::Int, n_hid::Int; sigma=0.01) =
    RBM(Float64, Normal, Bernoulli, n_vis, n_hid; sigma=sigma)

function Base.show{T,V,H}(io::IO, rbm::RBM{T,V,H})
    n_vis = size(rbm.vbias, 1)
    n_hid = size(rbm.hbias, 1)
    print(io, "RBM{$V,$H}($n_vis, $n_hid)")
end

## utils

function hid_means{T}(rbm::RBM, vis::Mat{T})
    p = rbm.W * vis .+ rbm.hbias
    return logistic(p)
end

function vis_means{T}(rbm::RBM, hid::Mat{T})
    p = rbm.W' * hid .+ rbm.vbias
    return logistic(p)
end

## samping

function sample{T}(::Type{Degenerate}, means::Mat{T})
    return means
end

function sample{T}(::Type{Bernoulli}, means::Mat{T})
    return map(T, float((rand(size(means)) .< means)))
end

function sample{T}(::Type{Gaussian}, means::Mat{T})
    sigma2 = 1                   # using fixed standard diviation
    samples = zeros(T,size(means))
    for j=1:size(means, 2), i=1:size(means, 1)
        samples[i, j] = T(rand(Normal(means[i, j], sigma2)))
    end
    return samples
end

function sample_hiddens{T,V,H}(rbm::AbstractRBM{T,V,H}, vis::Mat)
    means = hid_means(rbm, vis)
    return sample(H, means)
end

function sample_visibles{T,V,H}(rbm::AbstractRBM{T,V,H}, hid::Mat)
    means = vis_means(rbm, hid)
    return sample(V, means)
end

function gibbs(rbm::AbstractRBM, vis::Mat; n_times=1)
    v_pos = vis
    h_pos = sample_hiddens(rbm, v_pos)
    v_neg = sample_visibles(rbm, h_pos)
    h_neg = sample_hiddens(rbm, v_neg)
    for i=1:n_times-1
        v_neg = sample_visibles(rbm, h_neg)
        h_neg = sample_hiddens(rbm, v_neg)
    end
    return v_pos, h_pos, v_neg, h_neg
end

## scoring

function free_energy(rbm::RBM, vis::Mat)
    vb = sum(vis .* rbm.vbias, 1)

    fe_exp = 1 + exp(rbm.W * vis .+ rbm.hbias)
    tofinite!(fe_exp; nozeros=true)

    Wx_b_log = sum(log(fe_exp), 1)
    result = - vb - Wx_b_log

    return result
end

function score_samples(rbm::AbstractRBM, vis::Mat;
                          sample_size=10000)
    if issparse(vis)
        # sparse matrices may be infeasible for this operation
        # so using only little sample
        cols = rand(1:size(vis, 2), sample_size)
        vis = full(vis[:, cols])
    end
    n_feat, n_samples = size(vis)
    vis_corrupted = copy(vis)
    idxs = rand(1:n_feat, n_samples)
    for (i, j) in zip(idxs, 1:n_samples)
        vis_corrupted[i, j] = 1 - vis_corrupted[i, j]
    end

    fe = free_energy(rbm, vis)
    fe_corrupted = free_energy(rbm, vis_corrupted)
    fe_diff = fe_corrupted - fe
    tofinite!(fe_diff; nozeros=true)
    score_row =  n_feat * log(logistic(fe_diff))

    result = map(Float64, squeeze(score_row', 2))
    tofinite!(result)

    return result
end

function pseudo_likelihood(rbm::AbstractRBM, X)
    return mean(score_samples(rbm, X))
end

## gradient calculation

"""
Contrastive divergence sampler. Options:

 * n_gibbs - number of gibbs sampling loops (default: 1)
"""
function contdiv(rbm::AbstractRBM, vis::Mat, ctx::Dict)
    n_gibbs = @get(ctx, :n_gibbs, 1)
    v_pos, h_pos, v_neg, h_neg = gibbs(rbm, vis, n_times=n_gibbs)
    return v_pos, h_pos, v_neg, h_neg
end

"""
Persistent contrastive divergence sampler. Options:

 * n_gibbs - number of gibbs sampling loops
"""
function persistent_contdiv(rbm::AbstractRBM, vis::Mat, ctx::Dict)
    n_gibbs = @get(ctx, :n_gibbs, 1)
    persistent_chain = @get_array(ctx, :persistent_chain, size(vis), vis)
    if size(persistent_chain) != size(vis)
        # persistent_chain not initialized or batch size changed
        # re-initialize
        persistent_chain = vis
    end
    # take positive samples from real data
    v_pos, h_pos, _, _ = gibbs(rbm, vis)
    # take negative samples from "fantasy particles"
    persistent_chain, _, v_neg, h_neg = gibbs(rbm, vis, n_times=n_gibbs)
    return v_pos, h_pos, v_neg, h_neg
end

"""
Function for calculating gradient of negative log-likelihood of the data.
Options:

 * :sampler - sampler to use (default: persistent_contdiv)

Returns:

 * (dW, db, dc) - tuple of gradients for weights, visible and hidden biases,
                  respectively
"""
function gradient_classic{T}(rbm::RBM, vis::Mat{T}, ctx::Dict)
    sampler = @get_or_create(ctx, :sampler, persistent_contdiv)
    v_pos, h_pos, v_neg, h_neg = sampler(rbm, vis, ctx)
    dW = @get_array(ctx, :dW_buf, size(rbm.W), similar(rbm.W))
    n_obs = size(vis, 2)
    # same as: dW = ((h_pos * v_pos') - (h_neg * v_neg')) / n_obs
    gemm!('N', 'T', T(1 / n_obs), h_neg, v_neg, T(0.0), dW)
    gemm!('N', 'T', T(1 / n_obs), h_pos, v_pos, T(-1.0), dW)
    # gradient for vbias and hbias
    db = squeeze(sum(v_pos, 2) - sum(v_neg, 2), 2) ./ n_obs
    dc = squeeze(sum(h_pos, 2) - sum(h_neg, 2), 2) ./ n_obs
    return dW, db, dc
end

## updating

function grad_apply_learning_rate!{T,V,H}(rbm::RBM{T,V,H}, X::Mat,
                                          dtheta::Tuple, ctx::Dict)
    dW, db, dc = dtheta
    lr = T(@get(ctx, :lr, 0.1))
    # same as: dW *= lr
    scal!(length(dW), lr, dW, 1)
    scal!(length(db), lr, db, 1)
    scal!(length(dc), lr, dc, 1)
end

function grad_apply_momentum!{T}(rbm::RBM{T}, X::Mat,
                                 dtheta::Tuple, ctx::Dict)
    dW, db, dc = dtheta
    momentum = @get(ctx, :momentum, 0.9)
    dW_prev = @get_array(ctx, :dW_prev, size(dW), zeros(T, size(dW)))
    # same as: dW += momentum * dW_prev
    axpy!(momentum, dW_prev, dW)
end

function grad_apply_weight_decay!(rbm::RBM, X::Mat,
                                  dtheta::Tuple, ctx::Dict)
    # The decay penalty should drive all weights toward
    # zero by some small amount on each update.
    dW, db, dc = dtheta
    decay_kind = @get_or_return(ctx, :weight_decay_kind, nothing)
    decay_rate = @get(ctx, :weight_decay_rate,
                      throw(ArgumentError("If using :weight_decay_kind, weight_decay_rate should also be specified")))
    is_l2 = @get(ctx, :l2, false)
    if decay_kind == :l2
        # same as: dW -= decay_rate * W
        axpy!(-decay_rate, rbm.W, dW)
    elseif decay_kind == :l1
        # same as: dW -= decay_rate * sign(W)
        axpy!(-decay_rate, sign(rbm.W), dW)
    end

end

function grad_apply_sparsity!{T}(rbm::RBM{T}, X::Mat,
                              dtheta::Tuple, ctx::Dict)
    # The sparsity constraint should only drive the weights
    # down when the mean activation of hidden units is higher
    # than the expected (hence why it isn't squared or the abs())
    dW, db, dc = dtheta
    cost = @get_or_return(ctx, :sparsity_cost, nothing)
    target = @get(ctx, :sparsity_target, throw(ArgumentError("If :sparsity_cost is used, :sparsity_target should also be defined")))
    curr_sparsity = mean(hid_means(rbm, X))
    penalty = T(cost * (curr_sparsity - target))
    add!(dW, -penalty)
    add!(db, -penalty)
    add!(dc, -penalty)
end

function update_weights!(rbm::RBM, dtheta::Tuple, ctx::Dict)
    dW, db, dc = dtheta
    axpy!(1.0, dW, rbm.W)
    rbm.vbias += db
    rbm.hbias += dc
    # save previous dW
    dW_prev = @get_array(ctx, :dW_prev, size(dW), similar(dW))
    copy!(dW_prev, dW)
end

"""
Update RBM parameters using provided tuple `dtheta = (dW, db, dc)` of
parameter gradients. Before updating weights, following transformations
are applied to gradients:

 * learning rate (see `grad_apply_learning_rate!` for details)
 * momentum (see `grad_apply_momentum!` for details)
 * weight decay (see `grad_apply_weight_decay!` for details)
 * sparsity (see `grad_apply_sparsity!` for details)
"""
function update_classic!(rbm::RBM, X::Mat, dtheta::Tuple, ctx::Dict)
    # apply gradient updaters. note, that updaters all have
    # the same signature and are thus composable
    grad_apply_learning_rate!(rbm, X, dtheta, ctx)
    grad_apply_momentum!(rbm, X, dtheta, ctx)
    grad_apply_weight_decay!(rbm, X, dtheta, ctx)
    grad_apply_sparsity!(rbm, X, dtheta, ctx)
    # add gradient to the weight matrix
    update_weights!(rbm, dtheta, ctx)
end

## fitting

function fit_batch!(rbm::RBM, X::Mat, ctx = Dict())
    grad = @get_or_create(ctx, :gradient, gradient_classic)
    upd = @get_or_create(ctx, :update, update_classic!)
    dtheta = grad(rbm, X, ctx)
    upd(rbm, X, dtheta, ctx)
    return rbm
end

"""
Fit RBM to data `X`. Options that can be provided in the `opts` dictionary:

 * :n_epochs - number of full loops over data (default: 10)
 * :batch_size - size of mini-batches to use (default: 100)
 * :randomize - boolean, whether to shuffle batches or not (default: false)
 * :gradient - function to use for calculating parameter gradients
               (default: gradient_classic)
 * :update - function to use to update weights using calculated gradient
             (default: update_classic!)
 * :scorer - function to calculate how good the model is at the moment
             (default: pseudo_likelihood)
 * :reporter - type for reporting intermediate results using `report()`
               function (default: TextReporter)

Each function can additionally take other options, see their
docstrings/code for details.

NOTE: this function is incremental, so one can, for example, run it for
10 epochs, then inspect the model, then run it for 10 more epochs
and check the difference.
"""
function fit{T}(rbm::RBM{T}, X::Mat, opts = Dict{Any,Any}())
    @assert minimum(X) >= 0 && maximum(X) <= 1
    ctx = copy(opts)
    check_options(ctx)
    n_examples = size(X, 2)
    batch_size = @get(ctx, :batch_size, 100)
    batch_idxs = split_evenly(n_examples, batch_size)
    if @get(ctx, :randomize, false)
        batch_idxs = sample(batch_idxs, length(batch_idxs); replace=false)
    end
    n_epochs = @get(ctx, :n_epochs, 10)
    scorer = @get_or_create(ctx, :scorer, pseudo_likelihood)
    reporter = @get_or_create(ctx, :reporter, TextReporter())
    for epoch=1:n_epochs
        epoch_time = @elapsed begin
            for (batch_start, batch_end) in batch_idxs
                # BLAS.gemm! can't handle sparse matrices, so cheaper
                # to make it dense here
                batch = full(X[:, batch_start:batch_end])
                batch = ensure_type(T, batch)
                fit_batch!(rbm, batch, ctx)
            end
        end
        score = scorer(rbm, X)
        report(reporter, rbm, epoch, epoch_time, score)
    end
    return rbm
end

fit(rbm::RBM, X::Mat; opts...) = fit(rbm, X, Dict(opts))

## operations on learned RBM

"""Base data X through trained RBM to obtain compressed representation"""
function transform{T}(rbm::RBM{T}, X::Mat)
    return hid_means(rbm, ensure_type(T, X))
end

"""Given trained RBM and sample of visible data, generate similar items"""
function generate{T}(rbm::RBM{T}, vis::Vec; n_gibbs=1)
    return gibbs(rbm, reshape(ensure_type(T, vis), length(vis), 1); n_times=n_gibbs)[3]
end

function generate{T}(rbm::RBM{T}, X::Mat; n_gibbs=1)
    return gibbs(rbm, ensure_type(T, X); n_times=n_gibbs)[3]
end

"""
Get weight matrix of a trained RBM. Options:

 * transpose - boolean, whether to transpose weight matrix (convenient)
               or not (efficient). Default: true
"""
function coef(rbm::RBM; transpose=true)
    return if transpose rbm.W' else rbm.W end
end
# synonyms
weights = coef

"""Get biases of hidden units"""
hbias(rbm::RBM) = rbm.hbias

"""Get biases of visible units"""
vbias(rbm::RBM) = rbm.vbias