Custom header in custom UICollectionViewLayout

Question

Im trying to achieve this UI:

from this tutorial https://www.raywenderlich.com/4829472-uicollectionview-custom-layout-tutorial-pinterest

I'm already able to nest 2 collectionView with different scroll directions:

But now im stuck with adding 2 custom headerView for each section. I did take a look at this question and tried to modify my code according to it but it didn't work well:

What I got:

As you can see, only 1 header shows up (expected 2) and it is overlapped with the cell.

Here is the code for header:

   override func collectionView(_ collectionView: UICollectionView, viewForSupplementaryElementOfKind kind: String, at indexPath: IndexPath) -> UICollectionReusableView {
        let header = collectionView.dequeueReusableSupplementaryView(ofKind: kind, withReuseIdentifier: HeaderHomeView.identifier, for: indexPath) as? HeaderHomeView
        return header!
    }

//Inside custom UICollectionViewLayout

  override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        var visibleLayoutAttributes: [UICollectionViewLayoutAttributes] = []
        
        // Loop through the cache and look for items in the rect
        let sectionsCount = self.collectionView!.dataSource!.numberOfSections!(in: self.collectionView!)
        for section in 0..<sectionsCount {
            visibleLayoutAttributes.append(self.layoutAttributesForSupplementaryView(ofKind: UICollectionView.elementKindSectionHeader, at: NSIndexPath(item: 0, section: section) as IndexPath)!)
        }
        for attributes in cache {
            if attributes.frame.intersects(rect) {
                visibleLayoutAttributes.append(attributes)
            }
        }
        return visibleLayoutAttributes
    }
    override func layoutAttributesForSupplementaryView(ofKind elementKind: String, at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        let attributes = UICollectionViewLayoutAttributes(forSupplementaryViewOfKind: elementKind, with: indexPath as IndexPath)
        attributes.frame = CGRect(x: 0, y: 0, width: contentWidth, height: 100)
        
        return attributes
    }

Edit: I tried to print each UICollectionViewLayoutAttributes and it seems like I have to calculate the x,y for each header and I still haven't figured out.

<UICollectionViewLayoutAttributes: 0x7f833f93b3f0> index path: (<NSIndexPath: 0xb3260380861bce99> {length = 2, path = 0 - 0}); element kind: (UICollectionElementKindSectionHeader); frame = (0 0; 358 100); 

<UICollectionViewLayoutAttributes: 0x7f833f93b550> index path: (<NSIndexPath: 0xb3260380861bcf99> {length = 2, path = 1 - 0}); element kind: (UICollectionElementKindSectionHeader); frame = (0 0; 358 100); 

<UICollectionViewLayoutAttributes: 0x7f833f93ac60> index path: (<NSIndexPath: 0xb3260380861bce99> {length = 2, path = 0 - 0}); frame = (6 6; 346 100); 

<UICollectionViewLayoutAttributes: 0x7f833f939360> index path: (<NSIndexPath: 0xb3260380861bcf99> {length = 2, path = 1 - 0}); frame = (6 118; 167 292); 

<UICollectionViewLayoutAttributes: 0x7f833f9394c0> index path: (<NSIndexPath: 0xb3260380863bcf99> {length = 2, path = 1 - 1}); frame = (185 118; 167 328); 

Here is the full code of my custom UICollectionViewLayOut

class PinterestLayout: UICollectionViewLayout {
    // 1
    weak var delegate: PinterestLayoutDelegate?
    
    // 2
    var numberOfColumns = 2
    private let cellPadding: CGFloat = 6
    
    // 3
    private var cache: [UICollectionViewLayoutAttributes] = []
    
    // 4
    private var contentHeight: CGFloat = 0
    
    private var contentWidth: CGFloat {
        guard let collectionView = collectionView else {
            return 0
        }
        let insets = collectionView.contentInset
        return collectionView.bounds.width - (insets.left + insets.right)
    }
    
    // 5
    override var collectionViewContentSize: CGSize {
        return CGSize(width: contentWidth, height: contentHeight)
    }
    
    override func prepare() {
        // 1
        guard
            cache.isEmpty == true,
            let collectionView = collectionView
        else {
            return
        }
        // 2
        let columnWidth = contentWidth / CGFloat(numberOfColumns)
        var xOffset: [CGFloat] = []
        for column in 0..<numberOfColumns {
            xOffset.append(CGFloat(column) * columnWidth)
        }
        var column = 0
        var yOffset: [CGFloat] = .init(repeating: 0, count: numberOfColumns)
        
        // 3
        for section in 0..<collectionView.numberOfSections {
            for item in 0..<collectionView.numberOfItems(inSection: section) {
                let indexPath = IndexPath(item: item, section: section)
                let attributes: UICollectionViewLayoutAttributes
                let photoHeight = delegate?.collectionView(
                    collectionView,
                    heightForPhotoAtIndexPath: indexPath) ?? 180
                let height = cellPadding * 2 + photoHeight
                let frame: CGRect
                
                if section == 0 {
                    frame = CGRect(x: 0,
                                   y: yOffset[column],
                                   width: contentWidth,
                                   height: height)
                    
                    contentHeight = frame.maxY
                    for column in 0..<numberOfColumns {
                        yOffset[column] = frame.maxY
                    }
                    column = 0
                }   else {
                    // 4
                    frame = CGRect(x: xOffset[column],
                                   y: yOffset[column],
                                   width: columnWidth,
                                   height: height)
                    
                    contentHeight = max(contentHeight, frame.maxY)
                    yOffset[column] = yOffset[column] + height
                    column = column < (numberOfColumns - 1) ? (column + 1) : 0
                }
                
                let insetFrame = frame.insetBy(dx: cellPadding, dy: cellPadding)
                attributes = UICollectionViewLayoutAttributes(forCellWith: indexPath)
                attributes.frame = insetFrame
                cache.append(attributes)
            }
        }
    }
    
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        var visibleLayoutAttributes: [UICollectionViewLayoutAttributes] = []
        
        // Loop through the cache and look for items in the rect
        let sectionsCount = self.collectionView!.dataSource!.numberOfSections!(in: self.collectionView!)
        for section in 0..<sectionsCount {
            visibleLayoutAttributes.append(self.layoutAttributesForSupplementaryView(ofKind: UICollectionView.elementKindSectionHeader, at: NSIndexPath(item: 0, section: section) as IndexPath)!)
        }
        for attributes in cache {
            if attributes.frame.intersects(rect) {
                visibleLayoutAttributes.append(attributes)
            }

        }
        for each in visibleLayoutAttributes {
            print("\(each)\n")
        }
        return visibleLayoutAttributes
    }
    override func layoutAttributesForSupplementaryView(ofKind elementKind: String, at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        let attributes = UICollectionViewLayoutAttributes(forSupplementaryViewOfKind: elementKind, with: indexPath as IndexPath)
        attributes.frame = CGRect(x: 0, y: 0, width: contentWidth, height: 100)

        return attributes
    }

    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return cache[indexPath.item]
    }
}