mutable struct SesNbOperator{T} <: AbstractSesCell
    ses_op::AbstractSesCell
    theta_proxy::T
end



function SesNbOperator(state::NTuple{2, T}, dim::Integer = 1;
                       discount::II = 0.88f0, lb = 0.05f0, ub = 0.95f0,
                       theta::II = -1.0f0, ewma = false) where {I, II<:Union{I, AbstractVecOrMat{I}}, T<:Union{I, AbstractVecOrMat{I}}}
    _ses_op = SesOperator(state, dim,
                          discount = discount, lb = lb, ub = ub,
                          ewma = ewma
                          );

    szt = size(theta);
    numeric_theta = isempty(szt);
    thetaV = if numeric_theta
        fill(theta, (dim, 1))
    elseif szt[1] == dim 
        theta
    else
        throw(DimensionMismatch("argument `theta` has dimension $szt, expected ($dim, 1)."));
    end;
    theta_proxy = softplus_inverse.(thetaV);
    cell = SesNbOperator(_ses_op, theta_proxy);
    return cell;
end


@forward SesNbOperator.ses_op alpha
@forward SesNbOperator.ses_op restart!
@forward SesNbOperator.ses_op state0
@forward SesNbOperator.ses_op warmup_done
theta(op::SesNbOperator{T}) where {T} = softplus.(op.theta_proxy);



function Base.getproperty(obj::SesNbOperator, sym::Symbol)
    if sym in [:ses_op, :theta_proxy]
        return getfield(obj, sym);
    else
        return getproperty(obj.ses_op, sym);
    end
end



function Base.show(io::IO, c::SesNbOperator)
    print(io, "Negative Bin. Operator(\n");
    sz = size(c.alpha_proxy);
    print(io, "α: ", (if sz[1] * sz[2] == 1 c.alpha_proxy[1] else sz end), "\n");
    sz = size(c.theta_proxy);
    print(io, "θ: ", (if sz[1] * sz[2] == 1 c.theta_proxy[1] else sz end), "\n");
    print(io, "state0: ", c.state0);
    print(io, ")");
end



(op::SesNbOperator{<:AbstractVecOrMat{T}})(state, x::Union{I, AbstractVecOrMat{I}}) where {I, T} =
    op.ses_op(state, x);



@functor SesNbOperator



testmode!(op::SesNbOperator, mode=true) =
  (op.ses_op.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; op)



trainable(op::SesNbOperator) = (op.ses_op, op.theta_proxy,)