"""Plotting functions with an example below"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_plotting.ipynb.

# %% auto 0
__all__ = ['BivariateChoropleth', 'plot_bivariate_choropleth']

# %% ../nbs/03_plotting.ipynb 2
from .grid_palette import GridPalette, BivariateGridPalette, letters
import geopandas as gpd
import pandas as pd
import matplotlib.pyplot as plt
from string import ascii_uppercase    

# %% ../nbs/03_plotting.ipynb 4
class BivariateChoropleth:
    "Class for Bivariate Choropleth"
    def __init__(self, gdf:gpd.GeoDataFrame, 
                 x:str, #column name for the x part of the bivariate axis
                 y:str, #column name for the y part of the bivariate axis
                 grid_size:int=3, # grid size, used in pd.cut and GridPalette
                 gp=None, 
                 palette_name='blues2reds'):
        if not isinstance(gdf, gpd.GeoDataFrame):
            raise ValueError('Inputted frame is not a GeoDataFrame')
        
        self.gdf = gdf.loc[:, [x,y, 'geometry']]
        self.x, self.y = x,y
        self.grid_size = grid_size
        self.gdf = self.gdf.pipe(BivariateChoropleth.get_bivar_rank, x, y, grid_size)
        if gp is None:
            self.gp = BivariateGridPalette(palette_name=palette_name, grid_size=grid_size)
        else:
            self.gp = gp
        
    @staticmethod
    def get_bivar_rank(df,x,y,grid_size=3):
        "Create the bivariate ranks. These are of 1A, 3C etc., where y are numbers and x are letters"
        x_ranks = [letters[x] for x in range(grid_size)]
        y_ranks = list(range(1,grid_size+1))
        bivar_ranks = [str(y)+x for y in y_ranks for x in x_ranks]

        x_rank_col=x[:3]+'_rank'
        y_rank_col=y[:3]+'_rank'

        df[x_rank_col] = pd.cut(df[x], grid_size, labels=x_ranks)
        df[y_rank_col] = pd.cut(df[y], grid_size, labels=y_ranks)

        df['bivar_rank'] = df[y_rank_col].astype(str) + df[x_rank_col].astype(str)
        df['bivar_rank'] = pd.Categorical(df.bivar_rank, ordered=True, categories=bivar_ranks)

        return df
    
    def set_defaults(self):
        pass
    
    @staticmethod
    def add_grid_legend(ax=None, gp=None, grid_size=None, palette_name='blues2reds'):
        if grid_size is None:
            # is this a good idea?
            grid_size = globals().get('grid_size', 3)
        if gp is None:
            gp = BivariateGridPalette(palette_name=palette_name, grid_size=grid_size)
        if ax is None:
            ax = plt.gca()

        #TODO: allow this to be changed by user input
        cax = ax.inset_axes([1.1,0,0.2,0.2])
        cax.imshow(gp.color_matrix)
        cax.set_xticks([])
        cax.set_yticks([])
        cax.scatter(-0.5,-0.5, marker='^', color='k', clip_on=False)
        cax.scatter(grid_size-0.5,grid_size-0.5, marker='>', color='k', clip_on=False) 
        return cax
    
    def plot_choropleth_with_grid_legend(self):
        ax = self.gdf.plot(column='bivar_rank', cmap=self.gp.cmap)
        ax.set_axis_off()
        cax = BivariateChoropleth.add_grid_legend(ax, gp=self.gp, grid_size=self.grid_size)
        cax.set_xlabel(self.x)
        cax.set_ylabel(self.y)
        return ax, cax

# %% ../nbs/03_plotting.ipynb 5
def plot_bivariate_choropleth(
    data: gpd.GeoDataFrame,
    x: str,
    y: str,
    grid_size: int = None,  # Overrides grid_palette size if given
    grid_palette: GridPalette = None,  # Defaults to palette_name
    palette_name: str = 'blues2reds',  # Only used if grid_palette is None
    return_frame: bool = False
) -> plt.Axes:
    """Plotting function for bivariate choropleth."""
    
    if grid_palette is not None:
        if not isinstance(grid_palette, GridPalette):
            print("grid_palette is not GridPalette, using default.")
            grid_palette = None
        else:
            # Use grid_palette's size if grid_size is not provided
            if grid_size is None:
                grid_size = grid_palette.grid_size
            
            # Ensure grid_palette matches the desired grid_size
            if grid_size is not None and grid_palette.grid_size != grid_size:
                print(f"Resizing grid_palette from {grid_palette.grid_size} to {grid_size}.")
                grid_palette.resize(grid_size)
    
    # Final fallback if grid_size is still None
    if grid_size is None:
        grid_size = 3

    choropleth = BivariateChoropleth(
        data, x=x, y=y, grid_size=grid_size, gp=grid_palette, palette_name=palette_name
    )
    ax, cax = choropleth.plot_choropleth_with_grid_legend()

    if return_frame:
        return ax, cax, choropleth.gdf
    return ax, cax

