using GeoData, Test, Statistics, Dates, CFTime, Plots
import ArchGDAL, NCDatasets
using GeoData: name, mode, span, sampling, val, Ordered, metadata, bounds,
               FileArray, FileStack, NCDfile
include(joinpath(dirname(pathof(GeoData)), "../test/test_utils.jl"))

ncexamples = "https://www.unidata.ucar.edu/software/netcdf/examples/"
ncsingle = maybedownload(joinpath(ncexamples, "tos_O1_2001-2002.nc"))
ncmulti = maybedownload(joinpath(ncexamples, "test_echam_spectral.nc"))

stackkeys = (
    :abso4, :aclcac, :aclcov, :ahfcon, :ahfice, :ahfl, :ahfliac, :ahfllac,
    :ahflwac, :ahfres, :ahfs, :ahfsiac, :ahfslac, :ahfswac, :albedo, :albedo_nir,
    :albedo_nir_dif, :albedo_nir_dir, :albedo_vis, :albedo_vis_dif, :albedo_vis_dir,
    :alsobs, :alsoi, :alsol, :alsom, :alsow, :ameltdepth, :ameltfrac, :amlcorac,
    :ao3, :apmeb, :apmegl, :aprc, :aprl, :aprs, :aps, :az0i, :az0l, :az0w,
    :barefrac, :dew2, :drain, :evap, :evapiac, :evaplac, :evapwac, :fage, :friac,
    :geosp, :glac, :gld, :hyai, :hyam, :hybi, :hybm, :lsp, :q, :qres, :qvi, :relhum,
    :runoff, :sd, :seaice, :siced, :sicepdi, :sicepdw, :sicepres, :slm, :sn, :snacl,
    :snc, :sni, :snifrac, :snmel, :sofliac, :sofllac, :soflwac, :srad0, :srad0d,
    :srad0u, :sradl, :srads, :sradsu, :sraf0, :srafl, :srafs, :st, :svo, :t2max,
    :t2min, :temp2, :thvsig, :topmax, :tpot, :trad0, :tradl, :trads, :tradsu,
    :traf0, :trafl, :trafs, :trfliac, :trfllac, :trflwac, :tropo, :tsi, :tsicepdi,
    :tslm1, :tsurf, :tsw, :u10, :ustr, :ustri, :ustrl, :ustrw, :v10, :vdis, :vdisgw,
    :vstr, :vstri, :vstrl, :vstrw, :wimax, :wind10, :wl, :ws, :wsmx, :xi, :xivi,
    :xl, :xlvi
)

@testset "GeoArray" begin
    @time ncarray = GeoArray(ncsingle)

    @testset "open" begin
        @test all(open(A -> A[Y=1], ncarray) .=== ncarray[:, 1, :])
    end

    # GeoData.create(tempfile, eltype(ncarray), dims(ncarray);
        # keys=keys(ncarray), missingval=missingval(gdalarray)
    # )

    @testset "read" begin
        @time A = read(ncarray);
        @test A isa GeoArray
        @test parent(A) isa Array
        A2 = zero(A)
        @time read!(ncarray, A2);
        A3 = zero(A)
        @time read!(ncsingle, A3)
        @test all(A .=== A2)
        @test all(A .=== A3)
    end

    @testset "array properties" begin
        @test size(ncarray) == (180, 170, 24)
        @test ncarray isa GeoArray
        @test index(ncarray, Ti) == DateTime360Day(2001, 1, 16):Month(1):DateTime360Day(2002, 12, 16)
        @test index(ncarray, Y) == -79.5:89.5
        @test index(ncarray, X) == 1.0:2:359.0
        @test bounds(ncarray) == (
            (0.0, 360.0),
            (-80.0, 90.0),
            (DateTime360Day(2001, 1, 1), DateTime360Day(2003, 1, 1)),
        )
    end

    @testset "dimensions" begin
        @test ndims(ncarray) == 3
        @test length.(dims(ncarray)) == (180, 170, 24)
        @test dims(ncarray) isa Tuple{<:X,<:Y,<:Ti}
        @test refdims(ncarray) == ()
        # TODO detect the time span, and make it Regular
        modes = (
             Mapped(Ordered(), Explicit(vcat((0.0:2.0:358.0)', (2.0:2.0:360.0)')), Intervals(Center()), EPSG(4326), EPSG(4326)),
             Mapped(Ordered(), Explicit(vcat((-80.0:89.0)', (-79.0:90.0)')), Intervals(Center()), EPSG(4326), EPSG(4326)),
             Sampled(Ordered(), Explicit(
                 vcat(permutedims(DateTime360Day(2001, 1, 1):Month(1):DateTime360Day(2002, 12, 1)),
                      permutedims(DateTime360Day(2001, 2, 1):Month(1):DateTime360Day(2003, 1, 1)))
                ), Intervals(Center())
            )
        )
        @test val.(span(ncarray)) == val.(span.(modes))
        @test typeof(mode(ncarray)) == typeof(modes)
        @test bounds(ncarray) == ((0.0, 360.0), (-80.0, 90.0), (DateTime360Day(2001, 1, 1), DateTime360Day(2003, 1, 1)))
    end
    tempfile = tempname() * ".nc"

    @testset "other fields" begin
        @test ismissing(missingval(ncarray))
        @test metadata(ncarray) isa Metadata{NCDfile}
        @test name(ncarray) == :tos
    end

    @testset "indexing" begin
        @test ncarray[Ti(1)] isa GeoArray{<:Any,2}
        @test ncarray[Y(1), Ti(1)] isa GeoArray{<:Any,1}
        @test ncarray[X(1), Ti(1)] isa GeoArray{<:Any,1}
        @test ncarray[X(1), Y(1), Ti(1)] isa Missing
        @test ncarray[X(30), Y(30), Ti(1)] isa Float32
        # Russia
        @test ncarray[X(50), Y(100), Ti(1)] isa Missing
        # Alaska
        @test ncarray[Y(Near(64.2008)), X(Near(149.4937)), Ti(1)] isa Missing
        @test ncarray[Ti(2), X(At(59.0)), Y(At(-50.5))] == ncarray[30, 30, 2] === 278.47168f0
    end

    @testset "methods" begin
        @testset "mean" begin
            @test all(mean(ncarray; dims=Y) .=== mean(parent(ncarray); dims=2))
        end
        @testset "trim, crop, extend" begin
            a = read(ncarray)
            a[X(1:20)] .= missingval(a)
            trimmed = trim(a)
            @test size(trimmed) == (160, 169, 24)
            cropped = crop(a; to=trimmed)
            @test size(cropped) == (160, 169, 24)
            @test all(collect(cropped .=== trimmed))
            extended = extend(cropped; to=a)
            @test all(collect(extended .=== a))
        end
        @testset "mask and mask!" begin
            msk = read(ncarray)
            msk[X(1:100), Y([1, 5, 95])] .= missingval(msk)
            @test !all(ncarray[X(1:100)] .=== missingval(msk))
            masked = mask(ncarray; to=msk)
            @test all(masked[X(1:100), Y([1, 5, 95])] .=== missingval(msk))
            tempfile = tempname() * ".nc"
            cp(ncsingle, tempfile)
            @test !all(GeoArray(tempfile)[X(1:100), Y([1, 5, 95])] .=== missing)
            open(GeoArray(tempfile); write=true) do A
                mask!(A; to=msk, missingval=missing)
                # TODO: replace the CFVariable with a FileArray{NCDfile} so this is not required
                nothing
            end
            @test all(GeoArray(tempfile)[X(1:100), Y([1, 5, 95])] .=== missing)
            rm(tempfile)
        end
        @testset "mosaic" begin
            @time ncarray = GeoArray(ncsingle)
            A1 = ncarray[X(1:80), Y(1:100)]
            A2 = ncarray[X(50:150), Y(90:150)]
            tempfile = tempname() * ".nc"
            Afile = mosaic(first, read(A1), read(A2); missingval=missing, atol=1e-7, filename=tempfile)
            Amem = mosaic(first, A1, A2; missingval=missing, atol=1e-7)
            Atest = ncarray[X(1:150), Y(1:150)]
            Atest[X(1:49), Y(101:150)] .= missing
            Atest[X(81:150), Y(1:89)] .= missing
            @test all(Atest .=== Afile .=== Amem)
        end
        @testset "chunk_series" begin
            @test GeoData.chunk_series(ncarray) isa GeoSeries
            @test size(GeoData.chunk_series(ncarray)) == (1, 1, 1)
        end
        @testset "slice" begin
            @test_throws ArgumentError GeoData.slice(ncarray, Z)
            ser = GeoData.slice(ncarray, Ti)
            @test ser isa GeoSeries
            @test size(ser) == (24,)
            @test index(ser, Ti) == DateTime360Day(2001, 1, 16):Month(1):DateTime360Day(2002, 12, 16)
            @test bounds(ser) == ((DateTime360Day(2001, 1, 1), DateTime360Day(2003, 1, 1)),)
            A = ser[1]
            @test index(A, Y) == -79.5:89.5
            @test index(A, X) == 1.0:2:359.0
            @test bounds(A) == ((0.0, 360.0), (-80.0, 90.0))
        end
    end

    @testset "indexing with reverse lat" begin
        if !haskey(ENV, "CI") # CI downloads fail. But run locally
            ncrevlat = maybedownload("ftp://ftp.cdc.noaa.gov/Datasets/noaa.ersst.v5/sst.mon.ltm.1981-2010.nc")
            ncrevlatarray = GeoArray(ncrevlat; name=:sst, missingval=-9.96921f36)
            @test order(dims(ncrevlatarray, Y)) == Ordered(ReverseIndex(), ReverseArray(), ForwardRelation())
            @test ncrevlatarray[Y(At(40)), X(At(100)), Ti(1)] == missingval(ncrevlatarray)
            @test ncrevlatarray[Y(At(-40)), X(At(100)), Ti(1)] == ncrevlatarray[51, 65, 1] == 14.5916605f0
            @test val(span(ncrevlatarray, Ti)) == Month(1)
            @test val(span(ncrevlatarray, Ti)) isa Month # Not CompoundPeriod
        end
    end

    @testset "selectors" begin
        a = ncarray[X(At(21.0)), Y(Between(50, 52)), Ti(Near(DateTime360Day(2002, 12)))]
        @test bounds(a) == ((50.0, 52.0),)
        x = ncarray[X(Near(150)), Y(Near(30)), Ti(1)]
        @test x isa Float32
        dimz = X(Between(0.0, 360)), Y(Between(-80, 90)),
               Ti(Between(DateTime360Day(2001, 1, 1), DateTime360Day(2003, 01, 02)))
        @test size(ncarray[dimz...]) == (180, 170, 24)
        @test index(ncarray[dimz...]) == index(ncarray)
        nca = ncarray[Y(Between(-80, -25)), X(Between(0, 180)), Ti(Contains(DateTime360Day(2002, 02, 20)))]
        @test size(nca) == (90, 55)
        @test index(nca, Y) == index(ncarray[1:90, 1:55, 2], Y)
        @test all(nca .=== ncarray[1:90, 1:55, 14])
    end

    @testset "conversion to GeoArray" begin
        geoA = ncarray[X(1:50), Y(20:20), Ti(1)]
        @test size(geoA) == (50, 1)
        @test eltype(geoA) <: Union{Missing,Float32}
        @test geoA isa GeoArray{Union{Missing,Float32},2}
        @test dims(geoA) isa Tuple{<:X,<:Y}
        @test refdims(geoA) isa Tuple{<:Ti}
        @test metadata(geoA) == metadata(ncarray)
        @test ismissing(missingval(geoA))
        @test name(geoA) == :tos
    end

    @testset "write" begin
        @testset "to netcdf" begin
            # TODO save and load subset
            geoA = read(ncarray)
            @test size(geoA) == size(ncarray)
            filename = tempname() * ".nc"
            write(filename, geoA)
            saved = read(GeoArray(filename))
            @test size(saved) == size(geoA)
            @test refdims(saved) == refdims(geoA)
            @test missingval(saved) === missingval(geoA)
            @test map(metadata.(dims(saved)), metadata.(dims(GeoArray))) do s, g
                all(s .== g)
            end |> all
            @test metadata(saved) == metadata(geoA)
            @test all(metadata.(dims(saved)) == metadata.(dims(geoA)))
            @test GeoData.name(saved) == GeoData.name(geoA)
            @test all(mode.(dims(saved)) .!= mode.(dims(geoA)))
            @test all(order.(dims(saved)) .== order.(dims(geoA)))
            @test all(typeof.(span.(dims(saved))) .== typeof.(span.(dims(geoA))))
            @test all(val.(span.(dims(saved))) .== val.(span.(dims(geoA))))
            @test all(sampling.(dims(saved)) .== sampling.(dims(geoA)))
            @test typeof(dims(saved)) <: typeof(dims(geoA))
            @test index(saved, 3) == index(geoA, 3)
            @test all(val.(dims(saved)) .== val.(dims(geoA)))
            @test all(data(saved) .=== data(geoA))
            @test saved isa typeof(geoA)
            # TODO test crs
        end
        @testset "to gdal" begin
            gdalfilename = tempname() * ".tif"
            nccleaned = replace_missing(ncarray[Ti(1)], -9999.0)
            write(gdalfilename, nccleaned)
            gdalarray = GeoArray(gdalfilename)
            # gdalarray WKT is missing one AUTHORITY
            # @test_broken crs(gdalarray) == convert(WellKnownText, EPSG(4326))
            # But the Proj representation is the same
            @test convert(ProjString, crs(gdalarray)) == convert(ProjString, EPSG(4326))
            @test bounds(gdalarray) == (bounds(nccleaned)..., (1, 1))
            # Tiff locus = Start, Netcdf locus = Center
            @test reverse(index(gdalarray, Y)) .+ 0.5 ≈ index(nccleaned, Y)
            @test index(gdalarray, X) .+ 1.0  ≈ index(nccleaned, X)
            @test reverse(GeoArray(gdalarray); dims=Y()) ≈ nccleaned
        end
        @testset "to grd" begin
            nccleaned = replace_missing(ncarray[Ti(1)], -9999.0)
            write("testgrd.gri", nccleaned)
            grdarray = GeoArray("testgrd.gri");
            @test crs(grdarray) == convert(ProjString, EPSG(4326))
            @test bounds(grdarray) == (bounds(nccleaned)..., (1, 1))
            @test index(grdarray, Y) ≈ index(nccleaned, Y) .- 0.5
            @test index(grdarray, X) ≈ index(nccleaned, X) .- 1.0
            @test GeoArray(grdarray) ≈ reverse(nccleaned; dims=Y)
        end
    end

    @testset "show" begin
        sh = sprint(show, MIME("text/plain"), ncarray)
        # Test but don't lock this down too much
        @test occursin("GeoArray", sh)
        @test occursin("Y", sh)
        @test occursin("X", sh)
        @test occursin("Time", sh)
    end

    @testset "plot" begin
        ncarray[Ti(1:3:12)] |> plot
        ncarray[Ti(1)] |> plot
        ncarray[Y(100), Ti(1)] |> plot
    end

