decompositions.ts

import decompositionCoreByNameUnknown from "@openfisca/france-json/decompositions.json"
import waterfallsUnknown from "@openfisca/france-json/waterfalls.json"
export type {
  Decomposition as DecompositionCore,
  DecompositionByName as DecompositionCoreByName,
} from "@openfisca/json-model"
import type {
  Decomposition as DecompositionCore,
  DecompositionByName as DecompositionCoreByName,
  DecompositionReference,
  EntityByKey,
  Variable,
  VariableByName,
  Waterfall,
  WaterfallOptions,
} from "@openfisca/json-model"

import type { CalculationName } from "$lib/calculations"
import { reformChangesByName } from "$lib/reforms"
import type { Situation } from "$lib/situations"

export interface CalculationEvaluation {
  delta: number[]
  deltaAtVectorIndex: number
}

export type CalculationEvaluationByName = Partial<{
  [name in CalculationName]: CalculationEvaluation
}>

export interface Decomposition extends DecompositionCore {
  name: string
  open?: boolean
}

export type DecompositionByName = { [name: string]: Decomposition }

export interface Evaluation {
  calculationEvaluationByName: CalculationEvaluationByName
  fromOpenFisca?: boolean
}

export type EvaluationByName = { [name: string]: Evaluation }

export interface LatchkeyDataItem {
  aggregate?: Decomposition
  leaf: Decomposition
}

export interface VisibleDecomposition {
  decomposition: Decomposition
  depth: number
  rows: VisibleRow[]
  trunk: boolean
  variable?: Variable
  /// These children are not always the same as the one from the decomposition,
  /// because, the visible children of a decomposition, are not always the same
  /// as the children used to calculate the evaluation of the decomposition
  /// from OpenFisca variables.
  visibleChildren?: DecompositionReference[]
  visibleEvaluationByCalculationName: VisibleEvaluationByCalculationName
}

export interface VisibleEvaluation {
  delta: number[]
  deltaAtVectorIndex: number
  deltaSums: [number, number][]
  deltaSumsAtVectorIndex: [number, number]
}

export type VisibleEvaluationByCalculationName = Partial<{
  [name in CalculationName]: VisibleEvaluation
}>

export interface VisibleRow {
  calculationName: CalculationName
  deltaAtVectorIndex: number
  deltaSumsAtVectorIndex: [number, number]
}

export const decompositionCoreByName: DecompositionCoreByName =
  decompositionCoreByNameUnknown

export const decompositionCoreByNameByReformName: {
  [name: string]: DecompositionCoreByName
} = Object.fromEntries(
  Object.entries(reformChangesByName).map(([reformName, reformChanges]) => [
    reformName,
    patchDecompositionCoreByName(
      decompositionCoreByName,
      reformChanges.decompositions,
    ),
  ]),
)

export const waterfalls: Waterfall[] = waterfallsUnknown

// Note: Duplicates are removed from nonVirtualDecompositionsName, because a variable name
// may appear more than once in decomposition.
export const nonVirtualDecompositionsName = extractNonVirtualDecompositionsName(
  decompositionCoreByName,
  waterfalls,
)
export const nonVirtualDecompositionsNameByReformName: {
  [name: string]: string[]
} = Object.fromEntries(
  Object.entries(decompositionCoreByNameByReformName).map(
    ([reformName, reformDecompositionCoreByName]) => [
      reformName,
      extractNonVirtualDecompositionsName(
        reformDecompositionCoreByName,
        waterfalls,
      ),
    ],
  ),
)

export const decompositionsOptionsVariablesName = new Set<string>()
for (const decompositionCore of Object.values(decompositionCoreByName)) {
  if (decompositionCore.options === undefined) {
    continue
  }
  for (const options of decompositionCore.options) {
    if (options.waterfall !== undefined) {
      continue
    }
    for (const variableName of Object.keys(options)) {
      if (["else", "then"].includes(variableName)) {
        continue
      }
      decompositionsOptionsVariablesName.add(variableName)
    }
  }
}

export function buildDecompositionByNameFromCore(
  decompositionCoreByName: DecompositionCoreByName,
): DecompositionByName | undefined {
  return Object.fromEntries(
    Object.entries(decompositionCoreByName).map(([name, decompositionCore]) => [
      name,
      {
        ...decompositionCore,
        name,
      },
    ]),
  )
}

