Ball grid animation

BallGrid.kt

@OptIn(ExperimentalFoundationApi::class)
@Composable
fun AnimatedBallGrid(
    numBalls: Int = 32,
    numBallsPrRow: Int = 7,
) {
    // viewModel contains an observable list of balls
    val viewModel: BallsViewModel by remember { mutableStateOf(BallsViewModel(numBalls)) } // TODO: Use injection instead
    var animationSpeed: AnimationSpeed by remember { mutableStateOf(AnimationSpeed.Fast) }
    val ballSizeInclPadding = 44.dp // TODO: Hardcoded value... Do this better

    Column {
        LazyVerticalGrid(
            modifier = Modifier
                .width(ballSizeInclPadding * numBallsPrRow)
                .height(ballSizeInclPadding * (numBalls / numBallsPrRow.toFloat() + 0.5f).roundToInt()),
            cells = GridCells.Fixed(numBallsPrRow),
        ) {
            itemsIndexed(viewModel.balls) { index, item ->
                Ball(
                    number = item,
                    animationDelay = getDiagonal(index, numBallsPrRow) * animationSpeed.speedFactor,
                    onBallClick = { /* TODO */ }
                )
            }
        }
    }
}

@OptIn(ExperimentalAnimationApi::class)
@Composable
fun Ball(
    ball: BallModel,
    animationDelay: Long = 0L,
    onBallClick: (LottoBallValue) -> Unit
) {
    val entryAnimationSpec = tween<Dp>(durationMillis = 100, delayMillis = animationDelay.toInt())
    val ballSizeExtend by animateDpAsState(
        targetValue = if (ball.isVisible) 40.dp else 20.dp,
        animationSpec = if (ball.isVisible) entryAnimationSpec else spring(visibilityThreshold = Dp.VisibilityThreshold)
    )

    Box(
        modifier = Modifier
            .requiredSize(44.dp)
            .clickable { onBallClick(ball) },
        contentAlignment = Alignment.Center,
    ) {
        Box(
            modifier = Modifier
                .padding(2.dp)
                .clip(CircleShape)
                .background(Color.Red)
                .requiredSize(ballSizeExtend),
            contentAlignment = Alignment.Center
        ) {
            AnimatedVisibility(visible = (ballSizeExtend == 40.dp)) {
                Text(
                    text = ball.value?.toString() ?: "",
                    color = Color.White,
                    fontWeight = FontWeight.Bold,
                    modifier = Modifier.align(Alignment.Center)
                )
            }
        }
    }
}

/**
 * Will return on which diagonal (from left to right) an element at `index` in list
 * is on, given `numElementsPrRow`.
 *
 * E.g. element at index 0 belongs to diagonal 0.
 * Given 5 element pr row; element with index 5 (first element on second row) is on diagonal 1.
 */
fun getDiagonal(index: Int, numElementsPrRow: Int): Int {
    val row = index / numElementsPrRow
    val col = index % numElementsPrRow
    return col + row
}