// Fermi Calculator in JavaScript
"use strict";

// Constants
const N_SAMPLES = 100000; // Reduced from 10000 for better browser performance
const NORMAL90CONFIDENCE = 1.6448536269514727;

// Help message string
const HELP_MSG = "Fermi Calculator Help:\n" +
"  Operations: * / + -\n" +
"  Operands:\n" +
"    scalar (e.g., 2.5 or 3M for 3,000,000),\n" +
"    lognormal: two numbers representing low and high,\n" +
"    beta: use 'beta a b' or 'b a b',\n" +
"    mixtures: 'mx var1 weight var2 weight ...'\n" +
"  Commands:\n" +
"    help (h)         Show this help message\n" +
"    clear (c or .)   Clear stack\n" +
"    exit (e)         Exit calculator\n" +
"    stats (s)        Show sample statistics\n" +
"  Variable assignment:\n" +
"    =: varname       Assign current stack to variable\n" +
"    =. varname       Assign and clear stack\n";

// Distribution classes
class Scalar {
  constructor(value) {
    this.value = value;
  }
  sample() {
    return this.value;
  }
}

class Lognormal {
  constructor(low, high) {
    this.low = low;
    this.high = high;
  }
  sample() {
    let loglow = Math.log(this.low);
    let loghigh = Math.log(this.high);
    let mean = (loglow + loghigh) / 2;
    let sigma = (loghigh - loglow) / (2 * NORMAL90CONFIDENCE);
    let n = sampleNormal(mean, sigma);
    return Math.exp(n);
  }
}

class Beta {
  constructor(a, b) {
    this.a = a;
    this.b = b;
  }
  sample() {
    return sampleBeta(this.a, this.b);
  }
}

class FilledSamples {
  constructor(xs) {
    this.xs = xs;
  }
  
  sample() {
    return this.xs[Math.floor(Math.random() * this.xs.length)];
  }
  
  cleanup() {
    if (this.xs) {
      releaseSampleArray(this.xs);
      this.xs = null;
    }
  }
}

// Random sampling functions
function sampleNormal(mean, sigma) {
  // Box-Muller transform
  let u = Math.random();
  let v = Math.random();
  let z = Math.sqrt(-2.0 * Math.log(u)) * Math.cos(2 * Math.PI * v);
  return mean + sigma * z;
}

function sampleGamma(alpha) {
  if (alpha < 1) {
    return sampleGamma(1 + alpha) * Math.pow(Math.random(), 1 / alpha);
  }
  const d = alpha - 1/3;
  const c = 1 / Math.sqrt(9 * d);
  while (true) {
    let x = sampleNormal(0, 1);
    let v = 1 + c * x;
    if (v <= 0) continue;
    v = v * v * v;
    let u = Math.random();
    if (u < 1 - 0.0331 * Math.pow(x, 4)) return d * v;
    if (Math.log(u) < 0.5 * x * x + d * (1 - v + Math.log(v))) return d * v;
  }
}

function sampleBeta(a, b) {
  let ga = sampleGamma(a);
  let gb = sampleGamma(b);
  return ga / (ga + gb);
}

function sampleSerially(dist, n) {
  let xs = getSampleArray();
  for (let i = 0; i < n; i++) {
    xs[i] = dist.sample();
  }
  return xs;
}

function operateDistsAsSamples(dist1, dist2, op) {
  const xs = sampleSerially(dist1, N_SAMPLES);
  const ys = sampleSerially(dist2, N_SAMPLES);
  let zs = getSampleArray();
  
  try {
    for (let i = 0; i < N_SAMPLES; i++) {
      switch (op) {
        case "*":
          zs[i] = xs[i] * ys[i];
          break;
        case "/":
          if (ys[i] === 0) throw new Error("Division by zero");
          zs[i] = xs[i] / ys[i];
          break;
        case "+":
          zs[i] = xs[i] + ys[i];
          break;
        case "-":
          zs[i] = xs[i] - ys[i];
          break;
        default:
          throw new Error("Unknown operation");
      }
    }
    const result = new FilledSamples(zs);
    releaseSampleArray(xs);
    releaseSampleArray(ys);
    return result;
  } catch (e) {
    releaseSampleArray(xs);
    releaseSampleArray(ys);
    releaseSampleArray(zs);
    return { error: e.message };
  }
}

