#!/usr/bin/env python

from __future__ import print_function

import argparse
import math
import sys

import numpy
import emcee

import eb

import matplotlib

from nlfit import *

import fussyastronomer

# Figure size.
figsize = (8.0, 6.0)  # inches

# Deal with command-line arguments.
ap = argparse.ArgumentParser()
ap.add_argument("pset", help="parameter file")
ap.add_argument("vels", help="velocities")
ap.add_argument("-3", action="store_true", dest="triple", help="read and plot third star from TRICOR file")
ap.add_argument("-b", type=int, help="set number of burn-in samples")
ap.add_argument("-m", help="run MCMC storing output to file M")
ap.add_argument("-n", type=int, help="set number of MCMC samples to N")
ap.add_argument("-u", action="store_true", help="plot unfolded phase")
ap.add_argument("-w", type=int, help="set number of walkers")

if len(sys.argv) == 1:
  ap.print_help(sys.stderr)
  sys.exit(1)

args = ap.parse_args()

# Type 42 fonts needed for publication.
matplotlib.rcParams["pdf.fonttype"] = 42
matplotlib.rcParams["ps.fonttype"] = 42
matplotlib.rcParams["figure.subplot.left"] = 0.15
matplotlib.rcParams["figure.subplot.right"] = 0.85
matplotlib.rcParams["font.size"] = 14

# Make sure matplotlib doesn't try to use X11 if we're running
# in non-interactive mode.
if args.m is not None:
  matplotlib.use('Agg')

import matplotlib.pyplot as plt

# Fixed random number seed to make results reproducible.
numpy.random.seed(42)

# Read input file.
pset_names = {}
pset_lines = []

with open(args.pset, "r") as fp:
  for line in fp:
    # Remove comments and trim white space.
    ll = line.split("#", 1)
    ls = ll[0].strip()

    if ls == "":
      continue

    # Now split into fields.
    ll = ls.split()
    if len(ll) < 3:
      print("could not understand:", line, file=sys.stderr)
      sys.exit(1)

    name = ll[0]
    pset_names[name] = ll

    pset_lines.append(ll)

# Detect options.
mode_g = 0
mode_s = 0

if "gamma1" in pset_names and "gamma2" in pset_names:
  mode_g = 1

if "sigma1" in pset_names and "sigma2" in pset_names:
  mode_s = 1

# Build main parameter vector and metadata vectors.
parm = numpy.zeros(eb.NPAR+3+mode_g+mode_s, dtype=numpy.double)
parm_names = list(eb.parnames)
parm_texsym = list(eb.partexsym)
parm_units = list(eb.parunits)

npar = eb.NPAR

PAR_GAMMA1 = npar
npar += 1

if mode_g:
  parm_names.append("gamma1")
  parm_texsym.append("\\gamma_1")
else:
  parm_names.append("gamma")
  parm_texsym.append("\\gamma")

parm_units.append("km/s")

if mode_g:
  PAR_GAMMA2 = npar
  npar += 1

  parm_names.append("gamma2")
  parm_texsym.append("\\gamma_2")
  parm_units.append("km/s")
else:
  PAR_GAMMA2 = None

PAR_KTOT = npar
npar += 1

parm_names.append("K1+K2")
parm_texsym.append("(K_1+K_2)")
parm_units.append("km/s")

PAR_S1 = npar
npar += 1

if mode_s:
  parm_names.append("sigma1")
  parm_texsym.append("\\sigma_1")
else:
  parm_names.append("sigma")
  parm_texsym.append("\\sigma")

parm_units.append("km/s")

if mode_s:
  PAR_S2 = npar
  npar += 1

  parm_names.append("sigma2")
  parm_texsym.append("\\sigma_2")
  parm_units.append("km/s")
else:
  PAR_S2 = None

parm_units[eb.PAR_T0] = "BJD"

name2iparm = {}

for iparm, name in enumerate(parm_names):
  name2iparm[name] = iparm

