import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
import gradio as gr

def load_and_train_model():
    # Load data
    train_data = pd.read_csv('train.csv.zip')
    
    # Select important features
    important_features = ['Elevation', 'Horizontal_Distance_To_Roadways', 
                         'Horizontal_Distance_To_Fire_Points', 
                         'Horizontal_Distance_To_Hydrology',
                         'Vertical_Distance_To_Hydrology']
    
    X = train_data[important_features]
    y = train_data['Cover_Type']
    
    # Train model
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X, y)
    
    return model

def predict(elevation, distance_to_roads, distance_to_fire, distance_to_water, vertical_distance):
    input_data = np.array([[
        elevation,
        distance_to_roads,
        distance_to_fire,
        distance_to_water,
        vertical_distance
    ]])
    
    prediction = model.predict(input_data)[0]
    
    forest_types = {
        1: "Spruce/Fir",
        2: "Lodgepole Pine",
        3: "Ponderosa Pine",
        4: "Cottonwood/Willow",
        5: "Aspen",
        6: "Douglas-fir",
        7: "Krummholz"
    }
    
    confidence = max(model.predict_proba(input_data)[0]) * 100
    
    return f"{forest_types[prediction]} (Confidence: {confidence:.1f}%)"

# Train model at startup
print("Loading and training model...")
model = load_and_train_model()
print("Model training completed!")

# Create Gradio interface with improved inputs
inputs = [
    gr.Slider(
        minimum=1800, maximum=4000, step=10,
        label="Elevation (meters above sea level)",
        info="Height of the location above sea level",
        value=2500
    ),
    gr.Slider(
        minimum=0, maximum=8000, step=50,
        label="Distance to Nearest Road (meters)",
        info="How far is the nearest road?",
        value=500
    ),
    gr.Slider(
        minimum=0, maximum=8000, step=50,
        label="Distance to Fire Points (meters)",
        info="Distance to nearest fire ignition points",
        value=1000
    ),
    gr.Slider(
        minimum=0, maximum=1000, step=10,
        label="Distance to Water (meters)",
        info="Distance to nearest water source",
        value=200
    ),
    gr.Slider(
        minimum=-500, maximum=500, step=10,
        label="Vertical Distance to Water (meters)",
        info="Vertical distance to nearest water source (negative means below)",
        value=0
    )
]

output = gr.Text(
    label="Predicted Forest Type",
    info="Prediction with confidence level"
)

interface = gr.Interface(
    fn=predict,
    inputs=inputs,
    outputs=output,
    title="🌲 Forest Cover Type Predictor",
    description="""
    Predict the type of forest cover based on environmental features.
    Just adjust the sliders to match your location's characteristics!
    """,
    examples=[
        [2596, 510, 6279, 258, 0],
        [2300, 300, 1500, 100, -50],
        [3500, 1000, 2000, 500, 200]
    ],
    theme="soft"
)

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