import numpy as np
import jqdata
import datetime
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import math

# 平滑价格
# @inP [,,,,] 价格数组
class SmoothPrice:

    def __init__(self,inP:np.ndarray):
        self.__inP__ = inP
        self.__l_smooth__ = inP
        self.__q_smooth__ = inP
        self.__c_smooth__ = inP

    def __linearSmooth5__(self):
        _outP = []
        inP = self.__inP__
        n = len (inP)
        if n < 5:
            for i in inP:
                _outP.append (i)
        else:
            _outP.append ((3.0 * inP[0] + 2.0 * inP[1] + inP[2] - inP[4]) / 5.0)
            _outP.append ((4.0 * inP[0] + 3.0 * inP[1] + 2 * inP[2] + inP[3]) / 10.0)
            _i = 0
            for i in inP:
                if not _i == 0 and not _i == 1 and not _i == n - 2 and not _i == n - 1:
                    _outP.append ((inP[_i - 2] + inP[_i - 1] + inP[_i] + inP[_i + 1] + inP[_i + 2]) / 5.0)
                _i = _i + 1
            _outP.append ((4.0 * inP[n - 1] + 3.0 * inP[n - 2] + 2 * inP[n - 3] + inP[n - 4]) / 10.0)
            _outP.append ((3.0 * inP[n - 1] + 2.0 * inP[n - 2] + inP[n - 3] - inP[n - 5]) / 5.0)
        self.__l_smooth__ = _outP

    def __quadraticSmooth5__(self):
        _outP = []
        inP = self.__l_smooth__
        n = len (inP)
        if n < 5:
            for i in inP:
                _outP.append (i)
        else:
            _outP.append ((31.0 * inP[0] + 9.0 * inP[1] - 3.0 * inP[2] - 5.0 * inP[3] + 3.0 * inP[4]) / 35.0)
            _outP.append ((9.0 * inP[0] + 13.0 * inP[1] + 12 * inP[2] + 6.0 * inP[3] - 5.0 * inP[4]) / 35.0)
            _i = 0
            for i in inP:
                if not _i == 0 and not _i == 1 and not _i == n - 2 and not _i == n - 1:
                    _outP.append ((- 3.0 * (inP[_i - 2] + inP[_i + 2]) + 12.0 * (inP[_i - 1] + inP[_i + 1]) + 17 * inP[
                        _i]) / 35.0)
                _i = _i + 1
            _outP.append (
                (9.0 * inP[n - 1] + 13.0 * inP[n - 2] + 12.0 * inP[n - 3] + 6.0 * inP[n - 4] - 5.0 * inP[n - 5]) / 35.0)
            _outP.append (
                (31.0 * inP[n - 1] + 9.0 * inP[n - 2] - 3.0 * inP[n - 3] - 5.0 * inP[n - 4] + 3.0 * inP[n - 5]) / 35.0)
        self.__q_smooth__ = _outP

    def __cubicSmooth5__(self):
        _outP = []
        inP = self.__q_smooth__
        n = len (inP)
        if n < 5:
            for i in inP:
                _outP.append (i)
        else:
            _outP.append ((69.0 * inP[0] + 4.0 * inP[1] - 6.0 * inP[2] + 4.0 * inP[3] - inP[4]) / 70.0)
            _outP.append ((2.0 * inP[0] + 27.0 * inP[1] + 12.0 * inP[2] - 8.0 * inP[3] + 2.0 * inP[4]) / 35.0)
            _i = 0
            for i in inP:
                if not _i == 0 and not _i == 1 and not _i == n - 2 and not _i == n - 1:
                    _outP.append ((-3.0 * (inP[_i - 2] + inP[_i + 2]) + 12.0 * (inP[_i - 1] + inP[_i + 1]) + 17.0 * inP[
                        _i]) / 35.0)
                _i = _i + 1
            _outP.append (
                (2.0 * inP[n - 5] - 8.0 * inP[n - 4] + 12.0 * inP[n - 3] + 27.0 * inP[n - 2] + 2.0 * inP[n - 1]) / 35.0)
            _outP.append (
                (- inP[n - 5] + 4.0 * inP[n - 4] - 6.0 * inP[n - 3] + 4.0 * inP[n - 2] + 69.0 * inP[n - 1]) / 70.0)
        self.__c_smooth__ = _outP

    def smooth(self):
        self.__linearSmooth5__()
        self.__quadraticSmooth5__()
        self.__cubicSmooth5__()
        return self.__c_smooth__

def getstockPE(stock, days = 120,end_date=datetime.date.today()):
        _list = []
        jqv = valuation
        df = get_fundamentals_continuously(query(
            jqv.pe_ratio
        ).filter(
            jqv.code == stock
        ), end_date=end_date,count = days) 
    
        if not df.empty:
        #if  df._initialized:
            _list = df['pe_ratio'].values
            _list = [i[0] for i in _list] 
            _list_smooth = SmoothPrice(_list)
        return [_list,_list_smooth.smooth()]
    

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
fig = plt.figure (tight_layout=True,figsize=(12, 6))
gs = gridspec.GridSpec (1, 2)
        
_d11 = getstockPE("000651.XSHE")
ax = fig.add_subplot(gs[0, :])

ax.plot(range(len(_d11[0])), _d11[0], 'y')
ax.plot(range(len(_d11[1])), _d11[1], 'b')
_mean = sum(_d11[0]) / len(_d11[0])
ax.set_xlabel('PE(' + str(_mean) + ")")
ax.axhline(_mean, color='r', linestyle='--', label='mean-line')
ax.grid()