# %%
import torch
import torch.nn.functional as F
import pandas as pd
import numpy as np
from aiRNN import preprocessors

# %%
df = pd.read_csv("../all_second_lines.csv", header=None, names=["icao", "r", "t", "timestamp", "lat", "lon", "alt", "ias"])
_ = preprocessors.create_category_mappings(df)

# %%
df = preprocessors.map_categories(df, _)

# %%
X_a_1 = F.one_hot(torch.tensor(df["R_1_IDX"].values)).float()
X_a_2 = F.one_hot(torch.tensor(df["R_2_IDX"].values)).float()
X_t = F.one_hot(torch.tensor(df["T_IDX"].values)).float()

# %%
print(X_a_1.shape, X_a_2.shape, X_t.shape)

# %%
preprocessors.train_autoencoder(X_a_1=X_a_1, X_a_2=X_a_2, X_b=X_t, num_epochs=500, learning_rate=0.011)

# %%