import time
import numpy as np
import pandas as pd
from tqdm import tqdm
from ABBA.ABBA import ABBA
from sklearn.metrics import mean_squared_error

np.random.seed(1)

length = 5000
mu, sigma = 0, 1 # mean and standard deviation
data_df = np.zeros((100, length))
runtime = np.zeros((100, 2)) # symbols: 10, 100, 1000
accuracy = np.zeros((100, 2)) # symbols: 10, 100, 1000
k_symbols = [100, 1000]
TOL = 0.5

for i in tqdm(range(100)):
    ts = np.random.normal(mu, sigma, length)
    data_df[i] = ts
    
    st = time.time()
    abba = ABBA(min_k=k_symbols[0], max_k=k_symbols[0], tol=TOL, verbose=0, scl=1)
    string, centers = abba.transform(ts)
    reconstructed_ts = abba.inverse_transform(string, centers, ts[0])
    runtime[i, 0] = time.time() - st
    accuracy[i, 0] = mean_squared_error(ts, reconstructed_ts, squared=False)
    
    st = time.time()
    abba = ABBA(min_k=k_symbols[1], max_k=k_symbols[1], tol=TOL, verbose=0, scl=1)
    string, centers = abba.transform(ts)
    reconstructed_ts = abba.inverse_transform(string, centers, ts[0])
    runtime[i, 1] = time.time() - st
    accuracy[i, 1] = mean_squared_error(ts, reconstructed_ts, squared=False)

    
pd.DataFrame(data_df).to_csv("ts_test.csv", index=False)
pd.DataFrame(runtime).to_csv("runtime.csv", index=False)
pd.DataFrame(accuracy).to_csv("accuracy.csv", index=False)