import folium as fl
from streamlit_folium import st_folium
import streamlit as st
import googlemaps
import pandas as pd
from tensorflow.keras.models import load_model
import numpy as np
import pickle

st.title('Taxi Trip Duration Prediction 🚖')

# Load trained model and preprocessor
model = load_model('Taxi_Trip_Duration.h5', compile=False)
with open('preprocessor.pkl', 'rb') as f:
    preprocessor = pickle.load(f)

# Google Maps API client
gmaps_api_key = 'AIzaSyCmyJwPBalqt1djDKK09YY2iYL7_cVA7e8' 
gmaps = googlemaps.Client(key=gmaps_api_key)

def calculate_distance(pickup_lat, pickup_lng, dropoff_lat, dropoff_lng):
    directions_result = gmaps.directions((pickup_lat, pickup_lng), (dropoff_lat, dropoff_lng), mode="driving")
    distance = directions_result[0]['legs'][0]['distance']['value'] / 1000.0  # Distance in kilometers
    return distance

def predict_duration(features):
    predicted_duration = model.predict(features)[0][0]
    return predicted_duration

# Initialize session state variables
if 'pickup_location' not in st.session_state:
    st.session_state.pickup_location = None
if 'dropoff_location' not in st.session_state:
    st.session_state.dropoff_location = None
if 'last_clicked' not in st.session_state:
    st.session_state.last_clicked = None
if 'step' not in st.session_state:
    st.session_state.step = 0
if 'distance' not in st.session_state:
    st.session_state.distance = None
if 'predicted_duration' not in st.session_state:
    st.session_state.predicted_duration = None

# Function to save location based on map click
def save_location(location_type):
    if st.session_state[location_type] is None and st.session_state.last_clicked:
        st.session_state[location_type] = (st.session_state.last_clicked['lat'], st.session_state.last_clicked['lng'])
        st.session_state.last_clicked = None
        st.session_state.step += 1
        st.experimental_rerun()  # Trigger a rerun to update the state and UI

# Step 0: Select pickup location on the map
if st.session_state.step == 0:
    st.write('Please select the pickup location on the map 🚖')
    starting_point = (37.0, 35.3213)  # Adana coordinates
    m = fl.Map(location=starting_point, zoom_start=12)
    m.add_child(fl.LatLngPopup())
    map_data = st_folium(m, height=500)
    if map_data and 'last_clicked' in map_data and map_data['last_clicked']:
        st.session_state.last_clicked = map_data['last_clicked']
        if st.button("Save Pickup Location"):
            save_location('pickup_location')

# Step 1: Select dropoff location on the map
if st.session_state.step == 1 and st.session_state.pickup_location is not None:
    st.write('Pickup location selected. Now select the dropoff location 📍')
    m = fl.Map(location=st.session_state.pickup_location, zoom_start=12)
    fl.Marker(st.session_state.pickup_location, popup="Pickup Location").add_to(m)
    m.add_child(fl.LatLngPopup())
    map_data = st_folium(m, height=500)
    if map_data and 'last_clicked' in map_data and map_data['last_clicked']:
        st.session_state.last_clicked = map_data['last_clicked']
        if st.button("Save Dropoff Location"):
            save_location('dropoff_location')

# Step 2: Enter additional trip details and display prediction
if st.session_state.step == 2 and st.session_state.pickup_location is not None and st.session_state.dropoff_location is not None:
    st.write('Dropoff location selected. Please enter additional trip details.')

    passenger_count = st.slider('Passenger Count', min_value=1, max_value=6, value=1)
    hour = st.slider('Hour of Day', min_value=0, max_value=23, value=12)
    day = st.slider('Day of Month', min_value=1, max_value=31, value=15)
    weekday = st.selectbox('Weekday', ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday'])

    weekday_dict = {'Monday': 0, 'Tuesday': 1, 'Wednesday': 2, 'Thursday': 3, 'Friday': 4, 'Saturday': 5, 'Sunday': 6}
    weekday_numeric = weekday_dict[weekday]

    # Calculate distance between pickup and dropoff locations
    pickup_lat, pickup_lng = st.session_state.pickup_location
    dropoff_lat, dropoff_lng = st.session_state.dropoff_location
    distance = calculate_distance(pickup_lat, pickup_lng, dropoff_lat, dropoff_lng)
    st.session_state.distance = distance

    # Prepare input features for prediction
    features = np.array([[pickup_lat, pickup_lng, dropoff_lat, dropoff_lng, passenger_count, distance, hour, day, weekday_numeric]])
    features_scaled = preprocessor.transform(features)

    # Predict trip duration
    predicted_duration = predict_duration(features_scaled)
    st.session_state.predicted_duration = predicted_duration

    st.write(f'Predicted Trip Duration: {predicted_duration:.2f} seconds')

    # Show map with pickup and dropoff locations and route
    m = fl.Map(location=st.session_state.pickup_location, zoom_start=12)
    fl.Marker(st.session_state.pickup_location, popup="Pickup Location").add_to(m)
    fl.Marker(st.session_state.dropoff_location, popup="Dropoff Location").add_to(m)
    fl.PolyLine([st.session_state.pickup_location, st.session_state.dropoff_location], color="blue").add_to(m)
    st_folium(m, height=500)