import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import Normalize
from IPython.display import display, HTML

import numpy as np
from matplotlib.cm import ScalarMappable, get_cmap
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter

def highlighter(
    items: list,
    values: list,
    color_range: tuple = (-3.5, 3.5),
    cmap_name: str = "coolwarm_r"
):
    """
    Highlight a list of text items with background colors based on corresponding values.

    This function generates a visual HTML display of text items with background colors 
    that reflect the magnitude of their associated values, mapped to a color range using 
    a specified matplotlib colormap.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    items : list of str
        A list of strings to be displayed with color highlights.
    values : list of float
        A list of numerical values corresponding to each item. These values are used 
        to determine the background color intensity.
    color_range : tuple of float, optional (default = (-3.5, 3.5))
        The range (min, max) used for normalizing values before applying the colormap.
    cmap_name : str, optional (default = "coolwarm_r")
        The name of the matplotlib colormap used to map values to colors.

    ---------------------------------------------------------------------------
    OUTPUT:
    None
        Displays an HTML snippet in a Jupyter notebook cell with the highlighted text items.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    # Setup normalization and colormap
    norm = Normalize(vmin=color_range[0], vmax=color_range[1])
    cmap = cm.get_cmap(cmap_name)

    def get_color(score):
        rgba = cmap(norm(score))
        return f"rgba({int(rgba[0]*255)}, {int(rgba[1]*255)}, {int(rgba[2]*255)}, 1.0)"

    # Build HTML
    highlighted_html = ""
    for item, val in zip(items, values):
        color = get_color(val)
        highlighted_html += (
            f'<span style="background-color: {color}; padding: 2px; border-radius: 4px;">{item}</span> '
        )

    display(HTML(f"""
<div style='line-height: 1.8; white-space: normal; word-wrap: break-word; max-width: 100%;'>
    {highlighted_html}
