import { vec3 } from "gl-matrix";
import { assert } from "../../../../core/assert.js";
import { ImpostorBaker } from "../octahedral/ImpostorBaker.js";
import { ImpostorCaptureType } from "../octahedral/ImpostorCaptureType.js";
import { read_atlas_to_samples } from "./bake/read_atlas_to_samples.js";
import { voxelize_samples } from "./bake/voxelize_samples.js";
import { VoxelImpostorDescription } from "./VoxelImpostorDescription.js";


/**
 * Compute an axis-aligned bounding box over every input mesh's positions
 * after applying its world transform. We need this in addition to the
 * bounding sphere (which the underlying ImpostorBaker already computes)
 * because the voxel grid is axis-aligned, not spherical — a sphere bound
 * would waste cells in the corners.
 *
 * @param {{mesh:ShadedGeometry, transform:mat4}[]} objects
 * @returns {{min:number[], max:number[]}}
 */
function compute_bounding_box(objects) {
    const min = [Infinity, Infinity, Infinity];
    const max = [-Infinity, -Infinity, -Infinity];

    for (let k = 0; k < objects.length; k++) {
        const object = objects[k];

        const positions = object.mesh.geometry.getAttribute('position');
        const arr = positions.array;
        const m = object.transform;

        for (let i = 0; i < positions.count; i++) {
            const ix = i * positions.itemSize;
            const x = arr[ix];
            const y = arr[ix + 1];
            const z = arr[ix + 2];

            // Manual mat4 * vec3 (point) — Float32Array-friendly.
            const wx = m[0] * x + m[4] * y + m[8]  * z + m[12];
            const wy = m[1] * x + m[5] * y + m[9]  * z + m[13];
            const wz = m[2] * x + m[6] * y + m[10] * z + m[14];

            if (wx < min[0]) min[0] = wx;
            if (wy < min[1]) min[1] = wy;
            if (wz < min[2]) min[2] = wz;
            if (wx > max[0]) max[0] = wx;
            if (wy > max[1]) max[1] = wy;
            if (wz > max[2]) max[2] = wz;
        }
    }

    return { min, max };
}


/**
 * Voxel ("Nanite voxel") impostor baker.
 *
 * Inspiration: the Witcher 4 / UE5.6+ "Nanite Foliage" tech demo, which
 * swaps alpha-tested cards for a near-pixel-sized voxel representation
 * at distance. Each "voxel" is a point with a world-space size, carrying
 * a baked G-buffer (position, normal, albedo, ORM) — no alpha. Targeted
 * at very low point counts (128–2048), where each voxel projects to
 * around one screen pixel at the impostor's design viewing distance.
 *
 * Bake strategy: control the point count by **atlas resolution**, not
 * by varying the voxel grid resolution at a fixed atlas. Reason: varying
 * the voxel grid at fixed atlas biases each cell's centroid toward
 * whichever views happened to land more atlas pixels in that cell — a
 * systematic directional bias that grows with coarser grids. Re-baking
 * the atlas at the right resolution avoids that, because each bake
 * produces fresh, evenly-distributed samples in the target frequency
 * range.
 *
 * Search algorithm:
 *
 *   PHASE 1 — coarse halving. Start at `resolution`, halve atlas
 *   resolution each iteration until the point count fits budget.
 *   O(log) iterations, each a full GPU bake.
 *
 *   PHASE 2 — binary-search refinement. Between the last over-budget
 *   resolution and the first under-budget one, binary-search in
 *   steps of `frames` (the smallest atlas-resolution granularity
 *   that keeps each per-frame tile integer-sized). Pick the
 *   largest atlas resolution that still fits — that gives the most
 *   points up to budget = best detail.
 *
 * Each iteration:
 *   1. Re-bake atlas at candidate resolution via {@link ImpostorBaker}.
 *   2. Read back samples on CPU.
 *   3. Voxelise at grid_res = atlas_res / frames (matched 1:1 so each
 *      atlas pixel maps to roughly one voxel cell).
 *   4. Count, compare to budget, advance search.
 *
 * Re-baking costs a GPU pass per iteration (~250 draws for 16² frames),
 * so a full search is ~10 bakes ≈ a couple seconds on modern hardware.
 * Acceptable for one-shot offline impostor prep where quality
 * dominates.
 */
export class VoxelImpostorBaker {

    /**
     * @type {THREE.WebGLRenderer|null}
     * @private
     */
    _renderer = null;

