import { applyFilter } from "./fbp.js";

export async function generateSinogram(
  imageUrl,
  angles = 180,
  drawAngleCallback = null
) {
  const image = await loadImage(imageUrl);
  const size = Math.max(image.width, image.height);
  const projections = [];

  const canvas = Object.assign(document.createElement("canvas"), {
    width: size,
    height: size,
  });
  const ctx = canvas.getContext("2d");

  for (let angle = 0; angle < angles; angle++) {
    const theta = (angle * Math.PI) / angles;

    if (drawAngleCallback) drawAngleCallback(theta);
    // await new Promise((r) => setTimeout(r, 0.01)); // Removed for performance

    ctx.clearRect(0, 0, size, size);
    ctx.save();
    ctx.translate(size / 2, size / 2);
    ctx.rotate(theta);
    ctx.drawImage(image, -image.width / 2, -image.height / 2);
    ctx.restore();

    const { data } = ctx.getImageData(0, 0, size, size);

    const projection = [];
    for (let x = 0; x < size; x++) {
      let sum = 0;
      for (let y = 0; y < size; y++) {
        const i = (y * size + x) * 4;
        sum += data[i]; // grayscale from red channel
      }
      projection.push(sum / size);
    }
    projections.push(projection);
  }

  const isHorizontal = angles >= size;

  // Create rotated canvas accordingly
  const sinogramCanvas = Object.assign(document.createElement("canvas"), {
    width: isHorizontal ? size : angles,
    height: isHorizontal ? angles : size,
  });
  const sinCtx = sinogramCanvas.getContext("2d");
  const imgData = sinCtx.createImageData(
    sinogramCanvas.width,
    sinogramCanvas.height
  );

  for (let angle = 0; angle < angles; angle++) {
    for (let x = 0; x < size; x++) {
      const val = projections[angle][x];
      const px = isHorizontal ? x : angle;
      const py = isHorizontal ? angle : x;
      const i = (py * sinogramCanvas.width + px) * 4;
      imgData.data[i + 0] = val;
      imgData.data[i + 1] = val;
      imgData.data[i + 2] = val;
      imgData.data[i + 3] = 255;
    }
  }

  sinCtx.putImageData(imgData, 0, 0);
  return {
    sinogramUrl: sinogramCanvas.toDataURL(),
    projections: projections, // Return the raw projections
  };
}

