from cam_server.pipeline.data_processing import functions, processor
import numpy as np
import time
from skimage.feature import peak_local_max

def process_image(image, pulse_id, timestamp, x_axis, y_axis, parameters, bsdata=None):
    y_cog = x_cog = 0
    start = time.time()
    
    '''Input parameters'''

    cts_theory = 544  # Thoretical charge cloud per photon @ 850 eV
    # cts_theory = 452  #Thoretical charge cloud per photon @ 710 eV
    # cts_theory = 320  #Thoretical charge cloud per photon @ 530 eV
    size = 5  # Max spread of an event (5x5 pixels)
    n_std = 5  # Low threshold as numer of std dev of noise peak
    high_th = 600  # High threshold

    '''Bkg subtraction'''
    #bkg_avg = np.loadtxt(save_folder + 'Bkg.txt')  # 2D matrix same size as the data
    #frame_0 = data - bkg_avg
    frame_0 = image

    ''''''

    #TODO
    #'''Threshold Pixelwise'''
    #bkg_std = np.loadtxt('std.txt')  # 2D matrix same size as the data
    #frame = frame_0.copy()
    #frame[(frame_0 < n_std * bkg_std)] = 0
    #frame[frame_0 > high_th] = 0
    frame = frame_0
    

    ''''''

    '''Single photon counting'''

    evt_list = np.array([])
    tot_evt = 0

    xx_c, yy_c = np.meshgrid(np.arange(size), np.arange(size))
    time1 = time.time()-start; start = time.time()
    coords = peak_local_max(frame, min_distance=6)
    time2 = time.time()-start; start = time.time()
    for e in range(len(coords)):

        y0, x0 = coords.astype(int)[e]

        extr = frame[(y0 - int((size - 1) * 0.5)):(y0 + int((size - 1) * 0.5) + 1),
               (x0 - int((size - 1) * 0.5)):(x0 + int((size - 1) * 0.5) + 1)]

        if np.sum(extr) > (cts_theory * 0.5):
            tot_evt += 1

            '''Find center of gravity'''
            height, width = extr.shape
            total_mass = np.sum(extr)
            x, y = np.meshgrid(range(width), range(height))
            x_cog = np.sum(extr * x) / total_mass
            y_cog = np.sum(extr * y) / total_mass
            x_cog = x0 - int(size * 0.5) + x_cog
        y_cog = y0 - int(size * 0.5) + y_cog
        ''''''

        evt_list = np.append(evt_list, np.array([int(tot_evt), x_cog, y_cog]))
    time3 = time.time()-start; start = time.time()
    evt_list = np.reshape(evt_list[:tot_evt*3], (tot_evt, 3))
    time4 = time.time()-start; start = time.time()
    return {"evt_list": evt_list, "projection":np.sum(frame, axis=1), "time1": time1, "time2": time2, "time3": time3, "time4": time4}