Probability and Combinatorics Functions

import {
  // Combinatorics
  combinations, factorial, gamma, lgamma, multinomial, permutations,
  // Special combinatorics  
  bellNumbers, catalan, composition, stirlingS2,
  // Random functions
  pickRandom, random, randomInt,
  // Distance/divergence
  kldivergence,
  // Set operations (related to combinatorics)
  setCartesian, setDifference, setDistinct, setIntersect, 
  setIsSubset, setMultiplicity, setPowerset, setSize,
  setSymDifference, setUnion
} from 'mathjs'

combinations(n: MathType, k: MathType): MathType

// "n choose k" - number of ways to choose k items from n items
combinations(5, 2) // 10 (ways to choose 2 from 5)
combinations(10, 3) // 120
combinations(4, 4) // 1 (only one way to choose all items)
combinations(4, 0) // 1 (only one way to choose none)

// Large numbers using BigNumber
combinations(bignumber('100'), bignumber('50')) // Very large result

// Properties
combinations(n, k) === combinations(n, subtract(n, k)) // Symmetry
combinations(n, 0) === 1 // Base case
combinations(n, n) === 1 // Base case

// Pascal's triangle relationship  
combinations(5, 2) === add(combinations(4, 1), combinations(4, 2)) // 10 = 4 + 6

permutations(n: MathType, k?: MathType): MathType

// Number of ways to arrange k items from n items
permutations(5, 2) // 20 (5 * 4)
permutations(5, 5) // 120 (5!)
permutations(10, 3) // 720 (10 * 9 * 8)

// Without second argument: n!
permutations(5) // 120 (same as factorial(5))
permutations(0) // 1

// Relationship to combinations
permutations(n, k) === multiply(combinations(n, k), factorial(k))

// Circular permutations
function circularPermutations(n) {
  return divide(factorial(n), n) // (n-1)!
}
circularPermutations(4) // 6 (ways to arrange 4 people around a table)

factorial(n: MathType): MathType

// Standard factorial: n! = n * (n-1) * ... * 2 * 1
factorial(0) // 1 (by definition)
factorial(1) // 1  
factorial(5) // 120
factorial(10) // 3628800

// Large factorials with BigNumber
factorial(bignumber('50')) // Very large result, maintains precision

// Double factorial (not built-in, custom implementation)
function doubleFactorial(n) {
  if (smaller(n, 0)) throw new Error('Negative input')
  if (smallerEq(n, 1)) return 1
  return multiply(n, doubleFactorial(subtract(n, 2)))
}
doubleFactorial(8) // 384 (8 * 6 * 4 * 2)
doubleFactorial(9) // 945 (9 * 7 * 5 * 3 * 1)

// Subfactorial (derangements)
function subfactorial(n) {
  if (equal(n, 0)) return 1
  if (equal(n, 1)) return 0
  return multiply(subtract(n, 1), add(subfactorial(subtract(n, 1)), subfactorial(subtract(n, 2))))
}

multinomial(a: MathCollection): MathType

// Multinomial coefficient: n! / (k1! * k2! * ... * km!)
// where n = k1 + k2 + ... + km

multinomial([2, 3, 4]) // 1260 = 9! / (2! * 3! * 4!)
multinomial([1, 1, 1]) // 6 = 3! / (1! * 1! * 1!)

// Number of ways to partition n objects into groups of sizes k1, k2, ..., km
// Example: Ways to divide 9 people into groups of 2, 3, and 4
multinomial([2, 3, 4]) // 1260 ways

// Relationship to binomial coefficients (special case)
multinomial([k, subtract(n, k)]) === combinations(n, k)

bellNumbers(n: MathType): MathType

// Number of ways to partition a set of n elements
bellNumbers(0) // 1 (empty set has one partition)
bellNumbers(1) // 1 (one element: {1})
bellNumbers(2) // 2 (two elements: {{1},{2}} or {{1,2}})
bellNumbers(3) // 5 ({{1},{2},{3}}, {{1,2},{3}}, {{1,3},{2}}, {{2,3},{1}}, {{1,2,3}})
bellNumbers(4) // 15
bellNumbers(5) // 52

// Bell's triangle (similar to Pascal's triangle)
function bellTriangle(rows) {
  const triangle = [[1]]
  
  for (let i = 1; i < rows; i++) {
    const row = [triangle[i-1][triangle[i-1].length - 1]]
    
    for (let j = 1; j <= i; j++) {
      row.push(add(row[j-1], triangle[i-1][j-1]))
    }
    
    triangle.push(row)
  }
  
  return triangle
}

catalan(n: MathType): MathType

// nth Catalan number: C_n = (2n)! / ((n+1)! * n!)
catalan(0) // 1
catalan(1) // 1
catalan(2) // 2  
catalan(3) // 5
catalan(4) // 14
catalan(5) // 42

