#!/usr/bin/env python3

# author J.Beale

"""
# aim
analyses and selects crystals based on their reported resolution from crystfel

# usage
python stream_random.py -s <path to stream>
                        -o output file names
                        -p plots histogram of images resolution
                        -r selects all images with higher resolution than value

# output
either
- histogram of images by resolution
- .stream file of selected images
"""

# modules
import re
import argparse
import pandas as pd
import numpy as np
import os
import matplotlib.pyplot as plt

def extract_chunks( input_file ):

    # setup
    chunk_df = pd.DataFrame()
    image_no = []
    chunks = []
    hits = []
    collect_lines = False
    # Open the input file for reading
    with open(input_file, 'r') as f:
        for line in f:

            # Check for the start condition
            if line.startswith('----- Begin chunk -----'):
                hit = False
                collect_lines = True
                chunk_lines = []
            if collect_lines:
                chunk_lines.append(line)
            
            # find image_no
            if line.startswith( "Event:" ):
                image_search = re.findall( r"Event: //(\d+)", line )
                image = int(image_search[0])
                image_no.append( image )
            
            # is there a hit in chunk
            if line.startswith( "Cell parameters" ):
                hit = True

            if line.startswith('----- End chunk -----'):
                collect_lines = False  # Stop collecting lines
                chunks.append( chunk_lines )
                hits.append( hit )
    
    chunk_df[ "chunks" ] = chunks
    chunk_df[ "image_no" ] = image_no
    chunk_df[ "hit" ] = hits

    return chunk_df

def scrub_res( line ):

    # get resolution
    try:
        pattern = r"diffraction_resolution_limit\s=\s\d\.\d+\snm\^-1\sor\s(\d+\.\d+)\sA"
        res = re.search( pattern, line ).group(1)
    except AttributeError as e:
        res = np.nan

    return float( res )

def extract_xtals( chunk ):

    # setup
    xtals = []
    resolutions = []
    collect_crystal_lines = False
    # Open the input file for reading
    for line in chunk:

        # Check for the xtals start condition
        if line.startswith('--- Begin crystal'):
            collect_crystal_lines = True
            xtal_lines = []
        if collect_crystal_lines:
            xtal_lines.append(line)
        if line.startswith('--- End crystal\n'):
            collect_crystal_lines = False  # Stop collecting lines
            xtals.append( xtal_lines )
        if line.startswith( "diffraction_resolution_limit" ):
            res = scrub_res( line )
            resolutions.append( res )

    return xtals, resolutions

def extract_header( chunk ):

    # setup
    header = []
    collect_header_lines = False
    # Open the input file for reading
    for line in chunk:

        # Check for the xtals start condition
        if line.startswith('----- Begin chunk -----'):
            collect_header_lines = True
            header_lines = []
        if collect_header_lines:
            header_lines.append(line)
        if line.startswith('End of peak list'):
            collect_header_lines = False  # Stop collecting lines
            header.append( header_lines )

    return header

def get_header( header, input_file ):

    if header == "geom":
        start_keyword = "----- Begin geometry file -----"
        end_keyword = "----- End geometry file -----"
    if header == "cell":
        start_keyword = "----- Begin unit cell -----"
        end_keyword = "----- End unit cell -----"

    # setup
    collect_lines = False
    headers = []

    # Open the input file for reading
    with open(input_file, 'r') as f:
        for line in f:
            # Check for the start condition
            if line.strip() == start_keyword:
                collect_lines = True
                headers_lines = []
            # Collect lines between start and end conditions
            if collect_lines:
                headers_lines.append(line)
            # Check for the end condition
            if line.strip() == end_keyword:
                collect_lines = False  # Stop collecting lines
                headers.append(headers_lines)
    
    return headers[0]

def write_to_file( geom, cell, chunk_header, crystals, output_file ):

    # Write sections with matching cell parameters to the output file
    with open(output_file, 'w') as out_file:
        out_file.write('CrystFEL stream format 2.3\n')
        out_file.write('Generated by CrystFEL 0.10.2\n')
        out_file.writelines(geom)
        out_file.writelines(cell)
        for crystal, header in zip( crystals, chunk_header ):
            out_file.writelines( header )
            out_file.writelines( crystal )
            out_file.writelines( "----- End chunk -----\n" )

def sort_xtals( chunk_df ):

    # extract xtals
    xtal_df = pd.DataFrame()
    counter = 0
    for index, row in chunk_df.iterrows():

        chunk, hit, image_no = row[ "chunks" ], row[ "hit" ], row[ "image_no" ]

        if hit:
        
            # find xtals and header
            header = extract_header( chunk )
            xtals, resolutions = extract_xtals( chunk )

            # make header same length as xtals
            header = header*len(xtals)

            # concat results
            xtal_df_1 = pd.DataFrame()
            xtal_df_1[ "header" ] = header
            xtal_df_1[ "xtals" ] = xtals
            xtal_df_1[ "image_no" ] = image_no
            xtal_df_1[ "resolution" ] = resolutions 
            xtal_df = pd.concat( ( xtal_df, xtal_df_1 ) )

            # add count and print every 1000s
            counter = counter + len(xtals)
            if counter % 1000 == 0:
                print( counter, end='\r' )
    print( "done" )

    # sort by image no and reindex
    xtal_df = xtal_df.sort_values( by=[ "image_no" ] )
    xtal_df = xtal_df.reset_index( drop=True )

    return xtal_df

def plot_res_histogram( res_df, res_median, res_q1, res_q3 ):

    # calculate relative numbers of bins
    bin_range = 0.25
    res_min = res_df.min().values[0]
    res_max = res_df.max().values[0]

    q1_bins = round( ( res_q1 - res_min )/bin_range )
    q2_bins = round( ( res_median - res_q1 )/bin_range )
    q3_bins = round( ( res_q3 - res_median )/bin_range )
    q4_bins = round( ( res_max - res_q3 )/bin_range )

    # cut data by quantile
    df_q1 = res_df[ res_df.resolution <= res_q1 ]
    df_q2 = res_df[ ( res_df.resolution > res_q1 ) & ( res_df.resolution <= res_median ) ]
    df_q3 = res_df[ ( res_df.resolution > res_median ) & ( res_df.resolution <= res_q3 ) ]
    df_q4 = res_df[ res_df.resolution > res_q3 ]

    # plot histogram of resolution
    fig, axs = plt.subplots()

    axs.hist( df_q1, bins=q1_bins, rwidth=1, stacked=True, color="blue", label="x<q1" )
    axs.hist( df_q2, bins=q2_bins, rwidth=1, stacked=True, color="red", label="q1<x<q2")
    axs.hist( df_q3, bins=q3_bins, rwidth=1, stacked=True, color="green", label="q2<x<q3" )
    axs.hist( df_q4, bins=q4_bins, rwidth=1, stacked=True, color="purple", label="q3<x<q4" )

    axs.axvline( x=res_median, color="black", linestyle="dashed", label="median = {0}".format( res_median ) )
    axs.set_xlabel( "resolution" )
    axs.set_ylabel( "frequency" )
    axs.legend()
    bin_size = round( ( len( df_q1 ) + len( df_q2 ) + len( df_q3 ) + len( df_q4 ) )/4 )
    axs.text( 0.98, 0.7,
             "images per quartile = {0}".format( bin_size ),
             ha="right",
             va="top",
             transform=axs.transAxes 
             )

    fig.tight_layout()
    plt.show()

def main( input_file, output, plotter, resolution ):

    # get geom and cell file headers
    print( "getting header info from .stream file" )
    geom = get_header( "geom", input_file )
    cell = get_header( "cell", input_file )
    print( "done" )
    
    # extract chunks
    print( "finding chucks" )
    chunk_df = extract_chunks( input_file )
    # display no. of chunks
    print( "found {0} chunks".format( len(chunk_df) ) )
    # remove rows without xtals
    chunk_df = chunk_df.loc[chunk_df.hit, :]
    print( "found {0} hits (not including multiples)".format( len(chunk_df) ) )
    print( "done" )

    print( "sorting xtals from chunks" )
    xtal_df = sort_xtals( chunk_df )
    print( "done" )

    print( "calculate stats" )
    res_df = xtal_df[ [ "resolution" ] ]
    res_median = res_df.median().values[0]
    res_q1 = res_df.quantile( 0.25 ).values[0]
    res_q3 = res_df.quantile( 0.75 ).values[0]
    print( "median resolution and range = {0} ({1}-{2})".format( res_median, res_q1, res_q3 ) )
    print( "done" )

    if plotter == True:
        print( "plot image resolution histogram" )
        plot_res_histogram( res_df, res_median, res_q1, res_q3  )
    
    if resolution:

        print( "finding images with resolution greater than {0}".format( resolution ) )
        select_df = xtal_df[ xtal_df[ "resolution" ] <= resolution ]
        print( "done" )

        print( "writing {0} to output file".format( len( select_df ) ) )
        crystals = select_df.xtals.to_list()
        chunk_header = select_df.header.to_list()
        output_file = "{0}.stream".format( output )
        write_to_file( geom, cell, chunk_header, crystals, output_file )
        print( "done" )
    
    else:
        print( "no output. please use either the plot-histogram or resolution functions of the script" )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-s",
        "--stream",
        help="input stream file",
        required=True,
        type=os.path.abspath
    )
    parser.add_argument(
        "-o",
        "--output",
        help="output stream file name. '.stream will be added'",
        type=str,
        default="selected"
    )
    parser.add_argument(
        "-p",
        "--plot_histogram",
        help="plots a histogram of the crystfel calculated resolutions for inspection",
        type=bool,
        default=False
    )
    parser.add_argument(
        "-r",
        "--resolution",
        help="upper resolution limit. Will take all images below this.",
        type=float
    )
    args = parser.parse_args()

    # run main
    main( args.stream, args.output, args.plot_histogram, args.resolution )