import BN from "bn.js";
import { multiplyWad, divideWad } from "./wad.js";
import { LN_2_WAD } from "../constants.js";

/**
 * Exponential and Logarithm Module
 *
 * Precise exp and ln using BN arithmetic, matching on-chain Rust implementation.
 */

const WAD_BN = new BN("1000000000000000000");
const LN_2_WAD_BN = new BN(LN_2_WAD.toString());
const MAX_EXP_INPUT = WAD_BN.mul(new BN(40));

/**
 * Compute exp(x) using Taylor series (12 iterations)
 * @param x - Input in WAD format
 * @returns e^x in WAD format
 */
export function expWad(x: BN): BN {
  if (x.gte(MAX_EXP_INPUT)) {
    throw new Error("expWad: input too large (x >= 40 * WAD)");
  }

  if (x.isZero()) {
    return WAD_BN.clone();
  }

  if (x.isNeg()) {
    const expPositive = expWad(x.abs());
    if (expPositive.isZero()) {
      throw new Error("expWad: division by zero");
    }
    return divideWad(WAD_BN, expPositive);
  }

  let sum = WAD_BN.clone();
  let term = WAD_BN.clone();

  for (let n = 1; n <= 12; n++) {
    term = multiplyWad(term, x);
    term = term.div(new BN(n));

    if (term.isZero()) break;

    sum = sum.add(term);
  }

  return sum;
}

/**
 * Compute ln(x) using range reduction and Newton-Raphson (3 iterations)
 * @param x - Input in WAD format (must be >= WAD)
 * @returns ln(x) in WAD format
 */
export function lnWad(x: BN): BN {
  if (x.isZero() || x.isNeg()) {
    throw new Error("lnWad: input must be positive");
  }

  if (x.lt(WAD_BN)) {
    throw new Error("lnWad: input must be >= WAD (value >= 1)");
  }

  if (x.eq(WAD_BN)) {
    return new BN(0);
  }

  let result = new BN(0);
  let val = x.clone();

  const TWO = new BN(2);
  const twoWad = WAD_BN.mul(TWO);

  while (val.gte(twoWad)) {
    val = val.div(TWO);
    result = result.add(LN_2_WAD_BN);
  }

  const yNumerator = val.sub(WAD_BN);
  const yDenominator = val.add(WAD_BN);
  const y = divideWad(yNumerator, yDenominator);
  let approxLnVal = y.mul(TWO);

  for (let i = 0; i < 3; i++) {
    const expApprox = expWad(approxLnVal);
    if (expApprox.isZero()) {
      throw new Error("lnWad: expWad returned zero");
    }

    if (expApprox.gt(val)) {
      const diff = expApprox.sub(val);
      const term = multiplyWad(diff, WAD_BN);
      const correction = divideWad(term, expApprox);
      approxLnVal = approxLnVal.sub(correction);
    } else {
      const diff = val.sub(expApprox);
      const term = multiplyWad(diff, WAD_BN);
      const correction = divideWad(term, expApprox);
      approxLnVal = approxLnVal.add(correction);
    }
  }

  result = result.add(approxLnVal);
  return result;
}