from EtaInterpolationFunctions import *
import matplotlib
matplotlib.use('Agg')

import numpy as np
from array import array
import matplotlib.pyplot as plt
import h5py
from multiprocessing import Pool
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--runname', '-r', type=str,  default='2603MaxIV_Edge3Filters_pos0_12keV', choices=[
    '2603MaxIV_Edge3Filters_pos0_12keV', '2603MaxIV_FlatField3Filters_pos0_12keV',
    '2603MaxIV_Edge2Filters_12keV', '2603MaxIV_Flat2Filters_12keV',
    '2603MaxIV_Edge1Filters_12keV', '2603MaxIV_Flat1Filters_12keV'
    ], help='Name of the run, used to locate the cluster files and save the results')
args = parser.parse_args()
RUNNAME = args.runname


### cluster files
clusterFiles = [f'/home/xie_x1/MLXID/DataProcess/Samples/{RUNNAME}/1Photon_CS3_chunk{i}.h5' for i in range(16)]

NEtaBins = 201
### generate eta map
def get_eta_hist2D(clusterFile):
    with h5py.File(clusterFile, 'r') as f:
        _eta_hist2D = np.zeros((NEtaBins, NEtaBins)) ### shape (201, 201), X, Y, according build_xy_lut_Rosenblatt's convention
        clusters = f['clusters'][:] ### shape (N, 3, 3), N, Y, X
        sum_clusters = np.sum(clusters, axis=(1, 2)) ### shape (N,)
        position_weights = np.array([-1, 0, 1]) + 0.5

        etaX =  (np.sum(clusters, axis=1) * position_weights).sum(axis=-1) / sum_clusters
        etaY =  (np.sum(clusters, axis=2) * position_weights).sum(axis=-1) / sum_clusters

        etaX = np.clip(etaX, 0, 1)
        etaY = np.clip(etaY, 0, 1)

        _eta_hist2D += np.histogram2d(etaX, etaY, bins=NEtaBins, range=[[0, 1], [0, 1]])[0] ### X, Y convention according to build_xy_lut_Rosenblatt's convention
    return _eta_hist2D

print('Building eta histogram from clusters...')
eta_hist2D = np.zeros((NEtaBins, NEtaBins))
with Pool(16) as pool:
    results = pool.map(get_eta_hist2D, clusterFiles)
    for res in results:
        eta_hist2D += res

# generate U, V lookup tables
U_tab, V_tab = build_xy_lut_Rosenblatt(eta_hist2D) ### or using build_xy_lut_DoubleCDF

Roi = [1, 100, 1, 100]
interpolationBins = 10

### slightly better LUTs from strictly selected clusters in /home/xie_x1/MLXID/EtaInterpolation/Examples/etaInterpolation_edge_12keV.ipynb
# U_tab = np.load('/home/xi/
clusterFiles = [f'/home/xie_x1/MLXID/DataProcess/Samples/{RUNNAME}/1Photon_CS3_chunk{i}.h5' for i in range(16)]

def reconstruct_position(clusterFile):
    _image = np.zeros(((Roi[1]-Roi[0])*interpolationBins, (Roi[3]-Roi[2])*interpolationBins))
    _subpixel_positions = np.zeros((interpolationBins, interpolationBins))
    with h5py.File(clusterFile, 'r') as f:
        clusters = f['clusters'][:] ### shape (N, 3, 3), N, Y, X
        referencePoints = f['referencePoint'][:] ### x, y

        sum_clusters = np.sum(clusters, axis=(1, 2)) ### shape (N,)
        position_weights = np.array([-1, 0, 1]) + 0.5 

        etaX =  (np.sum(clusters, axis=1) * position_weights).sum(axis=-1) / sum_clusters 
        etaY =  (np.sum(clusters, axis=2) * position_weights).sum(axis=-1) / sum_clusters 

        etaX = np.clip(etaX, 0, 1)
        etaY = np.clip(etaY, 0, 1)

        x_frac, y_frac = bilinear_xy_lookup(etaX, etaY, U_tab, V_tab)

        x_ref = referencePoints[:, 0]
        y_ref = referencePoints[:, 1]

        x = x_ref + x_frac + clusters.shape[1]//2 
        y = y_ref + y_frac + clusters.shape[2]//2

        _image += np.histogram2d(
            y, x, 
            bins=(Roi[1]-Roi[0])*interpolationBins, 
            range=[[Roi[0], Roi[1]], [Roi[2], Roi[3]]]
        )[0]
        _subpixel_positions += np.histogram2d(
            y%1, x%1,
            bins=interpolationBins, 
            range=[[0, 1], [0, 1]]
        )[0]
    return (_image, _subpixel_positions)

print('Reconstructing positions from clusters and accumulating hits...')
reconstructed_image = np.zeros(((Roi[1]-Roi[0])*interpolationBins, (Roi[3]-Roi[2])*interpolationBins))
subpixel_positions = np.zeros((interpolationBins, interpolationBins))
with Pool(16) as pool:
    results = pool.map(reconstruct_position, clusterFiles)
    for res in results:
        reconstructed_image += res[0]
        subpixel_positions += res[1]

plt.figure(figsize=(16, 8), dpi=300)
plt.subplot(1, 2, 1)
plt.imshow(reconstructed_image, cmap='viridis', origin='lower')
plt.colorbar()
plt.subplot(1, 2, 2)
plt.imshow(subpixel_positions, cmap='viridis', origin='lower')
plt.colorbar()

np.save(f'Results/{RUNNAME}_fromClusters.npy', reconstructed_image)