/**
 * Lethality rules applied to an offspring genotype distribution.
 *
 * Data-driven so the biology is easy to adjust:
 *  - SpSp (Schecke homozygous): PRENATAL-LETHAL. Embryos die and are resorbed,
 *    so they never appear as live young -> removed from the distribution, which
 *    is then renormalised. Triggers a ScheckeLethal warning carrying the
 *    fraction of young lost (~1/4 for Schecke × Schecke).
 *  - ReRe (Rex homozygous): SEMI-LETHAL. Reduced viability but not modelled as a
 *    hard removal (some survive) -> kept in the distribution, warning only.
 */
import { add, divide, toString, ZERO, type Fraction } from './fraction'
import type { LocusKey } from './loci'
import type { Genotype } from './genotype'
import type { DistEntry } from './punnett'
import { GeneticsWarningCode, type GeneticsWarning } from './warnings'

interface LethalRule {
  readonly locus: LocusKey
  readonly allele: string
  /** 'lethal' = remove from live-birth distribution; 'semi' = warn only. */
  readonly kind: 'lethal' | 'semi'
  readonly warning: GeneticsWarningCode
}

const LETHAL_RULES: readonly LethalRule[] = [
  { locus: 'Sp', allele: 'Sp', kind: 'lethal', warning: GeneticsWarningCode.ScheckeLethal },
  { locus: 'Sls', allele: 'Sl', kind: 'lethal', warning: GeneticsWarningCode.SlsLethal },
  { locus: 'Re', allele: 'Re', kind: 'semi', warning: GeneticsWarningCode.RexSemiLethal },
]

function isHomozygous(g: Genotype, locus: LocusKey, allele: string): boolean {
  return g[locus][0] === allele && g[locus][1] === allele
}

export interface LethalityResult {
  /** Live-birth distribution (lethal genotypes removed, renormalised). */
  readonly distribution: DistEntry<Genotype>[]
  readonly warnings: GeneticsWarning[]
}

/**
 * Apply lethality to a (merged) genotype distribution.
 * The input probabilities are assumed to sum to 1.
 */
export function applyLethality(dist: DistEntry<Genotype>[]): LethalityResult {
  const warnings: GeneticsWarning[] = []

  // Total probability mass of fully-lethal genotypes (for the "young lost" stat).
  let lethalMass: Fraction = ZERO
  const survivors: DistEntry<Genotype>[] = []

  for (const entry of dist) {
    const hardLethal = LETHAL_RULES.some(
      (r) => r.kind === 'lethal' && isHomozygous(entry.value, r.locus, r.allele),
    )
    if (hardLethal) {
      lethalMass = add(lethalMass, entry.probability)
    } else {
      survivors.push(entry)
    }
  }

  // Renormalise survivors over the surviving mass.
  const survivingMass = survivors.reduce<Fraction>((acc, e) => add(acc, e.probability), ZERO)
  const distribution =
    survivingMass.num === 0
      ? survivors
      : survivors.map((e) => ({ value: e.value, probability: divide(e.probability, survivingMass) }))

  // One lethal warning PER lethal rule that actually removed young (so SpSp ->
  // ScheckeLethal and S(l)S(l) -> SlsLethal are reported distinctly).
  for (const rule of LETHAL_RULES) {
    if (rule.kind !== 'lethal') continue
    const mass = dist.reduce<Fraction>(
      (acc, e) => (isHomozygous(e.value, rule.locus, rule.allele) ? add(acc, e.probability) : acc),
      ZERO,
    )
    if (mass.num > 0) {
      warnings.push({
        code: rule.warning,
        detail: {
          youngLostFraction: toString(mass),
          youngLostPercent: Number(((mass.num / mass.den) * 100).toFixed(2)),
        },
      })
    }
  }

  // Semi-lethal: warn if any surviving genotype is homozygous for a semi-lethal allele.
  for (const rule of LETHAL_RULES) {
    if (rule.kind !== 'semi') continue
    const affected = distribution.reduce<Fraction>(
      (acc, e) => (isHomozygous(e.value, rule.locus, rule.allele) ? add(acc, e.probability) : acc),
      ZERO,
    )
    if (affected.num > 0) {
      warnings.push({
        code: rule.warning,
        detail: {
          affectedFraction: toString(affected),
          affectedPercent: Number(((affected.num / affected.den) * 100).toFixed(2)),
        },
      })
    }
  }

  // Superschecke: surviving young carrying BOTH spotting markers (Sp present and
  // S(l) present) are very-high-white and deafness-prone — info warning.
  const superMass = distribution.reduce<Fraction>((acc, e) => {
    const hasSp = e.value.Sp.includes('Sp')
    const hasSl = e.value.Sls.includes('Sl')
    return hasSp && hasSl ? add(acc, e.probability) : acc
  }, ZERO)
  if (superMass.num > 0) {
    warnings.push({
      code: GeneticsWarningCode.SuperscheckeDeaf,
      detail: {
        affectedFraction: toString(superMass),
        affectedPercent: Number(((superMass.num / superMass.den) * 100).toFixed(2)),
      },
    })
  }

  return { distribution, warnings }
}