import React, { useState, useEffect, useRef } from 'react'
import * as d3 from 'd3'
import { useTheme } from '../context/themeContext'

import MODELS from '../utils/models'
import DEVICES from '../utils/devices'

type Precision = 'fp32' | 'fp16' | 'int8' | 'int4'

interface ModelSizeBarChartProps {
  modelSize: number // in GB
  largestModelSize: number // largest model in full precision (fp32)
  modelPrecision: Precision // enum of fp32, fp16, int8, int4
  deviceMemorySet: boolean // true if device memory is set
}

interface InferenceRuntimeLineChartProps {
  availableMemory: AvailableMemory // in GB
  memoryPerInput: number // in GB
}

interface LineChartData {
  seqLength: number
  batchSize: number
}

interface AvailableMemory {
  int4: number
  int8: number
  fp16: number
  fp32: number
}

// Table containing the mapping of backends to precisions
const BackendPrecisionTable = () => {
  return (
    <table className='table-auto border-collapse'>
      <thead>
        <tr>
          <th className='table-cell px-4 py-2'>Backend</th>
          <th className='table-cell px-4 py-2'>GPU</th>
          <th className='table-cell px-4 py-2'>CPU</th>
          <th className='table-cell px-4 py-2'>Accuracy</th>
        </tr>
      </thead>
      <tbody>
        <tr>
          <td className='table-cell px-4 py-2'>fast</td>
          <td className='table-cell px-4 py-2'>16</td>
          <td className='table-cell px-4 py-2'>16</td>
          <td className='table-cell px-4 py-2'>⭐⭐⭐</td>
        </tr>
        <tr>
          <td className='table-cell px-4 py-2'>compress-fast</td>
          <td className='table-cell px-4 py-2'>4</td>
          <td className='table-cell px-4 py-2'>8</td>
          <td className='table-cell px-4 py-2'>⭐⭐</td>
        </tr>
        <tr>
          <td className='table-cell px-4 py-2'>compress</td>
          <td className='table-cell px-4 py-2'>4</td>
          <td className='table-cell px-4 py-2'>4</td>
          <td className='table-cell px-4 py-2'>⭐</td>
        </tr>
        <tr>
          <td className='table-cell px-4 py-2'>baseline</td>
          <td className='table-cell px-4 py-2'>16</td>
          <td className='table-cell px-4 py-2'>16</td>
          <td className='table-cell px-4 py-2'>⭐⭐⭐</td>
        </tr>
      </tbody>
    </table>
  )
}