# Derived parameters.
dermap = [eb.PAR_E, eb.PAR_OMEGA, eb.PAR_A, eb.PAR_MTOT, eb.PAR_M1, eb.PAR_M2]

dernames = ["e", "omega", "a sin i", "(M_1+M_2) sin^3 i", "M_1 sin^3 i", "M_2 sin^3 i", "Tperi"]
dertexsym = ["e", "\\omega", "a \\sin i", "(M_1+M_2) \\sin^3 i", "M_1 \\sin^3 i", "M_2 \\sin^3 i", "T_{\\rm peri}"]
derunits = ["", "deg", "AU", "Msol", "Msol", "Msol", "BJD"]

nder = len(dernames)

# Initialize a couple of parameters we don't use to sensible values
# to avoid floating point exceptions in the model calculation.
parm[eb.PAR_RASUM] =   0.001  # dummy value
parm[eb.PAR_RR]    =   0.001  # dummy value

# Supply a default for error scaling parameter(s).
parm[PAR_S1]       =   0.15   # km/s

if mode_s:
  parm[PAR_S2] = parm[PAR_S1]

# Load data.
if args.triple:
  vels = numpy.genfromtxt(args.vels,
                          dtype={"names": ("bjd", "v1", "v2", "v3", "h"),
                                 "formats": ("f8", "f8", "f8", "f8", "f8") },
                          loose=False)

  v3 = vels["v3"]
else:
  vels = numpy.genfromtxt(args.vels,
                          dtype={"names": ("bjd", "v1", "v2", "h"),
                                 "formats": ("f8", "f8", "f8", "f8") },
                          loose=False)

  v3 = None

tbase = math.floor(vels["bjd"][0])

bjd = vels["bjd"] - tbase
v1 = vels["v1"]
v2 = vels["v2"]
hh = vels["h"]

nobs = len(bjd)

# Bring in fit parameters.
plist = []
plist_names = []
plist_texsym = []
plist_units = []
fixed = []

for ll in pset_lines:
  name = ll[0]
  value = float(ll[1])
  vary = int(ll[2])

  if len(ll) > 3:
    sigma = float(ll[3])
  else:
    sigma = None

  if name in name2iparm:
    iparm = name2iparm[name]

    # Set value.
    if iparm == eb.PAR_T0:
      parm[iparm] = value - tbase
    else:
      parm[iparm] = value

    # Add to parameter list if we're varying.
    if vary > 0:
      plist.append(iparm)
      plist_names.append(name)
      plist_texsym.append(parm_texsym[iparm])
      plist_units.append(parm_units[iparm])

      if vary == 1:  # 1=in LM and MC, 2=in MC only
        fixed.append(False)
      elif vary == 2:
        fixed.append(True)
      else:
        print("unknown vary value:", vary, file=sys.stderr)
        sys.exit(1)

  else:
    print("unknown parameter:", name, file=sys.stderr)
    sys.exit(1)

# Convert these to numpy arrays.
plist = numpy.array(plist, dtype=int)
fixed = numpy.array(fixed, dtype=bool)

# Store this for later.
errguess1 = parm[PAR_S1]

if mode_s:
  errguess2 = parm[PAR_S2]
else:
  errguess2 = errguess1

# Format arrays for fitting.
bjdfit = numpy.empty([2*nobs], dtype=numpy.double)
vfit = numpy.empty_like(bjdfit)
e_vfit = numpy.empty_like(bjdfit)

bjdfit[0:nobs] = bjd
bjdfit[nobs:2*nobs] = bjd

vfit[0:nobs] = v1
vfit[nobs:2*nobs] = v2

e_vfit[0:nobs] = errguess1 / hh
e_vfit[nobs:2*nobs] = errguess2 / hh

typ = numpy.empty_like(vfit, dtype=numpy.uint8)
typ[0:nobs] = eb.OBS_VRAD1
typ[nobs:2*nobs] = eb.OBS_VRAD2

