import numpy as np
import os

"""
Explanation: Data Extraction
Joeri R. Hermans                  
Departement of Data Science & Knowledge Engineering        
Maastricht University, The Netherlands        
In this notebook we mainly consider the extraction of the data from numpy files, and preprocess them to a more desireable format.
End of explanation
"""


data_directory = "../data/"
track_types = os.listdir(data_directory)

# List the track types.
for e in track_types:
    print(e)

columns = []

# The noise event needs to be considered at a later time.
track_types.remove('noise')

# Fetch the columns names from the events.
for e in track_types:
    track_files = os.listdir(data_directory + e)
    track_file = track_files[0]
    data = np.load(data_directory + e + "/" + track_file)
    columns.append(data.dtype.names)

# Check if the columns have the same attributes.
num_columns = len(columns)
c = columns[0]
equal = True
for i in range(1, num_columns):
    equal &= (c == columns[i])
    if not equal:
        break
        
print("Equal attributes among all events: " + str(equal))

"""
Explanation: Extraction
End of explanation
"""


simulated_data = [t for t in track_types if 'SIM' in t]
real_data = [t for t in track_types if 'SIM' not in t]

"""
Explanation: One might have noticed that the data that was provided to us contains simulated data. In order to have a clean comparison of simulated and real data, we maintain two lists.
End of explanation
"""


schema = {
    "namespace": "trackml.cern.ch",
    "type": "record",
    "name": "Tracks",
    "fields": [
        {"name": "track_id", "type": "int"},
        {"name": "collision_id", "type": "string"}
        {"name": "run", "type": "int"},
        {"name": "event", "type": "int"},
        {"name": "luminosity", "type": "int"},
        {"name": "track_type", "type": "string"},
        {"name": "charge", "type": "int"},
        {"name": "ndof", "type": "int"},
        {"name": "qoverp", "type": "double"},
        {"name": "theta", "type": "double"},
        {"name": "dxy", "type": "double"},
        {"name": "d0", "type": "double"},
        {"name": "dsz", "type": "double"},
        {"name": "dz", "type": "double"},
        {"name": "p", "type": "double"},
        {"name": "pt", "type": "double"},
        {"name": "px", "type": "double"},
        {"name": "py", "type": "double"},
        {"name": "pz", "type": "double"},
        {"name": "eta", "type": "double"},
        {"name": "phi", "type": "double"},
        {"name": "vx", "type": "double"},
        {"name": "vy", "type": "double"},
        {"name": "vz", "type": "double"},
        {"name": "num_background_hits", "type": "int"},
        {"name": "num_track_hits", "type": "int"},
        # Define the structure of the background hits.
        {"name": "background_hits", "type": {
            "type": "array",
            "items": {
                "type": "record",
                "name": "background_hit",
                "fields": [
                    {"name": "x", "type": "double"},
                    {"name": "y", "type": "double"},
                    {"name": "z", "type": "double"}
                ]
            }
        }},
        # Define the structure of the track hits.
        {"name": "track_hits", "type": {
            "type": "array",
            "items": {
                "type": "record",
                "name": "track_hit",
                "fields": [
                    {"name": "x", "type": "double"},
                    {"name": "y", "type": "double"},
                    {"name": "z", "type": "double"}
                ]
            }
        }}
    ]
}

"""
Explanation: Schema definition
We store the data in the Avro format, which is a binary serialized row-based format. However, before we do that, we first need to specify a schema in order to describe the contents of the file-format.
Note: Collisions are grouped by: (run, event, luminosity)
End of explanation
"""


