import pandas as pd
from jqdata import *
from datetime import date, timedelta
import math

# 获取所有股票当天 / 近期的 PE, PB, ROE
q = query(
    valuation.code,
    valuation.pe_ratio,
    valuation.pb_ratio
)
df = get_fundamentals(q)
df = df[df["pe_ratio"] > 0]
df = df[df["pe_ratio"] < 10]
df = df[df["pb_ratio"] > 0]
df = df[df["pb_ratio"] < 1.5]

def DividendRatio(security_list,end_date,count=1):
    '''查询股息率(日更新)
    输入:股票池,截止日期,获取数量
    输出:panel结构,单位:1'''
    trade_days = get_trade_days(end_date=end_date,count = count)
    security_list.sort()
    secu_list = [x[:6] for x in security_list]
    code_df = jy.run_query(query(
         jy.SecuMain.InnerCode,jy.SecuMain.SecuCode,
    #     jy.SecuMain.ChiName,jy.SecuMain.CompanyCode
        ).filter(
        jy.SecuMain.SecuCode.in_(secu_list),jy.SecuMain.SecuCategory==1).order_by(jy.SecuMain.SecuCode))
    code_df['code'] = security_list
    df = jy.run_query(query(
#         jy.LC_DIndicesForValuation    #得到整表
        jy.LC_DIndicesForValuation.InnerCode,
                jy.LC_DIndicesForValuation.TradingDay,
                 jy.LC_DIndicesForValuation.DividendRatio,
                ).filter(jy.LC_DIndicesForValuation.InnerCode.in_(code_df.InnerCode),
                        jy.LC_DIndicesForValuation.TradingDay.in_(trade_days)
                        ))
    f_df = df.merge(code_df,on='InnerCode').set_index(['TradingDay','code']).drop(['InnerCode','SecuCode'],axis=1)
    panel = f_df.to_panel()
    return panel

df = df.sort_values(["code"], ascending=[True])

yesterday = date.today() - timedelta(1)
dividend_ratio_panel = DividendRatio(df.code.tolist(), yesterday.strftime("%Y-%m-%d"))
dividend_ratio = dividend_ratio_panel.major_xs(yesterday.strftime("%Y-%m-%d"))["DividendRatio"].tolist()

df["dividend_ratio"] = dividend_ratio
df = df.dropna()
df = df[df["dividend_ratio"] > 0.03]

industry_df = finance.run_query(
        query(
            finance.STK_COMPANY_INFO.code,
            finance.STK_COMPANY_INFO.industry_id
        ).filter(
            finance.STK_COMPANY_INFO.code.in_(df["code"].tolist())
        )
    )
df["industry_id"] = industry_df["industry_id"].tolist()

df_sorted = df.sort_values(["pe_ratio", "pb_ratio", "dividend_ratio"], ascending=[True, True, False])

def get_industry_id(code):
    df = finance.run_query(
        query(
            finance.STK_COMPANY_INFO.industry_id
        ).filter(
            finance.STK_COMPANY_INFO.code == code
        )
    )
    
    return df["industry_id"][0]

def get_six_stocks(sorted_stock_list, max_stock):
    stock_list_info = []
    stock_list = []
    occur_dict = {}
    for loc in range(sorted_stock_list.shape[0]):
        field = get_industry_id(sorted_stock_list.iloc[loc]["code"])
        if field in occur_dict:
            if occur_dict[field] >= max_stock * 0.3:
                continue
            else:
                occur_dict[field] += 1
                stock_list_info.append(sorted_stock_list.iloc[loc])
                stock_list.append(sorted_stock_list.iloc[loc]["code"])
        else:
            stock_list_info.append(sorted_stock_list.iloc[loc])
            stock_list.append(sorted_stock_list.iloc[loc]["code"])
            occur_dict[field] = 1
        if len(stock_list) >= max_stock:
            break
    return stock_list_info, stock_list

stock_list_info, stock_list = get_six_stocks(df_sorted, 10)

stock_list

# 将 sum_price 中的钱平均分到这些股票上面
def buy_stock(stock_list, sum_price):
    # 获取 stock_list 中股票的当日价格
    price_dict = history(1, unit='1d', field='avg', security_list=stock_list, df=False, skip_paused=False, fq='pre')
    for stock in price_dict.keys():
        price_dict[stock] = price_dict[stock][0]
    # 将 sum_price 平均分配到这些股票上
    sum_price_dict = {}
    per_price = sum_price / len(price_dict)
    for stock in price_dict.keys():
        sum_price_dict[stock] = math.floor(per_price / price_dict[stock] / 100)
    sum_value = 0
    for stock in sum_price_dict.keys():
        sum_value += sum_price_dict[stock] * price_dict[stock]*100
    res_value = sum_price - sum_value
    return price_dict, sum_price_dict, res_value

(datetime.strptime("2018-09-12", "%Y-%m-%d").date()- timedelta(1)).strftime("%Y-%m-%d")

dividend_ratio_panel.major_xs(dividend_ratio_panel.major_axis[0])

df["pe_ratio_rank"] = df["pe_ratio"].rank(ascending=0,method='max')

df.sort_values(["pe_ratio_rank"])