#-------------------------------------------------------------------------#
#-------------------------------------------------------------------------#
#Program: MME Calculator
#Purpose: Calculate MME using 4 formulas based on inputted Rx information
#Author: Bethany DiPrete
#        diprete@email.unc.edu
#Date modified: April 12, 2021
#Conversions from: 
#   https://www.cms.gov/Medicare/Prescription-Drug-Coverage/PrescriptionDrugCovContra/Downloads/Oral-MME-CFs-vFeb-2018.pdf
#-------------------------------------------------------------------------#
#-------------------------------------------------------------------------#


#-------------------------------------------------------------------------#
#Initialize
#-------------------------------------------------------------------------#

import numpy as np
import pandas as pd
import math as math

#-------------------------------------------------------------------------#
#Example Inputs
#-------------------------------------------------------------------------#

#window length for definition 3. Currently set to 90 days
window_length = 90

#Input Rx dates, days' supply, quantity, form, strength, ingredient
rx_start = pd.to_datetime(["2020-08-01", "2020-08-03","2020-09-27"])
days_supply = (7,5,7)
quantity = (7,5,7)
form = ("Tablet", "Tablet", "Tablet")
strength = (20, 20, 100)
ingredient = ("Oxycodone", "Hydrocodone", "Tramadol")


#-------------------------------------------------------------------------#
#Calculate MME using 4 definitions
#-------------------------------------------------------------------------#

#Function to calculate MME using each definition
def mme_calc(rx_start, 
             days_supply, 
             ingredient, 
             strength, 
             quantity, 
             form, 
             window_length = 90):

    pat_rx = pd.DataFrame({"rx_start": rx_start,
                       "days_supply": days_supply,
                       "ingredient": ingredient,
                       "strength": strength,
                       "quantity": quantity,
                       "form": form})
    
    mme_conditions = [
        (pat_rx.ingredient=="Butorphanol"),
        (pat_rx.ingredient=="Codeine"),
        (pat_rx.ingredient=="Dihydrocodeine"),
        ((pat_rx.ingredient=="Fentanyl") & 
             pat_rx.form.isin(["Buccal","Tablet","Lozenge"])),
        ((pat_rx.ingredient=="Fentanyl") & 
             (pat_rx.form.isin(["Film","Oral Spray"]))),
        ((pat_rx.ingredient=="Fentanyl") & (pat_rx.form=="Nasal Spray")),
        ((pat_rx.ingredient=="Fentanyl") & (pat_rx.form=="Patch")),
        (pat_rx.ingredient=="Hydrocodone"),
        (pat_rx.ingredient=="Hydromorphone"),
        (pat_rx.ingredient=="Levorphanol tartrate"),
        (pat_rx.ingredient=="Meperidine hydrochloride"),
        ((pat_rx.ingredient=="Methadone") & (pat_rx.strength.le(20))),
        ((pat_rx.ingredient=="Methadone") & (pat_rx.strength>=21) & (pat_rx.strength<=40)),
        ((pat_rx.ingredient=="Methadone") & (pat_rx.strength>=41) & (pat_rx.strength<=60)),
        ((pat_rx.ingredient=="Methadone") & (pat_rx.strength>=61) & (pat_rx.strength<=80)),
        (pat_rx.ingredient=="Morphine"),
        (pat_rx.ingredient=="Opium"),
        (pat_rx.ingredient=="Oxycodone"),
        (pat_rx.ingredient=="Oxymorphone"),
        (pat_rx.ingredient=="Pentazocine"),
        (pat_rx.ingredient=="Tapentadol"),
        (pat_rx.ingredient=="Tramadol")
        ]
    mme_conversions = [7, 0.15, 0.25, 0.13, 0.18, 0.16, 2.4, 1, 4, 11, 0.1, 4, 
                       8, 10, 12, 1, 1, 1.5, 3, 0.37, 0.4, 0.1]
    
    #Add conversion factors
    pat_rx['conversion']= np.select(mme_conditions, mme_conversions)
    
    #Calculate MME/day for each Rx
    pat_rx['MME']=(pat_rx['quantity']/pat_rx['days_supply'])*pat_rx['strength']*pat_rx['conversion']

    #Convert to days with a prescription, rather than just Rx start & end dates
    pat_rx['date'] = [pd.date_range(s, e, freq='d') for s, e in
              zip(pat_rx['rx_start'], 
                  pat_rx['rx_start']+pd.to_timedelta((pat_rx['days_supply']-1), unit="d"))]
    pat_rx = pat_rx.explode('date').drop(['rx_start'], axis=1)
    pat_rx = pat_rx.sort_values(by='date')

    #Daily MME and number of prescriptions on each day
    pat_rx = pat_rx.groupby('date').agg(mmed=('MME', 'sum'), rx_num=('date', 'count'))

    #Calculate MME using each definition
    def_1 = math.ceil(sum(pat_rx['mmed'])/sum(pat_rx['rx_num']))
    def_2 = math.ceil(sum(pat_rx['mmed'])/len(pat_rx))
    def_3 = math.ceil(sum(pat_rx['mmed'])/window_length)
    def_4 = math.ceil(max(pat_rx['mmed']))
    
    #Output object with the calculation for each definition
    return pd.Series([def_1, def_2, def_3, def_4], 
                     index = ["MME_Def1", "MME_Def2", "MME_Def3", "MME_Def4"],
                     dtype = 'object')



mme = mme_calc(rx_start = rx_start, 
               days_supply = days_supply, 
               ingredient = ingredient, 
               strength = strength, 
               quantity = quantity, 
               form = form)