import pandas as pd
import folium
from folium.plugins import MarkerCluster, HeatMap
import plotly.graph_objects as go
import plotly.express as px
from geopy.geocoders import Nominatim
from geopy.exc import GeocoderInsufficientPrivileges
import re
import streamlit as st
import time

# Streamlit title and description
st.title("米其林餐廳指南分析")
st.write("分析餐廳數據，可視化區域分佈，並在地圖上顯示位置和餐廳數量熱力圖。")

# Read data from Google Sheets
sheet_id = "1xUfnD1WCF5ldqECI8YXIko1gCpaDDCwTztL17kjI42U"
df = pd.read_csv(f"https://docs.google.com/spreadsheets/d/{sheet_id}/export?format=csv")

# Print column names and first few rows
st.write("資料框的列名：", df.columns.tolist())
st.write("資料預覽：")
st.dataframe(df.head())

# Initialize Nominatim geocoder
geolocator = Nominatim(user_agent="my_unique_app/3.0")

# Function to extract region (區域) from the address using regex
def extract_region(address):
    match = re.search(r'(.*?)區|縣|市', address)
    if match:
        return match.group(0)
    else:
        return "Unknown"

# Function to get latitude and longitude with caching
@st.cache_data
def get_lat_lon(district):
    try:
        location = geolocator.geocode(f"台南市{district}")
        if location:
            time.sleep(1)  # Delay to avoid rate limiting
            return location.latitude, longitude
    except GeocoderInsufficientPrivileges:
        st.error("地理編碼器遇到權限問題，請稍後再試。")
    return None, None

# Assuming we have a column that represents the region or can be used to derive it
# If we don't have such a column, we'll need to skip this part
if '區域' in df.columns:
    region_column = '區域'
elif '地址' in df.columns:
    df['區域'] = df['地址'].apply(extract_region)
    region_column = '區域'
else:
    st.error("無法找到區域資訊，某些分析將無法進行。")
    region_column = None

# Group the data by region and count the number of restaurants
if region_column:
    region_group = df.groupby(region_column).size().reset_index(name='Count')

    # Create hierarchical data for the Sunburst chart
    region_group['Total'] = 'All Regions'  # Add a root level
    hierarchical_data = region_group[['Total', region_column, 'Count']]

    # Plot interactive Sunburst chart
    sunburst = go.Figure(go.Sunburst(
        labels=hierarchical_data[region_column].tolist() + hierarchical_data['Total'].tolist(),
        parents=hierarchical_data['Total'].tolist() + [''],
        values=hierarchical_data['Count'].tolist() + [hierarchical_data['Count'].sum()],
        branchvalues="total",
        hovertemplate='<b>%{label}</b><br>餐廳數量: %{value}<extra></extra>',
        maxdepth=2,
    ))

    sunburst.update_layout(
        title="餐廳分佈（點擊可放大查看）",
        title_x=0.5,
        title_font=dict(size=24, family="Arial"),
        height=600,
        margin=dict(t=50, b=50, l=0, r=0)
    )

    st.subheader("餐廳分佈（Sunburst 圖）")
    st.plotly_chart(sunburst, use_container_width=True)

    # Plot bar chart with custom colors and labels
    bar_chart = go.Figure(go.Bar(
        x=region_group[region_column],
        y=region_group["Count"],
        text=region_group["Count"],
        textposition='auto',
        marker=dict(color=px.colors.qualitative.Set2)
    ))

    bar_chart.update_layout(
        title="各區域餐廳數量",
        title_x=0.5,
        title_font=dict(size=24, family="Arial"),
        height=400,
        margin=dict(t=50, b=50, l=50, r=50),
        xaxis_title="區域",
        yaxis_title="餐廳數量",
        xaxis=dict(tickangle=-45)
    )
    st.subheader("各區域餐廳數量（條形圖）")
    st.plotly_chart(bar_chart)

# Display a map using Folium if we have latitude and longitude
if '緯度' in df.columns and '經度' in df.columns:
    st.subheader("餐廳位置地圖（含數量熱力圖）")

    # Create map centered around the mean latitude and longitude
    center_lat = df['緯度'].mean()
    center_lon = df['經度'].mean()
    m = folium.Map(location=[center_lat, center_lon], zoom_start=12)

    # Add marker cluster to the map
    marker_cluster = MarkerCluster().add_to(m)

    # Prepare data for heatmap
    heat_data = []

    for index, row in df.iterrows():
        if pd.notnull(row["緯度"]) and pd.notnull(row["經度"]):
            folium.Marker(
                location=[row["緯度"], row["經度"]],
                popup=f"{row.get('店名', 'Unknown')}",
                tooltip=row.get('地址', 'Unknown')
            ).add_to(marker_cluster)
            heat_data.append([row["緯度"], row["經度"], 1])  # Weight of 1 for each restaurant

    # Add heatmap layer
    HeatMap(heat_data, radius=15, blur=10, max_zoom=1, name="餐廳數量熱力圖").add_to(m)

    # Add layer control
    folium.LayerControl().add_to(m)

    # Display the map in Streamlit
    st.components.v1.html(m._repr_html_(), height=600)
else:
    st.error("無法顯示地圖，因為缺少緯度和經度資訊。")

# Save the DataFrame to CSV with UTF-8 encoding
csv_file = "restaurants_data.csv"
df.to_csv(csv_file, encoding="utf-8-sig", index=False)

# Display download button for the CSV
st.download_button(
    label="下載餐廳數據 CSV 檔案",
    data=open(csv_file, "rb").read(),
    file_name=csv_file,
    mime="text/csv"
)