Flat tax on gross income

Fixes #4830

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    added:
+    - Flat tax on gross income.

policyengine_us/parameters/gov/contrib/ubi_center/flat_tax/flat_tax_on_gross_income.yaml

@@ -0,0 +1,7 @@
+description: The flat tax rate is applied on gross income if this is true.
+values:
+  0000-01-01: false
+metadata:
+  unit: bool
+  label: Flat tax on gross income
+  period: year

policyengine_us/tests/policy/contrib/ubi_center/flat_tax.yaml

@@ -0,0 +1,19 @@
+- name: Flat tax applies on AGI
+  period: 2023
+  input:
+    gov.contrib.ubi_center.flat_tax.flat_tax_on_gross_income: false
+    gov.contrib.ubi_center.flat_tax.rate: 0.1
+    positive_agi: 2_000
+    irs_gross_income: 5_000
+  output:
+    flat_tax: 200
+
+- name: Flat tax applies on gross income
+  period: 2023
+  input:
+    gov.contrib.ubi_center.flat_tax.flat_tax_on_gross_income: true
+    gov.contrib.ubi_center.flat_tax.rate: 0.1
+    positive_agi: 2_000
+    irs_gross_income: 5_000
+  output:
+    flat_tax: 500

policyengine_us/variables/contrib/ubi_center/flat_tax.py

@@ -6,10 +6,12 @@ class flat_tax(Variable):
     entity = TaxUnit
     label = "Flat tax"
     unit = USD
-    documentation = "Flat income tax on federal AGI."
+    documentation = "Flat income tax on federal AGI or gross income."
     definition_period = YEAR
 
     def formula(tax_unit, period, parameters):
-        agi = tax_unit("positive_agi", period)
+        income = tax_unit("positive_agi", period)
         p = parameters(period).gov.contrib.ubi_center.flat_tax
-        return p.rate * agi
+        if p.flat_tax_on_gross_income:
+            income = add(tax_unit, period, ["irs_gross_income"])
+        return p.rate * income