import streamlit as st
from streamlit_folium import st_folium
import folium
from folium.plugins import Draw
import pandas as pd
from shapely.geometry import Polygon, Point
import numpy as np

st.set_page_config(layout="wide", page_title="Multiplex Coop Map Filter")

st.title("🗺️ Multiplex Coop Housing Filter")
st.write("Draw a polygon on the map to filter the data points within it. Use the form below to apply additional filters based on property attributes.")

# --- 1. Create a Sample DataFrame with more attributes ---
@st.cache_data
def load_sample_data():
    num_points = 100
    data = {
        'id': range(1, num_points + 1),
        'name': [f'Property {i}' for i in range(1, num_points + 1)],
        'latitude': np.random.uniform(34.03, 34.07, num_points),
        'longitude': np.random.uniform(-118.28, -118.21, num_points),
        'zn_type': np.random.choice(['Residential (0)', 'Residential Apartment (101)', 'Commercial Residential (6)'], num_points),
        'zn_area': np.random.randint(200, 2000, num_points), # Lot Area in Sq Metres
        'fsi_total': np.round(np.random.uniform(0.5, 3.0, num_points), 2), # Floor Space Index
        'prcnt_cver': np.random.randint(20, 70, num_points), # Building Percent Coverage
        'height_metres': np.round(np.random.uniform(5, 30, num_points), 1), # Height in Metres
        'stories': np.random.randint(2, 10, num_points) # Number of Stories
    }
    df = pd.DataFrame(data)
    return df

df = load_sample_data()

# Initialize filtered_df with the full dataframe
filtered_df = df.copy()

# --- 2. Initialize the Folium Map with Drawing Tools ---
# Center the map around the sample data (e.g., Los Angeles area)
m = folium.Map(location=[df['latitude'].mean(), df['longitude'].mean()], zoom_start=12)

# Add drawing tools
draw = Draw(
    export=True,
    filename="drawn_polygon.geojson",
    position="topleft",
    draw_options={
        "polyline": False,
        "rectangle": False,
        "circlemarker": False,
        "circle": False,
        "marker": False,
        "polygon": {
            "allowIntersection": False,
            "drawError": {
                "color": "#e0115f",
                "message": "Oups!",
            },
            "shapeOptions": {
                "color": "#ef233c",
                "fillOpacity": 0.5,
            },
        },
    },
    edit_options={"edit": False, "remove": True},
)
m.add_child(draw)

# Add all data points to the map initially (these will be updated after filtering)
for idx, row in df.iterrows():
    folium.CircleMarker(
        location=[row['latitude'], row['longitude']],
        radius=5,
        color='blue',
        fill=True,
        fill_color='blue',
        fill_opacity=0.7,
        tooltip=(
            f"ID: {row['id']}<br>Name: {row['name']}<br>Zoning: {row['zn_type']}<br>"
            f"Area: {row['zn_area']} m²<br>FSI: {row['fsi_total']}<br>"
            f"Coverage: {row['prcnt_cver']}%<br>Height: {row['height_metres']}m<br>"
            f"Stories: {row['stories']}"
        )
    ).add_to(m)

st.subheader("Draw a Polygon on the Map")
output = st_folium(m, width=1000, height=600, returned_objects=["all_draw_features"])

polygon_drawn = False
shapely_polygon = None
polygon_coords = None

if output and output["all_draw_features"]:
    polygons = [
        feature["geometry"]["coordinates"]
        for feature in output["all_draw_features"]
        if feature["geometry"]["type"] == "Polygon"
    ]

    if polygons:
        polygon_coords = polygons[-1][0] # Get the last drawn polygon's coordinates
        # Shapely Polygon expects (lon, lat) tuples, Folium gives (lat, lon)
        shapely_polygon = Polygon([(lon, lat) for lat, lon in polygon_coords])
        polygon_drawn = True

        # Apply spatial filter
        filtered_df = df[
            df.apply(
                lambda row: shapely_polygon.contains(Point(row['longitude'], row['latitude'])),
                axis=1
            )
        ].copy() # Use .copy() to avoid SettingWithCopyWarning
        st.success(f"Initially filtered {len(filtered_df)} points within the drawn polygon.")
    else:
        st.info("Draw a polygon on the map to spatially filter points.")
else:
    st.info("Draw a polygon on the map to spatially filter points.")

# --- 3. Attribute Filtering Form ---
st.subheader("Filter Property Attributes")

