"use client";

import { useMemo, useState } from "react";

const BIT_DEPTHS = [2, 4, 6, 8, 16] as const;
const FOCUS_WEIGHT = 0.156347;
const EXAMPLE_WEIGHTS = [0.156347, -0.3734, 0.7234] as const;

type BitDepth = (typeof BIT_DEPTHS)[number];

type QuantizedWeight = {
  bf16Word?: string;
  error: number;
  original: number;
  reconstructed: number;
  stored: number | string;
};

function clamp(value: number, min: number, max: number) {
  return Math.min(max, Math.max(min, value));
}

function formatFloat(value: number, decimals = 6) {
  return value.toFixed(decimals);
}

function toBfloat16(value: number) {
  const floatView = new Float32Array(1);
  const intView = new Uint32Array(floatView.buffer);

  floatView[0] = value;
  const current = intView[0] ?? 0;
  const leastSignificantBit = (current >> 16) & 1;
  const roundingBias = 0x7fff + leastSignificantBit;
  const rounded = (current + roundingBias) & 0xffff0000;

  intView[0] = rounded;

  return {
    reconstructed: floatView[0] ?? value,
    word: `0x${(rounded >>> 16).toString(16).padStart(4, "0")}`,
  };
}

function getSharedScale(bitDepth: Exclude<BitDepth, 16>) {
  const maxMagnitude = Math.max(
    ...EXAMPLE_WEIGHTS.map((value) => Math.abs(value)),
  );
  const qmax = Math.pow(2, bitDepth - 1) - 1;

  return {
    qmax,
    scale: maxMagnitude / qmax,
  };
}

function quantizeWeight(value: number, bitDepth: BitDepth): QuantizedWeight {
  if (bitDepth === 16) {
    const bf16 = toBfloat16(value);
    return {
      bf16Word: bf16.word,
      error: bf16.reconstructed - value,
      original: value,
      reconstructed: bf16.reconstructed,
      stored: bf16.word,
    };
  }

  const { scale, qmax } = getSharedScale(bitDepth);
  const stored = clamp(Math.round(value / scale), -qmax, qmax);
  const reconstructed = stored * scale;

  return {
    error: reconstructed - value,
    original: value,
    reconstructed,
    stored,
  };
}

function getShortExplanation(bitDepth: BitDepth) {
  if (bitDepth === 16) {
    return "BF16 keeps much more of the original value because it is still a floating-point format.";
  }

  if (bitDepth === 2) {
    return "At 2 bits, many nearby weights are forced into the same tiny set of buckets.";
  }

  if (bitDepth === 4) {
    return "At 4 bits, the model has more buckets to work with, so the approximation gets better.";
  }

  if (bitDepth === 6) {
    return "At 6 bits, the rounded result is usually noticeably closer to the original.";
  }

  return "At 8 bits, the rounded result is often fairly close to the original weight.";
}

export function QuantizationExplorer() {
  const [bitDepthIndex, setBitDepthIndex] = useState(0);

  const bitDepth = BIT_DEPTHS[bitDepthIndex] ?? BIT_DEPTHS[0];
  const scaleSummary =
    bitDepth === 16 ? null : getSharedScale(bitDepth as Exclude<BitDepth, 16>);

  const focusWeight = useMemo(
    () => quantizeWeight(FOCUS_WEIGHT, bitDepth),
    [bitDepth],
  );

  return (
    <div className="quantization-explorer" data-widget-enhanced="true">
      <div className="quantization-explorer__header">
        <p className="quantization-explorer__eyebrow">Single Weight View</p>
        <h3>See one stored value become an approximation</h3>
        <p className="quantization-explorer__lede">
          The original weight is <code>{formatFloat(FOCUS_WEIGHT)}</code>. Lower
          precision stores a rougher version of it.
        </p>
      </div>

      <div className="quantization-explorer__controls">
        <div className="quantization-explorer__slider-card">
          <label
            className="quantization-explorer__slider-label"
            htmlFor="single-quant-depth"
          >
            Precision:{" "}
            <strong>
              {bitDepth === 16 ? "16-bit BF16" : `${bitDepth}-bit`}
            </strong>
          </label>
          <input
            id="single-quant-depth"
            type="range"
            min={0}
            max={BIT_DEPTHS.length - 1}
            step={1}
            value={bitDepthIndex}
            onChange={(event) => setBitDepthIndex(Number(event.target.value))}
          />
          <div className="quantization-explorer__tick-row" aria-hidden="true">
            {BIT_DEPTHS.map((depth) => (
              <span key={depth}>{depth}</span>
            ))}
          </div>
        </div>
      </div>

      <div className="quantization-explorer__focus-grid quantization-explorer__focus-grid--four">
        <div className="quantization-explorer__focus-card">
          <span className="quantization-explorer__focus-label">
            Original weight
          </span>
          <strong className="quantization-explorer__focus-value">
            {formatFloat(focusWeight.original)}
          </strong>
        </div>

        <div className="quantization-explorer__focus-card">
          <span className="quantization-explorer__focus-label">
            {bitDepth === 16 ? "Stored BF16 word" : "Stored bucket"}
          </span>
          <strong className="quantization-explorer__focus-value">
            {String(focusWeight.stored)}
          </strong>
        </div>

        <div className="quantization-explorer__focus-card">
          <span className="quantization-explorer__focus-label">
            {bitDepth === 16 ? "Decoded float" : "Scaled back to float"}
          </span>
          <strong className="quantization-explorer__focus-value">
            {formatFloat(focusWeight.reconstructed)}
          </strong>
        </div>

        <div className="quantization-explorer__focus-card">
          <span className="quantization-explorer__focus-label">
            Absolute error
          </span>
          <strong className="quantization-explorer__focus-value">
            {formatFloat(Math.abs(focusWeight.error))}
          </strong>
        </div>
      </div>

      <p className="quantization-explorer__helper">
        {bitDepth === 16 ? (
          <>
            <strong>How to read this:</strong> BF16 stores a 16-bit word and
            then decodes it back into an approximate float.
          </>
        ) : (
          <>
            <strong>How to read this:</strong> the stored bucket is a small
            integer. We multiply it by the scale to get the approximate float.
          </>
        )}{" "}
        {getShortExplanation(bitDepth)}
      </p>

      <div className="quantization-explorer__formula">
        <span className="quantization-explorer__formula-label">
          {bitDepth === 16 ? "BF16 decoding" : "Bucket math"}
        </span>
        <code>
          {bitDepth === 16
            ? `${focusWeight.bf16Word} -> ${formatFloat(focusWeight.reconstructed)}`
            : `${focusWeight.stored} x ${formatFloat(scaleSummary?.scale ?? 0)} = ${formatFloat(focusWeight.reconstructed)}`}
        </code>
      </div>
    </div>
  );
}
