import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from ravix.modeling.parse_formula import parse_formula
from typing import Union, Optional, Tuple, Literal
import matplotlib.patches as mpatches

def plot_cor(
    formula: Union[str, pd.DataFrame],
    data: Optional[pd.DataFrame] = None,
    style: Literal[1, 2, 3, 4] = 1,
    title: str = "Correlation Plot",
    xlab: str = "",
    ylab: str = "",
    figsize: Tuple[float, float] = (10, 8),
    show: bool = True,
    **kwargs
) -> Optional[Tuple[plt.Figure, plt.Axes]]:
    """
    Create a correlation matrix heatmap with various visualization styles.
    
    This function generates correlation matrix visualizations with four different
    display types, ranging from detailed annotations to clean visual representations.
    Automatically handles formulas or DataFrames and computes correlations for
    numeric variables only.
    
    Parameters
    ----------
    formula : str or pd.DataFrame
        Formula specifying variables (e.g., "Y ~ X1 + X2") or DataFrame.
        If DataFrame provided, all numeric columns are used.
    data : pd.DataFrame, optional
        DataFrame containing variables when formula is provided.
    style : {1, 2, 3, 4}, default=1
        Visualization style:
        
        Style 1: Split triangle display (default)
            - Diagonal: Black squares
            - Upper triangle: Color-coded squares (no numbers)
            - Lower triangle: Bold numbers (no color)
        
        Style 2: Circle magnitude display
            - Same as Style 1 but upper triangle shows circles sized by correlation magnitude
            - Diagonal: Black squares
            - Lower triangle: Bold numbers (no color)
        
        Style 3: Clean color-only display
            - No annotations, only color-coded squares
            - Both triangles show colors
            - Diagonal included with colors
        
        Style 4: Full annotation display
            - All cells show correlation values (2 decimals)
            - Color-coded background
            - Diagonal included
    title : str, default="Correlation Plot"
        Main title of the plot.
    xlab : str, default=""
        X-axis label.
    ylab : str, default=""
        Y-axis label.
    figsize : tuple, default=(10, 8)
        Figure size as (width, height) in inches.
    show : bool, default=True
        Display plot immediately; if False, return (fig, ax) for manual control
    **kwargs : dict
        Additional keyword arguments passed to sns.heatmap().
        Common options: cmap, linewidths, linecolor, cbar_kws, annot_kws, cbar
    
    Returns
    -------
    Optional[Tuple[plt.Figure, plt.Axes]]
        None if show=True, else (fig, ax) tuple
        
        WARNING: If show=False, caller is responsible for closing the figure
        to prevent memory leaks. Use plt.close(fig) when done.
    
    Examples
    --------
    >>> # Style 1: Circle magnitude in upper triangle with values in lower (default)
    >>> plot_cor("mpg ~ hp + wt + disp + drat", data=mtcars)
    
    >>> # Style 2: Colors in upper and values in lower
    >>> plot_cor(mtcars[["mpg", "hp", "wt", "disp"]], style=2)
    
    >>> # Style 3: Clean color-only display
    >>> plot_cor("mpg ~ hp + wt + disp", data=mtcars, style=3)
    
    >>> # Style 4: Full annotations
    >>> plot_cor(mtcars, style=4, figsize=(12, 10))
    
    >>> # Custom colormap and no colorbar
    >>> plot_cor(mtcars, style=1, cmap="RdBu_r", cbar=False)
    
    >>> # Manual figure management
    >>> fig, ax = plot_cor(mtcars, style=3, show=False)
    >>> ax.set_title("Custom Title")
    >>> plt.savefig("correlation_matrix.png")
    >>> plt.close(fig)  # Important!
    
    Notes
    -----
    - Only numeric columns are included in correlation computation
    - Diagonal represents correlation of variable with itself (always 1.0)
    - Color scale: Red (negative), White (zero), Blue (positive)
    - Style 1 and 2 provide asymmetric information (visual + numerical)
    - Style 3 is best for presentations (clean, no clutter)
    - Style 4 is most detailed (all correlations visible)
    - Categorical variables are automatically excluded
    
    Raises
    ------
    ValueError
        If no numeric columns found in data or invalid style specified.
    
    See Also
    --------
    plots : Scatter plot matrix for visual correlation assessment
    """
    # Font size settings
    TITLE_FONTSIZE = 14
    LABEL_FONTSIZE = 12
    TICK_FONTSIZE = 12
    
    # Validate style
    if style not in [1, 2, 3, 4]:
        raise ValueError(
            f"Invalid style '{style}'. "
            f"Must be 1, 2, 3, or 4."
        )
    
    # Handle DataFrame input
    if isinstance(formula, pd.DataFrame):
        data = formula
        formula = None
    
    # Parse formula if provided
    if formula is not None:
        formula = formula + "+0"
        Y_out, X_out = parse_formula(formula, data, drop_first=False)
        Y_name = Y_out.name
        # Combine Y and X data for the correlation matrix
        data = pd.concat([pd.Series(Y_out, name=Y_name), X_out], axis=1)
    
    # Keep only numeric columns
    numeric_data = data.select_dtypes(include=[np.number])
    
    if numeric_data.empty:
        raise ValueError(
            "No numeric columns found in the data. "
            "Cannot compute correlation matrix."
        )
    
    # Calculate the correlation matrix
    corr_matrix = numeric_data.corr()
    
    # Set default colormap if not provided
    if 'cmap' not in kwargs:
        kwargs['cmap'] = 'RdBu'
    
    # Style-specific plotting
    if style == 1:
        # Style 1: Split triangle - upper (color only), lower (numbers), diagonal (black)
        return _plot_type1(corr_matrix, title, xlab, ylab, figsize, show, **kwargs)
    
    elif style == 2:
        # Style 2: Circles in upper triangle sized by magnitude
        return _plot_type2(corr_matrix, title, xlab, ylab, figsize, show, **kwargs)
    
    elif style == 3:
        # Style 3: Clean color-only display
        return _plot_type3(corr_matrix, title, xlab, ylab, figsize, show, **kwargs)
    
    elif style == 4:
        # Style 4: Full annotations with 2 decimals
        return _plot_type4(corr_matrix, title, xlab, ylab, figsize, show, **kwargs)




