#* Functions for modelling the impact of hurricanes on power grids.
#*------------------------------------------------------------------------------

function sim_impact(
        seg_data::Dict{String,<:Any};
        γ::Float64
    )

    tl_failures = Array{Tuple{String,Int64},1}() # array for line failures
    
    ### Go through overhead transmission lines and simulate until failure
    for (i, tl) in seg_data 
        for (f, windloads) in tl["windloads"] # go through wind frames
            p_fail = min.(γ * windloads, 1)
            vals = rand(tl["N_seg"])
            if any(vals .< p_fail) # line gets destroyed
                push!(tl_failures, (i, f)) # save line ID and frame
                break
            end
        end
    end

    return tl_failures
end

#*------------------------------------------------------------------------------

#=
Calculates wind loads (wind force / breaking force) for overhead transmission lines in the network data dictionary and a given .nc file containing wind data. The overhead transmission lines are identified according to the chosen mode ((:sum or :single, see get_underground_tl) and divided into segments between towers. The results are returned in form of a dictionary (see calc_overhead_tl_segments).
=#
function calc_overhead_tl_windloads(
        network_data::Dict{String,<:Any},
        windfile::String, # .nc file containing wind data
        d_twrs = 161.0;
        mode = :sum # :sum or :single
    )

    ### Read wind data
    wind_lons, wind_lats, wind_speeds = get_windfield(windfile)

    ### Identify overhead lines and divide them into segments between towers
    overhead_tl_segments = calc_overhead_tl_segments(
        network_data, d_twrs, mode=mode
    )

    ### Assign location indices to segments that identify the local wind speed
    _assign_loc_indices!(overhead_tl_segments, wind_lons, wind_lats)

    ### Go through time steps in the wind data and calculate wind loads
    for f in 1:size(wind_speeds, 3)
        wind_frame = @view wind_speeds[:,:,f]
        _calc_seg_windloads!(overhead_tl_segments, wind_frame, f)
    end

    return overhead_tl_segments # dictionary with all information
end

#*------------------------------------------------------------------------------

#= 
Returns the longitude and latitude coordinates of the wind field together with the whole time series of wind speeds.
=#
function get_windfield(filepath::String)
    wind_lons = round.(ncread(filepath, "lon"), digits=6)
    wind_lats = round.(ncread(filepath, "lat"), digits=6)
    wind_speeds = ncread(filepath, "wind")

    return wind_lons, wind_lats, wind_speeds
end

function get_windfield(filepath::String, frame::Int64)
    wind_lons, wind_lats, wind_speeds = get_windfield(filepath)
    return wind_lons, wind_lats, wind_speeds[:, :, frame]
end

#=
Function that returns the number of frames (time steps) in the given .nc data file and the maximum wind speed.
=#
function get_winddata(filepath::String)
    wind = ncread(filepath, "wind")
    N_frames = size(wind, 3)
    max_ws = maximum(wind)
    return N_frames, max_ws
end

#*------------------------------------------------------------------------------

#=
Identifies overhead transmission lines in the network data dictionary according to the chosen mode (:sum or :single, see get_underground_tl) and divides them into segments between towers. The towers are assumed to be d_twrs (in m) apart and the resulting segments have a length (l_seg) close to d_twrs. The resulting segments are returned in form of a dictionary with entries explained in the code below.
=#
function calc_overhead_tl_segments(
        network_data::Dict{String,<:Any}, 
        d_twrs = 161.0;
        mode = :sum # :sum or :single
    )

    ### Get (string) indices of underground transmission lines
    underground_tl = get_underground_tl(network_data, mode=mode)
    
    unique_overhead_tl = [
        i for (i, b) in network_data["branch"] 
        if b["transformer"] == false && i ∉ underground_tl 
        && b["source_id"][end] == "1 " # parallel circuits excluded
    ] # string indices of unique overhead transmission lines

    seg_data = Dict{String,Dict{String,Any}}() # dictionary for segment data

    for i in unique_overhead_tl
        branch = network_data["branch"][i]
        f_bus, t_bus = branch["f_bus"], branch["t_bus"]
        lon_f = network_data["bus"]["$f_bus"]["bus_lon"]
        lat_f = network_data["bus"]["$f_bus"]["bus_lat"]
        lon_t = network_data["bus"]["$t_bus"]["bus_lon"]
        lat_t = network_data["bus"]["$t_bus"]["bus_lat"]

        L, N_seg, l_seg, seg_lons, seg_lats = _calc_tl_segments(
            lon_f, lat_f, lon_t, lat_t, d_twrs
        ) # data for transmission line segments

        ### Add data to dictionary
        seg_data[i] = Dict(
            "f_bus" => f_bus, # index (Int64) of from-bus
            "lon_f" => lon_f, # longitude of from-bus
            "lat_f" => lat_f, # latitude of from-bus
            "t_bus" => t_bus, # index (Int64) of to-bus
            "lon_t" => lon_t, # longitude of to-bus
            "lat_t" => lat_t, # latitude of to-bus
            "L" => L, # calculated length of transmission line
            "N_seg" => N_seg, # number of segments between towers
            "l_seg" => l_seg, # actual length of segments (close to d_twrs)
            "seg_lons" => seg_lons, # longitudes of segment-half-way points
            "seg_lats" => seg_lats, # latitudes of segment-half-way points
            "windloads" => Dict{Int64,Array{Float64,1}}() # for wind loads
        )
    end

    return seg_data
end

#= 
Function that divides a transmission line (tl) between a "from-bus" with coordinates (lon_f, lat_f) in degrees and a "to-bus" (lon_t, lat_t) into segments of wanted length d_twrs (average distance between transmission towers in meter) and returns the half-way points (longitude and latitude coordinates) of all segments together with length L of the transmission line (calculated using the haversine formula), the number of segments (N_twrs+1), and the actual length of segments l_seg.
=#
function _calc_tl_segments(
        lon_f::Float64, lat_f::Float64, 
        lon_t::Float64, lat_t::Float64,
        d_twrs = 161.0
    )

    R = 6371000 # earth radius
    λf, φf, λt, φt = deg2rad.([lon_f, lat_f, lon_t, lat_t])

    ### Calculate haversine distance between "from-" and "to-bus"
    h = sin((φt-φf)/2)^2 + cos(φf) * cos(φt) * sin((λt-λf)/2)^2 # haversine
    δ = 2 * atan(sqrt(h), sqrt(1-h)) # angular distance in radians
    L = R * δ # transmission line length

    N_twrs = round(L/d_twrs) # number of transmission towers
    l_seg = L/N_twrs # actual length of line segments

    ### Calulate the fractions of the line at which the half-way points lie
    f = 1/N_twrs
    f_seg = [(i+0.5)*f for i in 0:N_twrs]

    a = sin.((1 .- f_seg)*δ)/sin(δ)
    b = sin.((f_seg*δ))/sin(δ)
    x = a*cos(φf)*cos(λf) + b*cos(φt)*cos(λt)
    y = a*cos(φf)*sin(λf) + b*cos(φt)*sin(λt)
    z = a*sin(φf) + b*sin(φt)

    ### Calculate longitudes and latitudes of all half-way points in radians
    λ_seg = atan.(y, x) 
    φ_seg = atan.(z, sqrt.(x.^2 + y.^2))

    return L, Int64(N_twrs+1), l_seg, rad2deg.(λ_seg), rad2deg.(φ_seg)
end


#*------------------------------------------------------------------------------

#=
Assigns location indices to transmission line segments that help to identify local wind speeds for the calculation of wind loads. The location indices are added to the data dictionary seg_data.
=#
function _assign_loc_indices!(
        seg_data::Dict{String,<:Any},
        wind_lons::Array{Float64,1}, # longitudes in wind field
        wind_lats::Array{Float64,1} # latitudes in wind field
    )

    ### Go through overhead transmission lines
    for tl_data in collect(values(seg_data))
        N_seg = tl_data["N_seg"]
        lon_indices = zeros(Int64, N_seg)
        lat_indices = zeros(Int64, N_seg)
        ### Go through segments and identify their position in the wind field
        for i in 1:N_seg
            seg_lon = tl_data["seg_lons"][i]
            seg_lat = tl_data["seg_lats"][i]
            lon_indices[i] = argmin(abs.(wind_lons .- seg_lon))
            lat_indices[i] = argmin(abs.(wind_lats .- seg_lat))
        end
        tl_data["loc_indices"] = [loc for loc in zip(lon_indices, lat_indices)]
    end

    return nothing
end

#*------------------------------------------------------------------------------

#=
Calculates wind loads for all overhead transmission line segments in the segment data dictionary containing location indices (see _assign_loc_indices!) and a given time step (frame) of the wind field. The results are added to seg_data.
=#
function _calc_seg_windloads!(
        seg_data::Dict{String,<:Any}, 
        wind_frame::SubArray, # slice of wind speed time series
        f::Int64 # wind frame number
    )

    ### Go through overhead transmission lines
    for (key, tl_data) in seg_data
        windloads = zeros(tl_data["N_seg"])
        ### Go through unique segment location indices and calculate wind loads
        for loc in unique(tl_data["loc_indices"])
            ws = wind_frame[CartesianIndex(loc)] # local wind speed
            wl = _calc_windload(ws, tl_data["l_seg"]) # local wind load
            ### Add this wind load to all segments with these location indices
            for i in findall(l -> l == loc, tl_data["loc_indices"])
                windloads[i] = wl
            end
        end
        tl_data["windloads"][f] = windloads # add result to dictionary
    end
    
    return nothing
end

#= 
Function that calculates the failure probabilities of transmission line segments according to the local wind speed and the length of the segments.
=#
function _calc_windload(v::Float64, l::Float64)
    ### Parameters in wind force design equation
    Q, k_z, C_f, d, F_brk = [0.613, 1.17, 0.9, 0.03, 152130]

    ### Gust response factor for line segment
    E_W = 4.9*sqrt(0.005)*(33/70)^(1/7)
    G_WRF = 1 + 2.7*E_W/sqrt(1 + 0.8/67.056*l)

    ### Convert wind speeds to an effective height of h = 70ft = 21.336m
    v *= 1.16 # conversion factor ln(21.336/0.1)/ln(100) ≈ 1.16 from 10m to h

    ### Calculate wind force according to ASCE design equation
    F_wind = Q * k_z * C_f * G_WRF * d * l * v^2
    ### Return wind load on all segment
    return F_wind / F_brk
end