// Distribution arithmetic functions
function multiplyLogDists(l1, l2) {
  let logmean1 = (Math.log(l1.low) + Math.log(l1.high)) / 2;
  let logstd1 = (Math.log(l1.high) - Math.log(l1.low)) / (2 * NORMAL90CONFIDENCE);
  let logmean2 = (Math.log(l2.low) + Math.log(l2.high)) / 2;
  let logstd2 = (Math.log(l2.high) - Math.log(l2.low)) / (2 * NORMAL90CONFIDENCE);
  let logmean_product = logmean1 + logmean2;
  let logstd_product = Math.sqrt(logstd1 * logstd1 + logstd2 * logstd2);
  let h = logstd_product * NORMAL90CONFIDENCE;
  let loglow = logmean_product - h;
  let loghigh = logmean_product + h;
  return new Lognormal(Math.exp(loglow), Math.exp(loghigh));
}

function multiplyDists(oldDist, newDist) {
  if (oldDist instanceof Lognormal && newDist instanceof Lognormal) {
    return multiplyLogDists(oldDist, newDist);
  }
  if (oldDist instanceof Scalar && newDist instanceof Scalar) {
    return new Scalar(oldDist.value * newDist.value);
  }
  if (oldDist instanceof Scalar && oldDist.value === 1) {
    return newDist;
  }
  return operateDistsAsSamples(oldDist, newDist, "*");
}

function divideDists(oldDist, newDist) {
  if (oldDist instanceof Lognormal && newDist instanceof Lognormal) {
    if (newDist.low === 0 || newDist.high === 0) return { error: "Division by zero" };
    return multiplyLogDists(oldDist, new Lognormal(1 / newDist.high, 1 / newDist.low));
  }
  if (oldDist instanceof Scalar && newDist instanceof Scalar) {
    if (newDist.value === 0) return { error: "Division by zero" };
    return new Scalar(oldDist.value / newDist.value);
  }
  return operateDistsAsSamples(oldDist, newDist, "/");
}

function addDists(oldDist, newDist) {
  return operateDistsAsSamples(oldDist, newDist, "+");
}

function subtractDists(oldDist, newDist) {
  return operateDistsAsSamples(oldDist, newDist, "-");
}

function operateDists(oldDist, newDist, op) {
  if (!oldDist || !newDist) {
    return { error: "Invalid distribution" };
  }
  
  try {
    switch (op) {
      case "*":
        return multiplyDists(oldDist, newDist);
      case "/":
        return divideDists(oldDist, newDist);
      case "+":
        return addDists(oldDist, newDist);
      case "-":
        return subtractDists(oldDist, newDist);
      default:
        return { error: "Unknown operation" };
    }
  } catch (e) {
    return { error: e.message };
  }
}

// Global REPL state
let state = {
  currentDist: new Scalar(1),
  vars: {}
};

// Helper: parse a number with optional suffix (%, K, M, B, T)
function parseNumber(token) {
  let multiplier = 1;
  const lastChar = token[token.length - 1];
  if (lastChar === '%') multiplier = 0.01;
  else if (lastChar === 'K') multiplier = 1_000;
  else if (lastChar === 'M') multiplier = 1_000_000;
  else if (lastChar === 'B') multiplier = 1_000_000_000;
  else if (lastChar === 'T') multiplier = 1_000_000_000_000;
  let numStr = token;
  if ("KMBT%".includes(lastChar)) {
    numStr = token.slice(0, -1);
  }
  let num = parseFloat(numStr);
  if (isNaN(num)) return null;
  return num * multiplier;
}