// Bar chart for model footprint (standard version)
function ModelSizeBarChart({
  modelSize,
  largestModelSize,
  modelPrecision,
  deviceMemorySet,
}: ModelSizeBarChartProps) {
  const { theme } = useTheme()
  const chartRef = useRef<SVGSVGElement>(null)

  const width = 600
  const height = 50

  useEffect(() => {
    d3.select(chartRef.current).selectAll('*').remove()

    const svg = d3.select(chartRef.current).attr('width', width).attr('height', height)

    const xScale = d3.scaleLinear().domain([0, largestModelSize]).range([0, width])

    if (modelSize > largestModelSize) {
      svg
        .append('rect')
        .attr('x', 0)
        .attr('y', 0)
        .attr('width', width)
        .attr('height', height)
        .attr('fill', 'transparent')
        .style('stroke', theme === 'dark' ? '#f9fafb' : '#181f26')
        .style('stroke-dasharray', '4, 4')
        .style('stroke-width', '2px')
      svg
        .append('text')
        .attr('x', width / 2)
        .attr('y', height / 2)
        .attr('text-anchor', 'middle')
        .attr('alignment-baseline', 'middle')
        .attr('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
        .text('Out of Memory')
    } else {
      svg
        .append('rect')
        .attr('x', 0)
        .attr('y', 0)
        .attr('width', xScale(modelSize))
        .attr('height', height)
        .attr('fill', chooseColor(modelPrecision))

      if (deviceMemorySet) {
        svg
          .append('rect')
          .attr('x', xScale(modelSize))
          .attr('y', 0)
          .attr('width', xScale(largestModelSize - modelSize))
          .attr('height', height)
          .attr('fill', 'transparent')
          .style('stroke', chooseColor(modelPrecision))
          .style('stroke-width', '2px')
      }
    }
  }, [modelSize, largestModelSize, modelPrecision, deviceMemorySet, theme])

  function chooseColor(precision: Precision) {
    const colors = {
      fp32: '#e45f5b',
      fp16: '#ffc068',
      int8: '#71cce9',
      int4: '#383d95',
    }
    return colors[precision] || 'gray'
  }

  return <svg ref={chartRef}></svg>
}

// Line chart for standard inference runtime
function InferenceRuntimeLineChart({
  availableMemory,
  memoryPerInput,
}: InferenceRuntimeLineChartProps) {
  const { theme } = useTheme()
  const chartRef = useRef(null)
  const maxSeqLength = 4096
  const maxBatchSize = 128

  useEffect(() => {
    const margin = { top: 20, right: 20, bottom: 50, left: 50 }
    const width = 600 - margin.left - margin.right
    const height = 400 - margin.top - margin.bottom
    const precisions = [
      { name: 'FP32', color: '#e45f5b' },
      { name: 'FP16', color: '#ffc068' },
      { name: 'INT8', color: '#71cce9' },
      { name: 'INT4', color: '#383d95' },
    ]

    const svg = d3.select(chartRef.current)
    svg.selectAll('*').remove()

    const xScale = d3.scaleLinear().domain([0, maxSeqLength]).range([0, width])

    const yScale = d3.scaleLinear().domain([0, maxBatchSize]).range([height, 0])

    const xAxis = d3.axisBottom(xScale)
    const yAxis = d3.axisLeft(yScale)
    svg
      .attr('width', width + margin.left + margin.right)
      .attr('height', height + margin.top + margin.bottom)
      .append('g')
      .attr('transform', `translate(${margin.left}, ${margin.top})`)

    svg
      .append('g')
      .attr('transform', `translate(${margin.left}, ${height + margin.top})`)
      .call(xAxis)

    svg.append('g').attr('transform', `translate(${margin.left}, ${margin.top})`).call(yAxis)

    svg
      .append('text')
      .attr('transform', `translate(${width / 2 + margin.left}, ${height + margin.top + 40})`)
      .style('text-anchor', 'middle')
      .attr('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
      .text('Sequence Length')

    svg
      .append('text')
      .attr('transform', `rotate(-90)`)
      .attr('y', 0)
      .attr('x', 0 - height / 2 - margin.top)
      .attr('dy', '1em')
      .style('text-anchor', 'middle')
      .attr('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
      .text('Batch Size')

    const legend = svg
      .append('g')
      .attr('class', 'legend')
      .attr('transform', `translate(${width - 20}, 20)`)

    precisions.forEach((precision, index) => {
      const legendItem = legend.append('g').attr('transform', `translate(0, ${index * 30})`)

      legendItem
        .append('rect')
        .attr('x', 10)
        .attr('y', 10)
        .attr('width', 10)
        .attr('height', 10)
        .style('fill', precision.color)

      legendItem
        .append('text')
        .attr('x', 30)
        .attr('y', 16)
        .text(precision.name)
        .style('font-size', '16px')
        .style('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
        .attr('alignment-baseline', 'middle')
    })
    legend
      .append('rect')
      .attr('class', 'legend-box')
      .attr('width', 80)
      .attr('height', precisions.length * 30)
      .style('fill', 'none')
      .style('stroke-width', '1px')
      .style('stroke', theme === 'dark' ? '#f9fafb' : '#181f26')

    const tooltip = d3.select('#tooltip')

    for (const [precision, memory] of Object.entries(availableMemory)) {
      const sequenceLengths = d3
        .range(1, maxSeqLength, 1)
        .map((seqLength) => {
          return { seqLength, batchSize: memory / (seqLength * memoryPerInput) }
        })
        .filter((seqLength) => seqLength.batchSize <= maxBatchSize)
        .filter((seqLength) => seqLength.batchSize > 1 && seqLength.seqLength > 1)

      const lineGroup = svg
        .append('g')
        .attr('transform', `translate(${margin.left}, ${margin.top})`)
      const line = d3
        .line<LineChartData>()
        .x((d) => xScale(d.seqLength))
        .y((d) => yScale(d.batchSize))
        .curve(d3.curveBasis)

      lineGroup
        .append('path')
        .datum(sequenceLengths)
        .attr('fill', 'none')
        .attr('stroke', chooseColor(precision as Precision))
        .attr('stroke-width', 4)
        .attr('d', line)
        .on('mouseover', () => {
          tooltip.style('opacity', 1)
          tooltip.style('background-color', theme === 'dark' ? '#181f26' : '#f9fafb')
        })
        .on('mousemove', (event) => {
          tooltip.selectAll('text').remove()
          const [x, y] = d3.pointer(event)
          const xValue = xScale.invert(x)
          const yValue = yScale.invert(y)
          tooltip
            .html(`Sequence Length: ${xValue.toFixed(0)}<br/>Batch Size: ${yValue.toFixed(0)}`)
            .attr('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
            .style('left', event.pageX + 10 + 'px')
            .style('top', event.pageY + 10 + 'px')
        })
        .on('mouseout', () => {
          tooltip.style('opacity', 0)
        })
    }
  }, [availableMemory, memoryPerInput, theme])

  function chooseColor(precision: Precision) {
    const colors = {
      fp32: '#e45f5b',
      fp16: '#ffc068',
      int8: '#71cce9',
      int4: '#383d95',
    }
    return colors[precision] || 'gray'
  }

  return (
    <>
      <div id='tooltip'></div>
      <svg ref={chartRef} width={600} height={400} />
    </>
  )
}

// Prefill Chunking Model Size Bar Chart
function PrefillChunkingModelSizeBarChart({
  modelSize,
  largestModelSize,
  modelPrecision,
  deviceMemorySet,
  activationMemorySize,
}: {
  modelSize: number // in GB
  largestModelSize: number // largest model in full precision (fp32)
  modelPrecision: Precision // enum of fp32, fp16, int8, int4
  deviceMemorySet: boolean // true if device memory is set
  activationMemorySize: number // additional memory for activations in GB
}) {
  const { theme } = useTheme()
  const chartRef = useRef<SVGSVGElement>(null)

  const width = 600
  const height = 50

  useEffect(() => {
    d3.select(chartRef.current).selectAll('*').remove()

    const svg = d3.select(chartRef.current).attr('width', width).attr('height', height)

    const xScale = d3.scaleLinear().domain([0, largestModelSize]).range([0, width])

    if (modelSize + activationMemorySize > largestModelSize) {
      svg
        .append('rect')
        .attr('x', 0)
        .attr('y', 0)
        .attr('width', width)
        .attr('height', height)
        .attr('fill', 'transparent')
        .style('stroke', theme === 'dark' ? '#f9fafb' : '#181f26')
        .style('stroke-dasharray', '4, 4')
        .style('stroke-width', '2px')
      svg
        .append('text')
        .attr('x', width / 2)
        .attr('y', height / 2)
        .attr('text-anchor', 'middle')
        .attr('alignment-baseline', 'middle')
        .attr('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
        .text('Out of Memory')
    } else {
      svg
        .append('rect')
        .attr('x', 0)
        .attr('y', 0)
        .attr('width', xScale(modelSize))
        .attr('height', height)
        .attr('fill', chooseColor(modelPrecision))

      svg
        .append('rect')
        .attr('x', xScale(modelSize))
        .attr('y', 0)
        .attr('width', xScale(activationMemorySize))
        .attr('height', height)
        .attr('fill', '#a4b8e0')

      if (deviceMemorySet) {
        svg
          .append('rect')
          .attr('x', xScale(modelSize + activationMemorySize))
          .attr('y', 0)
          .attr('width', xScale(largestModelSize - (modelSize + activationMemorySize)))
          .attr('height', height)
          .attr('fill', 'transparent')
          .style('stroke', chooseColor(modelPrecision))
          .style('stroke-width', '2px')
      }
    }
  }, [modelSize, largestModelSize, modelPrecision, deviceMemorySet, activationMemorySize, theme])

  function chooseColor(precision: Precision) {
    const colors = {
      fp32: '#e45f5b',
      fp16: '#ffc068',
      int8: '#71cce9',
      int4: '#383d95',
    }
    return colors[precision] || 'gray'
  }

  return <svg ref={chartRef}></svg>
}

// Prefill Chunking Inference Runtime Line Chart
function PrefillChunkingInferenceRuntimeLineChart({
  availableMemory,
  memoryPerInput,
  activationMemorySize,
}: {
  availableMemory: AvailableMemory // in GB
  memoryPerInput: number // in GB
  activationMemorySize: number // in GB
}) {
  const { theme } = useTheme()
  const chartRef = useRef(null)
  const maxSeqLength = 4096
  const maxBatchSize = 128

  useEffect(() => {
    const margin = { top: 20, right: 20, bottom: 50, left: 50 }
    const width = 600 - margin.left - margin.right
    const height = 400 - margin.top - margin.bottom
    const precisions = [
      { name: 'FP32', color: '#e45f5b' },
      { name: 'FP16', color: '#ffc068' },
      { name: 'INT8', color: '#71cce9' },
      { name: 'INT4', color: '#383d95' },
    ]

    const svg = d3.select(chartRef.current)
    svg.selectAll('*').remove()

    const xScale = d3.scaleLinear().domain([0, maxSeqLength]).range([0, width])

    const yScale = d3.scaleLinear().domain([0, maxBatchSize]).range([height, 0])

    const xAxis = d3.axisBottom(xScale)
    const yAxis = d3.axisLeft(yScale)
    svg
      .attr('width', width + margin.left + margin.right)
      .attr('height', height + margin.top + margin.bottom)
      .append('g')
      .attr('transform', `translate(${margin.left}, ${margin.top})`)

    svg
      .append('g')
      .attr('transform', `translate(${margin.left}, ${height + margin.top})`)
      .call(xAxis)

    svg.append('g').attr('transform', `translate(${margin.left}, ${margin.top})`).call(yAxis)

    svg
      .append('text')
      .attr('transform', `translate(${width / 2 + margin.left}, ${height + margin.top + 40})`)
      .style('text-anchor', 'middle')
      .attr('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
      .text('Sequence Length')

    svg
      .append('text')
      .attr('transform', `rotate(-90)`)
      .attr('y', 0)
      .attr('x', 0 - height / 2 - margin.top)
      .attr('dy', '1em')
      .style('text-anchor', 'middle')
      .attr('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
      .text('Batch Size')

    const legend = svg
      .append('g')
      .attr('class', 'legend')
      .attr('transform', `translate(${width - 20}, 20)`)

    precisions.forEach((precision, index) => {
      const legendItem = legend.append('g').attr('transform', `translate(0, ${index * 30})`)

      legendItem
        .append('rect')
        .attr('x', 10)
        .attr('y', 10)
        .attr('width', 10)
        .attr('height', 10)
        .style('fill', precision.color)

      legendItem
        .append('text')
        .attr('x', 30)
        .attr('y', 16)
        .text(precision.name)
        .style('font-size', '16px')
        .style('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
        .attr('alignment-baseline', 'middle')
    })
    legend
      .append('rect')
      .attr('class', 'legend-box')
      .attr('width', 80)
      .attr('height', precisions.length * 30)
      .style('fill', 'none')
      .style('stroke-width', '1px')
      .style('stroke', theme === 'dark' ? '#f9fafb' : '#181f26')

    const tooltip = d3.select('#tooltip')

    for (const [precision, memory] of Object.entries(availableMemory)) {
      const sequenceLengths = d3
        .range(1, maxSeqLength, 1)
        .map((seqLength) => {
          return {
            seqLength,
            batchSize: (memory - activationMemorySize) / (seqLength * memoryPerInput),
          }
        })
        .filter((seqLength) => seqLength.batchSize <= maxBatchSize)
        .filter((seqLength) => seqLength.batchSize > 1 && seqLength.seqLength > 1)

      const lineGroup = svg
        .append('g')
        .attr('transform', `translate(${margin.left}, ${margin.top})`)
      const line = d3
        .line<LineChartData>()
        .x((d) => xScale(d.seqLength))
        .y((d) => yScale(d.batchSize))
        .curve(d3.curveBasis)

      lineGroup
        .append('path')
        .datum(sequenceLengths)
        .attr('fill', 'none')
        .attr('stroke', chooseColor(precision as Precision))
        .attr('stroke-width', 4)
        .attr('d', line)
        .on('mouseover', () => {
          tooltip.style('opacity', 1)
          tooltip.style('background-color', theme === 'dark' ? '#181f26' : '#f9fafb')
        })
        .on('mousemove', (event) => {
          tooltip.selectAll('text').remove()
          const [x, y] = d3.pointer(event)
          const xValue = xScale.invert(x)
          const yValue = yScale.invert(y)
          tooltip
            .html(`Sequence Length: ${xValue.toFixed(0)}<br/>Batch Size: ${yValue.toFixed(0)}`)
            .attr('fill', theme === 'dark' ? '#f9fafb' : '#181f26')
            .style('left', event.pageX + 10 + 'px')
            .style('top', event.pageY + 10 + 'px')
        })
        .on('mouseout', () => {
          tooltip.style('opacity', 0)
        })
    }
  }, [availableMemory, memoryPerInput, activationMemorySize, theme])

  function chooseColor(precision: Precision) {
    const colors = {
      fp32: '#e45f5b',
      fp16: '#ffc068',
      int8: '#71cce9',
      int4: '#383d95',
    }
    return colors[precision] || 'gray'
  }

  return (
    <>
      <div id='tooltip'></div>
      <svg ref={chartRef} width={600} height={400} />
    </>
  )
}

// Prefill Chunking Calculator
const PrefillChunkingCalculator = ({
  deviceMemory,
  modelParams,
  hiddenSize,
  numLayers,
  batchSize,
  seqLength,
  maxChunkSize,
  intermediateSize,
}: {
  deviceMemory: number
  modelParams: number
  hiddenSize: number
  numLayers: number
  batchSize: number | null
  seqLength: number | null
  maxChunkSize: number | null
  intermediateSize: number | null
}) => {
  // Function to calculate memory usage based on precision and model parameters
  function calculateMemory(params: number, precision: 'fp32' | 'fp16' | 'int8' | 'int4'): number {
    const paramSize = { fp32: 4, fp16: 2, int8: 1, int4: 0.5 }
    return params * paramSize[precision] // in GB
  }

  function calculateMemoryPerInput(hiddenSize: number, numLayers: number) {
    const memoryPerInput = 4 * hiddenSize * numLayers
    return memoryPerInput / 1_000_000_000 // in GB
  }

  function calculateMemoryNeededForActivations(
    maxChunkSize: number,
    intermediateSize: number,
    hiddenSize: number,
  ) {
    const activationMemorySize = maxChunkSize * 2 * Math.max(2 * intermediateSize, 4 * hiddenSize)
    return activationMemorySize / 1_000_000_000 // in GB
  }

  const memoryPerInput = calculateMemoryPerInput(hiddenSize, numLayers)
  const activationMemorySize = calculateMemoryNeededForActivations(
    maxChunkSize!,
    intermediateSize!,
    hiddenSize,
  )

  return (
    <>
      {/* Model Footprint with Prefill Chunking */}
      <div>
        <div className='chart mb-8'>
          <div className='flex flex-col items-center'>
            <div className='text-2xl'>Model Footprint with Prefill Chunking</div>
          </div>
          <div>
            <div className='chart-row'>
              <div className='chart-row-title'>FP32</div>
              <PrefillChunkingModelSizeBarChart
                modelSize={calculateMemory(modelParams, 'fp32')}
                largestModelSize={deviceMemory || calculateMemory(modelParams, 'fp32')}
                modelPrecision='fp32'
                deviceMemorySet={deviceMemory !== null && deviceMemory > 0}
                activationMemorySize={activationMemorySize}
              />
              <div className='chart-row-size ml-8'>
                {(calculateMemory(modelParams, 'fp32') + activationMemorySize).toFixed(2)}{' '}
                {deviceMemory !== null && deviceMemory > 0 ? `/ ${deviceMemory} ` : null}GB
              </div>
            </div>

            <div className='chart-row my-8'>
              <div className='chart-row-title'>FP16</div>
              <PrefillChunkingModelSizeBarChart
                modelSize={calculateMemory(modelParams, 'fp16')}
                largestModelSize={deviceMemory || calculateMemory(modelParams, 'fp16')}
                modelPrecision='fp16'
                deviceMemorySet={deviceMemory !== null && deviceMemory > 0}
                activationMemorySize={activationMemorySize}
              />
              <div className='chart-row-size ml-8'>
                {(calculateMemory(modelParams, 'fp16') + activationMemorySize).toFixed(2)}{' '}
                {deviceMemory !== null && deviceMemory > 0 ? `/ ${deviceMemory} ` : null}GB
              </div>
            </div>

            <div className='chart-row my-8'>
              <div className='chart-row-title'>INT8</div>
              <PrefillChunkingModelSizeBarChart
                modelSize={calculateMemory(modelParams, 'int8')}
                largestModelSize={deviceMemory || calculateMemory(modelParams, 'int8')}
                modelPrecision='int8'
                deviceMemorySet={deviceMemory !== null && deviceMemory > 0}
                activationMemorySize={activationMemorySize}
              />
              <div className='chart-row-size ml-8'>
                {(calculateMemory(modelParams, 'int8') + activationMemorySize).toFixed(2)}{' '}
                {deviceMemory !== null && deviceMemory > 0 ? `/ ${deviceMemory} ` : null}GB
              </div>
            </div>

            <div className='chart-row my-8'>
              <div className='chart-row-title'>INT4</div>
              <PrefillChunkingModelSizeBarChart
                modelSize={calculateMemory(modelParams, 'int4')}
                largestModelSize={deviceMemory || calculateMemory(modelParams, 'int4')}
                modelPrecision='int4'
                deviceMemorySet={deviceMemory !== null && deviceMemory > 0}
                activationMemorySize={activationMemorySize}
              />
              <div className='chart-row-size ml-8'>
                {(calculateMemory(modelParams, 'int4') + activationMemorySize).toFixed(2)}{' '}
                {deviceMemory !== null && deviceMemory > 0 ? `/ ${deviceMemory} ` : null}GB
              </div>
            </div>
          </div>
        </div>
        <div className='chart'>
          <div className='flex flex-col items-center'>
            <div className='text-2xl'>
              Maximum Batch Size / Sequence Length with Prefill Chunking
            </div>
          </div>
          <PrefillChunkingInferenceRuntimeLineChart
            availableMemory={{
              int4: deviceMemory - calculateMemory(modelParams, 'int4'),
              int8: deviceMemory - calculateMemory(modelParams, 'int8'),
              fp16: deviceMemory - calculateMemory(modelParams, 'fp16'),
              fp32: deviceMemory - calculateMemory(modelParams, 'fp32'),
            }}
            memoryPerInput={memoryPerInput}
            activationMemorySize={activationMemorySize}
          />
        </div>
      </div>
    </>
  )
}

// Calculator page with toggle feature
const Calculator = () => {
  // Model Parameters (billions)
  const [modelParams, setModelParams] = useState<number | null>(null)
  const [hiddenSize, setHiddenSize] = useState<number | null>(null)
  const [numLayers, setNumLayers] = useState<number | null>(null)

  // Device Memory (GB)
  const [deviceMemory, setDeviceMemory] = useState<number | null>(null)

  // Inputs
  const [batchSize, setBatchSize] = useState<number | null>(null)
  const [seqLength, setSeqLength] = useState<number | null>(null)

  // Prefill Chunking Inputs
  const [maxChunkSize, setMaxChunkSize] = useState<number | null>(null)
  const [intermediateSize, setIntermediateSize] = useState<number | null>(null)

  // Toggle between standard calculator and prefill chunking calculator
  const [isPrefillChunking, setIsPrefillChunking] = useState<boolean>(false)

  // Toggle between model selection and custom model input
  const [modelSelectionTab, setModelSelectionTab] = useState<boolean>(true)

  // Toggle between device selection and custom device input
  const [deviceSelectionTab, setDeviceSelectionTab] = useState<boolean>(true)

  // Calculate model memory
  function calculateMemory(params: number, precision: Precision) {
    const paramSize = { fp32: 4, fp16: 2, int8: 1, int4: 0.5 }
    return params * paramSize[precision] // in GB
  }

  // Calculate memory per input (sequence length and batch size)
  function calculateMemoryPerInput(hiddenSize: number, numLayers: number) {
    const memoryPerInput = 4 * hiddenSize * numLayers
    return memoryPerInput / 1_000_000_000 // in GB
  }

  // Calculate maximum batch size / sequence length
  function calculateMaxInputSize(
    deviceMemory: number,
    modelParams: number,
    hiddenSize: number,
    numLayers: number,
    precision: Precision,
    inputSize: number,
  ) {
    const memoryPerInput = calculateMemoryPerInput(hiddenSize, numLayers)
    const availableMemory = deviceMemory - calculateMemory(modelParams, precision)
    return Math.floor(availableMemory / (memoryPerInput * inputSize))
  }

  // Check if memory is valid for batch size / sequence length combination
  function calculateMemoryValid(
    deviceMemory: number,
    modelParams: number,
    hiddenSize: number,
    numLayers: number,
    precision: Precision,
    batchSize: number,
    seqLength: number,
  ) {
    const memoryPerInput = calculateMemoryPerInput(hiddenSize, numLayers)
    const availableMemory = deviceMemory - calculateMemory(modelParams, precision)
    return availableMemory >= memoryPerInput * batchSize * seqLength
  }

  return (
    <div className='flex flex-col items-center justify-center min-h-screen px-4'>
      {/* Toggle Button */}
      <div className='mb-4 flex space-x-4'>
        <button
          className={`${
            !isPrefillChunking ? 'calculator-input-tab-active' : 'calculator-input-tab'
          }`}
          onClick={() => setIsPrefillChunking(false)}
        >
          Standard Calculator
        </button>
        <button
          className={`${
            isPrefillChunking ? 'calculator-input-tab-active' : 'calculator-input-tab'
          }`}
          onClick={() => setIsPrefillChunking(true)}
        >
          Calculator with Prefill Chunking
        </button>
      </div>

      {/* Model Memory Calculator */}
      <div className='w-full max-w-4xl'>
        <div className='text-4xl mb-4 text-center'>Model Memory Calculator</div>
        <div className='mb-6 text-center'>
          Use our Model Memory Calculator to help you estimate the memory footprint of your model
          for different precisions and the maximum batch size / sequence length combination you can
          run on your device.
        </div>

        {/* Model and Device Selection */}
        <div className='grid grid-cols-1 sm:grid-cols-2 gap-4 mb-6'>
          {/* Model Selection */}
          <div className='calculator-input-box'>
            <div className='text-2xl calculator-input-title'>Model</div>
            <div className='calculator-input-content'>
              <div className='mb-2'>
                <button
                  className={`${
                    modelSelectionTab ? 'calculator-input-tab-active' : 'calculator-input-tab'
                  }`}
                  onClick={() => setModelSelectionTab(true)}
                >
                  Model Selection
                </button>
                <button
                  className={`${
                    modelSelectionTab ? 'calculator-input-tab' : 'calculator-input-tab-active'
                  }`}
                  onClick={() => setModelSelectionTab(false)}
                >
                  Custom Model
                </button>
              </div>
              <div>
                {modelSelectionTab ? (
                  <>
                    <label htmlFor='model'>Select a Model</label>
                    <select
                      id='model'
                      className='calculator-select'
                      onChange={(e) => {
                        setModelParams(Number(e.target.value))
                        setHiddenSize(
                          Number(
                            e.target.options[e.target.selectedIndex].getAttribute(
                              'data-hiddenSize',
                            ),
                          ),
                        )
                        setNumLayers(
                          Number(
                            e.target.options[e.target.selectedIndex].getAttribute('data-numLayers'),
                          ),
                        )
                      }}
                    >
                      <option value=''>None selected</option>
                      {MODELS.map((model) => (
                        <option
                          key={model.name}
                          value={model.params}
                          data-hiddenSize={model.hidden_size}
                          data-numLayers={model.num_hidden_layers}
                        >
                          {model.name}
                        </option>
                      ))}
                    </select>
                  </>
                ) : (
                  <>
                    <label htmlFor='modelParams'>Model Parameters (in billions)</label>
                    <input
                      type='number'
                      id='modelParams'
                      className='calculator-input mb-2'
                      placeholder='e.g. 7 (for LLaMA-7B)'
                      value={modelParams || ''}
                      min={0}
                      onChange={(e) => setModelParams(Number(e.target.value))}
                    />
                    <label htmlFor='hiddenSize'>Hidden Size</label>
                    <input
                      type='number'
                      id='hiddenSize'
                      className='calculator-input mb-2'
                      placeholder='e.g. 4096 (for LLaMA-7B)'
                      value={hiddenSize || ''}
                      min={1}
                      onChange={(e) => setHiddenSize(Number(e.target.value))}
                    />
                    <label htmlFor='numLayers'>Number of Layers</label>
                    <input
                      type='number'
                      id='numLayers'
                      className='calculator-input'
                      placeholder='e.g. 32 (for LLaMA-7B)'
                      value={numLayers || ''}
                      min={1}
                      onChange={(e) => setNumLayers(Number(e.target.value))}
                    />
                  </>
                )}
              </div>
            </div>
          </div>

          {/* Device Selection */}
          <div className='calculator-input-box'>
            <div className='text-2xl calculator-input-title'>Device</div>
            <div className='calculator-input-content'>
              <div className='mb-2'>
                <button
                  className={`${
                    deviceSelectionTab ? 'calculator-input-tab-active' : 'calculator-input-tab'
                  }`}
                  onClick={() => {
                    setDeviceSelectionTab(true)
                    setDeviceMemory(null)
                  }}
                >
                  Device Selection
                </button>
                <button
                  className={`${
                    deviceSelectionTab ? 'calculator-input-tab' : 'calculator-input-tab-active'
                  }`}
                  onClick={() => {
                    setDeviceSelectionTab(false)
                    setDeviceMemory(null)
                  }}
                >
                  Custom Device
                </button>
              </div>
              <div>
                {deviceSelectionTab ? (
                  <>
                    <label htmlFor='device'>Select a Device</label>
                    <select
                      id='device'
                      className='calculator-select'
                      onChange={(e) => setDeviceMemory(Number(e.target.value))}
                    >
                      <option value=''>None selected</option>
                      {DEVICES.map((device) => (
                        <option key={device.name} value={device.size}>
                          {device.name}
                        </option>
                      ))}
                    </select>
                  </>
                ) : (
                  <>
                    <label htmlFor='deviceMemory'>Device RAM (in GB)</label>
                    <input
                      type='number'
                      id='deviceMemory'
                      className='calculator-input'
                      placeholder='e.g. 24'
                      value={deviceMemory || ''}
                      min={0}
                      onChange={(e) => setDeviceMemory(Number(e.target.value))}
                    />
                  </>
                )}
              </div>
            </div>
          </div>

          <div className='calculator-box'>
            <div className='text-2xl ml-5 mb-4'>Backend Precision Table</div>
            <div className='ml-5 mb-4'>
              <BackendPrecisionTable />
            </div>
            <div className='ml-5'>
              This table shows the precision used by each Takeoff backend for CPUs and GPUs, as well
              as their accuracy preservation.
            </div>
          </div>

          <div className='calculator-box'>
            <div className='text-2xl ml-5 mb-4'>Input parameters</div>
            <div className='ml-5 mb-4'>
              <strong>Sequence Length</strong>: The combined length of input tokens and output
              tokens. To restrict the maximum sequence length for inference on Takeoff, use the API
              parameters <code>prompt_new_tokens</code> for input tokens and{' '}
              <code>max_new_tokens</code> for output tokens when making a request.
            </div>
            <div className='ml-5'>
              <strong>Batch Size</strong>: The number of sequences that can be processed in
              parallel. To set a maximum batch size for inference on Takeoff, set the environment
              variable <code>TAKEOFF_MAX_BATCH_SIZE</code> to your desired value.
            </div>
          </div>
        </div>

        {/* Prefill Chunking Settings */}
        {isPrefillChunking && (
          <div className='calculator-input-box mb-6'>
            <div className='text-2xl calculator-input-title'>Prefill Chunking Settings</div>
            <div className='calculator-input-content'>
              <label htmlFor='maxChunkSize'>Max Chunk Size</label>
              <input
                type='number'
                id='maxChunkSize'
                className='calculator-input mb-2'
                placeholder='e.g. 512'
                value={maxChunkSize || ''}
                min={1}
                onChange={(e) => setMaxChunkSize(Number(e.target.value))}
              />
              <label htmlFor='intermediateSize'>Intermediate Size</label>
              <input
                type='number'
                id='intermediateSize'
                className='calculator-input'
                placeholder='e.g. 2048'
                value={intermediateSize || ''}
                min={1}
                onChange={(e) => setIntermediateSize(Number(e.target.value))}
              />
            </div>
          </div>
        )}

        {/* Charts Section */}
        {isPrefillChunking ? (
          <PrefillChunkingCalculator
            deviceMemory={deviceMemory!}
            modelParams={modelParams!}
            hiddenSize={hiddenSize!}
            numLayers={numLayers!}
            batchSize={batchSize}
            seqLength={seqLength}
            maxChunkSize={maxChunkSize}
            intermediateSize={intermediateSize}
          />
        ) : (
          hiddenSize &&
          numLayers &&
          deviceMemory &&
          modelParams && (
            <>
              {/* Model Footprint Chart */}
              <div className='chart mb-8'>
                <div className='text-2xl text-center mb-4'>Model Footprint</div>
                <div className='space-y-8'>
                  <div className='chart-row'>
                    <div className='chart-row-title'>FP32</div>
                    <ModelSizeBarChart
                      modelSize={calculateMemory(modelParams, 'fp32')}
                      largestModelSize={deviceMemory || calculateMemory(modelParams, 'fp32')}
                      modelPrecision='fp32'
                      deviceMemorySet={deviceMemory !== null && deviceMemory > 0}
                    />
                    <div className='chart-row-size ml-8'>
                      {calculateMemory(modelParams, 'fp32')}{' '}
                      {deviceMemory ? `/ ${deviceMemory} ` : null}GB
                    </div>
                  </div>

                  {/* FP16 */}
                  <div className='chart-row'>
                    <div className='chart-row-title'>FP16</div>
                    <ModelSizeBarChart
                      modelSize={calculateMemory(modelParams, 'fp16')}
                      largestModelSize={deviceMemory || calculateMemory(modelParams, 'fp16')}
                      modelPrecision='fp16'
                      deviceMemorySet={deviceMemory !== null && deviceMemory > 0}
                    />
                    <div className='chart-row-size ml-8'>
                      {calculateMemory(modelParams, 'fp16')}{' '}
                      {deviceMemory ? `/ ${deviceMemory} ` : null}GB
                    </div>
                  </div>

                  {/* INT8 */}
                  <div className='chart-row'>
                    <div className='chart-row-title'>INT8</div>
                    <ModelSizeBarChart
                      modelSize={calculateMemory(modelParams, 'int8')}
                      largestModelSize={deviceMemory || calculateMemory(modelParams, 'int8')}
                      modelPrecision='int8'
                      deviceMemorySet={deviceMemory !== null && deviceMemory > 0}
                    />
                    <div className='chart-row-size ml-8'>
                      {calculateMemory(modelParams, 'int8')}{' '}
                      {deviceMemory ? `/ ${deviceMemory} ` : null}GB
                    </div>
                  </div>

                  {/* INT4 */}
                  <div className='chart-row'>
                    <div className='chart-row-title'>INT4</div>
                    <ModelSizeBarChart
                      modelSize={calculateMemory(modelParams, 'int4')}
                      largestModelSize={deviceMemory || calculateMemory(modelParams, 'int4')}
                      modelPrecision='int4'
                      deviceMemorySet={deviceMemory !== null && deviceMemory > 0}
                    />
                    <div className='chart-row-size ml-8'>
                      {calculateMemory(modelParams, 'int4')}{' '}
                      {deviceMemory ? `/ ${deviceMemory} ` : null}GB
                    </div>
                  </div>
                </div>
              </div>

              {/* Maximum Batch Size / Sequence Length Chart */}
              <div className='chart mb-8'>
                <div className='text-2xl text-center mb-4'>
                  Maximum Batch Size / Sequence Length
                </div>
                <div className='flex flex-row items-left'>
                  <InferenceRuntimeLineChart
                    availableMemory={{
                      int4: deviceMemory - calculateMemory(modelParams, 'int4'),
                      int8: deviceMemory - calculateMemory(modelParams, 'int8'),
                      fp16: deviceMemory - calculateMemory(modelParams, 'fp16'),
                      fp32: deviceMemory - calculateMemory(modelParams, 'fp32'),
                    }}
                    memoryPerInput={calculateMemoryPerInput(hiddenSize, numLayers)}
                  />
                  <div className='chart-side-panel ml-4 pt-4'>
                    <div className='mb-2'>
                      Memory/token:{' '}
                      {(calculateMemoryPerInput(hiddenSize, numLayers) * 1_000_000).toFixed(0)} KB
                    </div>
                    <label htmlFor='batchSize'>Batch Size</label>
                    <input
                      type='number'
                      id='batchSize'
                      className='side-panel-input mb-2'
                      value={batchSize || ''}
                      min={1}
                      onChange={(e) => setBatchSize(Number(e.target.value))}
                    />
                    <label htmlFor='seqLength'>Sequence Length</label>
                    <input
                      type='number'
                      id='seqLength'
                      className='side-panel-input'
                      value={seqLength || ''}
                      min={1}
                      onChange={(e) => setSeqLength(Number(e.target.value))}
                    />
                    <div className='mt-4'>
                      {!batchSize && !seqLength ? (
                        <div>
                          Input a batch size or sequence length to see the maximum batch size or
                          sequence length you can run on your device.
                        </div>
                      ) : null}
                      {batchSize && !seqLength ? (
                        <>
                          <div>Max Sequence Lengths:</div>
                          <div>
                            FP32:{' '}
                            <strong>
                              {calculateMaxInputSize(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'fp32',
                                batchSize,
                              ) > 0
                                ? calculateMaxInputSize(
                                    deviceMemory,
                                    modelParams,
                                    hiddenSize,
                                    numLayers,
                                    'fp32',
                                    batchSize,
                                  )
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            FP16:{' '}
                            <strong>
                              {calculateMaxInputSize(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'fp16',
                                batchSize,
                              ) > 0
                                ? calculateMaxInputSize(
                                    deviceMemory,
                                    modelParams,
                                    hiddenSize,
                                    numLayers,
                                    'fp16',
                                    batchSize,
                                  )
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            INT8:{' '}
                            <strong>
                              {calculateMaxInputSize(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'int8',
                                batchSize,
                              ) > 0
                                ? calculateMaxInputSize(
                                    deviceMemory,
                                    modelParams,
                                    hiddenSize,
                                    numLayers,
                                    'int8',
                                    batchSize,
                                  )
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            INT4:{' '}
                            <strong>
                              {calculateMaxInputSize(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'int4',
                                batchSize,
                              ) > 0
                                ? calculateMaxInputSize(
                                    deviceMemory,
                                    modelParams,
                                    hiddenSize,
                                    numLayers,
                                    'int4',
                                    batchSize,
                                  )
                                : 'Out of Memory'}
                            </strong>
                          </div>
                        </>
                      ) : null}
                      {!batchSize && seqLength ? (
                        <>
                          <div>Max Batch Sizes:</div>
                          <div>
                            FP32:{' '}
                            <strong>
                              {calculateMaxInputSize(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'fp32',
                                seqLength,
                              ) > 0
                                ? calculateMaxInputSize(
                                    deviceMemory,
                                    modelParams,
                                    hiddenSize,
                                    numLayers,
                                    'fp32',
                                    seqLength,
                                  )
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            FP16:{' '}
                            <strong>
                              {calculateMaxInputSize(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'fp16',
                                seqLength,
                              ) > 0
                                ? calculateMaxInputSize(
                                    deviceMemory,
                                    modelParams,
                                    hiddenSize,
                                    numLayers,
                                    'fp16',
                                    seqLength,
                                  )
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            INT8:{' '}
                            <strong>
                              {calculateMaxInputSize(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'int8',
                                seqLength,
                              ) > 0
                                ? calculateMaxInputSize(
                                    deviceMemory,
                                    modelParams,
                                    hiddenSize,
                                    numLayers,
                                    'int8',
                                    seqLength,
                                  )
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            INT4:{' '}
                            <strong>
                              {calculateMaxInputSize(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'int4',
                                seqLength,
                              ) > 0
                                ? calculateMaxInputSize(
                                    deviceMemory,
                                    modelParams,
                                    hiddenSize,
                                    numLayers,
                                    'int4',
                                    seqLength,
                                  )
                                : 'Out of Memory'}
                            </strong>
                          </div>
                        </>
                      ) : null}
                      {batchSize && seqLength ? (
                        <>
                          <div>Total Memory Usage:</div>
                          <div>
                            FP32:{' '}
                            <strong>
                              {calculateMemoryValid(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'fp32',
                                batchSize,
                                seqLength,
                              )
                                ? (
                                    calculateMemory(modelParams, 'fp32') +
                                    calculateMemoryPerInput(hiddenSize, numLayers) *
                                      batchSize *
                                      seqLength
                                  ).toFixed(2) + ' GB'
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            FP16:{' '}
                            <strong>
                              {calculateMemoryValid(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'fp16',
                                batchSize,
                                seqLength,
                              )
                                ? (
                                    calculateMemory(modelParams, 'fp16') +
                                    calculateMemoryPerInput(hiddenSize, numLayers) *
                                      batchSize *
                                      seqLength
                                  ).toFixed(2) + ' GB'
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            INT8:{' '}
                            <strong>
                              {calculateMemoryValid(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'int8',
                                batchSize,
                                seqLength,
                              )
                                ? (
                                    calculateMemory(modelParams, 'int8') +
                                    calculateMemoryPerInput(hiddenSize, numLayers) *
                                      batchSize *
                                      seqLength
                                  ).toFixed(2) + ' GB'
                                : 'Out of Memory'}
                            </strong>
                          </div>
                          <div>
                            INT4:{' '}
                            <strong>
                              {calculateMemoryValid(
                                deviceMemory,
                                modelParams,
                                hiddenSize,
                                numLayers,
                                'int4',
                                batchSize,
                                seqLength,
                              )
                                ? (
                                    calculateMemory(modelParams, 'int4') +
                                    calculateMemoryPerInput(hiddenSize, numLayers) *
                                      batchSize *
                                      seqLength
                                  ).toFixed(2) + ' GB'
                                : 'Out of Memory'}
                            </strong>
                          </div>
                        </>
                      ) : null}
                    </div>
                  </div>
                </div>
              </div>
            </>
          )
        )}
      </div>
    </div>
  )
}

export default Calculator