def _plot_type1(corr_matrix, title, xlab, ylab, figsize, show, **kwargs):
    """Type 1: Circles sized by magnitude in upper triangle."""
    n = len(corr_matrix)
    
    # Create figure
    fig, ax = plt.subplots(figsize=figsize)
    
    # Create a modified correlation matrix:
    # - Lower triangle: 0 (white, for bold numbers to show up)
    # - Upper triangle: 0 (white, for circles to show up)
    # - Diagonal: 0 (will be painted black)
    display_matrix = corr_matrix.copy()
    for i in range(n):
        for j in range(n):
            if i <= j:  # Upper triangle and diagonal
                display_matrix.iloc[i, j] = 0  # Set to 0 for white
            elif i > j:  # Lower triangle
                display_matrix.iloc[i, j] = 0  # Set to 0 for white
    
    # Create annotation matrix - only show numbers in lower triangle
    annot_matrix = np.empty((n, n), dtype=object)
    for i in range(n):
        for j in range(n):
            if i > j:  # Lower triangle
                annot_matrix[i, j] = f'{corr_matrix.iloc[i, j]:.2f}'
            else:
                annot_matrix[i, j] = ''
    
    # Setup colormap - red=negative, blue=positive
    cmap_name = kwargs.pop('cmap', 'RdBu')
    cmap = sns.color_palette(cmap_name, as_cmap=True)
    
    # Setup kwargs for heatmap
    heatmap_kwargs = {k: v for k, v in kwargs.items() if k not in ['annot', 'fmt', 'annot_kws']}
    heatmap_kwargs.update({'vmin': -1, 'vmax': 1, 'center': 0, 'square': True,
                           'linewidths': 0.5, 'cbar_kws': {"shrink": 1.0}})
    
    # Draw the heatmap with bold annotations in lower triangle
    sns.heatmap(display_matrix, annot=annot_matrix, fmt='',
                annot_kws={'weight': 'bold', 'color': 'black'},
                cmap=cmap, ax=ax, **heatmap_kwargs);
    
    # Paint the diagonal black
    _ = [ax.add_patch(plt.Rectangle((i, i), 1, 1, fill=True, facecolor='black', 
                                    linewidth=0.5, edgecolor='gray')) 
         for i in range(n)]
    
    # Add circles to upper triangle ON WHITE background
    _ = [ax.add_patch(mpatches.Circle((j + 0.5, i + 0.5), 
                                      abs(corr_matrix.iloc[i, j])*.5,
                                      facecolor=plt.cm.RdBu((corr_matrix.iloc[i, j] + 1) / 2),
                                      ec='black', linewidth=0.5))
         for i in range(n) for j in range(i + 1, n)]
    
    # Set labels
    ax.set_title(title, fontsize=14, fontweight='bold')
    if xlab:
        ax.set_xlabel(xlab, fontsize=12)
    else:
        ax.set_xlabel('')
    if ylab:
        ax.set_ylabel(ylab, fontsize=12)
    else:
        ax.set_ylabel('')
    ax.tick_params(labelsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    
    plt.tight_layout()
    if show:
        plt.show()
        plt.close(fig)
        return None
    return fig, ax

def _plot_type2(corr_matrix, title, xlab, ylab, figsize, show, **kwargs):
    """Type 2: Split triangle with upper color-only, lower bold numbers (no color), black diagonal."""
    n = len(corr_matrix)
    
    # Create annotation matrix - only show numbers in lower triangle
    annot_matrix = np.empty((n, n), dtype=object)
    for i in range(n):
        for j in range(n):
            if i > j:  # Lower triangle
                annot_matrix[i, j] = f'{corr_matrix.iloc[i, j]:.2f}'
            else:
                annot_matrix[i, j] = ''
    
    # Setup colormap - red=negative, blue=positive
    cmap_name = kwargs.pop('cmap', 'RdBu')
    cmap = sns.color_palette(cmap_name, as_cmap=True)
    
    # Create figure
    fig, ax = plt.subplots(figsize=figsize)
    
    # Create a modified correlation matrix with lower triangle set to 0 (white)
    display_matrix = corr_matrix.copy()
    for i in range(n):
        for j in range(n):
            if i > j:  # Lower triangle
                display_matrix.iloc[i, j] = 0  # Set to 0 for white color
    
    # Setup kwargs for heatmap
    heatmap_kwargs = {k: v for k, v in kwargs.items() if k not in ['annot', 'fmt', 'annot_kws']}
    heatmap_kwargs.update({
        'vmin': -1, 
        'vmax': 1, 
        'center': 0, 
        'square': True,
        'linewidths': 0.5, 
        'linecolor': 'gray',
        'cbar_kws': {"shrink": 1.0}
    })
    
    # Draw the heatmap with bold annotations in lower triangle
    sns.heatmap(display_matrix, annot=annot_matrix, fmt='', cmap=cmap, 
                annot_kws={'weight': 'bold', 'color': 'black'},
                ax=ax, **heatmap_kwargs);
    
    # Now paint the diagonal black
    _ = [ax.add_patch(plt.Rectangle((i, i), 1, 1, fill=True, facecolor='black', 
                                    linewidth=0.5, edgecolor='gray')) 
         for i in range(n)]
    
    # Set labels
    ax.set_title(title, fontsize=14, fontweight='bold')
    if xlab:
        ax.set_xlabel(xlab, fontsize=12)
    else:
        ax.set_xlabel('')
    if ylab:
        ax.set_ylabel(ylab, fontsize=12)
    else:
        ax.set_ylabel('')
    ax.tick_params(labelsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    
    plt.tight_layout()
    if show:
        plt.show()
        plt.close(fig)
        return None
    return fig, ax


def _plot_type3(corr_matrix, title, xlab, ylab, figsize, show, **kwargs):
    """Type 3: Clean color-only display, no annotations."""
    fig, ax = plt.subplots(figsize=figsize)
    
    # Setup kwargs - default to RdBu (red=negative, blue=positive)
    if 'cmap' not in kwargs:
        kwargs['cmap'] = 'RdBu'
    
    heatmap_kwargs = {k: v for k, v in kwargs.items()}
    heatmap_kwargs.update({'annot': False, 'vmin': -1, 'vmax': 1, 'center': 0,
                           'square': True, 'linewidths': 0.5, 
                           'cbar_kws': {"shrink": 1.0}})
    
    sns.heatmap(corr_matrix, ax=ax, **heatmap_kwargs);
    
    # Set labels
    ax.set_title(title, fontsize=14, fontweight='bold')
    if xlab:
        ax.set_xlabel(xlab, fontsize=12)
    else:
        ax.set_xlabel('')
    if ylab:
        ax.set_ylabel(ylab, fontsize=12)
    else:
        ax.set_ylabel('')
    ax.tick_params(labelsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    
    plt.tight_layout()
    if show:
        plt.show()
        plt.close(fig)
        return None
    return fig, ax


def _plot_type4(corr_matrix, title, xlab, ylab, figsize, show, **kwargs):
    """Type 4: Full annotations with 2 decimals everywhere."""
    fig, ax = plt.subplots(figsize=figsize)
    
    # Create annotation matrix with 2 decimals
    annot_matrix = corr_matrix.map(lambda x: f'{x:.2f}')
    
    # Setup kwargs - default to RdBu (red=negative, blue=positive)
    if 'cmap' not in kwargs:
        kwargs['cmap'] = 'RdBu'
    
    heatmap_kwargs = {k: v for k, v in kwargs.items() if k not in ['annot', 'fmt']}
    heatmap_kwargs.update({'vmin': -1, 'vmax': 1, 'center': 0, 'square': True,
                           'linewidths': 0.5, 'cbar_kws': {"shrink": 1.0}})
    
    sns.heatmap(corr_matrix, annot=annot_matrix, fmt='', ax=ax, **heatmap_kwargs);
    
    # Set labels
    ax.set_title(title, fontsize=14, fontweight='bold')
    if xlab:
        ax.set_xlabel(xlab, fontsize=12)
    else:
        ax.set_xlabel('')
    if ylab:
        ax.set_ylabel(ylab, fontsize=12)
    else:
        ax.set_ylabel('')
    ax.tick_params(labelsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    
    plt.tight_layout()
    if show:
        plt.show()
        plt.close(fig)
        return None
    return fig, ax
