import streamlit as st
import requests
import folium
from streamlit_folium import st_folium
import pandas as pd
import plotly.graph_objs as go
import branca.colormap as cm

# Set page layout to wide
st.set_page_config(layout="wide", page_title="Real-Time Wind Data Dashboard")

# Function to fetch GeoJSON data
@st.cache_data(ttl=300)
def fetch_geojson_data(url):
    response = requests.get(url)
    return response.json()

# Function to calculate wind statistics
def calculate_wind_stats(features):
    gust_speeds = [feature['properties']['10-Minute Maximum Gust(km/hour)'] for feature in features if
                   feature['properties']['10-Minute Maximum Gust(km/hour)'] is not None]
    mean_speeds = [feature['properties']['10-Minute Mean Speed(km/hour)'] for feature in features if
                   feature['properties']['10-Minute Mean Speed(km/hour)'] is not None]

    if not gust_speeds:
        return None, None, None, None
    avg_gust = sum(gust_speeds) / len(gust_speeds)
    min_gust = min(gust_speeds)
    max_gust = max(gust_speeds)
    avg_mean_speed = sum(mean_speeds) / len(mean_speeds) if mean_speeds else None
    return avg_gust, min_gust, max_gust, avg_mean_speed

# Function to convert wind direction to degrees
def mean_wind_direction_to_degrees(direction):
    directions = {
        'North': 0, 'Northeast': 45, 'East': 90, 'Southeast': 135,
        'South': 180, 'Southwest': 225, 'West': 270, 'Northwest': 315
    }
    return directions.get(direction, 0)

# Fetch GeoJSON data
url = 'https://csdi.vercel.app/weather/wind'
geo_data = fetch_geojson_data(url)

# Calculate wind statistics
avg_gust, min_gust, max_gust, avg_mean_speed = calculate_wind_stats(geo_data['features'])

# Create a map centered on a specific location
map_center = [22.35473034278638, 114.14827142452518]  # Coordinates of Hong Kong
my_map = folium.Map(location=map_center, zoom_start=10.35, tiles='CartoDB positron')

# Create a colormap for wind speed with limited width
colormap = cm.LinearColormap(colors=['#000000', '#0066eb', '#ff3d77', '#eb0000'],
                             vmin=0, vmax=30)
my_map.add_child(colormap)

# Function to calculate arrow size based on wind speed
def get_arrow_size(speed):
    if speed is None:
        return 20
    return max(20, min(50, speed * 2))

# Add the GeoJSON data to the map with arrow markers
for feature in geo_data['features']:
    coordinates = feature['geometry']['coordinates']
    mean_wind_direction = feature['properties']['10-Minute Mean Wind Direction(Compass points)']
    mean_speed = feature['properties']['10-Minute Mean Speed(km/hour)']

    # Skip plotting if wind direction is null
    if mean_wind_direction is None:
        continue

    # Calculate rotation angle for wind direction
    rotation_angle = mean_wind_direction_to_degrees(mean_wind_direction)

    # Calculate arrow size based on wind speed
    arrow_size = get_arrow_size(mean_speed)

    # Determine color based on wind speed
    color = colormap(mean_speed) if mean_speed is not None else 'gray'

    # Create an arrow marker for wind direction
    folium.Marker(
        location=[coordinates[1], coordinates[0]],
        icon=folium.DivIcon(html=f"""
            <div style="
                width: {arrow_size}px; height: {arrow_size}px;
                display: flex; align-items: center; justify-content: center;
                transform: rotate({rotation_angle}deg);
                ">
                <svg width="{arrow_size}" height="{arrow_size}" viewBox="0 0 24 24" fill="none" stroke="{color}" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
                    <line x1="12" y1="5" x2="12" y2="19"></line>
                    <polyline points="5 12 12 5 19 12"></polyline>
                </svg>
            </div>
        """),
        popup=folium.Popup(f"""
            <b>{feature['properties']['Automatic Weather Station']}</b><br>
            Direction: {mean_wind_direction}<br>
            Speed: {mean_speed} km/h<br>
            Max Gust: {feature['properties']['10-Minute Maximum Gust(km/hour)']} km/h
        """, max_width=300)
    ).add_to(my_map)

