import torch
from typing import Tuple


def build_param_list(
        input_dim: int,
        output_dim: int,
        device: torch.device,
        hidden_units: tuple = (64, 64),
        requires_grad: bool = True
) -> Tuple[list, list]:
    """
    Build parameter list of the neural network

    Parameters
    ----------
    input_dim: int
        dimension of the input of the neural network
    output_dim: int
        dimension of the output of the neural network
    device: torch.device
        cpu or cuda
    hidden_units: tuple
        hidden units of the neural network
    requires_grad: bool
        whether the parameters need to be trained

    Returns
    -------
    weights: a list of torch.Tensor
        weights of each layer of the neural network
    biases: a list of torch.Tensor
        biases of each layer of the neural network
    """
    weights = []
    biases = []
    units = input_dim
    for next_units in hidden_units:
        weights.append(torch.zeros(next_units, units, requires_grad=requires_grad).to(device))
        biases.append(torch.zeros(next_units, requires_grad=requires_grad).to(device))
        units = next_units
    weights.append(torch.zeros(output_dim, units, requires_grad=requires_grad).to(device))
    biases.append(torch.zeros(output_dim, requires_grad=requires_grad).to(device))
    return weights, biases