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

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))

def get_center(region):
    boundary = geocode_to_region_gdf(region)
    return boundary.centroid.iloc[0]

def get_points(region, size):
    boundary = geocode_to_region_gdf(region)
    return boundary["geometry"].sample_points(size=size).explode()

poland_center = get_center("Poland")
my_map = folium.Map(location = (54.50272, 20.51721), zoom_start = 5)
st.set_page_config(layout="wide")

marker_cluster = MarkerCluster().add_to(my_map)

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, 50, 5)
tallin_points_num = st.sidebar.slider('Number of points for Tallin, Estonia', 0, 50, 5)
pskov_points_num = st.sidebar.slider('Number of points for Pskov', 0, 50, 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_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)

st_data = folium_static(my_map, width=1200, height=800)