export function buildVisibleDecompositions(
  decompositionByName: DecompositionByName,
  entityByKey: EntityByKey,
  evaluationByName: EvaluationByName,
  situation: Situation,
  variableSummaryByName: VariableByName,
  waterfall: Waterfall,
  showNulls: boolean,
  vectorLength: number,
  year: number,
): VisibleDecomposition[] {
  const visibleDecompositions: VisibleDecomposition[] = []
  buildVisibleDecompositions1(
    decompositionByName,
    entityByKey,
    evaluationByName,
    situation,
    variableSummaryByName,
    waterfall.name,
    waterfall.root,
    showNulls,
    vectorLength,
    0,
    false,
    true,
    {},
    visibleDecompositions,
    year,
  )
  return visibleDecompositions
}

function buildVisibleDecompositions1(
  decompositionByName: DecompositionByName,
  entityByKey: EntityByKey,
  evaluationByName: EvaluationByName,
  situation: Situation,
  variableSummaryByName: VariableByName,
  waterfallName: string,
  name: string,
  showNulls: boolean,
  vectorLength: number,
  depth: number,
  negate: boolean,
  trunk = true,
  deltaSumsPreviousByCalculationName: Partial<{
    [calculationName in CalculationName]: number[]
  }>,
  visibleDecompositions: VisibleDecomposition[],
  year: number,
): number {
  const decomposition = decompositionByName[name]
  if (decomposition === undefined) {
    return -1
  }
  const evaluation = evaluationByName[name]
  if (evaluation === undefined) {
    return -1
  }

  let hidden = decomposition.hidden
  // Handle waterfall options for "hidden" attribute.
  for (const options of decomposition.options ?? []) {
    if (options.waterfall !== undefined) {
      if (
        options.then?.hidden !== undefined &&
        (options as WaterfallOptions).waterfall.includes(waterfallName)
      ) {
        hidden = options.then.hidden ?? undefined
      } else if (
        options.else?.hidden !== undefined &&
        !(options as WaterfallOptions).waterfall.includes(waterfallName)
      ) {
        hidden = options.else.hidden ?? undefined
      }
    }
  }
  // Handle variables options for "hidden" attribute.
  for (const options of decomposition.options ?? []) {
    if (options.waterfall === undefined) {
      if (
        options.then?.hidden !== undefined ||
        options.else?.hidden !== undefined
      ) {
        let allVariablesMatch = true
        for (const [variableName, variableValues] of Object.entries(options)) {
          if (["else", "then"].includes(variableName)) {
            continue
          }
          const variable = variableSummaryByName[variableName]
          const entity = entityByKey[variable.entity]
          const entitySituation = situation[entity.key_plural]
          if (entitySituation === undefined) {
            allVariablesMatch = false
            break
          }
          let variableMatch = false
          for (const population of Object.values(entitySituation)) {
            if (
              variableValues.includes(
                population[variableName]?.[year] ?? variable.default_value,
              )
            ) {
              variableMatch = true
            }
          }
          if (!variableMatch) {
            allVariablesMatch = false
            break
          }
        }
        if (options.then?.hidden !== undefined && allVariablesMatch) {
          hidden = options.then.hidden ?? undefined
        } else if (options.else?.hidden !== undefined && !allVariablesMatch) {
          hidden = options.else.hidden ?? undefined
        }
      }
    }
  }

  let visibleDecompositionIndex = -1
  const showNode =
    (showNulls && !hidden) ||
    Object.values(evaluation.calculationEvaluationByName).some(
      (deltaEvaluation) =>
        deltaEvaluation.delta.some((deltaItem) => deltaItem !== 0),
    )
  if (showNode) {
    const visibleDecomposition = {
      decomposition,
      depth,
      trunk,
      visibleEvaluationByCalculationName: Object.fromEntries(
        Object.entries(evaluation.calculationEvaluationByName).map(
          ([calculationName, { delta, deltaAtVectorIndex }]) => {
            if (negate) {
              delta = delta.map((deltaItem) => -deltaItem)
              deltaAtVectorIndex = -deltaAtVectorIndex
            }
            return [
              calculationName,
              {
                delta,
                deltaAtVectorIndex: deltaAtVectorIndex,
              } as VisibleEvaluation,
            ]
          },
        ),
      ),
    } as VisibleDecomposition

    let visibleChildren = decomposition.children
    for (const options of decomposition.options ?? []) {
      if (options.waterfall !== undefined) {
        if (
          (options as WaterfallOptions).then?.children !== undefined &&
          (options as WaterfallOptions).waterfall.includes(waterfallName)
        ) {
          visibleChildren =
            (options as WaterfallOptions).then.children ?? undefined
        } else if (
          (options as WaterfallOptions).else?.children !== undefined &&
          !(options as WaterfallOptions).waterfall.includes(waterfallName)
        ) {
          visibleChildren =
            (options as WaterfallOptions).else.children ?? undefined
        }
      }
    }
    if (visibleChildren !== undefined) {
      visibleDecomposition.visibleChildren = visibleChildren
      // if ((trunk || depth < 1) && !decomposition.open) {
      if (trunk && !decomposition.open) {
        decomposition.open = true
      }
    } else if (decomposition.open) {
      delete decomposition.open
    }

    let childrenDepth = depth
    if (!trunk) {
      visibleDecompositionIndex = visibleDecompositions.length
      visibleDecompositions.push(visibleDecomposition)
      childrenDepth = depth + 1
    }
    if (decomposition.open && visibleChildren !== undefined) {
      let beforeChildrenVisibleDecompositionLength =
        visibleDecompositions.length
      let childDeltaSumsPreviousByCalculationName =
        deltaSumsPreviousByCalculationName
      for (const [childIndex, childReference] of visibleChildren.entries()) {
        const childVisibleDecompositionIndex = buildVisibleDecompositions1(
          decompositionByName,
          entityByKey,
          evaluationByName,
          situation,
          variableSummaryByName,
          waterfallName,
          childReference.name,
          showNulls,
          vectorLength,
          childrenDepth,
          negate ? !childReference.negate : childReference.negate,
          trunk &&
            visibleDecompositions.length ===
              beforeChildrenVisibleDecompositionLength,
          childDeltaSumsPreviousByCalculationName,
          visibleDecompositions,
          year,
        )
        if (childVisibleDecompositionIndex < 0) {
          continue
        }
        const childVisibleDecomposition =
          visibleDecompositions[childVisibleDecompositionIndex]
        childDeltaSumsPreviousByCalculationName = Object.fromEntries(
          Object.entries(
            childVisibleDecomposition.visibleEvaluationByCalculationName,
          ).map(([calculationName, childVisibleEvaluation]) => [
            calculationName,
            childVisibleEvaluation.deltaSums.map((itemValue) => itemValue[1]),
          ]),
        )
      }
    }
    if (trunk) {
      visibleDecompositionIndex = visibleDecompositions.length
      visibleDecompositions.push(visibleDecomposition)
    }

    const vectorIndex = situation.slider?.vectorIndex ?? 0
    for (const [calculationName, visibleEvaluation] of Object.entries(
      visibleDecomposition.visibleEvaluationByCalculationName,
    )) {
      const deltaSumsPrevious =
        deltaSumsPreviousByCalculationName[calculationName] ??
        new Array(vectorLength).fill(0)
      const deltaSums = deltaSumsPrevious.map((previousItemValue, index) => [
        previousItemValue,
        previousItemValue + visibleEvaluation.delta[index],
      ]) as [number, number][]
      visibleEvaluation.deltaSums = deltaSums
      visibleEvaluation.deltaSumsAtVectorIndex =
        vectorIndex < deltaSums.length ? deltaSums[vectorIndex] : [0, 0]
    }

    const variable = variableSummaryByName[name]
    if (variable !== undefined) {
      visibleDecomposition.variable = variable
    }

    const visibleEvaluationByCalculationName =
      visibleDecomposition.visibleEvaluationByCalculationName
    const rows = (visibleDecomposition.rows = [])

    const lawVisibleEvaluation = visibleEvaluationByCalculationName.law
    if (lawVisibleEvaluation !== undefined) {
      const lawRow: VisibleRow = {
        calculationName: "law",
        deltaAtVectorIndex: lawVisibleEvaluation.deltaAtVectorIndex,
        deltaSumsAtVectorIndex: lawVisibleEvaluation.deltaSumsAtVectorIndex,
      }
      rows.push(lawRow)

      const billVisibleEvaluation = visibleEvaluationByCalculationName.bill
      let previousRow: VisibleRow
      let previousVisibleEvaluation: VisibleEvaluation
      if (billVisibleEvaluation === undefined) {
        previousRow = lawRow
        previousVisibleEvaluation = lawVisibleEvaluation
      } else {
        if (
          billVisibleEvaluation.deltaAtVectorIndex ===
          lawVisibleEvaluation.deltaAtVectorIndex
        ) {
          lawRow.deltaSumsAtVectorIndex =
            billVisibleEvaluation.deltaSumsAtVectorIndex
          previousRow = lawRow
        } else {
          lawRow.deltaSumsAtVectorIndex = [
            billVisibleEvaluation.deltaSumsAtVectorIndex[0],
            billVisibleEvaluation.deltaSumsAtVectorIndex[0] +
              lawVisibleEvaluation.deltaAtVectorIndex,
          ]
          const billRow: VisibleRow = {
            calculationName: "bill",
            deltaAtVectorIndex: billVisibleEvaluation.deltaAtVectorIndex,
            deltaSumsAtVectorIndex:
              billVisibleEvaluation.deltaSumsAtVectorIndex,
          }
          rows.push(billRow)
          previousRow = billRow
        }
        previousVisibleEvaluation = billVisibleEvaluation
      }

      const amendmentVisibleEvaluation =
        visibleEvaluationByCalculationName.amendment
      if (amendmentVisibleEvaluation !== undefined) {
        // Law + bill + amendement

        if (
          amendmentVisibleEvaluation.deltaAtVectorIndex ===
          previousVisibleEvaluation.deltaAtVectorIndex
        ) {
          previousRow.deltaSumsAtVectorIndex =
            amendmentVisibleEvaluation.deltaSumsAtVectorIndex
        } else {
          previousRow.deltaSumsAtVectorIndex = [
            amendmentVisibleEvaluation.deltaSumsAtVectorIndex[0],
            amendmentVisibleEvaluation.deltaSumsAtVectorIndex[0] +
              previousVisibleEvaluation.deltaAtVectorIndex,
          ]
          const amendmentRow: VisibleRow = {
            calculationName: "amendment",
            deltaAtVectorIndex: amendmentVisibleEvaluation.deltaAtVectorIndex,
            deltaSumsAtVectorIndex:
              amendmentVisibleEvaluation.deltaSumsAtVectorIndex,
          }
          rows.push(amendmentRow)
        }
      }
    }
  }

  return visibleDecompositionIndex
}

