content: """
  worker classes
  
  @author: Jev Kuznetsov
  Licence: GPL v2
  """
  
  __docformat__ = 'restructuredtext'
  
  import os
  import logger as logger
  import yahooFinance as yahoo
  from functions import returns, rank
  from datetime import date
  from pandas import DataFrame, Series
  import numpy as np
  import pandas as pd
  import matplotlib.pyplot as plt
  
  class Symbol(object):
      ''' 
      Symbol class, the foundation of Trading With Python library,
      This class acts as an interface to Yahoo data, Interactive Brokers etc 
      '''
      def __init__(self,name):
          self.name = name
          self.log = logger.getLogger(self.name)
          self.log.debug('class created.')
          
          self.dataDir = os.getenv("USERPROFILE")+'\\twpData\\symbols\\'+self.name
          self.log.debug('Data dir:'+self.dataDir)    
          self.ohlc = None # historic OHLC data
       
      def downloadHistData(self, startDate=(2010,1,1),endDate=date.today().timetuple()[:3],\
                      source = 'yahoo'):
          ''' 
          get historical OHLC data from a data source (yahoo is default)
          startDate and endDate are tuples in form (d,m,y)
          '''
          self.log.debug('Getting OHLC data')
          self.ohlc = yahoo.getHistoricData(self.name,startDate,endDate)
      
         
      def histData(self,column='adj_close'):
          '''
          Return a column of historic data.
  
          Returns
          -------------        
          df : DataFrame
          '''
          s = self.ohlc[column]
          return DataFrame(s.values,s.index,[self.name])
      
      @property
      def dayReturns(self):
          ''' close-close returns '''
          return (self.ohlc['adj_close']/self.ohlc['adj_close'].shift(1)-1)
          #return DataFrame(s.values,s.index,[self.name])
  
  class Portfolio(object):
      def __init__(self,histPrice,name=''):
          """
          Constructor
          
          Parameters
          ----------
          histPrice : historic price
          
          """
          self.histPrice = histPrice
          self.params = DataFrame(index=self.symbols)
          self.params['capital'] = 100*np.ones(self.histPrice.shape[1],dtype=np.float)
          self.params['last'] = self.histPrice.tail(1).T.ix[:,0]
          self.params['shares'] = self.params['capital']/self.params['last']
          self.name= name
          
          
      def setHistPrice(self,histPrice):
          self.histPrice = histPrice
      
      def setShares(self,shares):
          """ set number of shares, adjust capital
          shares: list, np array or Series
          """
          
          if len(shares) != self.histPrice.shape[1]:
              raise AttributeError('Wrong size of shares vector.')
          self.params['shares'] = shares
          self.params['capital'] = self.params['shares']*self.params['last']
      
      def setCapital(self,capital):
          """ Set target captial, adjust number of shares """
          if len(capital) != self.histPrice.shape[1]:
              raise AttributeError('Wrong size of shares vector.')
          self.params['capital'] = capital
          self.params['shares'] = self.params['capital']/self.params['last']
      
      
      def calculateStatistics(self,other=None):
          ''' calculate spread statistics, save internally '''
          res = {}
          res['micro'] = rank(self.returns[-1],self.returns)
          res['macro'] = rank(self.value[-1], self.value)
  
          res['last'] = self.value[-1]
          
          if other is not None:
              res['corr'] = self.returns.corr(returns(other))
          
          return   Series(res,name=self.name)    
          
      @property
      def symbols(self):
          return self.histPrice.columns.tolist() 
      
      @property   
      def returns(self):
          return (returns(self.histPrice)*self.params['capital']).sum(axis=1)
       
      @property
      def value(self):
          return (self.histPrice*self.params['shares']).sum(axis=1)
      
      def __repr__(self):
          return ("Portfolio %s \n" % self.name ) + str(self.params)
          #return ('Spread %s :' % self.name ) + str.join(',',
          #        ['%s*%.2f' % t  for t in zip(self.symbols,self.capital)])
          
          
      
  class Spread(object):
      ''' 
      Spread class, used to build a spread out of two symbols.    
      '''
      
      def __init__(self,stock,hedge,beta=None):
          ''' init with symbols or price series '''
          
          if isinstance(stock,str) and isinstance(hedge,str):
              self.symbols = [stock,hedge]
              self._getYahooData()
          elif isinstance(stock,pd.Series) and isinstance(hedge,pd.Series):
              self.symbols = [stock.name,hedge.name]
              self.price = pd.DataFrame(dict(zip(self.symbols,[stock,hedge]))).dropna()
          else:
              raise ValueError('Both stock and hedge should be of the same type, symbol string or Series')
      
          # calculate returns
          self.returns = self.price.pct_change()
          
          if beta is not None:
              self.beta = beta
          else:
              self.estimateBeta()
              
              
          # set data
          self.data = pd.DataFrame(index = self.symbols)
          self.data['beta'] = pd.Series({self.symbols[0]:1., self.symbols[1]:-self.beta})
      
      def calculateShares(self,bet):
          ''' set number of shares based on last quote '''
          if 'price' not in self.data.columns:
              print 'Getting quote...'
              self.getQuote()
          self.data['shares'] = bet*self.data['beta']/self.data['price']    
          
      
      def estimateBeta(self,plotOn=False):
          """ linear estimation of beta """
          x = self.returns[self.symbols[1]] # hedge
          y = self.returns[self.symbols[0]] # stock
          
          #avoid extremes
          low = np.percentile(x,20)
          high = np.percentile(x,80)
          iValid = (x>low) & (x<high)
          
          x = x[iValid]
          y = y[iValid]
          
          if plotOn:
              plt.plot(x,y,'o')
              plt.grid(True)
          
          iteration = 1
          nrOutliers = 1
          while iteration < 3 and nrOutliers > 0 :
              
              (a,b) = np.polyfit(x,y,1)
              yf = np.polyval([a,b],x)
              
              err = yf-y
              idxOutlier = abs(err) > 3*np.std(err)
              nrOutliers =sum(idxOutlier)
              beta = a
              #print 'Iteration: %i beta: %.2f outliers: %i' % (iteration,beta, nrOutliers)
              x = x[~idxOutlier]
              y = y[~idxOutlier]
              iteration += 1
  
      
          if plotOn:
              yf = x*beta
              plt.plot(x,yf,'-',color='red')
              plt.xlabel(self.symbols[1])
              plt.ylabel(self.symbols[0])
          
          self.beta = beta
          return beta
      
      @property
      def spread(self):
          ''' return daily returns of the pair '''
          return (self.returns*self.data['beta']).sum(1)
  
      
      def getQuote(self):
          ''' get current quote from yahoo '''
          q = yahoo.getQuote(self.symbols)
          self.data['price'] = q['last']
          
      def _getYahooData(self, startDate=(2007,1,1)):
          """ fetch historic data """
          data = {}        
          for symbol in self.symbols:
              print 'Downloading %s' % symbol
              data[symbol]=(yahoo.getHistoricData(symbol,sDate=startDate)['adj_close'] )
             
          self.price = pd.DataFrame(data).dropna()
     
          
      def __repr__(self):
          return 'Spread 1*%s & %.2f*%s ' % (self.symbols[0],-self.beta,self.symbols[1])
          
      @property
      def name(self):
          return str.join('_',self.symbols)
  
  
          
  if __name__=='__main__':
      
      
      s = Spread(['SPY','IWM'])
      
          