import talib as ta
import numpy as np
import pandas as pd
import jqdata as jq
import requests
import re
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from io import BytesIO

'''定义公共变量和函数'''

# 获取交易日数据
today = datetime.date.today()
trade_days = jq.get_trade_days(end_date=today, count=21)
today = trade_days[-1]             # 当天
yesterday = trade_days[-2]         # 上一个交易日
previous_day_20 = trade_days[-21]  # 前20个交易日
last_end_date = datetime.date(today.year-1, 12, 31)                    # 上年度最后一天
last_end_date = jq.get_trade_days(end_date=last_end_date, count=1)[0]  # 上年度最后一个交易日

# 获取涨跌幅函数
def get_stock_rise(stock_list, start_date, end_date):
    # 获取当日未停牌股票
    panel = get_price(stock_list, end_date=end_date, fields=['paused'], count=1)
    df = panel['paused'].T
    df = df[df.iloc[:,0] == 0]
    stock_list = list(df.index)

    # 获取收盘价数据
    panel = get_price(stock_list, start_date=start_date, end_date=end_date, fields=['close'])
    df = panel['close'].T

    # 计算涨跌幅
    df['rise'] = (df[end_date] - df[start_date]) / df[start_date]
    df = df.sort_values(by='rise', ascending=False)
    
    return df

# 匹配行业和概念名称
def match_text(industry_name):
    r  = requests.get(r'https://www.joinquant.com/data/dict/plateData')
    tmp = r.content.decode('utf-8')
    matches = re.findall(r"<h[\d]{1,} id=\"" + industry_name + "\">.*?" + industry_name + ".*?</h[\d]{1,}>.*?<table>(.*?)</table>", tmp, re.S)
    content = matches[0]
    tr_matches = re.findall(r"<tr>.*?<td>(.*?)</td>.*?<td>(.*?)</td>.*?<td>(.*?)</td>.*?</tr>", content, re.S)
    for code, name, start_date in tr_matches:
        yield code, name, start_date

# 获取板块和概念数据
def get_block_list(textlist):
    df = map(lambda x:pd.DataFrame([[code, name, start_date] for code, name, start_date in match_text(x)], columns=['code','name','start_date']), textlist)
    df_block = pd.concat(df,axis=0).reset_index(drop=1)
    
    # 增加涨跌幅列
    df_block['rise'] = 0
    
    return df_block

# 绘制行业相较基准涨跌幅分布图
def get_block_rise_map(df_rise):
    df_rise = df_rise.sort_index(ascending=True)
    plt.figure(figsize = (25, 25))
    x = df_rise.rise_20_ex * 100
    y = df_rise.rise_year_ex * 100
    
    # 添加当天涨跌标记
    txt = []
    for index in df_rise.index:
        name = df_rise.name[index] + (' ↑' if df_rise.rise[index] > 0 else '↓')
        txt.append(name) 
        
    # 设置标签
    for i in range(len(x)):
        plt.annotate(txt[i], xy=(x[i], y[i]), xytext=(x[i]+0.1, y[i]-0.15))

    # 设置坐标轴范围
    plt.xlim((-15, 15))
    plt.ylim((-30, 30))

    # 设置坐标轴名称
    plt.xlabel('20天相对涨幅')
    plt.ylabel('年度相对涨幅')

    # 设置坐标轴位置
    ax = plt.gca()
    ax.spines['right'].set_color('none')  
    ax.spines['top'].set_color('none') 
    ax.spines['bottom'].set_position(('data', 0))  
    ax.spines['left'].set_position(('data',0)) 

    # 生成网格线
    plt.grid(linestyle='--')

    # 生成离散点
    plt.scatter(x, y)
    plt.show()
    

'''统计盘面数据'''

# 获取涨跌幅数据
stock_list = list(get_all_securities(['stock']).index)
df = get_stock_rise(stock_list, yesterday, today)

# 涨跌幅统计
df_positive = df[df['rise'] > 0]
df_negative = df[df['rise'] < 0]
df_zero = df[df['rise'] == 0]
df_positive_5 = df[df['rise'] >= 0.05]
df_negative_5 = df[df['rise'] <= -0.05]
rise_median = df['rise'].median()

# 获取首日上市新股
# 因首日上市新股涨幅统计为nan，需在涨幅统计时加上新股的数量
df_new = np.isnan(df.rise)
df_new = df_new[df_new == True]

# 获取两市成交额
index_sh = '000001.XSHG' #上证综指
index_sz = '399106.XSHE' #深证综指
panel = get_price([index_sh,index_sz], start_date=yesterday, end_date=today, fields=['money'])
df = panel['money'].T

# 成交额统计
money_today = df[today].sum() / 1e+8
money_rise = (df[today].sum() - df[yesterday].sum()) / 1e+8

# 复盘总结
review = '（%s）\n● 复盘\n' % today
review += '今日两市上涨%s家、下跌%s家、平盘%s家，涨幅在5%%以上%s家、跌幅在5%%以上%s家，涨幅中位数%.2f%%。\n' % (
        (len(df_positive) + len(df_new)), 
        len(df_negative),
        len(df_zero), 
        (len(df_positive_5) + len(df_new)), 
        len(df_negative_5), 
        rise_median*100)

review += '两市成交额%.2f亿，较上一交易日%s%.2f亿。\n' % (
        money_today, 
        ('增加' if money_rise > 0 else '减少'), 
        abs(money_rise))

