Repository: regro/regolith (language: python)

Issue:
ENH: attestations -> helper
implement simon's attestations.py as a helper


Discussion / hints:

Fix patch:
diff --git a/regolith/helper.py b/regolith/helper.py
index 2d0a6becd..617fb96e9 100644
--- a/regolith/helper.py
+++ b/regolith/helper.py
@@ -1,5 +1,6 @@
 """Generic builder."""
 
+from regolith.helpers import attestationshelper as attestations
 from regolith.helpers import a_manurevhelper as a_manurev
 from regolith.helpers import a_proprevhelper as a_proprev
 from regolith.helpers import a_grppub_readlisthelper as a_gprl
@@ -28,6 +29,7 @@
 from regolith.helpers import u_todohelper as u_todo
 
 HELPERS = {
+    "attestations": (attestations.AttestationsHelper, attestations.subparser),
     "a_expense": (a_expense.ExpenseAdderHelper, a_expense.subparser),
     "a_grppub_readlist": (a_gprl.GrpPubReadListAdderHelper, a_gprl.subparser),
     "a_manurev": (a_manurev.ManuRevAdderHelper, a_manurev.subparser),
diff --git a/regolith/schemas.py b/regolith/schemas.py
index f5edb3b4b..15b65fbca 100644
--- a/regolith/schemas.py
+++ b/regolith/schemas.py
@@ -231,6 +231,20 @@
             for i in range(1, 11)
         ],
         "payee": "scopatz",
+        "reimbursements": [
+            {
+            "amount": 500,
+            "date": "tbd",
+            "submission_date": "tbd",
+            "where": "Columbia"
+            },
+            {
+            "amount": 1000,
+            "date": "2019-09-15",
+            "submission_date": "2019-09-05",
+            "where": "Columbia"
+            }
+        ],
         "project": "Cyclus",
         "overall_purpose": "testing the databallectionsse",
         "notes": "this expense was used to get the work done",


