#!/usr/bin/env python

from mpl_toolkits.basemap import Basemap
import numpy as np
import matplotlib.pyplot as plt
import getopt, sys
from glob import glob

from matplotlib.patches import Polygon
from string import ascii_uppercase

#import scipy as sp # has to be 0.18!!!

import lib_fat_gradplot as lfg

FS = 11
pref_def = 'gradient'

def main(argv):
    '''Basic plotting and copying of files.'''

    help_line = '''
  A program to help give info about gradients that are actually used
  in a study, both per subject and summarizing across a group.
  
  Currently, spherical Voronoi tesselations and area calculations
  are used to quantify goodness-of-spread (i.e., homogeneity of
  gradient spread across a representative sphere).  The python
  spherical Voronoi package has been written and made available by:
  TJ Reddy
  https://github.com/tylerjereddy/py_sphere_Voronoi
  following quantitative work:
  Manuel Caroli, Pedro Machado Manhaes de Castro, Sebastien Loriot,
  Olivier Rouiller, Monique Teillaud, Camille Wormser (2009).
  Robust and Efficient Delaunay triangulations of points on or
  close to a sphere. Research Report RR-7004, INRIA.
  
  (v1.2, Apr. 2016, by PA Taylor-- still early days)

* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *

  OUTPUT

    + PREFIX_STAT.txt :some basic stats of spread-of-gradients;
                       current info includes (one line per subject):
                       o Ngrad: Number of grads.
                       o devi: deviation of Voronoi cell area 
                           relative to mean cell area, [0,1]; smaller  
                           number means more constant spread.
                       o maxA: maximum cell area.
                       o meanA: mean cell area.
                       o stdA: stdev of cell area.
                       o filename: grad file from whence information
                           came

    + PREFIX_MAP.jpg  :map of all grads for one or more subjects. Shows
                       overlaps within a default tolerance of 5 deg
                       radius.

* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *

  COMMANDS

    -h, --help         :everyone needs a little help, sometime.
    -g, --grads        :enter one or more grad files to analyze; can use
                        wildcard characters (probably use quotes, if you 
                        do so).
    -p, --prefix       :give output prefix; default is "'''+pref_def+'''".

    -r, --refset       :user can input a reference set of gradients

    -H, --Hold_image   :switch to hold the Python-produced image on the
                        output screen until a key has been hit; it puts
                        a 'raw_input()' line in, if you are curious  
                        (default:  not to do so -> meaning the image 
                        flashes briefly when running from a commandline, 
                        and not from, for example, ipython). Even without
                        this switch used, the image can be saved.
    -F, --Fig_off      :switch if you *don't* want a matrix figure output
                        (default is to save one).
    -t, --type_file=TT :Can select from a full range of image formats:
                        jpg (default), pdf, png, tif, etc. (whatever
                        your computer will allow).           
    -d, --dpi_file=DPI :set resolution (dots per inch) of output image 
                        (default = 80).

    -e, --eps_sep=EE   :interval for how far apart b-values can be in order
                        to be lumped into the same group (technically, the 
                        epsilon value for DBSCAN).
                        Default: if max |g| < 2, then EE=0.01; else, EE=10 
                        (basically guessing whether unit-magnitude gradients
                        have been entered, or ones that will be of order 
                        1000; guesses are made based on ref, if available;
                        else, all grads are processed together for clumps).


* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * 

  EXAMPLES:

   # can input many grads and get a summary over a group
   fat_grad_plot.py -g 'GRADS_*.dat' -p 'ctrl_grp'

   # can also just focus on one subject
   fat_grad_plot.py -g 'GRADS_30.dat' -p 'subj_01' -H

* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * 

'''

    file_grad_glob = ''     # get several grads
    file_prefix = pref_def  # def. output prefix 
    file_refset = ''        # optional set of ref grads
    DO_hold_image = 0       # keep showing on screen
    DO_PLOT = 1             # save images as jpg, etc.
    MATDPI = 80
    FTYPE = 'jpg'
    DB_EPS = -1

    try:
        opts, args = \
         getopt.getopt( argv,
                        "hHFt:d:g:p:r:e:",
                        [ "help",
                          "Hold_image",
                          "--Fig_off",
                          "--type_file",
                          "--dpi_file"
                          "grads=",
                          "prefix=",
                          "refset=",
                          "eps_sep="
                      ])

    except getopt.GetoptError:
        print "** ERROR: please wrong input-- please check help file:"
        print "\n\t  fat_grad_plot.py -h \n"
        # print help_line
        sys.exit(2)
    for opt, arg in opts:
        if opt in ("-h", "--help"):
            print help_line
            sys.exit()
        elif opt in ("-g", "--grads"):
            file_grad_glob = arg
        elif opt in ("-H", "--Hold_image"):
            DO_hold_image = 1
        elif opt in ("-F", "--Fig_off"):
            DO_PLOT = 0
        elif opt in ("-t", "--type_file"):
            FTYPE = arg
        elif opt in ("-d", "--dpi_file"):
            MATDPI = int(arg)
        elif opt in ("-p", "--prefix"):
            file_prefix = arg
        elif opt in ("-r", "--refset"):
            file_refset = arg
        elif opt in ("-e", "--eps_sep"):
            if float(arg)<=0:
                print "** ERROR: can't have a negative epsilon!"
                print "   See the '-help' about '-e, --eps_sep'."
                sys.exit(111)
            DB_EPS = float(arg)

    if ( file_grad_glob == '' ):
	print "** ERROR: missing a necessary grad file(s) input."
        print "\t Need to use either '-g' or '--grads='."
        print "\t Use 'fat_grad_plot.py -h' for viewing helpfile."
        sys.exit()
    if ( file_prefix == pref_def ):
	print "*+ Warning: no prefix entered, so one will be chosen:"
        print "\t ", pref_def

    return file_grad_glob, file_prefix, DO_hold_image, DO_PLOT, \
        FTYPE, MATDPI, file_refset, DB_EPS

