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_REMAINDER = "1st Month Remainder"
+PI_2ND_REMAINDER = "2nd Month Remainder"
+###
+
+
 ### Invoice field names
 INVOICE_DATE_FIELD = "Invoice Month"
 PROJECT_FIELD = "Project - Allocation"
@@ -46,43 +55,54 @@ 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)
     except FileNotFoundError:
         print("Applying credit 0002 failed. Old PI file does not exist")
         sys.exit(1)
 
-    return old_pi_dict
+    return old_pi_df
 
 
-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 dump_old_pis(old_pi_file, old_pi_df: pandas.DataFrame):
+    old_pi_df.to_csv(old_pi_file, index=False)
 
 
-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 add_invoice_month(invoice_month, month_to_add):
+    """Given an invoice month string, i.e "2023-02", and an integer,
+    add as many months to the invoice month, and return the sum in same
+    format"""
+    dt = datetime.datetime.strptime(invoice_month, "%Y-%m")
+    new_month = dt.month - 1 + month_to_add
+    new_year = dt.year + new_month // 12
+    new_dt = datetime.datetime(year=new_year, month=new_month % 12 + 1, day=1)
+    return new_dt.strftime("%Y-%m")
 
 
 def get_invoice_bucket():
@@ -335,40 +355,72 @@ def validate_pi_names(dataframe):
 
 def apply_credits_new_pi(dataframe, old_pi_file):
     new_pi_credit_code = "0002"
-    new_pi_credit_amount = 1000
 
     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["Initial Credits"].empty or pandas.isna(
+        new_pi_credit_amount := invoice_pis["Initial Credits"].iat[0]
+    ):
+        new_pi_credit_amount = 1000
+
+    print(f"New PI Credit set at {new_pi_credit_amount} for {invoice_month}")
+
+    mask = (old_pi_df[PI_FIRST_MONTH] == add_invoice_month(invoice_month, -1)) & (
+        ~old_pi_df[PI_PI_FIELD].isin(current_pi_set)
+    )
+    old_pi_df.loc[mask, PI_2ND_REMAINDER] = old_pi_df.loc[mask, PI_1ST_REMAINDER]
 
-    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)
 
-        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(old_pi_df.loc[old_pi_df[PI_PI_FIELD] == pi]) == 0:
+                    pi_entry = [pi, invoice_month, new_pi_credit_amount, None, None]
+                    old_pi_df = pandas.concat(
+                        [
+                            pandas.DataFrame([pi_entry], columns=old_pi_df.columns),
+                            old_pi_df,
+                        ],
+                        ignore_index=True,
+                    )
+
+                remaining_credit = new_pi_credit_amount
+                remainder_field = PI_1ST_REMAINDER
+            elif pi_age == 1:
+                remaining_credit = old_pi_df.loc[
+                    old_pi_df[PI_PI_FIELD] == pi, PI_1ST_REMAINDER
+                ].iat[0]
+                remainder_field = PI_2ND_REMAINDER
+
             for i, row in pi_projects.iterrows():
-                project_cost = row[COST_FIELD]
-                applied_credit = min(project_cost, remaining_credit)
+                if remaining_credit == 0:
+                    dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD]
+                else:
+                    project_cost = row[COST_FIELD]
+                    applied_credit = min(project_cost, remaining_credit)
 
-                dataframe.at[i, CREDIT_FIELD] = applied_credit
-                dataframe.at[i, CREDIT_CODE_FIELD] = new_pi_credit_code
-                dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD] - applied_credit
-                remaining_credit -= applied_credit
+                    dataframe.at[i, CREDIT_FIELD] = applied_credit
+                    dataframe.at[i, CREDIT_CODE_FIELD] = new_pi_credit_code
+                    dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD] - applied_credit
+                    remaining_credit -= applied_credit
 
-                if remaining_credit == 0:
-                    break
+            old_pi_df.loc[
+                old_pi_df[PI_PI_FIELD] == pi, remainder_field
+            ] = remaining_credit
 
-    dump_old_pis(old_pi_file, old_pi_dict)
+    dump_old_pis(old_pi_file, old_pi_df)
 
     return dataframe