import copy


def update_state_prior_dirichlet(
    pD, qs, lr=1.0, factors="all"
):
    """
    Update Dirichlet parameters of the initial hidden state distribution 
    (prior beliefs about hidden states at the beginning of the inference window).

    Parameters
    -----------
    pD: ``numpy.ndarray`` of dtype object
        Prior Dirichlet parameters over initial hidden state prior (same shape as ``qs``)
    qs: 1D ``numpy.ndarray`` or ``numpy.ndarray`` of dtype object
        Marginal posterior beliefs over hidden states at current timepoint
    lr: float, default ``1.0``
        Learning rate, scale of the Dirichlet pseudo-count update.
    factors: ``list``, default "all"
        Indices (ranging from 0 to ``n_factors - 1``) of the hidden state factors to include 
        in learning. Defaults to "all", meaning that factor-specific sub-vectors of ``pD``
        are all updated using the corresponding hidden state distributions.
    
    Returns
    -----------
    qD: ``numpy.ndarray`` of dtype object
        Posterior Dirichlet parameters over initial hidden state prior (same shape as ``qs``), after having updated it with state beliefs.
    """

    num_factors = len(pD)

    qD = copy.deepcopy(pD)
   
    if factors == "all":
        factors = list(range(num_factors))

    for factor in factors:
        idx = pD[factor] > 0 # only update those state level indices that have some prior probability
        qD[factor][idx] += (lr * qs[factor][idx])
       
    return qD