import pandas as pd
import plotly.express as px
import streamlit as st
import geopandas as gpd

# --- Configuration and Constants ---
MISSING_VALUE_PLACEHOLDER = -1

# --- Helper for Color Scale ---
def _generate_custom_colorscale(min_val, max_val, red_gradient=False):
    if red_gradient:
        reversed_rdylgn_colors =  px.colors.sequential.Reds
        return reversed_rdylgn_colors
    else:
        custom_colorscale = [[0.0, "lightgrey"]]
        reversed_rdylgn_colors = px.colors.diverging.RdYlGn[::-1]

        normalized_min_actual = (min_val - MISSING_VALUE_PLACEHOLDER) / (
            max_val - MISSING_VALUE_PLACEHOLDER
        )
        normalized_max_actual = (max_val - MISSING_VALUE_PLACEHOLDER) / (
            max_val - MISSING_VALUE_PLACEHOLDER
        )

        num_steps = len(reversed_rdylgn_colors)
        for i, color in enumerate(reversed_rdylgn_colors):
            normalized_point = normalized_min_actual + (
                normalized_max_actual - normalized_min_actual
            ) * (i / (num_steps - 1))
            if normalized_point > 0.0:
                custom_colorscale.append([normalized_point, color])
        return sorted(custom_colorscale, key=lambda x: x[0])

def display_choropleth_map_country(
    df: pd.DataFrame,
    geojson_data: dict,
    metric_name: str = "prixm2moyen",
    metric_description: str = "Average Price per Square Meter",
    title: str = "Average Price per Square Meter in France (2015-2024)",
    height: int = 1000,
    width: int = 1400,
    red_gradient: bool = False,
):
    df_agg = (
        df
        .groupby(["code_departement", "annee"])[metric_name]
        .mean()
        .reset_index()
    )
    metric_min = df_agg[metric_name].min()
    metric_max = df_agg[metric_name].max()

    custom_colorscale = _generate_custom_colorscale(metric_min, metric_max, red_gradient)

    fig = px.choropleth_map(
        df_agg,
        geojson=geojson_data,
        locations="code_departement",
        featureidkey="properties.code",
        color=metric_name,
        labels={metric_name: metric_description},
        range_color=[metric_min, metric_max],
        color_continuous_scale=custom_colorscale,
        center={"lat": 46.603354, "lon": 1.888334},
        zoom=5,
        opacity=0.75,
        hover_name="code_departement",
        hover_data={metric_name: ":.1f", "annee": True},
        title=title,
        height=height,
        width=width,
        animation_frame="annee",
        animation_group="code_departement",
    )
    fig.update_traces(marker_line_width=0)
    if fig.layout.updatemenus:
        try:
            fig.layout.updatemenus[0].buttons[0].args[1]["frame"]["duration"] = 1500
            fig.layout.updatemenus[0].buttons[0].args[1]["transition"]["duration"] = 500
        except IndexError:
            st.warning("Could not set animation speed.")
    else:
        st.warning("No animation updatemenus found.")
    return fig


def display_choropleth_map_for_department(
    df: pd.DataFrame,
    department_code: str,
    geojson_data: dict,
    min_global_prixm2moyen: float,
    max_global_prixm2moyen: float,
    metric_name: str = "prixm2moyen",
    metric_description: str = "Average Price per Square Meter",
    title: str = "Average Price per Square Meter",
    height_graph: int = 1000,
    width_graph: int = 1400,
    red_gradient: bool = False,
):
    df_filtered = df[df["code_departement"] == department_code].copy()

    filtered_features = [
        feature
        for feature in geojson_data["features"]
        if feature["properties"]["code"].startswith(department_code)
    ]
    filtered_geojson = {"type": "FeatureCollection", "features": filtered_features}

    custom_colorscale = _generate_custom_colorscale(
        min_global_prixm2moyen, max_global_prixm2moyen, red_gradient
    )

    center_lat, center_lon, zoom_level = 46.603354, 1.888334, 6
    if filtered_features:
        department_gdf = gpd.GeoDataFrame.from_features(filtered_geojson["features"])
        if not department_gdf.empty:
            dissolved_department = department_gdf.dissolve()
            department_centroid = dissolved_department.geometry.centroid.iloc[0]
            center_lat, center_lon = department_centroid.y, department_centroid.x
            minx, miny, maxx, maxy = dissolved_department.total_bounds
            width, height = maxx - minx, maxy - miny
            if width < 0.2 and height < 0.2:
                zoom_level = 10
            elif width < 0.5 and height < 0.5:
                zoom_level = 9
            else:
                zoom_level = 8

    fig = px.choropleth_map(
        df_filtered,
        geojson=filtered_geojson,
        locations="code_commune_insee",
        featureidkey="properties.code",
        color=metric_name,
        labels={metric_name: metric_description},
        range_color=[min_global_prixm2moyen, max_global_prixm2moyen],
        # colorscale in red gradient
        color_continuous_scale=custom_colorscale,
        center={"lat": center_lat, "lon": center_lon},
        zoom=zoom_level,
        opacity=0.75,
        hover_name="code_commune_insee",
        hover_data={metric_name: ":.0f", "annee": True},
        title=title,
        height=height_graph,
        width=width_graph,
        animation_frame="annee",
        animation_group="code_commune_insee",
    )
    fig.update_traces(marker_line_width=0)
    if fig.layout.updatemenus:
        try:
            fig.layout.updatemenus[0].buttons[0].args[1]["frame"]["duration"] = 750
            fig.layout.updatemenus[0].buttons[0].args[1]["transition"]["duration"] = 250
        except IndexError:
            st.warning("Could not set animation speed.")
    else:
        st.warning("No animation updatemenus found.")
    return fig