import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import folium
from streamlit_folium import st_folium

# ----------------------------------------------------
# 1. Load data
# ----------------------------------------------------
@st.cache_data
def load_data():
    # Load daily and monthly CSV from local files (or a URL if needed)
    daily_df = pd.read_csv("daily_data_2013_2024.csv", parse_dates=["date"])
    monthly_df = pd.read_csv("monthly_data_2013_2024.csv")
    
    # If monthly_df also needs a 'date' column for plotting (like first day of month), you can create:
    # monthly_df["date"] = pd.to_datetime(monthly_df["year"].astype(str) + "-" + monthly_df["month"].astype(str) + "-01")

    return daily_df, monthly_df

daily_data, monthly_data = load_data()

# Pre-define your location dictionary so we can map lat/lon
LOCATIONS = {
    "Karagwe": {"lat": -1.7718, "lon": 30.9876},
    "Masasi":  {"lat": -10.7167, "lon": 38.8000},
    "Igunga":  {"lat": -4.2833, "lon": 33.8833}
}

# ----------------------------------------------------
# 2. Streamlit UI Layout
# ----------------------------------------------------
st.title("Malaria & Dengue Outbreak Analysis (2013–2024)")

st.sidebar.header("Filters & Options")

# Choose disease type to focus on
disease_choice = st.sidebar.radio("Select Disease", ["Malaria", "Dengue"])

# Choose data granularity
data_choice = st.sidebar.radio("Data Granularity", ["Monthly", "Daily"])

# Let user filter location(s)
location_list = list(LOCATIONS.keys())
selected_locations = st.sidebar.multiselect("Select Location(s)", location_list, default=location_list)

# For monthly data, let user select a year range
if data_choice == "Monthly":
    year_min = int(monthly_data["year"].min())
    year_max = int(monthly_data["year"].max())
    year_range = st.sidebar.slider("Select Year Range", 
                                   min_value=year_min, 
                                   max_value=year_max,
                                   value=(year_min, year_max))
# For daily data, let user select a date range
else:
    date_min = daily_data["date"].min()
    date_max = daily_data["date"].max()
    date_range = st.sidebar.date_input("Select Date Range", 
                                       [date_min, date_max],
                                       min_value=date_min,
                                       max_value=date_max)

# ----------------------------------------------------
# 3. Filter data based on user input
# ----------------------------------------------------
if data_choice == "Monthly":
    # Subset monthly data for selected locations
    df = monthly_data[monthly_data["location"].isin(selected_locations)].copy()
    # Filter year range
    df = df[(df["year"] >= year_range[0]) & (df["year"] <= year_range[1])]
    
    # Create a "date" column for monthly plotting (1st of each month)
    # This can help create a time-series for Plotly
    df["date"] = pd.to_datetime(df["year"].astype(str) + "-" + df["month"].astype(str) + "-01")
    
else:
    # Subset daily data
    df = daily_data[daily_data["location"].isin(selected_locations)].copy()
    # Filter date range
    df = df[(df["date"] >= pd.to_datetime(date_range[0])) & (df["date"] <= pd.to_datetime(date_range[1]))]

# ----------------------------------------------------
# 4. Interactive Plotly Time-Series
# ----------------------------------------------------
st.subheader(f"{data_choice} {disease_choice} Risk & Climate Parameters")

# Decide which columns are relevant
risk_col = "malaria_risk" if disease_choice == "Malaria" else "dengue_risk"

if data_choice == "Monthly":
    # We'll plot a line chart of risk, temperature, and rainfall vs. date
    fig = px.line(df, x="date", y=risk_col, color="location", 
                  title=f"{disease_choice} Risk Over Time ({data_choice})")
    fig.update_layout(yaxis_title="Risk (0–1)")
    
    st.plotly_chart(fig, use_container_width=True)
    
    # Optionally plot temperature / rainfall on another figure
    col1, col2 = st.columns(2)
    with col1:
        fig_temp = px.line(df, x="date", y="temp_avg", color="location",
                           title="Average Temperature (°C)")
        st.plotly_chart(fig_temp, use_container_width=True)
    with col2:
        # 'monthly_rainfall_mm' is total monthly rainfall
        fig_rain = px.line(df, x="date", y="monthly_rainfall_mm", color="location",
                           title="Monthly Rainfall (mm)")
        st.plotly_chart(fig_rain, use_container_width=True)

    # Show outbreak flags if focusing on monthly
    if disease_choice == "Malaria":
        flag_col = "malaria_outbreak"
    else:
        flag_col = "dengue_outbreak"

    outbreak_months = df[df[flag_col] == True]
    if not outbreak_months.empty:
        st.write(f"**Months with likely {disease_choice} outbreak:**")
        st.dataframe(outbreak_months[["location","year","month","temp_avg","humidity","monthly_rainfall_mm",flag_col]])
    else:
        st.write(f"No months meet the {disease_choice} outbreak criteria in this selection.")

