import type { Candle } from './types'

export function computeSMA(values: number[], period: number): (number | null)[] {
  const result: (number | null)[] = []
  for (let i = 0; i < values.length; i++) {
    if (i < period - 1) { result.push(null); continue }
    let sum = 0
    for (let j = i - period + 1; j <= i; j++) sum += values[j]
    result.push(sum / period)
  }
  return result
}

export function computeEMA(values: number[], period: number): (number | null)[] {
  const result: (number | null)[] = []
  if (values.length < period) return values.map(() => null)
  const k = 2 / (period + 1)
  let e = 0
  for (let i = 0; i < period; i++) e += values[i]
  e /= period
  for (let i = 0; i < values.length; i++) {
    if (i < period - 1) { result.push(null); continue }
    if (i === period - 1) { result.push(e); continue }
    e = values[i] * k + e * (1 - k)
    result.push(e)
  }
  return result
}

export function computeRSI(closes: number[], period = 14): (number | null)[] {
  const result: (number | null)[] = new Array(closes.length).fill(null)
  if (closes.length < period + 1) return result
  let avgGain = 0, avgLoss = 0
  for (let i = 1; i <= period; i++) {
    const d = closes[i] - closes[i - 1]
    if (d > 0) avgGain += d; else avgLoss -= d
  }
  avgGain /= period; avgLoss /= period
  result[period] = avgLoss === 0 ? 100 : 100 - 100 / (1 + avgGain / avgLoss)
  for (let i = period + 1; i < closes.length; i++) {
    const d = closes[i] - closes[i - 1]
    avgGain = (avgGain * (period - 1) + (d > 0 ? d : 0)) / period
    avgLoss = (avgLoss * (period - 1) + (d < 0 ? -d : 0)) / period
    result[i] = avgLoss === 0 ? 100 : 100 - 100 / (1 + avgGain / avgLoss)
  }
  return result
}

export function computeMACD(closes: number[]): { macd: (number | null)[]; signal: (number | null)[]; hist: (number | null)[] } {
  const ema12 = computeEMA(closes, 12)
  const ema26 = computeEMA(closes, 26)
  const macdLine: (number | null)[] = closes.map((_, i) =>
    ema12[i] !== null && ema26[i] !== null ? ema12[i]! - ema26[i]! : null
  )
  const macdValues = macdLine.filter((v) => v !== null) as number[]
  const signalRaw = computeEMA(macdValues, 9)
  // Align signal back
  const signal: (number | null)[] = new Array(closes.length).fill(null)
  let si = 0
  for (let i = 0; i < closes.length; i++) {
    if (macdLine[i] !== null) { signal[i] = signalRaw[si++] ?? null }
  }
  const hist: (number | null)[] = closes.map((_, i) =>
    macdLine[i] !== null && signal[i] !== null ? macdLine[i]! - signal[i]! : null
  )
  return { macd: macdLine, signal, hist }
}

export function computeBollinger(closes: number[], period = 20, mult = 2): { upper: (number | null)[]; middle: (number | null)[]; lower: (number | null)[] } {
  const middle = computeSMA(closes, period)
  const upper: (number | null)[] = []
  const lower: (number | null)[] = []
  for (let i = 0; i < closes.length; i++) {
    if (middle[i] === null) { upper.push(null); lower.push(null); continue }
    let variance = 0
    for (let j = i - period + 1; j <= i; j++) variance += (closes[j] - middle[i]!) ** 2
    const stddev = Math.sqrt(variance / period)
    upper.push(middle[i]! + mult * stddev)
    lower.push(middle[i]! - mult * stddev)
  }
  return { upper, middle, lower }
}

export function computeStochastic(candles: Candle[], kPeriod = 14, dPeriod = 3): { k: (number | null)[]; d: (number | null)[] } {
  const kArr: (number | null)[] = []
  for (let i = 0; i < candles.length; i++) {
    if (i < kPeriod - 1) { kArr.push(null); continue }
    let hh = -Infinity, ll = Infinity
    for (let j = i - kPeriod + 1; j <= i; j++) {
      if (candles[j].high > hh) hh = candles[j].high
      if (candles[j].low < ll) ll = candles[j].low
    }
    kArr.push(hh === ll ? 50 : ((candles[i].close - ll) / (hh - ll)) * 100)
  }
  const kVals = kArr.filter((v) => v !== null) as number[]
  const dRaw = computeSMA(kVals, dPeriod)
  const dArr: (number | null)[] = new Array(candles.length).fill(null)
  let di = 0
  for (let i = 0; i < candles.length; i++) {
    if (kArr[i] !== null) dArr[i] = dRaw[di++] ?? null
  }
  return { k: kArr, d: dArr }
}

export function computeATR(candles: Candle[], period = 14): (number | null)[] {
  const tr: number[] = []
  for (let i = 0; i < candles.length; i++) {
    if (i === 0) { tr.push(candles[i].high - candles[i].low); continue }
    const hl = candles[i].high - candles[i].low
    const hc = Math.abs(candles[i].high - candles[i - 1].close)
    const lc = Math.abs(candles[i].low - candles[i - 1].close)
    tr.push(Math.max(hl, hc, lc))
  }
  const result: (number | null)[] = new Array(candles.length).fill(null)
  if (tr.length < period) return result
  let atr = 0
  for (let i = 0; i < period; i++) atr += tr[i]
  atr /= period
  result[period - 1] = atr
  for (let i = period; i < tr.length; i++) {
    atr = (atr * (period - 1) + tr[i]) / period
    result[i] = atr
  }
  return result
}

export function computeOBV(candles: Candle[]): (number | null)[] {
  const result: (number | null)[] = [0]
  for (let i = 1; i < candles.length; i++) {
    const prev = result[i - 1] ?? 0
    if (candles[i].close > candles[i - 1].close) result.push(prev + candles[i].volume)
    else if (candles[i].close < candles[i - 1].close) result.push(prev - candles[i].volume)
    else result.push(prev)
  }
  return result
}