# modules
import pandas as pd
import subprocess
import os, errno
import regex as re
import numpy as np

def h5_sample( lst, sample ):

    # create sample of images from run
    # read h5.lst - note - removes // from imade column
    cols = [ "h5", "image" ]
    sample_df = pd.read_csv( lst, sep="\s//", engine="python", names=cols )

    # take defined sample
    sample_df = sample_df.sample( sample )

    # sort list
    sample_df = sample_df.sort_index()

    # re-add // to image columm
    sample_df[ "image" ] = "//" + sample_df.image.astype(str)

    # write sample to file
    sample_file = "h5_{0}_sample.lst".format( sample )
    sample_df.to_csv( sample_file, sep=" ", index=False, header=False )

    # return sample file name
    return sample_file

def geom_amend( lab6_geom_file, clen ):

    # read lab6 geom
    lab6_geom = open( lab6_geom_file, "r" )

    # use regex to find clen and replace with new
    # clen example => clen = 0.1217
    clen_geom = re.sub( "clen = 0\.\d+", "clen = {0}".format( clen ), lab6_geom.read() )

    # close lab6 geom file
    lab6_geom.close()

    # write new clen_geom to file
    clen_geom_file = "{0}.geom".format( clen )
    geom = open( clen_geom_file, "w" )
    geom.write( clen_geom )
    geom.close()

    # return clen_geom file name
    return clen_geom_file

def write_crystfel_run( clen, sample_h5_file, clen_geom_file, cell_file ):

    # crystfel file name
    cryst_run_file = "{0}_cryst_run.sh".format( clen )

    # write file
    run_sh = open( cryst_run_file, "w" )
    run_sh.write( "#!/bin/sh\n\n" )
    run_sh.write( "module purge\n" )
    run_sh.write( "module load crystfel/0.10.2\n" )
#    run_sh.write( "module use MX unstable\n" )
#    run_sh.write( "module load gcc/4.8.5 hdf5_serial/1.10.3 xds/20210205 DirAx/1.17 pinkindexer/2021.08\n" )
#    run_sh.write( "module load xgandalf/2021.08 fdip/2021.08 mosflm/7.3.0 crystfel/0.10.0 HDF5_bitshuffle/2018.05 HDF5_LZ4/2018.05 ccp4\n\n" )
    run_sh.write( "indexamajig -i {0} \\\n".format( sample_h5_file ) )
    run_sh.write( "  --output={0}.stream \\\n".format( clen ) )
    run_sh.write( "  --geometry={0}\\\n".format( clen_geom_file ) )
    run_sh.write( "  --pdb={0} \\\n".format( cell_file ) )
    run_sh.write( "  --indexing=xgandalf-latt-cell --peaks=peakfinder8 \\\n" )
    run_sh.write( "  --integration=rings-grad --tolerance=10.0,10.0,10.0,2,3,2 --threshold=10 --min-snr=5 --int-radius=2,3,6 \\\n" )
    run_sh.write( "  -j 36 --no-multi --no-retry --check-peaks --max-res=3000 --min-pix-count=1 --local-bg-radius=4 --min-res=85\n\n" )
    run_sh.close()

    # make file executable
    subprocess.call( [ "chmod", "+x", "{0}".format( cryst_run_file ) ] )

    # return crystfel file name
    return cryst_run_file

def main( lst, sample, lab6_geom_file, centre_clen, cell_file, steps, scan_name, step_size ):

    # set current working directory
    cwd = os.getcwd()

    # make sample list
    print( "making {0} sample of images".format( sample ) )
    sample_h5 = h5_sample( lst, sample)
    sample_h5_file = "{0}/{1}".format( cwd, sample_h5 )
    print( "done" )

    # make list of clen steps above and below the central clen
    print( "make clen array around {0}".format( centre_clen ) )
    step_range = step_size*steps
    bottom_step = centre_clen-step_range/2
    top_step = bottom_step+step_range
    step_range = np.arange( bottom_step, top_step, step_size )
    step_range = step_range.round( 4 ) # important - otherwise np gives your .99999999 instead of 1 somethimes
    print( "done" )

    # make directorys for results
    print( "begin CrystFEL anaylsis of different clens" )

    # loop to cycle through clen steps
    for clen in step_range:

        # move back to cwd
        os.chdir( cwd )

        print( "processing clen = {0}".format( clen ) )
        # define process directory
        proc_dir = "{0}/{1}/{2}".format( cwd, scan_name, clen )

        # make process directory
        try:
            os.makedirs( proc_dir )
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise
        
        # move to process directory
        os.chdir( proc_dir )

        # make geom file
        print( "amend .geom file" )
        clen_geom_file = geom_amend( lab6_geom_file, clen )
        print( "done" )

        # make crystfel run file
        print( "make crystfel file" )
        cryst_run_file = write_crystfel_run( clen, sample_h5_file, clen_geom_file, cell_file )
        print( "done" )       

        # run crystfel file
        print( "run crystFEL" )
        #subprocess.call( [ "./{0}".format( cryst_run_file ) ] )
        subprocess.call( [ "sbatch", "-p", "day", "--cpus-per-task=32", "--", "./{0}".format( cryst_run_file ) ] )
        print( "done" )

        #subprocess.call( [ "sbatch", "-p", "day", "--cpus-per-task=32", "--", "run{0}.sh".format( run.zfill(4) ) ] )

# variables
sample = 500
lst = "/sf/cristallina/data/p20590/work/process/jhb/detector_refinement/acq0001.JF17T16V01.dark.lst"
lab6_geom_file = "/sf/cristallina/data/p20590/work/process/jhb/detector_refinement/8M_p-op_c-op_p20590.geom"
centre_clen = 0.122 # in m
cell_file = "/sf/cristallina/data/p20590/work/process/jhb/detector_refinement/hewl.cell"
steps = 10
scan_name = "fine_scan"
step_size = 0.0005 # m - 0.001 = coarse scan, 0.0005 = fine

main( lst, sample, lab6_geom_file, centre_clen, cell_file, steps, scan_name, step_size )