import numpy as np
import json, os, glob, yaml, time
import argparse, pickle, h5py
from pathlib import Path
from sfdata import SFDataFile, SFDataFiles

from alvra_tools.load_data import *
from alvra_tools.channels import *
from alvra_tools.utils import *


EventChan = 'SAR-CVME-TIFALL4:EvtSet'
Detector  = 'JF06T32V07'
chIzero   = 'SAROP11-PBPS122:INTENSITY'

# This is for 16M!

#roi = [900, 1300, 850, 1250]    
roi = [800, 1400, 750, 1350]    

def get_acq_from_pth(path):
    import re
    return re.search(r"acq(\d+)", path).group(1).zfill(4)

def get_pumped(Eventcode):    
    FEL_raw  = Eventcode[:,13] #Event 13: changed from 12 on June 22
    Ppicker  = Eventcode[:,200]
    Laser    = Eventcode[:,18]
    Darkshot = Eventcode[:,21]
    Deltap_FEL = (1 / FEL_raw.mean()).round().astype(int)

    FEL = FEL_raw
    if Darkshot.mean()==0:
        laser_reprate = (1 / Laser.mean()).round().astype(int)
        index_light = np.logical_and.reduce((FEL, Laser))
        index_dark  = np.logical_and.reduce((FEL, np.logical_not(Laser)))
    else:
        laser_reprate = (Laser.mean() / Darkshot.mean()).round().astype(int) 
        index_light = np.logical_and.reduce((FEL, Laser, np.logical_not(Darkshot)))
        index_dark = np.logical_and.reduce((FEL, Laser, Darkshot))

    print ('Laser rep rate is {} Hz (delayed or dark)'.format(100 / laser_reprate))
    print ('Pump scheme is {}:1'.format(laser_reprate - 1))
    return (index_light, index_dark, Deltap_FEL)

def prepare_radial_profile(shape, x0, y0):
    y,x = np.indices(shape)
    rad = np.sqrt((x-x0)**2 + (y-y0)**2)
    rad = rad.astype(int).ravel()
    norm = np.bincount(rad)
    return rad, norm

def radial_profile(data, rad, norm):
    data = data.ravel()
    tbin = np.bincount(rad, data)
    rp = tbin / norm
    return rp

parser = argparse.ArgumentParser()
parser.add_argument("--input", "-i", required=True)
args = parser.parse_args()

filename = Path(args.input)
output_file = Path(str(filename).replace("/raw/", "/res/processed/").replace(".h5",".pp_PROCESSED.h5") )
output_file.parent.mkdir(parents=True, exist_ok=True, mode=0o775)

acq = get_acq_from_pth(args.input)
acq_meta_pth = Path(str(filename).replace("/data/ac", "/meta/ac")).parent
acq_info_pth = acq_meta_pth.joinpath(f"acq{acq}.json")

with open(acq_info_pth, "r") as f:
    acq_parameters = json.load(f)
    rbk_value = acq_parameters.get("scan_info", {}).get("scan_readbacks", 0)
    #xlabel = acq_parameters.get("scan_info", {}).get("name", None)
    #units = acq_parameters.get("scan_info", {}).get("units", None)

centerX = 1113 - roi[0] #int((roi[1] - roi[0])/2)
centerY = 1054 - roi[2] #int((roi[3] - roi[2])/2)

file2load = args.input.replace('BSDATA', "*")
print ('sleeping 30 seconds, ROI = {}'.format(roi))
time.sleep(30) 

#file2load = args.input.replace('JF06T32V07', "*")

with SFDataFiles(file2load) as step:
    subset = step[EventChan, chIzero, Detector]
    subset.drop_missing()
    Event_code = subset[EventChan].data
    index_light, index_dark, Deltap_FEL = get_pumped(Event_code)
    
    pids = subset[chIzero].pids
    pids_pump   = pids[index_light]
    pids_unpump = pids[index_dark]
    shifted_pids_pump   = pids_unpump + Deltap_FEL
    final_pids, indPump, indUnPump = np.intersect1d(pids_pump, shifted_pids_pump, return_indices=True)

    JF    = crop_roi(subset[Detector].data, roi)
    Izero = subset[chIzero].data
    rad, norm = prepare_radial_profile(np.shape(JF[0]), centerX, centerY)

    #rp = np.array([radial_profile(img, rad, norm) for img in JF])
    data_flat = JF.reshape(JF.shape[0], -1)
    tbin = np.array([np.bincount(rad, weights=data_flat[i], minlength=len(norm)) for i in range(data_flat.shape[0])])
    rp = tbin/norm

    before = len(rp)    

    area_rp = np.sum(rp, axis=1)
    q_low  = np.percentile(area_rp, 5)
    q_high = np.percentile(area_rp, 95)
    outliers = (area_rp < q_low) | (area_rp > q_high)

    rp = rp[~outliers]
    Izero = Izero[~outliers]
    data_flat = data_flat[~outliers]
    index_light = index_light[~outliers]
    index_dark  = index_dark[~outliers]

    print ("Survived {} out of {} shots".format(len(rp), before))
    
    rp = rp / np.sum(rp[:,300:350], axis=1)[:, None]
    #rp = rp / np.sum(rp, axis=1)[:, None]

    #rp_pump   = rp[indPump]#/Izero[indPump][:, np.newaxis]
    #rp_unpump = rp[indUnPump]#/Izero[indUnPump][:, np.newaxis]
    rp_pump   = rp[index_light]
    rp_unpump = rp[index_dark]
    i0_pump   = Izero[index_light]
    i0_unpump = Izero[index_dark]

    print (np.shape(data_flat))

    JF_pump   = np.sum(data_flat[index_light], axis=1)
    JF_unpump = np.sum(data_flat[index_dark], axis=1)

    print ("Loaded {} light and {} dark".format(len(rp_pump), len(rp_unpump)))
    
    rpp_avg = np.mean(rp_pump, axis=0)
    rpu_avg = np.mean(rp_unpump, axis=0)

    rp_pp_avg = rpp_avg - rpu_avg
    int2save = np.sum(np.abs(rp_pp_avg))
    #print (np.shape(rp_pp), int2save)
    
    #Izero = subset[chIzero].data    
    #JF_pump    = np.sum(JF[indPump].sum(axis=1).sum(axis=1))/Izero[indPump]
    #JF_unpump  = np.sum(JF[indUnPump].sum(axis=1).sum(axis=1))/Izero[indUnPump]
    
with h5py.File(output_file, 'w') as f:
    gr1 = f.create_group("data")
    gr1['Izero_pump'] = i0_pump
    gr1['Izero_unpump'] = i0_unpump
    gr1['Int_JF_pump'] = JF_pump
    gr1['Int_JF_unpump'] = JF_unpump 
    gr1['int_radial'] = int2save
    gr1['rp_pump']   = rpp_avg    
    gr1['rp_unpump'] = rpu_avg
    gr1['rp_pp_avg'] = rp_pp_avg

    gr2 = f.create_group("meta")
    gr2["readback_value"] = np.atleast_1d(rbk_value)
    #gr2["xlabel"] = xlabel
    #gr2["units"]  = units

#with h5py.File(output_file, 'w') as f:
#   f['cropped'] = cropped
#   f['intensity'] = intensity