// Applications:
// - Number of binary trees with n internal nodes
// - Number of ways to triangulate a polygon with n+2 sides  
// - Number of paths from (0,0) to (n,n) that don't cross the diagonal
// - Number of ways to parenthesize n+1 factors

// Recurrence relation: C_0 = 1, C_{n+1} = sum(C_i * C_{n-i}) for i=0 to n
function catalanRecurrence(n) {
  if (equal(n, 0)) return 1
  let sum = 0
  for (let i = 0; i < n; i++) {
    sum = add(sum, multiply(catalan(i), catalan(subtract(n, add(i, 1)))))
  }
  return sum
}

stirlingS2(n: MathType, k: MathType): MathType

// Number of ways to partition n objects into k non-empty subsets
stirlingS2(4, 2) // 7 (ways to partition 4 objects into 2 groups)
stirlingS2(5, 3) // 25
stirlingS2(n, 1) // 1 (only one way: all objects in one subset)
stirlingS2(n, n) // 1 (only one way: each object in its own subset)

// Relationship to Bell numbers
function bellFromStirling(n) {
  let sum = 0
  for (let k = 0; k <= n; k++) {
    sum = add(sum, stirlingS2(n, k))
  }
  return sum
}
bellFromStirling(4) === bellNumbers(4) // true

// Recurrence: S(n,k) = k*S(n-1,k) + S(n-1,k-1)

composition(n: MathType, k: MathType): MathType

// Number of ways to write n as a sum of k positive integers (order matters)
composition(4, 2) // 3 (4 = 1+3, 2+2, 3+1)
composition(5, 3) // 6 (5 = 1+1+3, 1+2+2, 1+3+1, 2+1+2, 2+2+1, 3+1+1)

// Formula: C(n,k) = C(n-1, k-1) where C is combinations
// This is because we place k-1 dividers among n-1 positions
composition(n, k) === combinations(subtract(n, 1), subtract(k, 1))

// Total compositions of n (into any number of parts)
function totalCompositions(n) {
  return pow(2, subtract(n, 1))
}
totalCompositions(4) // 8 (2^3)

gamma(n: MathType): MathType

// Gamma function: Γ(n) = (n-1)! for positive integers
gamma(1) // 1 = 0!
gamma(2) // 1 = 1!
gamma(3) // 2 = 2!
gamma(5) // 24 = 4!

// For non-integers
gamma(0.5) // √π ≈ 1.772
gamma(1.5) // 0.5 * √π ≈ 0.886
gamma(2.5) // 1.5 * 0.5 * √π ≈ 1.329

// Complex numbers supported
gamma(complex(1, 1)) // Complex result

// Relationship to factorial
gamma(add(n, 1)) === factorial(n) // For non-negative integers

// Reflection formula: Γ(z) * Γ(1-z) = π / sin(πz)
multiply(gamma(z), gamma(subtract(1, z))) === divide(pi, sin(multiply(pi, z)))

lgamma(n: MathType): MathType

// Natural logarithm of the gamma function (more numerically stable)
lgamma(5) // log(24) ≈ 3.178 = log(gamma(5))
lgamma(100) // Much more stable than log(gamma(100))

// Useful for large arguments where gamma(n) would overflow
lgamma(bignumber('1000')) // Computable, whereas gamma(1000) is huge

// Relationship
lgamma(n) === log(gamma(n)) // But lgamma is more numerically stable

// Stirling's approximation for large n
function stirlingApprox(n) {
  return add(
    multiply(subtract(n, 0.5), log(n)),
    multiply(-1, n),
    multiply(0.5, log(multiply(2, pi)))
  )
}
// lgamma(n) ≈ stirlingApprox(n) for large n

random(size?: number | number[], min?: number, max?: number): number | Matrix

// Single random number [0, 1)
random() // e.g., 0.7234...

// Random number in range [min, max)
random(5, 10) // e.g., 7.8234...

// Array of random numbers
random(5) // [0.1234, 0.5678, ...] (5 numbers)
random([2, 3]) // 2x3 matrix of random numbers

// Seeded random (for reproducibility)
import { create, all, config } from 'mathjs'
const math = create(all, {
  randomSeed: 'deterministic-seed'
})
math.random() // Reproducible sequence

randomInt(size?: number | number[], min?: number, max?: number): number | Matrix

// Single random integer [min, max)
randomInt(1, 7) // Dice roll: 1, 2, 3, 4, 5, or 6
randomInt(0, 2) // Coin flip: 0 or 1

// Array of random integers
randomInt(10, 1, 11) // 10 numbers from 1 to 10
randomInt([3, 4], 0, 100) // 3x4 matrix of integers [0, 100)

