use core:geometry;

/**
 * Layout for putting components in a grid-like pattern.
 *
 * The grid is automatically sized to fit the constraints dictated by the contained components. Each
 * row and column will have their own size. To fill all available space, specify which rows and
 * columns are allowed to grow using `expandRow` or `expandCol`. A weight may additionally be
 * specified to allow different rows/columns to grow at different ratios with respect to each other.
 * There are two modes that can be applied to the expansion. The default is to assign any extra
 * space (according to the weights) to the expandable rows and columns. This means that they will
 * have different sizes based on the minimum size, but the additional space is allocated fairly. The
 * second mode (enabled using `normalizeRows`/`normalizeCols`) is to base the extra assignment based
 * on the total size of the elements. That way, given equal weights, the layout will keep them at
 * equal size. This means that the minimum size of some components might be enlarged to keep the
 * distribution.
 *
 * Elements may also be positioned at absolute grid coordinates by specifying `row` and `col`,
 * starting at index 0. Additional rows and columns are then created to accomodate all
 * components. Elements may span multiple rows and column, as dictated by `rowspan` and `colspan`.
 *
 * Elements are added row by row, wrapping if a maximum number of columns have been specified using
 * `wrapCols`, or if `nextLine` is called.
 *
 * Elements are spaced according to `border`. There is no space around the outer edge of the layout.
 */
class Grid extends Layout {
	/**
	 * Information for components in grids.
	 */
	class Info extends Layout:Info {
		// Location in the grid.
		Nat row;

		// Location in the grid.
		Nat col;

		// Span number of rows from `row`.
		Nat rowspan;

		// Span number of columns from `col`.
		Nat colspan;

		// Initialize values.
		init(Nat row, Nat col) {
			init() {
				row = row;
				col = col;
				rowspan = 1;
				colspan = 1;
			}
		}
	}

	// Space between elements.
	Size border;

	// Set it.
	void border(Float w, Float h) {
		border = Size(w, h);
	}

	// Initialize.
	init() {
		init() {
			border = defaultBorder;
		}
	}

	// Set number of columns used for automatic wrapping.
	void wrapCols(Nat count) {
		wrapAt = count;
	}

	// Add a new child.
	Info add(Component c) : override {
		if ((wrapAt != 0) & (wrapAt == insertCol))
			nextLine();

		Info info(insertRow, insertCol++);
		children << Child(info, c);

		info;
	}

	// Go to the next line of the grid.
	void nextLine() {
		insertRow++;
		insertCol = 0;
	}

	// Set the weight for an individual column to 1.
	void expandCol(Nat col) {
		expandCol(col, 1);
	}

	// Set the weight for an individual column.
	void expandCol(Nat col, Nat weight) {
		while (colWeights.count <= col)
			colWeights << 0;
		colWeights[col] = weight;
	}

	// Set the weight for an individual row to 1.
	void expandRow(Nat row) {
		expandRow(row, 1);
	}

	// Set the weight for an individual row.
	void expandRow(Nat row, Nat weight) {
		while (rowWeights.count <= row)
			rowWeights << 0;
		rowWeights[row] = weight;
	}

	// Expand the current row.
	void expandRow() {
		expandRow(insertRow, 1);
	}

	// Expand the current column.
	void expandCol() {
		expandCol(insertCol, 1);
	}

	// Ignore differences in minimum size when allocating space for expandable rows.
	void normalizeRows() {
		rowNormalize = true;
	}

	// Ignore differences in minimum size when allocating space for expandable columns.
	void normalizeCols() {
		colNormalize = true;
	}

	// Compute the minimum size.
	Size minSize() {
		update();

		var row = rowMin();
		var col = colMin();

		Size r;
		for (x in row)
			r.h += x;
		for (x in col)
			r.w += x;

		r.w += border.w*(col.count.int.float - 1);
		r.h += border.h*(row.count.int.float - 1);
		r;
	}

	// Perform the layout.
	void layout() {
		update();

		Point offset = pos.p0;
		Size me = pos.size;

		var row = toOffsets(adjust(rowMin(), border.h, me.h, rowWeights), offset.y, border.h);
		var col = toOffsets(adjust(colMin(), border.w, me.w, colWeights), offset.x, border.w);

		for (c in children) {
			layout(c, row, col);
		}
	}

	// Traverse.
	void findAll(fn(Component)->void fn) : override {
		super:findAll(fn);
		for (c in children) {
			c.component.findAll(fn);
		}
	}

private:
	/**
	 * Children in the grid.
	 */
	value Child {
		Info info;
		Component component;
		Size minSize;

		init(Info info, Component c) {
			init() {
				info = info;
				component = c;
			}
		}

		// Get position and information from 'info' based from 'rows'(0) or 'cols'(1).
		Nat pos(Nat dim) {
			if (dim == 0)
				info.row;
			else
				info.col;
		}

		Nat span(Nat dim) {
			if (dim == 0)
				info.rowspan;
			else
				info.colspan;
		}
	}

	// Data about the children.
	Child[] children;

	// Current row and column for next insertion.
	Nat insertRow;
	Nat insertCol;

	// Wrap automatically at.
	Nat wrapAt;

