#!/usr/bin/env python
# File: dataset_libs.py
# Author: Kiri Wagstaff, 5/7/13
#
# Readers and plotters for LIBS (ChemCam) data sets
#
# Copyright 2013-2015, by the California Institute of Technology. ALL
# RIGHTS RESERVED.  United States Government Sponsorship
# acknowledged. Any commercial use must be negotiated with the Office
# of Technology Transfer at the California Institute of Technology.
#
# This software may be subject to U.S. export control laws and
# regulations.  By accepting this document, the user agrees to comply
# with all applicable U.S. export laws and regulations.  User has the
# responsibility to obtain export licenses, or other export authority
# as may be required before exporting such information to foreign
# countries or providing access to foreign persons.

import os, sys, pickle, fnmatch
import pylab, csv, math, copy
import numpy as np
from dataset import *
from log import printt

################### LIBS ##############
class LIBSData(Dataset):
  # Contains code needed to load, plot, and interpret 
  # LIBS (CSV) data file(s)

  def  __init__(self, inputname=None, initfilename=None,
                startsol=-1, endsol=-1, initpriorsols=False, 
                shotnoisefilt=0):
    """LIBSData(inputname="", sol=-1)

    Read in LIBS (ChemCam) data in CSV format from inputname.
    If inputname ends in .csv, treat it as a CSV file.
    If inputname ends in .pkl, treat it as a pickled file.
    Otherwise, treat it as a directory and look for a .pkl file inside;
    if not found, generate it with contents from all .csv files present.

    If present, also read in data from initfilename (must be .csv).
    This data will be used to initialize the DEMUD model.

    Optionally, specify a sol range (startsol-endsol) for data to analyze.
    Optionally, use data prior to startsol to initialize the model.
    Optionally, specify the width of a median filter to apply.
    """

    input_type = inputname[-3:]

    if input_type == 'csv':
      filename = inputname
      expname  = 'libs-' + \
          os.path.splitext(os.path.basename(filename))[0]
      #filename[filename.rfind('/')+1:filename.find('.')]
    elif input_type == 'pkl':
      if shotnoisefilt > 0:
        #filename = inputname[:-4] + ('-snf%d.pkl' % shotnoisefilt)
        filename = os.path.splitext(inputname)[0] + \
            ('-snf%d.pkl' % shotnoisefilt)
      else:
        filename = inputname
      expname  = 'libs-' + \
          os.path.splitext(os.path.basename(filename))[0]
      #filename[filename.rfind('/')+1:filename.find('.')]
    else:  # assume directory
      input_type = 'dir'
      #filename = inputname + '/libs-mean-norm.pkl'
      filename = os.path.join(inputname, 'libs-mean-norm.pkl')
      if shotnoisefilt > 0:
        #filename = filename[:-4] + ('-snf%d.pkl' % shotnoisefilt)
        filename = os.path.splitext(inputname)[0] + \
            ('-snf%d.pkl' % shotnoisefilt)
        #expname  = 'libs-' + inputname[inputname.rfind('/')+1:]
      expname  = 'libs-' + os.path.basename(inputname)

    Dataset.__init__(self, filename, expname, initfilename)

    printt('Reading %s data from %s.' % (input_type, self.filename))

    if input_type == 'dir' and not os.path.exists(filename):
      LIBSData.read_dir(inputname, filename, shotnoisefilt)
    
    self.readin(startsol, endsol, initpriorsols, shotnoisefilt)


  def  readin(self, startsol=-1, endsol=-1, initpriorsols=False, shotnoisefilt=0):
    """readin()
    
    Read in LIBS data from self.filename.
    Read in initialization data from self.initfilename.
    Normalize according to Nina's instructions.

    Optionally, specify a sol range (startsol-endsol) for data to analyze.
    Optionally, use data prior to startsol to initialize the model.
    Optionally, specify the width of a median filter to apply.
    """

    input_type = os.path.splitext(self.filename)[1][1:]

    self.data     = []
    self.initdata = []
    self.xlabel   = 'Wavelength (nm)'
    self.ylabel   = 'Intensity'

    if input_type == 'csv':
      (self.data, self.labels) = LIBSData.read_csv_data(self.filename)

      # Prune off first column (wavelengths)
      wavelengths = self.data[:,0]
      self.xvals  = wavelengths.reshape(-1,1)
      self.data   = self.data[:,1:]  # features x samples

      (self.data, self.xvals) = \
          LIBSData.prune_and_normalize(self.data, self.xvals, shotnoisefilt)
      
      (self.data, self.labels) = self.filter_data(self.data, self.labels)

    elif input_type == 'pkl':

      inf = open(self.filename, 'r')
      (self.data, self.labels, self.xvals) = pickle.load(inf)
      inf.close()

      # Temporary: until I re-run full extraction on shiva
      use = np.where(np.logical_and(self.xvals >= 270,
                                    self.xvals < 820))[0]
      self.xvals = self.xvals[use]
      self.data  = self.data[use,:]

      (self.data, self.labels) = self.filter_data(self.data, self.labels)

    else:  # Unknown format

      printt(' Error: Unknown input type for %s; no data read in' % \
             self.filename)

    # Read in the init data file, if present
    if self.initfilename != '':
      printt('Reading initialization data set from %s' % self.initfilename)
      (self.initdata, unused_labels) = LIBSData.read_csv_data(self.initfilename)

      # Prune off first column (wavelengths)
      wavelengths = self.initdata[:,0]
      self.initdata = self.initdata[:,1:] # features x samples
      (self.initdata, unused_xvals) = \
          LIBSData.prune_and_normalize(self.initdata, wavelengths, shotnoisefilt)
      print self.initdata.shape

      (self.initdata, unused_labels) = self.filter_data(self.initdata, unused_labels)
      print self.initdata.shape

    ########## Subselect by sol, if specified ##########
    if startsol > -1 and endsol >=-1:
      printt("Analyzing data from sols %d-%d only." % (startsol, endsol))
      current_sols  = [i for (i,s) in enumerate(self.labels) \
                       if (int(s.split('_')[0][3:]) >= startsol and \
                           int(s.split('_')[0][3:]) <= endsol)]
      if initpriorsols:
        previous_sols = [i for (i,s) in enumerate(self.labels) \
                         if int(s.split('_')[0][3:]) < startsol]
        printt("Putting previous sols' (before %d) data in initialization model." % startsol)
        # Concatenate initdata with data from all previous sols
        if self.initdata != []:
          print self.initdata.shape
          print self.data[:,previous_sols].shape
          self.initdata = np.hstack((self.initdata, self.data[:,previous_sols]))
        else:
          self.initdata = self.data[:,previous_sols]

      # Prune analysis data set to only include data from the sol of interest
      self.data   = self.data[:,current_sols]
      self.labels = self.labels[current_sols]


  def  filter_data(self, data, labels):
    """filter_data(data, labels)

    Filter out bad quality data, using criteria provided by Nina Lanza:
    1) Large, broad features (don't correspond to narrow peaks)
    2) Low SNR

    For each item thus filtered, write out a plot of the data
    with an explanation:
    1) Annotate in red the large, broad feature, or
    2) Annotate in text the SNR.

    Returns updated (filtered) data and label arrays.
    """

    n = data.shape[1]

    newdata = data
    remove_ind = []

    printt("Filtering out data with large, broad features.")
    #for i in [78]: # test data gap
    #for i in [1461]: # test broad feature
    #for i in [3400]: # test broad feature
    for i in range(n):
      waves     = range(data.shape[0])
      this_
data = data[waves,i]
      peak_ind  = this_data.argmax()
      peak_wave = self.xvals[waves[peak_ind]]

      # Set min peak to examine as 30% of max
      min_peak = 0.15 * this_data[peak_ind]

      # Track red_waves: indices of bands that contribute to deciding
      # to filter out this item (if indeed it is).
      # These same wavelengths will be removed from further consideration
      # regardless of filtering deci
sion
      red_waves = []
        
      # Iterate over peaks sufficiently big to be of interest
      while this