import { assert } from "../../assert.js";
import { complex_horner_eval } from "../complex/complex_horner_eval.js";
import { polynomial_root_bound_cauchy } from "./polynomial_root_bound_cauchy.js";

/*
 Find all roots of a univariate polynomial in the complex plane via
 Aberth–Ehrlich simultaneous iteration. Real and complex roots are isolated
 together; the caller separates real ones by inspecting the imaginary parts.

 Reference: Aberth (1973), "Iteration methods for finding all zeros of a
 polynomial simultaneously", Math. Comp. 27.

 The iteration is

   z_k ← z_k − w_k / (1 − w_k · Σ_{j≠k} 1/(z_k − z_j))

 with w_k = p(z_k) / p'(z_k). Initial points are placed on a circle of radius
 = Cauchy bound, with a fixed phase offset so an initial guess never lands
 exactly on a real root (which would make the pairwise sum singular at step
 zero). Convergence is quadratic once close.
*/

const MAX_SUPPORTED_DEGREE = 64;

const ABERTH_MAX_ITERATIONS = 80;
const ABERTH_CONVERGENCE_TOLERANCE_SQR = 1e-28;

const _DERIV = new Float64Array(MAX_SUPPORTED_DEGREE);
const scratch_v2 = new Float64Array(2);

/**
 * Find all complex roots of a univariate polynomial via Aberth–Ehrlich.
 *
 * The polynomial is given in ascending-power coefficient form:
 *   p(x) = coeffs[0] + coeffs[1]·x + … + coeffs[degree]·x^degree
 * The leading coefficient `coeffs[degree]` must be nonzero.
 *
 * Exactly `degree` complex numbers are written into `roots_re`/`roots_im` at
 * index `offset`, `offset + 1`, …, `offset + degree - 1`. The order they
 * appear in is not specified.
 *
 * @param {Float64Array|number[]} coeffs ascending-power, length ≥ degree + 1
 * @param {number} degree polynomial degree, ≥ 1, ≤ 64
 * @param {Float64Array|number[]} roots_re real parts, length ≥ offset + degree
 * @param {Float64Array|number[]} roots_im imaginary parts, length ≥ offset + degree
 * @param {number} offset
 * @param {number} [max_iterations=80] outer iteration cap. Lower values
 *   trade worst-case convergence accuracy for speed; the per-iteration
 *   convergence check still fires early in the common case. A low cap
 *   (~24-32) is appropriate for well-conditioned low-degree polynomials
 *   where Aberth typically converges in 15-25 iterations.
 */
export function polynomial_complex_roots_aberth_ehrlich(
    coeffs, degree,
    roots_re, roots_im, offset,
    max_iterations = ABERTH_MAX_ITERATIONS
) {
    assert.greaterThanOrEqual(degree, 1, 'degree');
    assert.greaterThanOrEqual(MAX_SUPPORTED_DEGREE, degree, 'degree exceeds MAX_SUPPORTED_DEGREE');
    assert.isNonNegativeInteger(offset, 'offset');

    // derivative: _DERIV[i] = (i + 1) * coeffs[i + 1], i = 0 .. degree - 1
    for (let i = 0; i < degree; i++) {
        _DERIV[i] = (i + 1) * coeffs[i + 1];
    }

    const radius = polynomial_root_bound_cauchy(coeffs, degree);

    // Initial points around a circle, with a small phase offset to avoid
    // landing exactly on a real root at iteration zero.
    const phase_offset = 0.5;
    const two_pi_over_n = (2 * Math.PI) / degree;

    for (let k = 0; k < degree; k++) {
        const angle = k * two_pi_over_n + phase_offset;
        roots_re[offset + k] = radius * Math.cos(angle);
        roots_im[offset + k] = radius * Math.sin(angle);
    }

    for (let iter = 0; iter < max_iterations; iter++) {
        let max_correction_sq = 0;

        for (let k = 0; k < degree; k++) {
            const zr = roots_re[offset + k];
            const zi = roots_im[offset + k];

            complex_horner_eval(scratch_v2, coeffs, degree, zr, zi);
            const pr = scratch_v2[0];
            const pi = scratch_v2[1];

            complex_horner_eval(scratch_v2, _DERIV, degree - 1, zr, zi);
            const dpr = scratch_v2[0];
            const dpi = scratch_v2[1];

            // w = p / p'
            const dp_norm_sq = dpr * dpr + dpi * dpi;

            if (dp_norm_sq === 0) {
                // division by 0 guard
                continue;
            }

            const inv_dp_sq = 1 / dp_norm_sq;
            const wr = (pr * dpr + pi * dpi) * inv_dp_sq;
            const wi = (pi * dpr - pr * dpi) * inv_dp_sq;

            // S = Σ_{j≠k} 1/(z_k − z_j) in complex
            let sr = 0, si = 0;
            for (let j = 0; j < degree; j++) {

                if (j === k) {
                    continue;
                }

                const dr = zr - roots_re[offset + j];
                const di = zi - roots_im[offset + j];
                const denom_sq = dr * dr + di * di;

                if (denom_sq === 0) {
                    continue;
                }
                const inv_d_sq = 1 / denom_sq;
                sr += dr * inv_d_sq;
                si += -di * inv_d_sq;
            }

            // correction = w / (1 − w·S)
            const ws_r = wr * sr - wi * si;
            const ws_i = wr * si + wi * sr;

            const denom_r = 1 - ws_r;
            const denom_i = -ws_i;

            const denom_sq2 = denom_r * denom_r + denom_i * denom_i;

            if (denom_sq2 === 0){
                // division by 0 guard
                continue;
            }

            const inv_denom_sq2 = 1 / denom_sq2;
            const corr_r = (wr * denom_r + wi * denom_i) * inv_denom_sq2;
            const corr_i = (wi * denom_r - wr * denom_i) * inv_denom_sq2;

            roots_re[offset + k] = zr - corr_r;
            roots_im[offset + k] = zi - corr_i;

            const corr_sq = corr_r * corr_r + corr_i * corr_i;

            if (corr_sq > max_correction_sq){
                max_correction_sq = corr_sq;
            }
        }

        if (max_correction_sq < ABERTH_CONVERGENCE_TOLERANCE_SQR) {
            break;
        }
    }
}
