from flask import Flask, render_template, request, redirect, url_for
import sqlite3
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import io
import base64
from datetime import datetime, timedelta
import os

app = Flask(__name__)

# Pharmacokinetics parameters (in minutes)
# ln(2) / (Half-life-in-minutes)
CAFFEINE_ABSORPTION_RATE = 0.05  # per minute (half-life ~15 min)
CAFFEINE_ELIMINATION_RATE = 0.0023  # per minute (half-life ~5 hours)
AMPHETAMINE_ABSORPTION_RATE = 0.03  # per minute (half-life ~20 min)
AMPHETAMINE_ELIMINATION_RATE = 0.0013  # per minute (half-life ~9 hours)

# Initialize database
def init_db():
    conn = sqlite3.connect('stimulants.db')
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS doses
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                 timestamp DATETIME,
                 type TEXT,
                 amount INTEGER)''')
    conn.commit()
    conn.close()

# Calculate stimulant amount at given time
def calculate_amount(doses, time, k_abs, k_elim):
    total = 0
    for dose in doses:
        dose_time = datetime.strptime(dose[0], '%Y-%m-%d %H:%M:%S')
        t = (time - dose_time).total_seconds() / 60  # convert to minutes
        if t >= 0:
            # One-compartment model with first-order absorption
            total += dose[1] * (k_abs / (k_abs - k_elim)) * (
                np.exp(-k_elim * t) - np.exp(-k_abs * t)
            )
    return total

# Generate graph for past 24 hours
def generate_graph():
    # Get current time and 24 hours ago
    now = datetime.now()
    start_time = now - timedelta(hours=24)
    
    # Create time points (every 10 minutes)
    times = [start_time + timedelta(minutes=i*10) for i in range(144)]
    
    # Get all doses from database
    conn = sqlite3.connect('stimulants.db')
    c = conn.cursor()
    c.execute("SELECT timestamp, amount FROM doses WHERE type='caffeine'")
    caffeine_doses = c.fetchall()
    c.execute("SELECT timestamp, amount FROM doses WHERE type='amphetamine'")
    amphetamine_doses = c.fetchall()
    conn.close()
    
    # Calculate amounts for each time point
    caffeine_levels = [calculate_amount(caffeine_doses, t, 
                                        CAFFEINE_ABSORPTION_RATE, 
                                        CAFFEINE_ELIMINATION_RATE) for t in times]
    amphetamine_levels = [calculate_amount(amphetamine_doses, t,
                                           AMPHETAMINE_ABSORPTION_RATE,
                                           AMPHETAMINE_ELIMINATION_RATE) for t in times]
    
    # Create plot
    plt.figure(figsize=(12, 6))
    plt.plot(times, caffeine_levels, 'b-', label='Caffeine (mg)')
    plt.plot(times, amphetamine_levels, 'r-', label='Amphetamine (mg)')
    plt.xlabel('Time')
    plt.ylabel('Estimated Amount (mg)')
    plt.title('Stimulant Levels in Body (Past 24 Hours)')
    plt.legend()
    plt.grid(True)
    
    # Format x-axis
    ax = plt.gca()
    ax.xaxis.set_major_locator(matplotlib.dates.HourLocator(interval=4))
    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%H:%M'))
    
    # Convert to base64 string
    img = io.BytesIO()
    plt.savefig(img, format='png')
    img.seek(0)
    graph_url = base64.b64encode(img.getvalue()).decode()
    plt.close()
    return graph_url

@app.route('/')
def index():
    # Calculate current amounts
    now = datetime.now()
    
    conn = sqlite3.connect('stimulants.db')
    c = conn.cursor()
    c.execute("SELECT timestamp, amount FROM doses WHERE type='caffeine'")
    caffeine_doses = c.fetchall()
    c.execute("SELECT timestamp, amount FROM doses WHERE type='amphetamine'")
    amphetamine_doses = c.fetchall()
    conn.close()
    
    caffeine_amount = calculate_amount(caffeine_doses, now,
                                      CAFFEINE_ABSORPTION_RATE,
                                      CAFFEINE_ELIMINATION_RATE)
    amphetamine_amount = calculate_amount(amphetamine_doses, now,
                                         AMPHETAMINE_ABSORPTION_RATE,
                                         AMPHETAMINE_ELIMINATION_RATE)
    
    # Generate graph
    graph_url = generate_graph()
    
    return render_template('index.html', 
                          caffeine_amount=round(caffeine_amount, 2),
                          amphetamine_amount=round(amphetamine_amount, 2),
                          graph_url=graph_url)

@app.route('/add_dose', methods=['POST'])
def add_dose():
    stim_type = request.form['type']
    amount = 10  # Fixed 10mg amount
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    
    conn = sqlite3.connect('stimulants.db')
    c = conn.cursor()
    c.execute("INSERT INTO doses (timestamp, type, amount) VALUES (?, ?, ?)",
              (timestamp, stim_type, amount))
    conn.commit()
    conn.close()
    
    return redirect(url_for('index'))

if __name__ == '__main__':
    init_db()
    app.run(debug=True)