import gradio as gr
import numpy as np
import pandas as pd


def load_health_data():
    # Load the CSV file into a pandas DataFrame
    df = pd.read_csv('healthdata.csv', header=2)
    # Replace null values with empty strings
    df = df.replace({np.nan: "", None: ""})
    return df

def get_device_data(device):
    df = load_health_data()
    
    # Common columns for all devices
    base_cols = ['Day', 'Date', 'Day of Week']
    
    if device == "AppleWatch + AutoSleep":
        # Get columns with AS suffix + base columns
        device_cols = base_cols + [col for col in df.columns if 'AS' in col]
    elif device == "8sleep Pod 4 Ultra":
        # Get columns with 8S suffix + base columns
        device_cols = base_cols + [col for col in df.columns if '8S' in col]
    elif device == "Oura":
        # Get columns with OU suffix + base columns
        device_cols = base_cols + [col for col in df.columns if 'OU' in col]
    elif device == "Whoop":
        # Get columns with WH suffix + base columns
        device_cols = base_cols + [col for col in df.columns if 'WH' in col]
    else:
        return df
    
    # Add the Average and Notes columns if they exist
    if 'Average' in df.columns:
        device_cols.append('Average')
    if 'Notes' in df.columns:
        device_cols.append('Notes')
    
    return df[device_cols]

def get_color_for_metric(val, metric_name):
    """Return a color based on the metric type and value"""
    try:
        val = float(val)
    except (ValueError, TypeError):
        return ""
        
    # Define thresholds for different metrics
    if "BPM" in metric_name:
        # Lower is better for BPM
        if val <= 50:
            return "background-color: rgba(0, 180, 0, 0.3)"  # Green
        elif val >= 54:
            return "background-color: rgba(180, 0, 0, 0.3)"  # Red
        else:
            return "background-color: rgba(180, 180, 180, 0.1)"  # Light gray for middle range
            
    elif "Score" in metric_name:
        # Higher is better for Score
        if val >= 85:
            return "background-color: rgba(0, 180, 0, 0.3)"  # Dark green
        elif val >= 75:
            return "background-color: rgba(0, 140, 0, 0.2)"  # Medium green
        elif val <= 70:
            return "background-color: rgba(180, 0, 0, 0.3)"  # Red
        else:
            return "background-color: rgba(180, 180, 180, 0.1)"  # Light gray
            
    elif "HRV" in metric_name:
        # Higher is better for HRV
        if val >= 80:
            return "background-color: rgba(0, 180, 0, 0.3)"  # Dark green
        elif val <= 55:
            return "background-color: rgba(180, 0, 0, 0.3)"  # Red
        else:
            return "background-color: rgba(180, 180, 180, 0.1)"  # Light gray
            
    elif "Deep" in metric_name:
        # Higher might be better for Deep sleep
        if val >= 1.0:
            return "background-color: rgba(0, 180, 0, 0.3)"  # Green
        elif val <= 0.2:
            return "background-color: rgba(180, 0, 0, 0.3)"  # Red
        else:
            return "background-color: rgba(180, 180, 180, 0.1)"  # Light gray
            
    elif "REM" in metric_name:
        # Higher might be better for REM
        if val >= 1.5:
            return "background-color: rgba(0, 180, 0, 0.3)"  # Green
        elif val <= 0.8:
            return "background-color: rgba(180, 0, 0, 0.3)"  # Red
        else:
            return "background-color: rgba(180, 180, 180, 0.1)"  # Light gray
            
    elif "Light" in metric_name:
        # Medium values might be ideal for Light sleep
        if 5.0 <= val <= 6.0:
            return "background-color: rgba(0, 180, 0, 0.3)"  # Green
        elif val >= 6.3 or val <= 4.5:
            return "background-color: rgba(180, 0, 0, 0.3)"  # Red
        else:
            return "background-color: rgba(180, 180, 180, 0.1)"  # Light gray
            
    return ""  # Default no color

def prepare_data_with_styling(df):
    data = df.values.tolist()
    headers = list(df.columns)
    
    # Create styling
    styling = []
    for row in data:
        row_styling = []
        for i, val in enumerate(row):
            if headers[i] in ['Day', 'Date', 'Day of Week', 'Notes'] or val == "":
                row_styling.append(["", ""])
            else:
                style = get_color_for_metric(val, headers[i])
                row_styling.append(["", f"width: var(--cell-width-3); left: auto; {style}"])
        styling.append(row_styling)
    
    return {
        "data": data,
        "headers": headers,
        "metadata": {
            "styling": styling
        }
    }

column_widths = [70,150,150,150,150]
# Create the Gradio application using Blocks
with gr.Blocks(title="@Andrej Karpathy's Sleep Data") as demo:
    gr.Markdown("# @karpathy's Sleep Data")
    gr.Markdown("This app displays Andrej Karpathy's sleep tracking data from multiple devices. See his [blog post](https://karpathy.bearblog.dev/finding-the-best-sleep-tracker/) to learn more about the data and the process of finding the best sleep tracker.")
    
    with gr.Tabs():
        with gr.TabItem("AppleWatch + AutoSleep"):
            df = get_device_data("AppleWatch + AutoSleep")
            gr.Dataframe(prepare_data_with_styling(df), show_search="search", column_widths=column_widths, show_copy_button=True, pinned_columns=1)
        
        with gr.TabItem("8sleep Pod 4 Ultra"):
            df = get_device_data("8sleep Pod 4 Ultra")
            gr.Dataframe(prepare_data_with_styling(df), show_search="search", column_widths=column_widths, show_copy_button=True, pinned_columns=1)
        
        with gr.TabItem("Oura"):
            df = get_device_data("Oura")    
            gr.Dataframe(prepare_data_with_styling(df), show_search="search", column_widths=column_widths, show_copy_button=True, pinned_columns=1)
        
        with gr.TabItem("Whoop"):
            df = get_device_data("Whoop")
            gr.Dataframe(prepare_data_with_styling(df), show_search="search", column_widths=column_widths, show_copy_button=True, pinned_columns=1)
        
        with gr.TabItem("All Data"):
            df = load_health_data()
            gr.Dataframe(prepare_data_with_styling(df), show_search="search", column_widths=column_widths, show_copy_button=True, pinned_columns=1)
if __name__ == "__main__":
    demo.launch()