import { Bin } from "models";
import { useMemo } from "react";
import {
  Color,
  ShaderMaterial,
} from "three";
import { useThree } from "@react-three/fiber";

import { NodeData } from "bin/3dView/Data/BuildNodeData";
import { colourFade } from "bin/3dView/utils/tempToColour";

interface Props{
    bin: Bin
    nodes: NodeData[]
    opacity?: number
    /**
     * Point opacity (lower = see deeper).
     */
    pointOpacity?: number
    /**
     * Point size in world units (scaled with your Bin3dView scale group).
     */
    pointSize?: number
    /**
     * Enables MSAA alpha coverage smoothing (WebGL2 + MSAA).
     * Helps look continuous without additive blending.
     */
    alphaToCoverage?: boolean
    /**
     * Density along Y (vertical slices).
     */
    ySlices?: number
    /**
     * Radial rings per slice.
     */
    radialRings?: number
    /**
     * Angular segments per ring.
     */
    thetaSegments?: number
    /**
     * Inset to avoid z-fighting with the shell.
     */
    wallInsetFactor?: number
    /**
     * Adds jittered samples inside each polar cell to better fill the volume.
     */
    samplesPerCell?: number
    /**
     * 0..1 jitter amount within a cell (0 = none).
     */
    jitter?: number
    /**
     * Enables screen-door alpha hashing. This fixes incorrect transparency sorting
     * (points popping in front when tilted) without additive blending.
     */
    alphaHash?: boolean
    /**
     * Makes in-threshold (green) points more transparent so hot/cold pockets show through.
     * 0..1 where 0 = invisible green, 1 = same opacity as out-of-threshold.
     */
    greenOpacityFactor?: number
    /**
     * Curves how strongly out-of-threshold points become visible.
     * >1 makes only strong deviations pop; <1 makes small deviations pop more.
     */
    deviationPower?: number
    // (reverted) extra perf knobs removed
}

/**
 * this is a work in progress, the heatmap generated may not be accurate so avoid using this component for now
 * @param props 
 * @returns 
 */
