import numpy as np
import pandas as pd
from nixtla import NixtlaClient
import matplotlib.pyplot as plt
from timegpt_sparseattack import GWN
from timegpt_sparseattack import TSA
from sklearn.preprocessing import StandardScaler
import time

# TimeGPT key identification
nixtla_client = NixtlaClient(
    api_key = '' # add your timegpt key
)

nixtla_client.validate_api_key()

# Data read

ds_root = 'dataset/'
#ds_name = 'ETTh1'
#ds_name = 'ETTh2'
ds_name = 'weather'
#ds_name = 'IstanbulTraffic'
#ds_name = 'exchange2'
ds_location = ds_root + ds_name + '.csv'


if ds_name == 'exchange':
    df = pd.read_csv(ds_location, header=None)
    #data = pd.Series(df.iloc[:,0].values, index=df.index)
    #std = df.iloc[:,0].std()
    #mean = df.iloc[:,0].mean()
    df['date'] = pd.date_range(start='1990-01-01', periods=len(df), freq='D')
    # Set the date column as the index
    df.set_index('date', inplace=True)
else:
    df = pd.read_csv(ds_location)
    #data = pd.Series(df['OT'].values, index=df['date'])
    #std = df['OT'].std()
    #mean = df['OT'].mean()


if ds_name == 'IstanbulTraffic':
    l_train = 48
    l_validation = 0
    l_test = 144
    freq = 'H'
elif ds_name == 'exchange':
    l_train = 4*960
    l_validation = 2*960
    l_test = 2*960
    freq = 'D'
elif ds_name == 'exchange2':
    l_train = 4*960
    l_validation = 2*960
    l_test = 2*960
    freq = 'D'
elif ds_name =='weather':
    l_train = 12 * 30 * 24
    l_validation = 4 * 30 * 24
    l_test = 5 * 30 * 24
    freq = '10min'
else:
    l_train = 12 * 30 * 24
    l_validation = 4 * 30 * 24
    l_test = 5 * 30 * 24
    freq = 'H'

train = df.iloc[0:l_train,:]
test = df.iloc[l_train + l_validation : l_train + l_validation + l_test,:]

# Statistics
z_score_flag = True
mean = train['OT'].mean()
std = train['OT'].std()
print(std)
print(mean)

# Setting
scale = mean * 0.02
epsilon = 0.1
tau = 9
print('scale:', scale)
print('eplison:', epsilon)

#save_name_train = ds_name +'_' + 'train.csv'
#train.to_csv(save_name_train, header=True)
#save_name_test = ds_name +'_' +'test.csv'
#test.to_csv(save_name_test, header=True)


time_col = 'date'
target_col = 'OT'
h = 48
historical_n = 96
model = 'TimeGPT_'


# prediction by timegpt
for i in range(10):
    input = test.iloc[i * historical_n : (i+1)*historical_n,:]
    output = test.iloc[(i+1) * historical_n: (i+1) * historical_n + h,:]
    train_ = input.copy()
    train_[target_col] = (train_[target_col] - mean)/std 
    test_ = output.copy()
    test_[target_col] = (test_[target_col] - mean)/std 
       
    
    timegpt_fcst_df = nixtla_client.forecast(df=train_, h=h, time_col=time_col, target_col=target_col, freq=freq, model='timegpt-1-long-horizon')
    timegpt_fcst_df.head()

    prediction = timegpt_fcst_df['TimeGPT']
    error = pd.Series(prediction.values - test_[target_col].values)
    mae_output = error.abs().mean()
    mse_output = (error**2).mean()
    print(str(i)+'_mae of timegpt prediction:',mae_output)
    print(str(i)+'_mse of timegpt prediction:',mse_output)
    pred_pure = pd.Series(prediction.values, index=test_['date'])
    save_name_pure = 'output/' + model + ds_name + '_pure_' + str(i) + '.csv'
    pred_pure.to_csv(save_name_pure, header=True)
    
    # prediction with GWN
    #input_gwn = GWN(input, scale, target_col)
    #input_gwn_ = input_gwn.copy()
    #input_gwn_[target_col] = (input_gwn_[target_col] - mean)/std 
    #timegpt_gwn_df = nixtla_client.forecast(df=input_gwn_, h=h, time_col=time_col, target_col=target_col, freq=freq, model='timegpt-1-long-horizon')
    #timegpt_gwn_df.head()

    #error_gwn_input = pd.Series(input_gwn[target_col].values - input[target_col].values)
    #mae_gwn_input = error_gwn_input.abs().mean()
    #print(str(i)+'_mae of timegpt input with gwn:',mae_gwn_input)
    #prediction_gwn = timegpt_gwn_df['TimeGPT']
    #error_gwn = pd.Series(prediction_gwn.values - test_[target_col].values)
    #mae_gwn = error_gwn.abs().mean()
    #print(str(i)+'_mae of timegpt prediction with gwn:',mae_gwn)
    
    #pred_gwn = pd.Series(prediction_gwn.values, index=test_['date'])
    #save_name_gwn = 'output/' + model + ds_name + '_gwn_' + str(i) + '.csv'
    #pred_gwn.to_csv(save_name_gwn, header=True)
#
    # predcition with TSA
    print('TSA begin!')
    start = time.time()
    input_tsa,s,f,r = TSA(input, nixtla_client, time_col, target_col, h, freq, mean, std, tau, epsilon)
    end = time.time()
    print(f'TSA end! Excution Time: {end - start:.2f} seconds')
    # save the poisoned input
    input_tsa_save = pd.Series(input_tsa[target_col].values, index=train_['date'])
    save_name_input_tsa = 'output/' + model + ds_name + '_tsa_input_' + str(tau) + '_' +str(i) + '.csv'
    input_tsa_save.to_csv(save_name_input_tsa, header=True)
   
    # feed poisoned input into the forecasting model
    input_tsa_ = input_tsa.copy()
    input_tsa_[target_col] = (input_tsa_[target_col] - mean)/std 
    timegpt_tsa_df = nixtla_client.forecast(df=input_tsa_, h=h, time_col=time_col, target_col=target_col, freq=freq, model='timegpt-1-long-horizon')
    timegpt_tsa_df.head()
#
    error_tsa_input = pd.Series(input_tsa[target_col].values - input[target_col].values)
    mae_tsa_input = error_tsa_input.abs().mean()
    print(str(i)+'_mae of timegpt input with tsa:',mae_tsa_input)
    prediction_tsa = timegpt_tsa_df['TimeGPT']
    error_tsa = pd.Series(prediction_tsa.values - test_[target_col].values)
    mae_tsa = error_tsa.abs().mean()
    mse_tsa = (error_tsa**2).mean()
    print(str(i)+'_mae of timegpt prediction with tsa:',mae_tsa)
    print(str(i)+'_mse of timegpt prediction with tsa:',mse_tsa)
    
    pred_tsa = pd.Series(prediction_tsa.values, index=test_['date'])
    save_name_tsa = 'output/' + model + ds_name + '_tsa_' + str(tau) + '_' + str(i) + '.csv'
    pred_tsa.to_csv(save_name_tsa, header=True)
   