Test patch:
diff --git a/regolith/helpers/attestationshelper.py b/regolith/helpers/attestationshelper.py
new file mode 100644
index 000000000..2fa7dca55
--- /dev/null
+++ b/regolith/helpers/attestationshelper.py
@@ -0,0 +1,222 @@
+import os
+
+from datetime import datetime, timedelta
+import dateutil.parser as date_parser
+import matplotlib.pyplot as plt
+import numpy as np
+
+from regolith.runcontrol import DEFAULT_RC, load_rcfile, filter_databases, \
+    connect_db
+from regolith.dates import month_to_int, get_dates
+from regolith.tools import (
+    fuzzy_retrieval,
+    all_docs_from_collection,
+    get_pi_id,
+    merge_collections_superior
+)
+from regolith.fsclient import _id_key
+from regolith.sorters import position_key
+from regolith.helpers.basehelper import DbHelperBase
+# print([k for k,v in chained_db['people'].items()])
+from pprint import pprint
+
+MONTH_COST = 3400
+
+def daterange(date1, date2):
+    return [date1 + timedelta(n) for n in range(int((date2 - date1).days) + 1)]
+
+def subparser(subpi):
+    subpi.add_argument("grant",
+                       help="grant id for the grant you want to find appointments")
+    subpi.add_argument("-b", "--begin_date",
+                       help="attestation period begins on this date, format YYYY-MM-DD")
+    subpi.add_argument("-e", "--end_date",
+                       help="attestation period ends on this date")
+    subpi.add_argument("--no-plot", action="store_true",
+                       help="suppress the plotting")
+    return subpi
+
+
+class AttestationsHelper(DbHelperBase):
+    """Helper for attestations"""
+    # btype must be the same as helper target in helper.py
+    btype = "attestation"
+    needed_colls = ['people', 'grants', 'proposals', 'expenses']
+
+    def construct_global_ctx(self):
+        """Constructs the global context"""
+        super().construct_global_ctx()
+        gtx = self.gtx
+        rc = self.rc
+        # if not rc.database:
+        #     rc.database = rc.databases[0]["name"]
+        gtx["people"] = sorted(
+            all_docs_from_collection(rc.client, "people"),
+            key=position_key,
+            reverse=True,
+        )
+        gtx["expenses"] = sorted(
+            all_docs_from_collection(rc.client, "expenses"), key=_id_key
+        )
+        gtx["grants"] = sorted(
+            all_docs_from_collection(rc.client, "grants"), key=_id_key
+        )
+        gtx["proposals"] = sorted(
+            all_docs_from_collection(rc.client, "proposals"), key=_id_key
+        )
+        gtx["all_docs_from_collection"] = all_docs_from_collection
+        gtx["float"] = float
+        gtx["str"] = str
+        gtx["zip"] = zip
+
+    def db_updater(self):
+        rc = self.rc
+        print(f"Instructions/Notes:\n"
+              f"  Quarters are: Q1 July thru Sept, Q2 Oct - Dec, Q3 Jan - Mar, Q4 Apr - Jun\n"
+              f"  Grad salaries are about ${MONTH_COST} per month"
+              )
+        grant_id = rc.grant
+
+        begin_date, end_date = None, None
+        print("Collecting Appointments for grant {}:".format(grant_id))
+        expenses = self.gtx['expenses']
+        people = self.gtx['people']
+        jg = self.gtx['grants']
+        proposals = self.gtx["proposals"]
+        grants = merge_collections_superior(proposals, jg, "proposal_id")
+        grant = fuzzy_retrieval(
+            grants,
+            ["name", "_id", "alias"],
+            grant_id)
+        if not grant:
+            raise ValueError(f"ERROR: grant {grant_id} not found in grants")
+        if rc.begin_date:
+            begin_date = date_parser.parse(rc.begin_date).date()
+        if rc.end_date:
+            end_date = date_parser.parse(rc.end_date).date()
+        if not begin_date:
+            begin_date = get_dates(grant)['begin_date']
+        if not end_date:
+            end_date = get_dates(grant)['end_date']
+        plot_date_list = daterange(begin_date, end_date)
+
+        months = (end_date - begin_date).days / 30.42
+
+        appts, begin, end = [], datetime(3070, 1, 1).date(), datetime(1970, 1, 1).date()
+        for person in people:
+            person_appts = person.get('appointments', None)
+            if person_appts:
+                for _id, p_appt in person_appts.items():
+                    grantid = p_appt.get('grant')
+                    if grantid == grant_id:
+                        loading = p_appt.get('loading')
+                        bd = get_dates(p_appt).get("begin_date")
+                        begin = min(begin, bd)
+                        ed = get_dates(p_appt).get("end_date")
+                        end = max(end, ed)
+                        months_on_grant = (ed - bd).days / 30.4 * loading
+                        appt = (person['_id'], bd, ed, loading, months_on_grant)
+                        appts.append(appt)
+        appts.sort(key=lambda x: (x[0], x[1]))
+        folks = []
+        for app in appts:
+            if app[1] < end_date:
+                if app[2] >= begin_date:
+                    print("{0}, from {1} to {2}, loading {3}. Total months: "
+                          "{4:6.2f}".format(app[0], app[1].strftime("%Y-%m-%d"),
+                                            app[2].strftime("%Y-%m-%d"), app[3], app[4]))
+            folks.append(app[0])
+        folks = list(set(folks))
+        plots = []
+        people_loadings = []
+        loadingc = np.zeros(len(plot_date_list))
+        for folk in folks:
+            fig, ax = plt.subplots()
+            loadinga = np.zeros(len(plot_date_list))
+            for app in appts:
+                if app[0] == folk:
+                    loadingl = []
+                    for day in plot_date_list:
+                        if app[1] <= day <= app[2]:
+                            loadingl.append(app[3])
+                        else:
+                            loadingl.append(0)
+                    loadinga = loadinga + np.array(loadingl)
+            loadingc = loadingc + loadinga
+
+            months, loadingm, accum, days = [plot_date_list[0]], [], 0, 0
+            for day, load in zip(plot_date_list, loadinga):
+                if day.day == 1 and days != 0:
+                    months.append(day)
+                    loadingm.append(accum * MONTH_COST / days)
+                    accum, days = 0, 0
+                accum = accum + load
+                days += 1
+            months.pop()
+
+            people_loadings.append((folk, loadinga, loadingm))
+            if not rc.no_plot:
+                ax.plot_date(plot_date_list, loadinga, ls='-', marker="", label=folk)
+                ax.set_xlabel('date')
+                ax.set_ylabel(f"loading for student {app[0]}")
+                ax.legend(loc='best')
+                fig.autofmt_xdate()
+                plots.append(fig)
+
+        if not rc.no_plot:
+            fig, ax = plt.subplots()
+            ax.plot_date(plot_date_list, loadingc, ls='-', marker="")
+
+        print(f"\n-----------\nLoadings by month\n------------")
+        index = 0
+        for month in months:
+            print(f"{month.isoformat()}:")
+            for person in people_loadings:
+                if person[2][index] > 0:
+                    print(f"    {person[0]}\tloading: {round(person[2][index], 2)}")
+            index += 1
+
+        print(f"\n----------------\nExpenses\n----------------")
+        expenses_on_grant = [expense for expense in expenses if
+                             grant_id in expense.get('grants')]
+
+        if len(expenses_on_grant) > 1:
+            expenses_on_grant.sort(key=lambda x: get_dates(x).get('end_date'))
+        for expense in expenses_on_grant:
+            # print(expense.get('overall_purpose'))
+            for reimb in expense.get('reimbursements'):
+                if reimb.get('amount') == 0:
+                    amt = 0
+                    for exp_item in expense.get('itemized_expenses', []):
+                        amt += exp_item.get('unsegregated_expense')
+                        amt += exp_item.get('prepaid_expense', 0)
+                    reimb['amount'] = amt
+
+        total_spend, month_spend, all_reimb_dates, all_reimb_amts = 0, 0, [], []
+        for e in expenses_on_grant:
+            reimb_amts = [round(i.get('amount'), 2) for i in e.get('reimbursements', [{}])]
+            reimb_dates = [get_dates(i).get('date', get_dates(e).get('end_date')) for i in
+                           e.get('reimbursements', [{}])]
+            all_reimb_dates.extend(reimb_dates)
+            all_reimb_amts.extend(reimb_amts)
+            total_spend += sum(reimb_amts)
+
+            for reim_date, amt in zip(reimb_dates, reimb_amts):
+                print(
+                    f"{reim_date} (reimb date), {get_dates(e).get('end_date')} (expense date): amount: "
+                    f"{amt}, ")
+            print(
+                f"  payee: {e.get('payee')} "
+                f"purpose: {e.get('overall_purpose')[:60]}")
+        for month in months:
+            if month >= begin_date:
+                month_spend = 0
+                for amt, dte in zip(all_reimb_amts, all_reimb_dates):
+                    if month.year == dte.year and month.month == dte.month:
+                        month_spend += amt
+                print(f"{month}: expenses monthly total = {month_spend}")
+
+        print(f"Total spend = {round(total_spend, 2)}")
+        for plot in plots:
+            plt.show()
+
diff --git a/tests/outputs/a_expense/expenses.yaml b/tests/outputs/a_expense/expenses.yaml
index 2f91f0834..b2247916d 100644
--- a/tests/outputs/a_expense/expenses.yaml
+++ b/tests/outputs/a_expense/expenses.yaml
@@ -135,4 +135,13 @@ test:
   overall_purpose: testing the databallectionsse
   payee: scopatz
   project: Cyclus
