from __future__ import  annotations

from typing import  Callable, Iterable, Optional, Tuple, Union
import  matplotlib.pyplot   as plt
import  numpy               as np

FilterFn    =   Callable[[object], object]

def plot_xy(df, x_col, y_col, label:Optional[str] = None, ax=None,
            xlim:       Optional[Tuple[float, float]]   = None,
            ylim:       Optional[Tuple[float, float]]   = None,
            marker:     str = "o",
            linestyle:  str = "-",):
    
    """
    Gráfica y vs x para un DataFrame.
    - ax: para superponer curvas
    - xlim/ylim: zoom visual (No filtrado datos)
    """

    if ax   is  None:
        _, ax =   plt.subplots(figsize=(7,5))
    
    ax.plot(df[x_col], df[y_col], marker=marker, linestyle=linestyle, label=label)
    ax.set_xlabel(x_col)
    ax.set_ylabel(y_col)
    ax.grid(True)

    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)

    if  label:
        ax.legend()
    
    return  ax

def comparar_rondas(data: dict, experimento: int, x_col: str, y_col: str,
                    rondas: Iterable[int]   =   (1,2),
                    filter_fn: Optional[FilterFn] = None,
                    xlim: Optional[Tuple[float, float]] = None,
                    ylim: Optional[Tuple[float, float]] = None,
                    title: Optional[str] = None,):
    """
    Compara el mismo experimento entre rondas:
    data[(ronda, experimento)] -> df

    filter_fn: funcion opcional que recibe df y devuelve df filtrado.
        Ejemplo: filter_fn = lambda df: apply_filter(df, row_start=10, row_end=40)
    """
    fig, ax = plt.subplots(figsize=(7,5))

    for ronda in rondas:
        df = data[(ronda, experimento)]
        if filter_fn is not None:
            df = filter_fn(df)

        plot_xy(df, x_col, y_col, label=f"Ronda {ronda}", ax=ax,
                xlim=xlim, ylim=ylim)
    
    ax.set_title(title or f"Experimento {experimento} - Comparación de Rondas")
    plt.show()
    return fig, ax

def comparar_experimentos(data: dict, ronda: int, x_col: str, y_col: str, 
                          experimentos: Iterable[int] = (1,2,3,4),
                          show: bool = False,
                          filter_fn: Optional[FilterFn] = None,
                          xlim: Optional[Tuple[float, float]] = None,
                          ylim: Optional[Tuple[float, float]] = None,
                          title: Optional[str] = None, 
):
    """
    Comparar varios experimentos dentro de una misma ronda.
    """

    fig, ax = plt.subplots(figsize=(7, 5))

    for exp in experimentos:
        df = data[(ronda, exp)]
        if filter_fn is not None:
            df  = filter_fn(df)
        
        plot_xy(df, x_col, y_col, label=f"Experimento {exp}", ax=ax, xlim=xlim, ylim=ylim)
    if title == None:
        ax.set_title(f"Ronda {ronda} - Comparación de Experimentos")
    elif title == "No":
        ax.set_title("")
    else:
        ax.set_title(title)


    if show is True:
        plt.show()
    return fig, ax

def plot_con_ajuste(df, x_col: str, y_col: str, label_datos: str = "Datos",
                    label_ajuste: str = "Ajuste lineal",
                    show: bool = True 
):
    """
    Grafica scatter + ajuste lineal (y = mx + b). Devuelve (m, b)
    """

    x = np.asarray(df[x_col].values, dtype=float)
    y = np.asarray(df[y_col].values, dtype=float)

    m, b = np.polyfit(x, y, 1)

    fig, ax = plt.subplots(figsize=(7, 5))
    ax.scatter(x, y, label=label_datos)
    ax.plot(x, m * x + b, label=f"{label_ajuste} (m={m:.4f})")

    ax.set_xlabel(x_col)
    ax.set_ylabel(y_col)
    ax.grid(True)
    ax.legend()

    if show:
        plt.show()
    
    return m, b, fig, ax

