From 59dd51e8ff6d49a56310d1deddfbe825dab6c236 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              | 53 ++++++++++---------
 2 files changed, 75 insertions(+), 27 deletions(-)

diff --git a/src/__tests__/data/reformDefinitionCode.test.js b/src/__tests__/data/reformDefinitionCode.test.js
index a73e7733e..35e28e1aa 100644
--- a/src/__tests__/data/reformDefinitionCode.test.js
+++ b/src/__tests__/data/reformDefinitionCode.test.js
@@ -211,7 +211,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);
   });
@@ -219,13 +219,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));
@@ -242,10 +252,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")',
+    );
+  });
   test("If dataset provided, return lines with dataset", () => {
     const output = getImplementationCode(
       "policy",
diff --git a/src/data/reformDefinitionCode.js b/src/data/reformDefinitionCode.js
index 02959ba6b..8a4b046f0 100644
--- a/src/data/reformDefinitionCode.js
+++ b/src/data/reformDefinitionCode.js
@@ -28,7 +28,7 @@ export function getReproducibilityCodeBlock(
       householdInput,
       earningVariation,
     ),
-    ...getImplementationCode(type, region, year, policy, dataset),
+    ...getImplementationCode(type, region, metadata, year, policy, dataset),
   ];
 }
 
@@ -160,10 +160,13 @@ export function getSituationCode(
 export function getImplementationCode(
   type,
   region,
+  metadata,
   timePeriod,
   policy,
   dataset,
 ) {
+  const countryId = metadata.economy_options.region[0].name;
+
   if (type !== "policy") {
     return [];
   }
@@ -176,34 +179,34 @@ export function getImplementationCode(
   const hasDatasetSpecified =
     Object.keys(DEFAULT_DATASETS).includes(dataset) || region === "enhanced_us";
 
-  let formattedDataset = null;
-  if (region === "enhanced_us") {
-    formattedDataset = "enhanced_cps_2024";
-  } else if (hasDatasetSpecified) {
-    formattedDataset = DEFAULT_DATASETS[dataset];
-  }
+  // TODO: Check into how DEFAULT_DATASETS works now post-rebase
+
+  const isState =
+    countryId === "us" &&
+    region !== "us" &&
+    region !== "enhanced_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='${formattedDataset}'`
-        : hasBaseline
-          ? `reform=baseline`
-          : hasDatasetSpecified
-            ? `dataset='${formattedDataset}'`
-            : ""
-    })`,
-    `reformed = Microsimulation(${
-      hasDatasetSpecified && hasReform
-        ? `reform=reform, dataset='${formattedDataset}'`
-        : hasReform
-          ? `reform=reform`
-          : hasDatasetSpecified
-            ? `dataset='${formattedDataset}'`
-            : ""
-    })`,
+    `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",