Merge pull request #71 from QuanMPhm/52.4/refactor_bu_internal

Refactored BU Internal Invoice

process_report/invoices/billable_invoice.py

@@ -1,21 +1,20 @@
 from dataclasses import dataclass
-from decimal import Decimal
 import logging
 import sys
 
 import pandas
 import pyarrow
 
-import process_report.invoices.invoice as invoice
-import process_report.util as util
+from process_report.invoices import invoice, discount_invoice
+from process_report import util
 
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 
 @dataclass
-class BillableInvoice(invoice.Invoice):
+class BillableInvoice(discount_invoice.DiscountInvoice):
     NEW_PI_CREDIT_CODE = "0002"
     INITIAL_CREDIT_AMOUNT = 1000
     EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"]
@@ -89,7 +88,7 @@ def _prepare(self):
         self.data = self._validate_pi_names(self.data)
         self.data[invoice.CREDIT_FIELD] = None
         self.data[invoice.CREDIT_CODE_FIELD] = None
-        self.data[invoice.BALANCE_FIELD] = Decimal(0)
+        self.data[invoice.BALANCE_FIELD] = self.data[invoice.COST_FIELD]
         self.old_pi_df = self._load_old_pis(self.old_pi_filepath)
 
     def _process(self):
@@ -143,14 +142,17 @@ def get_initial_credit_amount(
 
         current_pi_set = set(data[invoice.PI_FIELD])
         for pi in current_pi_set:
-            pi_projects = data[data[invoice.PI_FIELD] == pi]
+            credit_eligible_projects = data[
+                (data[invoice.PI_FIELD] == pi)
+                & ~(data[invoice.SU_TYPE_FIELD].isin(self.EXCLUDE_SU_TYPES))
+            ]
             pi_age = self._get_pi_age(old_pi_df, pi, self.invoice_month)
             pi_old_pi_entry = old_pi_df.loc[
                 old_pi_df[invoice.PI_PI_FIELD] == pi
             ].squeeze()
 
             if pi_age > 1:
-                for i, row in pi_projects.iterrows():
+                for i, row in credit_eligible_projects.iterrows():
                     data.at[i, invoice.BALANCE_FIELD] = row[invoice.COST_FIELD]
             else:
                 if pi_age == 0:
@@ -176,25 +178,16 @@ def get_initial_credit_amount(
                     )
                     credit_used_field = invoice.PI_2ND_USED
 
-                initial_credit = remaining_credit
-                for i, row in pi_projects.iterrows():
-                    if (
-                        remaining_credit == 0
-                        or row[invoice.SU_TYPE_FIELD] in self.EXCLUDE_SU_TYPES
-                    ):
-                        data.at[i, invoice.BALANCE_FIELD] = row[invoice.COST_FIELD]
-                    else:
-                        project_cost = row[invoice.COST_FIELD]
-                        applied_credit = min(project_cost, remaining_credit)
-
-                        data.at[i, invoice.CREDIT_FIELD] = applied_credit
-                        data.at[i, invoice.CREDIT_CODE_FIELD] = self.NEW_PI_CREDIT_CODE
-                        data.at[i, invoice.BALANCE_FIELD] = (
-                            row[invoice.COST_FIELD] - applied_credit
-                        )
-                        remaining_credit -= applied_credit
+                credits_used = self.apply_flat_discount(
+                    data,
+                    credit_eligible_projects,
+                    remaining_credit,
+                    invoice.CREDIT_FIELD,
+                    invoice.BALANCE_FIELD,
+                    invoice.CREDIT_CODE_FIELD,
+                    self.NEW_PI_CREDIT_CODE,
+                )
 
-                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]
                 ):

process_report/invoices/bu_internal_invoice.py

