import { useEffect, useState, useRef } from "react";
import * as tf from "@tensorflow/tfjs";
import Plot from "react-plotly.js";

import { MnistData } from "./mnist.js";
import { Cnn, train } from "./train.ts";
import type { TrainController, RunInfo, OptimizerParams } from "./train.ts";
import { Button, InputField, Tabs, Dropdown, Card } from "@elvis/ui";
import InfoViewer from "./InfoViewer.tsx";

const DEFAULT_ARCHITECTURE = `[conv2d filters=8 kernel=11 
stride=1 padding=1 activation=relu]

[maxpool size=2 stride=2]

[flatten]

[dense units=10 activation=softmax]`;

const isFirefox = navigator.userAgent.toLowerCase().includes("firefox");
await tf.setBackend(isFirefox ? "cpu" : "webgl");
await tf.ready();

export default function NetworkVisualizer() {
  const [dataset, setDataset] = useState<MnistData | null>(null);

  useEffect(() => {
    const loadData = async () => {
      const data = new MnistData();
      await data.load();
      setDataset(data);
      console.log("dataset loaded");
    }

    loadData();
  }, []);

  // architecture states
  const [architecture, setArchitecture] = useState(DEFAULT_ARCHITECTURE);
  const [optimizerType, setOptimizerType] = useState<string>('adam');
  const [optimizerParams, setOptimizerParams] = useState<OptimizerParams>({
    learningRate: '0.001',
    beta1: '0.9',
    beta2: '0.999',
    epsilon: '1e-8',
    batchSize: '32',
    epochs: '5',
  });
  const modelRef = useRef<Cnn | null>(null);
  const optimizerRef = useRef<tf.Optimizer | null>(null);

  function handleArchitectureChange(newArchitecture: string) {
    if (isTraining) {
      alert('Cannot change architecture while training is in progress.');
    } else {
      setArchitecture(newArchitecture);
    }
  }

  function handleOptimizerChange(newOptimizerType: string, newOptimizerParams: OptimizerParams) {
    if (isTraining) {
      alert('Cannot change optimizer settings while training is in progress.');
    } else {
      setOptimizerType(newOptimizerType);
      setOptimizerParams(newOptimizerParams);
    }
  }

  function handleSampleIndexChange() {
    trainController.current.sampleIndex += 1;
    refreshSampleInfo();
  }

 
  function resetModel() {
    if (!dataset) return;

    if (modelRef.current) {
      modelRef.current.dispose();
    }

    const cnn = new Cnn(architecture, dataset.numInputChannels);
    modelRef.current = cnn;
  }

  function resetOptimizer() {
    if (optimizerType === 'adam') {
      const learningRate = parseFloat(optimizerParams.learningRate);
      const beta1 = parseFloat(optimizerParams.beta1 || "0.9");
      const beta2 = parseFloat(optimizerParams.beta2 || "0.999");
      const epsilon = parseFloat(optimizerParams.epsilon || "1e-8");

      if (Number.isNaN(learningRate) || learningRate <= 0) {
        alert('Invalid learning rate for Adam optimizer.');
        return;
      }
      if (Number.isNaN(beta1) || beta1 < 0) {
        alert('Invalid beta1 for Adam optimizer.');
        return;
      }
      if (Number.isNaN(beta2) || beta2 < 0) {
        alert('Invalid beta2 for Adam optimizer.');
        return;
      }
      if (Number.isNaN(epsilon) || epsilon <= 0) {
        alert('Invalid epsilon for Adam optimizer.');
        return;
      }

      const opt = tf.train.adam(learningRate, beta1, beta2, epsilon);
      if (optimizerRef.current) {
        optimizerRef.current.dispose();
      }
      optimizerRef.current = opt;
    } else if (optimizerType === 'sgd') {
      const learningRate = parseFloat(optimizerParams.learningRate);
      if (Number.isNaN(learningRate) || learningRate <= 0) {
        alert('Invalid learning rate for SGD optimizer.');
        return;
      }

      const opt = tf.train.sgd(learningRate);
      if (optimizerRef.current) {
        optimizerRef.current.dispose();
      }
      optimizerRef.current = opt;
    } else {
      alert(`Unsupported optimizer type: ${optimizerType}`);
    }
  }

  // Reset the model only when its architecture or input shape changes.
  useEffect(() => {
    resetModel();
  }, [architecture, dataset]);

  // Optimizer settings can change without discarding the trained model.
  useEffect(() => {
    resetOptimizer();
  }, [optimizerType, optimizerParams]);

  // training states
  const [isTraining, setIsTraining] = useState<boolean>(false);
  const isTrainingRef = useRef<boolean>(false);
  const lossesRef = useRef<Array<number>>([]);
  const trainController = useRef<TrainController>({
    isPaused: false,
    stopRequested: false,
    sampleIndex: 0,
  });
  const infoRef = useRef<Array<RunInfo>>([]);

  function refreshSampleInfo() {
    if (!dataset || !modelRef.current) return;

    let nextInfo: RunInfo[] = [];
    const cleanupToken = tf.tidy(() => {
      const sample = dataset.getTestSample(trainController.current.sampleIndex);
      const { output, info } = modelRef.current!.forwardWithInfo(
        sample.xs.reshape([
          1,
          dataset.imageSize,
          dataset.imageSize,
          dataset.numInputChannels,
        ]) as tf.Tensor4D,
      );
      const probs = tf.softmax(output);
      info.push({
        type: 'output',
        output: probs.dataSync(),
        shape: probs.shape,
      });
      nextInfo = info;
      return tf.scalar(0);
    });
    cleanupToken.dispose();
    infoRef.current = nextInfo;
    updateTick();
  }

  useEffect(() => {
    refreshSampleInfo();
  }, [architecture, dataset]);

  // render timing
  const [, setTick] = useState<number>(0);

  function updateTick() {
    setTick((tick) => tick + 1);
  }

  async function startTraining() {
    if (!modelRef.current || !dataset || !optimizerRef.current || isTrainingRef.current) {
      return;
    }

    isTrainingRef.current = true;
    setIsTraining(true);
    trainController.current.isPaused = false;
    trainController.current.stopRequested = false;

    const batchSize = parseFloat(optimizerParams.batchSize);
    if (Number.isNaN(batchSize) || batchSize <= 0) {
      alert('Invalid batch size.');
      isTrainingRef.current = false;
      setIsTraining(false);
      return;
    }
    const epochs = parseFloat(optimizerParams.epochs);
    if (Number.isNaN(epochs) || epochs <= 0) {
      alert('Invalid number of epochs.');
      isTrainingRef.current = false;
      setIsTraining(false);
      return;
    }

    let lastTickUpdate = 0;
    if (!modelRef.current) return;
    if (!optimizerRef.current) return;

    try {
      await train(
        dataset,
        modelRef.current,
        optimizerRef.current,
        batchSize,
        epochs,
        trainController.current,
        (_epoch, _batch, loss, info) => {
          // lossesRef.current.push({ epoch, batch, loss });
          lossesRef.current.push(loss);
          console.log(loss);

          infoRef.current = info;

          // update tick every 50ms
          const now = performance.now();
          if (now - lastTickUpdate > 50) {
            lastTickUpdate = now;
            updateTick();
          }
        },
      );
    } finally {
      const wasStopped = trainController.current.stopRequested;
      isTrainingRef.current = false;
      setIsTraining(false);
      trainController.current.isPaused = false;
      trainController.current.stopRequested = false;
      if (!wasStopped) {
        alert('Training finished.');
      }
    }
  }

  function handleTrain() {
    if (isTrainingRef.current) {
      trainController.current.isPaused = false;
      return;
    }
    console.log('Starting training...');
    startTraining();
  }

  function handlePauseTraining() {
    console.log('Pausing training...');
    trainController.current.isPaused = true;
  }

  function handleStopTraining() {
    console.log('Stopping training...');
    trainController.current.stopRequested = true;
    trainController.current.isPaused = false;
  }

  async function waitUntilNotTraining() {
    return new Promise<void>((resolve) => {
      function check() {
        if (!isTrainingRef.current) {
          resolve();
        } else {
          requestAnimationFrame(check);
        }
      }
      check();
    });
  }

  async function handleResetTraining() {
    console.log('Resetting training...');
    handleStopTraining();

    await waitUntilNotTraining();
    console.log('Training stopped. Resetting model.');
    lossesRef.current = [];
    infoRef.current = [];
    resetModel();
    resetOptimizer();
    refreshSampleInfo();
  }

  return (
    <div className="flex flex-col min-h-0 h-full gap-3">
      <p className="text-sm text-gray-700">
        Configure and train a convolutional neural network (CNN) to classify handwritten digits from the MNIST dataset.
      </p>
      <div className="grid grid-cols-[2fr_1fr] min-h-0 flex-1 gap-12">
        <TrainingViewer 
          infoRef={infoRef}
          handleSampleIndexChange={handleSampleIndexChange}
        />
        <Sidebar 
          architecture={architecture}
          onArchitectureChange={handleArchitectureChange}

          optimizerType={optimizerType}
          optimizerParams={optimizerParams}
          onOptimizerChange={handleOptimizerChange}
          isTraining={isTraining}
          lossesRef={lossesRef}

          onTrain={handleTrain}
          onPauseTraining={handlePauseTraining}
          onResetTraining={handleResetTraining}
        />
      </div>
    </div>
  );
}

