import math


def z_dir_imp_factor(z, num_waves, total_wave_length, **kwargs):
    """ Returns the imperfection factor based on the z-direction.
    :param float z: Coordinate in length direction.
    :param int num_waves: Number of half-wave lengths, either 1 or 2.
    :param float total_wave_length: Range of the local imperfection along the length.
    :return float: The factor in the length direction.
    """
    if num_waves == 2:
        z_norm_factor = 0.3849  # to normalize h to maximum of 1.0
        h = math.sin(math.pi * z / total_wave_length) ** 2 * \
            math.cos(math.pi * z / total_wave_length) / z_norm_factor
    elif num_waves == 1:
        h = math.sin(math.pi * z / total_wave_length) ** 2
    else:
        raise ValueError('num_waves should be either 1 or 2')
    return h