Dwinzo_Demo/app/src/components/heatMapGenerator/baked/bakedHeatMap.tsx

import * as THREE from "three";
import { useEffect, useMemo, useRef, useCallback } from "react";
import { useThree } from "@react-three/fiber";
import { useHeatMapStore } from "../../../store/simulation/useHeatMapStore";
import * as CONSTANTS from "../../../types/world/worldConstants";


const RADIUS = 0.0025;
const OPACITY = 0.8;
const GROWTH_RATE = 20.0;

const BakedHeatMap = () => {
    const { bakedPoints } = useHeatMapStore();
    const materialRef = useRef<THREE.ShaderMaterial>(null);
    const meshRef = useRef<THREE.Mesh>(null);
    const { gl } = useThree();

    const height = CONSTANTS.gridConfig.size;
    const width = CONSTANTS.gridConfig.size;

    const createPointTexture = useCallback(
        (filteredPoints: typeof bakedPoints) => {
            if (filteredPoints.length === 0) return null;

            const data = new Float32Array(filteredPoints.length * 4);
            filteredPoints.forEach((p, i) => {
                const index = i * 4;
                data[index] = (p.points.x + width / 2) / width;
                data[index + 1] = (p.points.y + height / 2) / height;
                data[index + 2] = 0.3;
                data[index + 3] = 0.0;
            });

            const texture = new THREE.DataTexture(data, filteredPoints.length, 1, THREE.RGBAFormat, THREE.FloatType);
            texture.needsUpdate = true;
            return texture;
        },
        [width, height]
    );

    const uniformsRef = useRef({
        u_points: { value: null as THREE.DataTexture | null },
        u_count: { value: 0 },
        u_radius: { value: RADIUS },
        u_opacity: { value: OPACITY },
        u_growthRate: { value: GROWTH_RATE },
    });

    const renderHeatmapToImage = useCallback(
        (type: string) => {
            if (!meshRef.current) return null;

            const filteredPoints = bakedPoints.filter((p) => p.type === type);
            if (filteredPoints.length === 0) return null;

            const pointTexture = createPointTexture(filteredPoints);
            if (!pointTexture) return null;

            uniformsRef.current.u_points.value = pointTexture;
            uniformsRef.current.u_count.value = filteredPoints.length;

            const exportCamera = new THREE.OrthographicCamera(width / -2, width / 2, height / 2, height / -2, 0.1, 10);
            exportCamera.position.set(0, 1, 0);
            exportCamera.lookAt(0, 0, 0);

            const renderTarget = new THREE.WebGLRenderTarget(1024, 1024, {
                format: THREE.RGBAFormat,
                type: THREE.UnsignedByteType,
            });

            const tempScene = new THREE.Scene();
            tempScene.add(meshRef.current);

            gl.setRenderTarget(renderTarget);
            gl.render(tempScene, exportCamera);
            gl.setRenderTarget(null);

            const pixels = new Uint8Array(1024 * 1024 * 4);
            gl.readRenderTargetPixels(renderTarget, 0, 0, 1024, 1024, pixels);

            const canvas = document.createElement("canvas");
            canvas.width = 1024;
            canvas.height = 1024;
            const ctx = canvas.getContext("2d");

            if (ctx) {
                const imageData = ctx.createImageData(1024, 1024);
                imageData.data.set(pixels);
                ctx.putImageData(imageData, 0, 0);
                return canvas.toDataURL("image/png");
            }

            return null;
        },
        [gl, width, height, bakedPoints, createPointTexture]
    );

    const downloadImage = (base64: string, filename: string) => {
        const link = document.createElement("a");
        link.href = base64;
        link.download = filename;
        document.body.appendChild(link);
        link.click();
        document.body.removeChild(link);
    };

    const exportHeatmapAsPNG = useCallback(() => {
        const types = ["human", "vehicle"];

        const result = types.map((type) => {
            const image = renderHeatmapToImage(type);
            if (image) {
                downloadImage(image, `${type}-heatmap.png`);
            }
            return { type, image };
        });

        console.log("Exported Heatmaps:", result);
        return result;
    }, [renderHeatmapToImage]);

    const pointTexture = useMemo(() => createPointTexture(bakedPoints), [bakedPoints, createPointTexture]);

    useEffect(() => {
        uniformsRef.current.u_points.value = pointTexture;
        uniformsRef.current.u_count.value = bakedPoints.length;
    }, [pointTexture, bakedPoints.length]);

    return (
        <>
            <mesh ref={meshRef} rotation={[Math.PI / 2, 0, 0]} position={[0, 0.025, 0]}>
                <planeGeometry args={[width, height]} />
                <shaderMaterial
                    ref={materialRef}
                    transparent
                    depthWrite={false}
                    blending={THREE.AdditiveBlending}
                    uniforms={uniformsRef.current}
                    side={THREE.DoubleSide}
                    vertexShader={`
                        varying vec2 vUv;
                        void main() {
                            vUv = uv;
                            gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
                        }
                    `}
                    fragmentShader={`
                        uniform sampler2D u_points;
                        precision highp float;
                        uniform int u_count;
                        uniform float u_radius;
                        uniform float u_opacity;
                        uniform float u_growthRate;
                        varying vec2 vUv;

                        float gauss(float dist, float radius) {
                            return exp(-pow(dist / radius, 2.0));
                        }

                        void main() {
                            float intensity = 0.0;

                            for (int i = 0; i < 10000; i++) {
                                if (i >= u_count) break;
                                float fi = float(i) + 0.5;
                                float u = fi / float(u_count);

                                vec4 point = texture2D(u_points, vec2(u, 0.5));
                                vec2 pos = point.rg;
                                float strength = point.b;

                                float d = distance(vUv, pos);
                                intensity += strength * gauss(d, u_radius);
                            }

                            float normalized = clamp(intensity / max(u_growthRate, 0.0001), 0.0, 1.0);

                            vec3 color = vec3(0.0);
                            if (normalized < 0.33) {
                                color = mix(vec3(0.0, 0.0, 1.0), vec3(0.0, 1.0, 0.0), normalized / 0.33);
                            } else if (normalized < 0.66) {
                                color = mix(vec3(0.0, 1.0, 0.0), vec3(1.0, 1.0, 0.0), (normalized - 0.33) / 0.33);
                            } else {
                                color = mix(vec3(1.0, 1.0, 0.0), vec3(1.0, 0.0, 0.0), (normalized - 0.66) / 0.34);
                            }

                            gl_FragColor = vec4(color, normalized * u_opacity);
                        }
                    `}
                />
            </mesh>
        </>
    );
};

export default BakedHeatMap;