# L-M objective function.
def fit_func (trial, ymod):
  ourparm = numpy.copy(parm)
  ourparm[plist] = trial

  eb.model(ourparm, bjdfit, typ, 0, out=ymod)
  ymod *= ourparm[PAR_KTOT] / (1.0 + ourparm[eb.PAR_Q])
  ymod[0:nobs] *= ourparm[eb.PAR_Q]
  ymod[nobs:2*nobs] *= -1

  if mode_g:
    ymod[0:nobs] += ourparm[PAR_GAMMA1]
    ymod[nobs:2*nobs] += ourparm[PAR_GAMMA2]
  else:
    ymod += ourparm[PAR_GAMMA1]

def vecok (ourparm):
  ecc = math.hypot(ourparm[eb.PAR_ECOSW], ourparm[eb.PAR_ESINW])
  return(ecc < 1.0 and
         ourparm[PAR_S1] > 0 and
         (not mode_s or ourparm[PAR_S2] > 0))

# Function returning log of posterior probability for MCMC.
def logpost (trial, x, y, yerr):
  # Parameter vector.
  ourparm = numpy.copy(parm)
  ourparm[plist] = trial

  # Check ranges.
  if not vecok(ourparm):
#    print "not vecok for", ourparm
    return -numpy.inf

  # Likelihood.
  ymod = eb.model(ourparm, x, typ, 0)
  ymod *= ourparm[PAR_KTOT] / (1.0 + ourparm[eb.PAR_Q])
  ymod[0:nobs] *= ourparm[eb.PAR_Q]
  ymod[nobs:2*nobs] *= -1

  if mode_g:
    ymod[0:nobs] += ourparm[PAR_GAMMA1]
    ymod[nobs:2*nobs] += ourparm[PAR_GAMMA2]
  else:
    ymod += ourparm[PAR_GAMMA1]
  
  wt = numpy.empty_like(ymod)
  wt[0:nobs] = (hh / ourparm[PAR_S1])**2

  if mode_s:
    wt[nobs:2*nobs] = (hh / ourparm[PAR_S2])**2
  else:
    wt[nobs:2*nobs] = (hh / ourparm[PAR_S1])**2

  loglike = -0.5*(numpy.sum(wt*(y-ymod)**2 - numpy.log(wt)))

  # Prior: Modified Jeffreys on error scale parameter.
  logprior = numpy.log(ourparm[PAR_S1] + 0.1*errguess1)

  if mode_s:
    logprior += numpy.log(ourparm[PAR_S2] + 0.1*errguess2)

  # Convert to uniform prior in eccentricity.
  ecc = numpy.hypot(ourparm[eb.PAR_ECOSW], ourparm[eb.PAR_ESINW])
  if ecc > 0:
    logprior += numpy.log(ecc)

  return loglike - logprior

def plot_chain(samp, outfile):
  fig = plt.figure(figsize=figsize)

  for idim in range(ndim):
    plt.subplot(ndim, 1, idim+1)

    for iwalk in range(nwalk):
      plt.plot(samp.chain[iwalk,:,idim])

    plt.ylabel(plist_names[idim])

  fig.savefig(outfile)

ndim = len(plist)

vfinal, verr, chisq, ndof = nlfit(fit_func, parm[plist], fixed, vfit, e_vfit)

parm[plist] = vfinal

# Fit parameters.
print("L-M parameters:")

for iparm in range(ndim):
  value = vfinal[iparm]

  if plist[iparm] == eb.PAR_T0:
    value += tbase

  print("{0:<16} {1:14.6f} +/- {2:8.6f} {3}".format(plist_names[iparm],
                                                    value,
                                                    verr[iparm],
                                                    plist_units[iparm]))

print('chi squared =', chisq, 'ndof =', ndof)

