include("abstract.jl")
include("simple.jl")
include("negative_binomial.jl")
include("ssm.jl")

include("logging.jl")

# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# ====> Available Operators <====
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
available_cells = Dict{String, Type{<:AbstractSesCell}}(
    "ses" => SesOperator,
    "nb" => SesNbOperator
)



# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# ====> Class Properties <====
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@forward Recur.cell alpha
@forward Recur.cell theta
@forward Recur.cell increment_cell_step
@forward Recur.cell nb_loss
@forward Recur.cell warmup_done



# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# ====> 'SEScells' DataFrame preprocessing <====
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
function check_columns(df::DataFrame, op::Type{<:AbstractSesCell}; message::String = "")
    input_cols = propertynames(first(df));
    expected_cols = columns(op);
    _check_columns(expected_cols, input_cols, message = message);
end


# function add_state!(df::DataFrame)
#     check_columns(df, SesOperator, message = ":Qty is mandatory to create initial state");
#     transform!(df, :Qty => ByRow(qty -> (Float32(qty[2]), Float32(qty[1]))) => :H);
# end


# function add_labels!(df::DataFrame)
#     check_columns(df, SesOperator, message = ":Qty is mandatory to create initial state");
#     transform!(df, :Qty => ByRow(qty -> Float32.(qty[3:end])) => :Labels);
# end


# function create_model!(df::DataFrame, cell::Type{<:AbstractSesCell})
#     transform!(df, :H =>
#                     ByRow(h ->
#                         SesMap(cell(h))
#                         )
#                     => :Model);
# end


# function create_recurrent_model!(df::DataFrame, cell::Type{<:AbstractSesCell})
#     transform!(df, :H =>
#                     ByRow(h ->
#                         SesMap(
#                             Recur(cell(h))
#                         ))
#                     => :Model);
# end



function init_operator(cell::AbstractString, state::NTuple{2, T};
                       discount::Float32, lb::Float32 = 0.05f0, ub::Float32 = 0.95f0,
                       dim::Integer = 1, ewma::Bool = false) where {T}
    if !(cell in keys(available_cells))
        throw(ArgumentError("$cell is not a valid cell type. Use $(keys(available_cells))."))
    end
    op = available_cells[cell];

    if discount <= lb || discount >= ub
        throw(ArgumentError("$discount is out of range [$lb, $ub]."))
    end

    if dim < 1
        throw(ArgumentError("$dim has to be at least 1."))
    end

    return op(state,
              discount = discount, lb = lb, ub = ub,
              dim = dim, ewma = ewma);
end


increment_cell_step(cell::Type{<:AbstractSesCell}, inc::Number = 1) = cell.step += Int8(inc);



# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# ====> 'SEScells' Training <====
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
function op_loss(::SesOperator, loss, yhat::AbstractVecOrMat{I}, labels::AbstractVecOrMat{I}) where {I}
    ls = sum(loss.(yhat, labels));
    n = length(labels);
    return 1 / n * ls;
end
op_loss(::SesOperator, loss, yhat::I, labels::I) where {I} = loss(yhat, labels)
function op_loss(op::SesNbOperator, _, yhat::AbstractVecOrMat{I}, labels::AbstractVecOrMat{I}) where {I}
    ls = sum(op_loss.(op, loss, yhat, labels), 1);
    n = length(ls);
    return 1/n * ls
end
function op_loss(op::SesNbOperator, _, yhat::Float32, labels::Float32)
    theta = Ses.theta(op);
    # nb_r = theta .* yhat;
    nb_r = yhat * inv.(theta);
    r = softplus.(nb_r);
    # nb_p = theta ./ (1.0f0 .+ theta);
    nb_p = 1.0f0 ./ (1.0f0 .+ theta);
    p = sigmoid.(nb_p);
    ls = negative_binomial_loglikelihood.(r, p, labels)[:];
    return sum(ls)
end



function predict(op::SesNbOperator, horizon::Integer = 1, num_samples::Integer = 200)
    trajectories = [];
    for _ in 1:num_samples
        alpha = alpha(op);
        theta = theta(op);
        _, state = op.state;
        trajectory = [];
        for _ in 1:horizon
            tmpr = state .* inv.(theta);
            r = softplus.(tmpr);
            tmpp = theta ./ (1.0f0 .+ theta);
            p = sigmoid.(tmpp);
            label = NegativeBinomial(r, p);
            sample = smooth(op.ses_op, label, state, alpha);
            state = sample;
            push!(trajectory, sample);
        end
        push!(trajectories, trajectory);
    end
    return trajectories;
end



function nb_loss(op::SesNbOperator, yhats::AbstractVector{T}, labels::AbstractVector{T}) where {T}
    theta = Ses.theta(op);
    tmpr = yhats .* inv.(theta);
    r = softplus.(tmpr);
    tmpp = theta ./ (1.0f0 .+ theta);
    p = sigmoid.(tmpp);
    ls = negative_binomial_loglikelihood.(r, p, labels);
    return ls;
end
function nb_loss(op::SesNbOperator, yhat::Float32, label::Float32)
    theta = Ses.theta(op);
    nb_r = theta .* yhat;
    # nb_r = yhat * inv.(theta);
    r = softplus.(nb_r);
    nb_p = theta ./ (1.0f0 .+ theta);
    # nb_p = 1.0f0 ./ (1.0f0 .+ theta);
    p = sigmoid.(nb_p);
    ls = negative_binomial_loglikelihood.(r, p, label)[:];
    return ls;
end
