import pandas as pd

from typing import Any
from typing import List
from typing import Self
from typing import Dict
from typing import Type
from typing import Tuple
from matplotlib.axes import Axes
from pandas.core.frame import DataFrame

from .chart import JointChart
from .plotter import Line
from .plotter import Pareto
from .plotter import Plotter
from .plotter import Scatter
from .plotter import Violine
from .plotter import MeanTest
from .plotter import StripeLine
from .plotter import SkipSubplot
from .plotter import HideSubplot
from .plotter import Probability
from .plotter import GaussianKDE
from .plotter import BlandAltman
from .plotter import ParallelCoordinate

from ..anova import LinearModel
from ..strings import STR
from ..constants import KW
from ..constants import COLOR
from ..constants import ANOVA
from ..constants import CATEGORY


class ParameterRelevanceCharts(JointChart):
    """
    Provides a set of charts for visualizing the relevance of a linear
    regression model's parameters.

    The `ParameterRelevanceCharts` class takes a `LinearModel` instance 
    and generates a set of two charts for visualizing the relevance
    of the model's parameters:
    - Pareto chart of the parameter standardized effects
    - Pareto chart of the Sum of Squares for each parameter

    On the first chart, the standardized effect of each parameter is 
    visualized as a Pareto chart. The red line indicates the alpha
    level of significance.

    Parameters
    ----------
    linear_model : LinearModel
        The linear regression model whose parameters will be visualized.
    drop_intercept : bool, optional
        Whether to drop the intercept from the model, by default True.
    stretch_figsize : bool, optional
        If True, stretch the figure height and width based on the number of
        rows and columns, by default False.
    """
    __slots__ = ('lm')

    lm: LinearModel
    """The linear regression model whose parameters are visualized."""

    def __init__(
            self,
            linear_model: LinearModel,
            drop_intercept: bool = True,
            stretch_figsize: bool = False
            ) -> None:
        self.lm = linear_model
        effects =  self.lm.effects()
        if ANOVA.INTERCEPT in effects and drop_intercept:
            effects = effects.drop(ANOVA.INTERCEPT)
        data = (pd
            .concat([self.lm.anova('I'), effects], axis=1)
            .reset_index(drop=False)
            .rename(columns={'index': ANOVA.SOURCE}))

        super().__init__(
            source=data,
            target=(ANOVA.EFFECTS, ANOVA.TABLE_COLNAMES[1]),
            feature=ANOVA.SOURCE,
            target_on_y=False,
            ncols=2,
            nrows=1,
            stretch_figsize=stretch_figsize)
        
    def plot(self) -> Self: # type: ignore
        """Generates a set of two charts for visualizing the relevance 
        of the model's parameters:
        - Pareto chart of the parameter standardized effects
        - Pareto chart of the Sum of Squares for each parameter
        
        Returns
        -------
        Self: 
            The `ParameterRelevanceCharts` instance, for method chaining.

        Note
        ----
        The `plot()` method will generate two charts. The first chart
        will contain the Pareto chart of the parameter standardized
        effects. The second chart will contain the Pareto chart of the
        Sum of Squares for each parameter.

        The red line in the Pareto charts represents the threshold
        for the effect of a parameter. The threshold is calculated as
        the alpha risk where the parameter is not relevant.
        """
        super().plot(
            Pareto, no_percentage_line=True, skip_na='all')
        super().plot(
            Pareto, highlight=ANOVA.RESIDUAL, skip_na='all')
        
        self.axes[0, 0].axvline(
            self.lm.effect_threshold, **KW.SPECIFICATION_LINE)
        return self

    def label(self, info: bool | str = False, **kwds) -> Self: # type: ignore
        """Adds titles and labels to the charts generated by the 
        `plot()` method.
        
        Parameters
        ----------
        info : bool | str, optional
            If `True`, the method will add an informative subtitle to 
            the chart. If a string is provided, it will be used as the 
            subtitle, by default False.
        **kwds
            Additional keyword arguments to be passed to the `label()`
            method of the `JointChart` instance.
        
        Returns
        -------
        Self
            The `ParameterRelevanceCharts` instance, for method chaining.
        """      
        labels = dict(
                fig_title=STR['paramcharts_fig_title'],
                sub_title=STR['paramcharts_sub_title'],
                target_label=(STR['effects_label'], STR['ss_label']),
                feature_label=(STR['paramcharts_feature_label'], ''),
                info = info
            ) | kwds
        super().label(**labels) # type: ignore
        return self


