from pandas_composition import UserFrame, UserSeries
import pandas as pd
import pandas.io.data as pdd
import IPython
IPython.core.pylabtools.figsize(20, 15)

from matplotlib.finance import candlestick
from itertools import izip
def ohlc(df, width=0.3, secondary_y=False,
               *args, **kwargs):
    """
        Takes a df and plots a candlestick. 
    """
    dates = df.index.to_pydatetime() 
    times = date2num(dates)
    quotes = izip(times, df['Open'], df['Close'], df['High'], df['Low'])
    ax = gca()
    candlestick(ax, quotes, width=width, colorup='g')
    return ax

class StockData(UserFrame):
    def __init__(self, *args, **kwargs):
        symbol = kwargs.get('symbol', None)
        self.symbol = symbol
        
    def ohlc(self):
        " plot candlestick "
        ax = ohlc(self, label=self.symbol)
        ax.set_title(self.symbol)
        return ax
        
    def get_gaps(df, offset=0):
       	" when open is above or below the previous days range " 
    	gap_up = df.Open > (df.High.shift(1) + offset)
    	gap_down = df.Open < (df.Low.shift(1) - offset)
    	gaps = gap_up | gap_down
    	return Indicator(gaps, source=df, name="Gaps") 
    
    def ma(self, period):
        " moving average "
        ma = pd.rolling_mean(df.Close, period)
        ma = Indicator(ma, source=self, name='MA {period}'.format(period=period))
        return ma
        
class Indicator(UserSeries):
    def __init__(self, *args, **kwargs):
        source = kwargs.pop('source', None)
        self.source = source

    def plot(self, source_col='Close', *args, **kwargs):
        " plot the indicator. source_col only used for markers "
        df = self.source
        label = kwargs.get('label', None) or self.name # use Series.name
        kwargs['label'] = label
        
        if self.dtype == bool: # plot markers for bool
        	sig = self * df[source_col]
        	sig[sig == 0] = np.nan
    		plot_date(sig.index.to_pydatetime(), sig.values, *args, **kwargs)
      	else:
            plot(self.index.to_pydatetime(), self, *args, **kwargs)
      	legend()      

df = pdd.get_data_yahoo('AAPL')
df = StockData(df, symbol='AAPL')

df.tail(30).ohlc()
df.get_gaps().tail(30).plot('Open', color='pink', markersize=10)
df.ma(30).tail(30).plot()