function extractNonVirtualDecompositionsName(
  decompositionCoreByName: DecompositionCoreByName,
  waterfalls: Waterfall[],
): string[] {
  const nonVirtualDecompositionsName: string[] = []
  for (const { name: waterfallName, root } of waterfalls) {
    for (const [name, decomposition] of walkDecompositionsCore(
      decompositionCoreByName,
      waterfallName,
      root,
      true,
    )) {
      if (
        !decomposition.virtual &&
        !nonVirtualDecompositionsName.includes(name)
      ) {
        nonVirtualDecompositionsName.push(name)
      }
    }
  }
  return nonVirtualDecompositionsName
}

// export function* iterDecompositionAncestorsName(
//   decompositionByName: DecompositionByName,
//   name: string | undefined | null,
// ): Generator<string, void, unknown> {
//   if (name == null) {
//     return
//   }
//   const decomposition = decompositionByName[name]
//   if (decomposition === undefined) {
//     return
//   }
//   yield* iterDecompositionAncestorsName(
//     decompositionByName,
//     decomposition.parentName,
//   )
//   yield name
// }

// export function* iterDecompositionChildren(
//   decompositionByName: DecompositionByName,
//   decomposition: Decomposition,
// ): Generator<Decomposition, void, unknown> {
//   for (const childReference of decomposition.children ?? []) {
//     const child = decompositionByName[childReference.name]
//     if (child !== undefined) {
//       yield child
//     }
//   }
// }

