use super::LinOp;
use crate::internal_prelude::*;
use crate::perm;
use linalg::matmul::matmul;
use linalg::temp_mat_zeroed;

#[math]
fn iterate_lanczos<T: ComplexField>(A: &dyn LinOp<T>, H: MatMut<'_, T>, V: MatMut<'_, T>, start: usize, end: usize, par: Par, stack: &mut MemStack) {
	let mut V = V;
	let mut H = H;

	for j in start..end + 1 {
		let mut H = H.rb_mut().col_mut(j - 1);
		H.fill(zero());

		let (V, Vnext) = V.rb_mut().split_at_col_mut(j);
		let V = V.rb();

		let mut Vnext = Vnext.col_mut(0);
		A.apply(Vnext.rb_mut().as_mat_mut(), V.col(j - 1).as_mat(), par, stack);

		let (mut converged, _) = stack.collect(core::iter::repeat_n(false, j));

		let mut h = H.rb_mut().get_mut(..j);

		for i in 0..j {
			let r = V.col(i).adjoint() * Vnext.rb();
			zip!(Vnext.rb_mut(), V.col(i)).for_each(|unzip!(y, x): Zip!(&mut T, &T)| *y = *y - r * *x);
			if i + 1 == j {
				h[i] = r;
			}
		}

		let f = from_f64::<T::Real>(Ord::max(j, 8) as f64) * eps::<T::Real>();

		loop {
			let mut all_true = true;
			for i in 0..j {
				if !converged[i] {
					all_true = false;

					let r = V.col(i).adjoint() * Vnext.rb();
					zip!(Vnext.rb_mut(), V.col(i)).for_each(|unzip!(y, x): Zip!(&mut T, &T)| *y = *y - r * *x);

					if i + 1 == j {
						h[i] = h[i] + r;
					}

					converged[i] = abs(r) < f * Vnext.norm_l2();
				}
			}
			if all_true {
				break;
			}
		}

		let norm = Vnext.norm_l2();
		if norm > zero() {
			let norm_inv = recip(norm);
			zip!(&mut Vnext).for_each(|unzip!(v)| *v = mul_real(*v, norm_inv));
		} else {
			break;
		}

		H[j] = from_real(norm);
	}
}