+  reimbursements:
+    - amount: 500
+      date: tbd
+      submission_date: tbd
+      where: Columbia
+    - amount: 1000
+      date: '2019-09-15'
+      submission_date: '2019-09-05'
+      where: Columbia
   status: submitted
diff --git a/tests/outputs/a_presentation/expenses.yaml b/tests/outputs/a_presentation/expenses.yaml
index 67b120b11..13aeab6bf 100644
--- a/tests/outputs/a_presentation/expenses.yaml
+++ b/tests/outputs/a_presentation/expenses.yaml
@@ -186,4 +186,13 @@ test:
   overall_purpose: testing the databallectionsse
   payee: scopatz
   project: Cyclus
+  reimbursements:
+    - amount: 500
+      date: tbd
+      submission_date: tbd
+      where: Columbia
+    - amount: 1000
+      date: '2019-09-15'
+      submission_date: '2019-09-05'
+      where: Columbia
   status: submitted
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
index ebfb1cdfb..f5567c4ce 100644
--- a/tests/test_helpers.py
+++ b/tests/test_helpers.py
@@ -7,6 +7,34 @@
 
 dash = "-"
 helper_map = [
+    (["helper", "attestations", "dmref15", "-b", "2019-09-01", "-e", "2019-11-01", "--no-plot"],
+     "Instructions/Notes:\n"
+     "  Quarters are: Q1 July thru Sept, Q2 Oct - Dec, Q3 Jan - Mar, Q4 Apr - Jun\n"
+     "  Grad salaries are about $3400 per month\n"
+     "Collecting Appointments for grant dmref15:\n"
+     "scopatz, from 2019-09-01 to 2019-10-31, loading 0.75. Total months:   1.48\n"
+     "\n-----------\nLoadings by month\n------------\n"
+     "2019-09-01:\n"
+     "    scopatz\tloading: 2550.0\n"
+     "2019-10-01:\n"
+     "    scopatz\tloading: 2550.0\n"
+     "\n----------------\nExpenses\n----------------\n"
+     "2018-01-10 (reimb date), 2018-01-10 (expense date): amount: 500, \n"
+     "2019-09-15 (reimb date), 2018-01-10 (expense date): amount: 1000, \n"
+     "  payee: scopatz purpose: testing the databallectionsse\n"
+     "2019-09-01: expenses monthly total = 1000\n"
+     "2019-10-01: expenses monthly total = 0\n"
+     "Total spend = 1500\n"),
+    (["helper", "attestations", "sym", "--no-plot"],
+     "Instructions/Notes:\n"
+     "  Quarters are: Q1 July thru Sept, Q2 Oct - Dec, Q3 Jan - Mar, Q4 Apr - Jun\n"
+     "  Grad salaries are about $3400 per month\n"
+     "Collecting Appointments for grant sym:\n"
+     "\n-----------\nLoadings by month\n------------\n"
+     "2030-05-01:\n2030-06-01:\n2030-07-01:\n2030-08-01:\n2030-09-01:\n2030-10-01:\n2030-11-01:\n"
+     "\n----------------\nExpenses\n----------------\n"
+     "2030-05-01: expenses monthly total = 0\n2030-06-01: expenses monthly total = 0\n2030-07-01: expenses monthly total = 0\n2030-08-01: expenses monthly total = 0\n2030-09-01: expenses monthly total = 0\n2030-10-01: expenses monthly total = 0\n2030-11-01: expenses monthly total = 0\n"
+     "Total spend = 0\n"),
     (["helper", "a_proprev", "A. Einstein", "nsf", "2020-04-08", "-q",
       "Tess Guebre", "--status", "downloaded", "--title", "A flat world theory"],
      "A. Einstein proposal has been added/updated in proposal reviews\n"),