#!/usr/bin/env python3

# author J.Beale

"""
# aim
a quick a dirty script to see how your indexing went
will spit out useful parameters like the number images, indexing rate etc

# usage
python stream_stats <path to stream file>

# output
text in line - no files
chunks, crystals, indexing rate, mean res, obs, unit cell parameters
"""
import numpy as np
import pandas as pd
import regex as re
import sys

def count_chunks( stream ):

    # get number of chunks
    # example - ----- Begin chunk -----
    # count them
    try:
        pattern = r"-----\sBegin\schunk\s-----"
        chunks = re.findall( pattern, stream )
        if AttributeError:
            return len( chunks )
    except AttributeError:
        return np.nan

def scrub_cells( stream ):

    # get uc values from stream file
    # example - Cell parameters 7.71784 7.78870 3.75250 nm, 90.19135 90.77553 90.19243 deg
    # scrub clen and return - else nan
    try:
        pattern = r"Cell\sparameters\s(\d+\.\d+)\s(\d+\.\d+)\s(\d+\.\d+)\snm,\s(\d+\.\d+)\s(\d+\.\d+)\s(\d+\.\d+)\sdeg"
        cell_lst = re.findall( pattern, stream )
        xtals = len( cell_lst )
        if AttributeError:
            return cell_lst, xtals
    except AttributeError:
        return np.nan

def scrub_res( stream ):

    # get diffraction limit
    # example - diffraction_resolution_limit = 4.07 nm^-1 or 2.46 A
    # scrub res_lst or return np.nan
    try:
        pattern = r"diffraction_resolution_limit\s=\s\d\.\d+\snm\^-1\sor\s(\d\.\d+)\sA"
        res_lst = re.findall( pattern, stream )
        if AttributeError:
            return res_lst
    except AttributeError:
        return np.nan

def scrub_obs( stream ):

    # get number of reflections
    # example - num_reflections = 308
    # scrub reflections or return np.nan
    try:
        pattern = r"num_reflections\s=\s(\d+)"
        obs_lst = re.findall( pattern, stream )
        if AttributeError:
            return obs_lst
    except AttributeError:
        return np.nan

def main( stream_pwd ):

    print( "reading stream file" )
    # open stream file
    stream = open( stream_pwd, "r" ).read()
    print( "done" )

    print( "scrubing data" )
    # get total number chunks
    chunks = count_chunks( stream )

    # get list of cells
    cell_lst, xtals = scrub_cells( stream )

    # get list of cells
    res_lst = scrub_res( stream )

    # get list of cells
    obs_lst = scrub_obs( stream )
    print( "done" )

    # res_df
    cols = [ "a", "b", "c", "alpha", "beta", "gamma" ]
    df = pd.DataFrame( cell_lst, columns=cols )
    df[ "resolution" ] = res_lst
    df[ "obs" ] = obs_lst

    # convert all to floats
    df = df.astype(float)

    print( "calculating stream stats" )
    # stats
    index_rate = round( xtals/chunks*100, 2 )
    mean_res, std_res = round( df.resolution.mean(), 2 ), round( df.resolution.std(), 2 )
    median_res = df.resolution.median()
    mean_obs, std_obs = round( df.obs.mean(), 2 ), round( df.obs.std(), 2)
    mean_a, std_a = round( df.a.mean()*10, 2 ), round( df.a.std()*10, 2 )
    mean_b, std_b = round( df.b.mean()*10, 2 ), round( df.b.std()*10, 2 )
    mean_c, std_c = round( df.c.mean()*10, 2 ), round( df.c.std()*10, 2 )
    mean_alpha, std_alpha = round( df.alpha.mean(), 2 ), round( df.alpha.std(), 2 )
    mean_beta, std_beta = round(df.beta.mean(), 2 ), round( df.beta.std(), 2 )
    mean_gamma, std_gamma = round( df.gamma.mean(), 2 ), round( df.gamma.std(), 2 )

    # results
    print( "image = {0}".format( chunks ) )
    print( "crystals = {0}".format( xtals ) )
    print( "indexing rate = {0} %".format( index_rate ) )
    print( "mean resolution = {0} +/- {1} A".format( mean_res, std_res ) )
    print( "median resolution = {0} A".format( median_res  ) )
    print( "mean observations = {0} +/- {1}".format( mean_obs, std_obs ) )
    print( "mean a = {0} +/- {1} A".format( mean_a, std_a ) )
    print( "mean b = {0} +/- {1} A".format( mean_b, std_b ) )
    print( "mean c = {0} +/- {1} A".format( mean_c, std_c ) )
    print( "mean alpha = {0} +/- {1} A".format( mean_alpha, std_alpha ) )
    print( "mean beta = {0} +/- {1} A".format( mean_beta, std_beta ) )
    print( "mean gamma = {0} +/- {1} A".format( mean_gamma, std_gamma ) )

if __name__ == "__main__":
    stream_pwd = sys.argv[1]
    main( stream_pwd )