export default function Heatmap(props: Props){
  const {
    bin,
    nodes,
    opacity = 0.65, // kept for backward compatibility
    pointOpacity,
    pointSize = 5,
    ySlices = 22,
    radialRings = 16,
    thetaSegments = 28,
    wallInsetFactor = 0.99,
    samplesPerCell = 1,
    jitter = 0.75,
    alphaHash = true,
    greenOpacityFactor = 0.18,
    deviationPower = 0.6,
  } = props;

  useThree(); // keep fiber context available if needed later

  const sidewallHeight = bin.sidewallHeight();
  const hopperHeight = bin.hopperHeight() ?? 0;
  const sidewallBaseY = -sidewallHeight / 2;
  const hopperTipY = sidewallBaseY - hopperHeight;

  const inGrainNodes = useMemo(
    () => nodes.filter((n) => n.inGrain && !n.excluded),
    [nodes],
  );

  const maxGrainY = useMemo(() => {
    let maxY = -Infinity;
    for (const n of inGrainNodes) maxY = Math.max(maxY, n.position.y);
    return Number.isFinite(maxY) ? maxY : sidewallBaseY;
  }, [inGrainNodes, sidewallBaseY]);

  const topNodes = useMemo(
    () => nodes.filter((n) => n.topNode && n.inGrain && !n.excluded),
    [nodes],
  );

  const anchors = useMemo(
    () =>
      topNodes.map((n) => ({
        x: n.position.x,
        z: n.position.z,
        y: n.position.y + n.nodeSpacing * 0.5,
      })),
    [topNodes],
  );

  const wallY = useMemo(() => {
    if (anchors.length === 0) return -sidewallHeight / 2;
    return anchors.reduce((sum, a) => sum + a.y, 0) / anchors.length;
  }, [anchors, sidewallHeight]);

  const idwHeight = (
    x: number,
    z: number,
    inputAnchors: { x: number; z: number; y: number }[],
    power = 2,
  ): number => {
    let totalWeight = 0;
    let weightedY = 0;

    for (const a of inputAnchors) {
      const dx = x - a.x;
      const dz = z - a.z;
      const distSq = dx * dx + dz * dz;
      if (distSq < 0.001) return a.y;
      const w = 1 / Math.pow(distSq, power / 2);
      totalWeight += w;
      weightedY += a.y * w;
    }

    return weightedY / totalWeight;
  };

  const maxRadiusAtY = (y: number, maxR: number): number => {
    if (y >= sidewallBaseY) return maxR;
    if (hopperHeight <= 0 || y <= hopperTipY) return 0;
    const t = (y - hopperTipY) / hopperHeight; // 0..1
    return maxR * t;
  };

  const grainSurfaceY = (x: number, z: number, rNorm: number): number => {
    // If we don't have top nodes, use a flat surface at maxGrainY.
    if (anchors.length === 0) return maxGrainY;

    const rawY = idwHeight(x, z, anchors);

    // Match `GrainCableFill` outer-wall taper so switching isn't jarring.
    const edgeStart = 0.8;
    const blendT = Math.max(0, (rNorm - edgeStart) / (1 - edgeStart));
    const s = blendT * blendT * (3 - 2 * blendT);
    const y = rawY * (1 - s) + wallY * s;
    return Math.max(-sidewallHeight / 2, Math.min(sidewallHeight / 2, y));
  };

  const evaluateTemp = (px: number, py: number, pz: number): number | null => {
    if (inGrainNodes.length === 0) return null;

    // Inverse-distance weighted interpolation.
    // Keep power modest so the field stays smooth.
    const IDW_POWER = 2;
    let totalWeight = 0;
    let weightedSum = 0;

    for (const n of inGrainNodes) {
      const dx = px - n.position.x;
      const dy = py - n.position.y;
      const dz = pz - n.position.z;
      const distSq = dx * dx + dy * dy + dz * dz;
      const weight = distSq < 0.001 ? 1e6 : 1 / Math.pow(distSq, IDW_POWER / 2);
      totalWeight += weight;
      weightedSum += n.celcius * weight;
    }

    if (totalWeight === 0) return null;
    return weightedSum / totalWeight;
  };

  const tempToHeatColor = (temp: number): Color => {
    // Match your 2D/point visuals: green in-threshold, fade to red/blue as distance grows.
    const lower = bin.lowerTempThreshold();
    const upper = bin.upperTempThreshold();

    const GREEN = new Color("#52c41a");
    const BLUE = new Color("#3399ff");
    const RED = new Color("#ff4d4f");

    if (temp >= lower && temp <= upper) return GREEN;

    const distance = temp < lower ? lower - temp : temp - upper;
    const intensity = Math.min(1, distance / colourFade); // 0..1

    // Similar HSL shaping as `TempToColour`, but always returns a color.
    const minimumLightness = 0.3;
    const lightnessRange = 0.2;
    const minimumSaturation = 0.7;
    const saturationRange = 0.8;

    const hsl = { h: 0, s: 1, l: 1 };
    (temp < lower ? BLUE : RED).getHSL(hsl);

    const c = new Color();
    c.setHSL(
      hsl.h,
      saturationRange * intensity + minimumSaturation,
      lightnessRange * intensity + minimumLightness,
    );
    return c;
  };

  const tempToDeviation = (temp: number): number => {
    const lower = bin.lowerTempThreshold();
    const upper = bin.upperTempThreshold();
    if (temp >= lower && temp <= upper) return 0;
    const distance = temp < lower ? lower - temp : temp - upper;
    return Math.min(1, distance / colourFade);
  };

  const { positions, colors, deviations } = useMemo(() => {
    const binR = bin.diameter() / 2;
    // Important: points are rendered as *sprites*, so even if the center is inside the wall,
    // the visible circle can extend outside. Shrink the sampling radius by ~half pointSize
    // so the rendered splats stay within the bin.
    const maxR = Math.max(0, binR * wallInsetFactor - pointSize * 0.55);

    const y0 = hopperHeight > 0 ? hopperTipY : sidewallBaseY;
    const y1 = Math.max(y0, maxGrainY);

    const pos: number[] = [];
    const col: number[] = [];
    const dev: number[] = [];
    const tmpColor = new Color();

    const safeYSlices = Math.max(6, Math.floor(ySlices));
    const safeRings = Math.max(4, Math.floor(radialRings));
    const safeTheta = Math.max(12, Math.floor(thetaSegments));
    const safeSamples = Math.max(1, Math.floor(samplesPerCell));
    const j = Math.min(1, Math.max(0, jitter));
    // Deterministic "random" so the cloud doesn't shimmer every render.
    const rand01 = (seed: number) => {
      // xorshift32
      let x = seed | 0;
      x ^= x << 13;
      x ^= x >>> 17;
      x ^= x << 5;
      // convert to [0,1)
      return ((x >>> 0) % 1000000) / 1000000;
    };

    for (let yi = 0; yi < safeYSlices; yi++) {
      const ty = safeYSlices === 1 ? 0 : yi / (safeYSlices - 1);
      const y = y0 + (y1 - y0) * ty;

      const rAtY = maxRadiusAtY(y, maxR);
      if (rAtY <= 0.001) continue;

      for (let ring = 0; ring < safeRings; ring++) {
        for (let seg = 0; seg < safeTheta; seg++) {
          // Cell bounds in polar space
          const ring0 = ring / safeRings;
          const ring1 = (ring + 1) / safeRings;
          const r0 = Math.sqrt(ring0) * rAtY;
          const r1 = Math.sqrt(ring1) * rAtY;

          const theta0 = (seg / safeTheta) * Math.PI * 2;
          const theta1 = ((seg + 1) / safeTheta) * Math.PI * 2;

          for (let s = 0; s < safeSamples; s++) {
            const seed = yi * 73856093 + ring * 19349663 + seg * 83492791 + s * 2654435761;
            const u = rand01(seed);
            const v = rand01(seed ^ 0x9e3779b9);

            // Jitter inside the cell
            const rr = r0 + (r1 - r0) * (j === 0 ? 0.5 : (0.5 + (u - 0.5) * j));
            const tt = theta0 + (theta1 - theta0) * (j === 0 ? 0.5 : (0.5 + (v - 0.5) * j));

            const x = Math.cos(tt) * rr;
            const z = Math.sin(tt) * rr;

            const rNorm = rAtY <= 0 ? 0 : rr / rAtY;
            const surfaceY = grainSurfaceY(x, z, rNorm);
            if (y > surfaceY) continue;

            const temp = evaluateTemp(x, y, z);
            const d0 = temp == null ? 0 : tempToDeviation(temp);

            pos.push(x, y, z);
            const c = temp == null ? tmpColor.set("#52c41a") : tempToHeatColor(temp);
            col.push(c.r, c.g, c.b);
            dev.push(d0);
          }
        }
      }
    }

    return {
      positions: new Float32Array(pos),
      colors: new Float32Array(col),
      deviations: new Float32Array(dev),
    };
  }, [
    bin,
    wallInsetFactor,
    hopperHeight,
    hopperTipY,
    sidewallBaseY,
    maxGrainY,
    ySlices,
    radialRings,
    thetaSegments,
    anchors,
    wallY,
    inGrainNodes,
    samplesPerCell,
    jitter,
    deviationPower,
  ]);

  const alphaHashedMaterial = useMemo(() => {
    return new ShaderMaterial({
      transparent: !alphaHash,
      depthTest: true,
      depthWrite: alphaHash,
      uniforms: {
        uOpacity: { value: pointOpacity ?? opacity },
        uSize: { value: pointSize },
        uMaxRadius: { value: (bin.diameter() / 2) * wallInsetFactor },
        uSidewallBaseY: { value: -bin.sidewallHeight() / 2 },
        uHopperHeight: { value: bin.hopperHeight() ?? 0 },
        uAlphaHash: { value: alphaHash ? 1 : 0 },
        uGreenOpacityFactor: { value: Math.min(1, Math.max(0, greenOpacityFactor)) },
        uDeviationPower: { value: Math.max(0.05, deviationPower) },
      },
      vertexShader: `
        uniform float uSize;
        varying vec3 vWorldPos;
        varying vec3 vColor;
        varying float vDev;
        attribute vec3 color;
        attribute float deviation;
        void main() {
          vColor = color;
          vDev = deviation;
          vec4 world = modelMatrix * vec4(position, 1.0);
          vWorldPos = world.xyz;
          vec4 mvPosition = viewMatrix * world;
          float attn = 300.0 / max(1.0, -mvPosition.z);
          gl_PointSize = uSize * attn;
          gl_Position = projectionMatrix * mvPosition;
        }
      `,
      fragmentShader: `
        precision highp float;
        uniform float uOpacity;
        uniform float uMaxRadius;
        uniform float uSidewallBaseY;
        uniform float uHopperHeight;
        uniform float uAlphaHash;
        uniform float uGreenOpacityFactor;
        uniform float uDeviationPower;
        varying vec3 vColor;
        varying float vDev;
        varying vec3 vWorldPos;
        // interleaved gradient noise
        float ign(vec2 p) {
          return fract(52.9829189 * fract(dot(p, vec2(0.06711056, 0.00583715))));
        }
        void main() {
          // Hard clip pixels to bin radius at this Y (prevents splats outside wall).
          float y = vWorldPos.y;
          float sidewallBaseY = uSidewallBaseY;
          float hopperHeight = uHopperHeight;
          float hopperTipY = sidewallBaseY - hopperHeight;
          float maxR;
          if (y >= sidewallBaseY) {
            maxR = uMaxRadius;
          } else if (hopperHeight <= 0.0 || y <= hopperTipY) {
            maxR = 0.0;
          } else {
            float t = (y - hopperTipY) / hopperHeight;
            maxR = uMaxRadius * t;
          }
          float r = length(vWorldPos.xz);
          if (r > maxR) discard;

          vec2 p = gl_PointCoord - vec2(0.5);
          float d = length(p) * 2.0;
          float mask = smoothstep(1.0, 0.0, d);

          float dev = clamp(vDev, 0.0, 1.0);
          float devCurve = pow(dev, uDeviationPower);
          // 0 => green/in-threshold, 1 => strong deviation
          float localOpacityFactor = mix(uGreenOpacityFactor, 1.0, devCurve);
          float a = clamp(mask * uOpacity * localOpacityFactor, 0.0, 1.0);

          if (uAlphaHash > 0.5) {
            float n = ign(gl_FragCoord.xy);
            if (n > a) discard;
            gl_FragColor = vec4(vColor, 1.0);
          } else {
            gl_FragColor = vec4(vColor, a);
          }
        }
      `,
    });
  }, [alphaHash, bin, deviationPower, greenOpacityFactor, opacity, pointOpacity, pointSize, wallInsetFactor]);

  // Fallback: normal points (no OIT)
  return (
    <points renderOrder={2} material={alphaHashedMaterial ?? undefined}>
      <bufferGeometry>
        <bufferAttribute
          attach="attributes-position"
          array={positions}
          count={positions.length / 3}
          itemSize={3}
        />
        <bufferAttribute
          attach="attributes-color"
          array={colors}
          count={colors.length / 3}
          itemSize={3}
        />
        <bufferAttribute
          attach="attributes-deviation"
          array={deviations}
          count={deviations.length}
          itemSize={1}
        />
      </bufferGeometry>
    </points>
  );
}