end

@testset "Single file stack" begin
    @time ncstack = GeoStack(ncmulti)

    @testset "load ncstack" begin
        @test ncstack isa GeoStack
        @test ismissing(missingval(ncstack))
        @test metadata(ncstack) isa Metadata{NCDfile}
        @test dims(ncstack[:abso4]) == dims(ncstack, (X, Y, Ti))
        @test refdims(ncstack) == ()
        # Loads child as a regular GeoArray
        @test_throws ErrorException ncstack[:not_a_key]
        @test ncstack[:albedo] isa GeoArray{<:Any,3}
        @test ncstack[:albedo, 2, 3, 1] isa Float32
        @test ncstack[:albedo, :, 3, 1] isa GeoArray{<:Any,1}
        @test dims(ncstack[:albedo]) isa Tuple{<:X,<:Y,<:Ti}
        @test keys(ncstack) isa NTuple{131,Symbol}
        @test keys(ncstack) == stackkeys
        @test first(keys(ncstack)) == :abso4
        @test metadata(ncstack) isa Metadata{NCDfile}
        @test metadata(ncstack)["institution"] == "Max-Planck-Institute for Meteorology"
        @test metadata(ncstack, :albedo) isa Metadata{NCDfile}
        @test metadata(ncstack, :albedo)["long_name"] == "surface albedo"
        # Test some DimensionalData.jl tools work
        # Time dim should be reduced to length 1 by mean
        @test axes(mean(ncstack[:albedo, Y(1:20)] , dims=Ti)) ==
              (Base.OneTo(192), Base.OneTo(20), Base.OneTo(1))
        geoA = ncstack[:albedo][Ti(4:6), X(1), Y(2)]
        @test geoA == ncstack[:albedo, Ti(4:6), X(1), Y(2)]
        @test size(geoA) == (3,)
    end

    if VERSION > v"1.1-"
        @testset "copy" begin
            geoA = read(ncstack[:albedo]) .* 2
            copy!(geoA, ncstack, :albedo);
            # First wrap with GeoArray() here or == loads from disk for each cell.
            # we need a general way of avoiding this in all disk-based sources
            @test geoA == read(ncstack[:albedo])
        end
    end

    @testset "indexing" begin
        ncmultistack = GeoStack(ncsingle)
        @test dims(ncmultistack[:tos]) isa Tuple{<:X,<:Y,<:Ti}
        @test ncmultistack[:tos] isa GeoArray{<:Any,3}
        @test ncmultistack[:tos][Ti(1)] isa GeoArray{<:Any,2}
        @test ncmultistack[:tos, Y(1), Ti(1)] isa GeoArray{<:Any,1}
        @test ncmultistack[:tos, 8, 30, 10] isa Float32
    end

    @testset "Subsetting keys" begin
        smallstack = subset(ncstack, (:albedo, :evap, :runoff))
        @test keys(smallstack) == (:albedo, :evap, :runoff)
    end

    # This is slow. We combine read/save to reduce test time
    # And it seems the memory is not garbage collected??
    @testset "read and write" begin
        @time st = read(ncstack)
        @test st isa GeoStack
        @test st.data isa NamedTuple
        @test first(st.data) isa Array
        length(dims(st[:aclcac]))
        filename = tempname() * ".nc"
        write(filename, st);
        saved = GeoStack(GeoStack(filename))
        @test keys(saved) == keys(st)
        @test metadata(saved)["advection"] == "Lin & Rood"
        @test metadata(saved) == metadata(st) == metadata(ncstack)
        @test all(first(DimensionalData.layers(saved)) .== first(DimensionalData.layers(st)))
    end

    @testset "show" begin
        ncstack = view(GeoStack(ncmulti), X(7:99), Y(3:90));
        sh = sprint(show, MIME("text/plain"), ncstack)
        # Test but don't lock this down too much
        @test occursin("GeoStack", sh)
        @test occursin("Y", sh)
        @test occursin("X", sh)
        @test occursin("Ti", sh)
        @test occursin(":tropo", sh)
        @test occursin(":tsurf", sh)
        @test occursin(":aclcac", sh)
        @test occursin("test_echam_spectral.nc", sh)
    end

end

@testset "series" begin
    ncseries = GeoSeries([ncsingle, ncsingle], (Ti,); child=GeoStack)
    @testset "read" begin
        geoseries = read(ncseries)
        @test geoseries isa GeoSeries{<:GeoStack}
        @test geoseries.data isa Vector{<:GeoStack}
    end
    geoA = GeoArray(ncsingle; key=:tos)
    @test all(read(ncseries[Ti(1)][:tos]) .=== read(geoA))
end

nothing


module Zoo

using Photon

include("alexnet.jl")
export AlexNet

include("densenet.jl")
export DenseNet121, DenseNet161, DenseNet169, DenseNet201

include("vgg.jl")
export VGG11, VGG13, VGG16, VGG19

end


"""
```
function old_to_new_cloud(cloud::DSGE.Cloud)::SMC.Cloud
```
Converter between types.
"""
function old_to_new_cloud(cloud::DSGE.Cloud)::SMC.Cloud
    return SMC.Cloud(cloud.particles, cloud.tempering_schedule,
                     cloud.ESS, cloud.stage_index, cloud.n_Φ, cloud.resamples,
                     cloud.c, cloud.accept, cloud.total_sampling_time)
end

"""
```
function Cloud(c::ParticleCloud)
function Cloud(c::Cloud)
```
Returns a Cloud type.
"""
function SMC.Cloud(cloud::ParticleCloud)
    return SMC.Cloud(hcat(Matrix{Float64}(get_vals(cloud)'), get_loglh(cloud), get_logprior(cloud),
                      get_old_loglh(cloud), get_accept(cloud), get_weights(cloud)),
                 cloud.tempering_schedule, cloud.ESS, cloud.stage_index, cloud.n_Φ,
                 cloud.resamples, cloud.c, cloud.accept, cloud.total_sampling_time)
end
function SMC.Cloud(cloud::SMC.Cloud)
    return cloud
end

function DSGE.Cloud(cloud::SMC.Cloud)
    return DSGE.Cloud(cloud.particles, cloud.tempering_schedule,
                     cloud.ESS, cloud.stage_index, cloud.n_Φ, cloud.resamples,
                     cloud.c, cloud.accept, cloud.total_sampling_time)
end
function SMC.Cloud(cloud::DSGE.Cloud)
    return SMC.Cloud(cloud.particles, cloud.tempering_schedule,
                     cloud.ESS, cloud.stage_index, cloud.n_Φ, cloud.resamples,
                     cloud.c, cloud.accept, cloud.total_sampling_time)
end


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

using ArgParse
using Distributed, ClusterManagers
include("Utils.jl")

function parse_commandline()
  s = ArgParseSettings()
  @add_arg_table! s begin
    "--nj"
    help = "number of jobs"
    arg_type = Int
    default = 1
    "--conf"
    help = "configuration setup"
    arg_type = String
    default = "2a"
  end
  return parse_args(ARGS, s)
end

# parse command line and add workers
parsed_args = parse_commandline()
nj, setup = parsed_args["nj"], parsed_args["conf"]
feat_dir = get_feat_dir(setup)
check_env()
if ispath(feat_dir)
  error("Data already processed for this feature conf in $(feat_dir). Remove this folder to re-run feature extraction from scratch.")
end

if nj > 1
  addprocs_sge(nj;
               qsub_flags=split(ENV["CPU_CMD"]),
               wd=mktempdir(pwd()),
               exeflags="--project"
              )
end

@everywhere begin
  setup = $setup
  using DSP, MFCC, HMMGradients, FiniteStateTransducers
  using JLD2, LibSndFile, FileIO
  include("Utils.jl")
  include("WFSTs.jl")
  include("DatasetParsers.jl")
end

@everywhere function process_data(dataset_path,lexicon,L,H,Fs,subsample,
    uttID2file,uttID2text,feat_dir,set,nj)
  T = Float32

  uttID2feats  = Dict()
  uttID2phones = Dict()
  uttID2tr     = Dict{String,Vector{Pair{Vector{Int},Vector{Int}}}}()
  wsym = get_osym(L)

  for uttID in keys(uttID2file)
    # process audio
    x = uttID2file[uttID]
    x = get_feats(x)
    uttID2feats[uttID] = x
    x = feats_post(x)
    Nt = size(x,1)

    # process text
    text = uttID2text[uttID]
    uttID2phones[uttID] = text2phones(lexicon,text)
    text = split(text;keepempty=false)
    text = String.(vcat("<SIL>",[[ti,"<SIL>"] for ti in text]...)) #silence between every word
    S = linearfst(text,text, ones(typeofweight(L),length(text)), wsym, wsym)
    HLS = rm_eps!(H∘(L∘S))
    Nt2 = subsample == 1 ? Nt : ceil(Int,Nt/3)
    time2tr = wfst2tr(HLS,Nt2)
    uttID2tr[uttID] = HMMGradients.t2tr2t2IJ(time2tr)
  end
  if nj > 1
    q = joinpath(feat_dir,"q_split_$set")
    mkpath(q)
    JLD2.@save joinpath(q,"$(myid()).jld2") uttID2feats uttID2tr uttID2phones
  else
    return uttID2feats, uttID2tr, uttID2text, uttID2phones
  end
end

function prepare_data(dataset_path,lexicon,L,H,Fs,subsample,set,feat_dir,nj)
  println("Processing $set set with $nj jobs")
  uttID2file  = get_uttID2file(dataset_path,set)
  uttID2text  = get_uttID2text(uttID2file)
  if nj > 1
    # split utterances
    uttIDs = [keys(uttID2text)...]
    Nu = length(uttIDs)
    delta = div(Nu,nj)
    uttIDss = [uttIDs[1+(i-1)*delta:(i==nj ? Nu : i*delta)] for i = 1:nj]
    uttID2files = [filter(x -> x.first in uttIDs, uttID2file) for uttIDs in uttIDss]
    uttID2texts = [filter(x -> x.first in uttIDs, uttID2text) for uttIDs in uttIDss]
    pmap(
         uttID2filetext ->
         process_data(dataset_path,lexicon,L,H,Fs,subsample,
                      uttID2filetext[1],uttID2filetext[2],
                      feat_dir,set,nj),
         zip(uttID2files,uttID2texts)
        )
    uttID2feats  = Dict()
    uttID2phones = Dict()
    uttID2tr     = Dict{String,Vector{Pair{Vector{Int},Vector{Int}}}}()
    println("Merging files")
    q = joinpath(feat_dir,"q_split_$set")
    for id in workers()
      data = load(joinpath(q,"$id.jld2"))
      uttID2feats_nj, uttID2tr_nj, uttID2phones_nj =
      data["uttID2feats"], data["uttID2tr"], data["uttID2phones"]
      merge!(uttID2feats , uttID2feats_nj )
      merge!(uttID2phones, uttID2phones_nj)
      merge!(uttID2tr    , uttID2tr_nj    )
    end
    rm(q;recursive=true)
    return uttID2feats, uttID2tr, uttID2text, uttID2phones
  else
    process_data(dataset_path,lexicon,L,H,Fs,subsample,uttID2file,uttID2text,feat_dir,set,nj)
  end
end

###
@everywhere begin
  include("conf/$(setup)/feat_conf.jl")
  lexicon, ilexicon = get_lexicon()
  H, L = get_HL(lexicon)
end

dataset_path = ENV["TIDIGITS_PATH"]
T = @elapsed uttID2feats_train, uttID2tr_train, uttID2text_train, uttID2phones_train =
prepare_data(dataset_path,lexicon,L,H,Fs,subsample,"train",feat_dir,nj)
println("Done in $T sec")
T = @elapsed uttID2feats_test, uttID2tr_test, uttID2text_test, uttID2phones_test =
prepare_data(dataset_path,lexicon,L,H,Fs,subsample,"test",feat_dir,nj)
println("Done in $T sec")

if nj > 1
  t = rmprocs(workers())
  wait(t)
end

# test data in TIDIGITS has same size of train, so we repartition it
uttID2feats_all = merge(uttID2feats_train , uttID2feats_test )
uttID2tr_all    = merge(uttID2tr_train    , uttID2tr_test    )
uttID2text_all  = merge(uttID2text_train  , uttID2text_test  )
uttID2phones_all= merge(uttID2phones_train, uttID2phones_test)
uttIDs_all = [keys(uttID2text_all)...]
Nu = length(uttIDs_all)
idx_train, idx_test = round(Int,Nu*0.7), round(Int,Nu*0.9)
set2uttID  = Dict()
set2uttID["train"], set2uttID["test"], set2uttID["dev"] = uttIDs_all[1:idx_train], uttIDs_all[idx_train+1:idx_test], uttIDs_all[idx_test+1:end]

mkpath(feat_dir)
for set in ("train","test","dev")
  filename = set
  uttID2feats   = filter(x -> x.first in set2uttID[set], uttID2feats_all )
  uttID2tr      = filter(x -> x.first in set2uttID[set], uttID2tr_all    )
  uttID2text    = filter(x -> x.first in set2uttID[set], uttID2text_all  )
  uttID2phones  = filter(x -> x.first in set2uttID[set], uttID2phones_all)
  JLD2.@save joinpath(feat_dir,"$set.jld2") uttID2feats uttID2tr uttID2text uttID2phones
end


using Base
using Interpolations

# While the StateID struct and its helper functions may seem a little over-engineered, they
# primarily exist to document the iSLP and iLV quantities used by TOPBase

"""
    StateID(iSLP, iLV)

`StateID` represents a given state (electron configuration) of an ion species. The arguments are
taken from the TOPBase.

# Arguments
- `iSLP::UInt16`: An integer representing the quantum atomic numbers SLπ (total spin, orbital
  angular momentum, and parity). The value is given by `100*(2*S+1) + 10*L + π`. Note that TOPBase
  documentation refers to the parity quantum number as P or π, but for consistency with how TOPBase
  presents the data, we prefer P in our implementation code.
- `iLV::UInt8`: Distinguishes between configurations that have the same values of `iSLP`. TOPBase
  seems to index the configuration based on the energy level of the configuration (with `iLV = 1`
  corresponding to the lowest energy).

# Notes
The combination of this struct and an ion species (or neutral element) should be enough to fully
specify a given electron configuration. For example, consider He I:
- `iSLP = 100, iLV = 1` has the electron configuration: 1s²
- `iSLP = 100, iLV = 2` has the electron configuration: 1s 2s
- `iSLP = 100, iLV = 3` has the electron configuration: 1s 3s
- `iSLP = 111, iLV = 1` has the electron configuration: 1s 2p
- `iSLP = 111, iLV = 2` has the electron configuration: 1s 3p
"""
struct StateID
    iSLP::UInt16
    iLV::UInt8
    function StateID(iSLP, iLV)
        @assert 99 < iSLP < 992 # this upper limit is set based on the data source
        @assert (iSLP % 10) < 2 # trailing digit must be 0 or 1
        @assert 0 < iLV # the upper limit is set based on the data source
        new(iSLP, iLV)
    end
end

Base.:(<)(a::StateID, b::StateID) = (a.iSLP < b.iSLP) || ((a.iSLP == b.iSLP) && (a.iLV < b.iLV))
Base.show(io::IO, obj::StateID) = print(io, "StateID{iSLP = ", obj.iSLP, ", iLV = ", obj.iLV, "}")

_parity_vals = ("even", "odd") # predefine these to avoid allocations

function _quantum_prop_type(state_id::StateID)
    S_code = div(state_id.iSLP, 100) # hundreds place of state_id.iSLP
    twice_S = S_code - 1
    L = div(state_id.iSLP - S_code*100, 10) # tens place of state_id.iSLP
    parity = _parity_vals[(state_id.iSLP % 2) + 1] # 0 or 1 in ones place of state_id.iSLP
    (Int32(twice_S), Int32(L), parity)
end
get_S_quantum_num(state_id::StateID) = _quantum_prop_type(state_id)[1] / 2.0f0
get_L_quantum_num(state_id::StateID) = _quantum_prop_type(state_id)[2]
get_parity(state_id::StateID) = _quantum_prop_type(state_id)[3]

function get_statistical_weight(state_id::StateID)
    twice_S, L, _ = _quantum_prop_type(state_id)
    (twice_S+1)*(2*L+1)
end

"""
    StateProp(ion_energy_ryd, excitation_potential_ryd, statistical_weight, electron_config)

`StateProp` holds data associated with different electron configurations. It's used to hold
data read from the TOPBase cross-section queries and to compute the statistical weight.
"""
struct StateProp
    ion_energy_ryd::Float64 # energy with respect to ionization potential in Rydbergs
    excitation_potential_ryd::Float64 # energy above ground state in Rydbergs
    statistical_weight::UInt16
    electron_config::String
end

# define some generally useful functions related to parsing
function _parse_species_stateid(str::AbstractString, start_Z::Integer,
                                start_numElectrons::Integer, start_iSLP::Integer,
                                start_iLV::Integer, last_iLV::Integer)
    # parses the atomic_num, number of electrons, iSLP, and iLV entries of the table
    atomic_num = parse(UInt8, str[start_Z:start_numElectrons-1])
    num_electrons = parse(UInt8, str[start_numElectrons:start_iSLP-1])
    iSLP = parse(UInt16, str[start_iSLP:start_iLV-1])
    iLV = parse(UInt16, str[start_iLV:last_iLV])
    (Species(Formula(atomic_num), atomic_num - num_electrons), StateID(iSLP, iLV))
end

# Functionallity for reading photoionization cross sections. Since these tables can be massive, we
# use an iterator to do this. This is based on the design used in the Julia standard library for
# implementing the ``EachLine`` iterable object. That code used the following license:
#
# Copyright (c) 2009-2021: Jeff Bezanson, Stefan Karpinski, Viral B. Shah,
# and other contributors:
#
# https://github.com/JuliaLang/julia/contributors
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

const _TABLE_BREAK = " ================================================"
const _HEADER_LINE_LENGTH = length(_TABLE_BREAK)

struct EachPhotoIonSubtable{IOT <: IO}
    stream::IOT
    ondone::Function
end

"""
    each_photo_ion_subtable(filename::AbstractString)

Create an iterable `EachPhotoIonSubtable` object for parsing photoionization cross-sections.

The input file is organized as a series of subtables. Each subtable provides the photoionization
energies (make sure this isn't the energy of the photoionized electron) in Ryd, and the associated
cross section (in megabarns) for a different state.

Each entry in the `EachPhotoIonSubtable` iterable provides a named tuple with the keys:
- name: holds the species name
- state_id: holds a `StateID` instance
- photon_energy_ryd: the photoionization values (check), in Ryd, stored in an array of Float32 vals
- cross_section_MegaBarn: the associated photoionization cross-sections. This array should have the
  same shape and datatype as photon_energy_ryd

The underlying file is only closed after the `EachPhotoIonSubtable` is garbage collected.

#Notes
This implementation is not robust enough to deal with blank lines at the start or end of the file.

The table should have been produced by a
[TOPbase Photoionisation Cross Section query](http://cdsweb.u-strasbg.fr/topbase/xsections.html)
that uses the following query options:
- Only a single atomic number and a single number of electrons is included in the query (other
  choices may yield too many results, which causes TOPBase to truncate output).
- There isn't any limitation on the number the quantum numbers, the energy or the levels
- The query request levels ordered in the "Level order in each series"
"""
function each_photo_ion_subtable(filename::AbstractString)
    stream = open(filename, "r")
    # confirm that standard file header is present and move the position of the
    # stream past this point
    @assert readline(stream) == _TABLE_BREAK
    @assert readline(stream) == "       I  NZ  NE  ISLP  ILV        E(RYD)      NP"
    @assert readline(stream) == _TABLE_BREAK
    EachPhotoIonSubtable(stream, ()->close(stream))::EachPhotoIonSubtable
end

function Base.iterate(itr::EachPhotoIonSubtable, state=nothing)
    next_line = readline(itr.stream)

    if next_line == _TABLE_BREAK
        return (itr.ondone(); nothing)
    end

    # first parse the subtable's header
    species, state_id = _parse_species_stateid(next_line, 9, 13, 17, 23, 27)
    ion_energy_ryd = parse(Float32, next_line[28:41])

    num_entries = parse(UInt32, next_line[42:49])

    # this could probably be 32-bit.
    photon_energy_ryd = Array{Float32,1}(undef, num_entries)
    cross_section_MegaBarn = Array{Float32,1}(undef, num_entries)
    for i = 1:num_entries # this properly handles the case when num_entries == 0
        cur_line = readline(itr.stream)
        photon_energy_ryd[i] = parse(Float32, cur_line[1:14])
        cross_section_MegaBarn[i] = parse(Float32, cur_line[15:24])
    end

    # parse the species, and the level
    ((species = species, state_id=state_id, ion_energy_ryd = ion_energy_ryd,
      photon_energy_ryd_arr = photon_energy_ryd,
      cross_section_MegaBarn_arr = cross_section_MegaBarn), nothing)
end

# this and the following definition implement Julia's informal iterator interface
Base.eltype(::Type{<:EachPhotoIonSubtable}) =
    Tuple{String, StateID, Float32, Array{Float32,1}, Array{Float32,1}}
Base.IteratorSize(::Type{<:EachPhotoIonSubtable}) = SizeUnknown()

# In a couple tables, there's a place where the photon energy is listed out of order. This seems to
# be okay given that this is in a section of the table where the cross-section is zero.
const _species_with_problematic_subtables = (Species("He_I"), Species("C_I"), Species("C_II"),
                                             Species("Al_II"), Species("O_II"))

function _interpolator_from_subtable(table_photon_energy_ryd::Vector{Float32},
                                     table_cross_section_MBarn::Vector{Float32},
                                     species::Union{Species,Nothing} = nothing,
                                     extrapolation_bc = 0.0)
    if species in _species_with_problematic_subtables

        for j in 1:length(table_photon_energy_ryd)-1
            if table_photon_energy_ryd[j+1] < table_photon_energy_ryd[j]
                @assert j > 1
                @assert (j + 2) <= length(table_photon_energy_ryd)
                @assert table_photon_energy_ryd[j+2] >= table_photon_energy_ryd[j+1]
                @assert table_photon_energy_ryd[j+2] >= table_photon_energy_ryd[j]
                @assert table_photon_energy_ryd[j+1] >= table_photon_energy_ryd[j-1]
                @assert table_photon_energy_ryd[j] >= table_photon_energy_ryd[j-1]

                @assert table_cross_section_MBarn[j+2] == 0.0
                @assert table_cross_section_MBarn[j+1] == 0.0
                @assert table_cross_section_MBarn[j] == 0.0
                @assert table_cross_section_MBarn[j-1] == 0.0

                # swap the order of the 2 values
                temp = table_photon_energy_ryd[j]
                table_photon_energy_ryd[j] = table_photon_energy_ryd[j+1]
                table_photon_energy_ryd[j+1] = temp
            end
        end
    end
    LinearInterpolation(table_photon_energy_ryd, table_cross_section_MBarn,
                        extrapolation_bc=extrapolation_bc)
end

function _try_env_prefix_path(prefix_env_var, fname)
    prefix = get(ENV, prefix_env_var, nothing)
    if isnothing(prefix)
        e = ErrorException(string("The \"", prefix_env_var, "\" environment variable was not set"))
        throw(e);
    end
    path = joinpath(prefix, fname)
    if !isfile(path) # this properly handles the case where path is a symlink to a file
        throw(ErrorException(string("There is no file at ", path)));
    end
    path
end

# return the minimum iterator over cross_section subtables and the ionization energy of the ground
# state (according to the table)
function _get_cross_sec_data(species::Species,
                             cross_sec_file::Union{AbstractString,Nothing} = nothing)

    # get the photo-ionization subtable iterator:
    my_cross_sec_file = if isnothing(cross_sec_file)
        # determine the string representation of Species
        # we don't directly use repr(species) since the representation could change in the future
        # to not use roman numeral (e.g. if support is added for negatively charged species)
        formula_str = repr("text/plain", species.formula)
        roman_numeral = get_roman_numeral(species)
        fname = string(formula_str, "_", roman_numeral, ".txt")

        println("cross_sec_file isn't specified. Searching for file in the directory specified ",
                "by the KORG_OP_DIR environment variable.")
        _try_env_prefix_path("KORG_OP_DIR", fname)
    else
        cross_sec_file
    end

    # determine the ionization energy of the ground state (from the subtables)
    gs_ion_energy_ryd = begin # this could be optimized significantly!
        cur_min = floatmax(Float32)
        for (cur_species, _, ion_energy_ryd, _, _) in each_photo_ion_subtable(my_cross_sec_file)
            @assert (cur_species == species) # this will catch mistakes in the query that generated
                                             # the cross_section data file
            cur_min = min(ion_energy_ryd, cur_min)
        end
        cur_min
    end

    subtable_iter = each_photo_ion_subtable(my_cross_sec_file)
    subtable_iter, gs_ion_energy_ryd
end

@doc raw"""
    weighted_bf_cross_section_TOPBase(λ_vals, T_vals, species; extrapolation_bc = 0.0
                                      cross_sec_file = nothing,
                                      partition_func = partition_funcs[species])

Calculate the combined bound-free cross section (in cm²) of an ion species using data from the
Opacity Project. The result is the mean of the the bound-free cross section of each electron state,
weighted by their Boltzmann probabilities. It includes the LTE correction for stimulated emission.
The result is stored in a 2D array with a size given by `(length(λ_vals), length(T_vals))`.

# Arguments
- `λ_vals::AbstractVector`: wavelengths (in cm) to compute the cross sections at
- `T_vals::AbstractVector`: temperatures (in K) to compute the cross sections at
- `species::Species`: the starting atom or atomic ion of the bound-free reaction.

# Keyword Arguments
- `extrapolation_bc`: Specifies handling of extrapolation during linear interpolation (this is
  passed to `Interpolations.LinearInterpolation`).
- `cross_sec_file`: Optional keyword that specifies the path to the text file that holds the
  cross-section subtables. When `nothing` (the default), this function tries to load the file
  from `$KORG_OP_DIR/<atomic-symbol>_<ion-state>.txt`. `<atomic-symbol>` is replaced with the
  atomic symbol of the species (e.g. `H`, `He`, `Al`) and `<ion-state>` is replaced with the
  capital Roman numeral(s) specifying the ionization state.
- `partition_func`: Specifies the partition function for the current species. This should be a
  callable that accepts temperature as an argument. When this is `nothing`, it falls back to the
  default partition functions.

# Explanation
In more mathematical rigor, this function basically evaluates the following equation:

``\overline{\sigma_{\rm bf}}(\lambda,T) = \sum_{{\bf n}\in{\rm states}} \frac{ g_{\bf n}}{U(T)} e^{-E_{\bf n}/(k_B T)}\left(1 - e^{-h c/(\lambda k_B T)}\right) \sigma_{{\rm bf}, {\bf n}}(\lambda)``

In this equation:
- ``g_{\bf n}`` is the statistical weight of state ``{\bf n}``.
- ``U`` is the partition function.
- ``E_{\bf n}`` is the excitation potential of state ``{\bf n}``. In other words, its the energy
  difference between state ``{\bf n}`` and the ground state.
- ``\sigma_{\lambda,{\rm bf}, {\bf n}}`` is the bound-free cross-section (a.k.a. photo-ionization
  cross section) of state ``{\bf n}``.

The linear aborption coefficient for bound-free absorption is simply:
``\alpha_{\lambda,{\rm bf}} = n_{\rm tot} \overline{\sigma_{\rm bf}}(\lambda,T)``.

# Notes
In the future, it needs to be possible to pass an argument that adjusts the energy levels of
different electron configurations.
"""
function weighted_bf_cross_section_TOPBase(λ_vals::AbstractVector, T_vals::AbstractVector,
                                           species::Species; extrapolation_bc=0.0,
                                           cross_sec_file::Union{AbstractString,Nothing} = nothing,
                                           partition_func = partition_funcs[species])
    # Note: it would not take a lot of work to make this work with 32-bit floats (since the tables
    # only store 32-bit floats). If doing that, it might be wise to accept λ in units of ångstroms
    # (for precision purposes)

    # TODO: make it possible to pass in a dict mapping dict names to the various energy levels
    #       (to allow for overwriting the values in the table with empirical values)

    @assert !ismolecule(species) # species must be an atom or atomic ion

    # precompute Temperature-dependent constant
    inv_partition_func_val = 1.0./partition_func.(T_vals)

    β_Ryd = RydbergH_eV./(kboltz_eV .* T_vals)

    # convert λ_vals to photon energies
    photon_energies = (hplanck_eV * c_cgs / RydbergH_eV) ./ λ_vals

    # prepare the output array where results will be accumulated
    weighted_average = zeros(eltype(photon_energies), (length(photon_energies),length(T_vals)))

    cross_section_subtables, gs_ion_energy_ryd = _get_cross_sec_data(species, cross_sec_file)

    for subtable in cross_section_subtables
        if length(subtable.photon_energy_ryd_arr) == 0
            continue
        end

        cross_section_interp = _interpolator_from_subtable(subtable.photon_energy_ryd_arr,
                                                           subtable.cross_section_MegaBarn_arr,
                                                           species, extrapolation_bc)
        statistical_weight = get_statistical_weight(subtable.state_id)
        excitation_potential_ryd = subtable.ion_energy_ryd - gs_ion_energy_ryd

        # now iterate over Temperatures
        for j in 1:length(T_vals)

            # now compute the weighting of the current state under LTE
            weight = inv_partition_func_val[j] * statistical_weight * exp(-excitation_potential_ryd
                                                                          * β_Ryd[j])

            current_cross_section = ( cross_section_interp.(photon_energies) .*
                                      (1.0 .- exp.(-photon_energies.*β_Ryd[j])) )
            view(weighted_average, :, j) .+= (current_cross_section .* weight)
        end
    end

    # convert weighted_average units from megabarn to cm²
    weighted_average .*= 1e-18
    weighted_average
end

"""
    absorption_coef_bf_TOPBase(λ, T, ndens_species, species; kwargs...)

Computes the bound-free linear absorption coefficient, α, (in units of cm⁻¹) using data from the
Opacity Project (downloaded from TOPBase) and assuming LTE. The result is stored in a 2D array with
a size given by `(length(λ), length(T))`.

# Arguments
- `λ::AbstractVector`: wavelengths (in cm) to compute α at.
- `T::AbstractVector`: temperatures (in K) to compute α at.
- `ndens_species::AbstractVector`: number density (in cm⁻³) of the species at each temperature
  (this should have the same length as T).
- `species::Species`: the starting atom or atomic ion of the bound-free reaction.

# Keyword Arguments
This function accepts all the same keyword arguments as `weighted_bf_cross_section_TOPBase`.
"""
function absorption_coef_bf_TOPBase(λ::AbstractVector, T::AbstractVector,
                                    ndens_species::AbstractVector, species::Species; kwargs...)
    @assert size(T) == size(ndens_species)

    arr = weighted_bf_cross_section_TOPBase(λ, T, species; kwargs...)

    # to save time and memory, let's modify the values of arr in-place
    for (i,n) in enumerate(ndens_species)
        view(arr,:,i) .*= n
    end

    arr
end


module GeometricIntegratorsDiffEq

using Reexport
@reexport using DiffEqBase

using GeometricIntegrators

const warnkeywords =
    (:save_idxs, :d_discontinuities, :unstable_check, :save_everystep,
     :save_end, :initialize_save, :adaptive, :abstol, :reltol, :dtmax,
     :dtmin, :force_dtmin, :internalnorm, :gamma, :beta1, :beta2,
     :qmax, :qmin, :qsteady_min, :qsteady_max, :qoldinit, :failfactor,
     :maxiters, :isoutofdomain, :unstable_check,
     :calck, :progress, :tstops, :saveat, :dense)

function __init__()
    global warnlist = Set(warnkeywords)
end

include("algorithms.jl")
include("solve.jl")

export GeometricIntegratorAlgorithm, GIEuler, GIMidpoint, GIHeun, GIKutta,
       GIKutta, GIERK4, GIERK438, GIImplicitEuler, GIImplicitMidpoint,
       GIRadIIA2, GIRadIIA3, GISRK3, GIGLRK, GISymplecticEulerA, GISymplecticEulerB,
       GILobattoIIIAIIIB2, GILobattoIIIBIIIA2

end # module


using Oceananigans.Diagnostics: AbstractDiagnostic
using Oceananigans.Fields: FieldSlicer
using Oceananigans.OutputWriters: fetch_output
using Oceananigans.Utils: AbstractSchedule, prettytime

import Oceananigans: run_diagnostic!
import Oceananigans.Utils: TimeInterval
import Oceananigans.Fields: location

"""
    mutable struct AveragedTimeInterval <: AbstractSchedule

Container for parameters that configure and handle time-averaged output.
"""
mutable struct AveragedTimeInterval <: AbstractSchedule
                       interval :: Float64
                         window :: Float64
                         stride :: Int
    previous_interval_stop_time :: Float64
                     collecting :: Bool
end

"""
    AveragedTimeInterval(interval; window=interval, stride=1)

Returns a `schedule` that specifies periodic time-averaging of output.
The time `window` specifies the extent of the time-average, which
reoccurs every `interval`.

`output` is computed and accumulated into the average every `stride` iterations
during the averaging window. For example, `stride=1` computs output every iteration,
whereas `stride=2` computes output every other iteration. Time-averages with
longer `stride`s are faster to compute, but less accurate.

The time-average of ``a`` is a left Riemann sum corresponding to

`` ⟨a⟩ = 1/T \\int_{tᵢ-T}^T a \\mathrm{d} t ,``

where ``⟨a⟩`` is the time-average of ``a``, ``T`` is the time-window for averaging,
and the ``tᵢ`` are discrete times separated by the time `interval`. The ``tᵢ`` specify
both the end of the averaging window and the time at which output is written.

Example
=======

```jldoctest averaged_time_interval
using Oceananigans.OutputWriters: AveragedTimeInterval
using Oceananigans.Utils: year, years

schedule = AveragedTimeInterval(4years, window=1year)

# output
AveragedTimeInterval(window=1 year, stride=1, interval=4 years)
```

An `AveragedTimeInterval` schedule directs an output writer
to time-average its outputs before writing them to disk:

```jldoctest averaged_time_interval
using Oceananigans
using Oceananigans.OutputWriters: JLD2OutputWriter
using Oceananigans.Utils: minutes

model = NonhydrostaticModel(grid=RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1)))

simulation = Simulation(model, Δt=10minutes, stop_time=30years)

simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocities,
                                                          prefix = "averaged_velocity_data",
                                                          schedule = AveragedTimeInterval(4years, window=1year, stride=2))

# output
JLD2OutputWriter scheduled on TimeInterval(4 years):
├── filepath: ./averaged_velocity_data.jld2
├── 3 outputs: (:u, :v, :w) averaged on AveragedTimeInterval(window=1 year, stride=2, interval=4 years)
├── field slicer: FieldSlicer(:, :, :, with_halos=false)
├── array type: Array{Float32}
├── including: [:grid, :coriolis, :buoyancy, :closure]
└── max filesize: Inf YiB
```
"""
function AveragedTimeInterval(interval; window=interval, stride=1)

    window > interval && throw(ArgumentError("Averaging window $window is greater than the output interval $interval."))

    return AveragedTimeInterval(Float64(interval), Float64(window), stride, 0.0, false)
end

# Determines whether or not to call run_diagnostic
(schedule::AveragedTimeInterval)(model) =
    schedule.collecting || model.clock.time >= schedule.previous_interval_stop_time + schedule.interval - schedule.window

TimeInterval(schedule::AveragedTimeInterval) = TimeInterval(schedule.interval)

Base.copy(schedule::AveragedTimeInterval) = AveragedTimeInterval(schedule.interval, window=schedule.window, stride=schedule.stride)

"""
    WindowedTimeAverage{OP, R, FS} <: AbstractDiagnostic

An object for computing 'windowed' time averages, or moving time-averages
of a `operand` over a specified `window`, collected on `interval`.
"""
mutable struct WindowedTimeAverage{OP, R, FS} <: AbstractDiagnostic
                      result :: R
                     operand :: OP
           window_start_time :: Float64
      window_start_iteration :: Int
    previous_collection_time :: Float64
                field_slicer :: FS
                    schedule :: AveragedTimeInterval
end

"""
    WindowedTimeAverage(operand, model=nothing; schedule, field_slicer=FieldSlicer())

Returns an object for computing running averages of `operand` over `schedule.window` and
recurring on `schedule.interval`, where `schedule` is an `AveragedTimeInterval`.
During the collection period, averages are computed every `schedule.stride` iteration.

`operand` may be a `Oceananigans.Field` or a function that returns an array or scalar.

Calling `wta(model)` for `wta::WindowedTimeAverage` object returns `wta.result`.
"""
function WindowedTimeAverage(operand, model=nothing; schedule, field_slicer=FieldSlicer())

    output = fetch_output(operand, model, field_slicer)
    result = similar(output) # convert views to arrays
    result .= output # initialize `result` with initial output

    return WindowedTimeAverage(result, operand, 0.0, 0, 0.0, field_slicer, schedule)
end

# Time-averaging doesn't change spatial location
location(wta::WindowedTimeAverage) = location(wta.operand)

function accumulate_result!(wta, model)

    # Time increment:
    Δt = model.clock.time - wta.previous_collection_time

    # Time intervals:
    T_current = model.clock.time - wta.window_start_time
    T_previous = wta.previous_collection_time - wta.window_start_time

    # Accumulate left Riemann sum
    integrand = fetch_output(wta.operand, model, wta.field_slicer)

    @. wta.result = (wta.result * T_previous + integrand * Δt) / T_current

    # Save time of integrand collection
    wta.previous_collection_time = model.clock.time

    return nothing
end

function run_diagnostic!(wta::WindowedTimeAverage, model)

    if model.clock.iteration == 0 # initialize previous interval stop time
        wta.schedule.previous_interval_stop_time = model.clock.time
    end

    # Don't start collecting if we are *only* "initializing" run_diagnostic! at the beginning
    # of a Simulation.
    #
    # Note: this can be false at the zeroth iteration if interval == window (which
    # implies we are always collecting)

    unscheduled = model.clock.iteration == 0 &&
        model.clock.time < wta.schedule.previous_interval_stop_time + wta.schedule.interval - wta.schedule.window

    if unscheduled
        # This is an "unscheduled" call to run_diagnostic! --- which occurs when run_diagnostic!
        # is called at the beginning of a run (and schedule.interval != schedule.window).
        # In this case we do nothing.
    elseif !(wta.schedule.collecting)
        # run_diagnostic! has been called on schedule, but we are not currently collecting data.
        # Initialize data collection:

        # Start averaging period
        wta.schedule.collecting = true

        # Zero out result
        wta.result .= 0

        # Save averaging start time and the initial data collection time
        wta.window_start_time = model.clock.time
        wta.window_start_iteration = model.clock.iteration
        wta.previous_collection_time = model.clock.time

    elseif model.clock.time >= wta.schedule.previous_interval_stop_time + wta.schedule.interval
        # Output is imminent. Finalize averages and cease data collection.
        accumulate_result!(wta, model)

        # Averaging period is complete.
        wta.schedule.collecting = false

        # Reset the "previous" interval time, subtracting a sliver that presents overshoot from accumulating.
        wta.schedule.previous_interval_stop_time = model.clock.time - rem(model.clock.time, wta.schedule.interval)

    elseif mod(model.clock.iteration - wta.window_start_iteration, wta.schedule.stride) == 0
        # Collect data as usual
        accumulate_result!(wta, model)
    end

    return nothing
end

function (wta::WindowedTimeAverage)(model)

    # For the paranoid
    wta.schedule.collecting &&
        model.clock.iteration > 0 &&
        @warn "Returning a WindowedTimeAverage before the collection period is complete."

    return wta.result
end

Base.show(io::IO, schedule::AveragedTimeInterval) = print(io, summary(schedule))

Base.summary(schedule::AveragedTimeInterval) = string("AveragedTimeInterval(",
                                                      "window=", prettytime(schedule.window), ", ",
                                                      "stride=", schedule.stride, ", ",
                                                      "interval=", prettytime(schedule.interval),  ")")

show_averaging_schedule(schedule) = ""
show_averaging_schedule(schedule::AveragedTimeInterval) = string(" averaged on ", summary(schedule))

output_averaging_schedule(output::WindowedTimeAverage) = output.schedule


function ARMA(y,u,regx,regu,modelform)
# y = streams of outputs
# u = streams of inputs
# regx = number of regression steps in x
# regu = number of regression steps in u
# modelform = "linear","bilinear",etc
# returns ARMA parameters

params = [];

return params;


module TestGroupBy
include("preamble.jl")
using Transducers: DefaultInit

@testset begin
    @test foldl(right, GroupBy(string, Map(last), push!!), [1, 2, 1, 2, 3]) ==
          Dict("1" => [1, 1], "2" => [2, 2], "3" => [3])

    @test foldl(
        right,
        GroupBy(
            identity,
            opcompose(Map(last), Scan(+)),
            (_, x) -> x > 3 ? reduced(x) : x,
            nothing,
        ),
        [1, 2, 1, 2, 3],
    ) == Dict(2 => 4, 1 => 2)
end

@testset "post-groupby filtering" begin
    d = foldl(right, GroupBy(isodd, opcompose(Map(last), Filter(isodd)), +), 1:10)
    @test d == Dict(true => 25)
    @test d.state == Dict(true => 25, false => DefaultInit(+))
    @test valtype(d) <: Int
    @test valtype(d.state) <: Union{Int,typeof(DefaultInit(+))}
end

@testset "automatic asmonoid" begin
    @test foldl(right, GroupBy(identity, Map(first), (a, b) -> a + b), [1, 2, 1, 2, 3]) ==
          Dict(1 => 2, 2 => 4, 3 => 3)

    @test foldl(right, GroupBy(identity, (_, x) -> x), [1, 2, 1, 2, 3]) ==
          Dict(1 => 1 => 1, 2 => 2 => 2, 3 => 3 => 3)
end

@testset "GroupByViewDict" begin
    gd1 = foldl(right, GroupBy(string, Map(last), push!!), [1, 2, 1, 2, 3])
    gd2 = foldl(
        right,
        GroupBy(x -> gcd(x, 6), opcompose(Map(last), Filter(isodd)), push!!),
        1:10,
    )
    gd3 = foldl(
        right,
        opcompose(Filter(isodd), GroupBy(x -> gcd(x, 6), Map(last), push!!)),
        1:10,
    )
    gd4 = foldl(
        right,
        GroupBy(x -> gcd(x, 6), opcompose(Map(last), Unique(x -> gcd(x, 4))), push!!),
        1:10,
    )
    @test gd1 == Dict("1" => [1, 1], "2" => [2, 2], "3" => [3])
    @test gd2 == gd3
    @test gd4 == Dict(2 => [2, 4], 3 => [3], 6 => [6], 1 => [1])
    @test Dict(gd1) == Dict("1" => [1, 1], "2" => [2, 2], "3" => [3])
    @test Dict(gd2) == Dict(gd3)
    @test Dict(gd4) == Dict(2 => [2, 4], 3 => [3], 6 => [6], 1 => [1])
end

end  # module


let
    roadway = gen_straight_roadway(1, 1000.0)
    @test length(roadway.segments) == 1

    seg = roadway[1]
    @test length(seg.lanes) == 1

    lane = seg.lanes[1]
    @test isapprox(lane.width, DEFAULT_LANE_WIDTH)
    @test lane.curve[1] == CurvePt(VecSE2(0.0,0.0,0.0), 0.0)
    @test lane.curve[2] == CurvePt(VecSE2(1000.0,0.0,0.0), 1000.0)
end

let
    roadway = gen_straight_roadway(3, 500.0, lane_width=1.0)
    @test length(roadway.segments) == 1

    seg = roadway[1]
    @test length(seg.lanes) == 3

    lane = seg.lanes[1]
    @test isapprox(lane.width, 1.0)
    @test lane.curve[1] == CurvePt(VecSE2(0.0,0.0,0.0), 0.0)
    @test lane.curve[2] == CurvePt(VecSE2(500.0,0.0,0.0), 500.0)

    lane = seg.lanes[2]
    @test isapprox(lane.width, 1.0)
    @test lane.curve[1] == CurvePt(VecSE2(0.0,1.0,0.0), 0.0)
    @test lane.curve[2] == CurvePt(VecSE2(500.0,1.0,0.0), 500.0)

    lane = seg.lanes[3]
    @test isapprox(lane.width, 1.0)
    @test lane.curve[1] == CurvePt(VecSE2(0.0,2.0,0.0), 0.0)
    @test lane.curve[2] == CurvePt(VecSE2(500.0,2.0,0.0), 500.0)

    origin = VecSE2(1.0,2.0,deg2rad(90))
    roadway = gen_straight_roadway(1, 100.0, lane_width=1.0, origin=origin)
    @test length(roadway.segments) == 1
    lane = roadway[1].lanes[1]
    @test lane.curve[1].pos == origin
    @test isapprox(lane.curve[2].pos.x,         1.0, atol=1e-6)
    @test isapprox(lane.curve[2].pos.y,       102.0, atol=1e-6)
    @test isapprox(lane.curve[2].pos.θ, deg2rad(90), atol=1e-6)

    roadway = gen_straight_roadway(2, 100.0, lane_widths=[1.0, 2.0])
    @test length(roadway.segments) == 1
    @test isapprox(roadway[1].lanes[1].curve[1].pos.y, 0.0, atol=1e-6)
    @test isapprox(roadway[1].lanes[2].curve[1].pos.y, 0.5 + 1.0, atol=1e-6)
end

let
    roadway = gen_stadium_roadway(1, length=100.0, width=10.0,
                                  radius=25.0, ncurvepts_per_turn=2)
    @test length(roadway.segments) == 4

    seg = roadway[1]
    @test length(seg.lanes) == 1

    lane = seg.lanes[1]
    @test lane.curve[1] == CurvePt(VecSE2(100.0,0.0,0.0), 0.0)
    @test lane.curve[2] == CurvePt(VecSE2(125.0,25.0,π/2), 25*π/2)

    @test roadway[LaneTag(2,1)].curve[1] == CurvePt(VecSE2(125.0,35.0,π/2), 0.0)

    roadway = gen_stadium_roadway(4)
end


using DataFrames

function molar_mass(MI)
	sql = "	SELECT 	molecule_id, local_id, molar_mass
			FROM	isotopologues
			WHERE 	(molecule_id, local_id)
			IN 		(VALUES " * join(["(" * join("?"^length(t), ',') * ")" for t in MI], ',') * ")"
	result = query_local_db(sql, [i for t in MI for i in t])

	out = Dict{Tuple{Int, Int}, Float64}()
	for row in result
		out[(row.molecule_id, row.local_id)] = row.molar_mass
	end

	return out
end

## Default data for molecules and isotopologues for new databases
const molecules = DataFrame(id=Int[], formula=String[], name=String[])
push!(molecules,
    (1, "H2O", "Water"),
    (2, "CO2", "Carbon Dioxide"),
    (3,	"O3", "Ozone"),
    (4,	"N2O", "Nitrogen oxide"),
    (5,	"CO", "Carbon Monoxide"),
    (6,	"CH4", "Methane"),
    (7, "O2", "Oxygen"),
    (8,	"NO", "Nitric Oxide"),
    (9,	"SO2", "Sulfur Dioxide"),
    (10, "NO2", "Nitrogen Dioxide"),
    (11, "NH3", "Ammonia"),
    (12, "HNO3", "Nitric Acid"),
    (13, "OH", "Hydroxyl"),
    (14, "HF", "Hydrogen Fluoride"),
    (15, "HCl", "Hydrogen Chloride"),
    (16, "HBr", "Hydrogen Bromide"),
    (17, "HI", "Hydrogen Iodide"),
    (18, "ClO", "Chlorine Monoxide"),
    (19, "OCS", "Carbonyl Sulfide"),
    (20, "H2CO", "Formaldehyde"),
    (21, "HOCl", "Hypochlorous Acid"),
    (22, "N2", "Nitrogen"),
    (23, "HCN", "Hydrogen Cyanide"),
    (24, "CH3Cl", "Methyl Chloride"),
    (25, "H2O2", "Hydrogen Peroxide"),
    (26, "C2H2", "Acetylene"),
    (27, "C2H6", "Ethane"),
    (28, "PH3", "Phosphine"),
    (29, "COF2", "Carbonyl Fluoride"),
    (30, "SF6", "Sulfur Hexafluoride"),
    (31, "H2S", "Hydrogen Sulfide"),
    (32, "HCOOH", "Formic Acid"),
    (33, "HO2", "Hydroperoxyl"),
    (34, "O", "Oxygen Atom"),
    (35, "ClONO2", "Chlorine Nitrate"),
    (36, "NO+", "Nitric Oxide Cation"),
    (37, "HOBr", "Hypobromous Acid"),
    (38, "C2H4", "Ethylene"),
    (39, "CH3OH", "Methanol"),
    (40, "CH3Br", "Methyl Bromide"),
    (41, "CH3CN", "Acetonitrile"),
    (42, "CF4", "PFC-14"),
    (43, "C4H2", "Diacetylene"),
    (44, "HC3N", "Cyanoacetylene"),
    (45, "H2", "Hydrogen"),
    (46, "CS", "Carbon Monosulfide"),
    (47, "SO3", "Sulfur trioxide"),
    (48, "C2N2", "Cyanogen"),
    (49, "COCl2", "Phosgene"),
	(53, "CS2", "Carbon disulfide"),
	(55, "NF3", "Nitrogen trifluoride")
)

const isotopologues = DataFrame(
    molecule_id=Int[],
    global_id=Int[],
    local_id=Int[],
    formula=String[],
    afgl_code=Int[],
    abundance=Float64[],
    molar_mass=Float64[],
    q_t0=Float64[],
    g_i=Int[]
    )

push!(isotopologues,
    # 1: H2O
	(1, 1, 1, "H216O", 161, 0.997317, 18.010565, 174.58, 1),
	(1, 2, 2, "H218O", 181, 0.002000, 20.014811, 176.05, 1),
	(1, 3, 3, "H217O", 171, 3.718840e-4, 19.01478, 1052.14, 6),
	(1, 4, 4, "HD16O", 162, 3.106930e-4, 19.01674, 864.74, 6),
	(1, 5, 5, "HD18O", 182, 6.230030e-7, 21.020985, 875.57, 6),
	(1, 6, 6, "HD17O", 172, 1.158530e-7, 20.020956, 5226.79, 36),
	(1, 129, 7, "D216O", 262, 2.419700e-8, 20.022915, 1027.80, 1),

    # 2: CO2
	(2, 7, 1, "12C16O2", 626, 0.984204, 43.98983, 286.09, 1),
	(2, 8, 2, "13C16O2", 636, 0.011057, 44.993185, 576.64, 2),
	(2, 9, 3, "16O12C18O", 628, 0.003947, 45.994076, 607.81, 1),
	(2, 10, 4, "16O12C17O", 627, 7.339890e-4, 44.994045, 3542.61, 6),
	(2, 11, 5, "16O13C18O", 638, 4.434460e-5, 46.997431, 1225.46, 2),
	(2, 12, 6, "16O13C17O", 637, 8.246230e-6, 45.9974, 7141.32, 12),
	(2, 13, 7, "12C18O2", 828, 3.957340e-6, 47.998322, 323.42, 1),
	(2, 14, 8, "17O12C18O", 827, 1.471800e-6, 46.998291, 3766.58, 6),
	(2, 121, 9, "12C17O2", 727, 1.368470e-7, 45.998262, 10971.57, 1),
	(2, 15, 0, "13C18O2", 838, 4.446000e-8, 49.001675, 652.24, 2),
	(2, 120, 11, "18O13C17O", 837, 1.653540e-8, 48.001646, 7595.04, 12),
	(2, 122, 12, "13C17O2", 737, 1.537500e-9, 47.0016182378, 22120.47, 2),

    # 3: O3
	(3, 16, 1, "16O3", 666, 0.992901, 47.984745, 3483.71, 1),
	(3, 17, 2, "16O16O18O", 668, 0.003982, 49.988991, 7465.68, 1),
	(3, 18, 3, "16O18O16O", 686, 0.001991, 49.988991, 3647.08, 1),
	(3, 19, 4, "16O16O17O", 667, 7.404750e-4, 48.98896, 43330.85, 6),
	(3, 20, 5, "16O17O16O", 676, 3.702370e-4, 48.98896, 21404.96, 6),

    # 4: N2O
	(4, 21, 1, "14N216O", 446, 0.990333, 44.001062, 4984.90, 9),
	(4, 22, 2, "14N15N16O", 456, 0.003641, 44.998096, 3362.01, 6),
	(4, 23, 3, "15N14N16O", 546, 0.003641, 44.998096, 3458.58, 6),
	(4, 24, 4, "14N218O", 448, 0.001986, 46.005308, 5314.74, 9),
	(4, 25, 5, "14N217O", 447, 3.692800e-4, 45.005278, 30971.79, 54),

    # 5: CO
	(5, 26, 1, "12C16O", 26, 0.986544, 27.994915, 107.42, 1),
	(5, 27, 2, "13C16O", 36, 0.011084, 28.99827, 224.69, 2),
	(5, 28, 3, "12C18O", 28, 0.001978, 29.999161, 112.77, 1),
	(5, 29, 4, "12C17O", 27, 3.678670e-4, 28.99913, 661.17, 6),
	(5, 30, 5, "13C18O", 38, 2.222500e-5, 31.002516, 236.44, 2),
	(5, 31, 6, "13C17O", 37, 4.132920e-6, 30.002485, 1384.66, 12),

    # 6: CH4
	(6, 32, 1, "12CH4", 211, 0.988274, 16.0313, 590.48, 1),
	(6, 33, 2, "13CH4", 311, 0.011103, 17.034655, 1180.82, 2),
	(6, 34, 3, "12CH3D", 212, 6.157510e-4, 17.037475, 4794.73, 3),
	(6, 35, 4, "13CH3D", 312, 6.917850e-6, 18.04083, 9599.16, 6),

    # 7: O2
	(7, 36, 1, "16O2", 66, 0.995262, 31.98983, 215.73, 1),
	(7, 37, 2, "16O18O", 68, 0.003991, 33.994076, 455.23, 1),
	(7, 38, 3, "16O17O", 67, 7.422350e-4, 32.994045, 2658.12, 6),

    # 8: NO
	(8, 39, 1, "14N16O", 46, 0.993974, 29.997989, 1142.13, 3),
	(8, 40, 2, "15N16O", 56, 0.003654, 30.995023, 789.26, 2),
	(8, 41, 3, "14N18O", 48, 0.001993, 32.002234, 1204.44, 3),

    # 9: SO2
	(9, 42, 1, "32S16O2", 626, 0.945678, 63.961901, 6340.30, 1),
	(9, 43, 2, "34S16O2", 646, 0.041950, 65.957695, 6368.98, 1),

    # 10: NO2
	(10, 44, 1, "14N16O2", 646, 0.991616, 45.992904, 13577.48, 3),
	(10, 130, 2, "15N16O2", 656, 0.003646, 46.989938, 9324.70, 2),

    # 11: NH3
	(11, 45, 1, "14NH3", 4111, 0.995872, 17.026549, 1725.22, 3),
	(11, 46, 2, "15NH3", 5111, 0.003661, 18.023583, 1153.30, 2),

    # 12: HNO3
	(12, 47, 1, "H14N16O3", 146, 0.989110, 62.995644, 2.14e5, 6),
	(12, 117, 2, "H15N16O3", 156, 0.003636, 63.99268, 1.43e5, 4),

    # 13: OH
	(13, 48, 1, "16OH", 61, 0.997473, 17.00274, 80.35, 2),
	(13, 49, 2, "18OH", 81, 0.002000, 19.006986, 80.88, 2),
	(13, 50, 3, "16OD", 62, 1.553710e-4, 18.008915, 209.32, 3),

    # 14: HF
	(14, 51, 1, "H19F", 19, 0.999844, 20.006229, 41.47, 4),
	(14, 110, 2, "D19F", 29, 1.557410e-4, 21.012404, 115.91, 6),

    # 15: HCl
	(15, 52, 1, "H35Cl", 15, 0.757587, 35.976678, 160.65, 8),
	(15, 53, 2, "H37Cl", 17, 0.242257, 37.973729, 160.89, 8),
	(15, 107, 3, "D35Cl", 25, 1.180050e-4, 36.982853, 462.78, 12),
	(15, 108, 4, "D37Cl", 27, 3.773500e-5, 38.979904, 464.13, 12),

    # 16: HBr
	(16, 54, 1, "H79Br", 19, 0.506781, 79.92616, 200.17, 8),
	(16, 55, 2, "H81Br", 11, 0.493063, 81.924115, 200.23, 8),
	(16, 111, 3, "D79Br", 29, 7.893840e-5, 80.932336, 586.40, 12),
	(16, 112, 4, "D81Br", 21, 7.680160e-5, 82.930289, 586.76, 12),

    # 17: HI
	(17, 56, 1, "H127I", 17, 0.999844, 127.912297, 388.99, 12),
	(17, 113, 2, "D127I", 27, 1.557410e-4, 128.918472, 1147.06, 18),

    # 18: ClO
	(18, 57, 1, "35Cl16O", 56, 0.755908, 50.963768, 3274.61, 4),
	(18, 58, 2, "37Cl16O", 76, 0.241720, 52.960819, 3332.29, 4),

    # 19: OCS
	(19, 59, 1, "16O12C32S", 622, 0.937395, 59.966986, 1221.01, 1),
	(19, 60, 2, "16O12C34S", 624, 0.041583, 61.96278, 1253.48, 1),
	(19, 61, 3, "16O13C32S", 632, 0.010531, 60.970341, 2484.15, 2),
	(19, 62, 4, "16O12C33S", 623, 0.007399, 60.966371, 4950.11, 4),
	(19, 63, 5, "18O12C32S", 822, 0.001880, 61.971231, 1313.78, 1),
	(19, 135, 6, "16O13C34S", 634, 4.675080e-4, 62.966136, 2546.53, 2),

    # 20: H2CO
	(20, 64, 1, "H212C16O", 126, 0.986237, 30.010565, 2844.53, 1),
	(20, 65, 2, "H213C16O", 136, 0.011080, 31.01392, 5837.69, 2),
	(20, 66, 3, "H212C18O", 128, 0.001978, 32.014811, 2986.44, 1),

    # 21: HOCl
	(21, 67, 1, "H16O35Cl", 165, 0.755790, 51.971593, 19274.79, 8),
	(21, 68, 2, "H16O37Cl", 167, 0.241683, 53.968644, 19616.20, 8),

    # 22: N2
	(22, 69, 1, "14N2", 44, 0.992687, 28.006148, 467.10, 1),
	(22, 118, 2, "14N15N", 45, 0.007478, 29.003182, 644.10, 6),

    # 23: HCN
	(23, 70, 1, "H12C14N", 124, 0.985114, 27.010899, 892.20, 6),
	(23, 71, 2, "H13C14N", 134, 0.011068, 28.014254, 1830.97, 12),
	(23, 72, 3, "H12C15N", 125, 0.003622, 28.007933, 615.28, 4),

    # 24: CH3Cl
	(24, 73, 1, "12CH335Cl", 215, 0.748937, 49.992328, 57916.12, 4),
	(24, 74, 2, "12CH337Cl", 217, 0.239491, 51.989379, 58833.90, 4),

    # 25: H2O2
	(25, 75, 1, "H216O2", 1661, 0.994952, 34.00548, 9847.99, 1),

    # 26: C2H2
	(26, 76, 1, "12C2H2", 1221, 0.977599, 26.01565, 412.45, 1),
	(26, 77, 2, "H12C13CH", 1231, 0.021966, 27.019005, 1656.18, 8),
	(26, 105, 3, "H12C12CD", 1222, 3.045500e-4, 27.021825, 1581.84, 6),

    # 27: C2H6
	(27, 78, 1, "12C2H6", 1221, 0.976990, 30.04695, 70882.52, 1),
	(27, 106, 2, "12CH313CH3", 1231, 0.021953, 31.050305, 36191.80, 2),

    # 28: PH3
	(28, 79, 1, "31PH3", 1111, 0.999533, 33.997238, 3249.44, 2),

    # 29: COF2
	(29, 80, 1, "12C16O19F2", 269, 0.986544, 65.991722, 70028.43, 1),
	(29, 119, 2, "13C16O19F2", 369, 0.011083, 66.995083, 1.40e5, 2),

    # 30: SF6
	(30, 126, 1, "32S19F6", 29, 0.950180, 145.962492, 1.62e6, 1),

    # 31: H2S
	(31, 81, 1, "H232S", 121, 0.949884, 33.987721, 505.79, 1),
	(31, 82, 2, "H234S", 141, 0.042137, 35.983515, 504.35, 1),
	(31, 83, 3, "H233S", 131, 0.007498, 34.987105, 2014.94, 4),

    # 32: HCOOH
	(32, 84, 1, "H12C16O16OH", 126, 0.983898, 46.00548, 39132.76, 4),

    # 33: HO2
	(33, 85, 1, "H16O2", 166, 0.995107, 32.997655, 4300.39, 2),

    # 34: O
	(34, 86, 1, "16O", 6, 0.997628, 15.994915, 6.72, 1),

    # 35: ClONO2
	(35, 127, 1, "35Cl16O14N16O2", 5646, 0.749570, 96.956672, 4.79e6, 12),
	(35, 128, 2, "37Cl16O14N16O2", 7646, 0.239694, 98.953723, 4.91e6, 12),

    # 36: NO+
    (36, 87, 1, "14N16O+", 46, 0.993974, 29.997989, 311.69, 3),

    # 37: HOBr
	(37, 88, 1, "H16O79Br", 169, 0.505579, 95.921076, 28339.38, 8),
	(37, 89, 2, "H16O81Br", 161, 0.491894, 97.919027, 28237.98, 8),

    # 38: C2H4
	(38, 90, 1, "12C2H4", 221, 0.977294, 28.0313, 11041.54, 1),
	(38, 91, 2, "12CH213CH2", 231, 0.021959, 29.034655, 45196.89, 2),

    # 39: CH3OH
	(39, 92, 1, "12CH316OH", 2161, 0.985930, 32.026215, 70569.92, 2),

    # 40: CH3Br
	(40, 93, 1, "12CH379Br", 219, 0.500995, 93.941811, 83051.98, 4),
	(40, 94, 2, "12CH381Br", 211, 0.487433, 95.939764, 83395.21, 4),

    # 41: CH3CN
	(41, 95, 1, "12CH312C14N", 2124, 0.973866, 41.026549, 88672.19, 3),

    # 42: CF4
	(42, 96, 1, "12C19F4", 29, 0.988890, 87.993616, 1.21e5, 1),

    # 43: C4H2
	(43, 116, 1, "12C4H2", 2211, 0.955998, 50.01565, 9818.97, 1),

    # 44: HC3N
	(44, 109, 1, "H12C314N", 1224, 0.963346, 51.010899, 24786.84, 6),

    # 45: H2
	(45, 103, 1, "H2", 11, 0.999688, 2.01565, 7.67, 1),
	(45, 115, 2, "HD", 12, 3.114320e-4, 3.021825, 29.87, 6),

    # 46: CS
	(46, 97, 1, "12C32S", 22, 0.939624, 43.971036, 253.62, 1),
	(46, 98, 2, "12C34S", 24, 0.041682, 45.966787, 257.77, 1),
	(46, 99, 3, "13C32S", 32, 0.010556, 44.974368, 537.50, 2),
	(46, 100, 4, "12C33S", 23, 0.007417, 44.970399, 1022.97, 4),

    # 47: SO3
	(47, 114, 1, "32S16O3", 26, 0.943400, 79.95682, 7783.30, 1),

    # 48: C2N2
	(48, 123, 1, "12C214N2", 4224, 0.970752, 52.006148, 15582.44, 1),

    # 49: COCl2
	(49, 124, 1, "12C16O35Cl2", 2655, 0.566392, 97.9326199796, 1.48e6, 1),
	(49, 125, 2, "12C16O35Cl37Cl", 2657, 0.362235, 99.9296698896, 3.04e6, 16),

    # 53: CS2
	(53, 131, 1, "12C32S2", 222, 0.892811, 75.94414, 1352.60, 1),
	(53, 132, 2, "32S12C34S", 224, 0.079260, 77.93994, 2798.00, 1),
	(53, 133, 3, "32S12C33S", 223, 0.014094, 76.943256, 1107.00, 4),
	(53, 134, 4, "13C32S2", 232, 0.010310, 76.947495, 2739.70, 2),

	# 55: NF3
	(55, 136, 1, "14N19F3", 4999, 0.996337, 70.998284, 3.46e5, 3)
)


using Revise
using Plots
using Pkg
# Pkg.activate(".")
using RadiativeTransfer
using RadiativeTransfer.Architectures
using RadiativeTransfer.Absorption
using RadiativeTransfer.Scattering
using RadiativeTransfer.vSmartMOM
using RadiativeTransfer.SolarModel
using InstrumentOperator
using Interpolations
using Polynomials
using ForwardDiff
using Distributions
using NCDatasets

## Atmospheric Radiative Transfer

# Load parameters from file
parameters = vSmartMOM.parameters_from_yaml("RadiativeTransfer/test/helper/O2Parameters.yaml")
FT = Float64

# oco2_L1bScND_18688a_180105_B8100r_180206190633

oco_file = "/net/fluo/data1/group/oco2/L1bSc/oco2_L1bScGL_15258a_170515_B10003r_200214061601.h5"
oco_met_file = "/net/fluo/data1/group/oco2/L2Met/oco2_L2MetGL_26777a_190715_B10003r_200429213029.h5"
ils_file = "/home/rjeyaram/RadiativeTransfer/src/vSmartMOM/ils_oco2.json"

fp = 5
iOrbit = 4400
band = 1

sza_ = 32.4436
vza_ = [0.072]

function conv_spectra_local(m::VariableKernelInstrument, ν, spectrum; stride=1)
    # FT = eltype(m.ν_out)
    # Define grid where to perform convolution:

    # Padding at both sides required:
    off = ceil(Int, size(m.kernel, 1) / 2)
    ind = off:stride:(length(ν) - off)

    # knots where convolution will be applied to
    knots = view(ν, ind)
    te = LinearInterpolation(m.ν_out, Float32.(m.ind_out))
    spec_out = zeros(Real, length(knots));
    for i in eachindex(knots)
        # Simple first, nearest neighbor ILS
        ind_fraction = round(Int, te(knots[i]));
        kernel = view(m.kernel, :, ind_fraction)
        for j in eachindex(kernel)
            spec_out[i] += kernel[j] * spectrum[ind[i] + j]
        end
    end
    # Change this later to only perform conv around output grid!
    fin = LinearInterpolation(ν[ind], spec_out; extrapolation_bc=Interpolations.Flat())
    return fin(m.ν_out)
end;

# Runner is used to set AD fields as duals
function runner!(y, x, parameters=parameters, oco_file=oco_file,
                                              oco_met_file=oco_met_file,
                                              ils_file=ils_file, fp=fp, iOrbit=iOrbit,
                                              sza_=sza_, vza_=vza_)

    # Set parameters fields as the dual numbers
    parameters.brdf = [LambertianSurfaceScalar(x[1])]

    parameters.scattering_params.rt_aerosols[1].τ_ref = x[2];

    # parameters.scattering_params.rt_aerosols[1].aerosol.size_distribution = LogNormal(log(x[3]), log(x[4]), check_args=false)

    # parameters.scattering_params.rt_aerosols[1].aerosol.nᵣ = x[5];
    # parameters.scattering_params.rt_aerosols[1].aerosol.nᵢ = x[6];

    # parameters.scattering_params.rt_aerosols[1].p₀ = x[7];
    # parameters.scattering_params.rt_aerosols[1].σp = x[8];

    # Set profiles properly
    met = Dataset(oco_met_file);
    T_met = met.group["Meteorology"]["temperature_profile_met"][:,fp,iOrbit];
    ak = met.group["MeteorologyDiagnostics"]["ak"][:,fp,iOrbit];
    bk = met.group["MeteorologyDiagnostics"]["bk"][:,fp,iOrbit];
    p_surf = met.group["Meteorology"]["surface_pressure_met"][fp,iOrbit];
    p_half = (ak + bk * p_surf);
    p_half = vcat(p_half, p_surf);
    q = met.group["Meteorology"]["specific_humidity_profile_met"][:,fp,iOrbit];
    parameters.p = p_half / 100
    parameters.q = q # zeros(size(q))
    parameters.T = T_met
    parameters.sza = sza_
    parameters.vza = vza_

    @show parameters.sza, parameters.vza

    # @show parameters.sza
    # plot!(parameters.p)
    # return

    model = model_from_parameters(parameters);

    # Run the model to obtain reflectance matrix
    R = vSmartMOM.rt_run(model, i_band=1);

    # Produce black-body in wavenumber range
    T = 5777
    λ_grid = reverse(1e4 ./ parameters.spec_bands[1]) #collect(757:0.01:777)
    black_body = planck_spectrum_wl(T, λ_grid) * 2.1629e-05 * pi
    black_body = SolarModel.watts_to_photons(λ_grid, black_body)

    # Get solar transmittance spectrum
    solar_transmission = solar_transmission_from_file("/home/rjeyaram/RadiativeTransfer/src/solar_merged_20160127_600_26316_100.out", parameters.spec_bands[1])

    # Get outgoing solar radiation
    sun_out = reverse(solar_transmission) .* black_body

    # Apply Earth reflectance matrix
    earth_out = sun_out .* reverse(R[:])

    # Set up InstrumentOperator
    ils_Δ, ils_in, dispersion = InstrumentOperator.read_ils_table(oco_file, ils_file);

    # Define model grid:
    res = 0.001

    # Just consider the ILS within ± 0.35nm
    grid_x = -0.35e-3:res*1e-3:0.35e-3

    extended_dims = [fp,band] # Footprint, band

    # Re-interpolate I from ν_grid to new grid/resolution
    interp_I = LinearInterpolation(λ_grid, earth_out);
    wl = 757.5:res:771.0
    I_wl = interp_I(wl/1000)

    # Pixels to be used
    ind_out = collect(0:1015);

    # Eventual grid of OCO-2 for Band 1, FP 5:
    dispPoly = Polynomial(view(dispersion, :, extended_dims...))
    ν = Float32.(dispPoly.(1:1016))

    # Prepare ILS table:
    ils_pixel   = InstrumentOperator.prepare_ils_table(grid_x, ils_in, ils_Δ,extended_dims)
    oco2_kernel = VariableKernelInstrument(ils_pixel, ν, ind_out)

    # Convolve input spectrum with variable kernel
    I_conv = conv_spectra_local(oco2_kernel, wl./1e3, I_wl)

    y[:] = I_conv[:]

end

directory = "/net/fluo/data1/group/oco2/L1bSc/"
files_list = filter(x->endswith(x, ".h5") && startswith(x, "oco2_L1bScGL_26"), readdir(directory))

met_files_list = filter(x->endswith(x, ".h5"), readdir("/net/fluo/data1/group/oco2/L2Met/"))

I_convs_all = zeros(1016, length(files_list))
oco2_Abands_all = zeros(1016, length(files_list))

albedos = zeros(length(files_list))
cloud_levels = zeros(length(files_list))

optimal_x = zeros(8, length(files_list))

for i in 1:length(files_list)

    @show i

    oco_file = files_list[i]

    # try
        ocoData = Dataset(directory * oco_file)

        oco_met_file = "/net/fluo/data1/group/oco2/L2Met/" * filter(x->contains(x, oco_file[14:19]), met_files_list)[1]

        oco_file = directory * oco_file
        # @show oco_file
        # @show oco_met_file

        met = Dataset(oco_met_file)
        cloud_levels[i] = met.group["Meteorology"]["cloud_liquid_water_path_met"][fp,iOrbit]

        # println()

        # continue

        ## Getting an OCO spectrum for fun to fit:
        o2AbandSpectra = ocoData.group["SoundingMeasurements"]["radiance_o2"]
        SoundingGeometry = ocoData.group["SoundingGeometry"]
        vza = SoundingGeometry["sounding_zenith"]
        sza = SoundingGeometry["sounding_solar_zenith"]
        lat = SoundingGeometry["sounding_latitude"]
        lon = SoundingGeometry["sounding_longitude"]
        o2AbandSpectra = ocoData.group["SoundingMeasurements"]["radiance_o2"]
        oco2_Aband = o2AbandSpectra[:,fp,iOrbit]

        println("$(round(Float64(lat[fp,iOrbit]), digits=2)) $(round(Float64(lon[fp,iOrbit]), digits=4))")
        # continue

        sza_ = sza[fp,iOrbit]
        vza_ = [vza[fp,iOrbit]]

        x = FT[0.05,
        0.05]#,
        # 1.3,
        # 2.0,
        # 1.3,
        # 0.00000001,
        # 90000,
        # 5000.0]

        # Run FW model:
        # @time runner(x);
        I_conv = zeros(1016)
        dfdx = ForwardDiff.jacobian(runner!, I_conv, x);

        dfdx[:,2] .= 0;

        @show size(dfdx)

        y = oco2_Aband;
        Fx = I_conv;
        K = dfdx;

        # Regular least squares:
        dx = K \ (y-Fx)

        I_convs_all[:,i] = I_conv
        oco2_Abands_all[:,i] = oco2_Aband

        # optimal_x[:,i] = x+dx
        albedos[i] = x[1] + dx[1]

        # @show dx, x+dx

    # catch
    #     println("Error on " * oco_file)
    #     continue
    # end

end


@testset "Options" begin
    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    @test_logs (:error, r"Possible values are") Model(optimizer_with_attributes(
        CS.Optimizer,
        "lp_optimizer" => cbc_optimizer,
        "logging" => [],
        "traverse_strategy" => :KFS,
    ))
end


# simulate.jl: functions to simulate the model
# include("setup.jl")
function evolve!( pop::Population, generations::Int64 = 1 )
    # from ReputationSets https://github.com/tkessinger/reputation_sets
    # the main function to evolve the population
    # generations specifies the total number of generations for which to run the simulation
    for i::Int64 in 1:generations
        if pop.verbose println("\ninitiating generation $(pop.generation)") end

        # first, update payoffs
        if pop.verbose println("updating actions and payoffs") end
        update_payoffs_using_matrix!(pop)  # method that stores payoffs in a matrix

        # then, select a learner and a role model and perform update
        if pop.verbose println("") end
        if pop.verbose println("evolving, generation $(pop.generation)") end
        update_strategies_and_opinions_db!(pop)

        if pop.verbose println("") end
        if pop.verbose println("ending generation $(pop.generation)") end
        pop.generation += 1
    end
end

function update_payoffs_using_matrix!( pop::Population )
    # update individual payoffs by computing pairwise payoffs and storing them in a matrix

    # to make the code run faster, we only update payoffs in a given generation
    # if strategies or sets have changed. sets_changed and strategies_changed are
    # initialized as true so that payoffs are computed in the first round
    if pop.sets_changed
        # if (previous) learner's sets have changed, re-compute all payoffs
        if pop.verbose println("updating all payoffs because set_changed = $(pop.sets_changed)") end
        pop.payoffs_mat  .= zero(Float64) # reset all payoffs_mat
        pop.prev_actions .= zero(Int64)   # reset all prev_actions

        # then, recompute all payoffs
        for k in 1:pop.sets.M  # for each set
            for (i, j) in pop.sets.set_pairs[k]   # for each pair within each set
                compute_payoffs!( pop, i, j, k )  # compute payoffs for i & j when intereacting in set k
                                                  # this also updates pop.prev_actions accordingly
            end
        end

    elseif pop.strategies_changed
        # if (previous) learner's strategy has changed, recompute payoffs involving that individual
        if pop.verbose println("updating payoffs involving $(pop.prev_learner) because strategies_changed = $(pop.strategies_changed)") end
        pop.payoffs_mat[pop.prev_learner,:]    .= zero(Float64) # first, reset payoffs involving prev_learner
        pop.payoffs_mat[:,pop.prev_learner]    .= zero(Float64)
        pop.prev_actions[:,pop.prev_learner,:] .= zero(Int64)
        pop.prev_actions[:,:,pop.prev_learner] .= zero(Int64)

        # then, recompute payoffs involving prev_learner
        for k in 1:pop.sets.M  # for each set
            for (i, j) in pop.sets.set_pairs[k]  # for each pair within each set
                if i == pop.prev_learner || j == pop.prev_learner
                    compute_payoffs!( pop, i, j, k )  # compute payoffs for i & j when intereacting in set k
                                                      # this also updates pop.prev_actions accordingly
                end
            end
        end
    else
        if pop.verbose println("not updating payoffs b/c sets_changed = $(pop.sets_changed) and strategies_changed = $(pop.strategies_changed)") end
        return # no need to update prev_interactions or payoffs
    end

    # sum rows to get total individual payoffs
    pop.payoffs = [round( x, digits=13 ) for x in vec( sum(pop.payoffs_mat, dims=2) )]  # rounding avoids precision errors
    pop.prev_interactions = 2 * sum([length(x) for x in pop.sets.set_pairs])

    if pop.verbose println( "\npayoff matrix is now $(pop.payoffs_mat)" ) end
    if pop.verbose println( "payoffs are now $(pop.payoffs)" ) end
    if pop.verbose println( "there were $(pop.prev_interactions) interactions, $(sum(pop.prev_actions)) of which were C" ) end
end

function compute_payoffs!( pop::Population, i::Int64, j::Int64, k::Int64 )
    # function to compute payoffs for individuals i and j when they interact in set k
    pop.sim.i_action, pop.sim.j_action = -1, -1  # set to -1 to catch any errors in action assignment

    # actions depend on whether i and j agree on issue k or not
    if pop.sets.h[i,k] == 0 || pop.sets.h[j,k] == 0
        # if either i or j does not belong to set k, something's wrong
        # leave pop.sim.i_action, pop.sim.j_action as -1 so that the compute_payoffs! throws an error when assigning payoffs
        @warn("$i and $j should not be interacting in set $k"); return
    elseif pop.sets.h[i,k] == pop.sets.h[j,k]
        pop.sim.i_action, pop.sim.j_action = pop.strategies[i,1], pop.strategies[j,1] # if they agree, play s_agree
    else
        pop.sim.i_action, pop.sim.j_action = pop.strategies[i,2], pop.strategies[j,2] # if they disagree, play s_disagree
    end

    if pop.verbose
        println("\n\tupdating actions of $i and $j in set $k; agree on issue $k? $(pop.sets.h[i,k] == pop.sets.h[j,k])")
        println("\t\t$i has strats $(pop.strategies[i,:]); $i plays strat $(pop.sim.i_action), earns a payoff of $(pop.game.A[pop.sim.i_action+1, pop.sim.j_action+1])")
        println("\t\t$j has strats $(pop.strategies[j,:]); $j plays strat $(pop.sim.j_action), earns a payoff of $(pop.game.A[pop.sim.j_action+1, pop.sim.i_action+1])")
    end

    # add i and j's payoffs to the payoff matrix according to their strategies and the game matrix
    pop.payoffs_mat[i,j] += pop.game.A[pop.sim.i_action+1, pop.sim.j_action+1] # +1 because julia is 1-indexed
    pop.payoffs_mat[j,i] += pop.game.A[pop.sim.j_action+1, pop.sim.i_action+1]

    # track actions
    pop.prev_actions[k,i,j] = pop.sim.i_action
    pop.prev_actions[k,j,i] = pop.sim.j_action
end

function update_strategies_and_opinions_db!( pop::Population )
    # function to perform death-birth update
    if pop.verbose println("Strategies before updating:\t $(pop.strategies)") end
    if pop.verbose println("Set memberships before updating: $(pop.sets.h)") end

    # randomly choose a learner and store a list of non_learners
    pop.sim.learner = sample( 1:pop.sets.N )

    # compute the fitness of every individual in the population
    for i in 1:pop.sets.N pop.sim.fitnesses[i] = 1.0 + pop.game.β * pop.payoffs[i] end
    if any(x->x<0.0, pop.sim.fitnesses) throw("!!! negative fitness detected, aborting !!!"); end

    # choose a role model from the rest of the population, weighted by fitness
    # OLD: pop.sim.role_model = sample( pop.sim.non_learners, Weights(pop.sim.fitnesses) )
    pop.sim.role_model = sample( 1:pop.sets.N, Weights(pop.sim.fitnesses) )

    # learner updates
    if pop.verbose println("\n\trandomly chosen learner is $(pop.sim.learner), role_model is $(pop.sim.role_model)") end
    if pop.verbose println("\t\tall payoffs are $(pop.payoffs)") end
    if pop.verbose println("\t\tfitnesses are $(pop.sim.fitnesses)") end

    # compute probability of imitation: 1 if same party, 1-p if different parties
    pop.sets.affils[pop.sim.learner] == pop.sets.affils[pop.sim.role_model] ? imit_prob = 1.0 : imit_prob = 1-pop.game.ps[pop.sets.affils[pop.sim.learner]]

    if pop.verbose
        println("\t\tlearner and role model in the same party? $(pop.sets.affils[ pop.sim.learner ] == pop.sets.affils[ pop.sim.role_model ]), imitate with prob $(imit_prob)")
    end

    # update strategies and sets: imitation / mutation occurs with probability imit_prob
    if rand() < imit_prob
        if pop.verbose println("\twill consider imitating") end
        update_strategies!( pop ) # includes mutation
        update_opinions!( pop )   # includes mutation
        pop.num_imitating += 1
    else
        if pop.verbose println("\tnah won't imitate") end
        no_strategy_or_opinion_update!( pop )
        pop.num_not_imitating += 1
    end

    # track updates
    pop.prev_learner = pop.sim.learner
    if pop.sets_changed == true
        # h is updated above; need to update h_bar, h_base10, h_bar_base10, set_members, set_pairs
        pop.sets.set_members, pop.sets.set_pairs = set_members_and_pairs(pop.sets.h)  # update set_members and set_pairs
        pop.sets.h_bar        = [abs(x) for x in pop.sets.h]
        pop.sets.h_base10     = [vectobase3(pop.sets.h[i,:]) for i in 1:pop.sets.N]
        pop.sets.h_bar_base10 = [vectobase3(pop.sets.h_bar[i,:]) for i in 1:pop.sets.N]
        pop.num_sets_changed += 1
    end
    if pop.strategies_changed == true
        # strategies are updated above; need to update strategies_base10
        pop.strategies_base10 = [2 * pop.strategies[i,1] + pop.strategies[i,2] for i in 1:pop.sets.N]
        pop.num_strategies_changed += 1
    end
    if pop.strategies_changed && pop.sets_changed
        pop.num_both_changed += 1
    end

    if pop.verbose println("\nStrategies after updating:\t $(pop.strategies)") end
    if pop.verbose println("Set memberships after updating:\t $(pop.sets.h)") end
    if pop.verbose
        println("tracking: prev_learner = $(pop.prev_learner), strategies_changed = $(pop.strategies_changed), sets_changed = $(pop.sets_changed)")
    end
end

function update_strategies!( pop::Population )
    # function to update strategy of learner based on role_model
    pop.sim.strategy_copy = copy(pop.strategies[pop.sim.learner,:])  # save a copy to track changes
    # copy suffices here because the assignments below replace elements of pop.sets.h with Int64 (immutable)

    if rand() < pop.game.u
        # strategy mutation occurs
        pop.strategies[pop.sim.learner,:] = rand(collect(0:1), 1, 2) # conditional case
        pop.num_random_strategies += 1
        pop.num_random_sa_flipped += (pop.sim.strategy_copy[1] != pop.strategies[pop.sim.learner,1])
        pop.num_random_sd_flipped += (pop.sim.strategy_copy[2] != pop.strategies[pop.sim.learner,2])
        if pop.verbose println("\t\tmutating $(pop.sim.learner) to strategy $(pop.strategies[pop.sim.learner,:])") end
    else
        # learner imitates the strategy of role model
        pop.strategies[pop.sim.learner,:] = pop.strategies[pop.sim.role_model,:]
        if pop.verbose println("\t\tadopting $(pop.sim.role_model)'s strategy $(pop.strategies[pop.sim.role_model,:])") end
    end
    pop.strategies_changed = (pop.sim.strategy_copy != pop.strategies[pop.sim.learner,:]) # track changes
end

function update_opinions!( pop::Population )
    # function to update set memberships / opinions of learner based on role_model
    pop.sim.sets_copy = copy(pop.sets.h[pop.sim.learner,:])     # save a copy to track changes
    # copy suffices here because the assignments below replace elements of pop.sets.h with Int64 (immutable)

    # with probability v, set mutation occurs
    if rand() < pop.game.v
        if rand() < 1 - pop.game.ϵ * pop.game.ps[pop.sets.affils[pop.sim.learner]]
            # with probability 1-p*ϵ, select opinions without bias
            # select a random arrangement of K opinions
            pop.sets.h[pop.sim.learner,:] = shuffle( vcat( zeros(Int64, pop.sets.M-pop.sets.K), rand([-1,1], pop.sets.K) ) )

            pop.num_random_opinions += 1
            pop.num_random_opinions_flipped += (pop.sim.sets_copy != pop.sets.h[ pop.sim.learner,: ]) # new
            if pop.verbose println("\t\tmutating $(pop.sim.learner) to random opinions $(pop.sets.h[pop.sim.learner,:])") end
        else
            # with probability p*ϵ, select opinions with bias
            # select a random set of K issues, select opinions corresponding to party
            if pop.sets.affils[pop.sim.learner] == 1
                pop.sets.h[pop.sim.learner,:] = shuffle( vcat( zeros(Int64, pop.sets.M-pop.sets.K), -ones(Int64, pop.sets.K) ) )
            elseif pop.sets.affils[pop.sim.learner] == 2
                pop.sets.h[pop.sim.learner,:] = shuffle( vcat( zeros(Int64, pop.sets.M-pop.sets.K), ones(Int64, pop.sets.K) ) )
            end
            pop.num_biased_opinions += 1
            pop.num_biased_opinions_flipped += (pop.sim.sets_copy != pop.sets.h[ pop.sim.learner,: ]) # new
            if pop.verbose
                println("\t\t$(pop.sim.learner) is in party $(pop.sets.affils[ pop.sim.learner ]), mutating to biased opinions $(pop.sets.h[ pop.sim.learner,: ])")
            end
        end
    else
        # with probability 1-v, learner imitates the strategy of role model
        pop.sets.h[ pop.sim.learner,: ] = pop.sets.h[ pop.sim.role_model,: ]
        if pop.verbose println("\t\tadopting $( pop.sim.role_model )'s set membership $(pop.sets.h[ pop.sim.role_model,: ])") end
    end
    # h is updated above; need to update h_bar, sets_changed, set_members, set_pairs
    # pop.sets.set_members, pop.sets.set_pairs = set_members_and_pairs(pop.sets.h)  # update set_members and set_pairs # moved up
    pop.sets_changed = (pop.sim.sets_copy != pop.sets.h[ pop.sim.learner,: ]) # track changes
end

function no_strategy_or_opinion_update!( pop::Population )
    # function to keep track of changes when no imitation/mutation occurs
    pop.strategies_changed = false  # track changes
    pop.sets_changed       = false
end


"""
ITensors is a library for rapidly creating correct and efficient tensor network algorithms.

An ITensor is a tensor whose interface is independent of its memory layout. ITensor indices are objects which carry extra information and which 'recognize' each other (compare equal to each other).

The ITensor library also includes composable and extensible algorithms for optimizing and transforming tensor networks, such as matrix product state and matrix product operators, such as the DMRG algorithm.
"""
module ITensors

#####################################
# External packages
#
using Compat
using HDF5
using KrylovKit
using LinearAlgebra
using NDTensors
using PackageCompiler
using Pkg
using Printf
using Random
using StaticArrays
using TimerOutputs

#####################################
# NDTensors (definitions that will be moved to NDTensors
# module)
#
include("NDTensors/NDTensors.jl")

#####################################
# ContractionSequenceOptimization
#
include("ContractionSequenceOptimization/ContractionSequenceOptimization.jl")
using .ContractionSequenceOptimization

#####################################
# Directory helper functions (useful for
# running examples)
#
src_dir() = dirname(pathof(@__MODULE__))
pkg_dir() = joinpath(src_dir(), "..")
examples_dir() = joinpath(pkg_dir(), "examples")

#####################################
# Determine version and uuid of the package
#
function _parse_project_toml(field::String)
  return Pkg.TOML.parsefile(joinpath(pkg_dir(), "Project.toml"))[field]
end
version() = VersionNumber(_parse_project_toml("version"))
uuid() = Base.UUID(_parse_project_toml("uuid"))

#####################################
# Exports
#
include("exports.jl")

#####################################
# Imports
#
include("imports.jl")

#####################################
# Global Variables
#
include("global_variables.jl")

#####################################
# Index and IndexSet
#
include("smallstring.jl") # Not currently using in TagSet
include("readwrite.jl")
include("not.jl")
include("tagset.jl")
include("arrow.jl")
include("symmetrystyle.jl")
include("index.jl")
include("set_operations.jl")
include("indexset.jl")

#####################################
# ITensor
#
include("itensor.jl")
include("broadcast.jl")
include("decomp.jl")
include("iterativesolvers.jl")

#####################################
# QNs
#
include("qn/qn.jl")
include("qn/qnindex.jl")
include("qn/qnindexset.jl")
include("qn/qnitensor.jl")

#####################################
# MPS/MPO
#
include("mps/abstractmps.jl")
include("mps/deprecated.jl")
include("mps/mps.jl")
include("mps/mpo.jl")
include("mps/sweeps.jl")
include("mps/projmpo.jl")
include("mps/projmposum.jl")
include("mps/projmps.jl")
include("mps/projmpo_mps.jl")
include("mps/observer.jl")
include("mps/dmrg.jl")

#####################################
# Physics
#
include("physics/sitetype.jl")
include("physics/lattices.jl")
include("physics/site_types/generic_sites.jl")
include("physics/site_types/qubit.jl")
include("physics/site_types/spinhalf.jl")
include("physics/site_types/spinone.jl")
include("physics/site_types/fermion.jl")
include("physics/site_types/electron.jl")
include("physics/site_types/tj.jl")
include("physics/fermions.jl")
include("physics/autompo.jl")

#####################################
# Deprecations
#
include("deprecated.jl")

#####################################
# Argument parsing
#
include("argsdict/argsdict.jl")

#####################################
# Package compilation
#
include("packagecompile/compile.jl")

#####################################
# Developer tools, for internal
# use only
#
include("developer_tools.jl")

function __init__()
  return resize!(empty!(INDEX_ID_RNGs), Threads.nthreads()) # ensures that we didn't save a bad object
end

#####################################
# Precompile certain functions
# (generated from precompile/make_precompile.jl
# using SnoopCompile.jl)
#
#include("../precompile/precompile.jl")
#_precompile_()

end # module ITensors