
"""
Adapted from sandreza/Learning/sandbox/oceananigans_converter.jl
https://github.com/sandreza/Learning/blob/master/sandbox/oceananigans_converter.jl
"""

using JLD2, NetCDF

struct OceananigansData{𝒮, 𝒯, 𝒰, 𝒱}
    # initial conditions, 4
    T⁰::𝒮
    S⁰::𝒮
    U⁰::𝒮
    V⁰::𝒮

    # fields at each moment in time, 4
    T::𝒯
    S::𝒯
    U::𝒯
    V::𝒯

    # some second order statistics at each moment in time, 5
    wT::𝒯
    wS::𝒯
    uu::𝒯
    vv::𝒯
    ww::𝒯
    tke::𝒯

    # simulation constants, 8
    ρ::𝒰
    α::𝒰
    β::𝒰
    cᵖ::𝒰
    f⁰::𝒰
    g::𝒰
    L::𝒰
    κₑ::𝒰

    # time and grid, 2
    t::𝒮
    z::𝒮

    #top boundary condition data, see string for type of boundary condition, 4
    top_T::𝒰
    top_S::𝒰
    top_U::𝒰
    top_V::𝒰

    #bottom boundary condition data, see string for type of boundary condtion,4
    bottom_T::𝒰
    bottom_S::𝒰
    bottom_U::𝒰
    bottom_V::𝒰

    #info about the simulation, 1
    info::𝒱
end




"""
OceananigansData(filename)
# Description
- Constructor for Oceananigans data type. Loads data from LES
# Fields for the output are
    # initial conditions
    T⁰::𝒮
    S⁰::𝒮
    U⁰::𝒮
    V⁰::𝒮
    # fields at each moment in time
    T::𝒯
    S::𝒯
    U::𝒯
    V::𝒯
    # some second order statistics at each moment in time
    wT::𝒯
    wS::𝒯
    uu::𝒯
    vv::𝒯
    ww::𝒯
    tke::𝒯
    # simulation constants
    ρ::𝒰
    α::𝒰
    β::𝒰
    cᵖ::𝒰
    f⁰::𝒰
    g::𝒰
    # time and grid
    t::𝒮
    z::𝒮
    #top boundary condition data, see string for type of boundary condition
    top_T::𝒰
    top_S::𝒰
    top_U::𝒰
    top_V::𝒰
    #bottom boundary condition data, see string for type of boundary condtion
    bottom_T::𝒰
    bottom_S::𝒰
    bottom_U::𝒰
    bottom_V::𝒰
    #info about the simulation
    info::𝒱
"""
function ReadJLD2_OceananigansData(filename)
    les_data = jldopen(filename, "r")
    les_keys = keys(les_data)
    timeseries_keys = keys(les_data["timeseries"]["t"])

    # hold the entries for easy constructor creation
    container = []

    println(collect(les_data["grid"]["zC"]))

    # size of arrays
    Nz = length(collect(les_data["grid"]["zC"]))
    Nt = length(timeseries_keys)

    ## construct arrays
    #Initial Conditions
    T⁰ = zeros(Nz)
    S⁰ = zeros(Nz)
    U⁰ = zeros(Nz)
    V⁰ = zeros(Nz)

    #Timeseries
    T = zeros(Nz, Nt)
    S = zeros(Nz, Nt)
    U = zeros(Nz, Nt)
    V = zeros(Nz, Nt)
    t = zeros(Nt)

    println(Nz)

    #Second Order Statistics
    wT  = zeros(Nz, Nt)
    wS  = zeros(Nz, Nt)
    uu  = zeros(Nz, Nt)
    vv  = zeros(Nz, Nt)
    ww  = zeros(Nz, Nt)

    # grab arrays
    for j in 1:Nt
        # Fields
        key = timeseries_keys[j]
        println(length(les_data["timeseries"]["T"][key]))
        @. T[:,j] = les_data["timeseries"]["T"][key][2:(end-1)]
        @. S[:,j] = les_data["timeseries"]["S"][key][2:(end-1)]
        @. U[:,j] = les_data["timeseries"]["u"][key][2:(end-1)]
        @. V[:,j] = les_data["timeseries"]["v"][key][2:(end-1)]
        # Second Order Statistics
        @. wT[:,j] = les_data["timeseries"]["wT"][key][2:(end-1)]
        @. wS[:,j] = les_data["timeseries"]["wS"][key][2:(end-1)]
        @. uu[:,j] = les_data["timeseries"]["uu"][key][2:(end-1)]
        @. vv[:,j] = les_data["timeseries"]["vv"][key][2:(end-1)]
        @. ww[:,j] = les_data["timeseries"]["ww"][key][2:(end-1)]

        t[j] = les_data["timeseries"]["t"][key]
    end

    tke = 0.5*(uu.^2 .+ vv.^2 .+ ww.^2) .* 2e8

    # Set initial Conditions
    @. T⁰ = T[:,1]
    @. S⁰ = S[:,1]
    @. U⁰ = U[:,1]
    @. V⁰ = V[:,1]

    # Push initial conditions current stuff into container
    push!(container, T⁰, S⁰, V⁰, U⁰)
    # Push fields into container
    push!(container, T, S, U, V)
    # Push second order statistics into container
    push!(container, wT, wS, uu, vv, ww, tke)

    # Now grab parameter
    ρ = les_data["parameters"]["density"]
    α = les_data["buoyancy"]["equation_of_state"]["α"]
    β = les_data["buoyancy"]["equation_of_state"]["β"]
    cᵖ = les_data["parameters"]["specific_heat_capacity"]
    f⁰ = les_data["coriolis"]["f"]
    g = les_data["buoyancy"]["gravitational_acceleration"]
    L = les_data["grid"]["Lz"]
    κₑ = les_data["parameters"]["diffusivity_T"]

    # b = T .* (g * α)

    # Push parameters to container
    push!(container, ρ, α, β, cᵖ, f⁰, g, L, κₑ)

    # grab domain data
    z = collect(les_data["grid"]["zC"])

    # push
    push!(container, t, z)

    # now grab boundary condition data
    top_T = les_data["boundary_conditions"]["top"]["FT"]
    top_S = 0.0
    top_U = les_data["boundary_conditions"]["top"]["Fu"]
    top_V = 0.0
    #bottom boundary condition data, see string for type of boundary condtion
    bottom_T = les_data["boundary_conditions"]["bottom"]["dTdz"]
    bottom_S = 0.0
    bottom_U = 0.0
    bottom_V = 0.0

    # push to container
    push!(container, top_T, top_S, top_U, top_V, bottom_T, bottom_S, bottom_U, bottom_V)

    # Now construct types
    𝒮 = typeof(T⁰)
    𝒯 = typeof(T)
    𝒰 = typeof(α)
    𝒱 = typeof("string")

    # now create data string
    info_string = "The top boundary conditions are flux boundary conditions \n"
    info_string *= "The  bottom boundary condition for temperature is a gradient boundary condition \n"
    info_string *= "The grid data is assumed to be evenly spaced and a power of two \n"

    # push to container
    push!(container, info_string)
    #return container
    close(les_data)
    return OceananigansData{𝒮, 𝒯, 𝒰, 𝒱}(container...)
