using Flux
using CUDA: @allowscalar
using BSON: @save
using DataStructures

using ..Measures
using ..PLRNNs
using ..ObservationModels
using ..Utilities

"""
    loss(tfrec, XÃÉ, SÃÉ)

Performs a forward pass using the teacher forced recursion wrapper `tfrec` and
computes and return the loss w.r.t. data `XÃÉ`. Optionally external inputs `SÃÉ`
can be provided.
"""
function loss(
    tfrec::AbstractTFRecur,
    XÃÉ::AbstractArray{T, 3},
    ZÃÇ::AbstractArray{T, 3},
) where {T}
    Z = tfrec(XÃÉ, ZÃÇ)
    XÃÇ = tfrec.O(Z)
    return @views Flux.mse(XÃÇ, XÃÉ[:, :, 2:end])
end

function loss(
    tfrec::AbstractTFRecur,
    XÃÉ::AbstractArray{T, 3},
    ZÃÇ::AbstractArray{T, 3},
    SÃÉ::AbstractArray{T, 3},
) where {T}
    Z = tfrec(XÃÉ, ZÃÇ, SÃÉ)
    XÃÇ = tfrec.O(Z)
    return @views Flux.mse(XÃÇ, XÃÉ[:, :, 2:end])
end

function train_!(
    m::AbstractPLRNN,
    O::ObservationModel,
    ùíü::AbstractDataset,
    opt::Flux.Optimise.Optimiser,
    args::AbstractDict,
    save_path::String,
)
    # hypers
    E = args["epochs"]::Int
    M = args["latent_dim"]::Int
    S‚Çë = args["batches_per_epoch"]::Int
    S = args["batch_size"]::Int
    œÑ = args["teacher_forcing_interval"]::Int
    œÉ_noise = args["gaussian_noise_level"]::Float32
    TÃÉ = args["sequence_length"]::Int
    Œ∫ = args["MAR_ratio"]::Float32
    Œª‚Çò‚Çê·µ£ = args["MAR_lambda"]::Float32
    œÉ¬≤_scaling = args["D_stsp_scaling"]::Float32
    bins = args["D_stsp_bins"]::Int
    œÉ_smoothing = args["PSE_smoothing"]::Float32
    PE_n = args["PE_n"]::Int
    isi = args["image_saving_interval"]::Int
    ssi = args["scalar_saving_interval"]::Int
    exp = args["experiment"]::String
    name = args["name"]::String
    run = args["run"]::Int
    Œ± = args["gtf_alpha"]::Float32
    Œ±_method = args["gtf_alpha_method"]::String
    Œ≥ = args["gtf_alpha_decay"]::Float32
    Œª‚Çí = args["obs_model_regularization"]::Float32
    Œª‚Çó = args["lat_model_regularization"]::Float32
    partial_forcing = args["partial_forcing"]::Bool
    k = args["alpha_update_interval"]::Int
    use_gtf = args["use_gtf"]::Bool

    # data shape
    T, N = size(ùíü.X)

    # wheter external inputs are available
    ext_inputs = typeof(ùíü) <: ExternalInputsDataset

    # progress tracking
    prog = Progress(joinpath(exp, name), run, 20, E, 0.8)

    # decide on D_stsp scaling
    scal, stsp_name = decide_on_measure(œÉ¬≤_scaling, bins, N)

    # initialize stateful model wrapper
    z‚ÇÄ = similar(ùíü.X, M, S)
    if use_gtf
        tfrec = GTFRecur(m, O, z‚ÇÄ, Œ±)
        println(
            "Using GTF with initial Œ± = $Œ± and annealing method: $Œ±_method (Œ≥ = $Œ≥, k = $k)",
        )
        println("Partial forcing set to: $partial_forcing (N = $N, M = $M)")
    else
        tfrec = TFRecur(m, O, z‚ÇÄ, œÑ)
        println("Using sparse TF with œÑ = $œÑ")
        println("Partial forcing set to: $partial_forcing (N = $N, M = $M)")
    end

    # model parameters
    Œ∏ = Flux.params(tfrec)

    # initial Œ±
    Œ±_est = Œ±

    for e = 1:E
        # process a couple of batches
        t‚ÇÅ = time_ns()
        for s‚Çë = 1:S‚Çë
            # sample a batch
            XÃÉ, SÃÉ_exts = ext_inputs ? sample_batch(ùíü, TÃÉ, S) : (sample_batch(ùíü, TÃÉ, S), ())
            SÃÉ_exts = ext_inputs ? (SÃÉ_exts,) : SÃÉ_exts

            # add noise noise if noise level > 0
            œÉ_noise > zero(œÉ_noise) ? add_gaussian_noise!(XÃÉ, œÉ_noise) : nothing

            # precompute forcing signals
            ZÃÇ = estimate_forcing_signals(tfrec, XÃÉ)

            # Œ± estimation & annealing
            if s‚Çë % k == 0 && use_gtf
                Œ±_est = compute_Œ±(tfrec, @view(ZÃÇ[:, :, 2:end]), Œ±_method)
                if Œ±_est > tfrec.Œ±
                    tfrec.Œ± = Œ±_est
                else
                    tfrec.Œ± = Œ≥ * tfrec.Œ± + (1 - Œ≥) * Œ±_est
                end
            end

            # partial forcing
            @views ZÃÇ_subset = partial_forcing ? ZÃÇ[1:N, :, 2:end] : ZÃÇ[:, :, 2:end]

            # forward and backward pass
            grads = Flux.gradient(Œ∏) do
                L‚Çú·µ£ = loss(tfrec, XÃÉ, ZÃÇ_subset, SÃÉ_exts...)
                L·µ£ = regularization_loss(tfrec, Œ∫, Œª‚Çò‚Çê·µ£, Œª‚Çó, Œª‚Çí)
                return L‚Çú·µ£ + L·µ£
            end

            # keep W matrix offdiagonal for PLRNN, dendPLRNN
            keep_connectivity_offdiagonal!(tfrec.model, grads)

            # optimiser step
            Flux.Optimise.update!(opt, Œ∏, grads)

            # check for NaNs in parameters (exploding gradients)
            if check_for_NaNs(Œ∏)
                save_model(
                    [tfrec.model, tfrec.O],
                    joinpath(save_path, "checkpoints", "model_$e.bson"),
                )
                @warn "NaN(s) in parameters detected! \
                    This is likely due to exploding gradients. Aborting training..."
                return nothing
            end
        end
        t‚ÇÇ = time_ns()
        Œît = (t‚ÇÇ - t‚ÇÅ) / 1e9
        update!(prog, Œît, e)

        # plot trajectory
        if e % ssi == 0
            #@show tfrec.Œ±, cond(tfrec.O.B)
            # loss
            # sample a batch
            XÃÉ, SÃÉ_exts = ext_inputs ? sample_batch(ùíü, TÃÉ, S) : (sample_batch(ùíü, TÃÉ, S), ())
            SÃÉ_exts = ext_inputs ? (SÃÉ_exts,) : SÃÉ_exts

            # precompute forcing signals
            ZÃÇ = estimate_forcing_signals(tfrec, XÃÉ)
            # partial forcing
            @views ZÃÇ_subset = partial_forcing ? ZÃÇ[1:N, :, 2:end] : ZÃÇ[:, :, 2:end]
            L‚Çú·µ£ = loss(tfrec, XÃÉ, ZÃÇ_subset, SÃÉ_exts...)
            L·µ£ = regularization_loss(tfrec, Œ∫, Œª‚Çò‚Çê·µ£, Œª‚Çó, Œª‚Çí)

            # generated trajectory
            S_exts = ext_inputs ? (ùíü.S,) : ()
            X_gen =
                @allowscalar @views generate(tfrec.model, tfrec.O, ùíü.X[1, :], T, S_exts...)

            # move data to cpu for metrics and plotting
            X_cpu = ùíü.X |> cpu
            X_gen_cpu = X_gen |> cpu

            # metrics
            D_stsp = state_space_divergence(X_cpu, X_gen_cpu, scal)
            pse = power_spectrum_error(X_cpu, X_gen_cpu, œÉ_smoothing)
            pe = prediction_error(tfrec.model, tfrec.O, ùíü.X, PE_n, S_exts...)

            # progress printing
            scalars = gather_scalars(L‚Çú·µ£, L·µ£, D_stsp, stsp_name, pse, pe, PE_n)
            typeof(tfrec) <: GTFRecur ? scalars["Œ±"] = round(tfrec.Œ±, digits = 3) : nothing
            typeof(O) <: Affine ? scalars["cond(B)"] = round(cond(O.B), digits = 3) :
            nothing

            print_progress(prog, Œît, scalars)

            save_model(
                [tfrec.model, tfrec.O],
                joinpath(save_path, "checkpoints", "model_$e.bson"),
            )
            if e % isi == 0
                # plot
                plot_reconstruction(
                    X_gen_cpu,
                    X_cpu,
                    joinpath(save_path, "plots", "generated_$e.png"),
                )
            end
        end
    end
    return nothing