</div>
"""))

def color_map(gjson_file, color_covariate: list, color_range: tuple = (0,1), color_map='coolwarm', color_bar_label='Color Bar'):
    # Read geospatial data
    import geopandas as gpd
    gdf = gpd.read_parquet(gjson_file)
    
    # Validate length
    if len(color_covariate) != len(gdf):
        raise ValueError("Length of color_covariate must match number of rows in GeoDataFrame.")

    norm = Normalize(vmin=color_range[0], vmax=color_range[1])
    cmap = get_cmap(color_map)

    # Map normalized values to RGBA colors
    color_mapped = [cmap(norm(val)) for val in color_covariate]

    # Plot
    fig, ax = plt.subplots(figsize=(12, 8))
    gdf.plot(ax=ax,
             color=color_mapped,
             edgecolor='black',
             linewidth=0.15)

    # Colorbar setup
    sm = ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array(color_covariate)  # For colorbar scale
    cbar = plt.colorbar(sm, ax=ax, shrink=0.6)
    cbar.set_label(color_bar_label, rotation=270, labelpad=15)

    ax.set_axis_off()
    plt.tight_layout()
    plt.show()
    
    
def nice_plot(x, y, title='title', xlabel='xlabel', ylabel='ylabel', 
              color='blue', linewidth=2, showplot=True, xlim=None,
              ylim=None, figsize=(6,5), savename=None, fontsize=16,
              tick_fontsize=None, ax=None, marker = None, show_labels = False):
    # Create new figure/axes if none provided
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize, dpi=300)
    else:
        fig = ax.figure  # Use existing figure from provided ax

    if tick_fontsize is None:
        tick_fontsize = fontsize - 2

    # Set light grey background
    ax.set_facecolor('#f0f0f0')

    # Plot the data
    if not marker:
        ax.plot(x, y, color=color, linewidth=linewidth)
    else:
        ax.scatter(x, y, color=color, marker=marker)
        
    if show_labels:
        for i, (xi, yi) in enumerate(zip(x, y)):
            ax.annotate(f'label_{i}', (xi, yi), xytext=(5, 5), 
                    textcoords='offset points', fontsize=10, ha='left')

    # Grid, ticks, labels
    ax.grid(True, which='major', axis='both', color='white', linewidth=1)
    ax.tick_params(axis='both', colors='black', labelsize=tick_fontsize)
    ax.set_xlabel(xlabel, fontsize=fontsize)
    ax.set_ylabel(ylabel, fontsize=fontsize)
    ax.set_title(title, pad=10, fontsize=fontsize+2, weight='bold')
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)

    if savename and ax is None:
        fig.savefig(savename, format=savename[-3:], dpi=300, bbox_inches='tight')

    if showplot and ax is None:
        plt.show()

    return fig, ax


def nice_hist(data, xlabel = 'xlabel', ylabel = 'Frequency', title = 'title', 
              color = 'red', nbins = 20, edgecolor='black', xlim = None,
              ylim = None, figsize = (6,5), savename = None, showplot = True):
    plt.rcParams.update({'font.size': 14})

    fig, ax = plt.subplots(figsize=figsize, dpi=300)

    fig.patch.set_facecolor('white')
    ax.set_facecolor('#f0f0f0')

    ax.hist(data, 
        bins=nbins, 
        color=color, 
        edgecolor=edgecolor)

    # Add padding to the title for extra space
    ax.set_title(title, fontsize=18, weight='bold', pad=10)

    ax.set_xlabel(xlabel, fontsize=16, labelpad=12)
    ax.set_ylabel(ylabel, fontsize=16, labelpad=12)

    ax.tick_params(axis='both', labelsize=12)
    ax.grid(True, color='white', linestyle='-', linewidth=1)
    ax.set_axisbelow(True)

    ax.margins(x=0.05, y=0.1)     # Add space inside axes
    plt.tight_layout(pad=2.0)     # Add space outside axes
    
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    
    ax.yaxis.set_major_formatter(ScalarFormatter(useMathText=True))
    ax.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
    ax.yaxis.offsetText.set_fontsize(12)
    
    if showplot == True:
        plt.show()
    
    if savename:
        fig.savefig(savename, format=savename[-3:], dpi=300, bbox_inches='tight')

    return fig, ax 


def hist_scores(categories, scores, score_type='BLEU', model_font = 14, label_font = 30):
    """
    Create a bar chart with rotated labels on each bar.
    
    Parameters:
    -----------
    categories : list
        List of category names (e.g., model names)
    scores : list or array
        List of scores corresponding to each category
    score_type : str, optional
        Type of score being plotted (default: 'bleu')
    """
    values = np.array(scores)
    labels = categories
    
    # Create bar chart
    plt.figure(figsize=(20, 10))
    bars = plt.bar(categories, values)
    
    # Add rotated text labels on each bar
    for bar, label in zip(bars, labels):
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height/2,
                 label, ha='center', va='center',
                 fontweight='bold', color='white',
                 rotation=90, fontsize=model_font)
    
    # Remove x-axis labels
    plt.xticks([])
    
    plt.xlabel('Model', fontsize=label_font)
    plt.ylabel(score_type + ' score', fontsize=label_font) 
    plt.title(score_type + ' Scores', fontsize=label_font*(4/3))
    plt.show()

# Example usage:
# plot_scores(models, bleu_scores, 'bleu')


def nice_hist2(data, xlabel = 'xlabel', ylabel = 'Frequency', title = 'title', 
              color = 'red', nbins = 20, edgecolor='black', xlim = None,
              ylim = None, figsize = (6,5), savename = None, showplot = True,
              vlines = None, linewidth = 2, vline_fontsize = 12, legend_label = None):
    plt.rcParams.update({'font.size': 14})
    fig, ax = plt.subplots(figsize=figsize, dpi=300)
    fig.patch.set_facecolor('white')
    ax.set_facecolor('#f0f0f0')
    ax.hist(data, 
        bins=nbins, 
        color=color, 
        edgecolor=edgecolor)
    
    # Add vertical lines
    if vlines:
        for label, val in vlines.items():
            # Handle both formats: value or (value, color)
            if isinstance(val, tuple):
                xval, line_color = val
            else:
                xval = val
                line_color = 'black'
            
            ax.axvline(xval, color=line_color, linestyle='--', linewidth=linewidth)
            
            # Add text with arrow - positioned higher and further left
            x_range = ax.get_xlim()[1] - ax.get_xlim()[0]
            y_range = ax.get_ylim()[1] - ax.get_ylim()[0]
            
            ax.annotate(label, 
                       xy=(xval, ax.get_ylim()[1]*0.85),  # Arrow points here
                       xytext=(xval - x_range*0.15, ax.get_ylim()[1]*0.92),  # Text positioned here
                       fontsize=vline_fontsize, 
                       color=line_color,
                       ha='center',
                       va='center',
                       bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor=line_color, alpha=0.8),
                       arrowprops=dict(arrowstyle='->', color=line_color, lw=1.5, shrinkA=0, shrinkB=0))
    
    # Add legend annotation if provided
    if legend_label:
        x_range = ax.get_xlim()[1] - ax.get_xlim()[0]
        x_mid = (ax.get_xlim()[0] + ax.get_xlim()[1]) / 2
        
        ax.annotate(legend_label, 
                   xy=(x_mid, ax.get_ylim()[1]*0.5),  # Arrow points to middle of histogram
                   xytext=(x_mid - x_range*0.2, ax.get_ylim()[1]*0.7),  # Text positioned to the left
                   fontsize=vline_fontsize, 
                   color='black',
                   ha='center',
                   va='center',
                   bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='black', alpha=0.8),
                   arrowprops=dict(arrowstyle='->', color='black', lw=1.5, shrinkA=0, shrinkB=0))
    
    ax.set_title(title, fontsize=18, weight='bold', pad=10)
    ax.set_xlabel(xlabel, fontsize=16, labelpad=12)
    ax.set_ylabel(ylabel, fontsize=16, labelpad=12)
    ax.tick_params(axis='both', labelsize=12)
    ax.grid(True, color='white', linestyle='-', linewidth=1)
    ax.set_axisbelow(True)
    ax.margins(x=0.05, y=0.1)
    plt.tight_layout(pad=2.0)
    
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    
    ax.yaxis.set_major_formatter(ScalarFormatter(useMathText=True))
    ax.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
    ax.yaxis.offsetText.set_fontsize(12)
    
    if showplot == True:
        plt.show()
    
    if savename:
        fig.savefig(savename, format=savename[-3:], dpi=300, bbox_inches='tight')
    return fig, ax


def plot_heatmap(matrix, row_labels, cmap='viridis'):
    plt.imshow(matrix, cmap=cmap, aspect='auto')
    plt.colorbar()
    plt.yticks(range(len(row_labels)), row_labels)
    plt.title('Matrix Heatmap')
    plt.show()
