import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from jqdata import *
plt.style.use('fivethirtyeight')
pd.set_option('display.max_rows',10)

stock_all = get_all_securities(types=['stock'], date=None)
stock_list = stock_all.index.tolist()
len(stock_list)

input_date = '2019-05-10'

# 工具函数

def filter_st(stocks,date):
   
    st_data = get_extras('is_st',stocks, count=1,end_date=date)
    stocks = [stock for stock in stocks if not st_data[stock][0]]
    
    return stocks

def find_close_limit_stock(stocks,date):

    h_limit_stocks = []

    df = get_price(stocks, end_date=date,frequency='daily', fq='pre',\
                   fields=['close','volume','high_limit','low_limit'],count = 1).iloc[:,0,:]
    
    df.query('volume != 0',inplace=True)

    df['is_high_limit']=(df['close']==df['high_limit'])

    h_limit_stocks = [df.index[i] for i in range(len(df)) if (df.is_high_limit[i])]
     
    return  h_limit_stocks

#MAD
def filter_extreme_MAD(series,n): 
    median = series.quantile(0.50)
    new_median = ((series - median).abs()).quantile(0.50)
    max_range = median + n*new_median
    min_range = median - n*new_median
    return np.clip(series,min_range,max_range)

# change list in list to list 
func_flat_list = lambda l: [item for sublist in l for item in sublist]

stock_non_st_list= filter_st(stock_list,input_date)
len(stock_non_st_list)

trade_days = get_trade_days(end_date=input_date, count=250)

len(trade_days)

# 运行需要点时间
high_limit_list =[]

for date in trade_days:
    h = find_close_limit_stock(stock_non_st_list,date)
    high_limit_list.append(h)

limit_stocks_year_list = set(func_flat_list(high_limit_list))

len(limit_stocks_year_list)

stock_df = get_fundamentals(
    query(
        valuation.code, valuation.circulating_market_cap
    ).filter(
        valuation.code.in_(limit_stocks_year_list)
    ), date = input_date)

stock_df.dropna(how='any',inplace=True)
stock_df.shape

stock_df.circulating_market_cap.plot(figsize=(16,4),legend=True)
plt.axhline(stock_df.circulating_market_cap.mean(), color='red',linewidth=2,linestyle='--')
plt.show()

stock_df['circulating_market_cap'].describe()

market_df = stock_df.copy()
market_df['fs_circulating_market_cap'] = filter_extreme_MAD(stock_df['circulating_market_cap'],5)
market_df.drop_duplicates(subset='fs_circulating_market_cap',inplace=True)
market_df = market_df[market_df['circulating_market_cap'] == market_df['fs_circulating_market_cap']]
market_df.drop(['circulating_market_cap'],axis=1,inplace=True)
market_df.shape

sns.distplot(market_df.fs_circulating_market_cap,bins=20)

market_df.fs_circulating_market_cap.describe()