class pecData():
    ''' Class to store and manipulate a PEC Data '''
    def __init__(self, filename):
        if filename:
            self.loadData(filename)

    def loadData(self, filename):
        ''' Load Data from File and calculate the PEC Sum '''
        with open(filename) as pecFile:
            self.header = pecFile.readline().strip()
            self.rawData = [line.strip().split() for line in pecFile]
            
        self.segment = array([int(point[0]) for point in self.rawData])
        self.pec = array([int(point[1]) for point in self.rawData])
        self.pecSum = cumsum(self.pec[::-1])
        
    def totalSegments(self):
        ''' Return the total number of segments in the PEC Data '''
        return self.segment[-1]
        
    def sumAtSegment(self, segment):
        ''' Return the current sum at segment given in the argument '''
        return self.pecSum[segment]
        
    def plotData(self, raw = False, clear = False, newFig = False):
        ''' Plot PEC Sum Data.
            Modifiers: raw = True -> Plot Raw Data
                       clear = True -> Clear the current figure before plotting
                       newFig = True -> Creates new figure for plotting (overrides clear)
        '''
        if newFig:
            figure()
        elif clear:
            clf()
            
        if raw:
            plot(self.segment[::-1], self.pec, label=self.header)
        else:
            plot(self.segment[::-1], self.pecSum, label=self.header)
        showLabels(self, raw)
        
        
class pecList():
    ''' Class to store several PEC Data files and process them '''
    def __init__(self):
        self.list = {}
        self.update()
        
    def loadPEC(self, pecname, filename):
        ''' Add PEC Data file to list '''
        self.list[pecname] = pecData(filename)
        self.update()
        
    def removePEC(self, pecname):
        ''' Remove PEC Data file from list '''
        self.list.pop(pecname, None)
        self.update()
        
    def update(self):
        ''' Internal, update list '''
        self.size = len(self.list)
        self.keys = self.list.keys()
        self.items = self.list.values()
        
    def verifyPEC(self):
        ''' Internal, verifies consistence in segment number across different PEC Data files '''
        first = True
        self.segment = []
        self.update()
        for item in self.items:
            if len(item.segment) is not len(self.segment):
                if first:
                    self.segment = item.segment
                    first = False
                else:
                    return False
        return True
    
    def statistics(self, plot = True, plotRaw = False, plotAll = False):
        ''' Calculates the statistics of current PEC Data files in the list
            if plot, plot final results. plotRaw also plot raw data and plotAll plots original points as well
        '''
        if self.verifyPEC():
            pec = []
            pecSum = []
            for item in self.items:
                pec.append(item.pec)
                pecSum.append(item.pecSum)
                    
            self.pecAvg = [point.mean() for point in array(pec).transpose()]
            self.pecStd = [point.std() for point in array(pec).transpose()]
            
            self.pecSumAvg = [point.mean() for point in array(pecSum).transpose()]
            self.pecSumStd = [point.std() for point in array(pecSum).transpose()]

            if plot:
                self.plotAvg(newFig = True)
                
            if plotRaw:
                self.plotAvg(raw = True)
                
            if plotAll:
                for item in self.items:
                    item.plotData()
                    if plotRaw:
                        item.plotData(raw = True)
                
        else:
            print "Sorry, PEC Data differ in number of segments. Cannot take average!"
        
    def plotAvg(self, raw = False, clear = False, newFig = False):
        ''' Plot PEC Sum Data Average.
            Modifiers: raw = True -> Plot Raw Data Average
                       clear = True -> Clear the current figure before plotting
                       newFig = True -> Creates new figure for plotting (overrides clear)
        '''
        if newFig:
            figure()
        elif clear:
            clf()
            
        if raw:
            errorbar(self.segment[::-1], self.pecAvg, self.pecStd, label='Raw PEC Average')
        else:
            errorbar(self.segment[::-1], self.pecSumAvg, self.pecSumStd, label='PEC Sum Average')
        showLabels(self, raw)
             
                   
def showLabels(data, raw = False):
    ''' Internal, creates plot labels '''
    if raw:
        ylabel('PEC')
    else:
        xlim(data.segment[-1], data.segment[0])
        ylabel('PEC Sum')
    xlabel('Segment')
    title('PEC Train Data')
    text(180, -20, 'PEC Unit = 0.28125 arcsec')
    legend()