#! /usr/bin/env python
#
# util_NRBestOfIndex.py
#
# GOAL
#    Accepts ILE default .dat output (indx m1 m2 group param XXX).
#    Finds the *single best* case in that set.
#    ---> FIXME, SHOULD USE LOCAL QUADRATIC FIT <--
#
#
#   python ~/research-projects/MonteCarloMarginalizeCode/Code/util_NRBestOfIndex.py --fname RIT-Generic-Lmax2-vC2.indexed --group Sequence-RIT-Generic


import argparse
import sys
import numpy as np
import RIFT.lalsimutils as lalsimutils
import lalsimulation as lalsim
import lalframe
import lal
import functools
import gp

from multiprocessing import Pool
try:
    import os
    n_threads = int(os.environ['OMP_NUM_THREADS'])
    print " Pool size : ", n_threads
except:
    n_threads=1
    print " - No multiprocessing - "

try:
	import NRWaveformCatalogManager as nrwf
	hasNR =True
except:
	hasNR=False


###
### Load options
###

parser = argparse.ArgumentParser()
#parser.add_argument("--group",default=None)
parser.add_argument("--fname", default=None, help="Base output file for ascii text (.dat) and xml (.xml.gz)")
parser.add_argument("--verbose", action="store_true",default=False, help="Required to build post-frame-generating sanity-test plots")
parser.add_argument("--M-max-cut",type=float,default=1e5,help="Maximum mass to consider (e.g., if there is a cut on distance, this matters)")
parser.add_argument("--lnL-cut-up",default=5,type=float,help="Maximum amount (in lnL) a fit can increase lnL, based on what is calculated. Use a tighter threshold for a denser grid, and a wide threshold if you want to live dangerously.")
parser.add_argument("--sigma-cut",type=float,default=0.95,help="Eliminate points with large error from the fit.")
parser.add_argument("--fit", action="store_true",default=False, help="Local quadratic fit on best points")
parser.add_argument("--no-gp",action="store_true")
opts=  parser.parse_args()

if opts.verbose:
    True
    #lalsimutils.rosDebugMessagesContainer[0]=True   # enable error logging inside lalsimutils

if not opts.fname:
    sys.exit(0)



best_matches = {}
best_matches_masses ={}
full_spoke = {}

with open(opts.fname) as f:
 for line_str in f:
     # string split
     line = line_str.split()
#     if opts.verbose:
#         print "  Input ", line
     # Get NR data
     if float(line[1])+float(line[2]) > opts.M_max_cut:
         continue
     group = line[3]
     if not nrwf.internal_ParametersAreExpressions.has_key(group):
         continue
     param = None
     if nrwf.internal_ParametersAreExpressions[group]:
         param = eval(line[4])
     else:
         param = line[4]

     if len(line)<6:
         continue

     key = (group,param)

     failure_mode=False; sigma_here = 1e6
     # in case line abbreviated for some reason
     try:
         lnLhere = float(line[5])  
         sigma_here = float(line[6])
         npts_here = float(line[7])
     except:
         lnLhere = -10
         sigma_here = 0.1
         npts_here = 1
         failure_mode=True
         continue
     if sigma_here > opts.sigma_cut:
         failure_mode=True
         continue  # DO NOT RECORD ITEMS which are completely unconverged (one point). Insane answers likely.  (should ALSO have tunable cutoff on accuracy)
     if not failure_mode:
      if best_matches.has_key(key):
         if best_matches[key] < lnLhere:
                best_matches[key] = lnLhere
                best_matches_masses[key] = (float(line[1]),float(line[2]))
      else:
         best_matches[key] = line[5]
         best_matches_masses[key] = (float(line[1]),float(line[2]))

      # Add to spoke
      mtot = float(line[1])+ float(line[2])
      if full_spoke.has_key(key):
        full_spoke[key].append([mtot, lnLhere,sigma_here])
      else:
         full_spoke[key] = [[mtot,lnLhere, sigma_here]]

print " -----  BEST SINGLE POINT MATCHES ------ "  # unsorted
if  opts.fit:
  print "  +++ Using fit code +++ "
  for key in best_matches:
      full_spoke[key] = np.array(sorted(np.array(full_spoke[key]),key=lambda p: p[1]))
      lnLmaxHere = np.max(full_spoke[key][:,1])
      indx_peak = np.argmax(full_spoke[key][:,1])
      sigma_crit = full_spoke[key][indx_peak][2]
      n_required =np.max([4,np.min([10, 0.1*len(full_spoke[key])])])
      indx_crit =int(np.max([n_required-1,np.argmin( np.abs(full_spoke[key][:,1]-10))]))
      if opts.verbose:
	print " Spoke needs ", key, indx_crit+1
      reduced_spoke = full_spoke[key][-indx_crit:]
      mMin = np.min(reduced_spoke[:,0])
      mMax = np.max(reduced_spoke[:,0])
      weights = reduced_spoke[:,2]
      reduced_spoke[np.isnan(weights),2] = 0  # do not use NAN entries with errors
      reduced_spoke[np.isnan(weights),1] = 0  # do not use NAN entries with errors