if args.m is not None:
  # Init MC
  if args.w is not None:
    nwalk = args.w
  else:
    nwalk = 100

  samp = emcee.EnsembleSampler(nwalk, ndim, logpost, args=(bjdfit, vfit, e_vfit))

  cent = numpy.copy(vfinal)
  pert = 3*numpy.copy(verr)

  for iparm in range(ndim):
    jparm = plist[iparm]

    if fixed[iparm]:
      if jparm == PAR_GAMMA1 or jparm == PAR_KTOT:
        pert[iparm] = 1.0
      elif mode_g and jparm == PAR_GAMMA2:
        pert[iparm] = 1.0
      elif jparm == PAR_S1:
        pert[iparm] = 0.1 * errguess1
      elif mode_s and jparm == PAR_S2:
        pert[iparm] = 0.1 * errguess2
      else:
        pert[iparm] = 0.01

  pos = [None]*nwalk

  for iwalk in range(nwalk):
    trial = None

    for isim in range(100):
      trial = cent + pert * numpy.random.randn(ndim)

      trialparm = numpy.copy(parm)
      trialparm[plist] = trial

      isok = vecok(trialparm)
      if isok:
        break
      
    if not isok:
      raise RuntimeError("couldn't find valid starting point")

    pos[iwalk] = trial

  if args.n is not None:
    nsim = args.n
  else:
    nsim = 50000

  if args.b is not None:
    nburn = args.b
  else:
    nburn = nsim // 5
    if nburn < 1000:
      nburn = 1000

  print("MCMC using", nwalk, "walkers nburn =", nburn, "nsim =", nsim)
  print("Will produce", nburn*nwalk, "burn-in and", nsim*nwalk, "final samples")

  pos, prob, state = samp.run_mcmc(pos, nburn,
                                   rstate0=numpy.random.get_state())

  plot_chain(samp, args.m + "_burn.png")

  samp.reset()

  pos, prob, state = samp.run_mcmc(pos, nsim)

  plot_chain(samp, args.m + "_chains.png")

  print("Acceptance fraction:", samp.acceptance_fraction)

  try:
    print("Autocorrelation lengths:", samp.acor)
  except:
    pass

  ch = samp.flatchain

  numpy.savez(args.m,
              parm=parm,
              parm_names=parm_names,
              parm_texsym=parm_texsym,
              parm_units=parm_units,
              plist=plist,
              plist_names=plist_names,
              plist_texsym=plist_texsym,
              plist_units=plist_units,
              fixed=fixed,
              tbase=tbase,
              nwalk=nwalk,
              nburn=nburn,
              nsim=nsim,
              ch=ch)

  lfp = open(args.m + "_orbit.txt", "w")
  tfp = open(args.m + "_orbit.tex", "w")

  print("MC parameters:", file=lfp)
  print("MC parameters:")

  print("\\hline", file=tfp)
  print("MCMC parameters\\\\", file=tfp)
  print("\\hline", file=tfp)

  vmed = numpy.empty([ndim])
  vad = numpy.empty([ndim])
  vndp = numpy.empty([ndim], dtype=int)

  ecc_kh_ndp = None
  sigma_ndp = None

  for iparm in range(ndim):
    vec = ch[:,iparm]

    med = numpy.median(vec)
    ad = numpy.absolute(vec-med)

    err = numpy.percentile(ad, 68.3)

    ndp = fussyastronomer.ndp(err)

    vmed[iparm] = med
    vad[iparm] = err
    vndp[iparm] = ndp

    if plist[iparm] == eb.PAR_T0:
      med += tbase

    line = "{0:<16} {1:14.6f} +/- {2:8.6f} {3}".format(plist_names[iparm],
                                                       med,
                                                       err,
                                                       plist_units[iparm])

    print(line, file=lfp)
    print(line)

    # Special cases of n(dp) we want to be the same.
    if plist[iparm] == eb.PAR_ECOSW or plist[iparm] == eb.PAR_ESINW:
      if ecc_kh_ndp is None or ndp > ecc_kh_ndp:
        ecc_kh_ndp = ndp
    elif plist[iparm] == PAR_S1 or (PAR_S2 is not None and plist[iparm] == PAR_S2):
      if sigma_ndp is None or ndp > sigma_ndp:
        sigma_ndp = ndp

  # TeX output.
  for iparm in range(ndim):
    med = vmed[iparm]
    err = vad[iparm]

    if plist[iparm] == eb.PAR_T0:
      med += tbase

    if plist[iparm] == eb.PAR_ECOSW or plist[iparm] == eb.PAR_ESINW:
      ndp = ecc_kh_ndp
    elif plist[iparm] == PAR_S1 or (PAR_S2 is not None and plist[iparm] == PAR_S2):
      ndp = sigma_ndp
    else:
      ndp = vndp[iparm]

    parstr = "$" + plist_texsym[iparm] + "$"
    if plist_units[iparm] != "":
      utex = plist_units[iparm]

      if utex == "Msol":
        utex = "${\\rm M}_\\odot$"

      parstr += " (" + utex + ")"

    texline = "{0:<36} & ${1:s} \\pm {2:s}$ \\\\".format(parstr,
                                                         fussyastronomer.format(med, ndp),
                                                         fussyastronomer.format(err, ndp))

    print(texline, file=tfp)

  # Argument and phase of periastron need special treatment.
  parm[plist] = vmed

  theomega = math.degrees(math.atan2(parm[eb.PAR_ESINW], parm[eb.PAR_ECOSW]))
  if theomega < 0.0:  # wrap to conventional range [0,360)
    theomega += 360.0
  thephiperi = eb.phiperi(parm)

  # Derived parameters.
  nsamp = ch.shape[0]

  vder = numpy.empty([nsamp, nder])
  mkvec = numpy.empty([nsamp])

  for isamp in range(nsamp):
    thissamp = ch[isamp]
    parm[plist] = thissamp

    thisvder = eb.getvder(parm, parm[PAR_GAMMA1], parm[PAR_KTOT])

    # Change units of semimajor axis from Rsol to AU for consistency
    # with SB1 program outputs.
    thisvder[eb.PAR_A] *= eb.RSUN / eb.AU

    # Force correct revolution for omega by wrapping delta into (-pi,pi].
    domega = thisvder[eb.PAR_OMEGA] - theomega
    n = round(domega / 360.0)
    thisvder[eb.PAR_OMEGA] = theomega + domega - n * 360.0

    # Same for phiperi.
    phiperi = eb.phiperi(parm)

    dphi = phiperi - thephiperi
    n = round(dphi)
    phiperi = thephiperi + dphi - n

    tperi = parm[eb.PAR_T0] + parm[eb.PAR_P] * phiperi

    vder[isamp,0:len(dermap)] = thisvder[dermap]
    vder[isamp,len(dermap)] = tperi

  numpy.savez(args.m + "_der",
              dermap=dermap,
              dernames=dernames,
              dertexsym=dertexsym,
              derunits=derunits,
              tbase=tbase,
              nwalk=nwalk,
              nburn=nburn,
              nsim=nsim,
              vder=vder)

  print("Derived parameters:", file=lfp)
  print("Derived parameters:")

  print("\\hline", file=tfp)
  print("Derived parameters\\\\", file=tfp)
  print("\\hline", file=tfp)

  vdermed = numpy.empty([nder])
  vderad = numpy.empty([nder])
  vderndp = numpy.empty([nder], dtype=int)

  mass_ndp = None

  for ipar, name in enumerate(dernames):
    unit = derunits[ipar]

    vec = vder[:,ipar]

    med = numpy.median(vec)
    ad = numpy.absolute(vec-med)

    err = numpy.percentile(ad, 68.3)

    ndp = fussyastronomer.ndp(err)

    vdermed[ipar] = med
    vderad[ipar] = err
    vderndp[ipar] = ndp

    if name == "Tperi":
      med += tbase

    line = "{0:<18} {1:12.6f} +/- {2:8.6f} {3}".format(name, med, err, unit)

    if name == "e" and med < 5*err:
      # Also give upper limit.
      ul = numpy.percentile(vec, 95.0)
      line += " (95% CL < {0:.7f})".format(ul)

    print(line, file=lfp)
    print(line)

    # Special cases of n(dp) we want to be the same.
    if ipar < len(dermap) and (dermap[ipar] == eb.PAR_MTOT or dermap[ipar] == eb.PAR_M1 or dermap[ipar] == eb.PAR_M2):
      if mass_ndp is None or ndp > mass_ndp:
        mass_ndp = ndp

  # TeX output.
  for ipar in range(nder):
    unit = derunits[ipar]

    med = vdermed[ipar]
    err = vderad[ipar]

    if dernames[ipar] == "Tperi":
      med += tbase

    if ipar < len(dermap) and (dermap[ipar] == eb.PAR_MTOT or dermap[ipar] == eb.PAR_M1 or dermap[ipar] == eb.PAR_M2):
      ndp = mass_ndp
    else:
      ndp = vderndp[ipar]

    parstr = "$" + dertexsym[ipar] + "$"
    if unit != "":
      utex = unit

      if utex == "Msol":
        utex = "${\\rm M}_\\odot$"

      parstr += " (" + utex + ")"

    texline = "{0:<36} & ${1:s} \\pm {2:s}$ \\\\".format(parstr,
                                                         fussyastronomer.format(med, ndp),
                                                         fussyastronomer.format(err, ndp))

    print(texline, file=tfp)


  lfp.close()
  tfp.close()

  # Update parameters.
  parm[plist] = vmed
  vfinal = vmed
  verr = vad

  ndof = vfit.size - vmed.size

  if mode_s:
    print("Median rms(1) =", numpy.median(parm[PAR_S1]/hh), "km/s")
    print("Median rms(2) =", numpy.median(parm[PAR_S2]/hh), "km/s")
  else:
    print("Median rms =", numpy.median(parm[PAR_S1]/hh), "km/s")

