#import <UIKit/UIKit.h>
// 類的前置聲明
@class WaterfallFlowLayout;

@protocol WaterfallFlowLayoutDelegate <NSObject>
// 動態(tài)獲取 item 寬度
- (CGFloat) WaterfallFlowLayout:(WaterfallFlowLayout *) layout widthForItemAtIndexPath:(NSIndexPath *) indexPath;

@end

@interface WaterfallFlowLayout : UICollectionViewLayout

@property (nonatomic,assign) id <WaterfallFlowLayoutDelegate> delegate;

@property (nonatomic) NSInteger numberOfColumns;
@property (nonatomic) CGFloat minimumLineSpacing;
@property (nonatomic) CGFloat minimumInteritemSpacing;
@property (nonatomic) UIEdgeInsets sectionInset;
@end

#import "WaterfallFlowLayout.h"

@interface WaterfallFlowLayout ()
{
  // 用于記錄每一列布局到的寬度
  NSMutableArray * _widthOfColumns;
  // 用于保存所有item的屬性 (frame)
  NSMutableArray * _itemsAttributes;
}

@end


@implementation WaterfallFlowLayout
- (void) setNumberOfColumns:(NSInteger)numberOfColumns {
  if (_numberOfColumns != numberOfColumns) {
    _numberOfColumns = numberOfColumns;
    // 讓原有布局失效，需要重新布局
    [self invalidateLayout];
  }
}

- (void)setMinimumLineSpacing:(CGFloat)minimumLineSpacing {
  if (_minimumLineSpacing != minimumLineSpacing) {
    _minimumLineSpacing = minimumLineSpacing;
    [self invalidateLayout];
  }
}
- (void)setMinimumInteritemSpacing:(CGFloat)minimumInteritemSpacing {
  if (_minimumInteritemSpacing != minimumInteritemSpacing) {
    _minimumInteritemSpacing = minimumInteritemSpacing;
    [self invalidateLayout];
  }
}

- (void)setSectionInset:(UIEdgeInsets)sectionInset {
  if (!UIEdgeInsetsEqualToEdgeInsets(_sectionInset, sectionInset)) {
    _sectionInset = sectionInset;
    [self invalidateLayout];
  }
}

//重寫方法 1: 準(zhǔn)備布局
-(void)prepareLayout {
  [super prepareLayout];
  // 真正的布局在這里完成
  if (_itemsAttributes) {
    [_itemsAttributes removeAllObjects];
  }else {
    _itemsAttributes = [[NSMutableArray alloc] init];
  }
  if (_widthOfColumns) {
    [_widthOfColumns removeAllObjects];
  }else {
    _widthOfColumns = [[NSMutableArray alloc] init];
  }
  for (NSInteger i = 0; i < self.numberOfColumns; i++) {
    // 初始化每一列的寬度(默認為上邊距)
//    _heightOfColumns[i] = @(self.sectionInset.top);
    [_widthOfColumns addObject:@(self.sectionInset.left)];
  }
  // item的總數(shù)
  NSInteger count = [self.collectionView numberOfItemsInSection:0];

//  CGFloat itemWidth = (self.collectionView.frame.size.width - self.sectionInset.left - self.sectionInset.right - (_numberOfColumns-1) * _minimumInteritemSpacing )/_numberOfColumns;

  // 總的高度 (集合視圖的寬度)
  CGFloat totalHeight = self.collectionView.frame.size.height;
  // 有效的高度 (出去間隔及邊界)
  CGFloat validHeight = totalHeight - self.sectionInset.top - self.self.sectionInset.bottom - (self.numberOfColumns-1) * self.minimumInteritemSpacing;
  // 每一個item的高度
  CGFloat itemHeight = validHeight/self.numberOfColumns;


  // 設(shè)置item的默認寬度
  CGFloat itemWidth = itemHeight;
  for (NSInteger i = 0; i<count; i++) {
    // 最短列的下標(biāo)
    NSInteger index = [self indexOfShortestColumn];
    CGFloat originY = self.sectionInset.top + index * (itemHeight +self.minimumInteritemSpacing);
    CGFloat originX = [_widthOfColumns[index] floatValue];
    // 構(gòu)造 indexPath
    NSIndexPath * indexPath = [NSIndexPath indexPathForItem:i inSection:0];
    // 動態(tài)的獲取寬度
    if ([self.delegate respondsToSelector:@selector(WaterfallFlowLayout:widthForItemAtIndexPath:)]) {
      itemWidth = [self.delegate WaterfallFlowLayout:self widthForItemAtIndexPath:indexPath];
    }
    UICollectionViewLayoutAttributes * attr = [UICollectionViewLayoutAttributes layoutAttributesForCellWithIndexPath:indexPath];

    attr.frame = CGRectMake(originX, originY, itemWidth, itemHeight);
    // 保存 item 的屬性 到數(shù)組中
    [_itemsAttributes addObject:attr];
    // 更新布局到的一列(最短列) 的高度
    _widthOfColumns[index] = @(originX + itemWidth + self.minimumLineSpacing);
  }
  // 刷新顯示
  [self.collectionView reloadData];
}


