from datetime import datetime
import tables
from sapphire import download_data

STATIONS = [501, 510, 599]
START = datetime(2016, 5, 1)
END = datetime(2016, 5, 1, 3)
FILENAME = 'coinc_with_s599.h5'

data = tables.open_file(FILENAME, 'a')

station_groups = ['/s%d' % station for station in STATIONS]

for station, group in zip(STATIONS, station_groups):
    print("Station %d in group: %s." % (station, group),)
    if group not in data:
        print("Downloading data: ")
        download_data(data, group, station, START, END)
    else:
        print("group %s already in datafile." % group)

print(data)

from sapphire import CoincidencesESD
coin = CoincidencesESD(data, '/coincidences', station_groups, overwrite=True)
coin.search_and_store_coincidences(station_numbers=STATIONS)

print(data)

coinc_tabel = data.root.coincidences.coincidences
print(coinc_tabel)

coincidences = coinc_tabel.read()
coincidences[:10]

s501 = coincidences['s501']
s510 = coincidences['s510']
s599 = coincidences['s599']

print("Er zijn %d events van station 501 betrokken bij coincidenties tussen 501, 510 en/of 599" % sum(s501))

print("Er zijn %d coincidenties tussen 501, 510 en 599" % sum(s510 & s501 & s599))

sum(s501 & s599 & ~s510)

coinc_510_501 = coincidences.compress(s501 & s510)
coinc_510_501

print("Er zijn %d coincenties tussen 510 en 511 waarbij ook 599 betrokken is" % sum(coinc_510_501['s599']))

coinc_510_501_zonder599 = coincidences.compress(s501 & s510 & ~s599)
coinc_510_501_zonder599