#!/usr/bin/env python

from __future__ import print_function

import sys
import re
import math
import numpy
import os

# Import and set up astropy.io.fits or pyfits.  In order to read
# IRAF-style wavelength solutions properly, it needs to be configured
# not to strip header whitespace, which is done by setting the
# variable pyfits.conf.strip_header_whitespace to False.
try:
  import astropy.io.fits as pyfits
except ImportError:
  import pyfits

pyfits.conf.strip_header_whitespace = False

from chiron import *
from poly import *

if len(sys.argv) < 2:
  print("Usage:\t", sys.argv[0], "file [...]")
  sys.exit(1)

for infile in sys.argv[1:]:
  fp = pyfits.open(infile)

  bjd, zb, exptime, wave, flux, e_flux, blaze = chiron_read(fp)

  mp = fp[0]
  hdr = mp.header.copy(strip=True)

  indir, inbase = os.path.split(infile)

  outfile = 'i' + re.sub(r'^a', '', inbase)

  # The GAINS keyword is missing the closing quote, presumably because
  # it was truncated incorrectly by IDL in the pipeline.  pyfits
  # doesn't like this so we just remove it here.
  del hdr["GAINS"]

  hdr["CTYPE1"] = "MULTISPE"
  hdr["CTYPE2"] = "MULTISPE"
  hdr["CDELT1"] = 1.0
  hdr["CDELT2"] = 1.0
  hdr["CD1_1"] = 1.0
  hdr["CD2_2"] = 1.0
  hdr["WCSDIM"] = 2
  hdr["WAT0_001"] = "system=multispec"
  hdr["WAT1_001"] = "wtype=multispec label=Wavelength units=angstroms"

  nord, npix = wave.shape

  wat2kw = "wtype=multispec"

  for iord in range(nord):
    # IRAF [-1, 1] normalised pixel coordinate system for Legendre.
    xmin = 1
    xmax = npix

    xpix = numpy.arange(xmin, xmax+1, dtype=numpy.double)

    n = (2*xpix - xmax - xmin) / (xmax - xmin)

    # Wavelengths.
    ww = wave[iord,:]

    wmin = ww[0]
    dwavg = (ww[-1] - ww[0]) / npix

    # Fit.
    if nord > 65:
      degr = 6  # Walter
    else:
      degr = 3  # Yale

    ncoef = degr+1

    coef = numpy.polynomial.legendre.legfit(n, ww, degr, full=False)
    wfit = numpy.polynomial.legendre.legval(n, coef)

    resid = (ww-wfit) / ww

    medresid, sigresid = medsig(resid)

#    print "Residual", medresid, sigresid

    if sigresid > 1.0e-6:
      print("ERROR: residual is too large")
      plt.plot(n, resid)
      plt.show()
      sys.exit(0);

    # Form string.
    specn = "{0:d} {0:d} 2 {1:.17e} {2:.17e} {3:d} 0. 0. 0. 1. 0. 2 {4:d} {5:d} {6:d}".format(iord+1, wmin, dwavg, npix, ncoef, xmin, xmax)

    for icoef in range(ncoef):
      specn += " {0:.17e}".format(coef[icoef])

    wat2kw += " spec{0:d} = \"".format(iord+1) + specn + "\""

  # Chop into 68 char keywords.
  maxchars = 68

  nwat2kw = len(wat2kw) // maxchars
  if len(wat2kw) % maxchars:
    nwat2kw += 1

  for iwat2kw in range(nwat2kw):
    key = "WAT2_{0:03d}".format(iwat2kw+1)
    val = wat2kw[iwat2kw*maxchars:(iwat2kw+1)*maxchars]

    hdr[key] = val

  outhdus = [ pyfits.PrimaryHDU(flux, header=hdr) ]
  ohl = pyfits.HDUList(outhdus)
  ohl.writeto(outfile, clobber=True)
