# GP_FIT.PY
# Dominic Ford
# 27/04/2006

import gp_eval
import gp_datafile
from gp_autocomplete import *

import sys
import re
from math import *
from scipy.optimize import fmin

pass_function = ''
pass_vialist  = []
pass_xname    = ''
pass_vars     = {}
pass_funcs    = {}
pass_dataset  = []
pass_rows     = 0

xmin = 0 ; xmax = 0 ; xrange = 0
ymin = 0 ; ymax = 0 ; yrange = 0


# FIT_RESIDUAL(): Evaluates the residual between function and dataset

def fit_residual(x):
  global pass_function, pass_vialist, pass_xname, pass_vars, pass_funcs, pass_dataset, pass_rows
  global xmin, xmax, xrange, ymin, ymax, yrange

  local_vars  = pass_vars.copy()

  for i in range(len(pass_vialist)): # Set via variables
    local_vars[pass_vialist[i]] = x[i]

  accumulator = 0.0
  for i in range(pass_rows):
    if (xrange == 1):
      if (pass_dataset[i][0] < xmin): continue
      if (pass_dataset[i][0] > xmax): continue
    if (yrange == 1):
      if (pass_dataset[i][1] < ymin): continue
      if (pass_dataset[i][1] > ymax): continue
    local_vars[pass_xname] = pass_dataset[i][0]
    accumulator = accumulator + pow(pass_dataset[i][1] - gp_eval.gp_eval(pass_function,local_vars,pass_funcs), 2.0)

  return accumulator

# DIRECTIVE_FIT(): Implements the "fit" directive

def directive_fit(line, linelist, vars, funcs):
  global pass_function, pass_vialist, pass_xname, pass_vars, pass_funcs, pass_dataset, pass_rows
  global xmin, xmax, xrange, ymin, ymax, yrange

  # FIRST OF ALL, WE NEED TO READ THE INPUT PARAMETERS FROM THE COMMANDLINE

  datafile = ''
  using    = '1:2'
  every    = 1
  index    = -1 # GnuPlot API doesn't have an index value for "plot all indices", so I define it to be -1.
  xmin     = 0 ; xmax = 0 ; xrange = 0 # Reset range variables
  ymin     = 0 ; ymax = 0 ; yrange = 0
  vialist  = []
  state    = 0 # 0 = function not yet read ; 1 = after function read, last word was not keyword ; 2 = using read ; 3 = every read ; 4 = via read

  for i in range(1,len(linelist),1):
    item = linelist[i].strip()

    if (state < 2):
      if autocomplete(item, "using", 1): state = 2 ; continue
      if autocomplete(item, "every", 1): state = 3 ; continue
      if autocomplete(item, "index", 1): state = 4 ; continue
      if autocomplete(item, "via"  , 1): state = 5 ; continue

    if (state == 4): # set 'index' setting
      try:
        index = int(item)
      except:
        print "Error:" , sys.exc_info()[1], "(" , sys.exc_info()[0] , ")"
        print "'index' keyword should be followed by an integer"
        return # Error
      state = 1
      continue

    if (state == 3): # set 'every' setting
      try:
        every = int(item)
      except:
        print "Error:" , sys.exc_info()[1], "(" , sys.exc_info()[0] , ")"
        print "'every' keyword should be followed by an integer"
        return # Error
      state = 1
      continue

    if (state == 2): # read 'using' string
      using = item
      state = 1
      continue

    if (state == 5): # read via list
      commasplit = gp_eval.gp_split(item,",")
      for var in commasplit:
        var_clean = var.strip()
        test = re.match(r"""^[A-Za-z]\w*$""",var_clean)
        if (test == None):
          print "Illegal variable name '%s' in via. Must contain only alphanumeric characters."%var_clean
          return
        vialist.append(var_clean)
      continue

    if ((item[0] == '[') and (state == 0)): # set range (BEFORE function)
      while ((len(item) != 0) and (item[0] == '[')):
        test = re.match(r"""^\s*\[([^:\]]*)((:)|( *to *))([^:\]]*)\]\s*(.*)$""",item)
        if (test == None):
          print "Could not read range '%s'. Use format: [0:1]."%item
          return # Error
        else:
          try:
            if (xrange == 0): # set xrange
              xmin   = gp_eval.gp_eval(test.group(1),vars,funcs)
              xmax   = gp_eval.gp_eval(test.group(5),vars,funcs)
              xrange = 1 ; item = test.group(6) ; continue
            elif (yrange == 0): # set yrange
              ymin   = gp_eval.gp_eval(test.group(1),vars,funcs)
              ymax   = gp_eval.gp_eval(test.group(5),vars,funcs)
              yrange = 1 ; item = test.group(6) ; continue
            else:
              print "Too many ranges. Specify a maximum of two ranges, one for x, second for y."
              raise ValueError # User has specified too many ranges
          except:
            return # Error
      continue

    if (state == 0): # Receive function
      test = re.match("""^\s*[A-Za-z]\w*\(\s*([A-Za-z]\w*)\s*\)\s*$""",item)
      if (test == None):
        print "Badly formed function '%s'. Must be a function of only one variable"%item
        return # Error
      pass_function = item
      pass_xname    = test.group(1)
      state = 1
      continue

    if ((state == 1) and (item[0] == "'")):
      test = re.match(r"""^'(.*)'$""",item)
      if (test == None):
        print """Badly formed datafile name "%s"."""%item
        return # Error
      datafile = test.group(1)
      continue

    print "Syntax Error: Unrecognised word '%s'"%item
    return # Error

    # We have now read all of our commandline parameters, and are ready to start fitting

  try:
    (rows,columns,datagrid) = gp_datafile.gp_dataread(datafile, index, using, every, vars, funcs)[0]
  except:
    print "Error reading input datafile:" , sys.exc_info()[1], "(" , sys.exc_info()[0] , ")"
    return # Error

  if (columns < 2):
    print "Error: cannot fit to a one-dimensional dataset!"
    return # Error

  if (columns > 2):
    print "Warning: GnuPlot+ does not presently support error bars when fitting. Fit will assume uniform errors for now."

  pass_vialist = vialist
  pass_vars    = vars
  pass_funcs   = funcs
  pass_dataset = datagrid
  pass_rows    = rows

  x = []
  for i in range(len(vialist)): x.append(1.0)

  try:
    x = fmin(fit_residual, x)
  except:
    print "Numerical Error:" , sys.exc_info()[1], "(" , sys.exc_info()[0] , ")"
    return # Error

  print "Best fit parameters were:"
  for i in range(len(vialist)): # Set via variables
    o = x[i]
    if ((fabs(o) < 1e10) and (fabs(o) > 1e-6)): print "%s = %f"%(vialist[i],o)
    else:                                       print "%s = %e"%(vialist[i],o)
    vars[vialist[i]] = x[i]

  return # Done!!!