@@ -0,0 +1,70 @@
+from dataclasses import dataclass
+from decimal import Decimal
+
+import process_report.invoices.invoice as invoice
+import process_report.invoices.discount_invoice as discount_invoice
+
+
+@dataclass
+class BUInternalInvoice(discount_invoice.DiscountInvoice):
+    subsidy_amount: int
+
+    def _prepare(self):
+        def get_project(row):
+            project_alloc = row[invoice.PROJECT_FIELD]
+            if project_alloc.rfind("-") == -1:
+                return project_alloc
+            else:
+                return project_alloc[: project_alloc.rfind("-")]
+
+        self.data = self.data[
+            self.data[invoice.INSTITUTION_FIELD] == "Boston University"
+        ].copy()
+        self.data["Project"] = self.data.apply(get_project, axis=1)
+        self.data[invoice.SUBSIDY_FIELD] = Decimal(0)
+        self.data = self.data[
+            [
+                invoice.INVOICE_DATE_FIELD,
+                invoice.PI_FIELD,
+                "Project",
+                invoice.COST_FIELD,
+                invoice.CREDIT_FIELD,
+                invoice.SUBSIDY_FIELD,
+                invoice.BALANCE_FIELD,
+            ]
+        ]
+
+    def _process(self):
+        data_summed_projects = self._sum_project_allocations(self.data)
+        self.data = self._apply_subsidy(data_summed_projects, self.subsidy_amount)
+
+    def _sum_project_allocations(self, dataframe):
+        """A project may have multiple allocations, and therefore multiple rows
+        in the raw invoices. For BU-Internal invoice, we only want 1 row for
+        each unique project, summing up its allocations' costs"""
+        project_list = dataframe["Project"].unique()
+        data_no_dup = dataframe.drop_duplicates("Project", inplace=False)
+        sum_fields = [invoice.COST_FIELD, invoice.CREDIT_FIELD, invoice.BALANCE_FIELD]
+        for project in project_list:
+            project_mask = dataframe["Project"] == project
+            no_dup_project_mask = data_no_dup["Project"] == project
+
+            sum_fields_sums = dataframe[project_mask][sum_fields].sum().values
+            data_no_dup.loc[no_dup_project_mask, sum_fields] = sum_fields_sums
+
+        return data_no_dup
+
+    def _apply_subsidy(self, dataframe, subsidy_amount):
+        pi_list = dataframe[invoice.PI_FIELD].unique()
+
+        for pi in pi_list:
+            pi_projects = dataframe[dataframe[invoice.PI_FIELD] == pi]
+            self.apply_flat_discount(
+                dataframe,
+                pi_projects,
+                subsidy_amount,
+                invoice.SUBSIDY_FIELD,
+                invoice.BALANCE_FIELD,
+            )
+
+        return dataframe

process_report/invoices/discount_invoice.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass
+
+import pandas
+
+import process_report.invoices.invoice as invoice
+
+
+@dataclass
+class DiscountInvoice(invoice.Invoice):
+    """
+    Invoice class containing functions useful for applying discounts
+    on dataframes
+    """
+
+    @staticmethod
+    def apply_flat_discount(
+        invoice: pandas.DataFrame,
+        pi_projects: pandas.DataFrame,
+        discount_amount: int,
+        discount_field: str,
+        balance_field: str,
+        code_field: str = None,
+        discount_code: str = None,
+    ):
+        """
+        Takes in an invoice and a list of PI projects that are a subset of it,
+        and applies a flat discount to those PI projects. Note that this function
+        will change the provided `invoice` Dataframe directly. Therefore, it does
+        not return the changed invoice.
+
+        This function assumes that the balance field shows the remaining cost of the project,
+        or what the PI would pay before the flat discount is applied.
+
+        If the optional parameters `code_field` and `discount_code` are passed in,
+        `discount_code` will be comma-APPENDED to the `code_field` of projects where
+        the discount is applied
+
+        Returns the amount of discount used.
+
+        :param invoice: Dataframe containing all projects
+        :param pi_projects: A subset of `invoice`, containing all projects for a PI you want to apply the discount
+        :param discount_amount: The discount given to the PI
+        :param discount_field: Name of the field to put the discount amount applied to each project
+        :param balance_field: Name of the balance field
+        :param code_field: Name of the discount code field
+        :param discount_code: Code of the discount
+        """
+
+        def apply_discount_on_project(remaining_discount_amount, project_i, project):
+            remaining_project_balance = project[balance_field]
+            applied_discount = min(remaining_project_balance, remaining_discount_amount)
+            invoice.at[project_i, discount_field] = applied_discount
+            invoice.at[project_i, balance_field] = (
+                project[balance_field] - applied_discount
+            )
+            remaining_discount_amount -= applied_discount
+            return remaining_discount_amount
+
+        def apply_credit_code_on_project(project_i):
+            if code_field and discount_code:
+                if pandas.isna(invoice.at[project_i, code_field]):
+                    invoice.at[project_i, code_field] = discount_code
+                else:
+                    invoice.at[project_i, code_field] = (
+                        invoice.at[project_i, code_field] + "," + discount_code
+                    )
+
+        remaining_discount_amount = discount_amount
+        for i, row in pi_projects.iterrows():
+            if remaining_discount_amount == 0:
+                break
+            else:
+                remaining_discount_amount = apply_discount_on_project(
+                    remaining_discount_amount, i, row
+                )
+                apply_credit_code_on_project(i)
+
+        discount_used = discount_amount - remaining_discount_amount
+        return discount_used