#导入各种函数库包
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.finance as fin
#获取上证指数的行情数据
df = get_price('000001.XSHG',end_date='2017-5-25',frequency='1d',count=400)

#进行K线绘制先，整体看下走势情况
t=list(range(len(df)))
df['t']=pd.Series(t,index=df.index)
o=df['open'].values
c=df['close'].values
h=df['high'].values
l=df['low'].values
fig=plt.figure(figsize(20,10))
ax=fig.add_subplot(111)
ax.set_xlim(0,len(df))
ques=zip(t,o,c,h,l)
fin.candlestick_ochl(ax,ques,colorup='r',colordown='g',width=0.5)
plt.grid()

#进行K线的包含关系处理，再原数据表中增加label进行是否有效的标记，增加两列记录合并K线后的极值
label=[1]
h_jizhi,l_jizhi=h[0],l[0]
high_new=[h[0]]
low_new=[l[0]]
#标记Ｋ线的有效性，上升Ｋ线为１，下降为－１，无效为０
for i in t[1:]:
    #缠论上升K线的定义
    if h[i]>h_jizhi and l[i]>l_jizhi:
        label.append(1)
        h_jizhi=max(h[i],h[i-1])
        l_jizhi=max(l[i],l[i-1])
    #缠论下降K线的定义
    elif h[i]<h_jizhi and l[i]<l_jizhi:
        label.append(-1)
        h_jizhi=min(h[i],h[i-1])
        l_jizhi=min(l[i],l[i-1])
    #包涵K线的处理
    else:
        label.append(0)
        #下面的内容主要是进行极值的更新
        for j in label[::-1]:
            if j != 0:
                a=j
                break
        if a==1:
            h_jizhi=max(h[i],h_jizhi)
            l_jizhi=max(l[i],l_jizhi)
        else:
            h_jizhi=min(h[i],h_jizhi)
            l_jizhi=min(l[i],l_jizhi)
    high_new.append(h_jizhi)
    low_new.append(l_jizhi)
#把结果添加到df
df['label_k']=pd.Series(label,index=df.index)
df['high_k']=pd.Series(high_new,index=df.index)
df['low_k']=pd.Series(low_new,index=df.index)

#找出可以构成笔的K线位置
#step1 找出有效笔的起点位置
#得到有效笔的起点位置
##先将有效K线的label整理成列表l1
l1=[]
list1=df['label_k'].values
for i in range(len(list1)):
    if list1[i]!=0:
        l1.append(i)
##记录有效位置的起点
label_s_e=[]
for z in range(2,len(l1)):
    if (list1[l1[z]],list1[l1[z-1]],list1[[l1[z-2]]])==(1,1,1):
        label_s_e.append((l1[z-2],1))
    elif (list1[l1[z]],list1[l1[z-1]],list1[[l1[z-2]]])==(-1,-1,-1):
        label_s_e.append((l1[z-2],-1))
##取到第一个有效起点,并有a值进行方向的记录
label_se=[]
a=0
for i in range(1,len(label_s_e)):
    if a==0:
        label_se.append(label_s_e[i][0])
        if label_s_e[i][1]==1:
            a=1
        elif label_s_e[i][1]==-1:
            a=-1
    elif a==1:
        if label_s_e[i][1]==1:
            pass
        elif label_s_e[i][1]==-1:
            label_se.append(label_s_e[i][0])
            a=-1
    elif a==-1:
        if label_s_e[i][1]==1:
            label_se.append(label_s_e[i][0])
            a=1
        elif label_s_e[i][1]==-1:
            pass
##得到小趋势的起点位置后，找出每个区间的最高点最低点
#获得区间内最大最小值，次处找出的值是指在已经确定方向被终结后进行的搜索，主要用作画图了
b=1
new_label=[]
for i in range(1,len(label_se)):
    list_1=list(df['high'].values[label_se[i-1]:label_se[i]])
    if b==1:
        #返回列表中值最大的位置
        new_label.append(list_1.index(max(list_1))+label_se[i-1])
        b=-1
    elif b==-1:
        new_label.append(list_1.index(min(list_1))+label_se[i-1])
        b=1
#获取最值的价格，按收盘价算
new_label_p=[df['close'].values[i] for i in new_label]

#再次进行K线绘制先，整体看下走势情况，左为裸K，右为方向
t=list(range(len(df)))
df['t']=pd.Series(t,index=df.index)
o=df['open'].values
c=df['close'].values
h=df['high'].values
l=df['low'].values
fig=plt.figure(figsize=(20,10))
ax1=fig.add_subplot(221)
ax2=fig.add_subplot(222)
ax1.set_xlim(0,len(df))
ques=zip(t,o,c,h,l)
ax1.grid()
fin.candlestick_ochl(ax1,ques,colorup='r',colordown='g',width=0.5)
ax2.grid()
ax2.plot(new_label,new_label_p,color='r')