#      print " Fitting ", key, reduced_spoke
      z=[]; mBestGuess = 0; lnLBestGuess = 0;
      try:
          z = np.polyfit(reduced_spoke[:,0], reduced_spoke[:,1],2,w=1./(reduced_spoke[:,2])) # Note TERRIBLE documentation/convention in numpy https://docs.scipy.org/doc/numpy/reference/generated/numpy.polyfit.html
          mBestGuess = -0.5*z[1]/z[0]
          lnLBestGuess = z[2] -0.25*z[1]**2/z[0] 
      except:
          print " Interpolation failure (internal to polyfit, VERY UNUSUAL) for spoke ", key, " reverting to pointwise best "
          indxMax = np.argmax(reduced_spoke[:,2])
          mBestGuess = reduced_spoke[indxMax,0]
          lnLBestGuess = reduced_spoke[indxMax,2]
          z = [0,0,lnLBestGuess]
      print key, z[0], mBestGuess, lnLBestGuess, best_matches[key], sigma_crit
      if z[2]<0 and mBestGuess> mMin and mBestGuess < mMax:
        if lnLBestGuess < lnLmaxHere+opts.lnL_cut_up and lnLBestGuess > lnLmaxHere-5*sigma_crit:  # do not allow arbitrary extrapolation
          if opts.verbose:
              print " Replacing peak ", key, best_matches[key], " -> ", lnLBestGuess, " at mass ", mBestGuess
              if lnLBestGuess > lnLmaxHere+5:
                  print " VERY LARGE CHANGE FOR", key, lnLmaxHere, "->", lnLBestGuess
          best_matches[key] = lnLBestGuess 
          orig_mtot = best_matches_masses[key][0]+best_matches_masses[key][1]
          orig_m1 = best_matches_masses[key][0]
          orig_m2 = best_matches_masses[key][1]
          best_matches_masses[key]=(mBestGuess* orig_m1/orig_mtot, mBestGuess* orig_m2/orig_mtot)
        else:
	  print " Replacement rejected as out of range ", key, " reject ", lnLmaxHere, "->", lnLBestGuess, " : you probably need to rerun this spoke"
      if z[2]<0 and not ( mBestGuess> mMin and mBestGuess < mMax):
          print " PLACEMENT FAILURE: ", key, mBestGuess, " outside of ", [mMin,mMax]
      

      ## REVISE BEST FIT VIA GAUSSIAN PROCESS
      #   GP is much less problematic for very crazy values that are far above trend due to bad MC luck
      #   Has VERY significant implications for rankings overall - often downranking
      try:
       if not opts.no_gp:
          my_gp = gp.GaussianProcess1d(reduced_spoke, sigma0=0.1,sigmab=0.2,h=1)
          xvals_dense = np.linspace(mMin,mMax,1000)
          yvals_dense = my_gp.predict(xvals_dense)[:,0]
          indx_gp_best = np.argmax(yvals_dense)
          print " GP best fit ", xvals_dense[indx_gp_best], yvals_dense[indx_gp_best][0,0], "versus", mBestGuess, lnLmaxHere
          if  mBestGuess-2 < xvals_dense[indx_gp_best] < mBestGuess+2 and yvals_dense[indx_gp_best][0,0] < lnLBestGuess+5*sigma_crit:
              mBestGuess = xvals_dense[indx_gp_best]
              lnLmaxHere= yvals_dense[indx_gp_best][0,0]
              best_matches[key] = lnLmaxHere
              orig_mtot = best_matches_masses[key][0]+best_matches_masses[key][1]
              orig_m1 = best_matches_masses[key][0]
              orig_m2 = best_matches_masses[key][1]
              best_matches_masses[key]=(mBestGuess* orig_m1/orig_mtot, mBestGuess* orig_m2/orig_mtot)
              print "  ....using GP best fit  for ", key, best_matches[key], best_matches_masses[key]

      except:
          print " GP failure"


for key in best_matches:
  try:
    tmax =0
    if nrwf.internal_EstimatePeakL2M2Emission[key[0]].has_key(key[1]):
        tmax = nrwf.internal_EstimatePeakL2M2Emission[key[0]][key[1]]
    af = -1000
    Mf = -1000
    if nrwf.internal_WaveformMetadata[key[0]][key[1]].has_key('ChiFMagnitude'):
        af = nrwf.internal_WaveformMetadata[key[0]][key[1]]['ChiFMagnitude']
    if nrwf.internal_WaveformMetadata[key[0]][key[1]].has_key('MF'):
        Mf = nrwf.internal_WaveformMetadata[key[0]][key[1]]['MF']
    wfP = nrwf.WaveformModeCatalog(key[0],key[1], metadata_only=True)
    xi = wfP.P.extract_param('xi')
    print best_matches[key], key[0], str(key[1]).replace(' ',''),   best_matches_masses[key][0], best_matches_masses[key][1], tmax, xi, Mf, af
  except:
     print "Skipping ", key
