export type Vector<D extends number, I = number> = I[] & { length: D };
export type Matrix<R extends number, C extends number, I = number> = Vector<R, Vector<C, I>>;

export function vec<D extends number, I = number>(s: D, f: (i: number) => I): Vector<D, I> {
  return Array.from({ length: s }, (_, i) => f(i)) as Vector<D, I>;
}

export function mat<R extends number, C extends number, I = number>(r: R, c: C, f: (x: number, y: number) => number): Matrix<R, C, I> {
  return Array.from({ length: r }, (_, y) => Array.from({ length: c }, (_, x) => f(x, y))) as Matrix<R, C, I>;
}

export function sum<R extends number, C extends number>(l: Matrix<R, C>, r: Matrix<R, C>): Matrix<R, C> {
  const result: Matrix<R, C> = mat(l.length, l[0].length, () => 0);
  for (let i = 0; i < l.length; i++) {
    for (let j = 0; j < r[0].length; j++) {
      result[i][j] += l[i][j] + r[i][j];
    }
  }
  return result;
}

export function mul<R extends number, C extends number>(l: Matrix<R, C>, r: number): Matrix<R, C>;
export function mul<D0 extends number, D1 extends number>(l: Vector<D0>, r: Matrix<D0, D1>): Vector<D1>;
export function mul<R extends number, S extends number, C extends number>(l: Matrix<R, S>, r: Matrix<S, C>): Matrix<R, C>;
export function mul(l: Matrix<number, number> | Vector<number>, r: Matrix<number, number> | number): Matrix<number, number> | Vector<number> {
  if (typeof r === 'number') {
    const lm: Matrix<number, number> = l as Matrix<number, number>;
    const result = Array.from({ length: lm.length }, () => Array(lm[0].length).fill(0));
    for (let i = 0; i < lm.length; i++) {
      for (let j = 0; j < lm[i].length; j++) {
        result[i][j] += lm[i][j] * r;
      }
    }
    return result;
  } else if (!Array.isArray(l[0])) {
    const lv: Vector<number> = l as Vector<number>;
    const result = Array(r[0].length).fill(0);
    for (let j = 0; j < r[0].length; j++) {
      for (let i = 0; i < lv.length; i++) {
        result[j] += lv[i] * r[i][j];
      }
    }

    return result;
  } else {
    const lm: Matrix<number, number> = l as Matrix<number, number>;
    const result = Array.from({ length: lm.length }, () => Array(r[0].length).fill(0));
    for (let i = 0; i < lm.length; i++) {
      for (let j = 0; j < r[0].length; j++) {
        for (let k = 0; k < r.length; k++) {
          result[i][j] += lm[i][k] * r[k][j];
        }
      }
    }
    return result;
  }
}

export function transpose<R extends number, C extends number>(m: Matrix<R, C>): Matrix<C, R> {
  return mat(m[0].length, m.length, (x, y) => m[x][y]);
}

export function softmax(row: number[]): number[] {
  const max = Math.max(...row);
  const exps = row.map(v => Math.exp(v - max));
  const sum = exps.reduce((a, b) => a + b, 0);
  return exps.map(v => v / sum);
}

export function msoftmax<R extends number, C extends number>(m: Matrix<R, C>): Matrix<R, C> {
  return m.map(softmax) as Matrix<R, C>;
}

export function norm<R extends number, C extends number>(x: Matrix<R, C>, epsilon = 1e-5): Matrix<R, C> {
  return x.map(vec => {
    const mean = vec.reduce((sum, val) => sum + val, 0) / vec.length;
    const variance = vec.reduce((sum, val) => sum + (val - mean) ** 2, 0) / vec.length;
    return vec.map(val => (val - mean) / Math.sqrt(variance + epsilon));
  }) as Matrix<R, C>;
}

export function relu<D extends number>(vec: Vector<D>): Vector<D> {
  return vec.map(v => Math.max(0, v)) as Vector<D>;
}

export function mean<D extends number>(v: Vector<D>): number {
  return v.reduce((s, c) => s + c, 0) / v.length;
}

export function drm(data: number[], interval = 0.9): number {
  const sorted = [...data].sort((a, b) => a - b);
  const count = Math.floor(data.length * interval);
  let mi = 0;
  let mv = Infinity;

  for (let i = 0; i <= data.length - count; i++) {
    const range = sorted[i + count - 1] - sorted[i];
    if (range < mv) {
      mv = range;
      mi = i;
    }
  }

  return sorted.slice(mi, mi + count).reduce((a, b) => a + b, 0) / count;
}

export function randseq(n: number): number[] {
  const is = Array.from({ length: n }, (_, i) => i);
  for (let i = n - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * (i + 1));
    [is[i], is[j]] = [is[j], is[i]]; // swap
  }
  return is;
}

export function stddev(xs: number[]): number {
  const m = mean(xs);
  return Math.sqrt(xs.reduce((s, v) => s + (v - m) ** 2, 0) / xs.length);
}

export function ema(xs: number[], alpha = 0.1): number {
  let result = xs[0];
  for (let i = 1; i < xs.length; i++) {
    result = alpha * xs[i] + (1 - alpha) * result;
  }
  return result;
}

export function slope(xs: number[]): number {
  const xMean = (xs.length - 1) / 2;
  const yMean = mean(xs);

  let num = 0, den = 0;
  for (let i = 0; i < xs.length; i++) {
    const x = i - xMean;
    const y = xs[i] - yMean;
    num += x * y;
    den += x * x;
  }

  return num / den;
}