import os
import io
import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime

def plot_ketones(file):
    """
    Process uploaded file and create a ketone levels plot over time.

    Args:
        file: Uploaded file object (CSV or Excel)

    Returns:
        matplotlib figure as PIL Image
    """
    try:
        # Read the file
        if file is None:
            return None, "Please upload a file"

        # Get file extension
        filename = file.name if hasattr(file, 'name') else 'uploaded_file'

        # Read the file based on extension
        if filename.endswith('.xlsx') or filename.endswith('.xls'):
            df = pd.read_excel(file, engine='openpyxl')
        elif filename.endswith('.csv'):
            df = pd.read_csv(file)
        else:
            return None, "Please upload a CSV or Excel file"

        # Display column names for debugging
        print(f"Columns found: {df.columns.tolist()}")
        print(f"First few rows:\n{df.head()}")

        # Try to identify date and ketone columns (case-insensitive)
        date_col = None
        ketone_col = None

        for col in df.columns:
            col_lower = str(col).lower()
            if 'date' in col_lower or 'time' in col_lower:
                date_col = col
            if 'ketone' in col_lower or 'level' in col_lower or 'value' in col_lower or 'reading' in col_lower or 'sensor' in col_lower:
                ketone_col = col

        # If not found by name, assume first column is date and second is ketone level
        if date_col is None:
            date_col = df.columns[0]
        if ketone_col is None:
            ketone_col = df.columns[1] if len(df.columns) > 1 else df.columns[0]

        print(f"Using date column: {date_col}")
        print(f"Using ketone column: {ketone_col}")

        # Convert date column to datetime
        df[date_col] = pd.to_datetime(df[date_col], errors='coerce')

        # Remove rows with invalid dates or ketone values
        df = df.dropna(subset=[date_col, ketone_col])

        # Convert ketone values to numeric
        df[ketone_col] = pd.to_numeric(df[ketone_col], errors='coerce')
        df = df.dropna(subset=[ketone_col])

        # Sort by date
        df = df.sort_values(date_col)

        if len(df) == 0:
            return None, "No valid data found in the file. Please check your data format."

        # Create the plot
        fig, ax = plt.subplots(figsize=(12, 6))

        ax.plot(df[date_col], df[ketone_col], marker='o', linestyle='-',
                linewidth=2, markersize=6, color='#FF6B6B', label='Ketone Levels')

        # Formatting
        ax.set_xlabel('Date', fontsize=12, fontweight='bold')
        ax.set_ylabel('Ketone Level (mmol/L)', fontsize=12, fontweight='bold')
        ax.set_title('Ketone Levels Over Time', fontsize=14, fontweight='bold', pad=20)
        ax.grid(True, alpha=0.3, linestyle='--')
        ax.legend(loc='best')

        # Format x-axis dates
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        ax.xaxis.set_major_locator(mdates.AutoDateLocator())
        plt.xticks(rotation=45, ha='right')

        # Add some padding
        plt.tight_layout()

        # Save to bytes buffer
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        plt.close(fig)

        # Create summary stats
        summary = f"""
        **Data Summary:**
        - Total measurements: {len(df)}
        - Date range: {df[date_col].min().strftime('%Y-%m-%d')} to {df[date_col].max().strftime('%Y-%m-%d')}
        - Average ketone level: {df[ketone_col].mean():.2f} mmol/L
        - Min ketone level: {df[ketone_col].min():.2f} mmol/L
        - Max ketone level: {df[ketone_col].max():.2f} mmol/L
        """

        return buf, summary

    except Exception as e:
        error_msg = f"Error processing file: {str(e)}\n\nPlease ensure your file has:\n- A date/time column\n- A ketone level column\n- Valid numeric values"
        print(f"Error details: {e}")
        import traceback
        traceback.print_exc()
        return None, error_msg

# Create Gradio interface
with gr.Blocks(title="Ketone Level Plotter") as demo:
    gr.Markdown("""
    # 🔬 Ketone Level Plotter

    Upload a CSV or Excel file containing your ketone measurements to visualize trends over time.

    **Expected file format:**
    - First column: Date/Time (any common date format)
    - Second column: Ketone levels (numeric values in mmol/L)

    Or use column names containing 'date', 'time' for dates and 'ketone', 'level', or 'value' for measurements.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="Upload CSV or Excel file",
                file_types=[".csv", ".xlsx", ".xls"],
                type="filepath"
            )
            plot_button = gr.Button("Generate Plot", variant="primary", size="lg")

        with gr.Column(scale=2):
            image_output = gr.Image(label="Ketone Levels Plot", type="filepath")
            summary_output = gr.Markdown(label="Summary Statistics")

    gr.Markdown("""
    ---
    ### Tips:
    - Your file can be either CSV (.csv) or Excel (.xlsx, .xls)
    - The app will automatically detect date and ketone columns
    - Make sure dates are in a recognizable format (YYYY-MM-DD, MM/DD/YYYY, etc.)
    - Ketone levels should be numeric values
    """)

    # Connect the button to the function
    plot_button.click(
        fn=plot_ketones,
        inputs=[file_input],
        outputs=[image_output, summary_output]
    )

    # Also allow submission by uploading file
    file_input.upload(
        fn=plot_ketones,
        inputs=[file_input],
        outputs=[image_output, summary_output]
    )

if __name__ == "__main__":
    demo.launch()