from jqdata import macro
import numpy as np
import pylab as pl
import matplotlib.pyplot as plt   #导入画图模板
import pandas as pd               #导入时间序列模块

start = '2005-1-30'     #设置开始时间
end = pd.datetime.today()
dates =[]
gdp = []
shizhi = []

def qgdpshuju(times):                 #定义调取某季度GDP的数据
    q = query(macro.MAC_AREA_GDP_YEAR.gdp).filter(macro.MAC_AREA_GDP_YEAR.stat_year==times,
        macro.MAC_AREA_GDP_YEAR.area_code=='156')
    df = macro.run_query(q)
    return  float ('%.4f' % df['gdp'].values)


def q_sum_shizhi(time):          #定义求总市值的函数day_sum_shizhi_pj（）
    df = get_fundamentals(query(valuation.market_cap),date = time)
    day_sum_shizhi = sum(df)
    return float('%.4f' % day_sum_shizhi)

                    #计算每个季度的股票市值
for d in pd.date_range(start,end,freq='q'):
    a = str(d)
    dates.append(a[0:10])
    shizhi.append(q_sum_shizhi(a[0:10]))
    
                    #计算每个季度的GDP，数据做平滑处理
for d in range(2004,2018):
    if d == 2017:
        a = str(d)
        year_gdp = qgdpshuju(a)
        gdp.append(year_gdp)
       
        q1gdp = year_gdp + year_gdp * 0.068 / 4
        gdp.append(q1gdp)
        q2gdp = year_gdp + year_gdp * 0.068 / 2
        gdp.append(q2gdp)
        q3gdp = year_gdp + year_gdp * 0.067  * 0.75
        gdp.append(q3gdp)
    else:  
        a = str(d)
        year_gdp = qgdpshuju(a)
        gdp.append(year_gdp)
        b = qgdpshuju(str(d+1)) - qgdpshuju(a)
        q1gdp = b / 4 + year_gdp
        gdp.append(q1gdp)
        q2gdp = q1gdp + b / 4
        gdp.append(q2gdp)
        q3gdp = q2gdp + b / 4
        gdp.append(q3gdp)
del gdp[0]

                            #计算巴菲特指标
bafeite = []
for i in range(0,55):
    bafeite.append(shizhi[i] / gdp[i])

                   #画图
fig, axes = plt.subplots(figsize=(12,9))
axes.plot(dates,bafeite)
axes.set_title('巴菲特指标',fontsize=20)
axes.set_xlabel('时间',fontsize=16)
axes.set_ylabel('数值',fontsize=20)
xtic = np.arange(0,len(dates),2)
axes.set_xticks(xtic)
pl.xticks(rotation=45)
axes.grid(True)

def q_sum_shizhi(time):          #定义求总市值的函数day_sum_shizhi_pj（）
    df = get_fundamentals(query(valuation.market_cap),date = time)
    day_sum_shizhi = sum(df)
    return float('%.4f' % day_sum_shizhi)
time = '2018-9-30'
q_sum_shizhi(time)