%matplotlib inline
import numpy as np
import matplotlib.pyplot as pyplot
from km3net.kernels import QuadraticDifferenceSparse, PurgingSparse
import km3net.util as util
window_width = 1500
N,x,y,z,ct = util.get_real_input_data("sample1.txt")
print ("Read", N, "hits from file")

"""
Explanation: Comparison of criteria for correlating hits
This notebook shows how to use the KM3Net package to compare two different criteria for correlating L0 hits. 
The quadratic difference criterion checks whether the distance between two hits in space is large than the distance in time times the speed of light. The Match 3B criterion is a more restrictive criterion that also considers a maximum distance between two hits, based on a parameter called roadwidth which is the assumed maximim distance a photon can travel through seawater.
As usual we start with a number of imports and reading a sample input file.
End of explanation
"""


context, cc = util.init_pycuda()
qd_kernel = QuadraticDifferenceSparse(N, cc=cc)
purging = PurgingSparse(N, cc)

"""
Explanation: We also initialize the GPU, and instantiate the Python interfaces to the GPU codes to get the GPU kernels compiled.
End of explanation
"""


#obtain sparse matrix for Quadratic Difference and convert to dense
d_col_idx, d_prefix_sums, d_degrees, total_hits = qd_kernel.compute(x, y, z, ct)
matrix_qd = util.sparse_to_dense(d_prefix_sums, d_col_idx, N, total_hits)

#obtain correlation matrix for Match 3B criterion
correlations_3b = np.zeros((window_width,N), dtype=np.uint8)
correlations_3b = util.correlations_cpu_3B(correlations_3b, x, y, z, ct)
matrix_3b = util.get_full_matrix(correlations_3b)

"""
Explanation: The next step is to compute the correlation tables for both criteria using the sample input we read in earlier. The Quadratic Difference criterion is computed by the QuadraticDifferenceSparse object on the GPU, for the 3B criterion we use the correlations_cpu_3B function CPU implemented in Python.
End of explanation
"""


f, (ax1, ax2) = pyplot.subplots(nrows=1, ncols=2, sharex=True, sharey=True)
f.tight_layout()
ax1.set_title('quadratic difference')
ax2.set_title('match 3b')
ax1.set_adjustable('box-forced')
ax2.set_adjustable('box-forced')
ax1.imshow(matrix_qd, cmap=pyplot.cm.bone, interpolation='nearest')
ax2.imshow(matrix_3b, cmap=pyplot.cm.bone, interpolation='nearest')

"""
Explanation: Let's start our comparison by looking at the two correlation matrices we've just produced.
End of explanation
"""


pyplot.hist(np.sum(matrix_qd, axis=0), bins=100, alpha=0.5, label='quadratic difference')
pyplot.hist(np.sum(matrix_3b, axis=0), bins=100, alpha=0.5, label='match 3b')
pyplot.legend(loc='upper right')
pyplot.xlabel('# correlated hits per hit')
pyplot.ylabel('# of hits')

"""
Explanation: We can clearly see that the Match 3B criterion produces fewer correlations than the quadratic difference criterion. However, it may be more informative to look at the distributions of the number of correlated hits per hit.
End of explanation
"""


clique = purging.compute(*util.dense_to_sparse(matrix_qd))
print("Quadratic Difference resulted in clique of size", len(clique))
print(clique)

clique = purging.compute(*util.dense_to_sparse(matrix_3b))
print("Match 3B resulted in clique of size", len(clique))
print(clique)

"""
Explanation: We can clearly see that the number of correlated hits per hit is lower for all hits using the 3B criterion. The distribution is more narrow because there are more hits that have a similar 'degree'. 
However, to really see how the 3B criterion improves over quadratic difference we need to look at the cliques that are produced by the purging algorithm. Because the PurgingSparse GPU algorithm we've implemented operates on a sparsely stored correlation matrix we need to first convert the dense matrix to a sparse matrix.
End of explanation
"""