end



"""
OceananigansData(filename)
# Description
- Constructor for Oceananigans data type. Loads data from LES
# Fields for the output are
    # initial conditions
    T⁰::𝒮
    S⁰::𝒮
    U⁰::𝒮
    V⁰::𝒮
    # fields at each moment in time
    T::𝒯
    S::𝒯
    U::𝒯
    V::𝒯
    # some second order statistics at each moment in time
    wT::𝒯
    wS::𝒯
    uu::𝒯
    vv::𝒯
    ww::𝒯
    # simulation constants
    ρ::𝒰
    α::𝒰
    β::𝒰
    cᵖ::𝒰
    f⁰::𝒰
    g::𝒰
    # time and grid
    t::𝒮
    z::𝒮
    #top boundary condition data, see string for type of boundary condition
    top_T::𝒰
    top_S::𝒰
    top_U::𝒰
    top_V::𝒰
    #bottom boundary condition data, see string for type of boundary condtion
    bottom_T::𝒰
    bottom_S::𝒰
    bottom_U::𝒰
    bottom_V::𝒰
    #info about the simulation
    info::𝒱
"""
function ReadNetCDF_OceananigansData(filename)


    ncinfo(filename)
    x = Array(NetCDF.open(filename, "yC"))
    y = Array(NetCDF.open(filename, "yC"))
    z = Array(NetCDF.open(filename, "zC"))
    t = Array(NetCDF.open(filename, "time"))
    sim_day = t ./ 86400

    z = (z[2:end] + z[1:end-1]) / 2

    vectify(x) = [x[i] for i in 1:length(x)]
    b = vectify(NetCDF.open(filename, "b"))
    u = vectify(NetCDF.open(filename, "u"))
    v = vectify(NetCDF.open(filename, "v"))
    w = vectify(NetCDF.open(filename, "w"))
    ##

    # T = b./(g*α)
    T = b./(9.81 * 2e-4)

    # hold the entries for easy constructor creation
    container = []

    # size of arrays
    Nz = length(z)
    Nt = length(t)

    ## construct arrays
    #Initial Conditions
    T⁰ = zeros(Nz)
    S⁰ = zeros(Nz)
    U⁰ = zeros(Nz)
    V⁰ = zeros(Nz)
    #Timeseries
    T = zeros(Nz, Nt)
    S = zeros(Nz, Nt)
    U = zeros(Nz, Nt)
    V = zeros(Nz, Nt)
    t = zeros(Nt)

    #Second Order Statistics
    wT = zeros(Nz, Nt)
    wS = zeros(Nz, Nt)
    uu = zeros(Nz, Nt)
    vv = zeros(Nz, Nt)
    ww = zeros(Nz, Nt)

    # grab arrays
    for j in 1:Nt
        # Fields
        key = timeseries_keys[j]
        @. T[:,j] = les_data["timeseries"]["T"][key]
        @. S[:,j] = les_data["timeseries"]["S"][key]
        @. U[:,j] = les_data["timeseries"]["u"][key]
        @. V[:,j] = les_data["timeseries"]["v"][key]
        # Second Order Statistics
        @. wT[:,j] = les_data["timeseries"]["wT"][key]
        @. wS[:,j] = les_data["timeseries"]["wS"][key]
        @. uu[:,j] = les_data["timeseries"]["uu"][key]
        @. vv[:,j] = les_data["timeseries"]["vv"][key]
        @. ww[:,j] = les_data["timeseries"]["ww"][key]

        t[j] = les_data["timeseries"]["t"][key]
    end
    # Set initial Conditions
    @. T⁰ = T[:,1]
    @. S⁰ = S[:,1]
    @. U⁰ = U[:,1]
    @. V⁰ = V[:,1]

    # Push initial conditions current stuff into container
    push!(container, T⁰, S⁰, V⁰, U⁰)
    # Push fields into container
    push!(container, T, S, U, V)
    # Push second order statistics into container
    push!(container, wT, wS, uu, vv, ww)

    # Now grab parameter
    ρ = les_data["parameters"]["density"]
    α = les_data["buoyancy"]["equation_of_state"]["α"]
    β = les_data["buoyancy"]["equation_of_state"]["β"]
    cᵖ = les_data["parameters"]["specific_heat_capacity"]
    f⁰ = les_data["coriolis"]["f"]
    g = les_data["buoyancy"]["gravitational_acceleration"]
    L = les_data["grid"]["Lz"]

    # Push parameters to container
    push!(container, ρ, α, β, cᵖ, f⁰, g, L)

    # grab domain data
    z = collect(les_data["grid"]["zC"])

    # push
    push!(container, t, z)

    # now grab boundary condition data
    top_T = les_data["boundary_conditions"]["top"]["FT"]
    top_S = 0.0
    top_U = les_data["boundary_conditions"]["top"]["Fu"]
    top_V = 0.0
    #bottom boundary condition data, see string for type of boundary condtion
    bottom_T = les_data["boundary_conditions"]["bottom"]["dTdz"]
    bottom_S = 0.0
    bottom_U = 0.0
    bottom_V = 0.0

    # push to container
    push!(container, top_T, top_S, top_U, top_V, bottom_T, bottom_S, bottom_U, bottom_V)

    # Now construct types
    𝒮 = typeof(T⁰)
    𝒯 = typeof(T)
    𝒰 = typeof(α)
    𝒱 = typeof("string")

    # now create data string
    info_string = "The top boundary conditions are flux boundary conditions \n"
    info_string *= "The  bottom boundary condition for temperature is a gradient boundary condition \n"
    info_string *= "The grid data is assumed to be evenly spaced and a power of two \n"

    # push to container
    push!(container, info_string)
    return OceananigansData{𝒮, 𝒯, 𝒰, 𝒱}(container...)
end


# """
# avg(Φ, n)
# # Description
# - Average a field down by n.
# - Requires field to have evenly spaced points. Size of N leq length(Φ).
# - Furthermore requires
# # Arguments
# - `Φ` :(vector) The field, an array
# - `n` :(Int) number of grid points to average down to.
# # Return
# - `Φ2` :(vector) The field with values averaged, an array
# """
# function avg(Φ, n)
#   m = length(Φ)
#   scale = Int(floor(m/n))
#   if ( abs(Int(floor(m/n)) - m/n) > eps(1.0))
#       return error
#   end
#   Φ2 = zeros(n)
#   for i in 1:n
#       Φ2[i] = 0
#           for j in 1:scale
#               Φ2[i] += Φ[scale*(i-1) + j] / scale
#           end
#   end
#   return Φ2
# end