// Uniform distribution over integers
function uniformInt(n, trials = 1000) {
  const counts = new Array(n).fill(0)
  for (let i = 0; i < trials; i++) {
    counts[randomInt(0, n)]++
  }
  return counts
}

pickRandom(array: MathCollection, number?: number, weights?: MathCollection): MathType | MathCollection

const items = ['apple', 'banana', 'cherry', 'date', 'elderberry']

// Pick single random element
pickRandom(items) // e.g., 'cherry'

// Pick multiple elements (with replacement)
pickRandom(items, 3) // e.g., ['banana', 'apple', 'banana']

// Weighted random selection
const weights = [0.1, 0.3, 0.4, 0.15, 0.05] // Must sum to 1
pickRandom(items, 1, weights) // More likely to pick 'cherry' (weight 0.4)

// Multiple weighted picks
pickRandom(items, 10, weights) // 10 picks with given probabilities

// Lottery/sampling applications
function lottery(tickets, winners) {
  return pickRandom(tickets, winners) // Pick winners from ticket holders
}

// Bootstrap sampling
function bootstrap(data, samples = 1000) {
  return range(0, samples).map(() => pickRandom(data, data.length))
}

// Binomial probability mass function
function binomialPMF(k, n, p) {
  return multiply(combinations(n, k), multiply(pow(p, k), pow(subtract(1, p), subtract(n, k))))
}

// Poisson probability mass function
function poissonPMF(k, lambda) {
  return multiply(divide(pow(lambda, k), factorial(k)), exp(multiply(-1, lambda)))
}

// Geometric probability mass function
function geometricPMF(k, p) {
  return multiply(pow(subtract(1, p), subtract(k, 1)), p)
}

// Examples
binomialPMF(3, 10, 0.5) // P(X=3) in Binomial(10, 0.5)
poissonPMF(2, 3) // P(X=2) in Poisson(3)
geometricPMF(5, 0.2) // P(X=5) in Geometric(0.2)

// Standard normal PDF (approximation)
function normalPDF(x, mu = 0, sigma = 1) {
  const coefficient = divide(1, multiply(sigma, sqrt(multiply(2, pi))))
  const exponent = exp(divide(multiply(-0.5, pow(divide(subtract(x, mu), sigma), 2)), 1))
  return multiply(coefficient, exponent)
}

// Standard normal CDF (approximation using error function)
function normalCDF(x, mu = 0, sigma = 1) {
  const z = divide(subtract(x, mu), sigma)
  return multiply(0.5, add(1, erf(divide(z, sqrt(2)))))
}

kldivergence(q: MathCollection, p: MathCollection): MathType

// KL divergence: D_KL(P||Q) = sum(p_i * log(p_i / q_i))
// Measures how one probability distribution differs from another

const p = [0.5, 0.3, 0.2] // True distribution  
const q = [0.4, 0.4, 0.2] // Approximating distribution

kldivergence(p, q) // KL divergence from P to Q

// Properties:
// - D_KL(P||Q) >= 0 (Gibbs' inequality)
// - D_KL(P||Q) = 0 iff P = Q
// - Generally D_KL(P||Q) ≠ D_KL(Q||P) (not symmetric)

// Cross entropy
function crossEntropy(p, q) {
  return add(entropy(p), kldivergence(p, q))
}

// Entropy (custom implementation)
function entropy(p) {
  return multiply(-1, sum(p.map(x => multiply(x, log(x)))))
}

setUnion(a1: MathCollection, a2: MathCollection): MathCollection
setIntersect(a1: MathCollection, a2: MathCollection): MathCollection  
setDifference(a1: MathCollection, a2: MathCollection): MathCollection
setSymDifference(a1: MathCollection, a2: MathCollection): MathCollection

const A = [1, 2, 3, 4]
const B = [3, 4, 5, 6]

setUnion(A, B) // [1, 2, 3, 4, 5, 6]
setIntersect(A, B) // [3, 4]
setDifference(A, B) // [1, 2] (A - B)
setSymDifference(A, B) // [1, 2, 5, 6] (A ∪ B) - (A ∩ B)

// Principle of inclusion-exclusion
setSize(setUnion(A, B)) === add(
  setSize(A), 
  setSize(B), 
  multiply(-1, setSize(setIntersect(A, B)))
)

setCartesian(a1: MathCollection, a2: MathCollection): MathCollection
setPowerset(a: MathCollection): MathCollection
setDistinct(a: MathCollection): MathCollection

const A = [1, 2]
const B = ['a', 'b', 'c']

// Cartesian product
setCartesian(A, B) // [[1,'a'], [1,'b'], [1,'c'], [2,'a'], [2,'b'], [2,'c']]
setSize(setCartesian(A, B)) === multiply(setSize(A), setSize(B)) // 2 * 3 = 6