else:
  # Derived parameters for L-M model.
  thisvder = eb.getvder(parm, parm[PAR_GAMMA1], parm[PAR_KTOT])

  # Change units of semimajor axis from Rsol to AU for consistency
  # with SB1 program outputs.
  thisvder[eb.PAR_A] *= eb.RSUN / eb.AU

  phiperi = eb.phiperi(parm)

  tperi = parm[eb.PAR_T0] + parm[eb.PAR_P] * phiperi

  # Make vector.
  vder = numpy.empty([nder])

  vder[0:len(dermap)] = thisvder[dermap]
  vder[len(dermap)] = tperi

  print("Derived parameters:")

  for ipar, name in enumerate(dernames):
    unit = derunits[ipar]

    val = vder[ipar]

    if name == "Tperi":
      val += tbase

    line = "{0:<18} {1:12.6f} {2}".format(name, val, unit)

    print(line)

# Phases of observations.
tmp = (bjdfit - parm[eb.PAR_T0]) / parm[eb.PAR_P]
tmporig = (bjd - parm[eb.PAR_T0]) / parm[eb.PAR_P]

if args.u:
  phi = tmp
  phiorig = tmporig

  phimin = numpy.min(phi)
  phimax = numpy.max(phi)
  
  phimin -= 0.05*(phimax-phimin)
  phimax += 0.05*(phimax-phimin)