def plot_3D(df, x_col: str, y_col: str, z_col: str, 
            title: Optional[str] = None,):
    """
    Gráfica 3D: (x, y, z) como dispersión. 
    """

    from mpl_toolkits.mplot3d import Axes3D #noqa: F401

    fig = plt.figure(figsize=(8, 6))
    ax  = fig.add_subplot(111, projection="3d")

    ax.scatter(df[x_col], df[y_col], df[z_col])
    ax.set_xlabel(x_col)
    ax.set_ylabel(y_col)
    ax.set_zlabel(z_col)
    ax.set_title(title or f"3D: {x_col} vs {y_col} vs {z_col}")

    plt.show()
    return fig, ax

def plot_color_map(df, x_col: str, y_col: str, z_col: str,
                   title: Optional[str] = None, ):
    
    """
    Scatter 2D con tercera variable como color 
    """
    fig, ax = plt.subplots(figsize=(7, 5))
    sc = ax.scatter(df[x_col], df[y_col], c=df[z_col], cmap="viridis")

    cbar = plt.colorbar(sc, ax=ax)
    cbar.set_label(z_col)

    ax.set_xlabel(x_col)
    ax.set_ylabel(y_col)
    ax.grid(True)
    ax.set_title(title or f"{y_col} vs {x_col} (color={z_col})")

    plt.show()
    return fig, ax

def plot_dual_axis(df, x_col: str, y1_col: str, y2_col: str, ax=None,
                   label1: Optional[str] = None,
                   label2: Optional[str] = None,
                   title:  Optional[str] = None,
                   show: bool = True):
    
    """
    Gráfica con doble eje Y:
     y1 vs x (izquierda) y y2 vs x (derecha)
    """
    if ax is None:
        fig, ax1 = plt.subplots(figsize=(7, 5))
    else:
        ax1 =   ax
        fig =   ax1.figure

    ax2 = ax1.twinx()

    ax1.plot(df[x_col], df[y1_col], marker="o", linestyle="-", color="b",label= label1 or y1_col)
    ax2.plot(df[x_col], df[y2_col], marker="s", linestyle="--", color="r",label= label2 or y2_col)

    ax1.set_xlabel(x_col)
    ax1.set_ylabel(y1_col)
    ax2.set_ylabel(y2_col)


    ax1.grid(True)
    ax1.set_title(title or f"{y1_col} y {y2_col} vs {x_col}")

    # Leyendas Combinadas 
    lines, labels = [], []
    for a in (ax1, ax2):
        l, lab = a.get_legend_handles_labels()
        lines += l
        labels += lab
    ax1.legend(lines, labels, loc="upper left")
    
    if show is not True:
        plt.show()
    return fig, ax1, ax2

def marcar_zonas(ax, limites, labels=None, colors=None):
    """
    Marca zonas en una gráfica con líneas verticales.

    limites: lista de voltajes donde cambian las zonas
    lables:  nombre de cada zona 
    colors:  colores opcionales
    """

    if colors is None:
        colors  =   ["gray", "blue", "green", "red"]

    for i, v in enumerate(limites):

        ax.axvline(v, linestyle="--", color="black", alpha=0.6)

        if labels and i < len(labels):
            ax.text(
                v, 
                ax.get_ylim()[1]*0.65,
                labels[i],
                rotation=90,
                verticalalignment="top",
                horizontalalignment="right",
                fontsize=12
            )
def sombrear_zonas(ax, limites, colores=None, alpha=0.15):
    """
    Sombrea las zonas del experimento 
    """

    if colores is None:
        colores =   ["gray", "blue", "green", "red"]
    
    for i in range(len(limites)-1):
        ax.axvspan(
            limites[i],
            limites[i+1],
            color=colores[i],
            alpha=alpha
        )

