Added tracking of New-PI credit usage for second month

process_report/process_report.py

@@ -10,6 +10,15 @@
 import pyarrow
 
 
+### PI file field names
+PI_PI_FIELD = "PI"
+PI_FIRST_MONTH = "First Invoice Month"
+PI_INITIAL_CREDITS = "Initial Credits"
+PI_1ST_USED = "1st Month Used"
+PI_2ND_USED = "2nd Month Used"
+###
+
+
 ### Invoice field names
 INVOICE_DATE_FIELD = "Invoice Month"
 PROJECT_FIELD = "Project - Allocation"
@@ -51,24 +60,25 @@ def load_institute_map() -> dict:
     return institute_map
 
 
-def load_old_pis(old_pi_file):
-    old_pi_dict = dict()
-
+def load_old_pis(old_pi_file) -> pandas.DataFrame:
     try:
-        with open(old_pi_file) as f:
-            for pi_info in f:
-                pi, first_month = pi_info.strip().split(",")
-                old_pi_dict[pi] = first_month
+        old_pi_df = pandas.read_csv(
+            old_pi_file,
+            converters={
+                PI_INITIAL_CREDITS: Decimal,
+                PI_1ST_USED: Decimal,
+                PI_2ND_USED: Decimal,
+            },
+        )
     except FileNotFoundError:
         sys.exit("Applying credit 0002 failed. Old PI file does not exist")
 
-    return old_pi_dict
+    return old_pi_df
+
 
+def dump_old_pis(old_pi_file, old_pi_df: pandas.DataFrame):
+    old_pi_df.to_csv(old_pi_file, index=False)
 
-def dump_old_pis(old_pi_file, old_pi_dict: dict):
-    with open(old_pi_file, "w") as f:
-        for pi, first_month in old_pi_dict.items():
-            f.write(f"{pi},{first_month}\n")
 
 
 def load_alias(alias_file):
@@ -86,22 +96,29 @@ def load_alias(alias_file):
     return alias_dict
 
 
-def is_old_pi(old_pi_dict, pi, invoice_month):
-    first_invoice_month = old_pi_dict.get(pi, invoice_month)
-    if compare_invoice_month(first_invoice_month, invoice_month):
+def get_pi_age(old_pi_df: pandas.DataFrame, pi, invoice_month):
+    """Returns time difference between current invoice month and PI's first invoice month
+    I.e 0 for new PIs
+
+    Will raise an error if the PI'a age is negative, which suggests a faulty invoice, or a program bug"""
+    first_invoice_month = old_pi_df.loc[old_pi_df[PI_PI_FIELD] == pi, PI_FIRST_MONTH]
+    if first_invoice_month.empty:
+        return 0
+
+    month_diff = get_month_diff(invoice_month, first_invoice_month.iat[0])
+    if month_diff < 0:
         sys.exit(
             f"PI {pi} from {first_invoice_month} found in {invoice_month} invoice!"
         )
-    if compare_invoice_month(invoice_month, first_invoice_month):
-        return True
-    return False
+    else:
+        return month_diff
 
 
-def compare_invoice_month(month_1, month_2):
-    """Returns True if 1st date is later than 2nd date"""
+def get_month_diff(month_1, month_2):
+    """Returns a positive integer if month_1 is ahead in time of month_2"""
     dt1 = datetime.datetime.strptime(month_1, "%Y-%m")
     dt2 = datetime.datetime.strptime(month_2, "%Y-%m")
-    return dt1 > dt2
+    return (dt1.year - dt2.year) * 12 + (dt1.month - dt2.month)
 
 
 def get_invoice_bucket():
@@ -384,28 +401,57 @@ def fetch_s3_alias_file():
 
 def apply_credits_new_pi(dataframe, old_pi_file):
     new_pi_credit_code = "0002"
-    new_pi_credit_amount = 1000
+    INITIAL_CREDIT_AMOUNT = 1000
     EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"]
 
     dataframe[CREDIT_FIELD] = None
     dataframe[CREDIT_CODE_FIELD] = None
     dataframe[BALANCE_FIELD] = Decimal(0)
 
-    old_pi_dict = load_old_pis(old_pi_file)
+    old_pi_df = load_old_pis(old_pi_file)
 
-    current_pi_list = dataframe[PI_FIELD].unique()
+    current_pi_set = set(dataframe[PI_FIELD])
     invoice_month = dataframe[INVOICE_DATE_FIELD].iat[0]
+    invoice_pis = old_pi_df[old_pi_df[PI_FIRST_MONTH] == invoice_month]
+    if invoice_pis[PI_INITIAL_CREDITS].empty or pandas.isna(
+        new_pi_credit_amount := invoice_pis[PI_INITIAL_CREDITS].iat[0]
+    ):
+        new_pi_credit_amount = INITIAL_CREDIT_AMOUNT
+
+    print(f"New PI Credit set at {new_pi_credit_amount} for {invoice_month}")
 
-    for pi in current_pi_list:
+    for pi in current_pi_set:
         pi_projects = dataframe[dataframe[PI_FIELD] == pi]
+        pi_age = get_pi_age(old_pi_df, pi, invoice_month)
+        pi_old_pi_entry = old_pi_df.loc[old_pi_df[PI_PI_FIELD] == pi].squeeze()
 
-        if is_old_pi(old_pi_dict, pi, invoice_month):
+        if pi_age > 1:
             for i, row in pi_projects.iterrows():
                 dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD]
         else:
-            old_pi_dict[pi] = invoice_month
-            print(f"Found new PI {pi}")
-            remaining_credit = new_pi_credit_amount
+            if pi_age == 0:
+                if len(pi_old_pi_entry) == 0:
+                    pi_entry = [pi, invoice_month, new_pi_credit_amount, 0, 0]
+                    old_pi_df = pandas.concat(
+                        [
+                            pandas.DataFrame([pi_entry], columns=old_pi_df.columns),
+                            old_pi_df,
+                        ],
+                        ignore_index=True,
+                    )
+                    pi_old_pi_entry = old_pi_df.loc[
+                        old_pi_df[PI_PI_FIELD] == pi
+                    ].squeeze()
+
+                remaining_credit = new_pi_credit_amount
+                credit_used_field = PI_1ST_USED
+            elif pi_age == 1:
+                remaining_credit = (
+                    pi_old_pi_entry[PI_INITIAL_CREDITS] - pi_old_pi_entry[PI_1ST_USED]
+                )
+                credit_used_field = PI_2ND_USED
+
+            initial_credit = remaining_credit
             for i, row in pi_projects.iterrows():
                 if remaining_credit == 0 or row[SU_TYPE_FIELD] in EXCLUDE_SU_TYPES:
                     dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD]
@@ -418,7 +464,18 @@ def apply_credits_new_pi(dataframe, old_pi_file):
                     dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD] - applied_credit
                     remaining_credit -= applied_credit
 
-    dump_old_pis(old_pi_file, old_pi_dict)
+            credits_used = initial_credit - remaining_credit
+            if (pi_old_pi_entry[credit_used_field] != 0) and (
+                credits_used != pi_old_pi_entry[credit_used_field]
+            ):
+                print(
+                    f"Warning: PI file overwritten. PI {pi} previously used ${pi_old_pi_entry[credit_used_field]} of New PI credits, now uses ${credits_used}"
+                )
+            old_pi_df.loc[
+                old_pi_df[PI_PI_FIELD] == pi, credit_used_field
+            ] = credits_used
+
+    dump_old_pis(old_pi_file, old_pi_df)
 
     return dataframe