interface TrainingViewerProps {
  infoRef: React.RefObject<Array<RunInfo>>;
  handleSampleIndexChange: () => void;
}
function TrainingViewer({
  infoRef,
  handleSampleIndexChange,
}: TrainingViewerProps) {
  return (
    <div className="flex flex-col h-full overflow-auto gap-4 w-full min-h-0">
      <InfoViewer info={infoRef.current} onSampleIndexChange={handleSampleIndexChange} />
    </div>
  )
}

interface SidebarProps {
  architecture: string;
  onArchitectureChange: (newArchitecture: string) => void;

  optimizerType: string;
  optimizerParams: OptimizerParams;
  onOptimizerChange: (newOptimizerType: string, newOptimizerParams: OptimizerParams) => void;
  isTraining: boolean;
  lossesRef: React.RefObject<Array<number>>;

  onTrain: () => void;
  onPauseTraining: () => void;
  onResetTraining: () => void;
}

type TrainingInputRowProps = {
  label: string;
  value: string | undefined;
  onChange: (value: string) => void;
};

function TrainingInputRow({ label, value, onChange }: TrainingInputRowProps) {
  return (
    <div className="grid grid-cols-[7rem_minmax(0,1fr)] items-center gap-2">
      <span className="text-sm text-right text-gray-700">{label}</span>
      <div className="[&_label]:hidden [&_input]:py-1">
        <InputField label="" value={value ?? ""} onChange={onChange} />
      </div>
    </div>
  );
}

