mutable struct SesLevel{C} <:AbstractSesCell
    operator::C
end

mutable struct SesTrend{C} <:AbstractSesCell
    operator::C
end

mutable struct SesSeason{C} <:AbstractSesCell
    operator::C
end

mutable struct SesSSM{L, T, S}
    level::L
    trend::T
    trend_mode::Symbol
    season::S
    season_mode::Symbol
end



function SesLevel(state::NTuple{2, T};
                  discount::Float32 = 0.88f0, lb::Float32 = 0.05f0, ub::Float32 = 0.95f0,
                  dim::Integer = 1, ewma::Bool = false) where {T}
    _cell = init_operator("ses", state,
                          discount = discount, lb = lb, ub = ub,
                          dim = dim, ewma = ewma);
    return SesLevel(_cell);
end
function SesTrend(state::NTuple{2, T};
                  discount::Float32 = 0.88f0, lb::Float32 = 0.05f0, ub::Float32 = 0.95f0,
                  dim::Integer = 1, ewma::Bool = false) where {T}
    _cell = init_operator("ses", state,
                          discount = discount, lb = lb, ub = ub,
                          dim = dim, ewma = ewma);
    return SesTrend(_cell);
end
function SesSeason(state::NTuple{2, T};
                   discount::Float32 = 0.88f0, lb::Float32 = 0.05f0, ub::Float32 = 0.95f0,
                   dim::Integer = 1, ewma::Bool = false) where {T}
    _cell = init_operator("ses", state,
                          discount = discount, lb = lb, ub = ub,
                          dim = dim, ewma = ewma);
    return SesSeason(_cell);
end
function SesSSM(;level::AbstractString = "", level_state::NTuple{2, T} = (0.0f0, 0.0f0),
                init_alpha::Float32 = 0.88f0, alpha_lb::Float32 = 0.05f0, alpha_ub::Float32 = 0.95f0,
                trend::AbstractString = "none", trend_state::Union{NTuple{2, T}, Nothing} = nothing,
                init_beta::Float32 = 0.88f0, beta_lb::Float32 = 0.05f0, beta_ub::Float32 = 0.95f0,
                season::AbstractString = "none", season_state::Union{NTuple{2, T}, Nothing} = nothing,
                init_gamma::Float32 = 1.0f0 - init_alpha,
                dim::Integer = 1) where {T}
    _level_sym = Symbol(lowercase(level))
    _trend_sym = Symbol(lowercase(trend));
    _season_sym = Symbol(lowercase(season))
    if _level_sym == :none &&
       _trend_sym == :none &&
       _season_sym == :none
        throw(ArgumentError("At least one between `level`, `trend` or `season` has to be not `none`."))
    end
    
    level_op = nothing;
    if _level_sym != :none
        level = SesLevel(level_state,
                        discount = init_alpha, lb = alpha_lb, ub = alpha_ub,
                        dim = dim, ewma = true);
    end

    trend_op = nothing;
    if _trend_sym != :none
        if _trend_sym == :mul
            throw(ArgumentError("$_trend_sym not in available trends: ['none', 'add']"))
        end
        trend_op = SesTrend(trend_state,
                            discount = init_beta,
                            lb = beta_lb, ub = beta_ub,
                            dim = dim);
    end

    season_op = nothing;
    if _season_sym != :none
        season_op = SesSeason(season_state,
                              discount = init_gamma, ub = 0.095f0 - init_alpha,
                              dim = dim);
    end

    ssm = SesSSM(level_op, trend_op, _trend_sym, season_op, _season_sym)
    return ssm;
end



function Base.show(io::IO, m::SesSSM)
    print(io, "SSM(\n");
    print(io, "Level: ", (if !isnothing(m.level) m.level else "N/D" end), "\n");
    print(io, "Trend: ", (if !isnothing(m.trend) m.trend else "N/D" end), "\n");
    print(io, "Season: ", (if !isnothing(m.season) m.season else "N/D" end));
    print(io, ")");
end



(ssml::SesLevel{C})(state::NTuple{2, T}, x::I) where {C<:AbstractSesCell, I, T} =
    ssml.operator(state, x);
