#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import numpy as np
import numpy.random as npr
import matplotlib.pyplot as mpl
import matplotlib.ticker as mpt
import sys


# In[ ]:


#==============================================================================
# Create Mock Data
#==============================================================================

def generate_correlated_random_xydata(npoints=2500):
    x = np.linspace(0,3,npoints)+npr.normal(size=npoints)
    y = x+npr.normal(size=npoints)
    return x,y


# In[ ]:


#==============================================================================
#     Utility Functions
#==============================================================================

def find_level(histogram,level,nslices=1000):
    """
    This return the histogram value that contain level% of the data inside.
    Useful to construct isocontours for 2D data based on data counts.
    """
    HTOT = np.sum(histogram)
    HMAX = np.amax(histogram)
    HMIN = np.amin(histogram)
    slices = np.linspace(HMAX,HMIN,nslices)
    for i in range(nslices):
        if np.sum(histogram[histogram>slices[i]]) > level*HTOT:
            break
        else:
            continue
    return slices[i]


# In[ ]:


#==============================================================================
#  Double Hist Figure
#==============================================================================


histCOLOR = 'GoldenRod'
histCMAP = 'YlOrRd'
contourCMAP = 'YlGnBu_r'
pointColor = 'LimeGreen'

fig = mpl.figure(figsize=(12.5,14)) ## Create an empty figure, figsize allows to control extra space for added axes


# In[ ]:


## Add axes to the figure (where you will plot the data)
## [x_lowerCorner,y_lowerCorner,width,height]
## The full figure canvas goes from [0,1] both in x and y
## The use of sharex=axname will make the two axes linked in matplotlib interactive mode
## The same can be done with sharey=axname
axMain = fig.add_axes([0.1,0.1,0.65,0.65]) 
axHistX = fig.add_axes([0.1,0.75,0.65,0.20],sharex=axMain) 
axHistY = fig.add_axes([0.75,0.1,0.20,0.65],sharey=axMain) 

## Adjust label parameters for an axis to show (or not) the number labels
axHistX.tick_params(labelbottom='off')
axHistY.tick_params(labelleft='off')

## Generate random set of data
varx,vary = generate_correlated_random_xydata() 

## Create a 2D histogram from data
binx = np.linspace(-3,7,21)
biny= np.linspace(-3,7,21)
Histogram2D,Xedges,Yedges = np.histogram2d(varx,vary,bins=[binx,biny]) 
## Now we want to mask out all histogram bins where we have less than 10 data poin ts
masked_Histogram=np.ma.masked_where(Histogram2D<=10,Histogram2D,copy=False) 


## Compute levels to draw contours at 90,60,40 and 20% levels
levels = [0.90,0.60,0.40,0.20]
levels_contour = np.array([find_level(Histogram2D,l) for l in levels])

## Plot both histograms in the variables x and y
hx,ex,px=axHistX.hist(varx,bins=binx,color=histCOLOR,histtype='stepfilled')
hy,ey,py=axHistY.hist(vary,bins=biny,orientation='horizontal',color=histCOLOR,histtype='stepfilled')

## Adjust axis limits semi-automatically
axHistX.set_ylim(0,1.1*np.amax(hx))    
axHistY.set_xlim(0,1.1*np.amax(hy))



# In[ ]:


## Plot data points, 2D histogram and contour levels on the main axes
## the picker keyword allows for added interactive actions to your plot by selecting
## closer point to your mouse click and have accessible data to play with.
## The zorder keyword allows you to control which elements of the plot apper on top.
## The higher the value of zorder the more upfront the element will be. 
## This is all done in relative terms.
z=0
axMain.plot(varx,vary,'.',alpha=0.5,color=pointColor,zorder=z,picker=2) 
axMain.imshow(masked_Histogram.T,cmap=histCMAP,extent=(min(Xedges),max(Xedges),min(Yedges),max(Yedges)),aspect='auto',zorder=z+1,origin="lower",interpolation="nearest")
axMain.contour(Histogram2D.T,levels=levels_contour,extent=(min(Xedges),max(Xedges),min(Yedges),max(Yedges)),cmap=contourCMAP,origin='lower',zorder=z+10)

## This limits the number of main ticks on each axis which helps in reducing
## the confusing in cluttered axes
axHistY.xaxis.set_major_locator(mpt.MaxNLocator(3))
axHistX.yaxis.set_major_locator(mpt.MaxNLocator(3))
axMain.xaxis.set_major_locator(mpt.MaxNLocator(5))

## This is not useful here, but can be useful in future scenarios since we use loglog plots often
## This will make your axis have the number labels in scalar format and not in power law format
axMain.xaxis.set_major_formatter(mpt.ScalarFormatter(useOffset=False))
axMain.yaxis.set_major_formatter(mpt.ScalarFormatter(useOffset=False))

## Turn on the minorticks in the plot to look more 'scientific'. Helpful to read values off plots in an easier way
for eixo in  [axMain,axHistX,axHistY]:
    eixo.minorticks_on()

    
## Set your plot labels using mathtext, same as you would do in LaTex
axMain.set_xlabel(r'$var_x$')
axMain.set_ylabel(r'$var_y$')

## Connect your figure with keyboard action. In this case, pressing a key while 
## the matplotlib window is active will execute the function exit_code
## This in turn wil either close the plot or exit the program if you click
## 'q' or 'escape' respectively
def exit_code(event):
    if event.key=='escape':
        sys.exit()
    if event.key=='q':
        mpl.close('all')
fig.canvas.mpl_connect('key_press_event',exit_code)

## Connect yout figure to a picking event. This will allow you to select the 
## closest element (data point, line, etc.) and allow you to access the data
## This is good for exploratory plots where you need specific information to
## add to that shown in the plot
def on_pick(event):
    thisline = event.artist
    xdata, ydata = thisline.get_data()
    ind = event.ind
    print('You clicked at x=%.2f, y=%.2f'%(xdata[ind], ydata[ind]))
cid = fig.canvas.mpl_connect('pick_event', on_pick)

## Make the Figure pop up
mpl.show()


# In[ ]:





# In[ ]: