//声明LYPWaterFlowLayout为一个类
@class LYPWaterFlowLayout;
@protocol LYPWaterFlowLayoutDelegate <NSObject>
//必须实现的方法
@required
/**获取瀑布流每个元素的高度*/
- (CGFloat)waterFlowLayout:(LYPWaterFlowLayout *)waterFlowLayout heightForItemAtIndex:(NSInteger)index itemWith:(CGFloat)itemWith;
//可选实现的方法
@optional
/**获取瀑布流的列数*/
- (NSInteger)columnCountInWaterFlowLayout:(LYPWaterFlowLayout *)waterFlowLayout;
/**获取瀑布流列间距*/
- (CGFloat)columnMarginInWaterFlowLayout:(LYPWaterFlowLayout *)waterFlowLayout;
/**获取瀑布流的行间距*/
- (CGFloat)rowMarginInWaterFlowLayout:(LYPWaterFlowLayout *)waterFlowLayout;
/**获取瀑布流的内边距*/
- (UIEdgeInsets)edgeInsetsInWaterFlowLayout:(LYPWaterFlowLayout *)waterFlowLayout;
@end

@interface LYPWaterFlowLayout : UICollectionViewLayout
/**代理*/
@property (nonatomic, weak) id<LYPWaterFlowLayoutDelegate> delegate;
@end

/**默认的列数*/
static const NSInteger LYPDefaultColumnCount = 3;
/**默认每一列之间的间距*/
static const CGFloat LYPDefaultColumMargin = 10;
/**默认每一行之间的间距*/
static const CGFloat LYPDefaultRowMargin = 10;
/**默认边缘间距*/
static const UIEdgeInsets LYPDefaultEdgeInsets = {10, 10, 10, 10};

- (NSInteger)columnCount
{
  //判断代理是否实现了获取列数的可选方法
  if ([self.delegate respondsToSelector:@selector(columnCountInWaterFlowLayout:)])
  {
    //实现，返回通过代理设置的列数
    return [self.delegate columnCountInWaterFlowLayout:self];
  }
  else
  {
    //为实现，返回默认的列数
    return LYPDefaultColumnCount;
  }
}

/**所有cell的布局属性*/
@property (nonatomic, strong) NSMutableArray *attrsArray;
/**所有列的当前高度*/
@property (nonatomic, strong) NSMutableArray *columnHeights;

/**--attrsArray--懒加载*/
- (NSMutableArray *)attrsArray
{
  if (_attrsArray == nil)
  {
    _attrsArray = [NSMutableArray array];
  }
  return _attrsArray;
}
/**--columnHeights--懒加载*/
- (NSMutableArray *)columnHeights
{
  if (_columnHeights == nil)
  {
    _columnHeights = [NSMutableArray array];
  }
  return _columnHeights;
}

- (void)prepareLayout
{
  [super prepareLayout];

  /**清除之前跟布局相关的所有属性，重新设置新的布局*/
  //清除之前计算的所有列的高度
  [self.columnHeights removeAllObjects];
  //设置所有列的初始高度
  for (NSInteger i = 0; i<self.columnCount; i++)
  {
    self.columnHeights[i] = @(self.edgeInsets.top);
  }
  //清除之前所有的布局属性
  [self.attrsArray removeAllObjects];

  /**开始创建每一个cell对应的布局属性*/
  NSInteger count = [self.collectionView numberOfItemsInSection:0];
  for (NSInteger i = 0; i<count; i++)
  {
    NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:0];
    //获取indexPath位置cell对应的布局属性
    UICollectionViewLayoutAttributes *attrs = [self layoutAttributesForItemAtIndexPath:indexPath];
    //将indexPath位置的cell的布局属性添加到所有cell的布局属性数组中
    [self.attrsArray addObject:attrs];
  }
}

- (nullable NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect
{
  return self.attrsArray;
}
设置每一个cell的布局属性

- (nullable UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(nonnull NSIndexPath *)indexPath
{
  //获取indexPath位置的布局属性
  UICollectionViewLayoutAttributes *attrs = [UICollectionViewLayoutAttributes layoutAttributesForCellWithIndexPath:indexPath];

  /**设置cell布局属性的frame*/

  /***确定cell的尺寸***/
  //获取collectionView的宽度
  CGFloat collectionViewWidth = self.collectionView.frame.size.width;
  //cell宽度
  CGFloat width = ((collectionViewWidth - self.edgeInsets.left - self.edgeInsets.right - (self.columnCount - 1) * self.columMargin)) / self.columnCount;
  //cell高度
  CGFloat height = [self.delegate waterFlowLayout:self heightForItemAtIndex:indexPath.item itemWith:width];

  /***设置cell的位置***/
  NSInteger destColumn = 0;
  CGFloat minColumnHeight = [self.columnHeights[0] doubleValue];
  for (NSInteger i = 1; i<self.columnCount; i++)
  {
    CGFloat columnHeight = [self.columnHeights[i] doubleValue];
    if (minColumnHeight > columnHeight)
    {
      minColumnHeight = columnHeight;
      destColumn = i;
    }
  }
  //计算cell的位置
  CGFloat x = self.edgeInsets.left + destColumn * (width + self.columMargin);
  CGFloat y = minColumnHeight;
  //判断是不是第一行
  if (y != self.edgeInsets.top)
  {
    //若不是第一行，需要加上行间距
    y += self.rowMargin;
  }

  /**给cell的布局属性的frame赋值*/
  attrs.frame = CGRectMake(x, y, width, height);

  //更新最短那列的高度
  self.columnHeights[destColumn] = @(CGRectGetMaxY(attrs.frame));

  /**返回indexPath位置的cell的布局属性*/
  return attrs;
}

- (CGSize)collectionViewContentSize
{
  //获取最高的那一列的高度
  CGFloat maxColumnHeight = [self.columnHeights[0] doubleValue];
  for (NSInteger i = 1; i<self.columnCount; i++)
  {
    CGFloat columnHeight = [self.columnHeights[i] doubleValue];
    if (maxColumnHeight < columnHeight)
    {
      maxColumnHeight = columnHeight;
    }
  }
  //返回collectionView的contentSize，高度为最高的高度加上一个行间距
  return CGSizeMake(0, maxColumnHeight + self.rowMargin);
}