# PACKAGES
using CSV
using Ipopt
using JuMP
using Plots
using Statistics
using DataFrames

function generate_heat_demand(yearly_target::Int64)
    # This function returns a time series for the electricity consumption
    # of a heat pump given a total yearly consumption

    df = CSV.read(joinpath(@__DIR__, "temperature.csv"))
    hours = length(df[!,:T])

    # Degree days method
    heat_demand = zeros(Float64,hours)
    for d = 1:366
        daily_mean = 0
        for h = 1:24
            daily_mean = daily_mean + df[!,:T][d*24-24+h]
        end
        daily_mean = daily_mean/24
        for h = 1:24
            heat_demand[d*24-24+h] = max(17-daily_mean,0)
        end
    end
    total = sum(heat_demand)
    for h = 1:hours
        heat_demand[h] = (heat_demand[h]/total)*0.8+0.2*(1/hours)
    end

    # COP
    COP = zeros(Float64,hours)
    for h = 1:hours
        COP[h] = 2.77+0.155*df[!,:T][h]+0.00918*df[!,:T][h]^2
    end

    # Electricity demand
    demand = zeros(Float64,hours)
    for h = 1:hours
        demand[h] = heat_demand[h]/COP[h]
    end
    total = sum(demand)
    for h = 1:hours
        demand[h] = demand[h]*(1/total)*yearly_target
    end
    return demand
end

function run_model(enable_HP::Bool,enable_EV::Bool,enable_PV::Bool,power_tariff::Float64,power_interval::String,time_of_use::Bool)

    model = Model(with_optimizer(Ipopt.Optimizer, max_cpu_time=300.0))
    df = CSV.read(joinpath(@__DIR__, "data.csv"))
    hours = length(df[!,:Hour])

    ToU = zeros(Float64,hours)
    for i = 1:hours
        ToU[i] = 0.3356
        if time_of_use
            if i <= 3*30*24 || i >= 9*30*24
                if mod(i,24) >= 17 && mod(i,24) <= 20
                    ToU[i] = 0.8469
                end
            end
        end
    end

    @variable(model, pt[i = 1:hours] >= 0)
    @variable(model, CA[i = 1:hours] >= 0)
    @variable(model, WA[i = 1:hours] >= 0)
    @variable(model, HP[i = 1:hours] >= 0)
    @variable(model, 3 >= EV[i = 1:hours] >= 0)
    @variable(model, SC[i = 1:hours] >= 0)

    @NLconstraint(model, [i = 1:hours], CA[i] >= df[!,:CA_fixed][i])
    @NLconstraint(model, [p = 1:2928], sum(CA[i] for i in (p*3-2):(p*3)) == sum((df[!,:CA_fixed][i]+df[!,:CA_3hr][i]) for i in (p*3-2):(p*3)))

    @NLconstraint(model, [i = 1:hours], WA[i] >= df[!,:WA_fixed][i])
    @NLconstraint(model, [p = 1:366], sum(WA[i] for i in (p*24-23):(p*24)) == sum((df[!,:WA_fixed][i]+df[!,:WA_24hr][i]) for i in (p*24-23):(p*24)))

    if enable_HP
        demand_HP = generate_heat_demand(3200)
        HP_capacity = 1.1*maximum(demand_HP)
        @NLconstraint(model, [p = 1:2196], sum(HP[i] for i in (p*4-3):(p*4)) == sum(demand_HP[i] for i in (p*4-3):(p*4)))
        @NLconstraint(model, [i = 1:hours], HP[i] <= HP_capacity)
    else
        @NLconstraint(model, [i = 1:hours], HP[i] == 0)
        demand_HP = zeros(Int64,hours)
    end

    if enable_EV
        @NLconstraint(model, [i = 1:hours], EV[i] >= df[!,:EV_fixed][i])
        @NLconstraint(model, [p = 1:732], sum(EV[i] for i in (p*12-11):(p*12)) >= sum((df[!,:EV_fixed][i]+df[!,:EV_12hr][i]) for i in (p*12-11):(p*12)))
        @NLconstraint(model, [p = 1:366], sum(EV[i] for i in (p*24-23):(p*24)) >= sum((df[!,:EV_fixed][i]+df[!,:EV_12hr][i]+df[!,:EV_24hr][i]) for i in (p*24-23):(p*24)))
    else
        @NLconstraint(model, [i = 1:hours], EV[i] == 0)
    end

    if enable_PV
        @NLconstraint(model, [i = 1:hours], SC[i] <= df[!,:PV][i])
        @NLconstraint(model, [i = 1:hours], SC[i] <= EV[i]+HP[i]+CA[i]+WA[i]+df[!,:Other_fixed][i])
    else
        @NLconstraint(model, [i = 1:hours], SC[i] == 0)
    end

    if power_interval == "year"
        @NLconstraint(model, [i = 2:hours], pt[i] == pt[1])
    elseif power_interval == "month"
        month_duration = [31,29,31,30,31,30,31,31,30,31,30,31]
        month_sum = 0
        for m = 1:12
            @NLconstraint(model, [i = (month_sum*24+2):((month_sum+month_duration[m])*24)], pt[i] == pt[month_sum*24+1])
            month_sum += month_duration[m]
        end
    elseif power_interval == "week"
        for w = 1:52
            @NLconstraint(model, [i = ((w-1)*7*24+2):(w*7*24)], pt[i] == pt[((w-1)*7*24+1)])
        end
        @NLconstraint(model, [i = 8738:8784], pt[i] == pt[8737])
    elseif power_interval == "day"
        @NLconstraint(model, [i = ])
    end

    @NLconstraint(model, [i = 1:hours], CA[i]+WA[i]+HP[i]+EV[i]+df[!,:Other_fixed][i]-SC[i] <= pt[i])

    # OBJECTIVE
    @NLobjective(model, Min, sum(pt[i]*power_tariff+((CA[i]+WA[i]+HP[i]+EV[i]+df[!,:Other_fixed][i]-SC[i])*(0.884+0.08+df[!,:Spot_price][i]/1000) - (df[!,:PV][i]-SC[i])*df[!,:Spot_price][i]/1000) for i in 1:hours))

    # RESULTS
    optimize!(model)
    return objective_value(model), JuMP.value.(CA), JuMP.value.(WA), JuMP.value.(HP), JuMP.value.(EV), df[!,:Other_fixed], JuMP.value.(SC), df[!,:PV], demand_HP+df[!,:CA_fixed]+df[!,:CA_3hr]+df[!,:WA_fixed]+df[!,:WA_24hr]+df[!,:EV_fixed]+df[!,:EV_12hr]+df[!,:EV_24hr]+df[!,:Other_fixed], df[!,:Spot_price]/1000, JuMP.value.(pt)
end

obj, CA, WA, HP, EV, Other, SC, PV, without, spot, pt = run_model(true,true,true,0.28,"month",false)

total = CA+WA+HP+EV+Other-PV
#old = without-PV
sort!(total,rev=true)
#sort!(old,rev=true)

plotsize = 1:8784
plot(
    plotsize,
    [
        total[plotsize],
        pt[plotsize]
    ],
    ylim = (-10,10),
    label = ["Consumption" "Old"]
)