export async function reconstructImageFromSinogram(
  sinogramUrl,
  size = 256,
  onFrame = null,
  renderMode = "heatmap",
  filterType = "ramp"
) {
  const sinogramImage = await loadImage(sinogramUrl);
  const canvas = Object.assign(document.createElement("canvas"), {
    width: sinogramImage.width,
    height: sinogramImage.height,
  });
  const ctx = canvas.getContext("2d");
  ctx.drawImage(sinogramImage, 0, 0);
  const sinogramData = ctx.getImageData(
    0,
    0,
    sinogramImage.width,
    sinogramImage.height
  ).data;

  // Detect orientation
  let width, angles;
  const isVertical = sinogramImage.height > sinogramImage.width;
  if (isVertical) {
    angles = sinogramImage.width;
    width = sinogramImage.height;
  } else {
    angles = sinogramImage.height;
    width = sinogramImage.width;
  }

  size = width;
  const outputCanvas = Object.assign(document.createElement("canvas"), {
    width: size,
    height: size,
  });
  const outputCtx = outputCanvas.getContext("2d");
  const accum = new Float32Array(size * size);
  const center = size / 2;

  const rawProjectionsFromSinogram = []; // To store raw projections extracted from sinogram image
  const allFilteredProjections = []; // To store filtered projections

  for (let angle = 0; angle < angles; angle++) {
    const theta = (angle * Math.PI) / angles;

    let currentRawProjection = [];
    for (let x = 0; x < width; x++) {
      const i = isVertical
        ? (x * angles + angle) * 4 // transposed layout
        : (angle * width + x) * 4; // normal layout

      currentRawProjection.push(sinogramData[i]);
    }
    rawProjectionsFromSinogram.push(currentRawProjection); // Store the raw projection

    // Apply the selected filter
    const filteredProjection = applyFilter(
      [...currentRawProjection],
      filterType
    ); // Apply filter to a copy
    allFilteredProjections.push(filteredProjection); // Store the filtered projection

    for (let y = 0; y < size; y++) {
      for (let x = 0; x < size; x++) {
        const x0 = x - center;
        const y0 = center - y;
        const s = Math.round(
          x0 * Math.cos(theta) + y0 * Math.sin(theta) + width / 2
        );
        if (s >= 0 && s < width) {
          accum[y * size + x] += filteredProjection[s]; // Use the filtered projection for accumulation
        }
      }
    }

    if (onFrame) {
      let maxVal = 0;
      for (let i = 0; i < accum.length; i++) {
        if (accum[i] > maxVal) maxVal = accum[i];
      }
      const imageData = outputCtx.createImageData(size, size);
      for (let i = 0; i < accum.length; i++) {
        let val = accum[i] / maxVal;
        val = Math.min(1, Math.max(0, val));
        let r, g, b;
        if (renderMode === "grayscale") {
          const gray = Math.round(val * 255);
          r = g = b = gray;
        } else {
          [r, g, b] = getHeatmapColor(val);
        }
        imageData.data[i * 4 + 0] = r;
        imageData.data[i * 4 + 1] = g;
        imageData.data[i * 4 + 2] = b;
        imageData.data[i * 4 + 3] = 255;
      }

      outputCtx.putImageData(imageData, 0, 0);
      // await new Promise((r) => setTimeout(r, 1)); // Removed for performance
      onFrame(angle, outputCanvas.toDataURL());
    }
  }

  // Create filtered sinogram image
  const filteredSinogramCanvas = Object.assign(
    document.createElement("canvas"),
    {
      width: isVertical ? angles : width, // Same dimensions as original sinogram
      height: isVertical ? width : angles,
    }
  );
  const filteredSinCtx = filteredSinogramCanvas.getContext("2d");
  const filteredImgData = filteredSinCtx.createImageData(
    filteredSinogramCanvas.width,
    filteredSinogramCanvas.height
  );

  let maxFilteredVal = 0;
  for (const proj of allFilteredProjections) {
    for (const val of proj) {
      if (val > maxFilteredVal) maxFilteredVal = val;
    }
  }
  // Avoid division by zero if maxFilteredVal is 0
  if (maxFilteredVal === 0) maxFilteredVal = 1;

  for (let angle = 0; angle < angles; angle++) {
    const currentProjection = allFilteredProjections[angle];
    for (let x = 0; x < width; x++) {
      const val = (currentProjection[x] / maxFilteredVal) * 255; // Normalize and scale
      const px = isVertical ? angle : x;
      const py = isVertical ? x : angle;
      const i = (py * filteredSinogramCanvas.width + px) * 4;
      filteredImgData.data[i + 0] = val;
      filteredImgData.data[i + 1] = val;
      filteredImgData.data[i + 2] = val;
      filteredImgData.data[i + 3] = 255;
    }
  }
  filteredSinCtx.putImageData(filteredImgData, 0, 0);
  const filteredSinogramUrl = filteredSinogramCanvas.toDataURL();

  return {
    reconstructedUrl: outputCanvas.toDataURL(),
    filteredSinogramUrl: filteredSinogramUrl,
    rawProjectionsFromSinogram: rawProjectionsFromSinogram,
    allFilteredProjections: allFilteredProjections,
  };
}

// Heatmap mapping: blue → green → yellow → red
function getHeatmapColor(value) {
  const r = Math.min(255, Math.max(0, 255 * Math.min(1, 4 * (value - 0.75))));
  const g = Math.min(255, Math.max(0, 255 * (4 * Math.abs(value - 0.5) - 1)));
  const b = Math.min(255, Math.max(0, 255 * (1 - 4 * value)));
  return [r, g, b];
}

function loadImage(src) {
  return new Promise((resolve) => {
    const img = new Image();
    img.crossOrigin = "anonymous";
    img.onload = () => resolve(img);
    img.src = src;
  });
}

export async function convertToGrayscale(imageUrl) {
  const image = await loadImage(imageUrl);
  const canvas = Object.assign(document.createElement("canvas"), {
    width: image.width,
    height: image.height,
  });
  const ctx = canvas.getContext("2d");

  // Draw original image
  ctx.drawImage(image, 0, 0);

  // Get pixel data
  const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
  const data = imageData.data;

  // Convert to grayscale: set R, G, B to luminance
  for (let i = 0; i < data.length; i += 4) {
    const r = data[i];
    const g = data[i + 1];
    const b = data[i + 2];
    const luminance = 0.299 * r + 0.587 * g + 0.114 * b;
    data[i] = data[i + 1] = data[i + 2] = luminance;
  }

  ctx.putImageData(imageData, 0, 0);
  return canvas.toDataURL();
}