import pandas as pd
import matplotlib.pyplot as plt


def plot_historical_comparison(
    df_new: pd.DataFrame, df_hist: pd.DataFrame, compound: str, ax=None
) -> tuple[plt.Figure, plt.Axes]:
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))

    dt_column = ("-", "datetime")

    for data_type, df in zip(["Historical", "New"], [df_hist, df_new]):
        if data_type == "Historical":
            color = "blue"
        else:
            color = "red"

        serie = df[(compound, "conc")]
        dt = df[dt_column]
        if ("-", "type") in df.columns:
            mask_air = df[("-", "type")] == "air"
            serie = serie[mask_air]
            dt = dt[mask_air]

        ax.scatter(dt, serie, label=data_type, color=color, alpha=0.5, s=4)
    ax.set_title(compound)
    ax.set_xlabel("Date")
    ax.set_ylabel("Concentration (ppt)")
    ax.legend()
    return fig, ax


def plot_yearly_data(
    df: pd.DataFrame, compound: str, ax=None
) -> tuple[plt.Figure, plt.Axes]:
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))

    dt_column = ("-", "datetime")
    serie = df[(compound, "conc")]
    dt = df[dt_column]
    if ("-", "type") in df.columns:
        mask_air = df[("-", "type")] == "air"
        serie = serie[mask_air]
        dt = dt[mask_air]

    years = dt.dt.year.unique()
    x = dt.dt.day_of_year + dt.dt.hour / 24.0
    for year in years:
        mask_year = dt.dt.year == year
        ax.scatter(x[mask_year], serie[mask_year], label=str(year), alpha=0.5, s=4)

    ax.set_title(compound)
    ax.set_xlabel("Time of Year")
    ax.set_ylabel("Concentration (ppt)")

    # Add ticks with the mounths
    month_starts = pd.date_range(start="2024-01-01", end="2025-01-01", freq="MS")
    month_days = month_starts.dayofyear
    month_labels = month_starts.strftime("%b")
    ax.set_xticks(month_days)
    ax.set_xticklabels(month_labels)
    ax.legend()
    return fig, ax


def plot_yearly_plotly(
    df: pd.DataFrame,
    compound: str,
    df_new: pd.DataFrame | None = None,
    opacity: float = 0.5,
    size: int = 6,
) -> "plotly.graph_objs._figure.Figure":
    """Plot yearly data using plotly."""
    import plotly.express as px
    import plotly.graph_objects as go

    dt_column = ("-", "datetime")
    serie = df[(compound, "conc")]
    dt = df[dt_column]
    if ("-", "type") in df.columns:
        mask_air = df[("-", "type")] == "air"
        serie = serie[mask_air]
        dt = dt[mask_air]
    if ("-", "type") in df_new.columns:
        mask_air_new = df_new[("-", "type")] == "air"
        df_new = df_new[mask_air_new]

    x = dt.dt.day_of_year + dt.dt.hour / 24.0
    df_to_plot = pd.DataFrame(
        {
            "conc": serie.values,
            "year": dt.dt.year.values,
        },
        index=x.values,
    )
    # Break down by year, to have year as columns and conc as values
    df_to_plot = df_to_plot.pivot_table(
        index=df_to_plot.index, columns="year", values="conc"
    )
    fig = go.Figure()

    hover_template = "Timestamp: %{text}<br>Conc: %{y:.2f} ppt"

    kwargs = {
        "mode": "markers",
        "opacity": opacity,
        "marker": dict(size=size),
        "hovertemplate": hover_template,
    }

    for year in df_to_plot.columns:
        fig.add_trace(
            go.Scatter(
                x=df_to_plot.index,
                y=df_to_plot[year],
                name=str(year),
                zorder=-year,
                text=dt[dt.dt.year == year].dt.strftime("%y%m%d.%H%M"),
                **kwargs,
            )
        )
    x_values = pd.date_range(start="2024-01-01", end="2024-12-31", freq="MS")

    dt_new = df_new[dt_column]
    fig.add_trace(
        go.Scatter(
            x=dt_new.dt.dayofyear + dt_new.dt.hour / 24.0,
            y=df_new[(compound, "conc")],
            name="New Data",
            text=dt_new.dt.strftime("%y%m%d.%H%M"),
            **kwargs,
        )
    )
    fig.update_layout(
        xaxis_title="Time of Year",
        yaxis_title=f"{compound} (ppt)",
        xaxis=dict(
            tickmode="array",
            tickvals=x_values.dayofyear,
            ticktext=x_values.strftime("%b"),
        ),
    )

    return fig
