"""
tsPellencKit.py
Utilities for time-series visualisation of multi-machine multi-sensor data.
"""

import math
import logging
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime
from typing import Optional

logger = logging.getLogger(__name__)
# logging.basicConfig(level=logging.DEBUG)

def find_column_type(type_target: type, df: pd.DataFrame) -> list[str]:
    """
    Return column names whose first non-NaN value matches type_target.

    Example:
        str_cols = find_column_type(str, df_pivoted)
    """
    matches = [
        col for col in df.columns
        if not df[col].dropna().empty
        and type(df[col].dropna().iloc[0]) is type_target
    ]
    logger.info("find_column_type(%s) → %d column(s) found", type_target.__name__, len(matches))
    return matches


def get_y_min_max(series: pd.Series, margin: float = 0.2) -> tuple[float, float]:
    """
    Compute a (y_min, y_max) range for a numeric series.

    Example:
        y_min, y_max = get_y_min_max(clean_series)
    """
    clean = series.dropna()
    if clean.empty:
        return 0.0, 1.0
    lo, hi = clean.min(), clean.max()
    if lo == hi == 0:
        return 0.0, 1.0
    return lo - abs(lo) * margin, hi + abs(hi) * margin



def vis_metric(
    df: pd.DataFrame,
    metric_name: str,
    *,
    n_cols: int = 4,
    show_nan: bool = True,
    adaptive_x: bool = True,
    years_filter: Optional[list[int]] = None,
    period_filter: Optional[list[str]] = None,
    ref_year: int = 2000,
) -> plt.Figure:
    """
    Plot a metric across all machines, overlaying multiple years for a specific metric.

    Args:
        df:            MultiIndex DataFrame with levels (SerialNumber, time).
        metric_name:   Column to visualise.
        n_cols:        Subplots per row.
        show_nan:      Mark NaN timestamps at y_min with ×.
        adaptive_x:    Fit x-axis to data; if False, show full year.
        years_filter:  Restrict to these years only (None = all).
        period_filter: Keep only ['MM-DD', 'MM-DD'] date range.
        ref_year:      Common calendar year used for overlay.

    Example:
        fig = vis_metric(df_pivoted, "PORTEUR_VITESSE_AVANCEMENT",
                         show_nan=False, years_filter=[2025],
                         period_filter=["06-01", "08-01"])
        fig.savefig("speed.png", dpi=150)
    """
    unique_sn = df.index.get_level_values("SerialNumber").unique()
    n_rows = math.ceil(len(unique_sn) / n_cols)

    all_years = sorted(df.index.get_level_values("time").year.unique())
    if years_filter is not None:
        all_years = [y for y in all_years if y in years_filter]

    cmap = plt.get_cmap("tab10")
    year_color = {year: cmap(i) for i, year in enumerate(all_years)}

    fig, axes = plt.subplots(n_rows, n_cols, figsize=(24, 2 * len(unique_sn)), squeeze=False)

    for i, sn in enumerate(unique_sn):
        row, col = divmod(i, n_cols)
        series = df[metric_name].loc[sn]

        if period_filter is not None:
            mmdd = series.index.strftime("%m-%d")
            series = series[(mmdd >= period_filter[0]) & (mmdd <= period_filter[1])]

        ax = axes[row, col]
        ax.set_title(sn)
        y_min, y_max = get_y_min_max(series)

        for year in all_years:
            year_series = series[series.index.year == year]
            ref_idx = year_series.index - pd.DateOffset(years=year - ref_year)
            color = year_color[year]
            ax.plot(ref_idx, year_series, "o-", color=color, label=str(year))

            if show_nan:
                nan_times = series[series.isna()].index.map(lambda t: t.replace(year=ref_year))
                ax.scatter(nan_times, [y_min] * len(nan_times),
                           marker="x", color=color, s=36, label=f"NaN {year}")

        ax.set_ylim(y_min, y_max)
        ax.tick_params(axis="x", rotation=45)
        ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M"))
        ax.set_ylabel("Value")
        ax.legend(fontsize=7)

        if not adaptive_x:
            ax.set_xlim(datetime(ref_year, 1, 1), datetime(ref_year, 12, 31))

    for ax in axes.flat[len(unique_sn):]:
        fig.delaxes(ax)

    fig.suptitle(metric_name, fontsize=16)
    plt.tight_layout(rect=[0, 0, 1, 0.98])
    return fig