import pandas as pd  # stress hydrique and rendement, besoin en eau
import plotly.graph_objects as go
from typing import List
import plotly.express as px
from fuzzywuzzy import process
from data_pipelines.historical_weather_data import (
    download_historical_weather_data,
    aggregate_hourly_weather_data,
)

def concatenate_historic_forecast(historic, forecast, cols_to_keep):
    historic["period"] = "historique"
    forecast["period"] = "futur"
    historic = historic[cols_to_keep]
    forecast = forecast[cols_to_keep]
    full_data = pd.concat([historic, forecast])
    return full_data


def visualize_climate(
    moderate: pd.DataFrame,
    historic: pd.DataFrame,
    x_axis="year",
    column: str = "Precipitation (mm)",
    cols_to_keep: List[str] = [
        "Precipitation (mm)",
        "Near Surface Air Temperature (°C)",
        "Surface Downwelling Shortwave Radiation (W/m²)",
        "year",
        "period",
    ],
):
    concatenated_df = concatenate_historic_forecast(historic, moderate, cols_to_keep)
    concatenated_df = concatenated_df.sort_values(by=x_axis)  # Ensure order

    fig = go.Figure()
    # colors = {"historique": "blue", "forecast": "purple"}

    for condition_value in concatenated_df["period"].unique():
        segment = concatenated_df[concatenated_df["period"] == condition_value]
        fig.add_trace(
            go.Scatter(
                x=segment[x_axis],
                y=segment[column],
                mode="lines",
                name=f"{condition_value}",
                line=dict(color="blue" if condition_value == "futur" else "purple"),
            )
        )

    interpolation_df = concatenated_df[concatenated_df[x_axis] > 2023]
    fig.add_trace(
        go.Scatter(
            x=interpolation_df[x_axis],
            y=interpolation_df[column].interpolate(),
            mode="lines",
            line=dict(color="blue"),
            showlegend=False,
        ),
    )

    fig.update_layout(
        title=f"Historique et Forecast pour {column}",
        xaxis_title="Date",
        yaxis_title=column,
    )

    return fig


def aggregate_yearly(df, col_to_agg):
    df[col_to_agg] = df.groupby("year")[col_to_agg].transform("mean")
    return df


def generate_plots(
    moderate: pd.DataFrame,
    historic: pd.DataFrame,
    x_axes: List[str],
    cols_to_plot: List[str],
):
    plots = []
    for i, col in enumerate(cols_to_plot):
        plots.append(visualize_climate(moderate, historic, x_axes[i], col))
    return plots

if __name__ == "__main__":
    cols_to_keep = [
        "Precipitation (mm)",
        "Near Surface Air Temperature (°C)",
        "Surface Downwelling Shortwave Radiation (W/m²)",
        "year",
        "period",
    ]
    cols_to_plot = [
        "Precipitation (mm)",
        "Near Surface Air Temperature (°C)",
        "Surface Downwelling Shortwave Radiation (W/m²)",
    ]
    x_axes = ["year"] * len(cols_to_plot)

    latitude = 47
    longitude = 5
    start_year = 2000
    end_year = 2025
    cols_to_keep = [
        "Precipitation (mm)",
        "Near Surface Air Temperature (°C)",
        "Surface Downwelling Shortwave Radiation (W/m²)",
        "year",
        "period",
    ]
    cols_to_plot = [
        "Precipitation (mm)",
        "Near Surface Air Temperature (°C)",
        "Surface Downwelling Shortwave Radiation (W/m²)",
    ]
    x_axes = ["year"] * len(cols_to_plot)

    df = download_historical_weather_data(latitude, longitude, start_year, end_year)
    historic = aggregate_hourly_weather_data(df)
    historic = historic.reset_index()
    historic = historic.rename(
        columns={
            "precipitation": "Precipitation (mm)",
            "air_temperature_mean": "Near Surface Air Temperature (°C)",
            "irradiance": "Surface Downwelling Shortwave Radiation (W/m²)",
            "index": "time",
        }
    )
    historic["time"] = pd.to_datetime(historic["time"])
    historic = historic.sort_values("time")
    historic = historic[historic["time"] < "2025-01-01"]
    climate_data_path = "data/data_climate_test/final_climate_data.csv" # will use a similar func to download_historical_weather_data
    moderate = pd.read_csv(climate_data_path)
    moderate["time"] = pd.to_datetime(moderate["time"])
    moderate = moderate.sort_values("time")
    moderate["year"] = moderate["time"].dt.year
    historic["year"] = historic["time"].dt.year
    moderate["Precipitation (mm)"] = moderate["Precipitation (mm)"] * 3600
    for col in cols_to_plot:
        moderate = aggregate_yearly(moderate, col)
        historic = aggregate_yearly(historic, col)


    plots = generate_plots(moderate, historic, x_axes, cols_to_plot) # List of 3 plots to show