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_value(val, min_val, max_val):
    """Return a color based on where the value falls in the range"""
    if min_val == max_val:
        return ""  # No color if all values are the same
        
    normalized = (val - min_val) / (max_val - min_val)
    
    if normalized < 0.33:
        return ""  # Low values are blank
    elif normalized < 0.66:
        # Medium values are green
        intensity = int(255 * ((normalized - 0.33) / 0.33))
        return f"background-color: rgba(0, {intensity}, 0, 0.3)"
    else:
        # High values are red
        intensity = int(255 * ((normalized - 0.66) / 0.34))
        return f"background-color: rgba({intensity}, 0, 0, 0.3)"

def prepare_data_with_styling(df):
    data = df.values.tolist()
    headers = list(df.columns)
    
    # Get min/max for each numeric column
    col_ranges = {}
    for i, col in enumerate(headers):
        if col not in ['Day', 'Date', 'Day of Week', 'Notes']:
            try:
                values = [float(row[i]) for row in data if row[i] != "" and pd.notna(row[i])]
                if values:
                    col_ranges[i] = {'min': min(values), 'max': max(values)}
            except (ValueError, TypeError):
                continue
    
    # Create styling
    styling = []
    for row in data:
        row_styling = []
        for i, val in enumerate(row):
            if i in col_ranges and val != "" and pd.notna(val):
                try:
                    val_float = float(val)
                    style = get_color_for_value(val_float, col_ranges[i]['min'], col_ranges[i]['max'])
                    row_styling.append(["", f"width: var(--cell-width-3); left: auto; {style}"])
                except (ValueError, TypeError):
                    row_styling.append(["", ""])
            else:
                row_styling.append(["", ""])
        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()