class ResidualsCharts(JointChart):
    """
    Provides a set of charts for visualizing the residuals of a linear 
    regression model.

    The `ResidualsCharts` class takes a `LinearModel` instance and 
    generates a set of four charts:
    - Probability plot of the residuals
    - Gaussian kernel density estimate of the residuals
    - Scatter plot of the predicted values vs. the observed values
    - Line plot of the predicted values vs. the observed values

    The `plot()` method generates the charts, and the `label()` method 
    adds titles and labels to the charts.

    Parameters
    ----------
    linear_model : LinearModel
        The linear regression model whose residuals will be visualized.
    stretch_figsize : bool, optional
        If True, stretch the figure height and width based on the number of
        rows and columns, by default False.
    """
    __slots__ = ('lm')

    lm: LinearModel
    """The linear regression model whose residuals are visualized."""

    def __init__(
            self,
            linear_model: LinearModel,
            stretch_figsize: bool = False
            ) -> None:
        self.lm = linear_model
        super().__init__(
            source=self.lm.residual_data(),
            target=ANOVA.RESIDUAL,
            feature=('', '', ANOVA.PREDICTION, ANOVA.OBSERVATION),
            nrows=2,
            ncols=2,
            sharey=True,
            stretch_figsize=stretch_figsize)
        
    def plot(self) -> Self: # type: ignore
        """Generates a set of four charts for visualizing the residuals 
        of a linear regression model:
        - Probability plot of the residuals
        - Gaussian kernel density estimate of the residuals
        - Scatter plot of the predicted values vs. the observed values
        - Line plot of the predicted values vs. the observed values
        
        Returns
        -------
        Self: 
            The `ResidualsCharts` instance, for method chaining.
        """
        super().plot(Probability, show_fit_ci=True)
        super().plot(GaussianKDE)
        super().plot(Scatter)
        super().plot(Line, kw_call={'marker': 'o'})
        return self
    
    def stripes(self) -> Self: # type: ignore
        """Adds a line at position 0 for each subplot except for the 
        probability plot. This line represents the fit of the model."""
        fit = StripeLine(
            label=STR['fit'],
            position=0,
            orientation='horizontal',
            color=COLOR.PALETTE[0])
        for chart in self.itercharts():
            if chart == self.charts[0]:
                continue
            chart.stripes([fit])
        return self

    def label(self, info: bool | str = False, **kwds) -> Self: # type: ignore
        """Adds titles and labels to the charts generated by the 
        `plot()` method.
        
        Parameters
        ----------
        info : bool | str, optional
            If `True`, the method will add an informative subtitle to 
            the chart. If a string is provided, it will be used as the 
            subtitle, by default False.
        **kwds
            Additional keyword arguments to be passed to the `label()`
            method of the `JointChart` instance.
        
        Returns
        -------
        Self
            The `ResidualsCharts` instance, for method chaining.
        """
        sub_title = f'{self.lm.target} ~ {" + ".join(self.lm.effects().index)}'
        labels = dict(
                fig_title=STR['residcharts_fig_title'],
                sub_title=sub_title,
                target_label=STR['resid_name'],
                feature_label=(
                    STR['residcharts_flabel_quantiles'],
                    STR['residcharts_flabel_density'],
                     f'{STR["residcharts_flabel_predicted"]} {self.lm.target}',
                    STR['residcharts_flabel_observed']),
                info = info
            ) | kwds
        super().label(**labels) # type: ignore
        return self


class PairComparisonCharts(JointChart):
    """Provides a set of charts for visualizing the pairwise 
    comparison of two variables.
    
    Parameters
    ----------
    source : DataFrame
        The source data.
    target : str or Tuple[str]
        The target variable(s).
    feature : str or Tuple[str]
        The feature variable(s).
    identity : str
        Column name containing identities of each sample, must occur 
        once for each measurement.
    stretch_figsize : bool, optional
        If True, stretch the figure height and width based on the number of
        rows and columns, by default False.
    """

    __slots__ = ('identity')

    identity: str
    """Column name containing identities of each sample."""

    def __init__(
            self,
            source: DataFrame,
            target: str,
            feature: str,
            identity: str,
            stretch_figsize: bool = False
            ) -> None:
        self.identity = identity
        super().__init__(
            source=source,
            target=target,
            feature=feature,
            nrows=1,
            ncols=2,
            width_ratios=[4, 1],
            categorical_feature = (False, True),
            stretch_figsize=stretch_figsize)
    
    def plot(self) -> Self: # type: ignore
        """Generates a set of two charts for visualizing the difference
        between each pair.
        - BlandAltman where the difference is shown on the y-axis
        - ParallelCoordinate including a MeanTest and Violine plot, to 
        reflect the difference within the absolute values.
        
        Returns
        -------
        Self: 
            The `PairComparisonCharts` instance, for method chaining.
        """
        super().plot(BlandAltman, identity=self.identity, feature_axis='data')
        super().plot(
            ParallelCoordinate, identity=self.identity, show_points=False)
        super().plot(MeanTest, on_last_axes=True)
        super().plot(Violine, on_last_axes=True)
        return self

    def label(self, info: bool | str = False, **kwds) -> Self: # type: ignore
        """Adds titles and labels to the charts generated by the 
        `plot()` method.
        
        Parameters
        ----------
        info : bool | str, optional
            If `True`, the method will add an informative subtitle to 
            the chart. If a string is provided, it will be used as the 
            subtitle, by default False.
        **kwds
            Additional keyword arguments to be passed to the `label()`
            method of the `JointChart` instance.
        
        Returns
        -------
        Self
            The `PairComparisonCharts` instance, for method chaining.
        """
        labels = dict(
                fig_title=STR['paircharts_fig_title'],
                sub_title=STR['paircharts_sub_title'],
                target_label=(True, True),
                feature_label=(True, True),
                info = info
            ) | kwds
        super().label(**labels) # type: ignore
        return self


