#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jul 11 10:38:58 2019

@author: kantundpeterpan
"""

import peakutils
import matplotlib.pyplot as plt


class MassSpectrum():
    
    def __init__(self, array, thres=0.01, thres_abs=True, Isospec_input=False):
        self.raw = array
        self.mz = array[:,0]
        self.ints= array[:,1]
        
        if not Isospec_input:
            self.peaks_ind = peakutils.indexes(self.ints, thres=thres, thres_abs=thres_abs)
            self.peaks_ind_x = self.mz[self.peaks_ind]
            self.peaks_ind_y = self.ints[self.peaks_ind]
            self.peaks_interpolated_x = peakutils.interpolate(self.mz, self.ints, ind=self.peaks_ind)
        else:
            self.peaks_ind_x = self.mz
            self.peaks_ind_y = self.ints
        
    def plot_spectrum(self, ax = None, **kwargs):
        if ax:
            ax.plot(self.mz, self.ints, **kwargs)
        else:
            plt.plot(self.mz, self.ints, **kwargs)
            
    def plot_centroids(self, ax = None, zero_line=True, **kwargs):
        if not ax:
            ax = plt.gca()
        for m,i in zip(self.peaks_ind_x, self.peaks_ind_y):
            ax.vlines(m,0,i,**kwargs)
        if zero_line:
            ax.hlines(0,min(self.mz),max(self.mz), **kwargs)