import numpy as np
import pandas as pd
import time
import gradio as gr
import torch
import matplotlib.pyplot as plt
from PIL import Image
import cv2
import pygame
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import os  # Import the 'os' module

# Flag to track if pygame is initialized successfully
pygame_initialized = False

try:
    # Initialize pygame mixer for sound
    pygame.mixer.init()
    pygame_initialized = True
    print("Pygame mixer initialized successfully.")
except Exception as e:
    print(f"Error initializing pygame: {e}. Sound will be disabled.")

# Check if GPU is available
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    device = torch.device("cuda:0")
    print("Using GPU:", torch.cuda.get_device_name(0))
else:
    device = torch.device("cpu")
    print("Using CPU")

# Simulated dataset of finger pressure patterns and blood glucose readings
# In a real application, this would be replaced with actual training data
def generate_simulated_data(n_samples=1000):
    # Generate synthetic features that might correlate with blood glucose
    # In reality, these would be derived from screen pressure, touch area,
    # blood perfusion estimated from camera, etc.
    np.random.seed(42)

    # Feature 1: Simulated pressure values (0-100)
    pressure = np.random.normal(60, 15, n_samples).clip(20, 100)

    # Feature 2: Simulated touch area (mm²)
    touch_area = np.random.normal(150, 30, n_samples).clip(80, 220)

    # Feature 3: Simulated blood perfusion level (0-1)
    perfusion = np.random.normal(0.7, 0.15, n_samples).clip(0.3, 1.0)

    # Feature 4: Simulated tissue opacity (0-1)
    opacity = np.random.normal(0.5, 0.1, n_samples).clip(0.2, 0.8)

    # Generate glucose levels with some correlation to the features
    # Normal range: 70-140 mg/dL, with some values outside this range
    base_glucose = 100 + (pressure - 60) * 0.5 - (touch_area - 150) * 0.1 + (perfusion - 0.7) * 40 - (opacity - 0.5) * 30
    glucose = base_glucose + np.random.normal(0, 15, n_samples)
    glucose = glucose.clip(40, 400)  # Set realistic min/max values

    # Create a dataframe
    df = pd.DataFrame({
        'pressure': pressure,
        'touch_area': touch_area,
        'perfusion': perfusion,
        'opacity': opacity,
        'glucose': glucose
    })

    return df

# Simple PyTorch model for glucose prediction
class GlucosePredictor(torch.nn.Module):
    def __init__(self):
        super(GlucosePredictor, self).__init__()
        self.model = torch.nn.Sequential(
            torch.nn.Linear(4, 16),
            torch.nn.ReLU(),
            torch.nn.Linear(16, 32),
            torch.nn.ReLU(),
            torch.nn.Linear(32, 16),
            torch.nn.ReLU(),
            torch.nn.Linear(16, 1)
        )

    def forward(self, x):
        return self.model(x)

# Train the model
def train_model(df):
    # Split features and target
    X = df[['pressure', 'touch_area', 'perfusion', 'opacity']].values
    y = df['glucose'].values.reshape(-1, 1)

    # Scale the data
    scaler_X = StandardScaler()
    scaler_y = StandardScaler()

    X_scaled = scaler_X.fit_transform(X)
    y_scaled = scaler_y.fit_transform(y)

    # Split into train and test sets
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y_scaled, test_size=0.2, random_state=42)

    # Convert to PyTorch tensors
    X_train_tensor = torch.FloatTensor(X_train).to(device)
    y_train_tensor = torch.FloatTensor(y_train).to(device)
    X_test_tensor = torch.FloatTensor(X_test).to(device)
    y_test_tensor = torch.FloatTensor(y_test).to(device)

    # Initialize the model and move to GPU if available
    model = GlucosePredictor().to(device)

    # Loss function and optimizer
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # Training loop
    epochs = 100
    for epoch in range(epochs):
        # Forward pass
        y_pred = model(X_train_tensor)
        loss = criterion(y_pred, y_train_tensor)

        # Backward pass and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (epoch+1) % 10 == 0:
            print(f'Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}')

    # Evaluate the model
    model.eval()
    with torch.no_grad():
        y_pred_test = model(X_test_tensor)
        test_loss = criterion(y_pred_test, y_test_tensor)
        print(f'Test Loss: {test_loss.item():.4f}')

    return model, scaler_X, scaler_y

