#This code reads information from a file or set of files and plots it into a table
#Made with help from: https://towardsdatascience.com/simple-little-tables-with-matplotlib-9780ef5d0bc4

import numpy as np
import matplotlib.pyplot as plt

#This will put the data in order from smallest to largest
def selection_sort(x):
    for i in range(len(x)):
        swap = i + np.argmin(x[i:])
        (x[i], x[swap]) = (x[swap], x[i])
    return x

#The variables whose significance we want to compare (This is something you can change if you want to set your table up differently)
variables = ['MET', 'METetaBTrkJ1', 'METetaBTrkJ1etaBTrkJ2', 'METetaBTrkJ1etaBTrkJ2sumPtJets',
            'METetaBTrkJ1etaBTrkJ2sumPtJetsnSigJet', 'METetaBTrkJ1sumPtJetsnSigJet', 
            'METsumPtJetsnSigJet']

#Set what you'll put in the table
data = []
for i in range(len(variables)):
	f = open('evaluate_' + variables[i] + '.txt', 'r') #You can change the name to whatever file you're reading from
	lines = f.readlines()
	r = lines[-4].split('\t')[1] #lines[-4] tells the code whhere in the file you want to look, this will change depending on the setup of the file. The .split('\t') splits the line wherever there's a tab between values (here, I only cared about the second value after the split).
	h = r.split('\n')[0] #This splits up the line wherever a new one starts
	data.append(float(h))
	f.close()
        
data_sort = np.array(data) #Sets a new array, so we can keep the original un-sorted data
selection_sort(data_sort) #Orders the data from smallest to largest

color_sort = ['yellow', 'greenyellow', 'lightgreen', 'springgreen', 'mediumseagreen', 
              'green','darkgreen'] #Here's where you'll set the colors you want to associate with the sorted data values

color = []
for i in range(len(color_sort)):
        color.append(0)

#Set the colors from color_sort in the order of the original data
for i in range(len(data)):
    for j in range(len(data_sort)):
        if data_sort[j] == data[i]:
            color[i] = color_sort[j]

#Here's where you'll set up the figure 
plt.figure(linewidth=2,
           edgecolor='slategrey',
           facecolor='white'
           #figsize=(5,3)
          )

#Set up and fill the table
table1 = plt.table(cellText=[['MET', 'MET\netaBTrkJ1', 'MET\netaBTrkJ1\netaBTrkJ2', 'MET\netaBTrkJ1\netaBTrkJ2\nsumPtJets',
            'MET\netaBTrkJ1\netaBTrkJ2\nsumPtJets\nnSigJet', 'MET\netaBTrkJ1\nsumPtJets\nnSigJet',
            'MET\nsumPtJets\nnSigJet'], #In this case, this array would be the first row of the table
                             [data[0], data[1], data[2], data[3], data[4], data[5], data[6]]], #This is the second row of the table
          cellColours=[['slategrey', 'slategrey', 'slategrey', 'slategrey','slategrey',
                       'slategrey', 'slategrey'], #The colors for the first row of the table
                     [color[0], color[1], color[2], color[3], color[4], color[5], color[6]]], #Colors for the second row of the table
         loc = 'center')
table1.scale(1.25, 8) #Scale the table so the information fits in the cells and they take up as much space as possible. This will change based on the size of your table

#Here we'll get rid of the axes for the plot (since we're using matplotlib, the table is attached to a plot, but we omly want the table)
ax = plt.gca()
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)

#This gets rid of the box for the plot
plt.box(on = None)
plt.suptitle('Significance')

#Sets the font size for in the cells
table1.auto_set_font_size(False)
table1.set_fontsize(15)

plt.draw()
fig = plt.gcf()
plt.savefig('table.jpg',
            edgecolor='slategrey',
            facecolor='white',
            dpi=150
            )