    /**
     *
     * @param {THREE.WebGLRenderer} v
     */
    set renderer(v) {
        this._renderer = v;
    }

    /**
     * @param {object} params
     * @param {{mesh:ShadedGeometry, transform:mat4}[]} params.objects
     * @param {number} [params.point_budget]    target maximum point count.
     *      The baker may produce fewer if the input is too small for the
     *      target, but never more.
     * @param {ImpostorCaptureType} [params.capture_type]
     * @param {number} [params.frames]          atlas frame count per axis
     *      (total sampled views = frames²). Also the granularity of the
     *      atlas-resolution binary search: candidates must be multiples
     *      of `frames` so each per-frame tile stays integer-sized.
     *      Default 16.
     * @param {number} [params.resolution]      maximum atlas resolution.
     *      Upper bound of the search; the baker may pick a smaller value
     *      to fit the budget. Default 1024.
     * @returns {VoxelImpostorDescription}
     */
    bake({
             objects,
             point_budget = 16384,
             capture_type = ImpostorCaptureType.FullSphere,
             frames = 16,
             resolution = 1024
         }) {

        console.time('voxel bake');

        assert.defined(objects, 'objects');
        assert.isNonNegativeInteger(point_budget, 'point_budget');
        assert.greaterThanOrEqual(point_budget, 1, 'point_budget must be >= 1');

        if (this._renderer === null) {
            throw new Error('No renderer attached. Renderer is required for baking.');
        }

        const renderer = this._renderer;

        // Bounding box stays the same across re-bakes — compute once.
        const aabb = compute_bounding_box(objects);

        const octa_baker = new ImpostorBaker();
        octa_baker.renderer = renderer;

        // Minimum meaningful atlas resolution: 2 pixels per frame
        // (1 pixel per frame is rasteriser-deadtime; 2 gives some
        // content per view). Below this, the search is degenerate.
        const min_resolution = Math.max(frames * 2, frames);

        // Floor max resolution to a multiple of `frames` so the very
        // first iteration is already a valid candidate.
        const max_resolution = Math.floor(resolution / frames) * frames;
        if (max_resolution < min_resolution) {
            throw new Error(`resolution (=${resolution}) is too small for frames (=${frames}); need at least ${min_resolution}`);
        }

        /**
         * One bake-and-voxelise iteration at a given atlas resolution.
         * Each call is a full GPU bake + CPU readback + voxelise; the
         * atlas is disposed immediately after voxelising since the
         * point cloud carries everything we need.
         *
         * @param {number} atlas_res
         * @returns {{
         *   atlas_res: number,
         *   voxelised: {points: Float32Array, point_count: number, voxel_size: number},
         *   sample_count: number,
         *   sphere_radius: number,
         *   offset: number[],
         *   source_geometry_polygon_count: number,
         *   source_geometry_vertex_count: number,
         *   source_material_count: number,
         *   source_instance_count: number
         * }}
         */
        const bake_and_voxelise_at = (atlas_res) => {
            console.time(`voxel bake / iter ${atlas_res}`);

            const atlas = octa_baker.bake({
                objects,
                frames,
                resolution: atlas_res,
                type: capture_type
            });

            const { samples, sample_count } = read_atlas_to_samples({
                impostor: atlas,
                renderer
            });

            // grid_res = atlas pixels per frame, matched 1:1 so each
            // atlas pixel maps to roughly one voxel cell. With many
            // views (frames²), each cell still gets multiple samples
            // for noise-averaging — just fewer than at higher atlas
            // resolutions.
            const grid_resolution = atlas_res / frames;

            // Use the AABB plus a small margin (1% of bounding-sphere
            // radius) so depth-quantisation jitter at the AABB
            // boundary doesn't get rejected by the voxeliser.
            const margin = atlas.sphere_radius * 0.01;
            const bounds_min = [aabb.min[0] - margin, aabb.min[1] - margin, aabb.min[2] - margin];
            const bounds_max = [aabb.max[0] + margin, aabb.max[1] + margin, aabb.max[2] + margin];

            const voxelised = voxelize_samples({
                samples,
                sample_count,
                bounds_min,
                bounds_max,
                grid_resolution
            });

            const sphere_radius = atlas.sphere_radius;
            const offset = [atlas.offset[0], atlas.offset[1], atlas.offset[2]];
            const source_geometry_polygon_count = atlas.source_geometry_polygon_count;
            const source_geometry_vertex_count = atlas.source_geometry_vertex_count;
            const source_material_count = atlas.source_material_count;
            const source_instance_count = atlas.source_instance_count;

            // Free GPU MRT — we have everything we need on CPU now.
            atlas.rt.dispose();

            console.timeEnd(`voxel bake / iter ${atlas_res}`);
            console.log(`  atlas ${atlas_res} (grid ${grid_resolution}): ${voxelised.point_count} points from ${sample_count} samples`);

            return {
                atlas_res,
                voxelised,
                sample_count,
                sphere_radius,
                offset,
                source_geometry_polygon_count,
                source_geometry_vertex_count,
                source_material_count,
                source_instance_count
            };
        };


        // --- PHASE 1: coarse halving ---
        // Find the largest power-of-two-halving of max_resolution that
        // fits the budget, plus the smallest one that doesn't. These
        // bracket the optimal atlas resolution for phase 2.
        let coarse_under = null;
        let coarse_over_res = null;

        let res = max_resolution;

        while (true) {
            const result = bake_and_voxelise_at(res);

            if (result.voxelised.point_count <= point_budget) {
                coarse_under = result;
                break;
            }

            coarse_over_res = res;

            const next = Math.max(
                min_resolution,
                Math.floor((res / 2) / frames) * frames
            );

            if (next >= res) {
                // Halving would not make progress — accept the floor.
                return this._compose_result({
                    bake_result: bake_and_voxelise_at(min_resolution),
                    capture_type,
                    grid_resolution: min_resolution / frames,
                    source_atlas_frames: frames,
                    source_atlas_resolution: min_resolution
                });
            }

            res = next;
        }


        // --- PHASE 2: binary-search refinement ---
        // Between (coarse_over_res, coarse_under.atlas_res], find the
        // largest atlas resolution that still fits the budget. Step
        // granularity is `frames` so per-frame tiles stay integer-sized.
        let best = coarse_under;

        if (coarse_over_res !== null) {
            let lo = coarse_under.atlas_res;  // confirmed fits budget
            let hi = coarse_over_res;          // confirmed overflows

            while (hi - lo > frames) {
                // Round mid down to a multiple of `frames`.
                const mid_raw = Math.floor((lo + hi) / 2);
                const mid = Math.floor(mid_raw / frames) * frames;

                if (mid <= lo || mid >= hi) {
                    // No more integer-step candidates between lo and hi.
                    break;
                }

                const result = bake_and_voxelise_at(mid);

                if (result.voxelised.point_count <= point_budget) {
                    lo = mid;
                    best = result;
                } else {
                    hi = mid;
                }
            }
        }


        console.timeEnd('voxel bake');
        console.log(`VoxelImpostorBaker: settled on atlas ${best.atlas_res} → ${best.voxelised.point_count} points (budget ${point_budget})`);

        return this._compose_result({
            bake_result: best,
            capture_type,
            grid_resolution: best.atlas_res / frames,
            source_atlas_frames: frames,
            source_atlas_resolution: best.atlas_res
        });
    }


    /**
     * Pack the search result into a final VoxelImpostorDescription.
     * @private
     */
    _compose_result({
                        bake_result,
                        capture_type,
                        grid_resolution,
                        source_atlas_frames,
                        source_atlas_resolution
                    }) {
        const result = new VoxelImpostorDescription();

        result.capture_type = capture_type;
        result.sphere_radius = bake_result.sphere_radius;
        vec3.copy(result.offset, bake_result.offset);

        result.points = bake_result.voxelised.points;
        result.point_count = bake_result.voxelised.point_count;
        result.voxel_size = bake_result.voxelised.voxel_size;
        result.grid_resolution = grid_resolution;

        // Auto-tuned splat diameter — sqrt(2) covers the worst-case
        // oblique projection of a voxel cell (a 45°-to-view face is
        // sqrt(2)x wider in screen than its on-axis projection). Less-
        // oblique angles get overlap, depth-test resolves the
        // stacking. Render-time inward-shift (see shaders) keeps the
        // silhouette tight regardless.
        result.point_diameter = bake_result.voxelised.voxel_size * Math.SQRT2;

        result.source_atlas_resolution = source_atlas_resolution;
        result.source_atlas_frames = source_atlas_frames;
        result.source_geometry_polygon_count = bake_result.source_geometry_polygon_count;
        result.source_geometry_vertex_count = bake_result.source_geometry_vertex_count;
        result.source_material_count = bake_result.source_material_count;
        result.source_instance_count = bake_result.source_instance_count;

        return result;
    }
}