else:
  phi = tmp - numpy.floor(tmp)
  phiorig = tmporig - numpy.floor(tmporig)

  phimin = 0.0
  phimax = 1.0

print("Last point phi =", phi[-1])

# Plot.
phiplot = numpy.linspace(phimin, phimax, 1000)
typplot = numpy.empty_like(phiplot, dtype=numpy.uint8)

e_vfit[0:nobs] = parm[PAR_S1] / hh

if mode_s:
  e_vfit[nobs:2*nobs] = parm[PAR_S2] / hh
else:
  e_vfit[nobs:2*nobs] = parm[PAR_S1] / hh

typplot.fill(eb.OBS_VRAD1)
yplot1 = eb.model(parm, phiplot, typplot, eb.FLAG_PHI)
yplot1 = parm[PAR_GAMMA1] + yplot1 * parm[PAR_KTOT] * parm[eb.PAR_Q] / (1.0 + parm[eb.PAR_Q])

typplot.fill(eb.OBS_VRAD2)
yplot2 = eb.model(parm, phiplot, typplot, eb.FLAG_PHI)
if mode_g:
  yplot2 = parm[PAR_GAMMA2] - yplot2 * parm[PAR_KTOT] / (1.0 + parm[eb.PAR_Q])
else:
  yplot2 = parm[PAR_GAMMA1] - yplot2 * parm[PAR_KTOT] / (1.0 + parm[eb.PAR_Q])