function (ssmt::SesTrend{C})(state::NTuple{2, T}, x::I) where {C<:AbstractSesCell, I, T}
    
end
function (ssm::SesSSM{L, T, S})(x::I) where {L, S, T, I}

    _label = x;
    if !isnothing(ssm.season)
        _, _season = ssm.season.state;
        _label = if ssm.season_mode == :add
                _label - _season;
            else
                _label / _season;
            end
    end

    _plevel = 0.0f0;
    _ptrend = 0.0f0;
    if !isnothing(ssm.trend)
        _, _plevel = ssm.level.state;
        _, _ptrend = ssm.trend.state;
        ssm.level.state = (x, _plevel + _ptrend);
    end
    _level = ssm.level(_label);
    if ssm.level.step == inte_one
        increment_cell_step(ssm.level, 2);
    end

    _trend = 0.0f0;
    if !isnothing(ssm.trend)
        _trend = ssm.trend(_level - _plevel);
    end
    
    ### TODO: Lag operation is needed here. Code left as a placeholder
    _season = if ssm.season_mode == :add 0.0f0 else 1.0f0 end
    if !isnothing(ssm.season)
        _plb = (_plevel + _ptrend);
        _label = if ssm.season_mode == :add
            _label - _plb;
        else
            _label / _plb;
        end
        _season = ssm.season(_label)
    end

    yhat = _level + _trend;
    yhat = if ssm.season_mode == :add
        yhat + _season;
    else
        yhat * _season;
    end
    return yhat;
end



function reset!(ssm::SesSSM)
    if isa(ssm.level, AbstractVector)
        for op in ssm.level
            reset!(op);
        end
    else
        reset!(ssm.level);
    end

    if !isnothing(ssm.trend)
        if isa(ssm.trend, AbstractVector)
            for op in ssm.trend
                reset!(op);
            end
        else
            reset!(ssm.trend);
        end
    end

    if !isnothing(ssm.season)
        if isa(ssm.season, AbstractVector)
            for op in ssm.season
                reset!(op);
            end
        else
            reset!(ssm.season);
        end
    end
end



function restart!(ssm::SesSSM; kwargs...)
    if isa(ssm.level, AbstractVector)
        for op in ssm.level
            restart!(op);
        end
    else
        restart!(ssm.level);
    end

    if !isnothing(ssm.trend)
        if isa(ssm.trend, AbstractVector)
            for op in ssm.trend
                restart!(op);
            end
        else
            restart!(ssm.trend);
        end
    end

    if !isnothing(ssm.season)
        if isa(ssm.season, AbstractVector)
            for op in ssm.season
                restart!(op);
            end
        else
            restart!(ssm.season);
        end
    end
end



function testmode!(ssm::SesSSM, args...; kwargs...)
    if isa(ssm.level, AbstractVector)
        for op in ssm.level
            testmode!(op, args...; kwargs...);
        end
    else
        testmode!(ssm.level, args...; kwargs...);
    end

    if !isnothing(ssm.trend)
        if isa(ssm.trend, AbstractVector)
            for op in ssm.trend
                testmode!(op, args...; kwargs...);
            end
        else
            testmode!(ssm.trend, args...; kwargs...);
        end
    end

    if !isnothing(ssm.season)
        if isa(ssm.season, AbstractVector)
            for op in ssm.season
                testmode!(op, args...; kwargs...);
            end
        else
            testmode!(ssm.season, args...; kwargs...);
        end
    end
end



@functor SesLevel
@functor SesTrend
@functor SesSeason
@functor SesSSM



function trainable(ssm::SesSSM)
    t = (ssm.level,)

    if !isnothing(ssm.trend)
        t = (t..., ssm.trend)
    end
    
    if !isnothing(ssm.season)
        t = (t..., ssm.season)
    end

    return t
end
trainable(ssml::SesLevel) = (ssml.operator,)
trainable(ssmt::SesTrend) = (ssmt.operator,)
trainable(ssms::SesSeason) = (ssms.operator,)