import streamlit as st
import folium
from folium.plugins import MarkerCluster
import pandas as pd
from streamlit_folium import folium_static

from srai.regionalizers import geocode_to_region_gdf

# Define the function to create a map with markers
def create_map(points, color, marker_cluster):
    points.apply(lambda row: folium.Marker(location=[row.y, row.x], icon=folium.Icon(color=color)).add_to(marker_cluster))

# Define the function to get the center of the region
def get_center(region):
    boundary = geocode_to_region_gdf(region)
    return boundary.centroid.iloc[0]

# Define the function to get the points of the region
def get_points(region, size):
    boundary = geocode_to_region_gdf(region)
    return boundary["geometry"].sample_points(size=size).explode()

# Set up the map
map_center = (52.34260, 20.04503)
poland_center = get_center("Poland")
my_map = folium.Map(location = (poland_center.y, poland_center.x), zoom_start = 5)

marker_cluster = MarkerCluster().add_to(my_map)

# Set up the sliders
poland_points_num = st.sidebar.slider('Number of points for Poland', 0, 500, 350)
riga_points_num = st.sidebar.slider('Number of points for Riga, Latvia', 0, 25, 5)
tallin_points_num = st.sidebar.slider('Number of points for Tallin, Estonia', 0, 25, 5)
pskov_points_num = st.sidebar.slider('Number of points for Pskov', 0, 25, 25)

poland_points = get_points("Poland", poland_points_num)
pskov_points = get_points("Pskov", pskov_points_num)
riga_points = get_points("Riga", riga_points_num)
tallin_points = get_points("Tallin", tallin_points_num)

# Create the map with markers
create_map(poland_points, 'blue', marker_cluster)
create_map(riga_points, 'blue', marker_cluster)
create_map(tallin_points, 'blue', marker_cluster)
create_map(pskov_points, 'red', marker_cluster)

# Display the map
st_data = folium_static(my_map)