from clearml import Task
from src.data import DataProcessor, DataConfig
from src.trainers.trainer import Trainer
from src.utils.clearml import ClearMLHelper
from src.models import *
from src.losses import *
import torch
import numpy as np
from torch.nn import MSELoss, L1Loss
from datetime import datetime
import torch.nn as nn
from src.models.time_embedding_layer import TimeEmbedding
from src.models.diffusion_model import SimpleDiffusionModel
from src.trainers.diffusion_trainer import DiffusionTrainer


clearml_helper = ClearMLHelper(project_name="Thesis/NrvForecast")
task = clearml_helper.get_task(task_name="Diffusion Training")

# execute remotely
task.execute_remotely(queue_name="default", exit_process=True)
print("Running remotely")

#### Data Processor ####
data_config = DataConfig()
data_config.NRV_HISTORY = True
data_config.LOAD_HISTORY = True
data_config.LOAD_FORECAST = True

data_config.WIND_FORECAST = True
data_config.WIND_HISTORY = True

data_config.QUARTER = False
data_config.DAY_OF_WEEK = False

data_config.NOMINAL_NET_POSITION = True

data_config = task.connect(data_config, name="data_features")

data_processor = DataProcessor(data_config, path="", lstm=False)
data_processor.set_batch_size(8192)
data_processor.set_full_day_skip(True)

inputDim = data_processor.get_input_size()

model_parameters = {
    "epochs": 5000,
    "learning_rate": 0.0001,
    "hidden_sizes": [512, 512, 512],
    "time_dim": 64,
}

model_parameters = task.connect(model_parameters, name="model_parameters")

#### Model ####
model = SimpleDiffusionModel(96, model_parameters["hidden_sizes"], other_inputs_dim=inputDim[1], time_dim=model_parameters["time_dim"])

print("Starting training ...")

#### Trainer ####
trainer = DiffusionTrainer(model, data_processor, "cuda")
trainer.train(model_parameters["epochs"], model_parameters["learning_rate"], task)