with st.form("attribute_filters"):
    col1, col2 = st.columns(2)

    with col1:
        # Zoning Type
        all_zoning_types = ['All Resdidential Zoning (0, 101, 6)'] + sorted(df['zn_type'].unique().tolist())
        selected_zn_type = st.selectbox("Zoning Type", all_zoning_types, key="zn_type_select")

        # Lot Area in Sq Metres
        min_zn_area = st.number_input("Minimum Lot Area in Sq Metres", min_value=0, value=0, step=10, key="zn_area_input")

        # Floor Space Index (FSI)
        min_fsi_total = st.number_input("Minimum Floor Space Index (FSI)", min_value=0.0, value=0.0, step=0.1, format="%.2f", key="fsi_total_input")

    with col2:
        # Building Percent Coverage
        max_prcnt_cver = st.number_input("Maximum Building Percent Coverage (%)", min_value=0, value=100, step=1, key="prcnt_cver_input")

        # Height or Stories selection
        height_stories_option = st.radio(
            "Filter by",
            ("Height", "Stories"),
            index=0, # Default to Height
            key="height_stories_radio"
        )

        # Single input field for height/stories, label changes dynamically
        if height_stories_option == "Height":
            min_height_value = st.number_input("Minimum Height in Metres", min_value=0.0, value=0.0, step=0.1, format="%.1f", key="height_input")
        else: # Stories
            min_stories_value = st.number_input("Minimum Stories", min_value=0, value=0, step=1, key="stories_input")

    submitted = st.form_submit_button("Apply Attribute Filters")

    if submitted:
        # Apply attribute filters to the already spatially filtered_df
        if selected_zn_type != 'All Resdidential Zoning (0, 101, 6)':
            filtered_df = filtered_df[filtered_df['zn_type'] == selected_zn_type]

        if min_zn_area > 0:
            filtered_df = filtered_df[filtered_df['zn_area'] >= min_zn_area]

        if min_fsi_total > 0:
            filtered_df = filtered_df[filtered_df['fsi_total'] >= min_fsi_total]

        if max_prcnt_cver < 100: # Assuming 100% means no upper limit applied
            filtered_df = filtered_df[filtered_df['prcnt_cver'] <= max_prcnt_cver]

        if height_stories_option == "Height" and min_height_value > 0:
            filtered_df = filtered_df[filtered_df['height_metres'] >= min_height_value]
        elif height_stories_option == "Stories" and min_stories_value > 0:
            filtered_df = filtered_df[filtered_df['stories'] >= min_stories_value]

        st.success(f"Applied attribute filters. Total points after all filters: {len(filtered_df)}")
    else:
        # If form not submitted, the filtered_df remains as it was after spatial filtering
        st.info("Adjust filters and click 'Apply Attribute Filters'.")


# --- 4. Display Filtered Data on a New Map and as a Table ---
st.subheader("Filtered Data Points")

if not filtered_df.empty:
    # Create a new map to show only the filtered points
    # Adjust map center and zoom if filtered_df is very small or empty,
    # otherwise use the original map's center or the filtered_df's center.
    if len(filtered_df) > 0:
        filtered_map_center = [filtered_df['latitude'].mean(), filtered_df['longitude'].mean()]
        filtered_map_zoom = 14 if len(filtered_df) < 5 else 12
    else:
        filtered_map_center = [df['latitude'].mean(), df['longitude'].mean()]
        filtered_map_zoom = 12

    filtered_m = folium.Map(location=filtered_map_center, zoom_start=filtered_map_zoom)

    # Add the drawn polygon to the new map if it exists
    if polygon_drawn and polygon_coords:
        folium.Polygon(
            locations=polygon_coords, # Use original (lat,lon) for folium
            color="#ef233c",
            fill=True,
            fill_color="#ef233c",
            fill_opacity=0.5
        ).add_to(filtered_m)

    # Add filtered points to the new map
    for idx, row in filtered_df.iterrows():
        folium.CircleMarker(
            location=[row['latitude'], row['longitude']],
            radius=7,
            color='green',
            fill=True,
            fill_color='green',
            fill_opacity=0.8,
            tooltip=(
                f"ID: {row['id']}<br>Name: {row['name']}<br>Zoning: {row['zn_type']}<br>"
                f"Area: {row['zn_area']} m²<br>FSI: {row['fsi_total']}<br>"
                f"Coverage: {row['prcnt_cver']}%<br>Height: {row['height_metres']}m<br>"
                f"Stories: {row['stories']}"
            )
        ).add_to(filtered_m)

    st_folium(filtered_m, width=1000, height=500)

    st.subheader("Filtered Data Table")
    st.dataframe(filtered_df)

    # --- 5. Export Data Button ---
    csv = filtered_df.to_csv(index=False).encode('utf-8')
    st.download_button(
        label="Export Filtered Data to CSV",
        data=csv,
        file_name="multiplex_coop_filtered_data.csv",
        mime="text/csv",
    )

else:
    st.warning("No data points match the current filters. Try adjusting your criteria or drawing a different polygon.")

st.markdown("---")
st.markdown("This app demonstrates spatial filtering using a drawn polygon and attribute filtering based on the provided HTML structure.")