// Simple parser: remove comments and split into words
function parseInput(line) {
  let withoutComments = line.split("#")[0];
  return withoutComments.trim().split(/\s+/);
}

// Mixture: expects syntax "mx var weight var weight ..." using variables from state.vars
function sampleMixture(distributions, weights) {
  let xs = [];
  let sum = weights.reduce((a, b) => a + b, 0);
  let cum = [];
  let running = 0;
  for (let w of weights) {
    running += w / sum;
    cum.push(running);
  }
  for (let i = 0; i < N_SAMPLES; i++) {
    let p = Math.random();
    let chosen = distributions[distributions.length - 1];
    for (let j = 0; j < cum.length; j++) {
      if (p < cum[j]) { chosen = distributions[j]; break; }
    }
    xs.push(chosen.sample());
  }
  return new FilledSamples(xs);
}

// Pretty-printer helpers
function formatNumber(n) {
  let abs = Math.abs(n);
  if (abs >= 1e12) return (n / 1e12).toFixed(2) + "T";
  if (abs >= 1e9) return (n / 1e9).toFixed(2) + "B";
  if (abs >= 1e6) return (n / 1e6).toFixed(2) + "M";
  if (abs >= 1e3) return (n / 1e3).toFixed(2) + "K";
  if (abs < 0.0001) return n.toFixed(6);
  if (abs < 0.001) return n.toFixed(5);
  if (abs < 0.01) return n.toFixed(4);
  if (abs < 0.1) return n.toFixed(3);
  return n.toFixed(2);
}

function prettyFormatDist(dist) {
  if (dist instanceof Lognormal) {
    return "=> " + formatNumber(dist.low) + " " + formatNumber(dist.high);
  } else if (dist instanceof Beta) {
    return "=> beta " + formatNumber(dist.a) + " " + formatNumber(dist.b);
  } else if (dist instanceof Scalar) {
    return "=> scalar " + formatNumber(dist.value);
  } else if (dist instanceof FilledSamples) {
    let samples = dist.xs.slice().sort((a, b) => a - b);
    let lowIdx = Math.floor(N_SAMPLES * 0.05);
    let highIdx = Math.floor(N_SAMPLES * 0.95);
    return "=> " + formatNumber(samples[lowIdx]) + " " + formatNumber(samples[highIdx]) +
           " (" + N_SAMPLES + " samples)";
  }
  return String(dist);
}

function computeStats(dist) {
  let xs = sampleSerially(dist, N_SAMPLES);
  let n = xs.length;
  let mean = xs.reduce((a, b) => a + b, 0) / n;
  let stdev = Math.sqrt(xs.reduce((a, b) => a + Math.pow(b - mean, 2), 0) / n);
  xs.sort((a, b) => a - b);
  let pct = (p) => xs[Math.floor(n * p)];
  let statsStr = "Mean: " + mean.toFixed(4) + "\n";
  statsStr += "Stdev: " + stdev.toFixed(4) + "\n";
  statsStr += "ci 5%: " + formatNumber(pct(0.05)) + "\n";
  statsStr += "ci 95%: " + formatNumber(pct(0.95)) + "\n";
  return statsStr;
}

// Append text to the output div
function appendOutput(text) {
  const outputDiv = document.getElementById("output");
  outputDiv.textContent += text + "\n";
  outputDiv.scrollTop = outputDiv.scrollHeight;
}

