using SIIPExamples
using PowerSystems
using TimeSeries
using Dates
using TimeZones
using DataFrames
using CSV

println("downloading data...")
datadir = joinpath(dirname(dirname(pathof(SIIPExamples))), "US-System")
siip_data = joinpath(datadir, "SIIP")
if !isdir(datadir)
    mkdir(datadir)
    mkdir(siip_data)
    tempfilename =
        download("https://zenodo.org/record/3753177/files/USATestSystem.zip?download=1")
    SIIPExamples.unzip(SIIPExamples.os, tempfilename, datadir)
end

config_dir = joinpath(
    dirname(dirname(pathof(SIIPExamples))),
    "script",
    "2_PowerSystems_examples",
    "US_config",
)

interconnect = "Texas"
timezone = FixedTimeZone("UTC-6")
initial_time = ZonedDateTime(DateTime("2016-01-01T00:00:00"), timezone)

println("formatting data ...")
!isnothing(interconnect) && println("filtering data to include $interconnect ...")
gen = DataFrame(CSV.File(joinpath(datadir, "plant.csv")))
filter!(row -> row[:interconnect] == interconnect, gen)
gencost = DataFrame(CSV.File(joinpath(datadir, "gencost.csv")))
gen = innerjoin(gen, gencost, on = :plant_id, makeunique = true, validate = (false, false))

function make_pwl(gen::DataFrame, traunches = 2)
    output_pct_cols = ["output_point_" * string(i) for i in 0:traunches]
    hr_cols = ["heat_rate_incr_" * string(i) for i in 1:traunches]
    pushfirst!(hr_cols, "heat_rate_avg_0")
    columns =
        NamedTuple{Tuple(Symbol.(vcat(output_pct_cols, hr_cols)))}(repeat([Float64[]], 6))
    pwl = DataFrame(columns)
    for row in eachrow(gen)
        traunch_len = (1.0 - row.Pmin / row.Pmax) / traunches
        pct = [row.Pmin / row.Pmax + i * traunch_len for i in 0:traunches]
        #c(pct) = pct * row.Pmax * (row.GenIOB + row.GenIOC^2 + row.GenIOD^3)
        c(pct) = pct * row.Pmax * (row.c1 + row.c2^2) + row.c0 #this formats the "c" columns to hack the heat rate parser in PSY
        hr = [c(pct[1])]
        [push!(hr, c(pct[i + 1]) - hr[i]) for i in 1:traunches]
        push!(pwl, vcat(pct, hr))
    end
    return hcat(gen, pwl)
end

gen = make_pwl(gen);

gen[!, "fuel_price"] .= 1000.0;  #this formats the "c" columns to hack the heat rate parser in PSY

gen[:, :unit_type] .= "OT"
gen[:, :min_up_time] .= 0.0
gen[:, :min_down_time] .= 0.0
gen[:, :ramp_30] .= gen[:, :ramp_30] ./ 30.0 # we need ramp rates in MW/min
[
    gen[gen.type .== "wind", col] .= ["Wind", 0.0, 0.0][ix] for
    (ix, col) in enumerate([:unit_type, :min_up_time, :min_down_time])
]
[
    gen[gen.type .== "solar", col] .= ["PV", 0.0, 0.0][ix] for
    (ix, col) in enumerate([:unit_type, :min_up_time, :min_down_time])
]
[
    gen[gen.type .== "hydro", col] .= ["HY", 0.0, 0.0][ix] for
    (ix, col) in enumerate([:unit_type, :min_up_time, :min_down_time])
]
[
    gen[gen.type .== "ng", col] .= [4.5, 8][ix] for
    (ix, col) in enumerate([:min_up_time, :min_down_time])
]
[
    gen[gen.type .== "coal", col] .= [24, 48][ix] for
    (ix, col) in enumerate([:min_up_time, :min_down_time])
]
[
    gen[gen.type .== "nuclear", col] .= [72, 72][ix] for
    (ix, col) in enumerate([:min_up_time, :min_down_time])
]

idx_zero_pmin = [
    g.type in ["ng", "coal", "hydro", "nuclear"] && g.Pmin <= 0 for
    g in eachrow(gen[:, [:type, :Pmin]])
]
gen[idx_zero_pmin, :Pmin] = gen[idx_zero_pmin, :Pmax] .* 0.05

gen[:, :name] = "gen" .* string.(gen.plant_id)
CSV.write(joinpath(siip_data, "gen.csv"), gen)

bus = DataFrame(CSV.File(joinpath(datadir, "bus.csv")))
!isnothing(interconnect) && filter!(row -> row[:interconnect] == interconnect, bus)
zone = DataFrame(CSV.File(joinpath(datadir, "zone.csv")))
bus = leftjoin(bus, zone, on = :zone_id)
bustypes = Dict(1 => "PV", 2 => "PQ", 3 => "REF", 4 => "ISOLATED")
bus.bustype = [bustypes[b] for b in bus.type]
filter!(row -> row[:bustype] != PowerSystems.BusTypes.ISOLATED, bus)
bus.name = "bus" .* string.(bus.bus_id)
CSV.write(joinpath(siip_data, "bus.csv"), bus)

