# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias

"""High-level convenience functions for Aurora-GLM.

This module provides unified helper functions that work with any Aurora result:

- summary(result): Formatted model summary
- plot(result): Diagnostic plots
- compare(*results): Model comparison

These are also available at the package level:
>>> import aurora
>>> aurora.summary(result)
>>> aurora.plot(result)
>>> aurora.compare(model1, model2, model3)

Examples
--------
>>> from aurora import fit_glm, summary, plot, compare
>>>
>>> # Fit models
>>> result1 = fit_glm(X, y, family=Gaussian())
>>> result2 = fit_glm(X, y, family=Poisson())
>>>
>>> # Get summaries
>>> print(summary(result1))
>>>
>>> # Plot diagnostics
>>> plot(result1, kind='diagnostics')
>>>
>>> # Compare models
>>> compare(result1, result2, criterion='aic')
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Literal, Sequence

import numpy as np

if TYPE_CHECKING:
    pass


# =============================================================================
# Summary Function
# =============================================================================


def summary(
    result: Any,
    *,
    style: Literal["default", "brief", "detailed"] = "default",
    print_output: bool = True,
) -> str:
    """Generate a formatted summary of a fitted model result.

    This is a unified interface for summarizing any Aurora model result.
    It automatically dispatches to the appropriate summary method based
    on the result type.

    Parameters
    ----------
    result : ModelResult
        A fitted model result from fit_glm, fit_gam, fit_gamm, etc.
    style : {'default', 'brief', 'detailed'}, default='default'
        Summary style:
        - 'default': Standard coefficient table with fit statistics
        - 'brief': One-line summary with key metrics
        - 'detailed': Full output including diagnostics
    print_output : bool, default=True
        If True, print the summary. If False, only return the string.

    Returns
    -------
    summary_str : str
        Formatted summary string.

    Examples
    --------
    >>> from aurora import fit_glm, summary, Gaussian
    >>> result = fit_glm(X, y, family=Gaussian())
    >>> print(summary(result))
    ======================================================================
                          Linear Model Results
    ======================================================================
    No. Observations:        100    R-squared:          0.8542
    ...

    >>> # Brief summary
    >>> summary(result, style='brief')
    'GLM(n=100, converged=True, R²=0.854)'

    >>> # Detailed with diagnostics
    >>> summary(result, style='detailed')

    Notes
    -----
    The function works with any result object that has a `.summary()` method.
    For custom result classes, implement `.summary()` to integrate with this.

    See Also
    --------
    plot : Diagnostic plots for model results
    compare : Compare multiple models
    """
    # Check for summary method
    if not hasattr(result, "summary"):
        raise TypeError(
            f"Cannot summarize object of type {type(result).__name__}. "
            "Expected a model result with a .summary() method."
        )

    # Generate summary based on style
    if style == "default":
        summary_str = result.summary()
    elif style == "brief":
        summary_str = _brief_summary(result)
    elif style == "detailed":
        summary_str = _detailed_summary(result)
    else:
        raise ValueError(
            f"Unknown style '{style}'. Use 'default', 'brief', or 'detailed'."
        )

    if print_output:
        print(summary_str)

    return summary_str


def _brief_summary(result: Any) -> str:
    """Generate a one-line brief summary."""
    parts = []

    # Class name
    class_name = type(result).__name__
    if class_name.endswith("Result"):
        class_name = class_name[:-6]  # Remove 'Result' suffix
    parts.append(class_name)

    # Observations
    if hasattr(result, "n_obs_"):
        parts.append(f"n={result.n_obs_}")
    elif hasattr(result, "n_observations"):
        parts.append(f"n={result.n_observations}")

    # Convergence
    if hasattr(result, "converged_"):
        parts.append(f"converged={result.converged_}")
    elif hasattr(result, "converged"):
        parts.append(f"converged={result.converged}")

    # R² for linear models
    if hasattr(result, "r_squared"):
        r2 = result.r_squared
        if not np.isnan(r2):
            parts.append(f"R²={r2:.3f}")

    # Log-likelihood for mixed models
    if hasattr(result, "log_likelihood_"):
        parts.append(f"LL={result.log_likelihood_:.2f}")

    # AIC if available
    if hasattr(result, "aic"):
        aic = result.aic
        if not np.isnan(aic):
            parts.append(f"AIC={aic:.1f}")

    return f"{parts[0]}({', '.join(parts[1:])})"


def _detailed_summary(result: Any) -> str:
    """Generate detailed summary with diagnostics."""
    lines = []

    # Base summary
    lines.append(result.summary())
    lines.append("")

    # Diagnostics section
    lines.append("=" * 70)
    lines.append("Diagnostics")
    lines.append("=" * 70)

    # Residual statistics
    if hasattr(result, "residuals"):
        try:
            resid = result.residuals
            lines.append("Residuals:")
            lines.append(f"  Min:    {np.min(resid):>12.4f}")
            lines.append(f"  1Q:     {np.percentile(resid, 25):>12.4f}")
            lines.append(f"  Median: {np.median(resid):>12.4f}")
            lines.append(f"  3Q:     {np.percentile(resid, 75):>12.4f}")
            lines.append(f"  Max:    {np.max(resid):>12.4f}")
            lines.append("")
        except Exception:
            pass

    # Fit statistics
    if hasattr(result, "r_squared"):
        r2 = result.r_squared
        if not np.isnan(r2):
            lines.append(f"R-squared:          {r2:>12.4f}")

    if hasattr(result, "adj_r_squared"):
        adj_r2 = result.adj_r_squared
        if not np.isnan(adj_r2):
            lines.append(f"Adjusted R-squared: {adj_r2:>12.4f}")

    if hasattr(result, "aic"):
        aic = getattr(result, "aic", np.nan)
        if not np.isnan(aic):
            lines.append(f"AIC:                {aic:>12.2f}")

    if hasattr(result, "bic"):
        bic = getattr(result, "bic", np.nan)
        if not np.isnan(bic):
            lines.append(f"BIC:                {bic:>12.2f}")

    lines.append("=" * 70)

    return "\n".join(lines)


# =============================================================================
# Plot Function
# =============================================================================


def plot(
    result: Any,
    *,
    kind: Literal["diagnostics", "residuals", "qq", "smooth", "all"] = "diagnostics",
    ax: Any = None,
    **kwargs,
) -> Any:
    """Create diagnostic plots for a fitted model result.

    This is a unified interface for plotting any Aurora model result.
    It automatically creates appropriate plots based on the model type.

    Parameters
    ----------
    result : ModelResult
        A fitted model result from fit_glm, fit_gam, fit_gamm, etc.
    kind : {'diagnostics', 'residuals', 'qq', 'smooth', 'all'}, default='diagnostics'
        Type of plot:
        - 'diagnostics': Standard 4-panel diagnostic plot
        - 'residuals': Residuals vs fitted values
        - 'qq': Q-Q plot of residuals
        - 'smooth': Smooth function plots (GAM/GAMM only)
        - 'all': All available plots
    ax : matplotlib.axes.Axes, optional
        Axes to plot on. If None, creates new figure.
    **kwargs
        Additional arguments passed to the underlying plot function.

    Returns
    -------
    fig : matplotlib.figure.Figure or list of Figures
        The created figure(s).

    Examples
    --------
    >>> from aurora import fit_glm, plot, Gaussian
    >>> result = fit_glm(X, y, family=Gaussian())
    >>>
    >>> # Standard diagnostics
    >>> plot(result)
    >>>
    >>> # Just residuals
    >>> plot(result, kind='residuals')
    >>>
    >>> # For GAM, plot smooth terms
    >>> gam_result = fit_gam(X, y)
    >>> plot(gam_result, kind='smooth')

    See Also
    --------
    summary : Formatted model summary
    compare : Compare multiple models
    plot_diagnostics : Low-level diagnostic plots
    plot_smooth : Low-level smooth term plots
    """
    # Import visualization functions lazily
    from .visualization import (
        plot_diagnostics_panel,
        plot_all_smooths,
    )

    if kind == "diagnostics":
        return plot_diagnostics_panel(result, ax=ax, **kwargs)

    elif kind == "residuals":
        return _plot_residuals(result, ax=ax, **kwargs)

    elif kind == "qq":
        return _plot_qq(result, ax=ax, **kwargs)

    elif kind == "smooth":
        # Check if GAM result
        if hasattr(result, "smooth_terms") or hasattr(result, "smooth_info"):
            return plot_all_smooths(result, **kwargs)
        else:
            raise ValueError(
                "kind='smooth' requires a GAM or GAMM result with smooth terms."
            )

    elif kind == "all":
        figs = []
        figs.append(plot_diagnostics_panel(result, **kwargs))
        if hasattr(result, "smooth_terms") or hasattr(result, "smooth_info"):
            figs.append(plot_all_smooths(result, **kwargs))
        return figs

    else:
        raise ValueError(
            f"Unknown plot kind '{kind}'. "
            "Use 'diagnostics', 'residuals', 'qq', 'smooth', or 'all'."
        )


def _plot_residuals(result: Any, ax: Any = None, **kwargs) -> Any:
    """Plot residuals vs fitted values."""
    import matplotlib.pyplot as plt

    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 6))
    else:
        fig = ax.get_figure()

    fitted = result.fitted_values
    residuals = result.residuals

    ax.scatter(fitted, residuals, alpha=0.6, **kwargs)
    ax.axhline(y=0, color="red", linestyle="--", linewidth=1)
    ax.set_xlabel("Fitted values")
    ax.set_ylabel("Residuals")
    ax.set_title("Residuals vs Fitted")

    return fig


def _plot_qq(result: Any, ax: Any = None, **kwargs) -> Any:
    """Create Q-Q plot of residuals."""
    import matplotlib.pyplot as plt
    from scipy import stats

    if ax is None:
        fig, ax = plt.subplots(figsize=(6, 6))
    else:
        fig = ax.get_figure()

    residuals = result.residuals

    # Standardize residuals
    std_resid = (residuals - np.mean(residuals)) / np.std(residuals)

    stats.probplot(std_resid, dist="norm", plot=ax)
    ax.set_title("Normal Q-Q Plot")

    return fig


# =============================================================================
# Compare Function
# =============================================================================


def compare(
    *results: Any,
    criterion: Literal["aic", "bic", "loglik", "all"] = "aic",
    names: Sequence[str] | None = None,
    print_output: bool = True,
) -> dict[str, Any]:
    """Compare multiple fitted models.

    This function provides a unified interface for comparing different
    model specifications using information criteria or likelihood ratio tests.

    Parameters
    ----------
    *results : ModelResult
        Two or more fitted model results to compare.
    criterion : {'aic', 'bic', 'loglik', 'all'}, default='aic'
        Criterion for comparison:
        - 'aic': Akaike Information Criterion
        - 'bic': Bayesian Information Criterion
        - 'loglik': Log-likelihood
        - 'all': All available criteria
    names : sequence of str, optional
        Names for the models. If None, uses 'Model 1', 'Model 2', etc.
    print_output : bool, default=True
        If True, print the comparison table.

    Returns
    -------
    comparison : dict
        Dictionary with comparison results:
        - 'summary': Formatted comparison table (str)
        - 'criteria': DataFrame or dict with criterion values
        - 'best': Index of best model
        - 'delta': Differences from best model

    Examples
    --------
    >>> from aurora import fit_glm, compare, Gaussian, Poisson
    >>>
    >>> # Fit different models
    >>> m1 = fit_glm(X, y, family=Gaussian())
    >>> m2 = fit_glm(X, y, family=Poisson())
    >>> m3 = fit_glm(X[:, :2], y, family=Gaussian())  # Fewer predictors
    >>>
    >>> # Compare models
    >>> compare(m1, m2, m3, criterion='aic')
    ======================================================================
                            Model Comparison (AIC)
    ======================================================================
           Model        AIC      ΔAIC   Rank
    ----------------------------------------------------------------------
       Model 1     234.56      0.00      1   ← best
       Model 2     245.12     10.56      2
       Model 3     267.89     33.33      3
    ======================================================================

    >>> # With custom names
    >>> compare(m1, m2, names=['Linear', 'Poisson'], criterion='bic')

    Notes
    -----
    - Lower AIC/BIC indicates better fit (with parsimony)
    - For nested models, likelihood ratio tests can also be performed
    - Models should be fit on the same data for meaningful comparison

    See Also
    --------
    summary : Individual model summaries
    """
    if len(results) < 2:
        raise ValueError("compare() requires at least 2 models to compare.")

    # Generate names if not provided
    if names is None:
        names = [f"Model {i + 1}" for i in range(len(results))]
    elif len(names) != len(results):
        raise ValueError(
            f"Number of names ({len(names)}) must match "
            f"number of models ({len(results)})."
        )

    # Extract criteria values
    metrics = _extract_comparison_metrics(results, criterion)

    # Create comparison summary
    summary_str = _format_comparison_table(names, metrics, criterion)

    # Find best model
    primary_criterion = criterion if criterion != "all" else "aic"
    if primary_criterion in metrics:
        values = metrics[primary_criterion]
        if primary_criterion == "loglik":
            best_idx = int(np.argmax(values))  # Higher is better
        else:
            best_idx = int(np.argmin(values))  # Lower is better

        delta = np.array(values) - values[best_idx]
    else:
        best_idx = 0
        delta = np.zeros(len(results))

    if print_output:
        print(summary_str)

    return {
        "summary": summary_str,
        "criteria": metrics,
        "names": list(names),
        "best": best_idx,
        "best_model": names[best_idx],
        "delta": delta,
    }


def _extract_comparison_metrics(
    results: tuple[Any, ...], criterion: str
) -> dict[str, list[float]]:
    """Extract comparison metrics from results."""
    metrics: dict[str, list[float]] = {}

    # Always try to get n_obs and n_params
    n_obs_list = []
    n_params_list = []

    for result in results:
        # Number of observations
        if hasattr(result, "n_obs_"):
            n_obs_list.append(result.n_obs_)
        elif hasattr(result, "n_observations"):
            n_obs_list.append(result.n_observations)
        else:
            n_obs_list.append(np.nan)

        # Number of parameters
        if hasattr(result, "df_model"):
            n_params_list.append(result.df_model + 1)
        elif hasattr(result, "coefficients"):
            n_params_list.append(len(result.coefficients))
        else:
            n_params_list.append(np.nan)

    metrics["n_obs"] = n_obs_list
    metrics["n_params"] = n_params_list

    # Extract requested criteria
    criteria_to_extract = (
        ["aic", "bic", "loglik"] if criterion == "all" else [criterion]
    )

    for crit in criteria_to_extract:
        values = []
        for result in results:
            if crit == "aic":
                values.append(_get_aic(result))
            elif crit == "bic":
                values.append(_get_bic(result))
            elif crit == "loglik":
                values.append(_get_loglik(result))
            else:
                values.append(np.nan)
        metrics[crit] = values

    return metrics


def _get_aic(result: Any) -> float:
    """Extract or compute AIC."""
    if hasattr(result, "aic"):
        return result.aic

    # Try to compute from log-likelihood
    ll = _get_loglik(result)
    if np.isnan(ll):
        return np.nan

    if hasattr(result, "df_model"):
        k = result.df_model + 1
    elif hasattr(result, "coefficients"):
        k = len(result.coefficients)
    else:
        return np.nan

    return -2 * ll + 2 * k


def _get_bic(result: Any) -> float:
    """Extract or compute BIC."""
    if hasattr(result, "bic"):
        return result.bic

    # Try to compute from log-likelihood
    ll = _get_loglik(result)
    if np.isnan(ll):
        return np.nan

    if hasattr(result, "n_obs_"):
        n = result.n_obs_
    elif hasattr(result, "n_observations"):
        n = result.n_observations
    else:
        return np.nan

    if hasattr(result, "df_model"):
        k = result.df_model + 1
    elif hasattr(result, "coefficients"):
        k = len(result.coefficients)
    else:
        return np.nan

    return -2 * ll + k * np.log(n)


def _get_loglik(result: Any) -> float:
    """Extract log-likelihood."""
    if hasattr(result, "log_likelihood_"):
        return result.log_likelihood_
    if hasattr(result, "loglik"):
        return result.loglik
    if hasattr(result, "llf"):
        return result.llf
    return np.nan


def _format_comparison_table(
    names: Sequence[str], metrics: dict[str, list[float]], criterion: str
) -> str:
    """Format comparison table as string."""
    lines = []
    sep = "=" * 70

    lines.append(sep)
    lines.append(f"{'Model Comparison':^70}")
    lines.append(sep)

    # Determine which criteria to show
    criteria_to_show = ["aic", "bic", "loglik"] if criterion == "all" else [criterion]

    # Header
    header_parts = [f"{'Model':>15}"]
    for crit in criteria_to_show:
        header_parts.append(f"{crit.upper():>12}")
    header_parts.append(f"{'Rank':>8}")
    lines.append("  ".join(header_parts))
    lines.append("-" * 70)

    # Determine ranking (by first criterion)
    primary = criteria_to_show[0]
    if primary in metrics:
        values = np.array(metrics[primary])
        if primary == "loglik":
            order = np.argsort(-values)  # Higher is better
        else:
            order = np.argsort(values)  # Lower is better
        ranks = np.empty_like(order)
        ranks[order] = np.arange(1, len(order) + 1)
    else:
        ranks = np.arange(1, len(names) + 1)

    # Data rows
    for i, name in enumerate(names):
        row_parts = [f"{name:>15}"]
        for crit in criteria_to_show:
            if crit in metrics:
                val = metrics[crit][i]
                if np.isnan(val):
                    row_parts.append(f"{'N/A':>12}")
                else:
                    row_parts.append(f"{val:>12.2f}")
            else:
                row_parts.append(f"{'N/A':>12}")

        rank_str = str(int(ranks[i]))
        if ranks[i] == 1:
            rank_str += " ← best"
        row_parts.append(f"{rank_str:>8}")

        lines.append("  ".join(row_parts))

    lines.append(sep)

    return "\n".join(lines)


__all__ = [
    "summary",
    "plot",
    "compare",
]
