situations.ts

import type { Entity, EntityByKey, GroupEntity, Variable } from "@openfisca/ast"
import { getRolePersonsIdKey } from "@openfisca/ast"

export interface Axis {
  count: number
  index?: number
  max: number
  min: number
  name: string
  period?: string
}

export interface AxisDescription {
  entityKey: string
  populationId: string
  stepValue: number
  variableName: string
}

/// A population is either a group (a familly, etc) or a person.
export interface Population extends PopulationWithoutId {
  id: string
}

export interface PopulationWithoutId {
  name?: string
  [key: string]:
    | { [instant: string]: boolean | number | string | null } // variable value by instant
    | string // id & name
    | string[] // ids of persons (when key is a role and entity is not a person)
}

export interface Situation {
  [entityKeyPlural: string]: { [populationId: string]: PopulationWithoutId }
}

export type SituationWithAxes = Situation & {
  axes?: Axis[][]
}

export function buildTestCasesWithoutNonInputVariables(
  entityByKey: EntityByKey,
  inputInstantsByVariableNameArray: Array<{
    [name: string]: Set<string>
  }>,
  testCases: Situation[],
): Situation[] {
  const entities = Object.values(entityByKey)
  const situations = [...testCases]
  for (const [situationIndex, situation] of situations.entries()) {
    const inputInstantsByVariableName =
      inputInstantsByVariableNameArray[situationIndex]
    for (const entity of entities) {
      let entitySituation = situation[entity.key_plural]
      if (entitySituation === undefined) {
        continue
      }
      entitySituation = situation[entity.key_plural] = { ...entitySituation }
      const reservedKeys = getPopulationReservedKeys(entity)
      for (let [populationId, population] of Object.entries(
        entitySituation,
      ).sort(([populationId1], [populationId2]) =>
        populationId1.localeCompare(populationId2),
      )) {
        population = entitySituation[populationId] = { ...population }
        for (const [variableName, variableValueByInstant] of Object.entries(
          population,
        )) {
          if (reservedKeys.has(variableName)) {
            continue
          }
          const inputVariableValueByInstant: {
            [instant: string]: boolean | number | string | null
          } = {}
          const inputInstants =
            inputInstantsByVariableName[variableName] ?? new Set<string>()
          for (const [instant, variableValue] of Object.entries(
            variableValueByInstant,
          )) {
            if (!inputInstants.has(instant)) {
              // Remove calculated value.
              continue
            }
            inputVariableValueByInstant[instant] = variableValue
          }
          if (Object.keys(inputVariableValueByInstant).length > 0) {
            population[variableName] = inputVariableValueByInstant
          } else {
            delete population[variableName]
          }
        }
      }
    }
  }
  return testCases
}

export function getPopulationReservedKeys(entity: Entity): Set<string> {
  return new Set(
    entity.is_person
      ? ["name"]
      : [
          "name",
          ...(entity as GroupEntity).roles.map((role) =>
            getRolePersonsIdKey(role),
          ),
        ],
  )
}

export function getSituationVariableValue(
  entityByKey: EntityByKey,
  situation: Situation,
  variable: Variable,
  populationId: string,
  year: number,
): boolean | number | string {
  const entity = entityByKey[variable.entity]
  const entitySituation = situation[entity.key_plural]
  return (
    entitySituation?.[populationId]?.[variable.name]?.[year] ??
    variable.default_value
  )
}

export function indexOfSituationPopulationId(
  entity: Entity,
  situation: Situation,
  populationId: string,
): number {
  const entitySituation = situation[entity.key_plural]
  return Object.keys(entitySituation)
    .sort(([populationId1], [populationId2]) =>
      populationId1.localeCompare(populationId2),
    )
    .indexOf(populationId)
}

export function* iterSituationVariablesName(
  entityByKey: EntityByKey,
  situation: Situation,
): Generator<string, void, unknown> {
  const variablesName = new Set<string>()
  for (const entity of Object.values(entityByKey)) {
    const reservedKeys = getPopulationReservedKeys(entity)
    for (const population of Object.values(
      situation[entity.key_plural] ?? {},
    ) as PopulationWithoutId[]) {
      for (const variableName of Object.keys(population)) {
        if (reservedKeys.has(variableName)) {
          continue
        }
        if (!variablesName.has(variableName)) {
          yield variableName
          variablesName.add(variableName)
        }
      }
    }
  }
}

export function setSituationVariableValue(
  entityByKey: EntityByKey,
  situation: Situation,
  variable: Variable,
  populationId: string,
  year: number,
  value: boolean | number | string,
): Situation {
  const entity = entityByKey[variable.entity]
  const entitySituation = situation[entity.key_plural]
  return {
    ...situation,
    [entity.key_plural]: {
      ...(entitySituation ?? {}),
      [populationId]: {
        ...(entitySituation?.[populationId] ?? {}),
        [variable.name]: {
          ...((
            entitySituation?.[populationId] as {
              [key: string]: {
                [date: string]: boolean | number | string | null
              }
            }
          )?.[variable.name] ?? {}),
          [year - 3]: value ?? variable.default_value,
          [year - 2]: value ?? variable.default_value,
          [year - 1]: value ?? variable.default_value,
          [year]: value ?? variable.default_value,
        },
      },
    },
  }
}