import torch
from aiRNN import dataloader, models, losses
import pathlib
import pandas as pd

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")

results = []
for step in [1, 5, 10, 30]:
    dataset = dataloader.SaveDataset(torch.load("dataset.pt"), device="cuda", step=step)
    dataset_length = len(dataset)
    dataset, val_dataset = torch.utils.data.random_split(
        dataset,
        [int(0.8 * dataset_length), dataset_length - int(0.8 * dataset_length)],
        generator=torch.Generator().manual_seed(42)
    )
    print(f"Starting hyperparameter scan for step={step}")
    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)
                loss = -1.0
                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()
                        warm_up_steps = 450 // step
                        pred_steps = 150 // step
                        y_pred, _ = model(X_t, X_f, X_c, warm_up_steps, pred_steps)
                        loss = criterion(y_pred, y)
                        loss.backward()
                        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                        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']}")
                with torch.no_grad():
                    val_loss_history = []
                    for X_f, X_t, y, X_c in torch.utils.data.DataLoader(val_dataset, batch_size=256):
                        warm_up_steps = 450 // step
                        pred_steps = 150 // step
                        y_pred, _ = model(X_t, X_f, X_c, warm_up_steps, pred_steps)
                        v_loss = criterion(y_pred, y)
                        val_loss_history.append(v_loss.item())
                    loss = sum(val_loss_history) / len(val_loss_history)
                    print(f"Validation loss: {loss}")
                torch.save(model.state_dict(), f"{base_name}_hs{hidden_size}_rs{rnn_size}_step{step}.pt")
                results.append((base_name, hidden_size, rnn_size, step, loss))

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

# Save results to CSV
df = pd.DataFrame(results, columns=["Model", "Hidden Size", "RNN Size", "Step", "Final Loss"])
df.to_csv("hyperparameter_scan_results.csv", index=False)