#!/usr/bin/env python

from __future__ import print_function

import argparse
import math
import os
import re
import sys

import numpy
import scipy.optimize

import lfa

# Stop pyplot trying to use X.
import matplotlib
matplotlib.use('Agg')

import fftrv
import todcor

from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt

from makesky import *
from prepord import *
from read_spec import *

# Constants used by the script go here.

# Parameters.
PAR_ALPHA  = 0
PAR_VSINIA = 1
PAR_VSINIB = 2
NPARM      = 3

# Figure size based on pgplot.
figsize = (10.5, 7.8)  # inches

# Default velocity range.
velrange = 250  # km/s

# Limb darkening coefficients for rotational broadening kernel.
# These are from Phoenix mod. atm. for Teff=3224K, logg=5, Rc band,
# least-squares method, Claret et al. (2012) 2012A&A...546A..14C
ald = 0.4629
bld = 0.3212

def do_todcor(pdf, tmplaname, tmplbname, filename,
              tmpla_mbjd, tmpla_wave, tmpla_flux, tmpla_e_flux, tmpla_msk, tmpla_vbcv, tmpla_vrad,
              tmplb_mbjd, tmplb_wave, tmplb_flux, tmplb_e_flux, tmplb_msk, tmplb_vbcv, tmplb_vrad,
              mbjd, wave, flux, e_flux, msk, vbcv,
              vsinia, vsinib, alpha, order,
              restrict):

  # Extract order and clean.
  thistmpla_wave, thistmpla_flux, thistmpla_e_flux = prepord(order, tmpla_wave, tmpla_flux, tmpla_e_flux, tmpla_msk)
  thistmplb_wave, thistmplb_flux, thistmplb_e_flux = prepord(order, tmplb_wave, tmplb_flux, tmplb_e_flux, tmplb_msk)
  thiswave, thisflux, thise_flux = prepord(order, wave, flux, e_flux, msk)

  # Take off sky.
  ss = makesky(thistmpla_wave, thistmpla_flux, 4)

  thistmpla_flux -= ss

  ss = makesky(thistmplb_wave, thistmplb_flux, 4)

  thistmplb_flux -= ss

  thisrawflux = numpy.copy(thisflux)

  ss = makesky(thiswave, thisflux, 4)

  thisflux -= ss

  # Number of measurements (pixels).
  npix = len(thisflux)

  # Set up FFT sampling.  Double-lined solutions seem to benefit
  # from a bit less aggressive filtering, so reduce that.
  nbin = 32*npix
  filtl = int(npix/3)    # taper start 3 pix
  filth = int(npix/1.5)  # taper end 1.5 pix

  frv = fftrv.fftrv(nbin=nbin,
                    top_nrun=filtl,
                    nrun=filth,
                    pkfit=0,
                    t_emchop=emchop, s_emchop=emchop)

  tc = todcor.todcor(frv,
                     thistmpla_wave, thistmpla_flux,
                     thistmplb_wave, thistmplb_flux,
                     thiswave, thisflux,
                     vsinia, vsinib,
                     alpha,
                     ald, bld)

  va, vb, h = tc.run(minvela, maxvela,
                     minvelb, maxvelb,
                     restrict=restrict)

  if pdf is not None:
    fig = plt.figure(figsize=figsize)

    # Sampling: power of 2.
    ishft = int(math.log(nbin / 1024.0, 2))
    if ishft < 0:
      ishft = 0

    samp = 1<<ishft

    # Adjust endpoints to be on same grid.
    iamin = (tc.vel2ind(minvela) // samp) * samp + samp//2
    iamax = (tc.vel2ind(maxvela) // samp) * samp + samp//2
    ibmin = (tc.vel2ind(minvelb) // samp) * samp + samp//2
    ibmax = (tc.vel2ind(maxvelb) // samp) * samp + samp//2

    hh = tc.calc(iamin, iamax, samp,
                 ibmin, ibmax, samp)

    # Actual velocity range.
    amin = tc.ind2vel(iamin)
    amax = tc.ind2vel(iamax)
    bmin = tc.ind2vel(ibmin)
    bmax = tc.ind2vel(ibmax)

    plt.subplot(2, 2, 1)

    m = plt.imshow(hh.T,
                   interpolation='none',
                   extent=[amin, amax, bmax, bmin])
    plt.colorbar(m)
    plt.xlim(amin, amax)
    plt.ylim(bmin, bmax)
    plt.axvline(va, color='white')
    plt.axhline(vb, color='white')
    plt.ylabel("Secondary velocity (km/s)")

    # Full res around peak.
    iamin = tc.vel2ind(va) - 512
    iamax = tc.vel2ind(va) + 512
    ibmin = tc.vel2ind(vb) - 512
    ibmax = tc.vel2ind(vb) + 512

    hh = tc.calc(iamin, iamax, 1,
                 ibmin, ibmax, 1)

    amin = tc.ind2vel(iamin)
    amax = tc.ind2vel(iamax)
    bmin = tc.ind2vel(ibmin)
    bmax = tc.ind2vel(ibmax)

    plt.subplot(2, 2, 2)

    m = plt.imshow(hh.T,
                   interpolation='none',
                   extent=[amin, amax, bmax, bmin])
    plt.colorbar(m)
    plt.xlim(amin, amax)
    plt.ylim(bmin, bmax)
    plt.axvline(va, color='white')
    plt.axhline(vb, color='white')

    # Full res around peak.
    iamin = tc.vel2ind(va) - 64
    iamax = tc.vel2ind(va) + 64
    ibmin = tc.vel2ind(vb) - 64
    ibmax = tc.vel2ind(vb) + 64

    hh = tc.calc(iamin, iamax, 1,
                 ibmin, ibmax, 1)

    amin = tc.ind2vel(iamin)
    amax = tc.ind2vel(iamax)
    bmin = tc.ind2vel(ibmin)
    bmax = tc.ind2vel(ibmax)

    plt.subplot(2, 2, 3)

    m = plt.imshow(hh.T,
                   interpolation='none',
                   extent=[amin, amax, bmax, bmin])
    plt.colorbar(m)
    plt.xlim(amin, amax)
    plt.ylim(bmin, bmax)
    plt.axvline(va, color='white')
    plt.axhline(vb, color='white')
    plt.xlabel("Primary velocity (km/s)")
    plt.ylabel("Secondary velocity (km/s)")

    # Full res around peak.
    iamin = tc.vel2ind(va) - 16
    iamax = tc.vel2ind(va) + 16
    ibmin = tc.vel2ind(vb) - 16
    ibmax = tc.vel2ind(vb) + 16

    hh = tc.calc(iamin, iamax, 1,
                 ibmin, ibmax, 1)

    amin = tc.ind2vel(iamin)
    amax = tc.ind2vel(iamax)
    bmin = tc.ind2vel(ibmin)
    bmax = tc.ind2vel(ibmax)

    plt.subplot(2, 2, 4)

    m = plt.imshow(hh.T,
                   interpolation='none',
                   extent=[amin, amax, bmax, bmin])
    plt.colorbar(m)
    plt.xlim(amin, amax)
    plt.ylim(bmin, bmax)
    plt.axvline(va, color='white')
    plt.axhline(vb, color='white')
    plt.xlabel("Primary velocity (km/s)")

    pdf.savefig(fig)
    plt.close()

  return va, vb, h

def do_optimize(filelist, targlist, restrictlist, optlist, init_parm):

  nfit = len(optlist)

  # Standard bounds.
  std_lbound = numpy.array([ 0.01,   0.0,   0.0 ])
  std_ubound = numpy.array([ 1.0,  100.0, 100.0 ])

  # Make vector of fit parameters.
  vinit = init_parm[optlist]
  lbound = std_lbound[optlist]
  ubound = std_ubound[optlist]

  def fit_func(p):
    ssqh = 0

    trial_parm = numpy.copy(init_parm)
    trial_parm[optlist] = p

    l_x = numpy.empty([len(filelist)])
    l_y = numpy.empty([len(filelist)])
    l_w = numpy.empty([len(filelist)])

    for ifile, filename in enumerate(filelist):
      sp = targlist[ifile]

      xbest, ybest, hbest = do_todcor(None, tmplaname, tmplbname, filename,
                                      tmplasp.mbjd, tmplasp.wave, tmplasp.flux, tmplasp.e_flux, tmplasp.msk, tmplasp.vbcv, tmplasp.vrad,
                                      tmplbsp.mbjd, tmplbsp.wave, tmplbsp.flux, tmplbsp.e_flux, tmplbsp.msk, tmplbsp.vbcv, tmplbsp.vrad,
                                      sp.mbjd, sp.wave, sp.flux, sp.e_flux, sp.msk, sp.vbcv,
                                      trial_parm[PAR_VSINIA],
                                      trial_parm[PAR_VSINIB],
                                      trial_parm[PAR_ALPHA],
                                      order=rs.singleorder,
                                      restrict=restrictlist[ifile])

      l_x[ifile] = xbest
      l_y[ifile] = ybest
      l_w[ifile] = math.sqrt(1.0-hbest*hbest)

      ssqh += hbest*hbest

    print("Trial:", trial_parm, "ssqh =", ssqh, file=sys.stderr)

    return l_w

  res = scipy.optimize.least_squares(fit_func, vinit,
                                     bounds=(lbound, ubound),
                                     diff_step=0.01)

  vfit = res.x

  final_parm = numpy.copy(init_parm)
  final_parm[optlist] = vfit

  print("L_2/L_1 = {0:.4f}".format(final_parm[PAR_ALPHA]))
  print("vsinia  = {0:.4f}".format(final_parm[PAR_VSINIA]))
  print("vsinib  = {0:.4f}".format(final_parm[PAR_VSINIB]))

  return final_parm

def do_fixed(filelist, targlist, restrictlist, parm, dopdf=False):
  ssqh = 0

  donehdr = False
  
  for ifile, filename in enumerate(filelist):
    sp = targlist[ifile]
    
    if dopdf:
      basefile = stripname(filename)
      outfile = basefile + "_todcor.pdf"
      
      pdf = PdfPages(outfile)
    else:
      pdf = None

    # Velocities.
    xbest, ybest, hbest = do_todcor(pdf, tmplaname, tmplbname, filename,
                                    tmplasp.mbjd, tmplasp.wave, tmplasp.flux, tmplasp.e_flux, tmplasp.msk, tmplasp.vbcv, tmplasp.vrad,
                                    tmplbsp.mbjd, tmplbsp.wave, tmplbsp.flux, tmplbsp.e_flux, tmplbsp.msk, tmplbsp.vbcv, tmplbsp.vrad,
                                    sp.mbjd, sp.wave, sp.flux, sp.e_flux, sp.msk, sp.vbcv,
                                    parm[PAR_VSINIA],
                                    parm[PAR_VSINIB],
                                    parm[PAR_ALPHA],
                                    order=rs.singleorder,
                                    restrict=restrictlist[ifile])
    
    if pdf is not None:
      pdf.close()

    if not donehdr:
      print("#BJD-TDB     vrad1    vrad2    h         exptime")

      donehdr = True
      
    print("{0:12.4f} {1:8.3f} {2:8.3f} {3:9.6f} {4:7.2f}".format(lfa.ZMJD+sp.mbjd, xbest, ybest, hbest, sp.exptime))

    ssqh += hbest*hbest

  return ssqh

def get_range(s):
  ll = s.split(":")
  if len(ll) == 2:
    return map(float, ll)
  else:
    raise argparse.ArgumentError

ap = argparse.ArgumentParser()
ap.add_argument("alpha", help="light ratio star 2 / star 1 or search range")
ap.add_argument("vsinia", help="rotation of star 1")
ap.add_argument("vsinib", help="rotation of star 2")
ap.add_argument("tmpla", help="template for star 1 or @list of files to be stacked")
ap.add_argument("tmplb", help="template for star 2 or @list of files to be stacked")
ap.add_argument("filelist", metavar="file", nargs="+", help="spectrum file or @list of files to be stacked")
ap.add_argument("--alim", type=get_range, default=[-velrange, velrange], help="set velocity range star 1 --alim='min:max', default {0:d} to {1:d}".format(-velrange, velrange))
ap.add_argument("--blim", type=get_range, default=[-velrange, velrange], help="set velocity range star 2 --blim='min:max', default {0:d} to {1:d}".format(-velrange, velrange))
ap.add_argument("-a", help="automatically optimize parameters par1,par2,...")
ap.add_argument("-c", metavar="ra dec pmra pmdec plx vrad epoch", type=str, help="specify catalogue information for target")
ap.add_argument("-E", action="store_true", help="don't remove emission lines from spectrum")
ap.add_argument("-o", type=int, help="override order number used for analysis")
ap.add_argument("-p", action="store_true", help="save plots of 2-D correlation function")
ap.add_argument("-r", metavar="restrict", help="restrict sign of velocity difference")
ap.add_argument("-R", action="store_true", help="don't use cosmic rejection when stacking")

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

args = ap.parse_args()

# New read_spec structure.
rs = read_spec()

if args.o is not None:
  rs.overrideorder = args.o

emchop = not args.E

# Read templates.
tmplasp = rs.read_spec(args.tmpla, istmpl=1, wantstruct=True, doreject=not args.R)
tmplbsp = rs.read_spec(args.tmplb, istmpl=1, wantstruct=True, doreject=not args.R)

parm = numpy.empty([NPARM], dtype=numpy.double)

parm[PAR_VSINIA] = float(args.vsinia)
parm[PAR_VSINIB] = float(args.vsinib)

tmplaname = os.path.basename(args.tmpla)
tmplbname = os.path.basename(args.tmplb)

filelist = args.filelist
nf = len(args.filelist)

# Read targets.
targlist = [None] * nf
restrictlist = [None] * nf
for ifile, filename in enumerate(filelist):
  targlist[ifile] = rs.read_spec(filename, src_str=args.c, wantstruct=True, doreject=not args.R)
  if args.r is None:
    restrictlist[ifile] = None
  else:
    if ifile < len(args.r):
      restrictlist[ifile] = args.r[ifile]
    else:
      restrictlist[ifile] = None

minvela, maxvela = args.alim
minvelb, maxvelb = args.blim

m = re.match(r'^([^\-]+)\-([^\-]+)$', args.alpha)
if m:
  lrat1 = float(m.group(1))
  lrat2 = float(m.group(2))

  l_lrat = numpy.linspace(lrat1, lrat2, num=11, endpoint=True)
  l_ssqh = numpy.empty_like(l_lrat)
  
  for ilrat, thislrat in enumerate(l_lrat):
    print("L2/L1 =", thislrat)

    trial_parm = numpy.copy(parm)
    trial_parm[PAR_ALPHA] = l_lrat[ilrat]

    ssqh = do_fixed(filelist, targlist, restrictlist, trial_parm)
    l_ssqh[ilrat] = ssqh

  coef = numpy.polynomial.polynomial.polyfit(l_lrat, l_ssqh, 2)

  best_lrat = -0.5*coef[1]/coef[2]
  
  print("Best L2/L1 =", best_lrat, file=sys.stderr)
  
  plt.plot(l_lrat, l_ssqh, '.')
  
  xl = l_lrat[0]-0.05*(l_lrat[-1]-l_lrat[0])
  xh = l_lrat[-1]+0.05*(l_lrat[-1]-l_lrat[0])
  
  mod_x = numpy.linspace(xl, xh, 1001)
  mod_y = numpy.polynomial.polynomial.polyval(mod_x, coef)

  basefile = re.sub(r'_.*$', "", stripname(filelist[0]))
  outfile = basefile + "_todcor_lrat.pdf"

  print("Writing plot to", outfile)

  pdf = PdfPages(outfile)
  fig = plt.figure(figsize=figsize)

  plt.plot(mod_x, mod_y, color="blue")
  plt.plot(l_lrat, l_ssqh, 'o', color="black")
  plt.xlim(xl, xh)
  
  plt.axvline(best_lrat, color="black")
  
  plt.xlabel("L2/L1")
  plt.ylabel("Sum(h^2)")

  pdf.savefig(fig)
  plt.close()

  pdf.close()
else:
  parm[PAR_ALPHA]  = float(args.alpha)

  if args.a is not None:
    inplist = [ s.strip() for s in args.a.split(",") ]

    # List of pars we can optimize, used to clean up user inputs.
    pars = { "alpha":    PAR_ALPHA,
             "l2overl1": PAR_ALPHA,
             "vsinia":   PAR_VSINIA,
             "vsini1":   PAR_VSINIA,
             "vsinib":   PAR_VSINIB,
             "vsini2":   PAR_VSINIB }
    
    optlist = [ pars[s.lower()] for s in inplist ]

    parm = do_optimize(filelist, targlist, restrictlist, optlist, parm)

  ssqh = do_fixed(filelist, targlist, restrictlist, parm, args.p)

sys.exit(0)