def extract_hits(record):
    hits = []
    
    # Define the maximum number of hits (due to the given format).
    max_silicon_pixel_hits = 5
    max_silicon_strip_hits = 50
    
    # Extract hits from the silicon pixel detector.
    for i in range(0, max_silicon_pixel_hits):
        index = str(i)
        x = record['pix_' + index + '_x']
        y = record['pix_' + index + '_y']
        z = record['pix_' + index + '_z']
        if x == 0.0 and y == 0.0 and z == 0.0:
            break
        # Append the hit.
        hits.append({'x': x, 'y': y, 'z': z})
    # Extract hits from the silicon strip detector.
    for i in range(0, max_silicon_strip_hits):
        index = str(i)
        x = record['sis_' + index + '_x']
        y = record['sis_' + index + '_y']
        z = record['sis_' + index + '_z']
        if x == 0.0 and y == 0.0 and z == 0.0:
            break
        # Append the hit.
        hits.append({'x': x, 'y': y, 'z': z})
    
    return hits

def create_record(record, track_type):
    r = {}
    
    # Extract the detector hits.
    track_hits = extract_hits(record)
    background_hits = [] # No background hits
    # Extrac track parameters.
    r['track_id']    = int(record['TrackId'])
    r['run']         = int(record['run'])
    r['event']       = int(record['evt'])
    r['luminosity']  = int(record['lumi'])
    r['collision_id'] = str(r['run']) + "-" + str(r['event']) + "-" + str(r['luminosity'])
    r['track_type']  = track_type
    r['charge']      = int(record['charge'])
    r['ndof']        = int(record['ndof'])     
    r['qoverp']      = record['qoverp']
    r['theta']       = record['theta']
    r['dxy']         = record['dxy']
    r['d0']          = record['d0']
    r['dsz']         = record['dsz']
    r['dz']          = record['dz']
    r['p']           = record['p']
    r['pt']          = record['pt']
    r['px']          = record['px']
    r['py']          = record['py']
    r['pz']          = record['pz']
    r['eta']         = record['eta']
    r['phi']         = record['phi']
    r['vx']          = record['vx']
    r['vy']          = record['vy']
    r['vz']          = record['vz']
    # Add the detector hits.
    r['background_hits'] = background_hits
    r['num_background_hits'] = len(background_hits)
    r['track_hits'] = track_hits
    r['num_track_hits'] = len(track_hits)
    
    return r

def read_path(path):
    # Initialize the records buffer.
    records_buffer = []
    # Obtain the record type from the path.
    t = path.split('/')[-2]
    # Load the data from the specified path.
    data = np.load(path)
    for r in data:
        # Add the record to the buffer.
        records_buffer.append(create_record(r, t))
        
    return records_buffer

def write_records(records):
    for record in records:
        writer.write(record)
    writer.flush()

"""
Explanation: Data Construction
Now the schema is available for all track types described above, we can write the tracks to an output file which can be used on the cluster.
Utility function definitions
End of explanation
"""


from fastavro.writer import Writer
from fastavro.reader import SYNC_SIZE

from multiprocessing import Pool
from multiprocessing import Lock

from contextlib import closing

# Obtain all file paths, this will benefit the parallelism.
paths = []

for t in track_types:
    path = data_directory + t + '/'
    files = os.listdir(path)
    for f in files:
        paths.append(path + f)

records_buffer = []
buffer_size = 0
max_buffer_size = 100

num_paths = len(paths)
path_index = 0

print("Processed: " + str(path_index) + "/" + str(num_paths))
with open('tracks.avro', 'wb') as out:
    # Allocate a fastavro writer.
    writer = Writer(out, schema, sync_interval=10000 * SYNC_SIZE)
    # Process all paths, and flush buffer periodically.
    for p in paths:
        path_records = read_path(p)
        records_buffer.extend(path_records)
        buffer_size += 1
        path_index += 1
        if buffer_size >= max_buffer_size:
            print("Processed: " + str(path_index) + "/" + str(num_paths))
            write_records(records_buffer)
            del records_buffer
            records_buffer = []
            buffer_size = 0
    # Check if we have a non-empy buffer.
    if buffer_size > 0:
        write_records(records_buffer)
        del records_buffer
        records_buffer = []
        buffer_size = 0

"""
Explanation: Inserting the tracks
End of explanation
"""


