import matplotlib as mpl;mpl.use('Agg');import sys;from multiprocessing import Pool;import h5py as h5
import numpy as np;import pandas as pd;import matplotlib.pyplot as plt;import os
import time;from subprocess import call;import mplhep as hep;hep.styles.use('ATLAS')

cents=['0_5','5_10','10_20','20_30','30_40','40_50','50_60','60_70'][3:4]
colors=np.array(['#023047ff','#126782ff','#219ebcff','#43b4baff','#fda10dff','#fb8500ff','#db6202ff','#bb3e03ff','#ae2012ff','#941b10ff'])[[1,5,8]][:]
longdash = (0, (1, 2, 1, 2, 4, 2))

def PLOT(spe='', pid='',name='',marker='',scale=1):
    '''plot ptspectrum for specicy (pid, name)
    '''
    def plot_ptspec(hydro_path='./',exp_path='./',linestyle='solid',plot_exp=False,color='',scale=1):
        for i,cent in enumerate(cents[:]):
            if linestyle!='longdash':
                dat=np.loadtxt(os.path.join(hydro_path,cent,'dN_over_2pidYptdpt_mc_%s.dat'%pid))
                plt.plot(dat[:,0],dat[:,1]*scale,color=color,linestyle=linestyle)
            else:
                dat=np.loadtxt(os.path.join(hydro_path,cent,'dN_over_2pidYptdpt_mc_%s.dat'%pid))
                plt.plot(dat[:,0],dat[:,1]*scale,color=color,linestyle=longdash)

    hydro_path='./dat_w_err/nucleon/PbPb2760/'
    plot_ptspec(hydro_path=hydro_path,linestyle='solid', color=colors[0],scale=scale)
    hydro_path='./dat_w_err/nucleon_freezeout_0d4/PbPb2760/'
    plot_ptspec(hydro_path=hydro_path,linestyle='dashed',color=colors[1],scale=scale)
    hydro_path='./dat_w_err/nucleon_visc_1d5/PbPb2760/'
    plot_ptspec(hydro_path=hydro_path,linestyle='dotted',color=colors[2],scale=scale)
    
    dat=np.loadtxt(os.path.join(exp_path,'dNdPt_pbpb2760_%s_%s_exp.dat'%(cents[0],spe)),comments='#')
    plt.scatter(dat[:,0],dat[:,3]*scale,label=r'$%s$(*%s)'%(name,scale),marker=marker,edgecolor='k',facecolor='none')

plt.figure(figsize=(10,8))
exp_path='../exp_dat/data/dNdYptdpt_Alice'
scales=[10,1,0.1]
spes=['pion','kaon','proton'];pids=['211','321','2212'];names=['\pi^+','K^+','p'];markers=['*','s','o']
for k in range(3):
    PLOT(spes[k],pids[k],names[k],markers[k],scales[k])

plt.plot([],[],color=colors[0],linestyle='solid', label=r'$\varepsilon_{\rm frz}\approx$0.27 GeV/fm$^3$, $(\eta/s, \zeta/s)_{\text{Duke}}$')
plt.plot([],[],color=colors[1],linestyle='dashed',label=r'$\varepsilon_{\rm frz}\approx$0.40 GeV/fm$^3$, $(\eta/s, \zeta/s)_{\text{Duke}}$')
plt.plot([],[],color=colors[2],linestyle='dotted',label=r'$\varepsilon_{\rm frz}\approx$0.27 GeV/fm$^3$, $(\eta/s, \zeta/s)_{\text{Duke}}$(*1.5)')
plt.xlim(0,3);plt.xlabel(r'$p_T$(GeV)',loc='center');plt.yscale('log');plt.ylim(1e-3,1e6)
plt.ylabel(r'$(1/N_{ev}) (1/(2 \pi p_T)) d^2 N/dp_T dy (GeV^{-2})$',loc='center')
plt.text(0.2,1e-2,r'$|y|<0.5$',fontsize=20)
plt.legend(ncol=2,loc='upper right',columnspacing=1,fontsize=20)
plt.savefig('./PRD_pTspec_10_20.png',dpi=400,bbox_inches='tight')
plt.savefig('./PRD_pTspec_10_20.pdf',dpi=400,bbox_inches='tight')