#!/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 scrub_cells( line ):

    # get uc values from stream file
    # example - Cell parameters 7.71784 7.78870 3.75250 nm, 90.19135 90.77553 90.19243 deg
    pattern = r"Cell\sparameters\s(\d+\.\d+)\s(\d+\.\d+)\s(\d+\.\d+)\snm,\s(\d+\.\d+)\s(\d+\.\d+)\s(\d+\.\d+)\sdeg"
    a = re.search( pattern, line ).group(1)
    b = re.search( pattern, line ).group(2)
    c = re.search( pattern, line ).group(3)
    alpha = re.search( pattern, line ).group(4)
    beta = re.search( pattern, line ).group(5)
    gamma = re.search( pattern, line ).group(6)

    return [ a, b, c, alpha, beta, gamma ]

def scrub_res( stream ):

    # get diffraction limit
    # example - diffraction_resolution_limit = 4.07 nm^-1 or 2.46 A
    pattern = r"diffraction_resolution_limit\s=\s\d\.\d+\snm\^-1\sor\s(\d+\.\d+)\sA"
    res = re.search( pattern, stream ).group(1)
    return res

def scrub_obs( stream ):

    # get number of reflections
    # example - num_reflections = 308
    pattern = r"num_reflections\s=\s(\d+)"
    obs = re.search( pattern, stream ).group(1)
    return obs

def main( stream_pwd ):

    chunks = 0
    xtals = 0
    cells = []
    obs_list = []
    res_list = []

    print( "scrubing data" )
    # open stream file
    with open( stream_pwd ) as stream:
        for line in stream:
            
            # count chunks
            if line.startswith( "----- Begin chunk -----" ):
                chunks = chunks + 1

            # get cell
            if line.startswith( "Cell parameters" ):
                cell = scrub_cells( line )
                cells.append( cell )
                xtals = xtals + 1

            # get res
            if line.startswith( "diffraction_resolution_limit" ):
                res = scrub_res( line )
                res_list.append( res )

            # get obs
            if line.startswith( "num_reflections" ):
                obs = scrub_obs( line )
                obs_list.append( obs )

            if chunks % 1000 == 0:
                print( "scrubbed {0} chunks".format( chunks ), end='\r' )

    # res_df
    cols = [ "a", "b", "c", "alpha", "beta", "gamma" ]
    df = pd.DataFrame( cells, columns=cols )
    df[ "resolution" ] = res_list
    df[ "obs" ] = obs_list

    # 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 )