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

# import matplotlib.pyplot as plt   # our plotting module
import matplotlib as mpl
import pandas as pd               # the main HDF5 reader
import numpy as np                # must have
#import km3pipe as kp              # some KM3NeT related helper functions
import seaborn as sns              # beautiful statistical plots!
import sys, itertools, aa, collections
from ROOT import EventFile, Det, Timer
import glob
import os, sys, glob, numpy, matplotlib, ROOT
from matplotlib import pyplot
import csv, math
import matplotlib.pylab as plt
from numpy import sin, cos, pi, linspace
from numpy.random import randn
from scipy.signal import lfilter, lfilter_zi, filtfilt, butter

from matplotlib.pyplot import plot, legend, show, grid, figure, savefig
plt.rcParams['figure.figsize'] = 14, 10
from datetime import datetime, timedelta, time


EventFile.read_timeslices = True
EventFile.timeslice_treename = "KM3NET_TIMESLICE_SN"



# files on cc-lyon:
#filename="/sps/km3net/users/gdewasse/L0-files/KM3NeT_00000029_00002984.root"
#detxname="/pbs/throng/km3net/detectors/KM3NeT_00000029_20170920.detx"

# files on marantares2:
filename="/baie/huang/data/KM3NeT_00000049_00009050.root"
detxname="/baie/huang/data/KM3NeT_00000049_20200210.detx"
#filename="/baie/huang/data/L0-files/KM3NeT_00000029_00002984.root"
#detxname="/baie/huang/data/L0-files/KM3NeT_00000029_20170920.detx"

det = Det(detxname)
f = EventFile(filename) #3094 - 2984 #Feel free to add more runs!

df_pmt_channel_id = pd.DataFrame( [['A', 'A1', 22],
         ['B', 'B1', 14], ['B', 'B2', 19], ['B', 'B3', 25], ['B', 'B4', 24], ['B', 'B5', 26], ['B', 'B6', 18],
         ['C', 'C1', 13], ['C', 'C2', 21], ['C', 'C3', 29], ['C', 'C4', 28], ['C', 'C5', 20], ['C', 'C6', 17],
         ['D', 'D1', 12], ['D', 'D2', 15], ['D', 'D3', 23], ['D', 'D4', 30], ['D', 'D5', 27], ['D', 'D6', 16],
         ['E', 'E1', 10], ['E', 'E2', 6 ], ['E', 'E3', 3 ], ['E', 'E4', 2 ], ['E', 'E5', 1 ], ['E', 'E6', 11],
         ['F', 'F1', 9 ], ['F', 'F2', 8 ], ['F', 'F3', 4 ], ['F', 'F4', 0 ], ['F', 'F5', 5 ], ['F', 'F6', 7 ]],
         columns=('ring', 'PMT_position', 'channel_id'))

dict_pmt_channel_id = {22: 'A1' , 14:'B1',  19:'B2',  25:'B3', 24:'B4',  26:'B5', 18:'B6',
                                  13:'C1',  21:'C2',  29:'C3', 28:'C4',  20:'C5', 17:'C6',
                                  12:'D1',  15:'D2',  23:'D3', 30:'D4',  27:'D5', 16:'D6',
                                  10:'E1',  6 :'E2',  3 :'E3', 2 :'E4',  1 :'E5', 11:'E6',
                                  9 :'F1',  8 :'F2',  4 :'F3', 0 :'F4',  5 :'F5', 7 :'F6'}


R = []
R2 = []
R3 = []
Rsum = [] 
TIME = [] 
RR = []
RR2 = []
RR3 = []
RSUM = []
T = [] 
DATE = []
i = 0 
select_dom_id = 808982077   #if dom_id == 808982077 : #808432835:809544061:808982077:808997793  #some DOM IDs
print("events in file" , f.size())
print("events in file" , f.index)
for evt in f:
    if i == 0 :
        #print('time', evt.t.AsString())
        time_first_event = evt.t.GetSec() + 1e-9 * evt.t.GetNanoSec()
        
    i+=1
    if f.index > 50000: break
    f.get_rates( det ) 
    event_time = evt.t.GetSec() + 1e-9 * evt.t.GetNanoSec()
    # now, det knows the rate on each PMT
    #for dom_id, dom in det.doms :
    #    rates = [ pmt.rate for pmt in dom.pmts ]
    #    #print( dom_id , rates[:5] )
    dom = det.doms[select_dom_id]
    rates = [pmt.rate for pmt in dom.pmts]
    dict_pmt_channelID_rates = [pmt.channel_id:pmt.rate for pmt in dom.pmts ]
    R.append(rates[0])
    R2.append(rates[3])
    R3.append(rates[6])
    #print(event_time)
    TIME.append(event_time - time_first_event)
    Rsum.append(sum(rates))
    if max(TIME) - min(TIME) > 120: 
        DATE.append(evt.t.AsString())
        T.append(TIME)
        RSUM.append(Rsum)
        RR.append(R)
        RR2.append(R2)
        RR3.append(R3)
        R = []
        R2 = []
        R3 = []
        Rsum = []
        TIME = []
        time_first_event = evt.t.GetSec() + 1e-9 * evt.t.GetNanoSec()
    #print("\n\n\n")


print("TIME", TIME)
print("T", T)
print("R", R)
i = 0 
while i < len(RR):
    
    fig = plt.figure()
    
    plt.plot(T[i], RSUM[i], linestyle = 'None', marker = 'o')
    #v = [min(TIME), max(TIME), 190e3, 350e3]#220000, 310000]
    #plt.axis(v)
    plt.xlabel(r'Time [s]', size = 20)
    plt.ylabel(r'Rate [kHZ]', size = 20)
    plt.xticks([0,10,20,30, 40,50,60, 70,80, 90,100,110, 120], size=20)
    #plt.yticks([190e3,200e3,210e3, 220e3,230e3,240e3,250e3,260e3,270e3,280e3,290e3,300e3,310e3,320e3,330e3,340e3,350e3],['190', '200', '210','220','230','240','250','260','270','280','290','300','310', '320', '330', '340','350'],size=20)
    plt.grid()
    plt.savefig(("plots/"+str(DATE[i])+".png"))
    plt.savefig(("plots/"+str(DATE[i])+".pdf"))
    #fig.save()
    #plt.show()
    
    i+=1
    #plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
    #           ncol=2,fontsize = 15, mode="expand", borderaxespad=0.)


INFO_T = []
i = 0 
while i < len(T):
    j = 0 
    info_T = []
    while j < len(T[i]):
        info_T.append([T[i][j], RSUM[i][j], DATE[i]])
        j+=1
    INFO_T.append(info_T)
    i+=1

i = 0 
info_sorted = []
while i < len(INFO_T):
    info_sorted.append(sorted(INFO_T[i], key=lambda u: u[0], reverse=False))
    i+=1


i = 0
f = open('REINFORCE_DATA.txt', 'w')

#while i < 2:#len(info_sorted):
while i < len(info_sorted):
    j = 0
    f.write('New event \n')
    f.write(info_sorted[i][0][2])
    while j < len(info_sorted[i]):
        #print(info_sorted[i][j][0],info_sorted[i][j][1])
        L = [str(info_sorted[i][j][0]), ' ', str(info_sorted[i][j][1]), '\n']
        f.writelines(L) 
        j+=1
    i+=1
f.close()