Jetpack Compose Paging Group Items Under Header

@Composablefun SelectPhotoScreen(viewModel: SelectPhotoViewModel = viewModel(), albumId: String) {    val lazyImages = viewModel.fetchLocalImage(albumId = albumId).collectAsLazyPagingItems()    LazyColumn {        val itemCount = lazyImages.itemCount        // pre-process        var lastDateGroup: String? = null        // var images = mutableListOf<Gallery.LocalImage>()        val dayBatch = mutableMapOf<String, MutableList<Int>>()        val rowSize = 3        for (index in 0 until itemCount) {            // val nextImage =  lazyImages[index]            // prevent pageload            val nextImage =  lazyImages.peek(index)            if (nextImage != null) {                // format datetime to date string                val date = nextImage.created.format(DateTimeFormatter.ISO_LOCAL_DATE)                if (lastDateGroup == null || date != lastDateGroup) {                    dayBatch[date] = mutableListOf()                    lastDateGroup = date                }                dayBatch[lastDateGroup]?.add(index)            }        }        for ((dateGroup, imageIndexes) in dayBatch) {            stickyHeader(key = dateGroup) {                Text(                    text = dateGroup,                    color = MaterialTheme.colors.primary,                    style = MaterialTheme.typography.h6,                    modifier = Modifier.padding(vertical = 4.dp, horizontal = 16.dp)                )            }            item(key = "$dateGroup-images") {                // show 3 image per row                Grid(                    rowSize = rowSize,                    chunks = imageIndexes.chunked(rowSize)                ) { _, imageIndex, modifier ->                    // trigger pageload                    val localImage = lazyImages[imageIndex]                    PhotoItem(                        modifier = modifier,                        data = localImage?.uri                    )                }            }        }    }}