import torch
import os

from torch.utils.data import DataLoader

from .common.train_util import train_model
from .dataset.retina_dataset import RetinaDataset
from .model.super_retina import SuperRetina
import torch.optim as optim
import yaml
from torch.optim import lr_scheduler
import warnings

warnings.filterwarnings("ignore")

config_path = "./config/train.yaml"

if os.path.exists(config_path):
    with open(config_path) as f:
        config = yaml.safe_load(f)
else:
    raise FileNotFoundError("Config File doesn't Exist")

assert "MODEL" in config
assert "PKE" in config
assert "DATASET" in config
assert "VALUE_MAP" in config
train_config = {
    **config["MODEL"],
    **config["PKE"],
    **config["DATASET"],
    **config["VALUE_MAP"],
}

batch_size = train_config["batch_size"]
num_epoch = train_config["num_epoch"]
device = train_config["device"]
device = torch.device(device if torch.cuda.is_available() else "cpu")

dataset_path = train_config["dataset_path"]
data_shape = (train_config["model_image_height"], train_config["model_image_width"])

train_split_file = train_config["train_split_file"]
val_split_file = train_config["val_split_file"]
auxiliary = train_config["auxiliary"]
train_set = RetinaDataset(
    dataset_path,
    split_file=train_split_file,
    is_train=True,
    data_shape=data_shape,
    auxiliary=auxiliary,
)
val_set = RetinaDataset(
    dataset_path, split_file=val_split_file, is_train=False, data_shape=data_shape
)

load_pre_trained_model = train_config["load_pre_trained_model"]
pretrained_path = train_config["pretrained_path"]

model = SuperRetina(train_config, device=device)
if load_pre_trained_model:
    if not os.path.exists(pretrained_path):
        raise Exception("Pretrained model doesn't exist")
    checkpoint = torch.load(pretrained_path, map_location=device, weights_only=True)
    model.load_state_dict(checkpoint["net"])

optimizer = optim.Adam(model.parameters(), lr=1e-4)

dataloaders = {
    "train": DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=8),
    "val": DataLoader(val_set, batch_size=batch_size, shuffle=True, num_workers=8),
}

model = train_model(
    model,
    optimizer,
    dataloaders,
    device,
    num_epochs=num_epoch,
    train_config=train_config,
)