branch = DataFrame(CSV.File(joinpath(datadir, "branch.csv")))
branch = leftjoin(
    branch,
    DataFrames.rename!(bus[:, [:bus_id, :baseKV]], [:from_bus_id, :from_baseKV]),
    on = :from_bus_id,
)
branch = leftjoin(
    branch,
    DataFrames.rename!(bus[:, [:bus_id, :baseKV]], [:to_bus_id, :to_baseKV]),
    on = :to_bus_id,
)
!isnothing(interconnect) && filter!(row -> row[:interconnect] == interconnect, branch)
branch.name = "branch" .* string.(branch.branch_id)
branch.tr_ratio = branch.from_baseKV ./ branch.to_baseKV
CSV.write(joinpath(siip_data, "branch.csv"), branch)

dcbranch = DataFrame(CSV.File(joinpath(datadir, "dcline.csv")))
!isnothing(interconnect) && filter!(row -> row[:from_bus_id] in bus.bus_id, dcbranch)
!isnothing(interconnect) && filter!(row -> row[:to_bus_id] in bus.bus_id, dcbranch)
dcbranch.name = "dcbranch" .* string.(dcbranch.dcline_id)
dcbranch[:, :control_mode] .= "Power"
CSV.write(joinpath(siip_data, "dc_branch.csv"), dcbranch)

timeseries = []
ts_csv = ["wind", "solar", "hydro", "demand"]
plant_ids = Symbol.(string.(gen.plant_id))
for f in ts_csv
    println("formatting $f.csv ...")
    csvpath = joinpath(siip_data, f * ".csv")
    csv = DataFrame(CSV.File(joinpath(datadir, f * ".csv")))
    (category, name_prefix, label) =
        f == "demand" ? ("Area", "", "max_active_power") :
        ("Generator", "gen", "max_active_power")
    if !(:DateTime in names(csv))
        DataFrames.rename!(
            csv,
            (names(csv)[occursin.("UTC", String.(names(csv)))][1] => :DateTime),
        )
        #The timeseries data is in UTC, this converts it to a fixed UTC offset
        csv.DateTime =
            ZonedDateTime.(
                DateTime.(csv.DateTime, "yyyy-mm-dd HH:MM:SS"),
                timezone,
                from_utc = true,
            )
        delete!(csv, csv.DateTime .< initial_time)
        csv.DateTime = Dates.format.(csv.DateTime, "yyyy-mm-ddTHH:MM:SS")
    end
    device_names = f == "demand" ? unique(bus.zone_name) : gen.name
    for id in names(csv)
        colname = id
        if f == "demand"
            if Symbol(id) in Symbol.(zone.zone_id)
                colname = Symbol(zone[Symbol.(zone.zone_id) .== Symbol(id), :zone_name][1])
                DataFrames.rename!(csv, (id => colname))
            end
            sf = sum(bus[string.(bus.zone_id) .== id, :Pd])
        else
            if Symbol(id) in plant_ids
                colname = Symbol(gen[Symbol.(gen.plant_id) .== Symbol(id), :name][1])
                DataFrames.rename!(csv, (id => colname))
            end
            sf = maximum(csv[:, colname]) == 0.0 ? 1.0 : "Max"
        end
        if String(colname) in device_names
            push!(
                timeseries,
                Dict(
                    "simulation" => "DA",
                    "category" => category,
                    "module" => "InfrastructureSystems",
                    "type" => "SingleTimeSeries",
                    "component_name" => String(colname),
                    "name" => label,
                    "resolution" => 3600,
                    "scaling_factor_multiplier" => "get_max_active_power",
                    "scaling_factor_multiplier_module" => "PowerSystems",
                    "normalization_factor" => sf,
                    "data_file" => csvpath,
                ),
            )
        end
    end
    CSV.write(csvpath, csv)
end

timeseries_pointers = joinpath(siip_data, "timeseries_pointers.json")
open(timeseries_pointers, "w") do io
    PowerSystems.InfrastructureSystems.JSON3.write(io, timeseries)
end

println("parsing csv files...")
rawsys = PowerSystems.PowerSystemTableData(
    siip_data,
    100.0,
    joinpath(config_dir, "us_descriptors.yaml"),
    generator_mapping_file = joinpath(config_dir, "us_generator_mapping.yaml"),
)

println("creating System")
sys = System(rawsys; config_path = joinpath(config_dir, "us_system_validation.json"));
sys

to_json(sys, joinpath(siip_data, "sys.json"), force = true)