/**
 * Punnett engine: combine two parent genotypes into the offspring genotype
 * distribution, with exact fractional probabilities.
 *
 * Per locus, each parent contributes one of its two alleles with prob 1/2,
 * giving 4 equally-likely gametic combinations that we fold into unordered
 * (canonically-ordered) allele pairs. The full-genotype distribution is the
 * product across all loci.
 *
 * Wildcard ("?") alleles are expanded uniformly over the locus' allele set
 * before combining, so a parent known only by phenotype can still be paired.
 */
import { add, frac, multiply, ONE, type Fraction } from './fraction'
import { LOCI, LOCUS_ORDER, type LocusKey } from './loci'
import {
  canonicalPair,
  toDisplayString,
  WILDCARD,
  type AllelePair,
  type Genotype,
} from './genotype'

/** A probability distribution over outcomes of type T (keyed by a string). */
export interface DistEntry<T> {
  readonly value: T
  readonly probability: Fraction
}

/** Expand a (possibly wildcard) parent allele pair into weighted concrete alleles. */
function parentAlleleWeights(locus: LocusKey, pair: AllelePair): Map<string, Fraction> {
  const weights = new Map<string, Fraction>()
  const addWeight = (allele: string, w: Fraction) => {
    weights.set(allele, add(weights.get(allele) ?? frac(0, 1), w))
  }
  const alleles = LOCI[locus].alleles
  for (const a of pair) {
    if (a === WILDCARD) {
      // Unknown allele: uniform over the locus set, each contributing 1/2 of the gamete.
      const share = frac(1, 2 * alleles.length)
      for (const concrete of alleles) addWeight(concrete, share)
    } else {
      addWeight(a, frac(1, 2))
    }
  }
  return weights
}

/** Offspring allele-pair distribution at a single locus. */
export function combineLocus(
  locus: LocusKey,
  father: AllelePair,
  mother: AllelePair,
): Map<string, DistEntry<AllelePair>> {
  const fa = parentAlleleWeights(locus, father)
  const mo = parentAlleleWeights(locus, mother)
  const result = new Map<string, DistEntry<AllelePair>>()
  for (const [fAllele, fw] of fa) {
    for (const [mAllele, mw] of mo) {
      const pair = canonicalPair(locus, fAllele, mAllele)
      const key = pair[0] + pair[1]
      const p = multiply(fw, mw)
      const existing = result.get(key)
      result.set(key, {
        value: pair,
        probability: existing ? add(existing.probability, p) : p,
      })
    }
  }
  return result
}

/** Full offspring genotype distribution (Cartesian product across loci). */
export function combineGenotypes(
  father: Genotype,
  mother: Genotype,
): DistEntry<Genotype>[] {
  // Start with one certain empty partial genotype, then fold in each locus.
  let dist: DistEntry<Record<LocusKey, AllelePair>>[] = [
    { value: {} as Record<LocusKey, AllelePair>, probability: ONE },
  ]
  for (const locus of LOCUS_ORDER) {
    const locusDist = combineLocus(locus, father[locus], mother[locus])
    const next: DistEntry<Record<LocusKey, AllelePair>>[] = []
    for (const partial of dist) {
      for (const entry of locusDist.values()) {
        next.push({
          value: { ...partial.value, [locus]: entry.value },
          probability: multiply(partial.probability, entry.probability),
        })
      }
    }
    dist = next
  }
  return dist.map((e) => ({ value: e.value as Genotype, probability: e.probability }))
}

/** Merge distribution entries that share the same display string, summing probs. */
export function mergeByGenotype(dist: DistEntry<Genotype>[]): DistEntry<Genotype>[] {
  const merged = new Map<string, DistEntry<Genotype>>()
  for (const entry of dist) {
    const key = toDisplayString(entry.value)
    const existing = merged.get(key)
    merged.set(key, {
      value: entry.value,
      probability: existing ? add(existing.probability, entry.probability) : entry.probability,
    })
  }
  return Array.from(merged.values())
}