From 599aaeac58e0c876c651122dbe10177d449b0b0b Mon Sep 17 00:00:00 2001
From: Anthony Volk <anth.volk@gmail.com>
Date: Fri, 29 Nov 2024 18:30:55 +0100
Subject: [PATCH] fix: Display pooled 3-year CPS as dataset for state-level
 runs

---
 .../data/reformDefinitionCode.test.js         | 49 ++++++++++++++++-
 src/data/reformDefinitionCode.js              | 55 +++++++++++--------
 2 files changed, 80 insertions(+), 24 deletions(-)

diff --git a/src/__tests__/data/reformDefinitionCode.test.js b/src/__tests__/data/reformDefinitionCode.test.js
index 2c3bae72f..808c9f929 100644
--- a/src/__tests__/data/reformDefinitionCode.test.js
+++ b/src/__tests__/data/reformDefinitionCode.test.js
@@ -210,7 +210,7 @@ describe("Test getSituationCode", () => {
 });
 describe("Test getImplementationCode", () => {
   test("If not a policy type, return empty array", () => {
-    const output = getImplementationCode("household", "us", 2024);
+    const output = getImplementationCode("household", "us", metadataUS, 2024);
     expect(output).toBeInstanceOf(Array);
     expect(output.length).toBe(0);
   });
@@ -218,13 +218,23 @@ describe("Test getImplementationCode", () => {
     const output = getImplementationCode(
       "policy",
       "uk",
+      metadataUK,
       2024,
       baselinePolicyUK,
     );
     expect(output).toBeInstanceOf(Array);
+    // Ensure we don't code for reversed baseline and reform
     expect(output).not.toContain(
       "baseline = Microsimulation(reform=baseline_reform)",
     );
+    // Ensure we don't code for state-level runs
+    expect(output).not.toContain(
+      'baseline = Microsimulation(dataset="pooled_3_year_cps_2023")',
+    );
+    // And ensure we don't code for both simultaneously
+    expect(output).not.toContain(
+      'baseline = Microsimulation(reform=baseline_reform, dataset="pooled_3_year_cps_2023")',
+    );
   });
   test("If set baseline, return lines with baseline", () => {
     let testPolicy = JSON.parse(JSON.stringify(baselinePolicyUK));
@@ -241,10 +251,45 @@ describe("Test getImplementationCode", () => {
         id: 1,
       },
     };
-    const output = getImplementationCode("policy", "uk", 2024, testPolicy);
+    const output = getImplementationCode(
+      "policy",
+      "uk",
+      metadataUK,
+      2024,
+      testPolicy,
+    );
     expect(output).toBeInstanceOf(Array);
     expect(output).toContain("baseline = Microsimulation(reform=baseline)");
   });
+  test("If US state, return lines with pooled 3-year CPS dataset", () => {
+    let testPolicy = JSON.parse(JSON.stringify(baselinePolicyUS));
+    testPolicy = {
+      ...testPolicy,
+      reform: {
+        data: {
+          "sample.reform.item": {
+            "2020.01.01": true,
+            "2022.01.01": true,
+          },
+        },
+      },
+    };
+
+    const output = getImplementationCode(
+      "policy",
+      "ca",
+      metadataUS,
+      2024,
+      testPolicy,
+    );
+    expect(output).toBeInstanceOf(Array);
+    expect(output).toContain(
+      'baseline = Microsimulation(dataset="pooled_3_year_cps_2023")',
+    );
+    expect(output).toContain(
+      'reformed = Microsimulation(reform=reform, dataset="pooled_3_year_cps_2023")',
+    );
+  });
 });
 
 describe("Test sanitizeStringToPython", () => {
diff --git a/src/data/reformDefinitionCode.js b/src/data/reformDefinitionCode.js
index 6c20da3d2..c984275cd 100644
--- a/src/data/reformDefinitionCode.js
+++ b/src/data/reformDefinitionCode.js
@@ -27,7 +27,7 @@ export function getReproducibilityCodeBlock(
       householdInput,
       earningVariation,
     ),
-    ...getImplementationCode(type, region, year, policy),
+    ...getImplementationCode(type, region, metadata, year, policy),
   ];
 }
 
@@ -156,7 +156,15 @@ export function getSituationCode(
   return lines;
 }
 
-export function getImplementationCode(type, region, timePeriod, policy) {
+export function getImplementationCode(
+  type,
+  region,
+  metadata,
+  timePeriod,
+  policy,
+) {
+  const countryId = metadata.economy_options.region[0].name;
+
   if (type !== "policy") {
     return [];
   }
@@ -164,31 +172,34 @@ export function getImplementationCode(type, region, timePeriod, policy) {
   const hasBaseline = Object.keys(policy?.baseline?.data).length > 0;
   const hasReform = Object.keys(policy?.reform?.data).length > 0;
 
-  // Check if the region has a dataset specified
+  // Check if the region itself has a dataset specified (currently only enhanced_us)
+  // or if the region is a US state (which uses pooled 3-year CPS)
   const hasDatasetSpecified = Object.keys(DEFAULT_DATASETS).includes(region);
-  const dataset = hasDatasetSpecified ? DEFAULT_DATASETS[region] : "";
+  const isState =
+    countryId === "us" &&
+    region !== "us" &&
+    !Object.keys(DEFAULT_DATASETS).includes(region);
+
+  let dataset = "";
+  if (hasDatasetSpecified) {
+    dataset = DEFAULT_DATASETS[region];
+  } else if (isState) {
+    dataset = "pooled_3_year_cps_2023";
+  }
+
+  const datasetSpecifier = dataset ? `dataset="${dataset}"` : "";
+
+  const baselineSpecifier = hasBaseline ? "reform=baseline" : "";
+  const baselineComma = hasBaseline && dataset ? ", " : "";
+
+  const reformSpecifier = hasReform ? "reform=reform" : "";
+  const reformComma = hasReform && dataset ? ", " : "";
 
   return [
     "",
     "",
-    `baseline = Microsimulation(${
-      hasDatasetSpecified && hasBaseline
-        ? `reform=baseline, dataset='${dataset}'`
-        : hasBaseline
-          ? `reform=baseline`
-          : hasDatasetSpecified
-            ? `dataset='${dataset}'`
-            : ""
-    })`,
-    `reformed = Microsimulation(${
-      hasDatasetSpecified && hasReform
-        ? `reform=reform, dataset='${dataset}'`
-        : hasReform
-          ? `reform=reform`
-          : hasDatasetSpecified
-            ? `dataset='${dataset}'`
-            : ""
-    })`,
+    `baseline = Microsimulation(${baselineSpecifier}${baselineComma}${datasetSpecifier})`,
+    `reformed = Microsimulation(${reformSpecifier}${reformComma}${datasetSpecifier})`,
     `baseline_income = baseline.calculate("household_net_income", period=${timePeriod || defaultYear})`,
     `reformed_income = reformed.calculate("household_net_income", period=${timePeriod || defaultYear})`,
     "difference_income = reformed_income - baseline_income",