struct LevelTrend{L<:AbstractSesCell, T}
    level::Union{Recur{L}, Nothing}
    trend::Union{Recur{L}, Nothing}
    trend_mode::Int8
    theta_proxy::T
end



function LevelTrend(level::Recur{L}, trend::Recur{L};
                    mode = "add", theta::Float32 = -1.0f0) where {L<:AbstractSesCell}
    mode_sym = Symbol(mode);
    mode_int = Ses.to_intmode[mode_sym];
    thetaV = [theta];
    theta_proxy = softplus_inverse.(thetaV);
    return LevelTrend(level, trend, mode_int, theta_proxy);
end



function (lt::LevelTrend{L, T})(x::AbstractVecOrMat{I}) where {I, L, T}
    println("$x\nIN -->");
    [lt(xi) for xi in x];
end
function (lt::LevelTrend{L, T})(x::I) where {I, L, T}
    _llabel, _pplevel = lt.level.state;
    _, _ptrend = lt.trend.state;

    if !warmup_done(lt.level)
        _level = lt.level(x);
        return _level[1];
    end

    yhat = apply_trend(lt, _pplevel, _ptrend);
    # lt.level.state = (_llabel, _pplevel .+ _ptrend);
    lt.level.state = (_llabel, yhat);
    _level = lt.level(x);

    # _tlabel = if is_additive(lt.trend_mode) _level .- _pplevel
    #           else _level ./ _pplevel
    #           end;
    _tlabel = apply_trend(lt, _level, _pplevel, inverse=true);
    _trend = lt.trend(_tlabel);
    
    # yhat = if is_additive(lt.trend_mode) _pplevel .+ _ptrend else _pplevel .* _ptrend end;
    return yhat[1];
end



function apply_trend(lt::LevelTrend,
                     x::Union{I, AbstractVecOrMat{I}}, trend::Union{T, AbstractVecOrMat{T}};
                     inverse::Bool=false) where {I, T}
    _additive = Ses.is_additive(lt.trend_mode);
    _trend = if inverse && _additive
        .-trend;
    elseif inverse && !_additive
        inv.(trend);
    else
        trend
    end;
    return if _additive x .+ _trend else x .* _trend end;
end



@forward LevelTrend.level alpha
beta(lt::LevelTrend) = alpha(lt.trend);
theta(lt::LevelTrend) = softplus.(lt.theta_proxy);



function Base.getproperty(obj::LevelTrend, sym::Symbol)
    if sym in [:level, :trend, :trend_mode, :theta_proxy]
        return getfield(obj, sym);
    end

    if sym == :alpha_proxy
        return if !isnothing(obj.level) getproperty(obj.level.cell, :alpha_proxy) else nothing end;
    elseif sym == :beta_proxy
        return if !isnothing(obj.trend) getproperty(obj.trend.cell, :alpha_proxy) else nothing end;
    end

    if sym == :lstate
        return if !isnothing(obj.level) getproperty(obj.level, :state) else nothing end;
    elseif sym == :tstate
        return if !isnothing(obj.trend) getproperty(obj.trend, :state) else nothing end;
    end

    if sym == :lstate0
        return if !isnothing(obj.level) getproperty(obj.level.cell, :state0) else nothing end;
    elseif sym == :tstate0
        return if !isnothing(obj.trend) getproperty(obj.trend.cell, :state0) else nothing end;
    end

end



function Base.show(io::IO, lt::LevelTrend)
    print(io, "Level-Trend SSM(");
    szl = size(lt.alpha_proxy);
    szt = size(lt.beta_proxy);
    print(io, "\nLevel α: ", (if szl[1] * szl[2] == 1 lt.alpha_proxy[1] else szl end), ", state0: ", lt.lstate0);
    print(io, "\nTrend β: ", (if szt[1] * szt[2] == 1 lt.beta_proxy[1] else szt end), ", state0: ", lt.tstate0);
    print(io, ")");
end



@functor LevelTrend



function restart!(lt::LevelTrend; kwargs...)
    restart!(lt.level; kwargs...);
    restart!(lt.trend; kwargs...);
end



trainable(lt::LevelTrend) = (lt.level,lt.trend,lt.theta_proxy,)



function nb_loss(lt::LevelTrend, yhats::AbstractVector{T}, labels::AbstractVector{T}) where {T}
    theta = Ses.theta(lt);
    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