from extraction import split_into_3


def splice(data_dir, count):
    """
    Generalized version of splice_encreg_{single,dual,triple,quad}.
    """
    if not isinstance(count, int) or count < 1:
        raise ValueError(f"count must be a positive int (1,2,3,...). Got: {count}")

    layer_idx = count - 1
    norm_out = f"/encoder/layers.{layer_idx}/norm2/LayerNormalization_output_0"

    split_into_3(
        data_dir,
        [['hits', 'padding_mask'], ['/input_layer/Add_output_0', '/encoder/Where_output_0']],
        [['/input_layer/Add_output_0', '/encoder/Where_output_0'], [norm_out]],
        [[norm_out, '/encoder/Where_output_0'], ['1203']]
    )
