use wasm_bindgen::prelude::wasm_bindgen;

#[inline]
pub fn get_idx(w: usize, x: usize, y: usize) -> usize {
	y * w + x
}

#[inline]
pub fn get_x(w: usize, idx: usize) -> usize {
	idx % w
}

#[inline]
pub fn get_y(w: usize, idx: usize) -> usize {
	idx / w
}

#[inline]
pub fn get_diagonals(w: usize, board: &[u8], reverse: bool) -> Vec<Vec<u8>> {
	let h = get_y(w, board.len());
	(0..w + h - 1).into_iter()
		.map(|p| board.iter()
			.enumerate()
			.filter_map(|(idx, &v)| {
				let x = get_x(w, idx);
				let y = get_y(w, idx);
				if (x + y == p && !reverse) || (w - x - 1 + y == p && reverse) {
					Some(v)
				} else { None }
			}).collect()
		)
		.collect()
}

#[inline]
pub fn get_rows(w: usize, board: &[u8]) -> Vec<Vec<u8>> {
	board.chunks(w)
		.map(|c| c.to_vec())
		.collect()
}

#[inline]
pub fn get_columns(w: usize, board: &[u8]) -> Vec<Vec<u8>> {
	(0..w).into_iter()
		.map(|target_x| board.iter()
			.enumerate()
			.filter_map(|(idx, &v)| {
				let x = get_x(w, idx);
				if x == target_x {
					Some(v)
				} else { None }
			}).collect()
		)
		.collect()
}

#[inline]
pub fn test_chunks(chunks: &[Vec<u8>], min: usize) -> Option<u8> {
	chunks.iter()
		.find_map(|chunk| chunk.windows(min)
			.find_map(|window| {
				if window.iter()
					.all(|&n| n == window[0] && n != 0) {
					Some(window[0])
				} else { None }
			})
		)
}

#[wasm_bindgen]
pub fn find_winner(w: usize, board: &[u8], min: usize) -> u8 {
	let rows_test = test_chunks(&get_rows(w, board), min);
	if let Some(winner) = rows_test {
		return winner;
	}
	let cols_test = test_chunks(&get_columns(w, board), min);
	if let Some(winner) = cols_test {
		return winner;
	}
	let dia1_test = test_chunks(&get_diagonals(w, board, false), min);
	if let Some(winner) = dia1_test {
		return winner;
	}
	let dia2_test = test_chunks(&get_diagonals(w, board, true), min);
	if let Some(winner) = dia2_test {
		return winner;
	}
	0
}

#[wasm_bindgen]
pub fn count_empty(board: &[u8]) -> usize {
	board.iter().filter(|&n| *n == 0).count()
}

#[wasm_bindgen]
pub fn get_turn(me: u8, other: u8, first: bool, empty: usize) -> u8 {
	if empty % 2 == if first { 0 } else { 1 } {
		me
	} else {
		other
	}
}

#[wasm_bindgen]
pub fn mm_get_score(me: u8, other: u8, first: bool, w: usize, board: &[u8], min: usize, md: f64) -> f64 {
	mm_sub_get_score(me, other, first, w, &board.to_vec(), min, board.len() as f64 + 1.0, 0.0, 0.0, md)
}

#[wasm_bindgen]
pub fn mm_d_get_score(me: u8, other: u8, first: bool, w: usize, board: &[u8], min: usize, md: f64) -> f64 {
	mm_sub_get_score(me, other, first, w, &board.to_vec(), min, board.len() as f64 + 1.0, 0.0, 1.0, md)
}

// m = score (without depth) multiplier
// dm = depth multiplier
pub fn mm_sub_get_score(me: u8, other: u8, first: bool, w: usize, board: &Vec<u8>, min: usize, m: f64, depth: f64, dm: f64, md: f64) -> f64 {
	let winner = find_winner(w, board, min);
	if winner != 0 {
		return if winner == me { m - depth*dm } else { -m + depth*dm };
	}
	let empty = count_empty(board);
	if empty == 0 || depth > md {
		return 0.0;
	}
	let turn = get_turn(me, other, first, empty);
	let mut best = if turn == me { f64::MIN } else { f64::MAX };
	for (x, col) in get_columns(w, board).into_iter().enumerate() {
		let r = col.into_iter()
			.enumerate()
			.filter_map(|(idx, v)| {
				if v == 0 {
					Some(idx)
				} else { None }
			})
			.rev()
			.next();
		if let Some(y) = r {
			let copy = &mut board.clone();
			copy[get_idx(w, x, y)] = get_turn(me, other, first, empty);
			let value = mm_sub_get_score(me, other, first, w, copy, min, m, depth + 1.0, dm, md);
			best = if turn == me { best.max(value) } else { best.min(value) };
		}
	}
	best
}

// Algorithms:
// mm: minimax algorithm
// mm+d: minmax algorithm with depth
#[wasm_bindgen]
pub fn get_score(me: u8, other: u8, first: bool, w: usize, board: &[u8], min: usize, md: f64, algorithm: &str) -> f64 {
	match algorithm {
		"mm" => mm_get_score(me, other, first, w, board, min, md),
		"mm+d" => mm_d_get_score(me, other, first, w, board, min, md),
		_ => 0.0
	}
}

// Algorithms:
// mm: minimax algorithm
// mm+d: minmax algorithm with depth
#[wasm_bindgen]
pub fn predict(me: u8, other: u8, first: bool, w: usize, board: &[u8], min: usize, md: f64, algorithm: &str) -> usize {
	let (mut max_p, mut max_s) = (0, f64::MIN);
	let empty = count_empty(board);
	if empty == 0 {
		return 0;
	}
	let vec_board = board.to_vec();
	for (x, col) in get_columns(w, board).into_iter().enumerate() {
		let r = col.into_iter()
			.enumerate()
			.filter_map(|(idx, v)| {
				if v == 0 {
					Some(idx)
				} else { None }
			})
			.rev()
			.next();
		if let Some(y) = r {
			let copy = &mut vec_board.clone();
			copy[get_idx(w, x, y)] = get_turn(me, other, first, empty);
			let score = get_score(me, other, first, w, copy, min, md, algorithm);
			if score > max_s {
				(max_p, max_s) = (x, score);
			}
		}
	}

	max_p
}