col1, col2, col3 = st.columns([1.65, 2, 1.15])

with col1:
    if geo_data['features']:
        wind_directions = [feature['properties']['10-Minute Mean Wind Direction(Compass points)'] for feature in
                           geo_data['features']]
        direction_counts = {d: wind_directions.count(d) for d in
                            ['North', 'Northeast', 'East', 'Southeast', 'South', 'Southwest', 'West', 'Northwest']}

        # Prepare wind speeds for each direction
        direction_speeds = {d: [] for d in
                            ['North', 'Northeast', 'East', 'Southeast', 'South', 'Southwest', 'West', 'Northwest']}
        for feature in geo_data['features']:
            direction = feature['properties']['10-Minute Mean Wind Direction(Compass points)']
            speed = feature['properties']['10-Minute Mean Speed(km/hour)']
            if direction in direction_speeds and speed is not None:
                direction_speeds[direction].append(speed)

        # Calculate average wind speed for each direction
        average_speeds = {d: sum(speeds) / len(speeds) if speeds else 0 for d, speeds in direction_speeds.items()}

        # Plot wind direction rose with average wind speed
        fig = go.Figure()

        # Add polar bar for wind direction
        fig.add_trace(go.Barpolar(
            r=[direction_counts[d] for d in direction_counts.keys()],
            theta=list(direction_counts.keys()),
            name='Wind Direction Count',
            marker_color='#0008ff',
            opacity=0.5
        ))

        # Add radial bar for average wind speed
        fig.add_trace(go.Barpolar(
            r=list(average_speeds.values()),
            theta=list(average_speeds.keys()),
            name='Average Wind Speed',
            marker_color='#ff0019',  # Orange color for wind speed
            opacity=0.5,
            thetaunit='radians',  # Ensures radial bars are correctly positioned
            base=0  # Base of the radial bars starts from 0
        ))

        fig.update_layout(
            polar=dict(
                radialaxis=dict(
                    visible=False,
                    range=[0, max(direction_counts.values())]
                ),
                angularaxis=dict(
                    tickvals=list(direction_counts.keys()),
                    ticktext=['N', 'NE', 'E', 'SE', 'S', 'SW', 'W', 'NW'],
                    rotation=90,  # Rotate to make North the top
                    direction='clockwise'
                )
            ),
            width=500,
            height=380,
            title={'text': 'Wind Direction and Average Speed Rose Plot', 'font': {'size': 18}},
            legend={'x': 0.8, 'y': 0.95}
        )

        st.plotly_chart(fig, use_container_width=True)

    if avg_gust is not None:
        col1a, col1b = st.columns(2)
        with col1a:
            st.metric(label="Avg Max Gust (km/h)", value=f"{avg_gust:.2f}")
            st.metric(label="Min Max Gust (km/h)", value=f"{min_gust}")
        with col1b:
            st.metric(label="Max Max Gust (km/h)", value=f"{max_gust}")
            if avg_mean_speed is not None:
                st.metric(label="Avg Mean Speed (km/h)", value=f"{avg_mean_speed:.2f}")
    else:
        st.write("No valid wind data available to calculate statistics.")

    gust_speeds = [feature['properties']['10-Minute Maximum Gust(km/hour)'] for feature in geo_data['features'] if
                   feature['properties']['10-Minute Maximum Gust(km/hour)'] is not None]


with col3:
    table_data = [{
        'Weather Station': feature['properties']['Automatic Weather Station'],
        'Mean Wind Direction': feature['properties']['10-Minute Mean Wind Direction(Compass points)'],
        'Mean Speed(km/hour)': feature['properties']['10-Minute Mean Speed(km/hour)'],
        'Maximum Gust(km/hour)': feature['properties']['10-Minute Maximum Gust(km/hour)']
    } for feature in geo_data['features']]

    st.dataframe(pd.DataFrame(table_data), height=600)

with col2:
    # Display map
    st_folium(my_map, width=500, height=600)