import folium
import geopandas as gpd
import plotly.express as px
from branca.colormap import LinearColormap
from folium.plugins import HeatMap
from streamlit_folium import folium_static

import streamlit as st
from streamlit_extras.stylable_container import stylable_container
from streamlit_extras.add_vertical_space import add_vertical_space


with stylable_container(
    key="banner",
    css_styles="""
    img {
        width: 1800px;
        height: 400px;
        overflow: hidden;
        position: relative;
        object-fit: cover;
        border-radius: 20px; /* Adiciona bordas arredondadas */
        mask-image: linear-gradient(to bottom, rgba(0, 0, 0, 1), rgba(0, 0, 0, 0));
        -webkit-mask-image: linear-gradient(to bottom, rgba(0, 0, 0, 1), rgba(0, 0, 0, 0)); /* For Safari */
    }
    """,
):
    st.image("mp.jpg")

st.title("Mapas da área")
add_vertical_space(5)
st.markdown(""" ### :world_map: **UBS Flamengo: (IBGE 2022)** """)
add_vertical_space(5)


@st.cache_data
def load_data():
    """
    A function that loads and reads geojson data for UBS Flamengo and converts it to the specified coordinate reference system.
    """
    return gpd.read_file("flamengo_ibge2022.geojson").to_crs(epsg=4326)

gdf = load_data()
LATITUDE = -19.971591804
LONGITUDE = -44.057912815
lat = -19.96214
long = -44.05603



total_pop = gdf["POP"].sum()
col1, col2, col3 = st.columns([1, 1, 5])

with col1:
    # st.write(f"###### População Total: {total_pop:,}")
    map_type = st.selectbox("Tipo de mapa", ["População", "Densidade", "Heatmap"])

with col2:
    # st.write(f"###### Número de Setores Censitários: {len(gdf)}")
    base_map = st.selectbox("Mapa base", ["Cartodb Positron", "OpenStreetMap"])

with col3:
    total_pop = gdf["POP"].sum()
    st.write(
        f"### 👪 População Total: {total_pop:,} habitantes, dados do Censo 2022 IBGE"
    )
    st.write(f"### 🗺️ Número de Setores Censitários: {len(gdf)}")

add_vertical_space(5)
add_vertical_space(5)

col1, col2 = st.columns(2)
with col1:
    m = folium.Map(
        location=[LATITUDE, LONGITUDE], tiles=base_map, zoom_start=15
    )
    dns_p = "Densidade pop. (hab/km²) UBS Flamengo - IBGE 2022"
    if map_type in ["População", "Densidade"]:
        if map_type == "População":
            column = "POP"
            caption = "Pop. residente UBS Flamengo IBGE 2022"
        else:
            gdf["DENSIDADE"] = gdf["POP"] / gdf["AREA_KM2"]
            column = "DENSIDADE"
            caption = dns_p
        colorscale = px.colors.sequential.Viridis
        colormap = LinearColormap(
            colors=colorscale,
            vmin=gdf[column].min(),
            vmax=gdf[column].max(),
            caption=caption,
        )
        folium.GeoJson(
            gdf,
            style_function=lambda feature: {
                "fillColor": colormap(feature["properties"][column]),
                "color": "black",
                "weight": 1,
                "fillOpacity": 0.7,
            },
            highlight_function=lambda feature: {
                "fillColor": "#ffaf00",
                "color": "green",
                "weight": 3,
                "fillOpacity": 0.9,
            },
            tooltip=folium.features.GeoJsonTooltip(
                fields=["CD_SETOR", column, "AREA_KM2"],
                aliases=[
                    "Setor Censitário:",
                    f"{caption}:",
                    "Área (km²):",
                ],
                style=(
                    "background-color: white; color: #333333;"
                    "font-family: calibri; font-size: 12px;"
                    "padding: 10px;"
                ),
            ),
        ).add_to(m)
        colormap.add_to(m)

    elif map_type == "Heatmap":
        heat_data = [
            [
                row["geometry"].centroid.y,
                row["geometry"].centroid.x,
                row["POP"],
            ]
            for idx, row in gdf.iterrows()
        ]
        HeatMap(heat_data).add_to(m)
    folium.Marker(
        [lat, long],
        popup="UBS Flamengo",
        tooltip="UBS Flamengo",
        icon=folium.Icon(color="red", icon="info-sign"),
    ).add_to(m)

    STYLE_STATEMENT = (
        "<style>.leaflet-control-layers"
        "{ position: fixed; top: 10px; left: 50px; } </style>"
    )
    m.get_root().html.add_child(folium.Element(STYLE_STATEMENT))
    folium_static(m)

with col2:
    fig = px.bar(
        gdf,
        x="CD_SETOR",
        y="POP",
        title="Distribuição da População por Setor Censitário",
        color="POP",
        color_continuous_scale=px.colors.sequential.Viridis,
    )
    st.plotly_chart(fig)

    age_columns = [
        "POP_0A4",
        "POP_5A9",
        "POP_10A14",
        "POP_15A19",
        "POP_20A24",
        "POP_25A29",
        "POP_30A34",
        "POP_35A39",
        "POP_40A44",
        "POP_45A49",
        "POP_50A54",
        "POP_55A59",
        "POP_60A64",
        "POP_65A69",
        "POP_70A74",
        "POP_75A79",
        "POP_80A84",
        "POP_85A89",
        "POP_90A94",
        "POP_95A99",
        "POP_100OUMAIS",
    ]



add_vertical_space(10)


st.write('----')