end

function regularization_loss(
    tfrec::AbstractTFRecur,
    Œ∫::Float32,
    Œª‚Çò‚Çê·µ£::Float32,
    Œª‚Çó::Float32,
    Œª‚Çí::Float32,
)
    # latent model regularization
    L·µ£ = 0.0f0
    if Œ∫ > zero(Œ∫)
        L·µ£ += mar_loss(m, Œ∫, Œª‚Çò‚Çê·µ£)
    elseif Œª‚Çó > 0
        L·µ£ += regularize(tfrec.model, Œª‚Çó)
    end

    # observation model regularization
    L·µ£ += (Œª‚Çí > 0) ? regularize(tfrec.O, Œª‚Çí) : 0
    return L·µ£
end

function gather_scalars(L‚Çú·µ£, L·µ£, D_stsp, stsp_name, pse, pe, PE_n)
    scalars = OrderedDict("‚àëL" => L‚Çú·µ£ + L·µ£)
    if L·µ£ > 0.0f0
        scalars["L‚Çú·µ£"] = L‚Çú·µ£
        scalars["L·µ£"] = L·µ£
    end
    scalars["D‚Çõ‚Çú‚Çõ‚Çö $stsp_name"] = round(D_stsp, digits = 3)
    scalars["D‚Çï"] = round(pse, digits = 3)
    scalars["PE($PE_n)"] = pe
    return scalars
end