function patchDecompositionCoreByName(
  decompositionCoreByName: DecompositionCoreByName,
  patch: { [name: string]: DecompositionCore | null },
): DecompositionCoreByName {
  if (Object.keys(patch).length === 0) {
    return decompositionCoreByName
  }
  const patchedDecompositionCoreByName = { ...decompositionCoreByName }
  for (const [name, decompositionCorePatch] of Object.entries(patch)) {
    if (decompositionCorePatch === null) {
      // This case should not occur, because a reform should always
      // have all the decompositions of the original tax-benefit system.
      delete patchedDecompositionCoreByName[name]
    } else {
      patchedDecompositionCoreByName[name] = decompositionCorePatch
    }
  }
  return patchedDecompositionCoreByName
}

export function updateEvaluations(
  decompositionByName: DecompositionByName,
  evaluationByName: EvaluationByName,
  vectorIndex: number,
  vectorLength: number,
  waterfalls: Waterfall[],
): EvaluationByName {
  const newEvaluationByName = { ...evaluationByName }
  const updatedNames = new Set<string>()
  for (const { root } of waterfalls) {
    updateEvaluations1(
      decompositionByName,
      evaluationByName,
      root,
      newEvaluationByName,
      updatedNames,
      vectorIndex,
      vectorLength,
    )
  }
  return newEvaluationByName
}

