import streamlit as st
import folium
from streamlit_folium import folium_static
from folium.plugins import MarkerCluster
import time
import math

# Wide mode
st.set_page_config(layout="wide")

# Define California medical centers data
california_med_centers = [
    ('UCSF Medical Center', 37.7631, -122.4576, 'General medical and surgical', 'San Francisco'),
    ('Cedars-Sinai Medical Center', 34.0762, -118.3790, 'Heart specialty', 'Los Angeles'),
    ('Stanford Health Care', 37.4331, -122.1750, 'Teaching hospital', 'Stanford'),
    ('UCLA Medical Center', 34.0659, -118.4466, 'Research and teaching', 'Los Angeles'),
    ('Scripps La Jolla Hospitals', 32.8851, -117.2255, 'Multiple specialties', 'La Jolla'),
    ('Sharp Memorial Hospital', 32.7992, -117.1542, 'Trauma center', 'San Diego'),
    ('Huntington Hospital', 34.1330, -118.1475, 'Non-profit hospital', 'Pasadena'),
    ('Hoag Memorial Hospital', 33.6045, -117.8664, 'Community hospital', 'Newport Beach'),
    ('UCSD Medical Center', 32.7554, -117.1682, 'Academic health center', 'San Diego'),
    ('UC Davis Medical Center', 38.5539, -121.4554, 'Public academic health', 'Sacramento'),
    ('John Muir Medical Center', 37.9192, -122.0426, 'Heart care', 'Walnut Creek'),
    ('Santa Clara Valley Medical Center', 37.3121, -121.9769, 'County hospital', 'San Jose'),
    ('Kaiser Permanente San Francisco', 37.7741, -122.4179, 'Health maintenance organization', 'San Francisco'),
    ('City of Hope', 34.1285, -117.9665, 'Cancer center', 'Duarte'),
    ('UCI Medical Center', 33.7886, -117.8572, 'University hospital', 'Orange'),
    ('Good Samaritan Hospital', 34.0506, -118.2831, 'Private hospital', 'Los Angeles'),
    ('Los Angeles County General', 34.0581, -118.2917, 'Public hospital', 'Los Angeles'),
    ('California Pacific Medical Center', 37.7864, -122.4357, 'Private non-profit', 'San Francisco'),
    ('Sutter Roseville Medical Center', 38.7521, -121.2760, 'General medical and surgical', 'Roseville'),
    ('St. Joseph Hospital', 33.7821, -117.9188, 'Faith-based care', 'Orange')
]

# Initialize session state
if 'car_position' not in st.session_state:
    st.session_state.car_position = [37.7631, -122.4576]
if 'is_moving' not in st.session_state:
    st.session_state.is_moving = False
if 'start' not in st.session_state:
    st.session_state.start = california_med_centers[0]
if 'destination' not in st.session_state:
    st.session_state.destination = california_med_centers[1]
if 'start_time' not in st.session_state:
    st.session_state.start_time = None
if 'elapsed_time' not in st.session_state:
    st.session_state.elapsed_time = 0
if 'arrived' not in st.session_state:
    st.session_state.arrived = False

# Distance calculation (Haversine in miles)
def calculate_distance(lat1, lon1, lat2, lon2):
    R = 3958.8
    lat1_rad, lon1_rad = math.radians(lat1), math.radians(lon1)
    lat2_rad, lon2_rad = math.radians(lat2), math.radians(lon2)
    dlat = lat2_rad - lat1_rad
    dlon = lon2_rad - lon1_rad
    a = math.sin(dlat / 2)**2 + math.cos(lat1_rad) * math.cos(lat2_rad) * math.sin(dlon / 2)**2
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
    return R * c

# Create map
def create_map():
    m = folium.Map(location=st.session_state.car_position, zoom_start=13, tiles="CartoDB Positron")
    marker_cluster = MarkerCluster().add_to(m)
    for center in california_med_centers:
        folium.Marker(
            location=[center[1], center[2]],
            popup=f'<b>{center[0]}</b><br>{center[3]}<br>{center[4]}',
            icon=folium.Icon(color='red')
        ).add_to(marker_cluster)
    folium.Marker(
        location=st.session_state.car_position,
        popup='Hovercar',
        icon=folium.Icon(color='purple', icon='rocket', prefix='fa')
    ).add_to(m)
    return m

# Move hovercar with constant speed
def move_car(total_distance, step_time):
    target = [st.session_state.destination[1], st.session_state.destination[2]]
    current = st.session_state.car_position
    remaining_distance = calculate_distance(current[0], current[1], target[0], target[1])
    
    # Speed in miles per second
    total_time = 25.0  # Total trip time in seconds
    speed = total_distance / total_time  # miles per second
    step_distance = speed * step_time  # Distance per step
    
    # Calculate direction and move
    total_lat_diff = target[0] - current[0]
    total_lon_diff = target[1] - current[1]
    total_hyp = math.sqrt(total_lat_diff**2 + total_lon_diff**2)
    if total_hyp > 0:
        lat_step = (total_lat_diff / total_hyp) * step_distance / (calculate_distance(current[0], current[1], current[0] + 0.01, current[1]) / 0.01)
        lon_step = (total_lon_diff / total_hyp) * step_distance / (calculate_distance(current[0], current[1], current[0], current[1] + 0.01) / 0.01)
        st.session_state.car_position[0] += lat_step
        st.session_state.car_position[1] += lon_step
    
    # Check arrival
    if remaining_distance < 0.1:  # Within 0.1 miles
        st.session_state.car_position = target
        st.session_state.is_moving = False
        st.session_state.arrived = True

