import { assert } from "../../../assert.js";
import { v2_length } from "../../../geom/vec2/v2_length.js";
import { complex_add } from "../../complex/complex_add.js";
import { complex_div } from "../../complex/complex_div.js";
import { complex_mul } from "../../complex/complex_mul.js";
import { complex_sub } from "../../complex/complex_sub.js";

export function vec2(x, y) {
    return new Float64Array([x, y]);
}

function vec4(x, y, z, w) {
    return new Float64Array([x, y, z, w]);
}

export function cadd(a, b) {
    const r = vec2(0, 0);

    complex_add(r, a, b);

    return r;
}

function csub(a, b) {
    const r = vec2(0, 0);


    complex_sub(r, a, b);

    return r;
}

export function cmult(a, b) {
    const r = vec2(0, 0);


    complex_mul(r, a, b);

    return r;
}

export function cdiv(a, b) {
    const r = vec2(0, 0);


    complex_div(r, a, b);

    return r;
}

// Number of terms to include in infinite sums (truncation)
function terms_to_sum(z) {
    const size = v2_length(z[0], z[1]);
    return Math.ceil(size + 4.3 * Math.cbrt(size) + 1.0);
}

// Downward recurrence for A
function A_all_n(z, M) {
    const A = new Array(M + 2);
    A[M + 1] = vec2(0.0, 0.0);
    for (let n = M; n >= 0; --n) {
        const tmp = cdiv(vec2(n + 1.0, 0.0), z);
        A[n] = csub(tmp, cdiv(vec2(1.0, 0.0), cadd(tmp, A[n + 1])));
    }
    return A;
}

// Upward recurrences for B, psi_zeta, and R
function psi_zeta(n, z, A, old_B, old_psi_zeta) {
    if (n > 0) {
        const n_z = cdiv(vec2(n, 0.0), z);
        const tmp = cmult(csub(n_z, A[n - 1]), csub(n_z, old_B));
        old_psi_zeta = cmult(old_psi_zeta, tmp);
    }
    return old_psi_zeta;
}

/**
 * Calculates the B_n(z) coefficient for term 'n' via upward recurrence.
 * This is a helper function for LorenzMie_ab.
 *
 * @param {number} n - The current recurrence term index (n >= 0).
 * @param {vec2} z - The complex size parameter (x or y).
 * @param {Array<vec2>} A - Pre-computed log-derivatives A_n(z).
 * @param {vec2} old_B - The previous B coefficient (B_{n-1}).
 * @param {vec2} old_psi_zeta - The previous psi_zeta product term.
 * @returns {vec2} The new B_n coefficient (complex number).
 * @ignore
 */
function B(n, z, A, old_B, old_psi_zeta) {

    if (n > 0) {
        old_B.psi_zeta = psi_zeta(n, z, A, old_B, old_psi_zeta);
        old_B = cadd(A[n], cdiv(vec2(0.0, 1.0), old_B.psi_zeta));
    }

    return old_B;
}

/**
 * Calculates the R_n(z) coefficient (ratio psi_n / zeta_n) for term 'n'.
 * This is a helper function for LorenzMie_ab.
 *
 * @param {number} n - The current recurrence term index (n >= 0).
 * @param {vec2} z - The complex size parameter (x).
 * @param {Array<vec2>} A - Pre-computed log-derivatives A_n(z).
 * @param {vec2} B_n - The current B_n(z) coefficient.
 * @param {vec2} old_R - The previous R coefficient (R_{n-1}).
 * @returns {vec2} The new R_n coefficient (complex number).
 * @ignore
 */
function R(n, z, A, B_n, old_R) {
    if (n > 0) {
        const n_z = cdiv(vec2(n, 0.0), z);
        const tmp = cdiv(cadd(B_n, n_z), cadd(A[n], n_z));
        old_R = cmult(old_R, tmp);
    }
    return old_R;
}

/**
 * Calculates the Lorenz-Mie scattering coefficients a_n and b_n for a single term 'n'.
 * This implements the formulas from Figure 3 of Frisvad et al. 2007.
 *
 * @param {number} n - The current recurrence term index (n >= 0).
 * @param {number} size - Unitless size parameter x_vac = 2*pi*r / lambda_vac.
 * @param {vec2} n_p - Complex refractive index of the particle.
 * @param {vec2} n_med - Complex refractive index of the medium.
 * @param {Array<vec2>} A_p - Pre-computed log-derivatives A_n(y) for the particle.
 * @param {Array<vec2>} A_med - Pre-computed log-derivatives A_n(x) for the medium.
 * @param {vec2} old_B - The previous B_{n-1}(x) coefficient.
 * @param {vec2} old_R - The previous R_{n-1}(x) coefficient.
 * @param {vec2} old_psi_zeta - The previous psi_zeta product.
 * @returns {vec4} A vec4 [Re(a_n), Im(a_n), Re(b_n), Im(b_n)] with state properties (.old_B, .old_R, .old_psi_zeta) attached for the next iteration.
 * @ignore
 */