//重寫方法 2: 返回指定區(qū)域的item的屬性(frame)
- (NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect {
  NSMutableArray * array = [NSMutableArray array];

  for (UICollectionViewLayoutAttributes * attr in _itemsAttributes) {
    // 判斷兩個矩形是否有交集
    if (CGRectIntersectsRect(attr.frame, rect)) {
      [array addObject:attr];
    }
  }
  return array;
}


//重寫方法 3: 返回內(nèi)容的尺寸
-(CGSize)collectionViewContentSize {
  CGFloat height = self.collectionView.frame.size.height;
  NSInteger index = [self indexOfLongestColumn];
  CGFloat width = [_widthOfColumns[index] floatValue] + self.sectionInset.right - self.minimumLineSpacing;
  return CGSizeMake(width, height);
}


- (NSInteger) indexOfLongestColumn {
  NSInteger index = 0;
  for (NSInteger i = 0; i<_numberOfColumns; i++) {
    if ([_widthOfColumns[i] floatValue] > [_widthOfColumns[index] floatValue]) {
      index = i;
    }
  }

  return index;
}

- (NSInteger) indexOfShortestColumn {
  NSInteger index = 0;
  for (NSInteger i = 0; i<_numberOfColumns; i++) {
    if ([_widthOfColumns[i] floatValue] < [_widthOfColumns[index] floatValue]) {
      index = i;
    }
  }

  return index;
}
@end

#import "RootViewController.h"
#import "WaterfallFlowLayout.h"

@interface RootViewController () <UICollectionViewDataSource,UICollectionViewDelegateFlowLayout,WaterfallFlowLayoutDelegate>
{
  UICollectionView * _collectionView;
}
@end

@implementation RootViewController
- (void)dealloc {
  [_collectionView release];
  [super dealloc];

}

- (void)viewDidLoad {
  [super viewDidLoad];
  // 創(chuàng)建集合視圖
  [self createCollectionView];
}

- (UICollectionViewLayout *)createLayout {
#if 1
  WaterfallFlowLayout * layout = [[WaterfallFlowLayout alloc] init];
  layout.sectionInset = UIEdgeInsetsMake(20, 20, 20, 20);
  layout.minimumLineSpacing = 10;
  layout.minimumInteritemSpacing = 20;
  layout.numberOfColumns = 3;
  layout.delegate = self;
  [self performSelector:@selector(changeLayout:) withObject:layout afterDelay:3];

#else
  UICollectionViewFlowLayout * layout = [[UICollectionViewFlowLayout alloc] init];

  layout.minimumLineSpacing = 10;
  layout.itemSize = CGSizeMake(150, 100);
  layout.sectionInset = UIEdgeInsetsMake(10, 10, 10, 10);

#endif


  return [layout autorelease];
}
- (void)changeLayout:(WaterfallFlowLayout *)layout {
  layout.numberOfColumns = 3;
}


- (void)createCollectionView {
  CGRect frame = CGRectMake(0, 20, VIEW_WIDTH, VIEW_HEIGHT-20);
  _collectionView = [[UICollectionView alloc] initWithFrame:frame collectionViewLayout:[self createLayout]];

  _collectionView.backgroundColor = [UIColor cyanColor];
  // 設(shè)置代理
  _collectionView.delegate = self;
  _collectionView.dataSource = self;

  // 注冊cell 類型 及 復(fù)用標(biāo)識
  [_collectionView registerClass:[UICollectionViewCell class] forCellWithReuseIdentifier:@"cellId"];


  [self.view addSubview:_collectionView];
}

#pragma mark - UICollectionViewDataSource

- (NSInteger)collectionView:(UICollectionView *)collectionView numberOfItemsInSection:(NSInteger)section {
  return 102;
}

- (UICollectionViewCell *)collectionView:(UICollectionView *)collectionView cellForItemAtIndexPath:(NSIndexPath *)indexPath {
  UICollectionViewCell * cell = [collectionView dequeueReusableCellWithReuseIdentifier:@"cellId" forIndexPath:indexPath];

  UILabel * label = nil;
  NSArray * array = cell.contentView.subviews;
  if (array.count) {
    label = array[0];
  }else {
    label = [[UILabel alloc] init];
//    label.frame = cell.bounds;
    label.textAlignment = NSTextAlignmentCenter;
    label.font = [UIFont systemFontOfSize:50];
    [cell.contentView addSubview:label];
    [label release];
  }
  label.frame = cell.bounds;
  label.text = [NSString stringWithFormat:@"%ld",indexPath.item];
  label.textColor = [UIColor whiteColor];
  cell.backgroundColor = RandomColor;

  return cell;
}

- (CGSize) collectionView:(UICollectionView *)collectionView layout:(UICollectionViewLayout *)collectionViewLayout sizeForItemAtIndexPath:(NSIndexPath *)indexPath {
  return CGSizeMake( arc4random()%100+200, 110);
}

-(CGFloat) WaterfallFlowLayout:(WaterfallFlowLayout *)layout widthForItemAtIndexPath:(NSIndexPath *)indexPath{
  return arc4random()%150+50;
}

-(void)collectionView:(UICollectionView *)collectionView didSelectItemAtIndexPath:(NSIndexPath *)indexPath{
  NSLog(@"點擊了第 %ld 組，第 %ld 行",indexPath.section,indexPath.row);
}

- (void)didReceiveMemoryWarning {
  [super didReceiveMemoryWarning];
}


@end