from __future__ import annotations

from typing import Dict, Tuple, Optional, Union, List, TYPE_CHECKING

from jax import Array
import arviz as az
import numpyro
import xarray
import matplotlib.pyplot as plt
import pandas as pd

from catalax.mcmc.mcmc import MCMC
from catalax.mcmc.plotting import plot_corner, plot_ess, plot_mcse, plot_posterior

if TYPE_CHECKING:
    from catalax.model.model import Model
    from catalax.mcmc.mcmc import BayesianModel

try:
    import ipywidgets as widgets
    from IPython.display import display
    from IPython.core.getipython import get_ipython

    WIDGETS_AVAILABLE = True
except ImportError:
    WIDGETS_AVAILABLE = False


class HMCResults:
    """Results container for HMC sampling with integrated visualization methods.

    This class provides a comprehensive interface for analyzing and visualizing
    the results of Hamiltonian Monte Carlo (HMC) sampling. It wraps the MCMC
    object and provides convenient methods for diagnostics, plotting, and
    exporting results.

    The class integrates with ArviZ for advanced diagnostics and provides
    methods for generating various plots including corner plots, trace plots,
    posterior distributions, and credibility intervals.

    Attributes:
        mcmc: The MCMC object containing posterior samples
        bayesian_model: The Bayesian model used for sampling
        model: The original mechanistic model that was fit
    """

    def __init__(
        self,
        mcmc: MCMC,
        bayesian_model: BayesianModel,
        model: "Model",
    ):
        """Initialize HMC results.

        Args:
            mcmc: MCMC object containing samples and sampling metadata
            bayesian_model: Bayesian model used for sampling, containing
                prior definitions and likelihood specification
            model: Original mechanistic model that was fit to data
        """
        self.mcmc = mcmc
        self.bayesian_model = bayesian_model
        self.model = model

    def get_fitted_model(
        self,
        hdi_prob: float = 0.95,
        set_bounds: bool = False,
    ) -> "Model":
        """Get the fitted model with parameter estimates from posterior samples.

        Creates a new model instance with parameters set to posterior estimates.
        Parameter values are typically set to posterior means or medians, with
        uncertainty bounds optionally derived from highest density intervals.

        Args:
            hdi_prob: Probability mass for highest density interval calculation.
                Used to determine uncertainty bounds for parameters.
            set_bounds: Whether to set parameter bounds based on HDI. If True,
                parameter bounds will be updated to reflect posterior uncertainty.

        Returns:
            Model: New model instance with parameters updated from posterior samples.
                The returned model can be used for predictions with fitted parameters.
        """
        return self.model.from_samples(
            self.mcmc.get_samples(),
            hdi_prob=hdi_prob,
            set_bounds=set_bounds,
        )

    def get_samples(self) -> Dict[str, Array]:
        """Get posterior samples for all parameters.

        Returns the raw posterior samples as generated by the MCMC sampler.
        These samples represent the posterior distribution over model parameters
        and can be used for custom analysis or uncertainty quantification.

        Returns:
            Dict[str, Array]: Dictionary mapping parameter names to arrays of
                posterior samples. Each array has shape (num_samples, ...) where
                the parameter dimensions depend on the parameter structure.
        """
        return self.mcmc.get_samples()

    def print_summary(self):
        """Print MCMC summary statistics to console.

        Displays a tabular summary of the MCMC results including:
        - Parameter posterior means and standard deviations
        - Effective sample size (ESS) for each parameter
        - R-hat convergence diagnostics
        - Monte Carlo standard error estimates

        This provides a quick overview of sampling quality and parameter estimates.
        """
        self.mcmc.print_summary()

    def render_model(
        self,
        path: str,
        render_distributions: bool = False,
    ):
        """Render the Bayesian model structure as a graphical representation.

        Creates a visual representation of the model's probabilistic structure,
        showing the relationships between parameters, priors, and likelihood.
        Useful for understanding model structure and debugging.

        Args:
            path: File path where the rendered model graph will be saved.
                Should include appropriate file extension (e.g., '.pdf', '.png').
            render_distributions: Whether to include distribution shapes in the
                rendered graph. If True, adds visual representation of prior
                and likelihood distributions.

        Returns:
            The result of numpyro.render_model, typically None but may vary
            based on the rendering backend used.
        """
        return numpyro.render_model(
            self.bayesian_model,
            filename=path,
            render_distributions=render_distributions,
        )

    # Visualization methods from plotting.py
    def plot_corner(
        self,
        quantiles: Tuple[float, float, float] = (0.16, 0.5, 0.84),
        figsize: Optional[Tuple[float, float]] = None,
        backend: Optional[str] = None,
        show: bool = False,
        path: Optional[str] = None,
    ):
        """Plot corner plot of parameter correlations and marginal distributions.

        Creates a matrix plot showing pairwise parameter correlations in the
        lower triangle, marginal posterior distributions on the diagonal, and
        optionally scatter plots in the upper triangle. This is essential for
        understanding parameter dependencies and posterior structure.

        Args:
            quantiles: Three quantiles to display in marginal distributions,
                typically representing lower bound, median, and upper bound.
                Default (0.16, 0.5, 0.84) corresponds to 68% credible intervals.
            figsize: Figure size as (width, height) in inches. If None, uses default size.
            backend: Backend to use ('matplotlib' or 'bokeh'). If None, uses default.
            show: Whether to display the plot immediately using plt.show().
                If False, the figure is returned without displaying.
            path: Optional file path to save the plot. If provided, the figure
                will be saved to this location.

        Returns:
            matplotlib.figure.Figure or bokeh plot: The corner plot figure, or None if show=True.
                Can be used for further customization or saving.
        """

        f = plot_corner(self.mcmc, quantiles, figsize, backend)
        if path is not None:
            f.savefig(path)
        if show:
            plt.show()
            return None

        return plt.gcf()

    def plot_posterior(
        self,
        figsize: Optional[Tuple[float, float]] = None,
        backend: Optional[str] = None,
        show: bool = False,
        path: Optional[str] = None,
        **kwargs,
    ):
        """Plot posterior distributions for all parameters.

        Creates individual plots showing the posterior distribution for each
        model parameter. This provides a clear view of parameter uncertainty
        and can help identify multimodal distributions or convergence issues.

        Args:
            figsize: Figure size as (width, height) in inches. If None, uses default size.
            backend: Backend to use ('matplotlib' or 'bokeh'). If None, uses default.
            show: Whether to display the plot immediately using plt.show().
                If False, the figure is returned without displaying.
            path: Optional file path to save the plot. If provided, the figure
                will be saved to this location.
            **kwargs: Additional keyword arguments passed to arviz.plot_posterior.
                Common options include 'hdi_prob' for credible interval probability,
                'point_estimate' for central tendency measure, and 'ref_val' for
                reference values.

        Returns:
            matplotlib.figure.Figure or bokeh plot: The posterior plot figure, or None if show=True.
                Contains subplots for each parameter's posterior distribution.
        """

        f = plot_posterior(self.mcmc, self.model, figsize, backend, **kwargs)
        if path is not None:
            f.savefig(path)
        if show:
            plt.show()
            return None

        return plt.gcf()

    def plot_credibility_interval(
        self,
        initial_condition: Dict[str, float],
        time: Array,
        dt0: float = 0.1,
        figsize: Optional[Tuple[float, float]] = None,
        backend: Optional[str] = None,
        show: bool = False,
        path: Optional[str] = None,
    ):
        """Plot credibility intervals for model simulations.

        Generates model predictions using posterior parameter samples and plots
        credibility intervals around the mean trajectory. This shows prediction
        uncertainty arising from parameter uncertainty and is crucial for
        understanding model reliability.

        Args:
            initial_condition: Dictionary mapping species names to initial
                concentrations for the simulation. Must include all species
                defined in the model.
            time: Array of time points at which to evaluate the model.
                Should span the desired simulation duration.
            dt0: Time step size for numerical integration. Smaller values
                increase accuracy but require more computation.
            figsize: Figure size as (width, height) in inches. If None, uses default size.
            backend: Backend to use ('matplotlib' or 'bokeh'). If None, uses default.
                Note: Credibility interval plots currently only support matplotlib backend.
            show: Whether to display the plot immediately using plt.show().
                If False, the figure is returned without displaying.
            path: Optional file path to save the plot. If provided, the figure
                will be saved to this location.

        Returns:
            matplotlib.figure.Figure: The credibility interval plot figure,
                or None if show=True. Shows model trajectories with uncertainty bands.
        """
        from catalax.mcmc.plotting import plot_credibility_interval

        f = plot_credibility_interval(
            self.mcmc, self.model, initial_condition, time, dt0, figsize, backend
        )
        if path is not None:
            f.savefig(path)
        if show:
            plt.show()
            return None
        return plt.gcf()

    def plot_trace(
        self,
        figsize: Optional[Tuple[float, float]] = None,
        backend: Optional[str] = None,
        show: bool = False,
        path: Optional[str] = None,
        **kwargs,
    ):
        """Plot MCMC trace plots for convergence diagnostics.

        Shows the evolution of parameter values across MCMC iterations for
        each chain. Essential for diagnosing convergence issues, identifying
        burn-in periods, and detecting problematic sampling behavior such
        as poor mixing or trend.

        Args:
            figsize: Figure size as (width, height) in inches. If None, uses default size.
            backend: Backend to use ('matplotlib' or 'bokeh'). If None, uses default.
            show: Whether to display the plot immediately using plt.show().
                If False, the figure is returned without displaying.
            path: Optional file path to save the plot. If provided, the figure
                will be saved to this location.
            **kwargs: Additional keyword arguments passed to arviz.plot_trace.
                Common options include 'var_names' to select specific parameters,
                'compact' for layout style, and 'combined' for overlaying chains.

        Returns:
            matplotlib.figure.Figure or bokeh plot: The trace plot figure, or None if show=True.
                Contains time series plots for each parameter across all chains.
        """
        from catalax.mcmc.plotting import plot_trace

        f = plot_trace(self.mcmc, self.model, figsize, backend, **kwargs)
        if path is not None:
            f.savefig(path)
        if show:
            plt.show()

    def plot_forest(
        self,
        figsize: Optional[Tuple[float, float]] = None,
        backend: Optional[str] = None,
        show: bool = False,
        path: Optional[str] = None,
        **kwargs,
    ):
        """Plot forest plot of parameter distributions.

        Creates a horizontal plot showing credible intervals for all parameters
        in a compact format. Useful for comparing parameter estimates and
        uncertainties across many parameters simultaneously.

        Args:
            figsize: Figure size as (width, height) in inches. If None, uses default size.
            backend: Backend to use ('matplotlib' or 'bokeh'). If None, uses default.
            show: Whether to display the plot immediately using plt.show().
                If False, the figure is returned without displaying.
            path: Optional file path to save the plot. If provided, the figure
                will be saved to this location.
            **kwargs: Additional keyword arguments passed to arviz.plot_forest.
                Common options include 'hdi_prob' for credible interval probability,
                'combined' for pooling chains, and 'ess' for showing effective
                sample size.

        Returns:
            matplotlib.figure.Figure or bokeh plot: The forest plot figure, or None if show=True.
                Shows horizontal credible intervals for each parameter.
        """
        from catalax.mcmc.plotting import plot_forest

        f = plot_forest(self.mcmc, self.model, figsize, backend, **kwargs)
        if path is not None:
            f.savefig(path)

        if show:
            plt.show()
            return None

        return plt.gcf()

    def summary(self, hdi_prob: float = 0.95) -> Union[pd.DataFrame, xarray.Dataset]:
        """Generate comprehensive summary statistics for posterior samples.

        Computes and returns detailed summary statistics including central
        tendencies, credible intervals, effective sample sizes, and convergence
        diagnostics. This provides a comprehensive numerical summary of the
        MCMC results.

        Args:
            hdi_prob: Probability mass for highest density interval calculation.
                Determines the width of credible intervals in the summary.
                Default 0.95 corresponds to 95% credible intervals.

        Returns:
            Union[pd.DataFrame, xarray.Dataset]: Summary statistics table
                containing means, standard deviations, credible intervals,
                effective sample sizes, and R-hat diagnostics for each parameter.
                Format depends on the underlying implementation.
        """
        from catalax.mcmc.plotting import summary

        return summary(self.mcmc, hdi_prob)

    def plot_mcse(
        self,
        figsize: Optional[Tuple[float, float]] = None,
        backend: Optional[str] = None,
        show: bool = False,
        path: Optional[str] = None,
        **kwargs,
    ):
        """Plot Monte Carlo standard error diagnostics.

        Visualizes the Monte Carlo standard error (MCSE) for parameter estimates,
        which quantifies the uncertainty in posterior estimates due to finite
        sample size. Lower MCSE indicates more reliable estimates.

        Args:
            figsize: Figure size as (width, height) in inches. If None, uses default size.
            backend: Backend to use ('matplotlib' or 'bokeh'). If None, uses default.
            show: Whether to display the plot immediately using plt.show().
                If False, the figure is returned without displaying.
            path: Optional file path to save the plot. If provided, the figure
                will be saved to this location.
            **kwargs: Additional keyword arguments (currently unused but
                maintained for API consistency).

        Returns:
            matplotlib.figure.Figure or bokeh plot: The MCSE plot figure, or None if show=True.
                Shows Monte Carlo standard error for each parameter.
        """
        f = plot_mcse(self.mcmc, figsize, backend)
        if path is not None:
            f.savefig(path)
        if show:
            plt.show()
            return None
        return plt.gcf()

    def plot_ess(
        self,
        figsize: Optional[Tuple[float, float]] = None,
        backend: Optional[str] = None,
        show: bool = False,
        path: Optional[str] = None,
        **kwargs,
    ):
        """Plot effective sample size diagnostics.

        Visualizes the effective sample size (ESS) for each parameter, which
        measures how many independent samples the MCMC chains provide. Higher
        ESS indicates better sampling efficiency and more reliable estimates.

        Args:
            figsize: Figure size as (width, height) in inches. If None, uses default size.
            backend: Backend to use ('matplotlib' or 'bokeh'). If None, uses default.
            show: Whether to display the plot immediately using plt.show().
                If False, the figure is returned without displaying.
            path: Optional file path to save the plot. If provided, the figure
                will be saved to this location.
            **kwargs: Additional keyword arguments (currently unused but
                maintained for API consistency).

        Returns:
            matplotlib.figure.Figure or bokeh plot: The ESS plot figure, or None if show=True.
                Shows effective sample size for each parameter.
        """
        f = plot_ess(self.mcmc, figsize, backend)

        if path is not None:
            f.savefig(path)
        if show:
            plt.show()
            return None
        return plt.gcf()

    def to_arviz(self):
        """Convert MCMC results to ArviZ InferenceData format.

        Transforms the numpyro MCMC results into ArviZ's standardized format,
        enabling access to advanced diagnostic tools, visualization methods,
        and interoperability with other Bayesian analysis packages.

        Returns:
            arviz.InferenceData: MCMC results in ArviZ format containing
                posterior samples, sample statistics, and metadata. This format
                provides access to comprehensive Bayesian analysis tools and
                standardized visualization methods.
        """
        return az.from_numpyro(self.mcmc)

    def to_netcdf(self, path: str):
        """Save MCMC results to NetCDF format for persistent storage.

        Exports the complete MCMC results to a self-describing NetCDF file
        that preserves all sample data, metadata, and structure. This format
        is ideal for archiving results, sharing analyses, or loading results
        in different computing sessions.

        Args:
            path: File path where the NetCDF file will be saved. Should include
                the '.nc' extension. The file will contain all posterior samples,
                diagnostics, and metadata in a standardized format.

        Returns:
            None: Results are saved to the specified path. The file can later
                be loaded using arviz.from_netcdf() or similar tools.
        """
        return az.from_numpyro(self.mcmc).to_netcdf(path)

    def _is_jupyter_environment(self) -> bool:
        """Check if running in Jupyter environment.

        Returns:
            bool: True if running in Jupyter notebook/lab, False otherwise.
        """
        if not WIDGETS_AVAILABLE:
            return False

        try:
            # Check if we're in IPython and if it's a kernel
            ipython = get_ipython()
            if ipython is None:
                return False

            # Check for Jupyter-specific attributes
            return hasattr(ipython, "kernel") or ipython.__class__.__name__ in [
                "ZMQInteractiveShell",
                "TerminalInteractiveShell",
            ]
        except Exception:
            return False

    def _create_plot_tab(
        self,
        plot_name: str,
        plot_func,
        *args,
        figsize: Optional[Tuple[float, float]] = None,
        backend: Optional[str] = None,
        **kwargs,
    ):
        """Create a tab widget for a specific plot.

        Args:
            plot_name: Name of the plot for tab title
            plot_func: Function to call for generating the plot
            *args: Arguments to pass to plot function
            figsize: Figure size for the plot
            backend: Backend to use ('matplotlib' or 'bokeh')
            **kwargs: Additional keyword arguments for plot function

        Returns:
            ipywidgets.Output: Output widget containing the plot
        """
        output = widgets.Output()

        with output:
            try:
                # Add figsize and backend to kwargs if provided
                if figsize is not None:
                    kwargs["figsize"] = figsize
                if backend is not None:
                    kwargs["backend"] = backend

                plot_func(*args, **kwargs)

                # Only show plot if using matplotlib backend
                if backend != "bokeh":
                    plt.show()
            except Exception as e:
                print(f"Error creating {plot_name} plot: {str(e)}")

        return output

    def create_plots_widget(
        self,
        initial_condition: Optional[Dict[str, float]] = None,
        time: Optional[Array] = None,
        dt0: float = 0.1,
        figsize: Optional[Tuple[float, float]] = (10, 6),
        backend: Optional[str] = None,
        tabs_to_include: Optional[List[str]] = None,
    ) -> Optional[widgets.Widget]:
        """Create an interactive tabbed widget with all diagnostic plots.

        This method creates a comprehensive tabbed interface for exploring
        MCMC results with different diagnostic plots. Each tab contains
        a different type of visualization.

        Args:
            initial_condition: Initial conditions for credibility interval plot.
                If None, credibility interval tab will be excluded.
            time: Time points for credibility interval plot.
                If None, credibility interval tab will be excluded.
            dt0: Time step for credibility interval simulation.
            figsize: Default figure size for all plots.
            backend: Backend to use ('matplotlib' or 'bokeh'). If None, uses default.
            tabs_to_include: List of tab names to include. If None, includes all available.
                Available options: ['posterior', 'trace', 'corner', 'forest',
                'credibility', 'ess', 'mcse']

        Returns:
            ipywidgets.Tab: Tabbed widget with diagnostic plots, or None if widgets
                are not available or not in Jupyter environment.

        Example:
            >>> # Create widget with default settings
            >>> results.create_plots_widget()

            >>> # Create widget with specific tabs and backend
            >>> results.create_plots_widget(
            ...     figsize=(8, 6),
            ...     backend='bokeh',
            ...     tabs_to_include=['posterior', 'trace', 'corner']
            ... )
        """
        if not self._is_jupyter_environment():
            print("Interactive widgets are only available in Jupyter environments.")
            return None

        # Default tabs to include
        available_tabs = {
            "posterior": ("Posterior", self.plot_posterior),
            "trace": ("Trace", self.plot_trace),
            "corner": ("Corner", self.plot_corner),
            "forest": ("Forest", self.plot_forest),
            "ess": ("ESS", self.plot_ess),
            "mcse": ("MCSE", self.plot_mcse),
        }

        # Add credibility interval tab if data is provided
        if initial_condition is not None and time is not None:
            available_tabs["credibility"] = (
                "Credibility Interval",
                self.plot_credibility_interval,
            )

        # Filter tabs if specified
        if tabs_to_include is not None:
            available_tabs = {
                k: v for k, v in available_tabs.items() if k in tabs_to_include
            }

        # Create tab outputs
        tab_outputs = []
        tab_titles = []

        for tab_key, (title, plot_method) in available_tabs.items():
            tab_titles.append(title)

            # Special handling for credibility interval plot
            if tab_key == "credibility":
                output = self._create_plot_tab(
                    title,
                    plot_method,
                    initial_condition,
                    time,
                    dt0,
                    figsize=figsize,
                    backend=backend,
                    show=False,  # Don't show immediately in widget
                )
            else:
                output = self._create_plot_tab(
                    title,
                    plot_method,
                    figsize=figsize,
                    backend=backend,
                    show=False,  # Don't show immediately in widget
                )

            tab_outputs.append(output)

        # Create tabbed interface
        tab_widget = widgets.Tab(children=tab_outputs)

        # Set tab titles
        for i, title in enumerate(tab_titles):
            tab_widget.set_title(i, title)

        # Display the widget
        display(tab_widget)

        return tab_widget