function updateEvaluations1(
  decompositionByName: DecompositionByName,
  evaluationByName: EvaluationByName,
  name: string,
  newEvaluationByName: EvaluationByName,
  updatedNames: Set<string>,
  vectorIndex: number,
  vectorLength: number,
): Evaluation | undefined {
  if (updatedNames.has(name)) {
    return newEvaluationByName[name]
  }
  updatedNames.add(name)

  const decomposition = decompositionByName[name]
  if (decomposition === undefined) {
    return undefined
  }

  // Note: Don't use children defined in decompositions options here,
  // because those children are used for rendering, not to calculate
  // from OpenFisca variables.
  const children = decomposition.children
  const evaluation = evaluationByName[name]
  const newCalculationEvaluationByName: Partial<{
    [calculationName in CalculationName]: Partial<CalculationEvaluation>
  }> = {}
  if (children !== undefined) {
    for (const childReference of children) {
      const childEvaluation = updateEvaluations1(
        decompositionByName,
        evaluationByName,
        childReference.name,
        newEvaluationByName,
        updatedNames,
        vectorIndex,
        vectorLength,
      )
      // When variable is not (yet?) computed by OpenFisca, compute its delta
      // from the delta of its children.
      if (!evaluation?.fromOpenFisca && childEvaluation !== undefined) {
        for (const [
          calculationName,
          childCalculationEvaluation,
        ] of Object.entries(childEvaluation.calculationEvaluationByName)) {
          let newCalculationEvaluation =
            newCalculationEvaluationByName[calculationName]
          if (newCalculationEvaluation === undefined) {
            newCalculationEvaluation = newCalculationEvaluationByName[
              calculationName
            ] = {
              delta: new Array(vectorLength).fill(0),
            }
          }
          const delta = newCalculationEvaluation.delta
          if (childReference.negate) {
            delta.map(
              (_deltaAtIndex, index) =>
                (delta[index] -= childCalculationEvaluation.delta[index]),
            )
          } else {
            delta.map(
              (_deltaAtIndex, index) =>
                (delta[index] += childCalculationEvaluation.delta[index]),
            )
          }
        }
      }
    }
  }

  // If variable is computed by OpenFisca, recopy (and rescale if needed) its
  // computed delta
  if (evaluation?.fromOpenFisca) {
    for (const [calculationName, deltaEvaluation] of Object.entries(
      evaluation.calculationEvaluationByName,
    )) {
      newCalculationEvaluationByName[calculationName] = {
        delta:
          deltaEvaluation.delta.length === vectorLength
            ? deltaEvaluation.delta
            : new Array(vectorLength).fill(deltaEvaluation.deltaAtVectorIndex),
      }
    }
  }

  // Extract deltaAtIndex from delta.
  for (const [calculationName, newCalculationEvaluation] of Object.entries(
    newCalculationEvaluationByName,
  )) {
    const oldCalculationEvaluation =
      evaluation?.calculationEvaluationByName[calculationName]
    newCalculationEvaluation.deltaAtVectorIndex =
      vectorIndex < newCalculationEvaluation.delta.length
        ? newCalculationEvaluation.delta[vectorIndex]
        : oldCalculationEvaluation === undefined
        ? 0
        : oldCalculationEvaluation.deltaAtVectorIndex
  }

  const newEvaluation = {
    ...(evaluation ?? {}),
    calculationEvaluationByName:
      newCalculationEvaluationByName as CalculationEvaluationByName,
  }
  newEvaluationByName[name] = newEvaluation
  return newEvaluation
}

