paiai_ios/PaiAi/PaiaiUIKit/Reusable/UIKit/WaterfallFlowLayout/WaterfallFlowLayout.swift

//
//  WaterfallFlowLayout.swift
//  PaiAi
//
//  Created by FFIB on 2017/11/13.
//  Copyright © 2017年 FFIB. All rights reserved.
//

import UIKit

public final class WaterfallFlowLayout: UICollectionViewLayout {

    private var minColumn: Int = 0
    private var itemWidth: CGFloat = -1
    private var columnHeights: [CGFloat] = []
    private var minColumnHeight: CGFloat = 0
    private(set) var configuration = WaterfallFlowConfiguration()
    private var attributesArr: [UICollectionViewLayoutAttributes] = []
    private var isNeedLayout: Bool = true
    override public init() {
        super.init()
    }

    convenience init(configuration: WaterfallFlowConfiguration) {
        self.init()
        self.configuration = configuration
    }
    
    required init?(coder aDecoder: NSCoder) {
        super.init(coder: aDecoder)
    }

    override public func prepare() {
        super.prepare()
        initialize()
    }

    override public var collectionViewContentSize: CGSize {
        return calculateViewSize()
    }

    override public func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        guard attributesArr.count <= indexPath.row, isNeedLayout else {
            return attributesArr[indexPath.row]
        }

        let itemX = calculateItemX()
        let itemY = calculateItemY()
        let itemHeight = calculateItemHeight(indexPath: indexPath)
        let attributes = UICollectionViewLayoutAttributes(forCellWith: indexPath)
        attributes.frame = CGRect(x: itemX, y: itemY, width: itemWidth, height: itemHeight)

        setMinColumn(height: itemHeight)
        return attributes
    }

    override public func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return attributesArr
    }

    fileprivate func initialize() {
        guard collectionView?.numberOfSections == 1,
            let itemCount = collectionView?.numberOfItems(inSection: 0) else { return }

        let originIndex: Int
        if attributesArr.count != itemCount || itemCount == 0 || itemWidth == -1 || isNeedLayout {
            minColumn = 0
            originIndex = 0
            minColumnHeight = 0
            attributesArr.removeAll()
            calculateItemWidth()
            columnHeights = Array(repeating: 0, count: configuration.columnCount)
        } else {
            originIndex = attributesArr.count
        }
        
        if itemCount == 0 {
            columnHeights = []
            return
        }
        
        for i in originIndex..<itemCount {
            guard let attributes = layoutAttributesForItem(at: IndexPath(row: i, section: 0))
                else { continue }
            attributesArr.append(attributes)
        }
    }

    fileprivate func calculateViewSize() -> CGSize {
        guard let v = collectionView,
            let maxH = columnHeights.max() else { return CGSize.zero }
        return CGSize(width: v.bounds.width, height: maxH + configuration.rowSpace)
    }

    fileprivate func calculateItemX() -> CGFloat {
        return CGFloat(minColumn) * itemWidth + CGFloat(minColumn + 1) * configuration.columnSpace
    }

    fileprivate func calculateItemY() -> CGFloat {
        minColumnHeight += configuration.rowSpace
        return minColumnHeight
    }

    fileprivate func calculateItemWidth() {
        let width = collectionView?.bounds.width ?? 0
        let space = configuration.columnSpace * (CGFloat(configuration.columnCount + 1))
        itemWidth = (width - space) / CGFloat(configuration.columnCount)
    }

    fileprivate func calculateItemHeight(indexPath: IndexPath) -> CGFloat {
        guard let collectionView = collectionView,
            let delegate = collectionView.delegate as? UICollectionViewDelegateFlowLayout
            else { return 0 }

        let itemOriginSize = delegate.collectionView!(collectionView,
                                                      layout: self,
                                                      sizeForItemAt: indexPath)
        return itemWidth / itemOriginSize.width * itemOriginSize.height
    }

    fileprivate func setMinColumn(height: CGFloat) {
        minColumnHeight += height
        columnHeights[minColumn] = minColumnHeight
        (minColumn, minColumnHeight) = columnHeights.enumerated().min(by: { $0.1 < $1.1 }) ?? (0, 0)
    }
    
    /// called at collectionView reload
    func setNeedsLayout() {
        isNeedLayout = true
    }
}