# Function to play beep sound
def play_beep(glucose_value):
    """Play different sounds based on glucose level"""
    global pygame_initialized  # Access the global flag

    if not pygame_initialized:
        print("Pygame not initialized, skipping sound.")
        return  # Exit if pygame is not initialized

    try:
        # Define frequency and duration based on glucose level
        if glucose_value < 70:
            # Low glucose - urgent double beep (low frequency)
            freq = 300
            duration = 300
            # Play twice with a short pause
            for _ in range(2):
                sound_array = np.sin(2*np.pi*np.arange(44100)*freq/44100).astype(np.float32)
                sound = pygame.sndarray.make_sound(sound_array)
                sound.play()
                time.sleep(duration/1000)
                time.sleep(0.1)  # Pause between beeps
        elif glucose_value <= 140:
            # Normal glucose - standard beep (medium frequency)
            freq = 440
            duration = 400
            sound_array = np.sin(2*np.pi*np.arange(44100)*freq/44100).astype(np.float32)
            sound = pygame.sndarray.make_sound(sound_array)
            sound.play()
            time.sleep(duration/1000)
        elif glucose_value <= 200:
            # Elevated glucose - medium alert beep (higher frequency)
            freq = 600
            duration = 400
            sound_array = np.sin(2*np.pi*np.arange(44100)*freq/44100).astype(np.float32)
            sound = pygame.sndarray.make_sound(sound_array)
            sound.play()
            time.sleep(duration/1000)
        else:
            # High glucose - urgent beep (high frequency)
            freq = 800
            duration = 300
            # Play twice with a short pause
            for _ in range(3):
                sound_array = np.sin(2*np.pi*np.arange(44100)*freq/44100).astype(np.float32)
                sound = pygame.sndarray.make_sound(sound_array)
                sound.play()
                time.sleep(duration/1000)
                time.sleep(0.1)  # Pause between beeps

        print(f"Beep sound played for glucose level: {glucose_value}")
    except Exception as e:
        print(f"Error playing sound: {e}")

# Simulate camera usage for blood perfusion measurement
def activate_camera():
    """
    In a real application, this would activate the device's camera
    to measure blood perfusion in the fingertip.
    Here we just simulate this process.
    """
    print("Camera activated for blood perfusion analysis")

    # Attempt to use real camera if available (for demonstration)
    try:
        cap = cv2.VideoCapture(0)
        if cap.isOpened():
            ret, frame = cap.read()
            if ret:
                # In a real app, we would analyze this frame
                # For now, just display that we captured it
                print("Camera frame captured")
                # Here you'd analyze blood perfusion from the image
                # For example by looking at the red channel intensity

                # Convert to grayscale and use a placeholder image for demonstration
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img = Image.fromarray(frame_rgb)

                # Clean up
                cap.release()

                return img
        else:
            print("Could not open camera - using simulated data")

    except Exception as e:
        print(f"Camera error: {e} - using simulated data")

    # If camera isn't working, return a placeholder
    placeholder = np.ones((240, 320, 3), dtype=np.uint8) * 200
    # Add text to placeholder
    cv2.putText(placeholder, "Camera Simulation", (50, 120),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)

    return Image.fromarray(placeholder)

# Simulate finger press data collection
def collect_finger_data():
    """
    In a real application, this would collect data from:
    1. Screen pressure sensors
    2. Touch area measurement
    3. Camera-based blood perfusion estimation
    4. Optional: PPG (photoplethysmography) if available

    Here we're just simulating the data.
    """
    # Simulate data collection with some random values
    pressure = np.random.normal(60, 10)
    touch_area = np.random.normal(150, 20)
    perfusion = np.random.normal(0.7, 0.1)
    opacity = np.random.normal(0.5, 0.08)

    # Add some constraints to make values realistic
    pressure = max(20, min(100, pressure))
    touch_area = max(80, min(220, touch_area))
    perfusion = max(0.3, min(1.0, perfusion))
    opacity = max(0.2, min(0.8, opacity))

    return pressure, touch_area, perfusion, opacity

# Function to predict glucose level
def predict_glucose(model, scaler_X, scaler_y, features):
    # Scale the features
    features_scaled = scaler_X.transform(np.array(features).reshape(1, -1))

    # Convert to tensor and predict
    features_tensor = torch.FloatTensor(features_scaled).to(device)
    with torch.no_grad():
        prediction_scaled = model(features_tensor)

    # Convert back to original scale
    prediction = scaler_y.inverse_transform(prediction_scaled.cpu().numpy())

    return prediction[0][0]

