import matplotlib.pyplot as plt
import numpy as np
import math

def simulation( doses_per_day, dosage_per_day, blood_volume):
    # drug_dosage_start.py
    #
    # VARIABLE DICTIONARY DRUG DOSAGE MODEL:
    # --------------------------------------------------------------------
    #     Drug levels in the blood, for certain drugs, can be modeled by a pair of
    #     coupled first-order linear differential equations.  For compactness, let
    #     A= medicine_in_intestines and B = plasma_level.  Variable names are defined
    #     below.
    #     dA/dt = -absorption_rate*A + intake
    #     dB/dt = -excretion_rate*B + absorption_rate*A

    # Model parameters
    time_step = 0.25       # simulation time step in hours
    start_time = 0         # in hours
    end_time = 48          # in hours
    # doses_per_day = 2
    # dosage_per_day = 12000  #mg of the drug
    absorption_rate = 0.25 # absorption rate constant
    half_life = 6          # in hours
    # blood_volume = 4.6     # in liters
    medicinal_level = 800  # level at which drug becomes effective mg/l
    toxic_level = 1000     # level at which drug becomes toxic  mg/l

    # Derived constants
    N = int((end_time - start_time) / time_step)   # number of simulation steps
    dosage = dosage_per_day / doses_per_day        # amount of each dose
    dosage_interval = 24 / doses_per_day           # in hours
    excretion_rate = math.log(2) / half_life
    steps_between_doses = dosage_interval / time_step

    # Time-varying quantities, arrays with one value per time step
    # The syntax, "[0]*(N+1)," creates a one-dimensional array of length N+1 whose values are all zero.

    t = [0]*(N+1)                       # time in hours
    intake = [0]*(N+1)                  # dose at this time step
    medicine_in_intestines = [0]*(N+1)  # level of drug in intestines
    plasma_level = [0]*(N+1)            # level of drug in the blood
    plasma_concentration = [0]*(N+1)    # concentration of drug in blood
    absorption = [0]*(N+1)              # amount absorbed at this time step
    excretion = [0]*(N+1)               # amount excreted at this time step

    # Initialize variables - assume all others start at 0
    t[0] = start_time
    step_next_dose = 1         # give first dose at start time

    # Loop to calculate intake of drug over time 
    # Euler's method
    # this part calculates when each dose of the drug is added based on
    # the number of doses over a 24 hour period. Number of doses must be an
    # integer divisor of 24 for this to work.  Otherwise change the time step
    for i in range(N):
        t[i+1] = t[i] + time_step
        # Is it time for a dose?
        if i == step_next_dose:
            intake[i] = dosage
            step_next_dose = round(i + steps_between_doses, 0)
        else:
            intake[i] = 0
            
    #Create another loop that calculates the absorption, excretion,
    #plasma level, and plasma concentration over time
     
    for i in range(N):

        absorption[i] = time_step * absorption_rate * medicine_in_intestines[i]

        excretion[i] = time_step *  excretion_rate * plasma_level[i]

        plasma_level[i+1] = plasma_level[i] - excretion[i] + absorption[i]

        plasma_concentration[i+1] = plasma_level[i+1] / blood_volume

        medicine_in_intestines[i+1] = medicine_in_intestines[i] - absorption[i] + intake[i]
        

    # After the simulation has been run check to see that the last doseage
    # interval is in the safety region
    if all(test > medicinal_level  and test < toxic_level 
            for test in plasma_concentration[-1*int(steps_between_doses):-1]):
        return True
    else:
        return False



def count_mg():

    # Initialize list of safty conditions
    count = []
    count_low = []
    count_high = []
    mg = []
    mg_low = []
    mg_high = []


    # Loop through the range of doseage per day
    for i in range(1, 13):
        # For each count loop through the doses levels
        for j in np.linspace(1000, 20000, 100):
            # Check for the average, low, and high ranges
            check = simulation(i, j, 4.6)
            check_low = simulation(i, j, 4.6 - 0.2*4.6)
            check_high = simulation(i, j, 4.6*.2 + 4.6)

            # Add the safe conditions to the approriate list
            if check:
                count.append(i)
                mg.append(j)

            if check_low:
                count_low.append(i)
                mg_low.append(j)

            if check_high:
                count_high.append(i)
                mg_high.append(j)

    plt.plot(count, mg, 'o', label='Avg')
    plt.plot(count_low, mg_low, 'o', label='Low')
    plt.plot(count_high, mg_high, 'o', label='High')
    plt.xlim(0, 13)
    plt.xlabel('Count')
    plt.ylabel('mg')
    plt.legend(loc=0)
    plt.show()

def mg_vol(AvgBloodVol):
    # Initialize list of safty conditions
    count = []
    mg = []

    # Loop through the range of blood volumes
    for i in np.linspace(AvgBloodVol -.2*AvgBloodVol , AvgBloodVol  +.2*AvgBloodVol ):
        # For each blood volume loop through the doses
        for j in np.linspace(1000, 20000, 100):
            check = simulation(4, j, i)
            # If a safe condition is found then add it to the list
            if check:
                count.append(i)
                mg.append(j)

    plt.plot(count, mg, 'o')
    plt.axvline(AvgBloodVol, linestyle = 'dashed')
    plt.xlabel('Volume')
    plt.ylabel('mg')
    plt.show()

count_mg()
mg_vol(4.6)