fig = plt.figure(figsize=figsize)

plt.subplot(2, 1, 1)

plt.plot(phiplot, yplot1, color="blue")
plt.plot(phiplot, yplot2, color="green")

if args.triple:
  mean3 = numpy.average(v3, weights=hh*hh)
  resid3 = v3 - mean3

  print("gamma3 =", mean3, "km/s")

  if mode_s:
    tmp = resid3 * hh
    chisq = numpy.sum(tmp*tmp)
    err3 = math.sqrt(chisq / (hh.size-1))

    print("sigma3 =", err3, "km/s")
  else:
    err3 = parm[PAR_S1] / hh

  plt.axhline(mean3, color="red")

plt.errorbar(phi[0:nobs], vfit[0:nobs], yerr=e_vfit[0:nobs], fmt="o", color="blue")
plt.errorbar(phi[nobs:2*nobs], vfit[nobs:2*nobs], yerr=e_vfit[nobs:2*nobs], fmt="o", color="green")

if args.triple:
  plt.errorbar(phiorig, v3, yerr=err3, fmt="o", color="red")

plt.xlim(phimin, phimax)

plt.ylabel("Vrad (km/s)")

plt.subplot(2, 1, 2)

ymod = numpy.empty_like(vfit)
fit_func(vfinal, ymod)

resid = vfit-ymod

chisq = numpy.sum((resid / e_vfit)**2)

print("Final chi squared =", chisq, "ndof =", ndof)

veldiff = ymod[0:nobs] - ymod[nobs:2*nobs]

rstr = ""

for vel in veldiff:
  if vel > 0:
    rstr += "p"
  elif vel < 0:
    rstr += "n"
  else:
    rstr += "a"

print("String for todcor -r =", rstr)

rmin1 = numpy.min(resid[0:nobs]-e_vfit[0:nobs])
rmax2 = numpy.max(resid[nobs:2*nobs]+e_vfit[nobs:2*nobs])

if args.triple:
  rmin2 = numpy.min(resid[nobs:2*nobs]-e_vfit[nobs:2*nobs])
  rmax3 = numpy.max(resid3+err3)
  
  roff2 = 1.1*min(rmin1 - rmax2, rmin2 - rmax3)
  roff3 = 2*roff2
else:
  roff2 = 1.1*(rmin1 - rmax2)

plt.errorbar(phi[0:nobs], resid[0:nobs], yerr=e_vfit[0:nobs], fmt="o", color="blue")
plt.axhline(0, color="blue")
plt.errorbar(phi[nobs:2*nobs], roff2+resid[nobs:2*nobs], yerr=e_vfit[nobs:2*nobs], fmt="o", color="green")
plt.axhline(roff2, color="green")

if args.triple:
  plt.errorbar(phiorig, roff3+resid3, yerr=err3, fmt="o", color="red")
  plt.axhline(roff3, color="red")

plt.xlim(phimin, phimax)

plt.xlabel("Phase")
plt.ylabel("Residual + const (km/s)")

#plt.tight_layout()

if args.m is None:
  plt.show()
else:
  fig.savefig(args.m + "_orbit.pdf")
  fig.savefig(args.m + "_orbit.eps")
  plt.close()