print(review)
write_file('review_day.txt', review)

'''统计择时信号'''

# 初始化参数
n = 10
index_list = ('000001.XSHG', #上证指数
              '399001.XSHE', #深证成数
              '399005.XSHE', #中小板指
              '399006.XSHE', #创业板指
              '000016.XSHG', #上证50
              '000300.XSHG', #沪深300
              '000905.XSHG', #中证500
              '000852.XSHG', #中证1000
              '399678.XSHE', #深次新股
             )
df_timing = pd.DataFrame(index=index_list, columns=('name','signal','rise','rise_20','rise_year'))
review = '\n● 指数涨跌幅\n'

# 统计指数涨跌幅
for index in index_list:
    # 获取指数名称
    df_timing.loc[index, 'name'] = get_security_info(index).display_name
    
    # 计算今日涨幅
    df = get_price(index, end_date=today, fields=['close'], count=22)
    df_timing.loc[index, 'rise'] = (df.close[-1] - df.close[-2]) / df.close[-2]
    
    # 计算20天涨幅
    close_20 = (df.close[-20] + df.close[-21] + df.close[-22]) / 3
    df_timing.loc[index, 'rise_20'] = (df.close[-1] - close_20) / close_20
    
    # 计算年度涨幅
    df_last = get_price(index, end_date=last_end_date, fields='close', count=1)
    df_timing.loc[index, 'rise_year'] = (df.close[-1] - df_last.close[-1]) / df_last.close[-1]
    
    # 输出
    review += '%s：今日涨幅 %.2f%%，20天涨幅 %.2f%%，年度涨幅 %.2f%%\n' % (
        df_timing.name[index], 
        df_timing.rise[index]*100,
        df_timing.rise_20[index]*100,
        df_timing.rise_year[index]*100)
    
# 初始化择时信号
signal_today = 'wait'   # 今日择时信号
signal_20 = 'wait'      # 中期择时信号

# 涨幅为正的指数个数
positive_today = len(df_timing[df_timing.rise > 0])   # 今日涨幅为正的指数个数
positive_20 = len(df_timing[df_timing.rise_20 > 0])   # 20日前涨幅为正的指数个数

# 判断择时信号
if positive_today > 5:
    signal_today = 'yes'
elif positive_today < 4:
    signal_today = 'no'
    
if positive_20 > 5:
    signal_20 = 'yes'
elif positive_20 < 4:
    signal_20 = 'no'
    
# 输出
review += '\n● 择时信号\n今日信号：%s（%s↑%s↓），中期信号：%s（%s↑%s↓）\n' % (
    signal_today,
    positive_today,
    9 - positive_today,
    signal_20,
    positive_20,
    9 - positive_20)

print(review)
write_file('review_day.txt', review, append=True)

'''统计行业涨跌幅'''

# 获取行业数据
textlist = ["申万一级行业"]
df_block = get_block_list(textlist)
df_block['rise_20'] = 0
df_block['rise_20_ex'] = 0
df_block['rise_year'] = 0
df_block['rise_year_ex'] = 0

# 计算基准指数涨跌幅
benchmark = ['000300.XSHG']
rise_20_bm = get_stock_rise(benchmark, previous_day_20, today).rise[0]
rise_year_bm = get_stock_rise(benchmark, last_end_date, today).rise[0]

# 统计行业涨跌幅
for index in df_block.index:
    # 修改行业名称
    df_block.loc[index, 'name'] = df_block.name[index].split('I')[0]
    
    # 获取板块成份股
    stock_list = list(get_industry_stocks(df_block.code[index], today))
    
    # 获取板块今日涨跌幅
    df = get_stock_rise(stock_list, yesterday, today)
    df_block.loc[index, 'rise'] = df.rise.mean()
    
    # 获取板块20天涨跌幅
    df_20 = get_stock_rise(stock_list, previous_day_20, today)
    df_block.loc[index, 'rise_20'] = df_20.rise.mean()
    df_block.loc[index, 'rise_20_ex'] = df_20.rise.mean() - rise_20_bm
    
    # 获取板块年度涨跌幅
    df_year = get_stock_rise(stock_list, last_end_date, today)
    df_block.loc[index, 'rise_year'] = df_year.rise.mean()
    df_block.loc[index, 'rise_year_ex'] = df_year.rise.mean() - rise_year_bm

# 提取年度涨幅排名靠前和靠后的行业
df_block = df_block.sort_values(by='rise_year', ascending=False)
df_head = df_block.head(5)
df_tail = df_block.tail(5)

# 输出行业涨跌幅
review = '\n● 强势行业\n'
for index in df_head.index:
    review += '%s：今日涨幅 %.2f%%，20天涨幅 %.2f%%，年度涨幅 %.2f%%\n' % (
        df_head.name[index], 
        df_head.rise[index]*100,
        df_head.rise_20[index]*100,
        df_head.rise_year[index]*100)
    
review += '\n● 弱势行业\n'
for index in df_tail.index:
    review += '%s：今日涨幅 %.2f%%，20天涨幅 %.2f%%，年度涨幅 %.2f%%\n' % (
        df_tail.name[index], 
        df_tail.rise[index]*100,
        df_tail.rise_20[index]*100,
        df_tail.rise_year[index]*100)
    
print(review)
write_file('review_day.txt', review, append=True)

# 绘制行业相较基准涨跌幅分布图
get_block_rise_map(df_block)

df_block