########################################################################

if __name__=="__main__":
    DO_GROUP = 0
    np.set_printoptions(linewidth=200)
    print "\n"
    file_grad_glob, file_prefix, DO_hold_image, DO_PLOT, \
        FTYPE, MATDPI, file_refset, DB_EPS = main(sys.argv[1:])

    ### -----------------------------------------------------------
    ### initialize undergrid matrix

    # arrays and meshgrid for output
    my_lon, my_lat, x, y = lfg.Make_LongLat_Arr()

    ### -----------------------------------------------------------

    if file_grad_glob:
        list_all = glob(file_grad_glob)
        list_all.sort()
    else:
        print "** Error! Cannot read in matrix files."
        sys.exit(4)

    if len(list_all)==0:
        print "** Error! Cannot read in matrix files."
        print "\t Length of read in list is 0!\n\n"
        sys.exit(5)


    ### -----------------------------------------------------------
    ### Ref file

    if file_refset:
        tmp_ref = glob(file_refset)
        if len(tmp_ref) > 1:
            print "** Error! More than one refset matches input name:"
            print "\t "+tmp_ref+"\n\n"
            sys.exit(6)
            
        file_refset = tmp_ref[0]
        print "++ Found refset: ", file_refset
        my_refgrads_XYZ, my_refgrads_OI = lfg.Convert_CSV_to_arrays(file_refset)
        my_refgrads_Mxyz = lfg.SepMagnDir(my_refgrads_XYZ)
        my_refgrads_RPT = lfg.All_CtS(my_refgrads_Mxyz[:,1:])
        my_refvals = lfg.Calc_My_Vals(my_lon, my_lat, my_refgrads_RPT)
        my_refgrads_RPTneg = lfg.All_CtS(-my_refgrads_Mxyz[:,1:])
        my_refvals_neg = lfg.Calc_My_Vals(my_lon, my_lat, my_refgrads_RPTneg)
        # just use magnitudes
        rmm = lfg.Gen_bval_clust_levels_ref( my_refgrads_Mxyz[:,0:1],
                                             my_EPS = DB_EPS )
    else:  # !! not very efficient to do this stuff twice...
        print "++ No ref set to include. C'est la vie."
        all_Mxyz = np.zeros((0,4))
        for fff in list_all:
            my_grads_file = fff 
            my_grads_XYZ, my_grads_OI = lfg.Convert_CSV_to_arrays(my_grads_file)
            my_grads_Mxyz = lfg.SepMagnDir(my_grads_XYZ)
            all_Mxyz = np.concatenate((all_Mxyz,my_grads_Mxyz),axis=0)
        rmm = lfg.Gen_bval_clust_levels_ref( all_Mxyz[:,0:1],
                                             my_EPS = DB_EPS )

    # ==================================================================

    # prep to go through all
    Nbands = len(rmm) - 1
    quant_list = [[]]
    for i in range(Nbands-1):
        quant_list.append([])

    if len(list_all) > 1:
        DO_GROUP = 1
    
    if DO_GROUP:
        grp_arrs_Mxyz = [np.zeros((0,4))]*Nbands
    
    for ii in range(Nbands):
        letb = ascii_uppercase[ii]
        llim = rmm[ii]
        ulim = rmm[ii+1]

        llim_str = str(llim)
        ulim_str = str(ulim)
        if ulim >= lfg.BIG_VAL :
            ulim_str = "max"

        print "++ Processing band '%s':  (%s, %s]" % (letb, llim_str, ulim_str)

        if file_refset:
            #print "++ \t... using ref set: %s" % file_refset
            my_refs_XYZ, my_refs_OI = lfg.Convert_CSV_to_arrays(file_refset)
            my_refs_Mxyz = lfg.SepMagnDir(my_refs_XYZ)

            set1 = my_refs_Mxyz[:,0] > llim
            set2 = my_refs_Mxyz[:,0] <= ulim
            indi_ref_Mxyz = my_refs_Mxyz[set1*set2]
            indi_ref_OI = my_refs_OI[set1*set2]
            Nref = np.shape(indi_ref_Mxyz)[0]
            if not(Nref):
                indi_ref_Mxyz = np.zeros((0,4))
                indi_ref_OI = np.zeros(None) # null array
        else:
            #print "++ \t... no ref set to include. C'est la vie."
            indi_ref_Mxyz = np.zeros((0,4))
            indi_ref_OI = np.zeros(None) # null array
        for fff in list_all:
            my_grads_XYZ, my_grads_OI = lfg.Convert_CSV_to_arrays(fff)
            my_grads_Mxyz = lfg.SepMagnDir(my_grads_XYZ)

            set1 = my_grads_Mxyz[:,0] > llim
            set2 = my_grads_Mxyz[:,0] <= ulim

            # individual: goes for worldmap + polygons
            indi_arr_Mxyz = my_grads_Mxyz[set1*set2]
            Ngrad = np.shape(indi_arr_Mxyz)[0]
            fout_img = lfg.MakeName(file_prefix, fff, letb)

            if Ngrad:
                Bmean = np.mean(indi_arr_Mxyz[:,0])
                Bstd  = np.std(indi_arr_Mxyz[:,0])

                lfg.Proc_N_Print( indi_arr_Mxyz, indi_ref_Mxyz,
                                  fout_img, 
                                  FS, DO_PLOT, FTYPE, 
                                  MATDPI, DO_hold_image, 
                                  BMS=(Bmean, Bstd), 
                                  Rlist_OI=indi_ref_OI)
                VA, mA, sA, maxA, \
                    minA, VV, Vverts = \
                                       lfg.Do_Voronoi( indi_arr_Mxyz[:,1:],
                                                       fout_img,
                                                       1,
                                                       DO_PLOT,
                                                       FTYPE,
                                                       MATDPI,
                                                       FS, 
                                                       DO_hold_image, 
                                                       (Bmean, Bstd))
                subj_info = [fff, Ngrad, Bmean, Bstd, 
                             sA/mA, maxA, minA, mA, sA ]
                quant_list[ii].append(subj_info)
                if DO_GROUP:
                    # growing group list: goes for worldmap only
                    grp_arrs_Mxyz[ii] = np.concatenate( (grp_arrs_Mxyz[ii],
                                                         indi_arr_Mxyz),
                                                        axis=0)
            else:
                print "+* WARNING: no grads for this bval!"
                subj_info = [fff, 0,0,0,0,0,0,0,0]
                quant_list[ii].append(subj_info)

        if DO_GROUP:
            fout_img = lfg.MakeName(file_prefix, 'GROUP', letb)
            lfg.Proc_N_Print( grp_arrs_Mxyz[ii], indi_ref_Mxyz,
                              fout_img, 
                              FS, DO_PLOT, FTYPE, 
                              MATDPI, DO_hold_image)
        whatev = lfg.WriteOutGradStats(file_prefix+'_'+letb, quant_list[ii]) 
        #sys.exit('EARLY')

    print "DONE!"
