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

file_list = pathlib.Path("../../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 if (base_path / f).exists()]

if not pathlib.Path("dataset.pt").exists():
  dataset = dataloader.EvenlySpacedDataset(
      filepaths=file_list,
      n_input=30*15, # 15 minutes input
      n_output=30*5,  # 5 minutes output
      n_windows_per_file=7,
      step=1,
      feature_columns=("lat", "lon", "alt", "ias"),
      context_columns=("last_lat", "last_lon", "last_alt", "last_ias", "last_timestamp"),
      time_columns=("timestamp", "dt"),
      target_columns=("lat", "lon", "alt"),
      device="cuda",
  )
  dataset.save_entire_dataset("dataset.pt")

dataset = dataloader.SaveDataset(torch.load("dataset.pt"), device="cuda")
results = []

for base_name, base_model in [
    ("RNN", models.ThreeInputRNN),
    ("LSTM", models.ThreeInputLSTM),
    ("GRU", models.ThreeInputGRU),
]:
    for hidden_size in [16, 32, 64]:
        for rnn_size in [32, 64, 128]:
            model = base_model(
                time_in=2,
                feat_in=4,
                context_in=5,
                hidden_size=hidden_size,
                rnn_size=rnn_size,
                out_size=3,
                device="cuda",
            )
            print(f"Training {base_name} with hidden_size={hidden_size}, rnn_size={rnn_size}")
            optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
            criterion = losses.HaversineAltitudeLoss(alt_const=1e-3)
            for epoch in range(100):
                loss_history = []
                for X_f, X_t, y, X_c in torch.utils.data.DataLoader(dataset, batch_size=256, shuffle=True):
                    optimizer.zero_grad()
                    y_pred, _ = model(X_t, X_f, X_c, 450, 150)
                    loss = criterion(y_pred, y)
                    loss.backward()
                    optimizer.step()
                    loss_history.append(loss.item())
                loss = sum(loss_history) / len(loss_history)
                scheduler.step(loss)
                print(f"Epoch {epoch}: loss={loss}, lr={optimizer.param_groups[0]['lr']}")
            torch.save(model.state_dict(), f"{base_name}_hs{hidden_size}_rs{rnn_size}.pt")
            results.append((base_name, hidden_size, rnn_size, loss))

for r in results:
    print(f"Model: {r[0]}, hidden_size={r[1]}, rnn_size={r[2]} => final loss={r[3]}")