import numpy as np
import pandas as pd

from neuralprophet import set_log_level, set_random_seed, NeuralProphet, save, load
set_log_level('ERROR')

train_data = pd.read_csv('data/train_data.csv', index_col=0, parse_dates=[0])
sample_submission = pd.read_csv('data/submission.csv')

# ds: 年月
# ID: 野菜と地域
# y: 目的変数 (対数変換済み)
df = (
    train_data
    .pipe(lambda df: np.log1p(df))
    .rename_axis(index='ds')
    .reset_index()
    .melt(id_vars=['ds'], var_name='ID', value_name='y')
)

df.head()

# 野菜一覧
veg_list = train_data.columns.str.replace('_.+', '', regex=True).unique().to_list()

# 野菜ごとに分けてモデル化する。
for veg in veg_list:
    # 学習の分散が大きいので、
    # 乱数シードを変えて10回学習と予測を繰り返し、
    # その平均を最終的な予測値とする。
    for seed in range(10):
        print(f'veg: {veg}, seed: {seed}')
        set_random_seed(seed=seed)
        model = NeuralProphet(
            # 月次データのため、年間の季節性のみモデル化する。
            yearly_seasonality=True,
            weekly_seasonality=False,
            daily_seasonality=False,
            # トレンドと季節性を地域別にモデル化する。
            trend_global_local='local',
            season_global_local='local',
            # 自己回帰成分の次数。
            # 昨年度の価格を参照できる12以上の次数を設定すると予測精度が良くなる。
            n_lags=12,
            # 損失関数をMSEにするとコンペの評価指標と一致するため良さそうだが、
            # 試したところ過学習がひどくなったため、デフォルトのHuber lossを採用した。
            # loss_func='MSE'
        )
        df_veg = df.loc[lambda df: df['ID'].str.startswith(veg), :].copy()
        metrics = model.fit(df=df_veg, freq='MS', progress=None, learning_rate=0.01, epochs=20)
        # 野菜の数が多く学習に時間がかかるため、
        # 途中で中断できるよう学習が済んだ野菜のモデルを保存しておく。
        save(model, f'model_20230817_03/model_20230817_03_{veg}_{seed}.np')

forecast = []
for veg in veg_list:
    for seed in range(10):
        print(f'veg: {veg}, seed: {seed}')
        model = load(f'model_20230817_03/model_20230817_03_{veg}_{seed}.np')
        df_veg = df.loc[lambda df: df['ID'].str.startswith(veg), :].copy()
        df_future = model.make_future_dataframe(df_veg, n_historic_predictions=True, periods=1)
        forecast_ = model.predict(df_future)
        forecast_['seed'] = seed
        forecast.append(forecast_)
forecast = pd.concat(forecast, ignore_index=True)

submission = (
    forecast
    # 2019年12月の予測結果だけ抽出する。
    .loc[lambda df: (df['ds'].dt.year == 2019) & (df['ds'].dt.month == 12), ['ID', 'yhat1']]
    # Seed blending
    .groupby('ID', as_index=False)['yhat1'].mean()
    # 提出用にフォーマットを整える。
    .rename(columns={'ID': 'id', 'yhat1': 'y'})
    .merge(sample_submission[['id']], how='right', on='id')
    # 対数変換を元に戻す。
    .assign(y = lambda df: np.expm1(df['y']))
)

submission.to_csv('submission/submission_20230817_03.csv', index=False, header=True)

from matplotlib import pyplot as plt
import seaborn as sns
import japanize_matplotlib

veg = 'トマト'
seed = 0

model = load(f'model_20230817_03/model_20230817_03_{veg}_{seed}.np')
model.set_plotting_backend('matplotlib')

df_veg = df.loc[lambda df: df['ID'].str.startswith(veg), :].copy()
df_future = model.make_future_dataframe(df_veg, n_historic_predictions=True, periods=1)
forecast = model.predict(df_future)

# 実測値 (黒点) と予測値 (青線)
model.plot(forecast, df_name=f'{veg}_関東', figsize=(12, 4))

model.plot_parameters(df_name=f'{veg}_関東', figsize=(12, 8))

# 実測値 (黒点) と予測値 (青線)
model.plot(forecast, df_name=f'{veg}_中国', figsize=(12, 4))

model.plot_parameters(df_name=f'{veg}_中国', figsize=(12, 8))