# Create a fancy glucose meter display
def create_glucose_meter(glucose_value):
    # Define colors based on glucose range
    if glucose_value < 70:
        color = 'blue'  # Low
        status = 'LOW'
    elif glucose_value <= 140:
        color = 'green'  # Normal
        status = 'NORMAL'
    elif glucose_value <= 200:
        color = 'orange'  # Elevated
        status = 'ELEVATED'
    else:
        color = 'red'  # High
        status = 'HIGH'

    # Create a figure
    fig, ax = plt.subplots(figsize=(6, 4))

    # Create a gauge-like visualization
    ax.add_patch(plt.Rectangle((-1, -1), 2, 2, fc='lightgray', ec='gray'))

    # Create a gauge arc
    theta = np.linspace(-0.75 * np.pi, 0.75 * np.pi, 100)
    r = 0.8
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    ax.plot(x, y, 'k-', lw=2)

    # Create tick marks
    for i in range(40, 401, 40):
        angle = -0.75 * np.pi + (i - 40) / (400 - 40) * 1.5 * np.pi
        x_tick = (r + 0.1) * np.cos(angle)
        y_tick = (r + 0.1) * np.sin(angle)
        ax.plot([r * np.cos(angle), x_tick], [r * np.sin(angle), y_tick], 'k-', lw=1)
        ax.text(x_tick * 1.1, y_tick * 1.1, str(i), ha='center', va='center', fontsize=8)

    # Create the needle
    angle = -0.75 * np.pi + (glucose_value - 40) / (400 - 40) * 1.5 * np.pi
    ax.plot([0, 0.9 * np.cos(angle)], [0, 0.9 * np.sin(angle)], color=color, lw=3)

    # Add a center circle
    ax.add_patch(plt.Circle((0, 0), 0.05, fc=color, ec='k'))

    # Add text
    ax.text(0, -0.4, f"{glucose_value:.1f} mg/dL", ha='center', va='center', fontsize=14, fontweight='bold')
    ax.text(0, -0.6, status, ha='center', va='center', fontsize=12, color=color, fontweight='bold')

    # Remove axes
    ax.set_xlim(-1.2, 1.2)
    ax.set_ylim(-1.2, 1)
    ax.axis('off')
    ax.set_aspect('equal')

    # Add title
    ax.set_title('Digital Glucometer Reading', fontsize=16, pad=20)

    # Add disclaimer
    fig.text(0.5, 0.01, 'PROTOTYPE ONLY - NOT FOR MEDICAL USE',
             ha='center', va='bottom', fontsize=10, style='italic', color='gray')

    return fig
# Gradio interface for finger press simulation with yield for progressive updates
def finger_press_simulation():
    """
    This function simulates the process of measuring glucose from a finger press
    with progressive updates to show the process step by step.
    """
    # Step 1: Display a message that the camera is activating
    collection_info = "Activating camera for blood perfusion analysis..."
    camera_placeholder = np.ones((240, 320, 3), dtype=np.uint8) * 200
    cv2.putText(camera_placeholder, "Activating Camera...", (50, 120),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)
    camera_img = Image.fromarray(camera_placeholder)

    meter_fig = plt.figure(figsize=(6, 4))
    plt.text(0.5, 0.5, "Measuring...", ha='center', va='center', fontsize=18)
    plt.axis('off')

    yield collection_info, camera_img, meter_fig

    # Step 2: Activate the camera (simulated or real)
    time.sleep(1)
    collection_info = "Camera activated, analyzing blood perfusion..."
    try:
        camera_img = activate_camera()
    except Exception as e:
        print(f"Error activating camera: {e}.  Using placeholder image.")
        camera_placeholder = np.ones((240, 320, 3), dtype=np.uint8) * 200
        cv2.putText(camera_placeholder, "Camera Error", (50, 120),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)
        camera_img = Image.fromarray(camera_placeholder)

    yield collection_info, camera_img, meter_fig

    # Step 3: Collect finger data
    time.sleep(1)
    collection_info = "Collecting finger press data..."
    pressure, touch_area, perfusion, opacity = collect_finger_data()
    yield collection_info, camera_img, meter_fig

    # Step 4: Predict glucose level
    time.sleep(1)
    collection_info = "Analyzing data and predicting glucose level..."
    features = [pressure, touch_area, perfusion, opacity]
    glucose_value = predict_glucose(model, scaler_X, scaler_y, features)
    yield collection_info, camera_img, meter_fig

    # Step 5: Display final results
    time.sleep(1)
    collection_info = f"Measurement complete!\nPressure: {pressure:.1f}\nTouch Area: {touch_area:.1f}\nPerfusion: {perfusion:.2f}\nOpacity: {opacity:.2f}"
    meter_fig = create_glucose_meter(glucose_value)

    # Play sound based on glucose level
    try:
        play_beep(glucose_value)
    except Exception as e:
        print(f"Error playing sound: {e}. Skipping sound.")

    yield collection_info, camera_img, meter_fig

# Main execution
# [Previous code remains unchanged up to the main execution block]

if __name__ == "__main__":
    # Generate simulated data and train the model
    df = generate_simulated_data()
    model, scaler_X, scaler_y = train_model(df)

    # Create Gradio interface
    with gr.Blocks(title="Non-Invasive Glucose Monitor Prototype") as demo:
        gr.Markdown("# Non-Invasive Glucose Monitor Prototype")
        gr.Markdown("Press 'Measure Glucose' to simulate a finger-based glucose measurement")

        with gr.Row():
            with gr.Column(scale=1):
                measure_button = gr.Button("Measure Glucose")
                info_output = gr.Textbox(label="Measurement Process")
            with gr.Column(scale=1):
                camera_output = gr.Image(label="Camera Feed")
            with gr.Column(scale=1):
                meter_output = gr.Plot(label="Glucose Meter")

        gr.Markdown("*Note: This is a prototype simulation only - not for actual medical use*")

        # Connect the button to the simulation function (removed _js parameter)
        measure_button.click(
            fn=finger_press_simulation,
            inputs=None,
            outputs=[info_output, camera_output, meter_output]
        )

    # Launch the interface
    try:
        demo.launch()
    except Exception as e:
        print(f"Error launching Gradio: {e}")
        print("Attempting to launch in share mode (for debugging).")
        demo.launch(share=True)  # Try share mode as a fallback