import tushare as ts
import matplotlib.pyplot as plt
import matplotlib.finance as mpf
from matplotlib.pylab import date2num
import datetime
import pandas as pd

class tradelist_analysis(object):
    
    def __init__(self, path=None,symbol=None):
        self.path = path
        self.symbol = symbol
    
    # 交易次数统计
    def symbol_counts(self):
        syms = self.extract_symbols()
        return syms.value_counts()
    
    def plot_symbol_counts(self):
        syms = self.extract_symbols()
        syms.value_counts().plot()

    # 读取交易列表
    def get_all_df(self):
        df = pd.read_csv(self.path) # 导出的时候设置转为utf8无bom格式
        df = df[~(df['成交价']=='--')]
        return df
    
    # 提取股票代码
    def extract_symbols(self):
        df = self.get_all_df()
        pattern = r'[0-9]{6}.[A-Z]{4}'
        syms = df['标的'].str.findall(pattern)
        syms = syms.apply(lambda x : x[0])
        return syms

    # 股票代码添加后缀
    def trans_symbol(self):
        if len(self.symbol)==11:
            return self.symbol
      
        if int(self.symbol[0])==6:
            self.symbol = self.symbol+'.XSHG'
        else:
            self.symbol = self.symbol+'.XSHE'
        return self.symbol
            
    # 获取个股买入卖出信息
    def get_symbol_df(self):
        self.symbol = self.trans_symbol()
        df = self.get_all_df()
        df['日期']=pd.to_datetime(df['日期'])
        bs = df['标的'].str.contains(self.symbol)
        dfs = df[bs]
        return dfs

    # 获取行情数据
    def get_hist_data(self):
        self.symbol = self.trans_symbol()
        dfs = self.get_symbol_df()
        start = dfs['日期'].iloc[0]
        end = dfs['日期'].iloc[-1]
        
        hist_data = get_price(self.symbol,start_date = start-datetime.timedelta(5), end_date=end+datetime.timedelta(5), frequency='daily', fields=['open', 'high','low','close'])
        hist_data.reset_index(inplace=True)
        hist_data['index']=hist_data['index'].apply(lambda s:date2num(s))
        hist_data=hist_data.as_matrix()
        return hist_data
  
    # 创建子图
    def plot_candle(self):
        hist_data=self.get_hist_data()
        dfs = self.get_symbol_df()
        
        fig, ax = plt.subplots()
        fig.subplots_adjust(bottom=0.1)

        # 设置X轴刻度为日期时间
        ax.xaxis_date()
        ax.autoscale_view()
        plt.xticks(rotation=45)
        plt.yticks()
        plt.title("股票代码：xx")
        plt.xlabel("时间")
        plt.ylabel("股价（元）")
        mpf.candlestick_ohlc(ax,hist_data,width=1,colorup='r',colordown='green')
        plt.grid()

            
        for i in range(len(dfs)):
            dt = dfs['日期'].iloc[i]
            price = dfs['成交价'].iloc[i]
            #print dt,price
            if dfs['交易类型'].iloc[i]=='买':
                ax.annotate('buy%s'%price, xy=(dt, float(price)), xytext=(dt, float(price)-0.2),arrowprops=dict(facecolor='black', shrink=0.1),)
            else:
                ax.annotate('sell%s'%price, xy=(dt, float(price)), xytext=(dt, float(price)+0.2),arrowprops=dict(facecolor='yellow', shrink=0.1),)

ta = tradelist_analysis('transaction_1.csv','000498')
ta.plot_candle()

ta.get_symbol_df()