#[math]
pub fn partial_self_adjoint_eigen_imp<T: ComplexField>(
	eigvecs: MatMut<'_, T>,
	eigvals: &mut [T],

	A: &dyn LinOp<T>,
	v0: ColRef<'_, T>,
	min_dim: usize,
	max_dim: usize,
	n_eigval: usize,
	tol: T::Real,
	restarts: usize,
	par: Par,
	stack: &mut MemStack,
) -> usize {
	let n = A.nrows();

	let (mut H, stack) = temp_mat_zeroed::<T, _, _>(max_dim + 1, max_dim, stack);
	let mut H = H.as_mat_mut();
	let (mut V, stack) = temp_mat_zeroed::<T, _, _>(n, max_dim + 1, stack);
	let mut V = V.as_mat_mut();
	let (mut tmp, stack) = temp_mat_zeroed::<T, _, _>(n, max_dim, stack);
	let mut tmp = tmp.as_mat_mut();

	if max_dim < n {
		let mut active = 0usize;
		let (mut Q, stack) = temp_mat_zeroed::<T, _, _>(max_dim, max_dim, stack);
		let mut Q = Q.as_mat_mut();

		let (mut residual, stack) = temp_mat_zeroed::<T::Real, _, _>(max_dim, 1, stack);
		let mut residual = residual.as_mat_mut().col_mut(0);

		let f = v0.norm_l2();
		if f > min_positive() {
			let f = recip(f);
			zip!(V.rb_mut().col_mut(0), v0).for_each(|unzip!(y, x): Zip!(&mut T, &T)| *y = mul_real(*x, f));
		} else {
			let n0 = n as u32;
			let n1 = (n >> 32) as u32;

			let n = from_f64::<T>(n0 as f64) + from_f64::<T>(n1 as f64);
			let f = recip(sqrt(n));

			zip!(V.rb_mut().col_mut(0)).for_each(|unzip!(y)| *y = copy(f));
		}

		iterate_lanczos(A, H.as_mut(), V.as_mut(), 1, min_dim, par, stack);
		let mut k = min_dim;

		for _ in 0..restarts {
			iterate_lanczos(A, H.as_mut(), V.as_mut(), k + 1, max_dim, par, stack);
			let Hmm = copy(H[(max_dim, max_dim - 1)]);

			let n = max_dim - active;

			Q.fill(zero());
			Q.rb_mut().diagonal_mut().fill(one());

			let mut Q_slice = Q.rb_mut().get_mut(active..max_dim, active..max_dim);
			let mut H_slice = H.rb_mut().get_mut(active..max_dim, active..max_dim);

			{
				let (mut s, stack) = temp_mat_zeroed(n, 1, stack);
				let mut s = s.as_mat_mut().col_mut(0).as_diagonal_mut();
				if linalg::evd::self_adjoint_evd(H_slice.rb(), s.rb_mut(), Some(Q_slice.rb_mut()), par, stack, default()).is_err() {
					break;
				}
				H_slice.fill(zero());
				H_slice.rb_mut().diagonal_mut().copy_from(s.rb());
			}

			for j in 0..n {
				let mut idx = j;
				let mut max = zero::<T::Real>();
				for i in j..n {
					let v = abs(H_slice[(i, i)]);

					if v > max {
						max = v;
						idx = i;
					}
				}

				let i = idx;
				if i != j {
					let tmp = copy(H_slice[(i, i)]);
					H_slice[(i, i)] = copy(H_slice[(j, j)]);
					H_slice[(j, j)] = tmp;
					perm::swap_cols_idx(Q_slice.rb_mut(), i, j);
				}
			}

			let Hmm_abs = abs(Hmm);
			for j in 0..n {
				residual[active + j] = Hmm_abs * abs(Q[(max_dim - 1, active + j)]);
			}

			#[derive(Copy, Clone, PartialEq, Eq, Debug)]
			enum Group {
				Lock,
				Retain,
				Purge,
			}

			let mut groups = alloc::vec![Group::Purge; max_dim];

			let nev = n_eigval;

			let mut nlock = 0usize;
			for j in 0..nev {
				if residual[j] <= tol {
					groups[j] = Group::Lock;
					nlock += 1;
				} else {
					groups[j] = Group::Retain;
				}
			}

			let ideal_size = Ord::min(nlock + min_dim, (min_dim + max_dim) / 2);
			k = nev;

			for i in nev..max_dim {
				let group;
				if k < ideal_size && residual[i] > tol {
					group = Group::Retain;
					k += 1;
				} else {
					group = Group::Purge;
				}

				groups[i] = group;
			}

			let mut purge = 0usize;
			while purge < active && groups[purge] == Group::Lock {
				purge += 1;
			}

			let mut lo = 0usize;
			let mut mi = 0usize;
			let mut hi = 0usize;

			while hi < max_dim {
				match groups[hi] {
					Group::Lock => {
						if hi > lo {
							H[(lo, lo)] = copy(H[(hi, hi)]);
							residual[lo] = copy(residual[hi]);
							let (mut lo, hi) = Q.rb_mut().two_cols_mut(lo, hi);
							lo.copy_from(hi);
						}

						lo += 1;
						mi += 1;
						hi += 1;
					},
					Group::Retain => {
						if hi > mi {
							H[(mi, mi)] = copy(H[(hi, hi)]);
							let (mut mi, hi) = Q.rb_mut().two_cols_mut(mi, hi);
							mi.copy_from(hi);
						}
						mi += 1;
						hi += 1;
					},
					Group::Purge => {
						hi += 1;
					},
				}
			}

			let mut V_tmp = tmp.rb_mut().get_mut(.., purge..k);
			matmul(
				V_tmp.rb_mut(),
				Accum::Replace,
				V.rb().get(.., purge..max_dim),
				Q.rb().get(purge..max_dim, purge..k),
				one(),
				par,
			);
			V.rb_mut().get_mut(.., purge..k).copy_from(&V_tmp);

			let mut b_tmp = tmp.rb_mut().get_mut(0, ..);
			matmul(b_tmp.rb_mut(), Accum::Replace, H.rb().get(max_dim, ..), Q.rb(), one(), par);

			H.rb_mut().get_mut(max_dim, ..).copy_from(b_tmp);

			let (mut x, y) = V.rb_mut().two_cols_mut(k, max_dim);
			x.copy_from(&y);

			let (mut x, mut y) = H.rb_mut().two_rows_mut(k, max_dim);
			x.copy_from(&y);
			y.fill(zero());

			active = nlock;
			if nlock >= n_eigval {
				break;
			}
		}

		let n = active;

		let (mut norms, stack) = stack.make_with(n, |j| abs(H[(j, j)]));
		let (mut perm, stack) = stack.make_with(n, |j| j);
		let _ = stack;

		let perm = &mut *perm;
		let norms = &mut *norms;

		perm.sort_unstable_by(|&i, &j| {
			if norms[i] > norms[j] {
				core::cmp::Ordering::Less
			} else if norms[i] < norms[j] {
				core::cmp::Ordering::Greater
			} else {
				core::cmp::Ordering::Equal
			}
		});

		let mut eigvecs = eigvecs;
		for idx in 0..active {
			let j = perm[idx];
			eigvecs.rb_mut().col_mut(idx).copy_from(V.rb().col(j));
			eigvals[idx] = copy(H[(j, j)]);
		}

		active
	} else {
		panic!();
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::prelude::*;
	use crate::stats::prelude::*;
	use dyn_stack::MemBuffer;
	use equator::assert;

	#[test]
	fn test_small_cplx() {
		let rng = &mut StdRng::seed_from_u64(1);
		let n = 512;
		let n_eigval = 32;

		let mat = CwiseMatDistribution {
			nrows: n,
			ncols: n,
			dist: ComplexDistribution::new(StandardNormal, StandardNormal),
		};
		let col = CwiseColDistribution {
			nrows: n,
			dist: ComplexDistribution::new(StandardNormal, StandardNormal),
		};
		let A: Mat<c64> = mat.sample(rng);
		let A = &A + A.adjoint();

		let mut v0: Col<c64> = col.sample(rng);
		v0 /= v0.norm_l2();
		let A = A.as_ref();
		let v0 = v0.as_ref();

		let par = Par::Seq;
		let mem = &mut MemBuffer::new(crate::matrix_free::eigen::partial_eigen_scratch(&A, n_eigval, par, default()));
		let mut V = Mat::zeros(n, n_eigval);
		let mut w = vec![c64::ZERO; n_eigval];

		let n_converged = partial_self_adjoint_eigen_imp(
			V.rb_mut(),
			&mut w,
			&A,
			v0,
			32,
			64,
			n_eigval,
			f64::EPSILON * 128.0,
			1000,
			par,
			MemStack::new(mem),
		);
		assert!(w.iter().map(|x| x.norm()).is_sorted_by(|x, y| x >= y));
		assert!(n_converged == n_eigval);

		for j in 0..n_converged {
			assert!((A * V.col(j) - Scale(w[j]) * V.col(j)).norm_l2() < 1e-10);
		}
	}
}