// Process a user command line
function processCommand(line) {
  let words = parseInput(line);
  if (words.length === 0 || words[0] === "") return;
  
  const cmd = words[0].toLowerCase();
  
  if (cmd === "exit" || cmd === "e") {
    appendOutput("Exiting... (refresh page to restart)");
    document.getElementById("input").disabled = true;
    return;
  }
  
  if (cmd === "help" || cmd === "h") {
    appendOutput(HELP_MSG);
    return;
  }
  
  if (cmd === "clear" || cmd === "c" || words[0] === ".") {
    state.currentDist = new Scalar(1);
    appendOutput("Stack cleared.");
    return;
  }
  
  if (words[0] === "=:" && words.length === 2) {
    state.vars[words[1]] = state.currentDist;
    appendOutput(words[1] + " assigned.");
    return;
  }
  
  if (words[0] === "=." && words.length === 2) {
    state.vars[words[1]] = state.currentDist;
    appendOutput(words[1] + " assigned as: " + prettyFormatDist(state.currentDist));
    state.currentDist = new Scalar(1);
    return;
  }
  
  if (cmd === "stats" || cmd === "s") {
    appendOutput(computeStats(state.currentDist));
    return;
  }
  
  // Determine operator: if first word is one of the operators, use it; else default to "*"
  let op = "*";
  if (["*", "/", "+", "-"].includes(words[0])) {
    op = words[0];
    words = words.slice(1);
  }
  
  // Parse operand based on number of remaining tokens
  let newDist;
  if (words.length === 1) {
    let token = words[0];
    if (state.vars[token] !== undefined) {
      newDist = state.vars[token];
    } else {
      let value = parseNumber(token);
      if (value === null) {
        appendOutput("Error: Invalid operand '" + token + "'");
        return;
      }
      newDist = new Scalar(value);
    }
  } else if (words.length === 2) {
    let low = parseNumber(words[0]);
    let high = parseNumber(words[1]);
    if (low === null || high === null) {
      appendOutput("Error: Invalid numbers for lognormal");
      return;
    }
    newDist = new Lognormal(low, high);
  } else if (words.length === 3) {
    if (words[0].toLowerCase() === "beta" || words[0].toLowerCase() === "b") {
      let a = parseNumber(words[1]);
      let b = parseNumber(words[2]);
      if (a === null || b === null) {
        appendOutput("Error: Invalid numbers for beta distribution");
        return;
      }
      newDist = new Beta(a, b);
    } else {
      appendOutput("Error: Command not understood");
      return;
    }
  } else if (words.length >= 4 && words[0] === "mx") {
    if ((words.length - 1) % 2 !== 0) {
      appendOutput("Error: Mixture syntax incorrect");
      return;
    }
    let distributions = [];
    let mixWeights = [];
    for (let i = 1; i < words.length; i += 2) {
      let varName = words[i];
      let weight = parseNumber(words[i+1]);
      if (state.vars[varName] === undefined) {
        appendOutput("Error: Variable '" + varName + "' not defined for mixture");
        return;
      }
      distributions.push(state.vars[varName]);
      mixWeights.push(weight);
    }
    newDist = sampleMixture(distributions, mixWeights);
  } else {
    appendOutput("Error: Input not understood");
    return;
  }
  
  let oldDist = state.currentDist;
  let result = operateDists(state.currentDist, newDist, op);
  if (result.error) {
    appendOutput("Error: " + result.error);
    return;
  }
  
  // Clean up old distribution if it was FilledSamples
  if (oldDist instanceof FilledSamples) {
    oldDist.cleanup();
  }
  
  state.currentDist = result;
  appendOutput(prettyFormatDist(state.currentDist));
}

// Array pooling
const sampleArrayPool = [];
function getSampleArray() {
  return sampleArrayPool.pop() || new Array(N_SAMPLES);
}
function releaseSampleArray(arr) {
  arr.length = 0; // Clear the array
  sampleArrayPool.push(arr);
}

// UI event bindings
document.getElementById("submit-btn").addEventListener("click", function() {
  const inputEl = document.getElementById("input");
  const cmd = inputEl.value;
  appendOutput("> " + cmd);
  processCommand(cmd);
  inputEl.value = "";
});

document.getElementById("input").addEventListener("keydown", function(e) {
  if (e.key === "Enter") {
    e.preventDefault();
    document.getElementById("submit-btn").click();
  }
});

// On load welcome message
appendOutput("Fermi Calculator Ready. Type 'help' for instructions.\n");