export function updateEvaluationsVectorIndex(
  evaluationByName: EvaluationByName,
  vectorIndex: number,
): EvaluationByName {
  let changed = false
  const newEvaluationByName: EvaluationByName = {}
  for (const [name, evaluation] of Object.entries(evaluationByName)) {
    let evaluationChanged = false
    const newCalculationEvaluationByName = {
      ...evaluation.calculationEvaluationByName,
    }
    for (const [
      calculationName,
      { delta, deltaAtVectorIndex },
    ] of Object.entries(newCalculationEvaluationByName)) {
      if (vectorIndex < delta.length) {
        const newDeltaAtVectorIndex = delta[vectorIndex]
        if (newDeltaAtVectorIndex != deltaAtVectorIndex) {
          newCalculationEvaluationByName[calculationName] = {
            delta,
            deltaAtVectorIndex: newDeltaAtVectorIndex,
          }
          evaluationChanged = true
        }
      }
    }
    newEvaluationByName[name] = evaluationChanged
      ? {
          ...evaluation,
          calculationEvaluationByName: newCalculationEvaluationByName,
        }
      : evaluation
    if (evaluationChanged) {
      changed = true
    }
  }
  return changed ? newEvaluationByName : evaluationByName
}

// export function* walkDecompositions(
//   decompositionByName: DecompositionByName,
//   name: string,
//   depthFirst: boolean,
//   openOnly: boolean,
// ): Generator<Decomposition, void, unknown> {
//   const decomposition = decompositionByName[name]
//   if (decomposition === undefined) {
//     return
//   }
//   if (!depthFirst) {
//     yield decomposition
//   }
//   if (
//     decomposition.children !== undefined &&
//     (decomposition.open || !openOnly)
//   ) {
//     for (const childReference of decomposition.children) {
//       yield* walkDecompositions(
//         decompositionByName,
//         childReference.name,
//         depthFirst,
//         openOnly,
//       )
//     }
//   }
//   if (depthFirst) {
//     yield decomposition
//   }
// }

export function* walkDecompositionsCore(
  decompositionCoreByName: DecompositionCoreByName,
  waterfallName: string,
  name: string,
  depthFirst: boolean,
): Generator<[string, DecompositionCore], void, unknown> {
  const decompositionCore = decompositionCoreByName[name]
  if (decompositionCore === undefined) {
    return
  }
  if (!depthFirst) {
    yield [name, decompositionCore]
  }

  let children = decompositionCore.children
  for (const options of decompositionCore.options ?? []) {
    if (options.waterfall !== undefined) {
      if (
        (options as WaterfallOptions).then?.children !== undefined &&
        (options as WaterfallOptions).waterfall.includes(waterfallName)
      ) {
        children = (options as WaterfallOptions).then.children ?? undefined
      } else if (
        (options as WaterfallOptions).else?.children !== undefined &&
        !(options as WaterfallOptions).waterfall.includes(waterfallName)
      ) {
        children = (options as WaterfallOptions).else.children ?? undefined
      }
    }
  }
  if (children !== undefined) {
    for (const childReference of children) {
      yield* walkDecompositionsCore(
        decompositionCoreByName,
        waterfallName,
        childReference.name,
        depthFirst,
      )
    }
  }

  if (depthFirst) {
    yield [name, decompositionCore]
  }
}

export function* walkDecompositionsCoreName(
  decompositionCoreByName: DecompositionCoreByName,
  waterfallName: string,
  name: string,
  depthFirst: boolean,
): Generator<string, void, unknown> {
  for (const [decompositionName] of walkDecompositionsCore(
    decompositionCoreByName,
    waterfallName,
    name,
    depthFirst,
  )) {
    yield decompositionName
  }
}