else:
    # For daily data, plot daily risk
    fig = px.line(df, x="date", y=risk_col, color="location",
                  title=f"{disease_choice} Daily Risk Over Time (2013–2024)")
    fig.update_layout(yaxis_title="Risk (0–1)")
    st.plotly_chart(fig, use_container_width=True)
    
    # Similarly, temperature & daily rainfall
    col1, col2 = st.columns(2)
    with col1:
        fig_temp = px.line(df, x="date", y="temp_avg", color="location",
                           title="Daily Avg Temperature (°C)")
        st.plotly_chart(fig_temp, use_container_width=True)
    with col2:
        fig_rain = px.line(df, x="date", y="daily_rainfall_mm", color="location",
                           title="Daily Rainfall (mm)")
        st.plotly_chart(fig_rain, use_container_width=True)

# ----------------------------------------------------
# 5. Correlation Heatmap
# ----------------------------------------------------
st.subheader(f"Correlation Heatmap - {data_choice} Data")

# We'll pick relevant numeric columns
if data_choice == "Monthly":
    subset_cols = ["temp_avg", "humidity", "monthly_rainfall_mm", "malaria_risk", "dengue_risk"]
else:
    subset_cols = ["temp_avg", "humidity", "daily_rainfall_mm", "malaria_risk", "dengue_risk"]

corr_df = df[subset_cols].corr()
fig_corr = px.imshow(corr_df, text_auto=True, aspect="auto",
                     title="Correlation Matrix of Weather & Risk")
st.plotly_chart(fig_corr, use_container_width=True)

# ----------------------------------------------------
# 6. Interactive Map (Folium)
# ----------------------------------------------------
st.subheader("Interactive Map")
st.markdown(
    """
    **Note**: We only have 3 locations. Each marker popup shows some aggregated
    stats for the displayed data range.
    """
)

# Create a base map centered roughly in Tanzania
m = folium.Map(location=[-6.0, 35.0], zoom_start=6)

# We'll show monthly or daily aggregates in the popups
if data_choice == "Monthly":
    # For each location, let's gather monthly avg for the current df
    # Then we can show a simple summary in the popup
    for loc in selected_locations:
        loc_info = LOCATIONS[loc]
        loc_df = df[df["location"] == loc]
        
        if loc_df.empty:
            continue
        # Basic stats: average risk, average rainfall, etc
        avg_risk = loc_df[risk_col].mean()
        avg_temp = loc_df["temp_avg"].mean()
        avg_rain = loc_df["monthly_rainfall_mm"].mean()
        
        # Build popup HTML
        popup_html = f"""
        <b>{loc}</b><br/>
        Disease: {disease_choice}<br/>
        Avg Risk (in selection): {avg_risk:.2f}<br/>
        Avg Temp (°C): {avg_temp:.2f}<br/>
        Avg Rainfall (mm): {avg_rain:.2f}<br/>
        """
        folium.Marker(
            location=[loc_info["lat"], loc_info["lon"]],
            popup=popup_html,
            tooltip=f"{loc} ({disease_choice})"
        ).add_to(m)
else:
    # Daily data
    for loc in selected_locations:
        loc_info = LOCATIONS[loc]
        loc_df = df[df["location"] == loc]
        
        if loc_df.empty:
            continue
        avg_risk = loc_df[risk_col].mean()
        avg_temp = loc_df["temp_avg"].mean()
        avg_rain = loc_df["daily_rainfall_mm"].mean()
        
        popup_html = f"""
        <b>{loc}</b><br/>
        Disease: {disease_choice}<br/>
        Avg Risk (in selection): {avg_risk:.2f}<br/>
        Avg Temp (°C): {avg_temp:.2f}<br/>
        Avg Rain (mm/day): {avg_rain:.2f}<br/>
        """
        folium.Marker(
            location=[loc_info["lat"], loc_info["lon"]],
            popup=popup_html,
            tooltip=f"{loc} ({disease_choice})"
        ).add_to(m)

# Render Folium map in Streamlit
st_data = st_folium(m, width=700, height=500)