from typing import List
from typing import DefaultDict
from typing import Tuple
import copy


def laplace_smooth_param_counts(
    cmds: List[str],
    param_counts: DefaultDict[str, int],
    cmd_param_counts: DefaultDict[str, DefaultDict[str, int]],
    unk_token: str,
) -> Tuple[DefaultDict[str, int], DefaultDict[str, DefaultDict[str, int]]]:
    """
    Apply laplace smoothing to the input counts for the params.

    In particular, add 1 to each of the counts, including the unk_token. By including the
    unk_token, we can handle unseen params.

    Parameters
    ----------
    cmds: List[str]
        list of all the possible commands (including the unk_token)
    param_counts: DefaultDict[str, int]
        individual param counts
    cmd_param_counts: DefaultDict[str, DefaultDict[str, int]]
        param conditional on command counts
    unk_token: str
        dummy command to signify an unseen command (e.g. "##UNK##")

    Returns
    -------
    Tuple:
        individual param probabilities,
        param conditional on command probabilities

    """
    param_counts_ls = copy.deepcopy(param_counts)
    cmd_param_counts_ls = copy.deepcopy(cmd_param_counts)

    params: List[str] = list(param_counts.keys()) + [unk_token]
    for cmd in cmds:
        for param in params:
            if param in cmd_param_counts_ls[cmd] or param == unk_token:
                param_counts_ls[param] += 1
                cmd_param_counts_ls[cmd][param] += 1

    return param_counts_ls, cmd_param_counts_ls