# %%
import torch
from aiRNN import dataloader, models, losses
import pathlib

# %%
file_list = pathlib.Path("/home/lars/Documents/Studium/UiO/data_analysis/project3/Code/cpp/known_routes_and_aircraft.csv")
base_path = file_list.parent
file_list = file_list.read_text().splitlines()
file_list = [(base_path / f).resolve() for f in file_list]

# %%
dataset = dataloader.EvenlySpacedDataset(
    filepaths=file_list[:500],
    n_input=30*10, # 10 minutes input
    n_output=30*1,  # 1 minutes output
    n_windows_per_file=5,
    step=1,
    feature_columns=("lat", "lon", "alt", "ias"),
    context_columns=("last_lat", "last_lon", "last_alt", "last_ias"),
    time_columns=("timestamp", "dt"),
    target_columns=("lat", "lon", "alt"),
)

# %%
len(dataset)

# %%
test_model = models.ThreeInputRNN(
    time_in=2,
    feat_in=4,
    context_in=4,
    hidden_size=128,
    rnn_size=256,
    out_size=3,
)

# %%
# Example training
optimizer = torch.optim.Adam(test_model.parameters(), lr=1e-2)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
criterion = torch.nn.HuberLoss()
for epoch in range(50):
    losses = []
    for X_f, X_t, y, X_c in torch.utils.data.DataLoader(dataset, batch_size=256, shuffle=True):
        optimizer.zero_grad()
        y_pred, _ = test_model(X_t, X_f, X_c, 300, 30)
        loss = criterion(y_pred, y)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    loss = sum(losses) / len(losses)
    scheduler.step(loss)
    print(f"Epoch {epoch}: loss={loss}, lr={optimizer.param_groups[0]['lr']}")

# %%