'''
Descripttion: 
Author: Haixu He
Date: 2022-01-03 20:11:37
'''

from prophet import Prophet
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import argrelmax   # 求局部最大值

plt.rc('font', family='Times New Roman')
plt.rcParams['font.size'] = 15


def check_data(x):
    x[np.abs(x) < 0.05] = 0
    return x


def prophet_CD(df, changepoint):
    m = Prophet(
        n_changepoints=138,
        changepoint_range=1,
        changepoint_prior_scale=0.8,
    )
    m.add_seasonality(name='yearly', period=365, fourier_order=3)
    m.fit(df)
    forecast = m.predict(df)

    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(2, 1, 1)
    # m.plot_components(forecast)

    fig = m.plot(forecast, ax, xlabel='Date', ylabel='NDVI')
    ax.scatter(df['ds'].values[changepoint], df['y'].values[changepoint], color='red',
               label='Fire change point', zorder=5)

    ax.plot(forecast['ds'], forecast['trend'], c='r')

    # 确定变化点
    k = 3
    deltas = check_data(m.params['delta'])
    temp = np.abs(np.nanmean(deltas, axis=0)).copy()
    maxlist = argrelmax(temp, order=2)[0]  # 找change rate局部最大值
    top_k_idx = np.abs(deltas[0][maxlist]).argsort()[::-1][0:k]   # 寻找前k个最大值
    for jjj in maxlist[top_k_idx]:
        temp_j = jjj
        while 1:
            if deltas[0][temp_j] != 0:
                temp_j -= 1
            else:
                break
        ax.axvline(x=df['ds'][temp_j + 1], c='r', ls='--')

    ax.legend(['EVI', "Forecast curve", "Interval width = 0.9", "Fire change point",
               "Trend", "The change points based on change rate > 0.1"], loc=2, bbox_to_anchor=(1.02, 1.0),
              borderaxespad=0., fontsize=15)

    deltas = m.params['delta'].mean(0)

    ax = fig.add_subplot(212)
    ax.bar(range(len(deltas)), deltas, facecolor='#0072B2', edgecolor='#0072B2')
    ax.grid(True, which='major', c='gray', ls='-', lw=1, alpha=0.2)
    ax.set_ylabel('Rate change')
    ax.set_xlabel('Potential changepoint')
    plt.tight_layout()
    plt.show()


def run(file):
    df = pd.read_csv(file)
    ChangePoint = np.where(df['label1'].values == 1)[0][0]  # 已知变化点
    df['y'] = df['EVI']
    df['ds'] = pd.to_datetime(df['datetime'])
    del df['datetime']  # 删除A列，会就地修改
    del df['EVI']  # 删除A列，会就地修改
    prophet_CD(df, ChangePoint)


if __name__ == '__main__':
    run('data/ee-chart1.csv')
    run('data/ee-chart2.csv')