#import <UIKit/UIKit.h>
@class YJWaterFlowLayout;

@protocol YJWaterFlowLayoutDelegate <NSObject>
@required
-(CGFloat)waterFlowLayout:(YJWaterFlowLayout *)waterFlowLayout heightForItemAtIndexPath:(NSUInteger)index itemWidth:(CGFloat)itemWidth;

@optional
/** 列数*/
-(CGFloat)columnCountInWaterFlowLayout:(YJWaterFlowLayout *)waterFlowLayout;
/** 列间距*/
-(CGFloat)columnMarginInWaterFlowLayout:(YJWaterFlowLayout *)waterFlowLayout;
/** 行间距*/
-(CGFloat)rowMarginInWaterFlowLayout:(YJWaterFlowLayout *)waterFlowLayout;
/** 边缘之间的间距*/
-(UIEdgeInsets)edgeInsetInWaterFlowLayout:(YJWaterFlowLayout *)waterFlowLayout;

@end

@interface YJWaterFlowLayout : UICollectionViewLayout

/** delegate*/
@property (nonatomic, weak) id<YJWaterFlowLayoutDelegate> delegate;

@end

#import "YJWaterFlowLayout.h"
/** 默认的列数*/
static const NSInteger YJDefaultColumeCount = 3;
/** 每一列之间的间距*/
static const NSInteger YJDefaultColumeMargin = 10;
/** 每一行之间的间距*/
static const CGFloat YJDefaultRowMargin = 10;
/** 边缘之间的间距*/
static const UIEdgeInsets YJDefaultEdgeInset = {10, 10, 10, 10};

@interface YJWaterFlowLayout ()
/** 存放所有cell的布局属性*/
@property (strong, nonatomic) NSMutableArray *attrsArray;
/** 存放每一列的最大y值*/
@property (nonatomic, strong) NSMutableArray *columnHeights;
/** 每一行之间的间距*/
-(CGFloat)rowMargin;
/** 每一列之间的间距*/
-(CGFloat)columnMargin;
/** 列数*/
-(NSInteger)columnCount;
/** 边缘之间的间距*/
-(UIEdgeInsets)edgeInsets;

/** 内容的高度*/
@property (nonatomic, assign) CGFloat maxColumnHeight;

@end

@implementation YJWaterFlowLayout

#pragma mark 代理方法处理
-(CGFloat)rowMargin {
    if ([self.delegate respondsToSelector:@selector(rowMarginInWaterFlowLayout:)]) {
       return [self.delegate rowMarginInWaterFlowLayout:self];
    } else {
        return YJDefaultRowMargin;
    }
}

-(CGFloat)columnMargin {
    if ([self.delegate respondsToSelector:@selector(columnMarginInWaterFlowLayout:)]) {
        return [self.delegate columnMarginInWaterFlowLayout:self];
    } else {
        return YJDefaultColumeMargin;
    }
}

-(NSInteger)columnCount {
    if ([self.delegate respondsToSelector:@selector(columnCountInWaterFlowLayout:)]) {
        return [self.delegate columnCountInWaterFlowLayout:self];
    } else {
        return YJDefaultColumeCount;
    }
}

-(UIEdgeInsets)edgeInsets {
    if ([self.delegate respondsToSelector:@selector(edgeInsetInWaterFlowLayout:)]) {
        return [self.delegate edgeInsetInWaterFlowLayout:self];
    } else {
        return YJDefaultEdgeInset;
    }
}

#pragma mark - 懒加载
- (NSMutableArray *)columnHeights {
    if (!_columnHeights) {
        _columnHeights = [NSMutableArray array];
    }
    return _columnHeights;
}

-(NSMutableArray *)attrsArray {
    if (_attrsArray == nil) {
        _attrsArray = [NSMutableArray array];
    }
    return _attrsArray;
}

/** 初始化*/
-(void)prepareLayout {
    [super prepareLayout];
    
    //清除以前计算的所有高度
    self.maxColumnHeight = 0;
    [self.columnHeights removeAllObjects];
    for (NSInteger i = 0; i < self.columnCount; i++) {
        [self.columnHeights addObject:@(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];
        [self.attrsArray addObject:attrs];
    }
}

/** 决定cell的排布*/
-(NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect {
    return self.attrsArray;
}

/** 返回indexPath位置cell对应的布局属性*/
-(UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath {
    //设置布局属性
    UICollectionViewLayoutAttributes *attrs = [UICollectionViewLayoutAttributes  layoutAttributesForCellWithIndexPath:indexPath];
    //collectionView的宽度
    CGFloat collectionW = self.collectionView.frame.size.width;
    //设置布局属性的frame
    CGFloat w = (collectionW - self.edgeInsets.left - self.edgeInsets.right - (self.columnCount - 1) * self.columnMargin) / self.columnCount;
    CGFloat h = [self.delegate waterFlowLayout:self heightForItemAtIndexPath:indexPath.item itemWidth:w];
    
    //找出高度最短的那一列
    NSInteger destColumn = 0;
    CGFloat minColumnHeight = [self.columnHeights[0] doubleValue];
    for (NSInteger i = 1; i < self.columnCount; i++) {
        //取出第i列
        CGFloat columnHeight = [self.columnHeights[i] doubleValue];
        if (minColumnHeight > columnHeight) {
            minColumnHeight = columnHeight;
            destColumn = i;
        }
    }
    
    CGFloat x = self.edgeInsets.left + destColumn * (w + self.rowMargin);
    CGFloat y = minColumnHeight;
    if (y != self.edgeInsets.top) {
        if ([self.delegate respondsToSelector:@selector(rowMarginInWaterFlowLayout:)]) {
            y += [self.delegate rowMarginInWaterFlowLayout:self];
        } else {
            y += self.rowMargin;
        }
    }
    
    attrs.frame = CGRectMake(x, y, w, h);
    //更新最短那列的高度
    self.columnHeights[destColumn] = @(CGRectGetMaxY(attrs.frame));
    //记录内容的高度
    CGFloat columnHeight = [self.columnHeights[destColumn] doubleValue];
    if (self.maxColumnHeight < columnHeight) {
        self.maxColumnHeight = columnHeight;
    }
    
    return attrs;
}


-(CGSize)collectionViewContentSize {
    return CGSizeMake(0, self.maxColumnHeight + self.edgeInsets.bottom);
}

@end