function LorenzMie_ab(n, size, n_p, n_med, A_p, A_med, old_B, old_R, old_psi_zeta) {

    const x = vec2(size * n_med[0], size * n_med[1]);
    const B_n = B(n, x, A_med, old_B, old_psi_zeta);
    const R_n = R(n, x, A_med, B_n, old_R);
    const n_med_A_p = cmult(n_med, A_p[n]);
    const n_p_A_med = cmult(n_p, A_med[n]);
    const n_p_A_p = cmult(n_p, A_p[n]);
    const n_med_A_med = cmult(n_med, A_med[n]);
    const n_p_B_n = cmult(n_p, B_n);
    const n_med_B_n = cmult(n_med, B_n);
    const a = cmult(R_n, cdiv(csub(n_med_A_p, n_p_A_med), csub(n_med_A_p, n_p_B_n)));
    const b = cmult(R_n, cdiv(csub(n_p_A_p, n_med_A_med), csub(n_p_A_p, n_med_B_n)));


    const ab = vec4(a[0], a[1], b[0], b[1]);

    ab.old_psi_zeta = n > 0 ? old_B.psi_zeta : old_psi_zeta;
    ab.old_B = B_n;
    ab.old_R = R_n;

    return ab;
}

/**
 * Implementation of Lorenz-Mie theory.
 * Useful for calculating physically correct radiation scattering events for particles in a medium.
 *
 * @param {number} radius_m - Particle radius (in meters).
 * @param {number} wavelength_m - Wavelength of light in vacuum (in meters).
 * @param {vec2} n_p - Complex refractive index of the particle
 * @param {vec2} n_med - Complex refractive index of the medium (e.g., [1.0, 0.0] for air).
 * @returns {Float64Array} Mie coefficients a_n and b_n in form [a_n0.r, a_n0.i, b_n0.r, b_n0.i, ... a_nN.r, a_nN.i, b_nN.r, b_nN.i]
 *
 * @see "Computing the Scattering Properties of Participating Media Using Lorenz-Mie Theory", SIGGRAPH 2007 by Frisvad et al.
 * @see https://en.wikipedia.org/wiki/Mie_scattering
 * @author Alexander Goldring (c) 2025
 */
export function lorenz_mie_coefs(
    wavelength_m,
    radius_m,
    n_p,
    n_med
) {

    assert.isNumber(wavelength_m, 'wavelength_m');
    assert.isFinite(wavelength_m, 'wavelength_m');
    assert.greaterThan(wavelength_m, 0, 'wavelength_m');

    assert.isNumber(radius_m, 'radius_m');
    assert.isFinite(radius_m, 'radius_m');
    assert.greaterThan(radius_m, 0, 'radius_m');

    const psize = 2.0 * Math.PI * radius_m / wavelength_m;

    const x = vec2(psize * n_med[0], psize * n_med[1]);
    const y = vec2(psize * n_p[0], psize * n_p[1]);
    const M = terms_to_sum(x);
    const A_med = A_all_n(x, M);
    const A_p = A_all_n(y, M);

    const e_term = Math.exp(2.0 * x[1]);

    let prev_psi_zeta = vec2(0.5 * (1.0 - Math.cos(2.0 * x[0]) / e_term), -0.5 * Math.sin(2.0 * x[0]) / e_term);
    let prev_B = vec2(0.0, 1.0);
    let prev_R = vec2(0.5 * (1.0 - Math.cos(2.0 * x[0]) * e_term), 0.5 * Math.sin(2.0 * x[0]) * e_term);

    const ab = new Float64Array(M * 4);

    for (let n = 0; n < M; ++n) {

        const ab_n = LorenzMie_ab(
            n, psize, n_p, n_med, A_p, A_med,
            prev_B, prev_R, prev_psi_zeta
        );

        prev_psi_zeta = ab_n.old_psi_zeta;
        prev_B = ab_n.old_B;
        prev_R = ab_n.old_R;

        // plug the data in
        ab[4 * n] = ab_n[0];
        ab[4 * n + 1] = ab_n[1];
        ab[4 * n + 2] = ab_n[2];
        ab[4 * n + 3] = ab_n[3];

    }

    return ab;
}