import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from nixtla import NixtlaClient

def GWN(df, scale, target_col):
    l = len(df[target_col])
    noise = df.copy()
    noise[target_col] = noise[target_col] + np.random.normal(0, scale, l)
    return noise

def TSA(df, nixtla_client, time_col, target_col, h, freq, mean, std, tau, epsilon):
    l = len(df[target_col])
    s = []
    f = []
    for i in range(tau):
        index = 0
        r = 0
        temp = 0
        for i in range(int(h)):
            if i + (l-h) not in s: 
                df_1 = df.copy()
                df_2 = df.copy()
                df_ = df.copy()
                df_1[target_col] = (df_1[target_col] - mean)/std
                df_2[target_col] = (df_2[target_col] - mean)/std
                df_[target_col] = (df_[target_col] - mean)/std
                
                s_temp = s + [int(i + (l-h/2))]
                f_temp_1 = f + [epsilon]
                f_temp_2 = f + [-1*epsilon]
                
                df_1[target_col].iloc[s_temp] = df_1[target_col].iloc[s_temp] * (1+np.array(f_temp_1))
                df_2[target_col].iloc[s_temp] = df_2[target_col].iloc[s_temp] * (1+np.array(f_temp_2))
            
                timegpt_fcst_df_1 = nixtla_client.forecast(df=df_1, h=h, time_col=time_col, target_col=target_col, freq=freq, model='timegpt-1-long-horizon')
                pred_1 = timegpt_fcst_df_1['TimeGPT']
                timegpt_fcst_df_2 = nixtla_client.forecast(df=df_2, h=h, time_col=time_col, target_col=target_col, freq=freq, model='timegpt-1-long-horizon')
                pred_2 = timegpt_fcst_df_2['TimeGPT']
                timegpt_fcst_df = nixtla_client.forecast(df=df_, h=h, time_col=time_col, target_col=target_col, freq=freq, model='timegpt-1-long-horizon')
                pred_ = timegpt_fcst_df['TimeGPT']
            
                dis_1 = pd.Series(pred_1.values - pred_.values)
                dis_2 = pd.Series(pred_2.values - pred_.values)
            
                dis_1 = dis_1.abs().mean()
                dis_2 = dis_2.abs().mean()
            
                if dis_1>dis_2:
                    if dis_1>r:
                        r = dis_1
                        index = int(i + (l-h/2))
                        temp = epsilon
                elif dis_2>dis_1:
                    if dis_2>r:
                        r = dis_2
                        index = int(i + (l-h/2))
                        temp = -1* epsilon
        s.append(index)
        f.append(temp)
        
    print('s is :',s)
    print('direction is :', f)
    print('prediction error is:', r)
    noise = df.copy()
    noise[target_col].iloc[s] = noise[target_col].iloc[s] * (1+np.array(f))
    
    return noise,s,f,r
       
        
def DGA(df, test, nixtla_client, scale, time_col, target_col, h, freq, mean, std):
    l = len(df[target_col])
    u = (np.random.rand(l) - 0.5) * scale
    
    df_1 = df.copy()
    df_2 = df.copy()
    df_1[target_col] = (df_1[target_col] + u - mean)/std
    df_2[target_col] = (df_2[target_col] - mean)/std
    
    #df_target = pd.DataFrame({
    #time_col: test[time_col],          
    #target_col: np.random.normal(mean,std,h)})
    target = np.random.normal(0,1,h)
                              
    timegpt_fcst_df_1 = nixtla_client.forecast(df=df_1, h=h, time_col=time_col, target_col=target_col, freq=freq, model='timegpt-1-long-horizon')
    pred_1 = timegpt_fcst_df_1['TimeGPT']
    timegpt_fcst_df_2 = nixtla_client.forecast(df=df_2, h=h, time_col=time_col, target_col=target_col, freq=freq, model='timegpt-1-long-horizon')
    pred_2 = timegpt_fcst_df_2['TimeGPT']
    
    #print(pred_1)
    dis_1 = pd.Series(pred_1.values - target)
    dis_2 = pd.Series(pred_2.values - target)
    
    #dis_1 = pred_1 - df_target[target_col]
    #dis_2 = pred_2 - df_target[target_col]
    #print(dis_1)
    #print(dis_2)
    gradient = (dis_1.abs().sum()-dis_2.abs().sum()) / u
    #print(gradient)
    
    noise = df.copy()
    noise[target_col] = noise[target_col] + scale * np.sign(gradient)
    
    return noise



   