# This code is designed to reproduce Fig. S7 from 2014 Vigolo

import numpy as np
import h5py
import matplotlib.pyplot as plt

L = 0.0048
z_plane_of_symmetry = -0.0024
n_intervals = 30
max_creation_time = 1.5
fluid_run_name = 'mesh_38288_nodes'
fluid_run_name = 'mesh_66576_nodes'
snapshot_name = 't=0.2900000000000001_in_box'
snapshot_name = 't=1.0000000000000007_in_box'
snapshot_name = 't=2.0000000000000013_RADIUS=0.00024_in_box'
snapshot_name = 't=1.0_RADIUS=0.0002_in_box'
snapshot_name = 't=0.11_RADIUS=0.0002_in_box'
snapshot_name = 't=1.99_RADIUS=0.0002_in_box'
a = 0.00024 / L
# snapshot_name = 't=0.04_in_box'

def IntervalIndex(x, bottom, top, number_of_intervals):
    if x < bottom:
        x = bottom
    elif x > top:
        x = top
        # raise ValueError('The coordinate z = ' + str(x) + ' is not within range!')
    return int(x / (top - bottom) * number_of_intervals)

def RobustQuotient(a, b):
    if a == b:
        return 1
    else:
        return a / b

def TooNew(creation_time):
    return creation_time > max_creation_time

def CalculateProbabilityOfGettingTrapped():
    half_width = abs(z_plane_of_symmetry)
    z_init_intervals = [[0.5 * i / n_intervals, 0.5 * (i + 1) / n_intervals] for i in range(n_intervals)]
    z_midpoints = [0.5 * (I[1] + I[0]) for I in z_init_intervals]
    histogram = np.zeros(len(z_init_intervals))
    trapped_hist = np.zeros(len(z_init_intervals))

    with h5py.File('all_particles.hdf5', 'r') as f:
        Ids = f['/Id'][:]
        Z0s = f['/Z0'][:]
        creation_times = f['/TIME'][:]
        too_new_to_include = [TooNew(time) for time in creation_times]
        ids_too_new_to_include = dict(zip(Ids, too_new_to_include))
        zs = [abs(z - z_plane_of_symmetry) / L for z in Z0s]
        for z in Z0s:
            d = abs(z - z_plane_of_symmetry) / L
            histogram[IntervalIndex(d, 0., 0.5, n_intervals)] += 1

    with h5py.File('particles_snapshot.hdf5', 'r') as f:
        Ids_trapped = f['/' + fluid_run_name + '.hdf5/' + snapshot_name + '/Id'][:]
        Z0s_trapped = f['/' + fluid_run_name + '.hdf5/' + snapshot_name + '/Z0'][:]

        min_d =   float('inf')
        max_d = - float('inf')
        zs_trapped = [abs(z - z_plane_of_symmetry) / L for Id, z in zip(Ids_trapped, Z0s_trapped) if not ids_too_new_to_include[Id]]
        for Id, z in zip(Ids_trapped, Z0s_trapped):
            if not ids_too_new_to_include[Id]:
                d = abs(z - z_plane_of_symmetry) / L
                min_d = min(min_d, d)
                max_d = max(max_d, d)
                trapped_hist[IntervalIndex(d , 0., 0.5, n_intervals)] += 1

    histogram = [RobustQuotient(h, H) for h, H in zip(trapped_hist, histogram)]
    plt.plot(z_midpoints, histogram, '.-')
    axes = plt.gca()
    axes.set_xlim([min_d, max_d])
    axes.set_ylim([0,1])
    plt.show()

CalculateProbabilityOfGettingTrapped()
