situations.ts

import type {
  Entity,
  EntityByKey,
  GroupEntity,
  Variable,
} from "@openfisca/json-model"
import { getRolePersonsIdKey } from "@openfisca/json-model"

import { entityByKey } from "$lib/entities"
import testCasesCoreUnknown from "$lib/openfisca/test_cases.json"
import type { VariableValues } from "$lib/variables"
import { variableSummaryByName } from "$lib/variables"

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 type Situation = {
  [entityKeyPlural: string]: { [populationId: string]: PopulationWithoutId }
} & { sliders?: Slider[] }

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

export interface Slider {
  entity: string // Entity key
  id: string // Population ID
  max: number
  min: number
  name: string // Variable name
  // steps: number // Number of steps to use between min & max; currently always 101.
}

export const testCasesCore = testCasesCoreUnknown as unknown as Situation[]

export function buildTestCasesWithoutNonInputVariables(
  entityByKey: EntityByKey,
  inputInstantsByVariableNameArray: Array<{
    [name: string]: Set<string>
  }>,
  testCases: Situation[],
): Situation[] {
  const entities = Object.values(entityByKey)
  testCases = [...testCases]
  for (let [situationIndex, situation] of testCases.entries()) {
    situation = testCases[situationIndex] = { ...situation }
    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 {
  if (value == null) {
    value = variable.default_value
  }
  const entity = entityByKey[variable.entity]
  const entitySituation = situation[entity.key_plural]
  const existingValueByPeriod = entitySituation?.[populationId]?.[variable.name]
  if (
    existingValueByPeriod !== undefined &&
    existingValueByPeriod[year - 2] === value &&
    existingValueByPeriod[year - 1] === value &&
    existingValueByPeriod[year] === value
  ) {
    return situation
  }
  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 - 2]: value,
          [year - 1]: value,
          [year]: value,
        },
      },
    },
  }
}

export function updateTestCasesVariableValues(
  testCases: Situation[],
  variableName: string,
  variableValuesByName: { [name: string]: VariableValues },
  vectorIndexes: number[],
  year: number,
): Situation[] {
  const entityKey = variableSummaryByName[variableName]?.entity
  if (entityKey === undefined) {
    return testCases
  }
  const entity = entityByKey[entityKey]
  const values = variableValuesByName[variableName]
  if (values === undefined) {
    // Occurs for variables that are in test cases and are not calculated.
    // For example: variable "depcom_foyer".
    // For these variables, the situations shouldn't be updated.
    return testCases
  }

  let testCasesPopulationCount = 0
  for (const situation of testCases) {
    const entitySituation = situation[entity.key_plural] ?? {}
    const populationCount = Object.keys(entitySituation).length
    testCasesPopulationCount += populationCount
  }

  let populationIndex = 0
  return testCases.map((situation, situationIndex) => {
    const baseIndex = vectorIndexes[situationIndex] * testCasesPopulationCount
    const entitySituation = situation[entity.key_plural] ?? {}
    for (const [populationId, population] of Object.entries(
      entitySituation,
    ).sort(([populationId1], [populationId2]) =>
      populationId1.localeCompare(populationId2),
    )) {
      entitySituation[populationId] = {
        ...(population ?? {}),
        [variableName]: {
          ...((population[variableName] ?? {}) as {
            [date: string]: boolean | number | string | null
          }),
          [year]: values[baseIndex + populationIndex],
        },
      }
      populationIndex++
    }
    return {
      ...situation,
      [entity.key_plural]: entitySituation,
    }
  })
}