import sys
import numpy as np
from realtime import *

class Spectra(object):
    """Return a spectra object that can plot the total absorption spectrum or
    circular dichroism spectra for a given system. Can accept x,y,and z RT-TDDFT
    log files"""
    def __init__(self,x=None,y=None,z=None,s='abs',d=150,zp=None,auto=False):
        self.spectra = None
        self.frequency = None

        self.x = None
        self.y = None
        self.z = None

        self.spectra_type = s
        self.damp_const   = d
        self.zero_pad     = zp

        # Load all the RealTime objects
        self.directions = []
        if x:
            self.x = RealTime(x)
            self.directions.append('x')
        if y:
            self.y = RealTime(y)
            self.directions.append('y')
        if z:
            self.z = RealTime(z)
            self.directions.append('z')

        # Enforce consistent data lengths
        self.align_data()
 
        # Do the isotropic fourier transform 
        for q in self.directions:
            self.__dict__[q].fourier_tx(q,self.spectra_type,self.damp_const,
                self.zero_pad,auto=auto)

        self.spectra = np.zeros_like(self.__dict__[self.directions[0]].fourier)
        for q in self.directions:
            self.frequency = self.__dict__[q].frequency
            self.spectra += self.__dict__[q].fourier


    def plot(self,xlim=[0,15],ylim=None):
        toEV = 27.2114 
        import matplotlib.pyplot as plt
        ax = plt.subplot(111)
       
        ax.plot(self.frequency*toEV,self.spectra)
        ax.set_xlim(xlim)
        if not ylim:
            if self.spectra_type == 'abs':
                ax.set_ylim([0,4])
            elif self.spectra_type == 'ecd':
                ax.set_ylim([-400,400])
        plt.show()

    def align_data(self):
        lengths = []
        if self.x:
            lengths.append(self.x.min_length) 
        if self.y:
            lengths.append(self.y.min_length) 
        if self.z:
            lengths.append(self.z.min_length) 
        min_length = min(lengths)
        if self.x:
            self.x.truncate(self.x,min_length)
        if self.y:
            self.y.truncate(self.y,min_length)
        if self.z:
            self.z.truncate(self.z,min_length)

    def peaks(self,number=3):
        """ Return the peaks from the Fourier transform"""
        from scipy.signal import argrelextrema as pks
        idx = pks(self.spectra,np.greater,order=5)
        print "First "+str(number)+" peaks (eV) found: "
        for i in xrange(number): 
            print "{0:.2f}".format(self.frequency[idx][i]*27.2114)
        
            

         


if __name__ == '__main__':
    #spectra = Spectra(x='AuH-x2c_xx',
    #                  y='AuH-x2c_yy',
    #                  z='AuH-x2c_zz',
    #                  s='abs',auto=True)
    spectra = Spectra(x='hg',auto=True)
    spectra.peaks(10)
    spectra.plot()


@@ -2296,16 +2296,28 @@
 number=3
+,thresh=0.01
 ):%0A     
@@ -2424,16 +2424,91 @@
  as pks%0A
+        # find all peak indices %5Bidx%5D, and remove those below thresh %5Bjdx%5D%0A
         
@@ -2547,17 +2547,18 @@
 r,order=
-5
+10
 )%0A      
@@ -2563,56 +2563,242 @@
     
-print %22First %22+str(number)+%22 peaks (eV) found: %22
+jdx = np.where((self.spectra%5Bidx%5D %3E= thresh))%0A        kdx = idx%5B0%5D%5Bjdx%5B0%5D%5D # indices of peaks matching criteria%0A        print %22First %22+str(number)+%22 peaks (eV) found: %22%0A        if number %3E len(idx%5B0%5D):%0A            number = len(idx%5B0%5D)
 %0A   
@@ -2878,17 +2878,17 @@
 equency%5B
-i
+k
 dx%5D%5Bi%5D*2
@@ -3145,10 +3145,13 @@
 (x='
-hg
+mg-td
 ',au
@@ -3151,32 +3151,54 @@
 -td',auto=True)%0A
+    spectra.x.test()%0A#
     spectra.peak
@@ -3199,18 +3199,16 @@
 a.peaks(
-10
 )%0A    sp
@@ -3218,11 +3218,22 @@
 ra.plot(
+xlim=%5B0,30%5D
 )%0A%0A

modify peak finder to pick out major peaks