import BN from "bn.js";
import { toWad, multiplyWad, divideWad } from "./wad.js";
import { expWad, lnWad } from "./expLn.js";
import { PoolType, AffectedIndices } from "../types.js";
import { LN_4_WAD } from "../constants.js";

/**
 * LMSR (Logarithmic Market Scoring Rule) Module
 *
 * CRITICAL: resolveLmsrIndices must match programs/pit/src/math/lmsr.rs exactly
 */

/**
 * Resolve LMSR indices - maps token type to affected q_vector indices
 * Must match on-chain implementation in programs/pit/src/math/lmsr.rs EXACTLY
 *
 * High Pool (strikes ascending: [150, 160, 170]):
 * - HIT strike_index 0: [1,2,3] (S1,S2,S3)
 * - HIT strike_index 1: [2,3]   (S2,S3)
 * - HIT strike_index 2: [3]     (S3)
 * - MISS strike_index 0: [0]       (S0)
 * - MISS strike_index 1: [0,1]     (S0,S1)
 * - MISS strike_index 2: [0,1,2]   (S0,S1,S2)
 *
 * Low Pool (strikes descending: [160, 150, 140]):
 * - HIT strike_index 0: [1,2,3] (S1,S2,S3)
 * - HIT strike_index 1: [2,3]   (S2,S3)
 * - HIT strike_index 2: [3]     (S3)
 * - MISS strike_index 0: [0]       (S0)
 * - MISS strike_index 1: [0,1]     (S0,S1)
 * - MISS strike_index 2: [0,1,2]   (S0,S1,S2)
 */
export function resolveLmsrIndices(
  poolType: PoolType,
  strikeIndex: number,
  isHit: boolean
): AffectedIndices {
  // Validate strike index
  if (strikeIndex < 0 || strikeIndex > 2) {
    throw new Error(`Invalid strike index: ${strikeIndex}`);
  }

  // Both pools now use same mapping due to Low Pool strikes stored in descending order
  if (isHit) {
    switch (strikeIndex) {
      case 0:
        return { indices: [1, 2, 3, 0], count: 3 }; // S1, S2, S3
      case 1:
        return { indices: [2, 3, 0, 0], count: 2 }; // S2, S3
      case 2:
        return { indices: [3, 0, 0, 0], count: 1 }; // S3
      default:
        throw new Error(`Invalid strike index: ${strikeIndex}`);
    }
  } else {
    switch (strikeIndex) {
      case 0:
        return { indices: [0, 0, 0, 0], count: 1 }; // S0
      case 1:
        return { indices: [0, 1, 0, 0], count: 2 }; // S0, S1
      case 2:
        return { indices: [0, 1, 2, 0], count: 3 }; // S0, S1, S2
      default:
        throw new Error(`Invalid strike index: ${strikeIndex}`);
    }
  }
}

/**
 * Calculate LMSR cost: C(q) = b * ln(Σ exp(qi/b))
 * Uses Log-Sum-Exp trick for numerical stability
 */
export function lmsrCost(qVector: BN[], b: BN): BN {
  const scaledValues = qVector.map((q) => divideWad(q, b));
  const maxScaled = scaledValues.reduce((max, v) => (v.gt(max) ? v : max));

  let sumExp = new BN(0);
  for (const scaled of scaledValues) {
    const diff = scaled.sub(maxScaled);
    const expVal = expWad(diff);
    sumExp = sumExp.add(expVal);
  }

  const lnSum = lnWad(sumExp);
  return multiplyWad(b, maxScaled.add(lnSum));
}

/**
 * Simulate a trade and return cost/refund
 */
export function simulateTrade(
  qVector: BN[],
  b: BN,
  affectedIndices: AffectedIndices,
  amount: BN,
  isBuy: boolean
): { cost: BN; newQVector: BN[] } {
  const newQVector = [...qVector];
  const amountWad = toWad(amount);

  for (let i = 0; i < affectedIndices.count; i++) {
    const idx = affectedIndices.indices[i];
    if (isBuy) {
      newQVector[idx] = newQVector[idx].add(amountWad);
    } else {
      newQVector[idx] = newQVector[idx].sub(amountWad);
    }
  }

  const oldCost = lmsrCost(qVector, b);
  const newCost = lmsrCost(newQVector, b);
  const cost = isBuy ? newCost.sub(oldCost) : oldCost.sub(newCost);

  return { cost, newQVector };
}

/**
 * Calculate token price (marginal cost for 1 token)
 */
export function tokenPrice(
  qVector: BN[],
  b: BN,
  poolType: PoolType,
  strikeIndex: number,
  isHit: boolean
): BN {
  const affectedIndices = resolveLmsrIndices(poolType, strikeIndex, isHit);
  let totalProb = new BN(0);

  const expValues = qVector.map((q) => expWad(divideWad(q, b)));
  const sumExp = expValues.reduce((sum, v) => sum.add(v), new BN(0));

  if (sumExp.isZero()) return new BN(0);

  for (let i = 0; i < affectedIndices.count; i++) {
    const idx = affectedIndices.indices[i];
    const prob = divideWad(expValues[idx], sumExp);
    totalProb = totalProb.add(prob);
  }

  return totalProb;
}

/**
 * Calculate liquidity parameter b from total subsidy
 * Formula: b = (subsidy / 2) / ln(4)
 */
export function calculateLiquidityParameter(totalSubsidy: BN): BN {
  const poolSubsidy = totalSubsidy.div(new BN(2));
  const poolSubsidyWad = toWad(poolSubsidy);
  return divideWad(poolSubsidyWad, new BN(LN_4_WAD.toString()));
}