import type { Stroke, RecognitionCandidate } from './types'
import { strokesToImage28 } from './preprocess'

// Pure JS CNN inference - no ONNX Runtime needed
// Model: 4 conv layers (with ReLU, MaxPool, Dropout) + 2 linear layers

interface WeightTensor {
  dims: number[]
  data: Float32Array
}

interface Weights {
  [name: string]: WeightTensor
}

let weights: Weights | null = null
let loading = false
let initFailed = false

const LETTERS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

export function isReady(): boolean {
  return weights !== null
}

export async function init(): Promise<void> {
  if (weights || loading || initFailed) return
  loading = true

  try {
    const res = await fetch('/models/emnist-weights.bin')
    if (!res.ok) throw new Error('weights not found')
    const buf = await res.arrayBuffer()
    weights = parseWeights(buf)
    console.log('CNN engine ready (pure JS)')
  } catch (e) {
    console.warn('CNN engine init failed:', e)
    weights = null
    initFailed = true
  }
  loading = false
}

function parseWeights(buf: ArrayBuffer): Weights {
  const view = new DataView(buf)
  let offset = 0

  function readU32() { const v = view.getUint32(offset, true); offset += 4; return v }
  function readF32() { const v = view.getFloat32(offset, true); offset += 4; return v }

  const numTensors = readU32()
  const w: Weights = {}

  for (let t = 0; t < numTensors; t++) {
    const nameLen = readU32()
    const nameBytes = new Uint8Array(buf, offset, nameLen)
    const name = new TextDecoder().decode(nameBytes)
    offset += nameLen

    const ndims = readU32()
    const dims: number[] = []
    for (let d = 0; d < ndims; d++) dims.push(readU32())

    const dataLen = readU32()
    const data = new Float32Array(dataLen)
    for (let i = 0; i < dataLen; i++) data[i] = readF32()

    w[name] = { dims, data }
  }
  return w
}

// CNN operations

interface TensorPlane {
  data: Float32Array
  h: number
  w: number
}

function conv2d(input: Float32Array, inC: number, inH: number, inW: number,
                weight: Float32Array, bias: Float32Array, outC: number, kH: number, kW: number, pad: number): TensorPlane {
  const outH = inH + 2 * pad - kH + 1
  const outW = inW + 2 * pad - kW + 1
  const output = new Float32Array(outC * outH * outW)

  for (let oc = 0; oc < outC; oc++) {
    for (let oh = 0; oh < outH; oh++) {
      for (let ow = 0; ow < outW; ow++) {
        let sum = bias[oc]
        for (let ic = 0; ic < inC; ic++) {
          for (let kh = 0; kh < kH; kh++) {
            for (let kw = 0; kw < kW; kw++) {
              const ih = oh - pad + kh
              const iw = ow - pad + kw
              if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
                const inputIdx = ic * inH * inW + ih * inW + iw
                const weightIdx = oc * inC * kH * kW + ic * kH * kW + kh * kW + kw
                sum += input[inputIdx] * weight[weightIdx]
              }
            }
          }
        }
        output[oc * outH * outW + oh * outW + ow] = sum
      }
    }
  }
  return { data: output, h: outH, w: outW }
}

function relu(data: Float32Array): Float32Array {
  const out = new Float32Array(data.length)
  for (let i = 0; i < data.length; i++) out[i] = data[i] > 0 ? data[i] : 0
  return out
}

function maxpool2d(input: Float32Array, c: number, h: number, w: number, k: number): TensorPlane {
  const oh = Math.floor(h / k)
  const ow = Math.floor(w / k)
  const output = new Float32Array(c * oh * ow)
  for (let ch = 0; ch < c; ch++) {
    for (let i = 0; i < oh; i++) {
      for (let j = 0; j < ow; j++) {
        let max = -Infinity
        for (let ki = 0; ki < k; ki++) {
          for (let kj = 0; kj < k; kj++) {
            const idx = ch * h * w + (i * k + ki) * w + (j * k + kj)
            if (input[idx] > max) max = input[idx]
          }
        }
        output[ch * oh * ow + i * ow + j] = max
      }
    }
  }
  return { data: output, h: oh, w: ow }
}

function linear(input: Float32Array, weight: Float32Array, bias: Float32Array, outSize: number, inSize: number): Float32Array {
  const output = new Float32Array(outSize)
  for (let o = 0; o < outSize; o++) {
    let sum = bias[o]
    for (let i = 0; i < inSize; i++) {
      sum += input[i] * weight[o * inSize + i]
    }
    output[o] = sum
  }
  return output
}

function forward(input: Float32Array, w: Weights): Float32Array {
  // features.0: Conv2d(1, 32, 3, padding=1) + ReLU
  let { data, h, w: width } = conv2d(input, 1, 28, 28,
    w['features.0.weight'].data, w['features.0.bias'].data, 32, 3, 3, 1)
  data = relu(data)

  // features.2: Conv2d(32, 32, 3, padding=1) + ReLU
  ;({ data, h, w: width } = conv2d(data, 32, h, width,
    w['features.2.weight'].data, w['features.2.bias'].data, 32, 3, 3, 1))
  data = relu(data)

  // features.4: MaxPool2d(2)
  ;({ data, h, w: width } = maxpool2d(data, 32, h, width, 2))
  // features.5: Dropout2d - skip at inference

  // features.6: Conv2d(32, 64, 3, padding=1) + ReLU
  ;({ data, h, w: width } = conv2d(data, 32, h, width,
    w['features.6.weight'].data, w['features.6.bias'].data, 64, 3, 3, 1))
  data = relu(data)

  // features.8: Conv2d(64, 64, 3, padding=1) + ReLU
  ;({ data, h, w: width } = conv2d(data, 64, h, width,
    w['features.8.weight'].data, w['features.8.bias'].data, 64, 3, 3, 1))
  data = relu(data)

  // features.10: MaxPool2d(2)
  ;({ data, h, w: width } = maxpool2d(data, 64, h, width, 2))
  // features.11: Dropout2d - skip at inference

  // classifier.1: Linear(3136, 256) + ReLU
  data = linear(data, w['classifier.1.weight'].data, w['classifier.1.bias'].data, 256, 64 * h * width)
  data = relu(data)
  // classifier.3: Dropout - skip at inference

  // classifier.4: Linear(256, 26)
  data = linear(data, w['classifier.4.weight'].data, w['classifier.4.bias'].data, 26, 256)

  return data
}

export async function recognize(strokes: Stroke[]): Promise<RecognitionCandidate[]> {
  if (!weights) return []

  try {
    const imageData = strokesToImage28(strokes)
    const logits = forward(imageData, weights)

    // softmax
    const maxVal = Math.max(...logits)
    const exps = Array.from(logits).map(v => Math.exp(v - maxVal))
    const sum = exps.reduce((a, b) => a + b, 0)
    const probs = exps.map(e => e / sum)

    const candidates: RecognitionCandidate[] = probs.map((confidence, i) => ({
      letter: LETTERS[i],
      confidence,
    }))

    candidates.sort((a, b) => b.confidence - a.confidence)
    return candidates.slice(0, 5)
  } catch (e) {
    console.warn('CNN recognition failed:', e)
    return []
  }
}