class BivariateUnivariateCharts(JointChart):
    """Provides a set of charts for visualizing the relationship between
    a target variable and a feature variable.

    This class prepares a grid for drawing a bivariate plot with 
    marginal univariate plots:
    - Bottom left there is a large square axis for the bivariate
      diagram. For example, use a `LinearRegressionLine` plotter to show 
      the dependence of the target variable on the feature variable.
    - Top left there is a small square axis for the univariate
      plot of the feature variable.
    - Bottom right there is a small square axis for the univariate
      plot of the target variable.
    - Top right there is a small square axis hidden as soon a plot is 
      done.

    Parameters
    ----------
    source : DataFrame
        The source data.
    target : str
        The target (dependant) variable.
    feature : str
        The feature (independant) variable.
    hue : str, optional
        The hue variable.
    dodge_univariates : bool, optional
        Whether to dodge the univariate plots, by default False.
    categorical_feature_univariates : bool, optional
        Whether to treat the feature variable for univariate charts as
        categorical, by default False.
    ratios : List[float], optional
        The ratios of the bivariate axes height to the univariate axes 
        height, by default [4, 1].
    stretch_figsize : bool, optional
        If True, stretch the figure height and width based on the number of
        rows and columns, by default False.
    colors: Tuple[str, ...], optional
        Tuple of unique colors used for hue categories as hex or str. If
        not provided, the default colors will be used, by default ().
    """
    __slots__ = ('_top_right_hidden')

    _top_right_hidden: bool
    """Whether the top right axis is hidden."""

    def __init__(
            self,
            source: DataFrame,
            target: str,
            feature: str,
            hue: str = '',
            dodge_univariates: bool = False,
            categorical_feature_univariates: bool = False,
            ratios: List[float] = [3, 1],
            stretch_figsize: bool = False,
            colors: Tuple[str, ...] = ()
            ) -> None:
        assert len(ratios) == 2, ('ratios must be a list of two floats')
        self._top_right_hidden = False

        super().__init__(
            source=source,
            target=(
                feature, '',
                target, target),
            feature=(
                '', '',
                feature, ''),
            hue=hue,
            dodge=(
                dodge_univariates, False,
                False, dodge_univariates),
            categorical_feature=(
                categorical_feature_univariates, False,
                False, categorical_feature_univariates),
            target_on_y=(
                False, True,
                True, True),
            nrows=2,
            ncols=2,
            sharex='col',
            sharey='row',
            width_ratios=ratios,
            height_ratios=ratios[::-1],
            stretch_figsize=stretch_figsize,
            colors=colors or CATEGORY.PALETTE)
    
    @property
    def hidden_ax(self) -> Axes:
        """Get the top right axis (read-only)."""
        return self.axes[0, 1]
    
    @property
    def univariate_axs(self) -> List[Axes]:
        """Get the univariate axes (read-only)."""
        return [self.axes[0], self.axes[-1]]
    
    @property
    def bivariate_ax(self) -> Axes:
        """Get the bivariate axis (read-only)."""
        return self.axes[1, 0]
    
    def hide_top_right_ax(self) -> None:
        """Hides the top right axis if it is not already hidden."""
        if self._top_right_hidden:
            return

        HideSubplot(self.hidden_ax)()
        self._top_right_hidden = True
    
    def plot(
            self,
            plotter: Type[Plotter],
            kw_call: Dict[str, Any] = {},
            kw_where: Dict[str, Any] = {},
            on_last_axes: bool = False,
            **kwds
            ) -> Self:
        raise NotImplementedError(
            'Use the methods "plot_univariates" and "plot_bivariate"!')
    
    def plot_univariates(
            self,
            plotter: Type[Plotter],
            kw_call: Dict[str, Any] = {},
            kw_where: Dict[str, Any] = {},
            **kwds
            ) -> Self:
        """Plots the univariate plots on the top left and bottom right
        axes.

        Parameters
        ----------
        plotter : Type[Plotter]
            The plotter to use for the univariate plots.
        kw_call : Dict[str, Any]
            Additional keyword arguments for the plotter call method.
        kw_where : Dict[str, Any]
            Additional keyword arguments for the where method used to
            filter the data.
        on_last_axes : bool, optional
            If True, plot on the last axes in the grid. If False, plot 
            on the next axes in the grid, by default False.
        **kwds
            Additional keyword arguments to be passed to the super plot
            method.

        Returns
        -------
        Self
            The `BivariateUnivariateCharts` instance, for method 
            chaining.
        """
        self.hide_top_right_ax()
        for ax in self.axes:
            if ax in self.univariate_axs:
                super().plot(
                    plotter,
                    kw_call=kw_call,
                    kw_where=kw_where,
                    on_last_axes=False,
                    **kwds)
            else:
                super().plot(SkipSubplot)
        return self

    def plot_bivariate(
            self,
            plotter: Type[Plotter],
            kw_call: Dict[str, Any] = {},
            kw_where: Dict[str, Any] = {},
            **kwds) -> Self:
        """Plots the bivariate plot on the bottom left axis.

        Parameters
        ----------
        plotter : Type[Plotter]
            The plotter to use for the univariate plots.
        kw_call : Dict[str, Any]
            Additional keyword arguments for the plotter call method.
        kw_where : Dict[str, Any]
            Additional keyword arguments for the where method used to
            filter the data.
        **kwds
            Additional keyword arguments to be passed to the super plot
            method.

        Returns
        -------
        Self
            The `BivariateUnivariateCharts` instance, for method 
            chaining.
        """
        self.hide_top_right_ax()
        for ax in self.axes:
            if ax == self.bivariate_ax:
                super().plot(
                    plotter,
                    kw_call=kw_call,
                    kw_where=kw_where,
                    on_last_axes=False,
                    **kwds)
            else:
                super().plot(SkipSubplot)
        return self
    
    def label(
            self,
            *,
            fig_title: str = '',
            sub_title: str = '',
            feature_label: str | bool | Tuple = '', 
            target_label: str | bool | Tuple = '', 
            info: bool | str = False,
            axes_titles: Tuple[str, ...] = (),
            rows: Tuple[str, ...] = (),
            cols: Tuple[str, ...] = (),
            row_title: str = '',
            col_title: str = '') -> Self:
        """Add labels and titles to the chart.

        This method sets various labels and titles for the chart,
        including figure title, subplot title, axis labels, row and
        column titles, and additional information.

        Parameters
        ----------
        fig_title : str, optional
            The main title for the entire figure, by default ''.
        sub_title : str, optional
            The subtitle for the entire figure, by default ''.
        feature_label : str | bool | None, optional
            The label for the feature variable (x-axis), by default ''.
            If set to True, the feature variable name will be used.
            If set to False or None, no label will be added.
        target_label : str | bool | None, optional
            The label for the target variable (y-axis), by default ''.
            If set to True, the target variable name will be used.
            If set to False or None, no label will be added.
        info : bool | str, optional
            Additional information to display on the chart. If True,
            the date and user information will be automatically added at
            the lower left corner of the figure. If a string is
            provided, it will be shown next to the date and user,
            separated by a comma. By default, no additional information
            is displayed.
        axes_titles : Tuple[str, ...]
            Title for each Axes, by default ()
        rows: Tuple[str, ...], optional
            The row labels of the figure, by default ().
        cols: Tuple[str, ...], optional
            The column labels of the figure, by default ().
        row_title : str, optional
            The title of the rows, by default ''.
        col_title : str, optional
            The title of the columns, by default ''.

        Returns
        -------
        Self
            The `BivariateUnivariateCharts` instance, for method
            chaining.

        Notes
        -----
        This method allows customization of chart labels and titles to
        enhance readability and provide context for the visualized data.
        """
        super().label(
            fig_title=fig_title,
            sub_title=sub_title,
            feature_label=(
                '', '',
                feature_label, ''),
            target_label=(
                '', '',
                target_label, ''),
            info=info,
            axes_titles=axes_titles,
            rows=rows,
            cols=cols,
            row_title=row_title,
            col_title=col_title)
        return self


__all__ = [
    'ResidualsCharts',
    'ParameterRelevanceCharts',
    'PairComparisonCharts',
    'BivariateUnivariateCharts',
    ]
