<gh_stars>0
import numpy as np
from pyjet import cluster, DTYPE_PTEPM
from pyjet.testdata import get_event
import matplotlib.pyplot as plt
from matplotlib.pyplot import cm
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.colors import LogNorm

import glob
import os
import csv
import pandas as pd
import imageio



# define eta and phi ranges and number of bins along each axis
eta_min, eta_max = -5., 5.
extent = eta_min, eta_max, -np.pi, np.pi
bins = 200

import sys
WorkingDir=sys.argv[1]+'/'
C = int(sys.argv[2])

PATH = os.getcwd()+'/'
###########################################################
SavingDir=WorkingDir+"Plots/Pyjets/%sLep/" % C
DataDir=WorkingDir+"Data/IPPP/%sLep/Results%sLep/" % (C,C)
Cols = ['Pt','eta','phi','mass']
###########################################################

ClassesNames = []
events = glob.glob(DataDir+'*_pyjet.txt')
for i in range(len(events)):
	ClassesNames.append((events[i].split('/')[-1].split('_')[0]))
print(ClassesNames)


def LoadData(filename):
	D = 0	
	event = []
	with open(DataDir+filename+'_pyjet.txt') as f:
		lis = [line.split() for line in f]
		for i, x in enumerate(lis):
			if len(x) > 0:
				line = list(map(float, x))
				event.append((line[0],line[1],line[2],line[3]))
			if len(x) == 0:
				break
			D+=1
	data = np.array(event,np.dtype([('pT', 'f8'), ('eta', 'f8'), ('phi', 'f8'), ('mass', 'f8')]))		
	return data

def plotjets3clusters(eventOG,filename):
	# create regular grid of ghosts
	eta_edges = np.linspace(eta_min, eta_max, bins + 1)
	phi_edges = np.linspace(-np.pi, np.pi, bins + 1)
	eta = np.linspace(eta_min, eta_max, bins + 1)[:-1] + (eta_max - eta_min) / (2 * bins)
	phi = np.linspace(-np.pi, np.pi, bins + 1)[:-1] + (np.pi / bins)
	X, Y = np.meshgrid(eta, phi)
	ghosts = np.zeros(eta.shape[0] * phi.shape[0], dtype=DTYPE_PTEPM)
	ghosts['pT'] = 1e-8
	ghosts['eta'] = X.ravel()
	ghosts['phi'] = Y.ravel()

	# add ghosts to the event
	event = np.concatenate([eventOG, ghosts], axis=0)

	fig = plt.figure(figsize=(20, 8))

	ax = None
	for p, label in zip((-1, 0, 1), (r'anti-$k_t$', 'Cam/Aachen', '$k_t$')):
		# cluster
		sequence = cluster(event, R=1.0, p=p)
		sequencesmall = cluster(event, R=0.4, p=p)
		jets = sequence.inclusive_jets(ptmin=230)
		jetssmall = sequencesmall.inclusive_jets(ptmin=60) #May need to adjust

		# pick a different color for each jet
		colors = cm.seismic(np.linspace(0, 0.4, len(jets)))
		colorssmall = cm.seismic(np.linspace(0.6, 1.0, len(jetssmall)))
		cmap = LinearSegmentedColormap.from_list('cmap', colors, len(colors))
		cmapsmall = LinearSegmentedColormap.from_list('cmapsmall', colorssmall, len(colorssmall))
        
		ax = fig.add_subplot(1, 3, p + 2, sharey=ax)
		area = np.zeros((eta_edges.shape[0] - 1, phi_edges.shape[0] - 1),
                        dtype=np.float64)
		for ijet, jet in enumerate(jets):
			constit = jet.constituents_array()
			jetarea, _, _ = np.histogram2d(constit['eta'], constit['phi'],
                                           bins=(eta_edges, phi_edges))
			area += (jetarea > 0) * (ijet + 1)    

		# plot all areas
		ax.imshow(np.ma.masked_where(area == 0, area).T, cmap=cmap,
			extent=extent, aspect=(eta_max - eta_min) / (2*np.pi),
			interpolation='none', origin='lower')
        
		area = np.zeros((eta_edges.shape[0] - 1, phi_edges.shape[0] - 1),
                        dtype=np.float64)
		#Repeat for subjets
		for ijet, jet in enumerate(jetssmall):
			constit = jet.constituents_array()
			jetarea, _, _ = np.histogram2d(constit['eta'], constit['phi'],
                                           bins=(eta_edges, phi_edges))
			area += (jetarea > 0) * (ijet + 1)    

		# plot all areas
		ax.imshow(np.ma.masked_where(area == 0, area).T, cmap=cmapsmall,
			extent=extent, aspect=(eta_max - eta_min) / (2*np.pi),
			interpolation='none', origin='lower')

		# overlay original event
		particles = eventOG
		parts = ax.scatter(particles['eta'], particles['phi'], s=50,
			c=particles['pT'],cmap='Oranges',
			norm=LogNorm(),edgecolor='black',vmin = particles['pT'].min(),
			vmax=particles['pT'].max())
        
		ax.set_xlim(extent[:2])
		ax.set_ylim(extent[2:])
		if p == -1:
			ax.set_ylabel(r'$\phi$',fontsize=20)
			ax.set_xlabel(r'$\eta$',fontsize=20)
		if p != -1:
			ax.set_xlabel(r'$\eta$',fontsize=20)
		if p == 0:
			ax.set_title(filename)


		ax.text(0.95, 0.05, label,
			verticalalignment='bottom', horizontalalignment='right',
			transform=ax.transAxes,
			fontsize=20)
        
		ticks=[1e-1, 1, 1e1]
		cbar = plt.colorbar(parts)
		#cbar.ax.set_yticklabels([r'$10^{-1}$', r'$10^{0}$', r'$10^{1}$'])  # vertically oriented colorbar
		cbar.ax.set_ylabel(r'$p_T$ / GeV',fontsize=20)

	fig.subplots_adjust(hspace=0)
    
	plt.setp([a.get_yticklabels() for a in fig.axes[1:]], visible=True,fontsize=15)
	fig.tight_layout()
	fig.savefig(SavingDir+'%s.png' % filename)
	plt.close()

for i in range(len(ClassesNames)):
	filename = ClassesNames[i]
	print(filename)
	D = LoadData(filename)
	plotjets3clusters(D,filename)
