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

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def split_dataset(dataset, frac=0.8, seed=42):
    n = len(dataset)
    train_n = int(frac * n)
    return torch.utils.data.random_split(
        dataset,
        [train_n, n - train_n],
        generator=torch.Generator().manual_seed(seed),
    )

step = 10
base_name = "LSTM"
cls = models.ThreeInputLSTM
hidden_size = 16
rnn_size = 64
hidden_layers = 0
rnn_layers = 3
rnn_dropout = 0.0
altitude_weight = 1e-3
warm = 900
pred = 1600
start_offset = 30 * 60 - warm
end_offset = 30 * 60 - pred
base_ds = dataloader.SaveDataset(
        torch.load("long_dataset.pt"),
        step=step,
        start_offset=start_offset,
        end_offset=end_offset,
    )
train_ds, val_ds = split_dataset(base_ds)
model = cls(
            time_in=2,
            feat_in=4,
            context_in=5,
            hidden_size=hidden_size,
            rnn_size=rnn_size,
            out_size=3,
            hidden_layers=hidden_layers,
            rnn_layers=rnn_layers,
            rnn_dropout=rnn_dropout,
            device=DEVICE,
        )

model.load_state_dict(torch.load("LSTM_wu900_ps150_aw0.001.pt", map_location=torch.device('cpu')))
model.to(DEVICE)
model.eval()
with torch.no_grad():
    val_loader = torch.utils.data.DataLoader(val_ds, batch_size=64, collate_fn=dataloader.collate_to_cpu, num_workers=4)
    results = []
    for X_f, X_t, y, X_c in val_loader:
        X_f = X_f.to(DEVICE)
        X_t = X_t.to(DEVICE)
        y = y.to(DEVICE)
        if X_c is not None:
            X_c = X_c.to(DEVICE)
        y_pred, _ = model(X_t, X_f, X_c, warm // step, pred // step)
        results.append((y.cpu(), y_pred.cpu()))
    y_true = torch.cat([r[0] for r in results], dim=0)
    y_pred = torch.cat([r[1] for r in results], dim=0)
    np_y_all = np.concatenate(
        [y_true.numpy().reshape(-1, 3), y_pred.numpy().reshape(-1, 3)], axis=1
    )
    np.savetxt(
        f"predictions_{base_name}_wu{warm}_ps{pred}_aw{altitude_weight}.csv",
        np_y_all,
        delimiter=",",
        header="true_x,true_y,true_z,pred_x,pred_y,pred_z",
        comments="",
    )