using UnityEngine;
using UnityEngine.UI;
using System.Collections.Generic;

[AddComponentMenu("Layout/Wrap Layout Group")]
public class WrapLayoutGroup : LayoutGroup
{
    public float spacingX = 10;
    public float spacingY = 10;

    public override void CalculateLayoutInputHorizontal()
    {
        base.CalculateLayoutInputHorizontal();
        CalcAlongAxis(0);
    }

    public override void CalculateLayoutInputVertical()
    {
        CalcAlongAxis(1);
    }

    public override void SetLayoutHorizontal()
    {
        SetChildrenLayout(0);
    }

    public override void SetLayoutVertical()
    {
        SetChildrenLayout(1);
    }

    void CalcAlongAxis(int axis)
    {
        float width = rectTransform.rect.width;
        float curX = padding.left;
        float curY = padding.top;
        float maxHeightInRow = 0;

        for (int i = 0; i < rectChildren.Count; i++)
        {
            RectTransform child = rectChildren[i];
            float childWidth = LayoutUtility.GetPreferredSize(child, 0);
            float childHeight = LayoutUtility.GetPreferredSize(child, 1);

            if (curX + childWidth + padding.right > width)
            {
                curX = padding.left;
                curY += maxHeightInRow + spacingY;
                maxHeightInRow = 0;
            }

            curX += childWidth + spacingX;
            maxHeightInRow = Mathf.Max(maxHeightInRow, childHeight);
        }

        float totalHeight = curY + maxHeightInRow + padding.bottom;
        SetLayoutInputForAxis(totalHeight, totalHeight, -1, axis);
    }

    void SetChildrenLayout(int axis)
    {
        float width = rectTransform.rect.width;
        float curX = padding.left;
        float curY = padding.top;
        float maxHeightInRow = 0;

        for (int i = 0; i < rectChildren.Count; i++)
        {
            RectTransform child = rectChildren[i];
            float childWidth = LayoutUtility.GetPreferredSize(child, 0);
            float childHeight = LayoutUtility.GetPreferredSize(child, 1);

            if (curX + childWidth + padding.right > width)
            {
                curX = padding.left;
                curY += maxHeightInRow + spacingY;
                maxHeightInRow = 0;
            }

            SetChildAlongAxis(child, 0, curX, childWidth);
            SetChildAlongAxis(child, 1, curY, childHeight);

            curX += childWidth + spacingX;
            maxHeightInRow = Mathf.Max(maxHeightInRow, childHeight);
        }
    }
}