	// Weights for expanding rows/cols.
	Nat[] rowWeights;
	Nat[] colWeights;

	// Use total width/height instead of minimum width for rows/cols.
	Bool rowNormalize;
	Bool colNormalize;

	/**
	 * Volatile state that can be computed by the data in 'children'. Updated by calling 'refresh'.
	 */

	// Number of rows/cols.
	Nat rows;
	Nat cols;

	// Update our internal state from the data available to us.
	void update() {
		rows = rowWeights.count;
		cols = colWeights.count;

		for (i, c in children) {
			var info = c.info;
			rows = max(rows, info.row + info.rowspan);
			cols = max(cols, info.col + info.colspan);

			children[i].minSize = c.component.minSize();
		}

		while (rowWeights.count < rows)
			rowWeights << 0;
		while (colWeights.count < cols)
			colWeights << 0;
	}

	// Get minimum row heights
	Float[] rowMin() {
		Float[] result(rows, 0.0);
		if (rowNormalize) {
			Float minSize = normalizedMinSize(1, rowWeights);
			for (Nat i = 0; i < rows; i++)
				result[i] = minSize * rowWeights[i].float;
		}

		Nat width = 1;
		while (width > 0) {
			width = minStep(width, 0, result, rowWeights);
		}
		result;
	}

	// Get minimum column widths.
	Float[] colMin() {
		Float[] result(cols, 0);
		if (colNormalize) {
			Float minSize = normalizedMinSize(1, colWeights);
			for (Nat i = 0; i < cols; i++)
				result[i] = minSize * colWeights[i].float;
		}

		Nat width = 1;
		while (width > 0) {
			width = minStep(width, 1, result, colWeights);
		}
		result;
	}

	// Pick the proper dimension out of 's'.
	Float pick(Nat dim, Size s) {
		if (dim == 0)
			s.h;
		else
			s.w;
	}

	// Compute one step of the min-width process. Returns the next size to examine, or 0 if
	// none. Returns the next larger size to examine.
	Nat minStep(Nat current, Nat dim, Float[] result, Nat[] weights) {
		Nat next = 0;
		Float space = pick(dim, border);
		for (c in children) {
			Nat span = c.span(dim);

			if (span > current) {
				// Figure out the next step.
				if (next == 0) {
					next = span;
				} else {
					next = min(next, span);
				}
				continue;
			}

			if (span < current)
				continue;

			Float size = pick(dim, c.minSize);
			Nat pos = c.pos(dim);

			// See the current width of the desired columns.
			Float sum = 0;
			for (Nat x = 0; x < span; x++)
				sum += result[x + pos];
			sum += (span - 1).int.float * space;

			// ... and grow them if necessary.
			if (sum < size) {
				// To make it look nicer, we grow according to the weights. Otherwise it will
				// appear that the weights are not respected if two columns contain small things
				// and a large thing spanned over the two columns.
				Nat totalWeight = 0;
				for (Nat x = 0; x < span; x++)
					totalWeight += weights[x + pos];

				if (totalWeight == 0) {
					// Just equalize it. We have no other weights.
					Float grow = (size - sum) / span.int.float;
					for (Nat x = 0; x < span; x++)
						result[x + pos] += grow;
				} else {
					// Respect the weights.
					Float grow = (size - sum) / totalWeight.int.float;
					for (Nat x = 0; x < span; x++)
						result[x + pos] += grow * weights[x + pos].int.float;
				}
			}
		}
		next;
	}

	// Compute normalized minimum size for one unit of 'weight' so we can assign it to the elements later on.
	Float normalizedMinSize(Nat dim, Nat[] weights) {
		Float result = 0;
		for (c in children) {
			Nat pos = c.pos(dim);
			Nat span = c.span(dim);

			Nat weight = 0;
			for (Nat i = 0; i < span; i++)
				weight += weights[i + pos];

			if (weight > 0)
				result = result.max(pick(dim, c.minSize) / weight.float);
		}
		result;
	}


	// Adjust rows/colums to fit the specified width/height.
	Float[] adjust(Float[] elems, Float border, Float total, Nat[] weights) {
		Float min;
		for (x in elems)
			min += x;
		min += border*(elems.count.int.float - 1);

		Nat totalWeight;
		for (Nat i = 0; i < elems.count; i++)
			totalWeight += weights[i];

		Float remaining = total - min;
		Float adjust = 0;
		if (totalWeight > 0)
			adjust = remaining / (totalWeight.int.float);

		Float[] result;
		result.reserve(elems.count);

		for (i, x in elems) {
			result << (x + adjust*weights[i].int.float);
		}

		result;
	}

	// Convert to offsets.
	Float[] toOffsets(Float[] elems, Float offset, Float border) {
		Float[] result;
		result.reserve(elems.count + 1);
		Float sum = offset;
		for (x in elems) {
			result << sum;
			sum += x + border;
		}
		result << sum;
	}

	// Position a single child given the sizes of rows and columns.
	void layout(Child child, Float[] row, Float[] col) {
		var info = child.info;
		Point p0(col[info.col], row[info.row]);
		Point p1(col[info.col + info.colspan], row[info.row + info.rowspan]);
		p1 -= border;

		child.component.pos = Rect(p0, p1);
	}
}