import numpy as np
import matplotlib.pyplot as plt
from matplotlib.transforms import Affine2D
from matplotlib.markers import MarkerStyle
from IPython.display import display, Markdown, Latex
from matplotlib.ticker import MultipleLocator
import matplotlib.colors as colors

def plot_slice(X, Y, Z, vmin=None, vmax=None, skew_angle=90, ax=None,
              xlim=None, ylim=None,
              xticks=None, yticks=None,
              cbar=True,
              logscale=False,
              symlogscale=False,
              cmap='viridis',
              linthresh = 1,
              title=None,
             ):
    
    """
    This function plots a 2D heatmap of a given dataset (Z) with the possibility of applying logarithmic or symmetrical logarithmic scales to the colorbar. The plot can be sheared to create a trapezoidal view with a defined skew angle. Tick labels are displayed on all sides of the plot. The function takes the following parameters:

    Parameters
    ----------
    
    `X`: The x-coordinates of the data.
    
    `Y`: The y-coordinates of the data.
    
    `Z`: The 2D dataset to plot.
    
    `vmin`: The minimum value to plot in the dataset. If not provided, the minimum of the dataset will be used.
    
    `vmax`: The maximum value to plot in the dataset. If not provided, the maximum of the dataset will be used.
    
    `skew_angle`: The angle to shear the plot in degrees. Defaults to 90 degrees (no skewing).
    
    `ax`: An optional axis object to plot the heatmap onto.
    
    `xlim`: The limits of the x-axis. If not provided, the limits will be automatically set.
    
    `ylim`: The limits of the y-axis. If not provided, the limits will be automatically set.
    
    `xticks`: The major tick interval for the x-axis. If not provided, the function will use a default minor tick interval of 1.
    
    `yticks`: The major tick interval for the y-axis. If not provided, the function will use a default minor tick interval of 1.
    
    `cbar`: Whether to include a colorbar in the plot. Defaults to True.
    
    `logscale`: Whether to use a logarithmic color scale. Defaults to False.
    
    `symlogscale`: Whether to use a symmetrical logarithmic color scale. Defaults to False.
    
    `cmap`: The color map to use. Defaults to 'viridis'.
    
    `linthresh`: The linear threshold for the symmetrical logarithmic color scale. Defaults to 1.
    
    `title`: A string containing the title for the plot. If not provided, a default title will be used.

    Returns
    -------
    `p`: `matplotlib.collections.QuadMesh` object, similar to `matplotlib.pyplot.pcolormesh`.

    """
    
    Z = Z.transpose()

    # Display Markdown heading
    if title == "None":
        pass
    elif title != None:
        display(Markdown('### Figure - '+title))
    else:
        display(Markdown('### Figure'))
        
    # Inherit axes if user provides some
    if ax != None:
        ax=ax
        fig=ax.get_figure()
    # Otherwise set up some default axes
    else:
        fig = plt.figure()
        ax = fig.add_axes([0,0,1,1])
    
    # If limits not provided, use extrema
    if vmin==None:
        vmin=Z.min()
    if vmax==None:
        vmax=Z.max()

    # Plot data
    if symlogscale:
        p = ax.pcolormesh(X, Y, Z,
                          # vmin=vmin, vmax=vmax,
                          shading='auto',
                          norm=colors.SymLogNorm(linthresh=linthresh,vmin=-vmax, vmax=vmax),
                          cmap=cmap
                         )
    elif logscale:
        p = ax.pcolormesh(X, Y, Z,
                          # vmin=vmin, vmax=vmax,
                          shading='auto',
                          norm=colors.LogNorm(vmin=vmin, vmax=vmax),
                          cmap=cmap
                         )
    else:
        p = ax.pcolormesh(X, Y, Z,
                          vmin=vmin, vmax=vmax,
                          shading='auto',
                          cmap=cmap
                         )
    
    ## Transform data to new coordinate system if necessary
    # Correct skew angle
    skew_angle = 90-skew_angle
    # Create blank 2D affine transformation
    t = Affine2D()
    # Scale y-axis to preserve norm while shearing 
    t += Affine2D().scale(1,np.cos(skew_angle*np.pi/180))
    # Shear along x-axis
    t += Affine2D().skew_deg(skew_angle,0)
    # Return to original y-axis scaling
    t += Affine2D().scale(1,np.cos(skew_angle*np.pi/180)).inverted()
    ## Correct for x-displacement after shearing
    # If ylims provided, use those
    if ylim != None:
        # Set ylims
        ax.set(ylim=ylim)
        ymin,ymax = ylim
    # Else, use current ylims
    else:
        ymin,ymax = ax.get_ylim()
    # Use ylims to calculate translation (necessary to display axes in correct position)
    p.set_transform(t + Affine2D().translate(-ymin*np.sin(skew_angle*np.pi/180),0) + ax.transData)
    
    # Set x limits
    if xlim != None:
        xmin,xmax = xlim
    else:
        xmin,xmax = ax.get_xlim()
    ax.set(xlim=(xmin,xmax+(ymax-ymin)/np.tan((90-skew_angle)*np.pi/180)))
    
    # Correct aspect ratio for the x/y axes after transformation
    ax.set(aspect=np.cos(skew_angle*np.pi/180)) 
    
    # Add tick marks all around
    ax.tick_params(direction='in', top=True, right=True, which='both')
    
    # Set tick locations
    if xticks==None:
        # Add default minor ticks
        ax.xaxis.set_minor_locator(MultipleLocator(1))
    else:
        # Otherwise use user provided values
        ax.xaxis.set_major_locator(MultipleLocator(xticks))
        ax.xaxis.set_minor_locator(MultipleLocator(1))
    if yticks==None:
        # Add default minor ticks
        ax.yaxis.set_minor_locator(MultipleLocator(1))
    else:
        # Otherwise use user provided values
        ax.yaxis.set_major_locator(MultipleLocator(yticks))
        ax.yaxis.set_minor_locator(MultipleLocator(1))
    
    ## Apply transform to tick marks
    for i in range(0,len(ax.xaxis.get_ticklines())):
        # Tick marker
        m = MarkerStyle(3)
        line =  ax.xaxis.get_majorticklines()[i]
        if i%2:
            # Top ticks (translation here makes their direction="in")
            m._transform.set(Affine2D().translate(0,-1) + Affine2D().skew_deg(skew_angle,0))
            # This first method shifts the top ticks horizontally to match the skew angle. This does not look good in all cases.
            # line.set_transform(Affine2D().translate((ymax-ymin)*np.sin(skew_angle*np.pi/180),0) + line.get_transform())
            # This second method skews the tick marks in place and can sometimes lead to them being misaligned.
            line.set_transform(line.get_transform()) # This does nothing
        else:
            # Bottom ticks
            m._transform.set(Affine2D().skew_deg(skew_angle,0))
        
        line.set_marker(m)
        
    for i in range(0,len(ax.xaxis.get_minorticklines())):
        m = MarkerStyle(2)
        line =  ax.xaxis.get_minorticklines()[i]
        if i%2:
            m._transform.set(Affine2D().translate(0,-1) + Affine2D().skew_deg(skew_angle,0))
            # This first method shifts the top ticks horizontally to match the skew angle. This does not look good in all cases.
            # line.set_transform(Affine2D().translate((ymax-ymin)*np.sin(skew_angle*np.pi/180),0) + line.get_transform())
            # This second method skews the tick marks in place and can sometimes lead to them being misaligned.
            line.set_transform(line.get_transform()) # This does nothing
        else:
            m._transform.set(Affine2D().skew_deg(skew_angle,0))
        
        line.set_marker(m)
    
    if cbar:
        fig.colorbar(p)

    # Return the quadmesh object
    return p