// Power set (all subsets)
setPowerset([1, 2, 3]) // [[], [1], [2], [3], [1,2], [1,3], [2,3], [1,2,3]]
setSize(setPowerset(A)) === pow(2, setSize(A)) // 2^n subsets

// Remove duplicates
setDistinct([1, 2, 2, 3, 3, 3]) // [1, 2, 3]

setIsSubset(a1: MathCollection, a2: MathCollection): boolean
setMultiplicity(e: MathType, a: MathCollection): number
setSize(a: MathCollection): number

const A = [1, 2, 3]
const B = [1, 2, 3, 4, 5]

setIsSubset(A, B) // true (A ⊆ B)
setIsSubset(B, A) // false

setMultiplicity(2, [1, 2, 2, 2, 3]) // 3 (element 2 appears 3 times)
setSize([1, 2, 3, 4, 5]) // 5 (cardinality)

// Set equality
function setEqual(A, B) {
  return setSize(setSymDifference(A, B)) === 0
}

// Simulate coin flips
function coinFlips(n, p = 0.5) {
  return range(0, n).map(() => random() < p ? 1 : 0)
}

// Estimate π using Monte Carlo
function estimatePi(samples = 1000000) {
  let insideCircle = 0
  
  for (let i = 0; i < samples; i++) {
    const x = random(-1, 1)
    const y = random(-1, 1) 
    if (add(pow(x, 2), pow(y, 2)) < 1) {
      insideCircle++
    }
  }
  
  return multiply(4, divide(insideCircle, samples))
}

// Birthday paradox simulation
function birthdayParadox(people = 23, trials = 10000) {
  let matches = 0
  
  for (let trial = 0; trial < trials; trial++) {
    const birthdays = randomInt(people, 1, 366)
    const unique = setDistinct(birthdays)
    if (setSize(unique) < people) {
      matches++
    }
  }
  
  return divide(matches, trials)
}

// Expected value for discrete distribution
function expectedValue(values, probabilities) {
  return sum(values.map((val, i) => multiply(val, probabilities[i])))
}

// Variance for discrete distribution  
function varianceDiscrete(values, probabilities) {
  const mu = expectedValue(values, probabilities)
  return sum(values.map((val, i) => 
    multiply(probabilities[i], pow(subtract(val, mu), 2))
  ))
}

// Probability generating function
function probabilityGeneratingFunction(probabilities, z) {
  return sum(probabilities.map((p, k) => multiply(p, pow(z, k))))
}

// Moment generating function
function momentGeneratingFunction(values, probabilities, t) {
  return sum(values.map((val, i) => 
    multiply(probabilities[i], exp(multiply(t, val)))
  ))
}

// Classic problems using combinatorics

// Hypergeometric distribution
function hypergeometric(k, N, K, n) {
  // k successes in n draws from population of N with K successes
  return divide(
    multiply(combinations(K, k), combinations(subtract(N, K), subtract(n, k))),
    combinations(N, n)
  )
}

// Derangements probability
function derangementProb(n) {
  return divide(subfactorial(n), factorial(n))
}

// Matching problem (random permutation has fixed points)
function matchingProb(n, k) {
  // Probability that exactly k elements are in their original positions
  return divide(
    multiply(combinations(n, k), subfactorial(subtract(n, k))),
    factorial(n)
  )
}

// Use BigNumber for large factorials and combinations
const largeCombination = combinations(bignumber('1000'), bignumber('500'))

// Use lgamma for very large gamma function values
const largeGamma = exp(lgamma(bignumber('1000'))) // More stable than gamma(1000)

// Stirling's approximation for very large factorials
function stirlingFactorial(n) {
  return multiply(
    sqrt(multiply(2, pi, n)),
    pow(divide(n, e), n)
  )
}

// Use log-space calculations for small probabilities
function logBinomialPMF(k, n, p) {
  return add(
    lgamma(add(n, 1)),
    multiply(-1, lgamma(add(k, 1))),
    multiply(-1, lgamma(subtract(n, k, -1))),
    multiply(k, log(p)),
    multiply(subtract(n, k), log(subtract(1, p)))
  )
}

// Convert back to probability space
const logProb = logBinomialPMF(3, 10, 0.5)
const prob = exp(logProb)

// Chain operations with probability functions
const diceRolls = chain(range(1, 7))
  .map(face => divide(1, 6)) // Equal probability for each face
  .done()

// Functional approach to combinatorial calculations  
const pascalRow = n => 
  range(0, add(n, 1))
    .map(k => combinations(n, k))
    
const fibonacciFromBinomial = n =>
  range(0, add(n, 1))
    .map(k => combinations(subtract(n, k), k))
    .reduce((a, b) => add(a, b), 0)