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'][:-2]
colors=np.array(['#023047ff','#126782ff','#219ebcff','#43b4baff','#fda10dff','#fb8500ff','#db6202ff','#bb3e03ff','#ae2012ff','#941b10ff'])[[1,2,6,7,8]]
longdash = (0, (1, 2, 1, 2, 4, 2))
def plot_vn_eta_alice(paths=[],exp_path_alice='./', lines=[],labels=[]):
    plt.figure(figsize=(14,8));plt.subplots_adjust(wspace=0,hspace=0);rows=2;cols=round(len(cents)/2)

    for i,cent in enumerate(cents[:]):
        plt.subplot(rows,cols,i+1)
        for j,path in enumerate(paths):
            dat=np.loadtxt(os.path.join(path,cent,'v3_eta_charged_cumulant_etaref_-0.8_0.8.dat'))
            if lines[j]!='longdash':
                plt.plot(dat[:,0],dat[:,1],color=colors[j],linestyle=lines[j],label=labels[j] if i==0 else '')
                plt.fill_between(dat[:,0],dat[:,1]-dat[:,2],dat[:,1]+dat[:,2],alpha=0.3,color=colors[j],linestyle=lines[j])
            else:
                plt.plot(dat[:,0],dat[:,1],color=colors[j],linestyle=longdash,label=labels[j] if i==3 else '')
                plt.fill_between(dat[:,0],dat[:,1]-dat[:,2],dat[:,1]+dat[:,2],alpha=0.3,color=colors[j],linestyle=longdash)
                 
        file=os.path.join(exp_path_alice,'%s.csv'%cent)
        dat=np.loadtxt(file,delimiter=',')
        plt.errorbar(dat[34:51,0],dat[34:51,1],yerr=[-dat[34:51,3],dat[34:51,2]],fmt='*',label=cent.replace('_','-')+'% ',color=colors[0])
        plt.bar(dat[34:51,0], height=-dat[34:51,5]+dat[34:51,4], bottom=dat[34:51,1]+dat[34:51,5], width=0.05, align='center', facecolor='none', edgecolor=colors[0])

        plt.xlim(-5,5);plt.ylim(0,0.05)
        if i>2: 
            plt.xlabel(r'$\eta$',loc='center')
            plt.xticks([-4,-2,0,2,4])
        else:
            plt.xticks([-4,-2,0,2,4],['','','','',''])
        if i==0 or i==3: 
            plt.ylabel(r'$v_3\{2\}$',loc='center')
            plt.yticks([0,0.02,0.04],[0,0.02,0.04])
        else:
            plt.yticks([0,0.02,0.04],['','',''])
        plt.legend(ncol=1,fontsize=15,loc='upper left')
        if i==1:
            plt.text(2,0.043, r'$p_T$>0 GeV',fontsize=15)
            plt.text(2,0.038, r'$|\eta^{\rm ref}|<0.8$',fontsize=15)
    plt.savefig('./v3_vs_eta_Alice.png',dpi=400,bbox_inches='tight')
    plt.savefig('./v3_vs_eta_Alice.pdf',dpi=400,bbox_inches='tight')
    plt.close()

exp_path_alice='../exp_dat/zj/pbpb2760/vn_vs_eta_Alice'
p1='./dat_w_err/nucleon/PbPb2760/'
p2='./dat_w_err/nucleon_fluct/PbPb2760/'
p3='./dat_w_err/hotspots/PbPb2760/'
p4='./dat_w_err/hotspots_fluct/PbPb2760/'
p5='./dat_w_err/hotspots_fluctmore/PbPb2760/'
labels=['nucleon', r'nucleon fluct($\sigma$=0.637)', 'hotspots', r'hotspots fluct($\sigma$=0.637)', r'hotspots fluct($\sigma$=1.2)']
lines=['solid', 'dashed', 'dotted', 'dashdot', 'longdash']

plot_vn_eta_alice(paths=[p1,p2,p3,p4,p5],exp_path_alice=exp_path_alice,lines=lines,labels=labels)