import java.util.Iterator;
import java.lang.Iterable;

class Matrix<E> implements Iterable<E>{
	private Object[][] m = null;
	int rows = 0;
	int cols = 0;
	int nextElmRow = 0;
	int nextElmCol = 0;

	public Matrix(int rows, int cols) {
		// TODO: validate rows and cols

		m = new Object[rows][cols];
		this.rows = rows;
		this.cols = cols;
	}

	public void add(E elm) {
		// TODO: validate nextElmRow and nextElmCol 
		
		m[nextElmRow][nextElmCol] = elm;

		nextElmRow = (nextElmCol == cols - 1) ? nextElmRow + 1: nextElmRow;
		nextElmCol = (nextElmCol == cols - 1) ? 0 : nextElmCol + 1;
	}

	public void add(int row, int col, E elm) {
		// TODO: validate indices

		m[row][col] = elm;
	}

	public int numRows() {
		return rows;
	}

	public int numCols() {
		return cols;
	}

	@SuppressWarnings("unchecked")
	public E peek(int row, int col) {
		// TODO: validate row and col

		return (E) m[row][col];
	}

	public Iterator<E> iterator() {
		return new MatrixIterator();
	}

	public Iterator<E> diagonalIterator() {
		return new DiagonalIterator();
	}

	public Iterator<E> zigzagIterator() {
		return new ZigZagIterator();
	}

	class MatrixIterator implements Iterator<E> {
		private int nextElmRow = 0;
		private int nextElmCol = 0;
		private int count = 0;

		public boolean hasNext() {
			return count < (rows * cols);
		}

		@SuppressWarnings("unchecked")
		public E next() {
	//		System.out.printf("it: %d %d\n", nextElmRow, nextElmCol);
			count++;
			E temp = (E) m[nextElmRow][nextElmCol]; 

			nextElmRow = (nextElmCol == cols - 1) ? nextElmRow + 1: nextElmRow;
			nextElmCol = (nextElmCol == cols - 1) ? 0 : nextElmCol + 1;
	
			return temp;
		}
	}

	class DiagonalIterator implements Iterator<E> {
		int index = 0;
		int count = 0;

		public boolean hasNext() {
			return count < rows;
		}

		@SuppressWarnings("unchecked")
		public E next() {
			E temp = (E) m[index][index];

			index++;
			count++;
			return temp;
		}
	}

	public class ZigZagIterator implements Iterator<E> {
		int count = 0;
		int curRow = 0;
		int curCol = 0;
		Direction dir = Direction.NE;	

		public boolean hasNext() {
			return count < (rows * cols); 
		}

		@SuppressWarnings("unchecked")
		public E next() {
			E elm = (E) m[curRow][curCol];
			updatePosition();
			count++;
			return elm;
		}
		
		public int getCurRow() {
			return curRow;
		}
		
		public int getCurCol() {
			return curCol;
		}
	
		private boolean validPosition(int r, int c) {
			return (r >= 0 && r < rows && c >= 0 && c < cols);
		}

		private void updatePosition() {
			if (dir == Direction.NE) {
				if (validPosition(curRow - 1, curCol + 1)) {
					curRow--; curCol++;
				}
				else if (validPosition(curRow, curCol + 1)) {
					curCol++;
					dir = Direction.SW;
				}
				else {
					curRow++;
					dir = Direction.SW;
				} 
			}
			else if (dir == Direction.SW) {
				if (validPosition(curRow + 1, curCol - 1)) {
					curRow++; curCol--; 
				}
				else if (validPosition(curRow + 1, curCol)) {
					curRow++;
					dir = Direction.NE;
				}
				else {
					curCol++;
					dir = Direction.NE;
				} 
			}
		}			
	}

	enum Direction {
		SW, NE;
	}

}