# Calculate speed and ETA
total_distance = calculate_distance(st.session_state.start[1], st.session_state.start[2],
                                    st.session_state.destination[1], st.session_state.destination[2])
step_time = 0.5
total_time = 25.0  # Fixed total trip time
speed = total_distance / (total_time / 3600)  # mph

# Layout: 75% left, 25% right
left_col, right_col = st.columns([3, 1])

# Left column: Map (16x9, 2/3 larger)
with left_col:
    st.markdown("### 🚀 Hovercar Command Center 🛸")
    m = create_map()
    folium_static(m, width=2133, height=1200)  # 16:9 ratio, 2/3 larger than 1280x720

    # Controls
    col1, col2, col3 = st.columns(3)
    with col1:
        if st.button("🚀 Launch Hovercar"):
            st.session_state.is_moving = True
            st.session_state.start_time = time.time()
            st.session_state.arrived = False
    with col2:
        if st.button("🛑 Halt Hovercar"):
            st.session_state.is_moving = False
            st.session_state.start_time = None
    with col3:
        if st.button("🔄 Reset to Start"):
            st.session_state.car_position = [st.session_state.start[1], st.session_state.start[2]]
            st.session_state.is_moving = False
            st.session_state.start_time = None
            st.session_state.elapsed_time = 0
            st.session_state.arrived = False
            st.rerun()

    # Celebration on arrival
    if st.session_state.arrived:
        st.markdown("""
            <style>
            .confetti {
                position: absolute;
                width: 10px;
                height: 10px;
                background-color: #f00;
                animation: fall 3s infinite;
            }
            @keyframes fall {
                0% { transform: translateY(-100vh); }
                100% { transform: translateY(100vh); }
            }
            </style>
            <div class="confetti" style="left: 10%; background-color: #f00;"></div>
            <div class="confetti" style="left: 20%; background-color: #0f0;"></div>
            <div class="confetti" style="left: 30%; background-color: #00f;"></div>
            <div class="confetti" style="left: 40%; background-color: #ff0;"></div>
            <div class="confetti" style="left: 50%; background-color: #f0f;"></div>
            <div class="confetti" style="left: 60%; background-color: #0ff;"></div>
        """, unsafe_allow_html=True)
        st.success("🎉 Touchdown! Hovercar has landed!")

        # Easter egg: Secret message
        if st.session_state.destination[0] == "Stanford Health Care":
            st.warning("🦄 Easter Egg Unlocked: Hovercar found the secret unicorn base at Stanford!")

# Right column: Stats and two sub-columns for buttons
with right_col:
    st.markdown("### 📡 Mission Status")
    remaining_distance = calculate_distance(
        st.session_state.car_position[0], st.session_state.car_position[1],
        st.session_state.destination[1], st.session_state.destination[2]
    )
    st.write(f"**Start**: {st.session_state.start[0]} ({st.session_state.start[1]:.4f}, {st.session_state.start[2]:.4f})")
    st.write(f"**Destination**: {st.session_state.destination[0]} ({st.session_state.destination[1]:.4f}, {st.session_state.destination[2]:.4f})")
    st.write(f"**Total Distance**: {total_distance:.1f} miles")
    st.write(f"**Remaining Distance**: {remaining_distance:.1f} miles")
    st.write(f"**Speed**: {speed:.1f} mph")
    st.write(f"**ETA**: {total_time:.1f} seconds")
    if st.session_state.is_moving:
        elapsed = st.session_state.elapsed_time
        st.write(f"**Elapsed Time**: {elapsed:.1f} seconds")

    # Two sub-columns for start and destination
    start_col, dest_col = st.columns(2)
    
    with start_col:
        st.write("**🚦 Start Points**")
        for center in california_med_centers:
            if st.button(f"📍 {center[0]}", key=f"start_{center[0]}"):
                st.session_state.start = center
                st.session_state.car_position = [center[1], center[2]]
                st.session_state.is_moving = False
                st.session_state.start_time = None
                st.session_state.elapsed_time = 0
                st.session_state.arrived = False
                st.rerun()

    with dest_col:
        st.write("**🏁 Destinations**")
        for center in california_med_centers:
            if st.button(f"🎯 {center[0]}", key=f"dest_{center[0]}"):
                st.session_state.destination = center
                st.session_state.is_moving = False
                st.session_state.start_time = None
                st.session_state.elapsed_time = 0
                st.session_state.arrived = False
                st.rerun()

# Movement logic
if st.session_state.is_moving:
    move_car(total_distance, step_time)
    if st.session_state.start_time:
        st.session_state.elapsed_time = time.time() - st.session_state.start_time
    time.sleep(step_time)
    st.rerun()