'''
    ***** fba.py *****

Read filter basket data from CSV file generated by fba, and plot hole areas.

Usage:  python fba.py basket.png ... (other fba parameters)

Requirements:
1) Python matplotlib.pyplot (pip install matplotlib)
2) fba executable (fba on Linux, fba.exe on Windoze)
3) Correctly exposed filter basket image in PNG format

Author: John M. Weiss, Ph.D.
Written 200905.

Modifications:
'''

import sys, os, csv
import matplotlib.pyplot as plt

def read_data( filename ):
    ''' read_data( filename ): read COVID-19 data from CSV file'''

    # open input file and read data into matrix
    try:
        with open( filename ) as fin:
            lines = list( csv.reader( fin ) )
    except IOError:
        print( "Error: unable to open file:", filename )
        raise SystemExit( -2 )

    return lines

def plot_data( filename ):
    ''' plot_data(): generate scatter plot using matplotlib.pyplot '''

    # read data from CSV file
    lines = read_data( filename + '.csv' )

    # plot basket hole size distribution (columns 0/1 for area, 6/7 for diameter)
    x = [ float(line[0]) for line in lines[19:] ]
    y = [ int(line[1]) for line in lines[19:] ]
    plt.plot( x, y, color = 'b' )

    # add plot decorations
    plt.grid()
    plt.xlim( 0, max(x)+min(x) )
    plt.title( 'basket hole size distribution for ' + filename )
    plt.xlabel( 'hole area (pixels)' )
    plt.ylabel( 'number of holes' )

    # display plot
    plt.tight_layout()
    plt.draw()
    plt.waitforbuttonpress()
    plt.close()

def main( argv ):
    ''' main(): fake a main function '''

    # check for basket PNG file
    filename = argv[1]
    if not os.path.isfile( filename ):
        print( "Error: cannot open", filename )
        raise SystemExit( -2 )

    # run filter basket analysis
    os.system( 'fba ' + ' '.join(sys.argv[1:]) )

    # read data from CSV file and and plot basket hole size distribution
    plot_data( os.path.splitext( filename )[0] )

# call the "main" function
if __name__ == '__main__':
    if len( sys.argv ) < 2:
        print( "Usage: python fba.py basket.png ..." )
        raise SystemExit( -1 )

    main( sys.argv )