import { mat3, vec3 } from "gl-matrix";
import { assert } from "../../../../assert.js";
import { array_copy } from "../../../../collection/array/array_copy.js";
import { min2 } from "../../../../math/min2.js";
import { sinc } from "../../../../math/sinc.js";
import { SH3_COEFFICIENTS } from "./SH3_COEFFICIENTS.js";
import { sh3_rotate_bl } from "./sh3_rotate_bl.js";
import { sh_index } from "./sh_index.js";

/*
@see https://github.com/Bestmaker602/olacziy/blob/212b64ea5f1856b390cdf7629801243f76a4466d/libs/ibl/src/CubemapSH.cpp
 */

const scratch_sh3 = new Float32Array(9);

/**
 * SH from environment with high dynamic range (or high frequencies -- high dynamic range creates
 * high frequencies) exhibit "ringing" and negative values when reconstructed.
 * To mitigate this, we need to low-pass the input image -- or equivalently window the SH by
 * coefficient that tapper towards zero with the band.
 *
 * We use ideas and techniques from Stupid Spherical Harmonics (SH) Deringing Spherical Harmonics by Peter-Pike Sloan ( see https://www.ppsloan.org/publications/shdering.pdf )
 * @param {number} w cutoff
 * @param {number} l
 * @returns {number}
 */
function sinc_window(l, w) {
    if (l === 0) {
        return 1.0;
    } else if (l >= w) {
        return 0.0;
    }

    // we use a sinc window scaled to the desired window size in bands units
    // a sinc window only has zonal harmonics
    let x = (Math.PI * l) / w;

    x = sinc(x);

    // The convolution of a SH function f and a ZH function h is just the product of both
    // scaled by 1 / K(0,l) -- the window coefficients include this scale factor.

    // Taking the window to power N is equivalent to applying the filter N times
    return Math.pow(x, 4);
}

/**
 * Minimum values encoded on the spherical harmonics
 * Lets us find ringing artifacts, which are cases where harmonics encode a negative value
 * @param {number[]|Float32Array} input_sh3 sh3
 * @returns {number}
 */
function sh3_min(input_sh3) {
    // See "Deringing Spherical Harmonics" by Peter-Pike Sloan
    // https://www.ppsloan.org/publications/shdering.pdf


    // first this to do is to rotate the SH to align Z with the optimal linear direction
    const dir = vec3.fromValues(-input_sh3[3], -input_sh3[1], input_sh3[2]);

    vec3.normalize(dir, dir);

    const z_axis = vec3.create();
    vec3.negate(z_axis, dir);

    const x_axis = vec3.create();
    vec3.cross(x_axis, z_axis, vec3.fromValues(0, 1, 0));
    vec3.normalize(x_axis, x_axis);

    const y_axis = vec3.create();
    vec3.cross(y_axis, x_axis, z_axis);

    const M = mat3.create();
    array_copy(x_axis, 0, M, 0, 3);
    array_copy(y_axis, 0, M, 3, 3);
    array_copy(dir, 0, M, 6, 3);

    const f = new Float32Array(9);

    sh3_rotate_bl(f, input_sh3, M);
    // here we're guaranteed to have normalize(float3{ -f[3], -f[1], f[2] }) == { 0, 0, 1 }


    // Find the min for |m| = 2
    // ------------------------
    //
    // Peter-Pike Sloan shows that the minimum can be expressed as a function
    // of z such as:  m2min = -m2max * (1 - z^2) =  m2max * z^2 - m2max
    //      with m2max = A[8] * std::sqrt(f[8] * f[8] + f[4] * f[4]);
    // We can therefore include this in the ZH min computation (which is function of z^2 as well)
    const m2max = SH3_COEFFICIENTS[8] * Math.sqrt(f[8] * f[8] + f[4] * f[4]);

    // Find the min of the zonal harmonics
    // -----------------------------------
    //
    // This comes from minimizing the function:
    //      ZH(z) = (A[0] * f[0])
    //            + (A[2] * f[2]) * z
    //            + (A[6] * f[6]) * (3 * s.z * s.z - 1)
    //
    // We do that by finding where it's derivative d/dz is zero:
    //      dZH(z)/dz = a * z^2 + b * z + c
    //      which is zero for z = -b / 2 * a
    //
    // We also needs to check that -1 < z < 1, otherwise the min is either in z = -1 or 1
    //
    const a = 3 * SH3_COEFFICIENTS[6] * f[6] + m2max;
    const b = SH3_COEFFICIENTS[2] * f[2];
    const c = SH3_COEFFICIENTS[0] * f[0] - SH3_COEFFICIENTS[6] * f[6] - m2max;

    const zmin = -b / (2.0 * a);
    const m0min_z = a * zmin * zmin + b * zmin + c;
    const m0min_b = min2(a + b + c, a - b + c);

    const m0min = (a > 0 && zmin >= -1 && zmin <= 1) ? m0min_z : m0min_b;

    // Find the min for l = 2, |m| = 1
    // -------------------------------
    //
    // Note l = 1, |m| = 1 is guaranteed to be 0 because of the rotation step
    //
    // The function considered is:
    //        Y(x, y, z) = A[5] * f[5] * s.y * s.z
    //                   + A[7] * f[7] * s.z * s.x
    const d = SH3_COEFFICIENTS[4] * Math.sqrt(f[5] * f[5] + f[7] * f[7]);

    // the |m|=1 function is minimal in -0.5 -- use that to skip the Newton's loop when possible
    let minimum = m0min - 0.5 * d;

    if (minimum < 0) {
        // We could be negative, to find the minimum we will use Newton's method
        // See https://en.wikipedia.org/wiki/Newton%27s_method_in_optimization


        let dz;
        let z = -Math.SQRT1_2;   // we start guessing at the min of |m|=1 function
        do {
            minimum = func_(z, a, b, c, d); // evaluate our function
            dz = increment_(z, a, b, d); // refine our guess by this amount
            z = z - dz;
            // exit if z goes out of range, or if we have reached enough precision
        } while (Math.abs(z) <= 1 && Math.abs(dz) > 1e-5);

        if (Math.abs(z) > 1) {
            // z was out of range
            minimum = min2(func_(1, a, b, c, d), func_(-1, a, b, c, d));
        }
    }

    return minimum;
}

/**
 * this is the function we're trying to minimize
 * @param {number} x
 * @param {number} a
 * @param {number} b
 * @param {number} c
 * @param {number} d
 * @return {number}
 */
function func_(x, a, b, c, d) {
    // first term accounts for ZH + |m| = 2, second terms for |m| = 1
    return (a * x * x + b * x + c) + (d * x * Math.sqrt(1 - x * x));
}

/**
 * This is func' / func'' -- this was computed with Mathematica
 * @param {number} x
 * @param {number} a
 * @param {number} b
 * @param {number} d
 * @return {number}
 */
function increment_(x, a, b, d) {
    return (x * x - 1) * (d - 2 * d * x * x + (b + 2 * a * x) * Math.sqrt(1 - x * x))
        / (3 * d * x - 2 * d * x * x * x - 2 * a * Math.pow(1 - x * x, 1.5));
}

/**
 *
 * @param {number[]|Float32Array} f sh3
 * @param {number} cutoff
 * @param {number} numBands
 */
function windowing(f, cutoff, numBands) {

    for (let l = 0; l < numBands; l++) {

        const w = sinc_window(l, cutoff);

        f[sh_index(0, l)] *= w;

        for (let m = 1; m <= l; m++) {

            f[sh_index(-m, l)] *= w;
            f[sh_index(m, l)] *= w;

        }

    }

}

/**
 *
 * @param {number[]} output sh3
 * @param {number} output_offset
 * @param {number[]} input
 * @param {number} input_offset
 * @param {number} numBands
 * @param {number} channel_count
 * @param {number} [cutoff]
 */
function windowSH(
    output, output_offset,
    input, input_offset,
    numBands, channel_count,
    cutoff = 0
) {
    assert.isNonNegativeInteger(channel_count, 'channel_count');
    assert.greaterThan(channel_count, 0, 'channel_count must be greater than 0');

    assert.greaterThanOrEqual(cutoff, 0, 'cutoff must be >= 0');

    if (cutoff === 0) {
        // auto windowing (default)

        if (numBands > 3) {
            // auto-windowing works only for 1, 2 or 3 bands
            throw new Error("--sh-window=auto can't work with more than 3 bands. Disabling.");
        }

        cutoff = numBands * 4 + 1;// start at a large band

        // We need to process each channel separately
        for (let channel = 0; channel < channel_count; channel++) {

            for (let i = 0; i < numBands * numBands; i++) {
                // copy this channel
                scratch_sh3[i] = input[input_offset + i * channel_count + channel];
            }

            // find a cut-off band that works
            let l = numBands;
            let r = cutoff;

            for (let i = 0; i < 16 && l + 0.1 < r; i++) {

                const m = 0.5 * (l + r);

                windowing(scratch_sh3, m, numBands);

                if (sh3_min(scratch_sh3) < 0) {
                    r = m;
                } else {
                    l = m;
                }

            }

            // record minimum cutoff
            cutoff = min2(cutoff, l);
        }
    }

    array_copy(input, input_offset, output, output_offset, numBands * numBands * channel_count);

    for (let l = 0; l < numBands; l++) {
        // scale each band individually
        const w = sinc_window(l, cutoff);

        // write zonal harmonic
        for (let i = 0; i < channel_count; i++) {

            const j = sh_index(0, l);

            output[output_offset + j * channel_count + i] *= w;
        }

        // write the rest of the band
        for (let m = 1; m <= l; m++) {
            for (let i = 0; i < channel_count; i++) {

                const j0 = sh_index(-m, l);

                output[output_offset + j0 * channel_count + i] *= w;

                const j1 = sh_index(m, l);

                output[output_offset + j1 * channel_count + i] *= w;

            }
        }
    }

}

/**
 *
 * @param {number[]|Float32Array|Float64Array} result
 * @param {number} result_offset
 * @param {number[]|Float32Array|Float64Array} harmonics
 * @param {number} harmonics_offset
 * @param {number} dimension_count
 */
export function sh3_dering_optimize_positive(result, result_offset, harmonics, harmonics_offset, dimension_count = 3) {
    windowSH(result, result_offset, harmonics, harmonics_offset, 3, dimension_count, 0);
}
