#!/usr/bin/env python3

# authors J.Beale + lots of help from Alexander Steppke and Sven Augustin - love you both!

"""
# aim
 - -16M=varient for large detectors
make image file to input into pyFAI for initial detector beam-centre and detector distance calibration
refer to Cristallina8M-calibration for complete protocol
https://docs.google.com/document/d/1RoeUUogvRxX4M6uqGwkjf3dVJBabiMUx4ZxwcA5e9Dc/edit#

# protocol
take scan of LaB6
## IMPORTANT ##
- save image as photon-counts - in slic/run_control scale=beam energy
- detector_geometry=TRUE - saves detector panels in their correct orientation
## scan inputs ##
- <0.01 trans
- motor scan > 10 um per step
- 10 images per step, 100 steps
- use scan.json as input for this script

# usage
python convert-scan-for-pyfai.py -j <jugfrau-name> -s <path to scan file> -n <name of output file>

# output
creates a .npy file that can be loaded directly into pyFAI
"""

# modules
from matplotlib import pyplot as plt
import numpy as np
from sfdata import SFScanInfo
from tqdm import tqdm
import argparse

def convert_image( path_to_json, jungfrau, name ):

    # opens scan
    print( "opening scane" )
    scan = SFScanInfo( path_to_json )

    # steps in scane
    nsteps = len(scan)

    # define step ch and im_shape
    step = scan[0]
    ch = step[jungfrau]
    img_shape = ch[0].shape

    print("stepping through scan and averaging images at each step")
    # step through scan and average files from each positions
    imgs_shape = (nsteps, *img_shape)
    imgs = np.empty(imgs_shape)
    for i, subset in tqdm(enumerate(scan)): 

        # go through data in_batches so you don't run out of memory
        ch = subset[jungfrau]
        mean = np.zeros(img_shape)
        for _indices, batch in ch.in_batches(size=2):
            mean += np.mean(batch, axis=0)
        
        # take mean of means for batch opened data
        imgs[i] = mean

    print( "done" )
    # sum averaged imaged
    print( "final average" )
    mean_image = imgs.mean(axis=0)
    print("done")

    # output to file
    print( "saving to .npy = {0}".format( name ) )
    np.save( "{0}.npy".format( name ), mean_image )
    print( "done" )

    # create plot of summed, averaged scan
    fig, ax = plt.subplots()
    ax.imshow(mean_image, vmin=0, vmax=1000)
    plt.show()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-j",
        "--jungfrau",
        help="name of the jungfrau used, i.e., JF17T16V01 for Cristallina MX",
        type=str,
        default="JF17T16V01"
    )
    parser.add_argument(
        "-s",
        "--scan",
        help="path to json scan file",
        type=str,
        default="/sf/cristallina/data/p20590/raw/run0003/meta/scan.json"
    )
    parser.add_argument(
        "-n",
        "--name",
        help="name of output file",
        type=str,
        default="mean_scan"
    )
    args = parser.parse_args()

    convert_image( args.scan, args.jungfrau, args.name )