import {
  generateSinogram,
  reconstructImageFromSinogram,
  convertToGrayscale,
} from "./sinogram.js";

import { StepAccordion } from "./CTAccordion .js";

export const UploadImageComponent = {
  hasLoadedInitialImage: false,
  angleCount: 40,
  imageUrl:
    "https://upload.wikimedia.org/wikipedia/commons/e/e5/Shepp_logan.png",
  sinogramUrl: null,
  filteredSinogramUrl: null, // Added to store the filtered sinogram
  reconstructedUrl: null,
  defaultImageUrl:
    "https://upload.wikimedia.org/wikipedia/commons/e/e5/Shepp_logan.png",
  reconstructionFrames: [],
  currentFrameIndex: 0,
  manualAngleIndex: 0, // For the new angle control slider
  sinogramImageElement: null, // To store the sinogram img DOM element
  sinogramHighlightCanvas: null, // Canvas for sinogram highlight
  renderMode: "grayscale", // or "heatmap"
  filterType: "ramp", // Added to manage selected filter

  drawAngleOverlay(theta) {
    const canvas = this.overlayCanvas;
    if (!canvas || !this.imageElement) return;

    const ctx = canvas.getContext("2d");
    const w = canvas.width;
    const h = canvas.height;
    const cx = w / 2;
    const cy = h / 2;
    const len = Math.max(w, h);

    ctx.clearRect(0, 0, w, h);

    const dx = len * Math.cos(theta);
    const dy = len * Math.sin(theta);

    ctx.strokeStyle = "rgba(255,0,0,0.8)";
    ctx.lineWidth = 2;

    ctx.beginPath();
    ctx.moveTo(cx - dx, cy - dy);
    ctx.lineTo(cx + dx, cy + dy);
    ctx.stroke();
  },

  drawSinogramHighlight(angleIndex) {
    if (
      !this.sinogramHighlightCanvas ||
      !this.sinogramImageElement ||
      !this.sinogramImageElement.complete ||
      this.sinogramImageElement.naturalWidth === 0
    ) {
      if (this.sinogramHighlightCanvas) {
        const ctx = this.sinogramHighlightCanvas.getContext("2d");
        ctx.clearRect(
          0,
          0,
          this.sinogramHighlightCanvas.width,
          this.sinogramHighlightCanvas.height
        );
      }
      return;
    }

    const canvas = this.sinogramHighlightCanvas;
    const ctx = canvas.getContext("2d");
    const dispWidth = this.sinogramImageElement.width; // Displayed width of sinogram image
    const dispHeight = this.sinogramImageElement.height; // Displayed height of sinogram image
    const naturalW = this.sinogramImageElement.naturalWidth;
    const naturalH = this.sinogramImageElement.naturalHeight;

    if (canvas.width !== dispWidth) canvas.width = dispWidth;
    if (canvas.height !== dispHeight) canvas.height = dispHeight;

    ctx.clearRect(0, 0, dispWidth, dispHeight);

    if (this.angleCount <= 0 || naturalW === 0 || naturalH === 0) return;

    let anglesAlongImageHeight = true;

    // Infer orientation based on natural dimensions of the sinogram image vs angleCount
    // This logic assumes that one of the dimensions of the *natural* sinogram image data
    // (before it's scaled for display) corresponds to this.angleCount.
    if (naturalW === this.angleCount && naturalH !== this.angleCount) {
      anglesAlongImageHeight = false; // Angles are columns, highlight is vertical
    } else if (naturalH === this.angleCount && naturalW !== this.angleCount) {
      anglesAlongImageHeight = true; // Angles are rows, highlight is horizontal
    } else if (naturalH === this.angleCount && naturalW === this.angleCount) {
      // Square sinogram (angleCount might equal detector_size), assume angles along height by convention
      anglesAlongImageHeight = true;
    }
    // If neither naturalW nor naturalH strictly equals angleCount (and the other doesn't),
    // we stick with the default assumption (anglesAlongImageHeight = true).
    // This can happen if the sinogram image itself is generated with an aspect ratio
    // where neither dimension is exactly `this.angleCount` (e.g. if `generateSinogram` in sinogram.js
    // has `size` not equal to `angles` and the output canvas is `size` x `angles` or `angles` x `size`).
    // The key is that `this.angleCount` represents the number of projection angles.

    ctx.strokeStyle = "rgba(0, 255, 0, 0.7)"; // Green highlight

    if (anglesAlongImageHeight) {
      // Angles are effectively rows in the displayed sinogram image.
      // The number of such rows is this.angleCount.
      // We need to highlight the row corresponding to `angleIndex`.
      const highlightRowHeight = dispHeight / this.angleCount;
      const yPos = (angleIndex + 0.5) * highlightRowHeight;

      ctx.lineWidth = Math.max(1, highlightRowHeight * 0.7);
      ctx.beginPath();
      ctx.moveTo(0, yPos);
      ctx.lineTo(dispWidth, yPos);
      ctx.stroke();
    } else {
      // Angles are effectively columns in the displayed sinogram image.
      // The number of such columns is this.angleCount.
      const highlightColWidth = dispWidth / this.angleCount;
      const xPos = (angleIndex + 0.5) * highlightColWidth;

      ctx.lineWidth = Math.max(1, highlightColWidth * 0.7);
      ctx.beginPath();
      ctx.moveTo(xPos, 0);
      ctx.lineTo(xPos, dispHeight);
      ctx.stroke();
    }
  },

  isOverlayReady() {
    return (
      this.overlayCanvas &&
      this.imageElement &&
      this.imageElement.complete &&
      this.imageElement.naturalWidth > 0
    );
  },

  async loadAndProcess(url, isUploaded = false) {
    this.imageUrl = url;
    this.sinogramUrl = "loading";
    this.filteredSinogramUrl = null; // Reset filtered sinogram
    this.reconstructedUrl = null;
    this.manualAngleIndex = 0; // Reset manual angle on new image
    m.redraw();

    this.imageUrl = url;
    let finalUrl = url;

    if (isUploaded) {
      finalUrl = await convertToGrayscale(url);
      this.imageUrl = finalUrl;
    }

    this.sinogramUrl = await generateSinogram(
      finalUrl,
      this.angleCount,
      this.drawAngleOverlay.bind(this)
    );
    m.redraw();

    this.reconstructionFrames = [];
    this.currentFrameIndex = 0;

    const reconstructionResult = await reconstructImageFromSinogram(
      this.sinogramUrl,
      undefined,
      (angle, frameUrl) => {
        this.reconstructionFrames.push(frameUrl);
        this.currentFrameIndex = this.reconstructionFrames.length - 1;
        this.reconstructedUrl = frameUrl; // This is the final reconstructed image frame
        m.redraw();
      },
      this.renderMode,
      this.filterType
    );

    // After all frames are processed, the final reconstructedUrl and filteredSinogramUrl are available
    this.reconstructedUrl = reconstructionResult.reconstructedUrl;
    this.filteredSinogramUrl = reconstructionResult.filteredSinogramUrl;
    m.redraw(); // Ensure UI updates with the final urls
  },

  oninit() {
    this.loadAndProcessDebounced = debounce((url) => {
      this.loadAndProcess(url);
    }, 300);
  },

  view() {
    return m(
      "div",
      {
        class: "flex flex-col items-center min-h-screen bg-gray-100 py-10 px-4",
      },
      [
        // Header
        m("header", { class: "mb-10 text-center" }, [
          m(
            "h1",
            { class: "text-4xl font-bold text-gray-800 mb-2" },
            "Sinogram Generator"
          ),
          m(
            "p",
            { class: "text-gray-600 text-lg" },
            "Upload a grayscale image to simulate CT scan projections"
          ),
        ]),

        m(StepAccordion, { index: 0 }),

        // Upload Box
        m(
          "div",
          {
            class:
              "w-full max-w-lg border-4 border-dashed border-gray-400 bg-white rounded-xl p-6 text-center hover:bg-gray-50 cursor-pointer transition",
            ondragover: (e) => e.preventDefault(),
            ondrop: (e) => {
              e.preventDefault();
              const file = e.dataTransfer.files[0];
              if (file && file.type.startsWith("image/")) {
                const url = URL.createObjectURL(file);
                this.loadAndProcess(url, true);
              }
            },
            onclick: () => document.getElementById("fileInput").click(),
          },
          [
            m(
              "p",
              { class: "text-gray-500" },
              "Click or drag a grayscale image here"
            ),
            m("input", {
              id: "fileInput",
              type: "file",
              class: "hidden",
              accept: "image/*",
              onchange: (e) => {
                const file = e.target.files[0];
                if (file && file.type.startsWith("image/")) {
                  const url = URL.createObjectURL(file);
                  this.loadAndProcess(url);
                }
              },
            }),
          ]
        ),

        m(StepAccordion, { index: 1 }),

        // Image Preview
        m("div", { class: "relative mt-6 w-full max-w-md" }, [
          m("img", {
            src: this.imageUrl,
            class: "rounded shadow max-w-full h-auto mx-auto",
            onload: (e) => {
              this.imageElement = e.target;

              // Only start once both image and canvas are ready
              if (this.isOverlayReady() && !this.hasLoadedInitialImage) {
                this.hasLoadedInitialImage = true;
                this.loadAndProcess(this.imageUrl);
              }
            },
          }),
          m("canvas", {
            width: this.imageElement?.width || 0,
            height: this.imageElement?.height || 0,
            style: "position:absolute; top:0; left:0; pointer-events:none;",
            oncreate: ({ dom }) => {
              this.overlayCanvas = dom;

              // Trigger load if image was already ready
              if (this.isOverlayReady() && !this.hasLoadedInitialImage) {
                this.hasLoadedInitialImage = true;
                this.loadAndProcess(this.imageUrl);
              }
            },
          }),
        ]),

        // Angle Slider
        m("div", { class: "mt-6 w-full max-w-md" }, [
          m(
            "label",
            { class: "block text-sm font-medium text-gray-700 mb-1" },
            `Number of angles: ${this.angleCount}`
          ),
          m("input", {
            type: "range",
            min: 5,
            max: 360,
            value: this.angleCount,
            step: 1,
            class: "w-full",
            oninput: (e) => {
              this.angleCount = parseInt(e.target.value, 10);
              this.manualAngleIndex = 0; // Reset manual angle if angleCount changes
              // Redraw overlay for new manualAngleIndex=0 with new angleCount
              const newTheta = (0 * Math.PI) / this.angleCount; // Theta for index 0
              this.drawAngleOverlay(newTheta);
              if (this.sinogramHighlightCanvas) {
                // If sinogram highlight exists, update it too
                this.drawSinogramHighlight(this.manualAngleIndex);
              }
              this.loadAndProcessDebounced(this.imageUrl); // reprocess with new angle count
            },
          }),
        ]),

        // Manual Angle Slider for Preview
        m("div", { class: "mt-4 w-full max-w-md" }, [
          m(
            "label",
            { class: "block text-sm font-medium text-gray-700 mb-1" },
            `Preview Angle: ${this.manualAngleIndex} (θ ≈ ${(
              (this.manualAngleIndex * 180) /
              this.angleCount
            ).toFixed(1)}°)`
          ),
          m("input", {
            type: "range",
            min: 0,
            max: Math.max(0, this.angleCount - 1), // Ensure max is not negative if angleCount is 0 or 1
            value: this.manualAngleIndex,
            step: 1,
            class: "w-full",
            disabled: !this.hasLoadedInitialImage, // Disable if no image loaded
            oninput: (e) => {
              this.manualAngleIndex = parseInt(e.target.value, 10);
              const theta = (this.manualAngleIndex * Math.PI) / this.angleCount;
              this.drawAngleOverlay(theta);
              this.drawSinogramHighlight(this.manualAngleIndex); // Update sinogram highlight
            },
          }),
        ]),

        m(StepAccordion, { index: 2 }),

        // Sinogram
        this.sinogramUrl &&
          m("div", { class: "mt-10 w-full max-w-md text-center relative" }, [
            m(
              "h2",
              { class: "text-xl font-semibold text-gray-700 mb-4" },
              "Generated Sinogram"
            ),
            m("div", { class: "relative" }, [
              this.sinogramUrl === "loading"
                ? m("p", "Processing...")
                : m("img", {
                    src: this.sinogramUrl,
                    alt: "Sinogram",
                    class: "rounded shadow max-w-full h-auto mx-auto",
                    style: "width: 100%; max-height: 300px;",
                    onload: ({ target }) => {
                      this.sinogramImageElement = target;
                      // Set initial dimensions for highlight canvas based on loaded image
                      if (this.sinogramHighlightCanvas) {
                        this.sinogramHighlightCanvas.width = target.width;
                        this.sinogramHighlightCanvas.height = target.height;
                      }
                      this.drawSinogramHighlight(this.manualAngleIndex); // Draw initial highlight
                    },
                    onerror: () => {
                      // If sinogram fails to load, clear element
                      this.sinogramImageElement = null;
                      this.drawSinogramHighlight(this.manualAngleIndex); // Attempt to clear highlight
                    },
                  }),
              // Sinogram Highlight Canvas
              m("canvas", {
                style:
                  "position:absolute; top:0; left:0; pointer-events:none; max-width: 100%; max-height: 300px;",
                oncreate: ({ dom }) => {
                  this.sinogramHighlightCanvas = dom;
                  // Ensure dimensions are set if image is already loaded by now (race condition)
                  if (
                    this.sinogramImageElement &&
                    this.sinogramImageElement.complete
                  ) {
                    dom.width = this.sinogramImageElement.width;
                    dom.height = this.sinogramImageElement.height;
                    this.drawSinogramHighlight(this.manualAngleIndex);
                  }
                },
                // Set width/height to 0 initially or based on a placeholder if needed
                width: this.sinogramImageElement
                  ? this.sinogramImageElement.width
                  : this.overlayCanvas
                  ? this.overlayCanvas.width
                  : 300, // Default or match other canvas
                height: this.sinogramImageElement
                  ? this.sinogramImageElement.height
                  : this.overlayCanvas
                  ? this.overlayCanvas.height
                  : 300,
              }),
            ]),
          ]),

        // Filtered Sinogram Display
        this.filteredSinogramUrl &&
          m("div", { class: "mt-10 w-full max-w-md text-center" }, [
            m(
              "h2",
              { class: "text-xl font-semibold text-gray-700 mb-4" },
              "Filtered Sinogram"
            ),
            m("img", {
              src: this.filteredSinogramUrl,
              alt: "Filtered Sinogram",
              class: "rounded shadow max-w-full h-auto mx-auto",
              style: "width: 100%; max-height: 300px;",
            }),
          ]),

        m(StepAccordion, { index: 3 }),

        // Reconstructed
        this.reconstructedUrl &&
          m("div", { class: "mt-10 w-full max-w-md text-center" }, [
            m(
              "h2",
              { class: "text-xl font-semibold text-gray-700 mb-4" },
              "Reconstructed Image (Back Projection)"
            ),
            m("img", {
              src: this.reconstructionFrames[this.currentFrameIndex],
              alt: "Reconstructed",
              class: "rounded shadow max-w-full h-auto mx-auto",
            }),
            m("div", { class: "mt-6 w-full max-w-md text-center" }, [
              m(
                "label",
                { class: "text-sm text-gray-600 mr-2" },
                "Render style:"
              ),
              m(
                "select",
                {
                  value: this.renderMode,
                  onchange: (e) => {
                    this.renderMode = e.target.value;
                    this.loadAndProcess(this.imageUrl); // re-render using selected mode
                  },
                },
                [
                  m("option", { value: "heatmap" }, "Heatmap"),
                  m("option", { value: "grayscale" }, "Grayscale"),
                ]
              ),
            ]),

            m("div", { class: "mt-4 w-full max-w-md text-left" }, [
              m(
                "label",
                { class: "text-sm text-gray-600 mr-2" },
                "Filter Type:"
              ),
              m(
                "select",
                {
                  class: "p-1 border rounded text-sm",
                  value: this.filterType,
                  onchange: (e) => {
                    this.filterType = e.target.value;
                    this.loadAndProcess(this.imageUrl);
                  },
                },
                [
                  m("option", { value: "ramp" }, "Ramp"),
                  m("option", { value: "shepp-logan" }, "Shepp-Logan"),
                  m("option", { value: "hann" }, "Hann"),
                  m("option", { value: "hamming" }, "Hamming"),
                  m("option", { value: "cosine" }, "Cosine"),
                  m("option", { value: "none" }, "None"),
                ]
              ),
            ]),

            this.reconstructionFrames.length > 1 &&
              m("div", { class: "mt-4" }, [
                m("input", {
                  type: "range",
                  min: 0,
                  max: this.reconstructionFrames.length - 1,
                  value: this.currentFrameIndex,
                  step: 1,
                  oninput: (e) => {
                    this.currentFrameIndex = +e.target.value;
                    this.reconstructedUrl =
                      this.reconstructionFrames[this.currentFrameIndex];
                  },
                }),
                m(
                  "p",
                  { class: "text-sm text-gray-500 mt-1" },
                  `Angle ${this.currentFrameIndex + 1} / ${
                    this.reconstructionFrames.length
                  }`
                ),
              ]),
            m(StepAccordion, { index: 4 }),
          ]),
      ]
    );
  },
};

function debounce(fn, delay) {
  let timeout;
  return (...args) => {
    clearTimeout(timeout);
    timeout = setTimeout(() => fn(...args), delay);
  };
}