include("../Ses/Ses.jl");
using .Ses

using DataFrames, Dates
using Flux: @epochs, @forward, @functor, Chain, params, pullback, Recur, SkipConnection
using Flux.Losses: mse
using Flux.Optimise: AbstractOptimiser, ADAM

struct PES{M}
    inner_model::M
end



x_getsecond(mx::M, x::X) where {M<:Union{<:Real, <:AbstractVecOrMat}, X<:Union{<:NamedTuple, <:Tuple, <:AbstractVecOrMat}} =
    (mx, x[2]);
x_swapsplat(mx::M, x::X) where {M<:Union{<:Real, <:AbstractVecOrMat}, X<:Union{<:NamedTuple, <:Tuple, <:AbstractVecOrMat}} =
    (x..., mx);



function PES(cell::Type{C},
                 texplanatory::Dict, modes_dct::Dict,
                 state::NTuple{2, I} = (0.0f0, 0.0f0), dim::Integer = 1; kwargs...) where {C<:AbstractSesCell,
                                                                                           II,
                                                                                           I<:Union{II, AbstractVecOrMat{II}}}
    scaler_types = Ses.collect_values(Ses.to_scaler, texplanatory);
    modes = Ses.collect_values(modes_dct, texplanatory);
    frwrd = [];
    bckwrd = [];
    for (scalers, mode) in zip(scaler_types, modes)
        scaler = scalers[1](mode);
        inv_scaler = scalers[2](scaler);
        push!(frwrd, scaler);
        push!(bckwrd, inv_scaler);
    end
    frwrd = Chain(frwrd...);
    bckwrd = Chain(bckwrd...);
    c = Recur(cell(state, dim));
    first_elem = SelectIndex(1);

    im = Chain(
        scale = frwrd,
        smooth = SkipConnection(
            Chain(
                get_values = first_elem,
                cell = c
            ),
            x_getsecond
        ),
        rescale = Chain(
                apply_factors = bckwrd,
                drop_idxs = first_elem
            ),
    );
    return PES(im);
end



function Base.getproperty(obj::PES, sym::Symbol)
    if sym == :inner_model
        return getfield(obj, sym);
    else
        return getproperty(obj.inner_model, sym);
    end
end



function Base.show(io::IO, m::PES)
    print(io, "PES(\n");
    print(io, "Model: ", m.inner_model);
    print(io, ")");
end



(m::PES{M})(inpt::Union{AbstractVecOrMat{I}, Tuple, NamedTuple}) where {I, M} =
    m.inner_model(inpt);



trainable(m::PES{M}) where {M} = (m.inner_model,);


@forward PES.inner_model Base.getindex
init_log_file(m::PES{M}, args...; kwargs...) where {M} =
    Ses.init_log_file(m.inner_model, args...; kwargs...);
function init_log_file(m::PES{M},
                       location::String = default_log_location,
                       name::String = "") where {M}
    _cell = m[:smooth].layers[:cell];
    log_name = name * Ses.log_appendix(_cell);
    log_name = log_name * Ses.log_appendix(m[:scale]);
    return Ses.init_log_file(_cell, location, log_name);
end
# @forward PES.inner_model log_appendix
function logging(m::PES{M}, logger, appendices) where {M}
    cell = m[:smooth].layers[:cell];
    for (idx, appendix) in enumerate(appendices)
        Ses.logging(cell, logger, "$appendix", "", idx);
    end
    Ses.logging(m[:scale], logger, "scaler");
end
@functor PES



restart!(m::PES{M}; kwargs...) where {M} =
    Ses.restart!(m[:smooth].layers[:cell]; kwargs...);



save(m::PES{M}; location = default_log_location) where {M} =
    Ses.save_model(location, m);
