from __future__ import annotations

import numpy as np
import polars as pl
from rgram.base import BaseUtils
from typing import Callable, Literal, Sequence, Optional, Union, Any


class Regressogram(BaseUtils):
    """
    Regressogram

    Binned regression estimator for one or more features and targets.
    Predicts using binned aggregation with optional confidence intervals.

    Parameters
    ----------
    binning : {'dist', 'width', 'none', 'int'}, default='dist'
        Binning strategy.
    agg : callable, default=mean
        Aggregation function for y in each bin.
    ci : tuple of callables, optional
        Tuple of lower/upper confidence interval functions.
    n_bins : int, optional
        Number of bins for 'dist' binning. If None, automatically calculated
        using Freedman-Diaconis rule. Ignored for other binning strategies.

    Methods
    -------
    fit(data, x, y, hue=None, keys=None)
        Learn bin parameters from training data.
    predict(x, return_ci=False)
        Predict binned regression values at new x points.
        Returns array or tuple with optional confidence intervals.
    fit_predict(data, x, y, hue=None, keys=None, return_ci=False)
        Fit and predict on training x values.
    """

    def __init__(
        self,
        *,
        binning: Literal["dist", "width", "none", "int"] = "dist",
        agg: Callable[[pl.Expr], pl.Expr] = lambda x: x.mean(),
        ci: Optional[
            tuple[Callable[[pl.Expr], pl.Expr], Callable[[pl.Expr], pl.Expr]]
        ] = (
            lambda x: x.mean() - x.std(),
            lambda x: x.mean() + x.std(),
        ),
        n_bins: Optional[int] = None,
    ):
        """
        Construct a Regressogram instance.

        Parameters
        ----------
        binning : {'dist', 'width', 'none', 'int'}, default='dist'
            Binning strategy.
        agg : callable, default=mean
            Aggregation function for y in each bin.
        ci : tuple of callables, optional
            Tuple of lower/upper confidence interval functions.
        n_bins : int, optional
            Number of bins for 'dist' binning. If None, automatically calculated
            using Freedman-Diaconis rule. Ignored for other binning strategies.
        """
        self.binning = binning
        self.agg = agg
        self.ci = ci
        self.n_bins = n_bins

    def _learn_bin_params(self, data: pl.LazyFrame) -> None:
        x_min, x_max, q25, q75, n = (
            data.select(
                pl.col("x_val").min().alias("min"),
                pl.col("x_val").max().alias("max"),
                pl.col("x_val").quantile(0.25).alias("q25"),
                pl.col("x_val").quantile(0.75).alias("q75"),
                pl.len(),
            )
            .collect()
            .row(0)
        )

        self._x_min = x_min
        self._x_max = x_max

        if self.binning in ("int", "none"):
            self._min_bin = x_min
            self._max_bin = x_max

        if self.binning in ("width", "dist"):
            self._bin_width = 2 * (q75 - q25) / (n ** (1 / 3))

            # Handle edge case where bin_width is 0 (all x values identical)
            if self._bin_width == 0:
                self._bin_width = 1.0

            if self.binning == "dist":
                # Use user-specified n_bins or calculate from Freedman-Diaconis rule
                if self.n_bins is not None:
                    n_bins = max(1, self.n_bins)

                else:
                    n_bins = max(1, int((x_max - x_min) // self._bin_width))

                self._n_bins = n_bins

                self._min_bin = 0
                self._max_bin = n_bins - 1

            else:
                self._min_bin = 0
                self._max_bin = int((x_max - x_min) // self._bin_width)

    def _predict_bin_expr(self) -> pl.Expr:
        """
        Returns a Polars expression for binning x values.

        Returns
        -------
        pl.Expr
            The binning expression.
        """
        if self.binning == "width":
            bin_id = ((pl.col("x_val") - self._x_min) // self._bin_width).cast(int)

        elif self.binning == "dist":
            bin_id = (
                pl.col("x_val")
                .qcut(
                    quantiles=self._n_bins,
                    allow_duplicates=True,
                    include_breaks=True,
                )
                .struct.field("breakpoint")
                .rank(method="dense")
                .cast(int)
            )

        elif self.binning == "int":
            bin_id = pl.col("x_val").cast(int)

        elif self.binning == "none":
            return pl.col("x_val")

        else:
            raise ValueError(f"Unknown binning type: {self.binning}")

        # clip to edge bins
        return bin_id.clip(self._min_bin, self._max_bin)

    def fit(
        self,
        x: Union[str, Sequence[Any]],
        y: Union[str, Sequence[Any]],
        data: Union[pl.DataFrame, pl.LazyFrame, None] = None,
        hue: Optional[Union[str, Sequence[str]]] = None,
        keys: Optional[Union[str, Sequence[Any]]] = None,
    ) -> "Regressogram":
        """
        Learn bin parameters from training data.

        Supports flexible input similar to seaborn (e.g., kdeplot):
        - Provide DataFrame + column names (recommended for production)
        - Provide raw arrays/Series (convenient for exploration)

        Parameters
        ----------
        x : str or array-like
            Feature(s) to bin. Column name(s) if `data` provided, else array-like.
        y : str or array-like
            Target(s). Column name(s) if `data` provided, else array-like.
        data : pl.DataFrame, pl.LazyFrame, or None, optional
            Input data. If provided, x/y/hue/keys are treated as column names.
            If None, x/y/keys are treated as array-like values.
        hue : str or sequence of str, optional
            Optional grouping variable(s).
        keys : str or array-like, optional
            Additional grouping column(s).

        Returns
        -------
        self : object
            Fitted estimator.
        """
        data_lf, x_cols, y_cols, keys_cols = self._prepare_data(
            data=data, x=x, y=y, keys=keys
        )

        x_list = self._to_list(x_cols) or [x_cols]
        y_list = self._to_list(y_cols) or [y_cols]
        hue_list = self._to_list(hue) or [] if data is not None else []
        keys_list = self._to_list(keys_cols) or [keys_cols] if keys_cols else []

        idx_cols = (y_list or []) + (keys_list or []) + (hue_list or [])
        self.over_cols = ["x_var", "y_var"] + (hue_list or [])

        data = (
            data_lf.select(x_list + idx_cols)
            .unpivot(
                on=x_list, index=idx_cols, variable_name="x_var", value_name="x_val"
            )
            .unpivot(
                on=y_list,
                index=["x_val", "x_var"] + hue_list + keys_list,
                variable_name="y_var",
                value_name="y_val",
            )
            .filter(pl.col("x_var") != pl.col("y_var"))
            .with_columns([pl.col("y_val").cast(float)])
        )

        # learn bin parameters and assign bins
        self._learn_bin_params(data)
        data = data.with_columns(
            [self._predict_bin_expr().over(self.over_cols).alias("rgram_bin")]
        )

        # aggregate y values per bin
        data = data.with_columns(
            [
                self.agg(pl.col("y_val"))
                .over(self.over_cols + ["rgram_bin"])
                .alias("y_pred_rgram")
            ]
        )

        # Compute and store confidence intervals for each bin (if configured)
        select_cols = ["rgram_bin", "y_pred_rgram"]
        if self.ci:
            ci_cols = ["y_pred_rgram_lci", "y_pred_rgram_uci"]
            ci_exprs = [
                ci_calc(pl.col("y_val").fill_null(pl.col("y_val").mean()))
                .over(self.over_cols + ["rgram_bin"])
                .alias(alias)
                for ci_calc, alias in zip(self.ci, ci_cols)
            ]
            data = data.with_columns(ci_exprs)
            select_cols.extend(ci_cols)

        self._bin_to_y = data.select(select_cols).unique().collect()
        self._training_data = data.collect()

        return self

    def _get_full_predictions(self) -> pl.DataFrame:
        """Internal method to compute full predictions with all columns."""
        if not hasattr(self, "_training_data"):
            raise RuntimeError("You must call fit() before getting predictions.")

        data = self._training_data.lazy()

        if self.ci:
            ci_cols = ["y_pred_rgram_lci", "y_pred_rgram_uci"]
            ci_exprs = [
                ci_calc(pl.col("y_val").fill_null(pl.col("y_val").mean()))
                .over(self.over_cols + ["rgram_bin"])
                .alias(alias)
                for ci_calc, alias in zip(self.ci, ci_cols)
            ]

            data = data.with_columns(ci_exprs)

        cols_to_drop = []

        schema = data.collect_schema()
        cols_to_drop.extend([i for i in schema if i in ["x_var", "y_var"]])

        return data.sort(by=["x_val"]).drop(cols_to_drop).collect()

    def transform(self) -> pl.LazyFrame:
        """
        Return the full binned regression results (for backward compatibility).

        Returns full data with bin assignments and predictions.

        Returns
        -------
        pl.LazyFrame
            Full results including all columns. Call `.collect()` to materialize.
        """
        return self._get_full_predictions().lazy()

    def predict(
        self, x: Union[Sequence[float], pl.Series], return_ci: bool = False
    ) -> Union[np.ndarray, tuple]:
        """
        Predict binned regression values at new x points.

        Parameters
        ----------
        x : array-like or pl.Series
            New x values at which to predict.
        return_ci : bool, default=False
            If True, return confidence intervals along with predictions.
            Returns tuple (y_pred, y_ci_low, y_ci_high).
            If False, returns just the predictions array.

        Returns
        -------
        np.ndarray or tuple
            If return_ci=False: numpy array of predicted values (same length as x)
            If return_ci=True: tuple of (y_pred, y_ci_low, y_ci_high)
                y_pred: numpy array of predictions
                y_ci_low: numpy array of lower CI or None if ci not configured
                y_ci_high: numpy array of upper CI or None if ci not configured
        """
        if not hasattr(self, "_bin_to_y"):
            raise RuntimeError("Call fit() before predict().")

        lf = pl.DataFrame({"x_val": x}).lazy()

        lf = lf.with_columns(self._predict_bin_expr().alias("rgram_bin"))

        lf = lf.join(
            self._bin_to_y.lazy(),
            on="rgram_bin",
            how="left",
        )

        result_df = lf.select("y_pred_rgram").collect()
        y_pred = result_df["y_pred_rgram"].to_numpy()

        if not return_ci:
            return y_pred

        # Retrieve pre-computed CIs from _bin_to_y (stored at fit time)
        y_ci_low = None
        y_ci_high = None

        if self.ci:
            ci_cols = ["y_pred_rgram_lci", "y_pred_rgram_uci"]
            if ci_cols[0] in self._bin_to_y.columns:
                ci_result = lf.select(
                    ["y_pred_rgram_lci", "y_pred_rgram_uci"]
                ).collect()
                y_ci_low = ci_result["y_pred_rgram_lci"].to_numpy()
                y_ci_high = ci_result["y_pred_rgram_uci"].to_numpy()

        return y_pred, y_ci_low, y_ci_high

    def fit_predict(
        self,
        x: Union[str, Sequence[Any]],
        y: Union[str, Sequence[Any]],
        data: Union[pl.DataFrame, pl.LazyFrame, None] = None,
        hue: Optional[Union[str, Sequence[str]]] = None,
        keys: Optional[Union[str, Sequence[Any]]] = None,
        return_ci: bool = False,
    ) -> Union[np.ndarray, tuple]:
        """
        Fit and predict on training x values in one call.

        Parameters
        ----------
        x : str or array-like
            Feature(s) to bin. Column name(s) if `data` provided, else array-like.
        y : str or array-like
            Target(s). Column name(s) if `data` provided, else array-like.
        data : pl.DataFrame, pl.LazyFrame, or None, optional
            Input data. If provided, x/y/hue/keys are column names.
            If None, x/y/keys are array-like.
        hue : str or sequence of str, optional
            Optional grouping variable(s).
        keys : str or array-like, optional
            Additional grouping column(s).
        return_ci : bool, default=False
            If True, return confidence intervals along with predictions.

        Returns
        -------
        np.ndarray or tuple
            If return_ci=False: array of predictions at training x values
            If return_ci=True: tuple of (y_pred, y_ci_low, y_ci_high)
        """
        self.fit(data=data, x=x, y=y, hue=hue, keys=keys)

        # Get unique x values from training data for prediction
        x_train = (
            self._training_data.select("x_val")
            .unique()
            .sort("x_val")["x_val"]
            .to_numpy()
        )

        return self.predict(x_train, return_ci=return_ci)