function Sidebar({
  architecture,
  onArchitectureChange,
  optimizerType,
  optimizerParams,
  onOptimizerChange,
  isTraining,
  lossesRef,
  onTrain,
  onPauseTraining,
  onResetTraining,
}: SidebarProps) {
  const tabs = ["Architecture", "Train"];
  const [activeTab, setActiveTab] = useState<string>(tabs[0]);

  const [architectureDraft, setArchitectureDraft] = useState<string>(architecture);

  return (
    <Card className="flex flex-col h-full p-4 gap-2 overflow-auto">
      <Tabs tabs={tabs} activeTab={activeTab} onChange={setActiveTab} />
      <div className="flex flex-col p-4 gap-4 h-full overflow-auto">
        { isFirefox && (
          <p className="text-red-500">
            Warning: This demo may be quite slow on Firefox.
          </p>
        )}

        { activeTab === "Architecture" && (
          <>
            <InputField 
              label="Architecture" 
              value={architectureDraft}
              onChange={setArchitectureDraft}
              rows={15}
            />
            <Button
              label="Apply architecture"
              onClick={() => onArchitectureChange(architectureDraft)}
            />
          </>
        )}

        { activeTab === "Train" && (
          <div className="flex flex-col gap-2">
            <p className="text-sm">Training {isTraining ? "in progress" : "not in progress"}</p>
            <Plot
              data={[
                {
                  x: lossesRef.current.map((_, i) => i),
                  y: lossesRef.current,
                  mode: 'lines',
                  type: 'scatter',
                },
              ]}
              layout={{
                xaxis: { title: { text: 'Training steps' } },
                yaxis: { title: { text: 'Train loss' } },
                margin: { t: 15, r: 15, b: 40, l: 45 },
              }}
              className="w-full h-[220px]"
              config={{ responsive: true }}
            />
            <div className="grid grid-cols-[7rem_minmax(0,1fr)] items-center gap-2">
              <span className="text-sm text-right text-gray-700">Optimizer</span>
              <div className="[&_label]:hidden [&_select]:py-1">
                <Dropdown
                  label=""
                  options={["sgd", "adam"]}
                  activeOption={optimizerType}
                  onChange={(newOptimizerType) => onOptimizerChange(newOptimizerType, optimizerParams)}
                />
              </div>
            </div>
            { optimizerType === 'sgd' && (
              <TrainingInputRow
                label="Learning Rate"
                value={optimizerParams.learningRate}
                onChange={(newLearningRate) => onOptimizerChange(optimizerType, {...optimizerParams, learningRate: newLearningRate})}
              />
            )}
            { optimizerType === 'adam' && (
              <>
                <TrainingInputRow
                  label="Learning Rate"
                  value={optimizerParams.learningRate}
                  onChange={(newLearningRate) => onOptimizerChange(optimizerType, {...optimizerParams, learningRate: newLearningRate})}
                />
                <TrainingInputRow
                  label="Beta 1"
                  value={optimizerParams.beta1}
                  onChange={(newBeta1) => onOptimizerChange(optimizerType, {...optimizerParams, beta1: newBeta1})}
                />
                <TrainingInputRow
                  label="Beta 2"
                  value={optimizerParams.beta2}
                  onChange={(newBeta2) => onOptimizerChange(optimizerType, {...optimizerParams, beta2: newBeta2})}
                />
                <TrainingInputRow
                  label="Epsilon"
                  value={optimizerParams.epsilon}
                  onChange={(newEpsilon) => onOptimizerChange(optimizerType, {...optimizerParams, epsilon: newEpsilon})}
                />
              </>
            )}
            <TrainingInputRow
              label="Batch Size"
              value={optimizerParams.batchSize}
              onChange={(newBatchSize) => onOptimizerChange(optimizerType, {...optimizerParams, batchSize: newBatchSize})}
            />
            <TrainingInputRow
              label="Epochs"
              value={optimizerParams.epochs}
              onChange={(newEpochs) => onOptimizerChange(optimizerType, {...optimizerParams, epochs: newEpochs})}
            />
            <div className="grid grid-cols-3 gap-2 mt-2 [&_button]:px-2 [&_button]:py-1 [&_button]:text-xs">
              <Button label="Train" onClick={onTrain} />
              <Button label="Pause" onClick={onPauseTraining} />
              <Button label="Reset" onClick={onResetTraining} />
            </div>
          </div>
        )}
      </div>
    </Card>
  )
}