import streamlit as st
import geopandas as gpd
import pandas as pd
import os
import zipfile
import tempfile
from keplergl import KeplerGl
from shapely.geometry import Polygon, MultiPolygon
import pydeck as pdk

def main():
    polygons_zip_path = 'polygons-20240817T212638Z-001.zip'

    
    valid_polygon_gdfs = []

   
    with tempfile.TemporaryDirectory() as tmpdirname:
        with zipfile.ZipFile(polygons_zip_path, 'r') as zip_ref:
            zip_ref.extractall(tmpdirname)
        
        polygons_dir = os.path.join(tmpdirname, 'polygons')
        geojson_files = [f for f in os.listdir(polygons_dir) if f.endswith('.geojson')]

        geojson_files = geojson_files[:7]

        for file in geojson_files:
            try:
                gdf = gpd.read_file(os.path.join(polygons_dir, file)).to_crs(epsg=4326)

                if gdf.empty:
                    st.warning(f"Skipping empty file: {file}")
                    continue

                gdf = gdf[gdf.geometry.apply(lambda geom: isinstance(geom, (Polygon, MultiPolygon)))]

                if gdf.empty:
                    st.warning(f"No valid geometries found in file: {file}")
                    continue

                valid_polygon_gdfs.append(gdf)
            
            except Exception as e:
                st.error(f"Error processing file {file}: {str(e)}")
                continue

    
    if valid_polygon_gdfs:
        tab1, tab2 = st.tabs(["Editable Point based map", "Pydeck Heatmap"])

        combined_polygon_gdf = gpd.GeoDataFrame(pd.concat(valid_polygon_gdfs, ignore_index=True))

        with tab1:
            kepler_data = {
                "polygons": combined_polygon_gdf
            }

            kepler_map = KeplerGl(height=600)

            kepler_map.add_data(data=kepler_data["polygons"], name="Polygons")

            config = {
            "version": "v1",
            "config": {
                "visState": {
                    "layers": [
                        {
                            "id": "polygon-layer",
                            "type": "geojson",
                            "config": {
                                "dataId": "Polygons",
                                "label": "Polygons",
                                "color": [255, 0, 0],
                                "highlightColor": [252, 242, 26, 255],
                                "columns": {
                                    "geojson": "geometry"
                                },
                                "isVisible": True,
                                "visConfig": {
                                    "opacity": 0.8,
                                    "thickness": 1.5,
                                    "strokeColor": [255, 0, 0],
                                    "colorRange": {
                                        "name": "Custom",
                                        "type": "sequential",
                                        "category": "Custom",
                                        "colors": [
                                            "#000004", "#3b0f70", "#8c2981", "#de4968", "#fe9f6d", "#fcfdbf"  # Magma
                                        ]
                                    },
                                    "filled": True,
                                    "stroked": True,
                                    "enable3d": False,
                                    "wireframe": False
                                },
                                "visualChannels": {
                                    "colorField": {
                                        "name": "class",  
                                        "type": "real"
                                    },
                                    "colorScale": "quantile"
                                }
                            }
                        }
                    ]
                },
                "mapStyle": {
                    "styleType": "satellite",  
                    "topLayerGroups": {},     
                    "visibleLayerGroups": {
                        "label": True,
                        "road": True,
                        "border": False,
                        "building": True,
                        "water": True,
                        "land": True,
                        "3d building": False
                    },
                    "threeDBuildingColor": [9.665468314072013, 17.18305478057247, 31.1442867897876],
                    "mapStyles": {}
                },
                "mapState": {
                    "bearing": 0,
                    "dragRotate": True,
                    "latitude": combined_polygon_gdf.geometry.centroid.y.mean(),
                    "longitude": combined_polygon_gdf.geometry.centroid.x.mean(),
                    "pitch": 0,
                    "zoom": 10
                },
                "interactionConfig": {
                    "brush": {
                        "enabled": True
                    },
                    "tooltip": {
                        "enabled": True,
                        "fieldsToShow": {
                            "Polygons": [
                                {
                                    "name": "class", 
                                    "format": None
                                }
                            ]
                        }
                    }
                },
                "animationConfig": {
                    "enabled": True
                }
            }
        }

            kepler_map = KeplerGl(height=600, config=config)
            kepler_map.add_data(data=kepler_data["polygons"], name="Polygons")
            kepler_map.save_to_html(file_name="kepler_map.html")
            st.components.v1.html(open("kepler_map.html", 'r').read(), height=600)
            st.sidebar.title("Summary Stats")
            st.sidebar.write(combined_polygon_gdf.describe())

        with tab2:
            combined_polygon_gdf['centroid'] = combined_polygon_gdf.geometry.centroid
            combined_polygon_gdf['longitude'] = combined_polygon_gdf.centroid.x
            combined_polygon_gdf['latitude'] = combined_polygon_gdf.centroid.y
            data = combined_polygon_gdf[['longitude', 'latitude', 'class']]
            heatmap_layer = pdk.Layer(
                "HeatmapLayer",
                data=data,
                get_position='[longitude, latitude]',
                get_weight="class",
                radius_pixels=50,
            )

            view_state = pdk.ViewState(
                latitude=combined_polygon_gdf.geometry.centroid.y.mean(),
                longitude=combined_polygon_gdf.geometry.centroid.x.mean(),
                zoom=10,
                bearing=0,
                pitch=0
            )

            r = pdk.Deck(
                layers=[heatmap_layer],
                initial_view_state=view_state,
                map_style='mapbox://styles/mapbox/satellite-v9'
            )

            st.pydeck_chart(r)
    else:
        st.error("No valid GeoDataFrames could be combined.")

if __name__ == "__main__":
    main()