import gradio as gr
import spaces
from codecarbon import EmissionsTracker

# Import necessary libraries
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, accuracy_score
import pandas as pd
import numpy as np

# Let's create a sample dataset (you can replace this with your own data)
def create_sample_data():
    np.random.seed(42)
    n_samples = 10000
    
    # Create features (X)
    X = np.random.randn(n_samples, 4)  # 4 features
    
    # Create target (y) - binary classification
    y = (X[:, 0] + X[:, 1] + X[:, 2] > 0).astype(int)
    
    return X, y

# Get data (replace this with your data loading code)
X, y = create_sample_data()
tracker = EmissionsTracker()

@spaces.GPU
def submit(username):

    tracker.start()

    tracker.start_task("train_model")
    # Split the data into training and testing sets
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
        
    # Initialize the model
    rf_model = RandomForestClassifier(
        n_estimators=1000,
        max_depth=5,
        random_state=42
    )

    # Train the model
    print("Training the model...")
    rf_model.fit(X_train, y_train)

    training_emissions = tracker.stop_task()

    tracker.start_task("inference")
    rf_model.predict(X_test)
    inference_emissions = tracker.stop_task()

    emissions = inference_emissions.emissions
    energy = inference_emissions.energy_consumed

    return [emissions, energy]

# Update the interface configuration
demo = gr.Interface(
    fn=submit,
    inputs=gr.Textbox(label="Username"),
    outputs=[
        gr.Number(label="Emissions (kgCO2eq)", precision=6),
        gr.Number(label="Energy Consumed (kWh)", precision=6)
    ],
    title="Carbon Emissions Tracker",
    description="Track the carbon emissions and energy consumption of model training and inference."
)

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch()