class ViewController: UIViewController {
    @IBOutlet var collectionView: UICollectionView!

    var colors: [UIColor] = [.red, .blue, .purple, .orange, .cyan, .green, .magenta, .lightGray, .blue, .green, .systemPink]

    override func viewDidLoad() {
        super.viewDidLoad()
        let layout = CustomLayout(unitSize: CGSize(width: 106, height: 88))
        collectionView.setCollectionViewLayout(layout, animated: true)
    }

    override func viewDidLayoutSubviews() {
        super.viewDidLayoutSubviews()
    }
}

extension ViewController: UICollectionViewDataSource {
    func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        return 10
    }

    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        let cell = collectionView.dequeueReusableCell(withReuseIdentifier: "cellId", for: indexPath) as! CustomCell

        cell.label.text = String(indexPath.row)
        cell.backgroundColor = colors[indexPath.row % colors.count]
        return cell
    }
}



class CustomLayout: UICollectionViewLayout {

    private var cellLayouts: [IndexPath: UICollectionViewLayoutAttributes] = [:]
    private var smallUnitSize: CGSize = .zero

    convenience init(unitSize: CGSize) {
        self.init()
        self.smallUnitSize = unitSize
    }

    override func prepare() {
        cellLayouts.removeAll()
        guard let numberOfRows = collectionView?.numberOfItems(inSection: 0) else { return }
        for aRow in 0..<numberOfRows {
            let indexPath = IndexPath(item: aRow, section: 0)
            let attributes = UICollectionViewLayoutAttributes(forCellWith: indexPath)

            let frame: CGRect
            switch indexPath.row {
                case 0:
                    frame = CGRect(x: 0, y: 0,
                                   width: smallUnitSize.width * 2, height: smallUnitSize.height * 2)
                case 1:
                    frame = CGRect(origin: CGPoint(x: 0, y: smallUnitSize.height * 2),
                                   size: smallUnitSize)
                case 2:
                    frame = CGRect(origin: CGPoint(x: smallUnitSize.width, y: smallUnitSize.height * 2),
                                   size: smallUnitSize)
                case 3:
                    frame = CGRect(origin: CGPoint(x: smallUnitSize.width * 2, y: 0),
                                   size: smallUnitSize)
                case 4:
                    frame = CGRect(origin: CGPoint(x: smallUnitSize.width * 2, y: smallUnitSize.height),
                                   size: smallUnitSize)
                case 5:
                    frame = CGRect(origin: CGPoint(x: smallUnitSize.width * 2, y: smallUnitSize.height * 2),
                                   size: smallUnitSize)
                default:
                    let rest = indexPath.row - 5
                    frame = CGRect(x: smallUnitSize.width * 3 * CGFloat(rest), y: 0,
                                   width: smallUnitSize.width * 3, height: smallUnitSize.height * 3)
            }
            attributes.frame = frame
            cellLayouts[indexPath] = attributes
        }
//        cellLayouts.sorted(by: { $0.key < $1.key }).forEach({ print("\($0.key): \($0.value.frame)") })
    }

    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return cellLayouts.compactMap { $0.value.frame.intersects(rect) ? $0.value : nil }
    }

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

    override var collectionViewContentSize: CGSize {
        guard let numberOfRows = collectionView?.numberOfItems(inSection: 0) else { return .zero }
        switch numberOfRows {
            case 0...5:
                return CGSize(width: smallUnitSize.width * 3, height: smallUnitSize.height * 3)
            default:
                let rest = numberOfRows - 6
                return CGSize(width: (smallUnitSize.width * 3) * CGFloat((1 + rest)), height: smallUnitSize.height * 3)
        }
    }

    override func shouldInvalidateLayout(forBoundsChange newBounds: CGRect) -> Bool {
        return true
    }
}