import json

from pathlib import Path

import torch

import torch.nn as nn

from torch.optim import Adam

from tqdm import tqdm

from argparse import Namespace

from data import GIImage, GIImageDataset, GIImageDataLoader, train_valid_split_cases

from utils import load_yaml, set_seed

from model import UNet

from evaluation import evaluate

def train(config: Namespace):

    exp_path = Path(config.save_path) / config.exp_name

    exp_path.mkdir(parents=True, exist_ok=True)

    set_seed(config.seed)

    model = UNet(n_classes=len(GIImage.organs))

    train_cases, valid_cases = train_valid_split_cases(config.input_path, config.valid_size)

    train_set = GIImageDataset(image_path=config.input_path, label_path=config.label_path, cases=train_cases)

    valid_set = GIImageDataset(image_path=config.input_path, label_path=config.label_path, cases=valid_cases)

    train_loader = GIImageDataLoader(

        model=model,

        dataset=train_set,

        batch_size=config.batch_size,

        shuffle=True,

        input_resolution=config.input_resolution,

        padding_mode=config.padding_mode

    ).get_data_loader()

    valid_loader = GIImageDataLoader(

        model=model,

        dataset=valid_set,

        batch_size=config.batch_size,

        shuffle=False,

        input_resolution=config.input_resolution,

        padding_mode=config.padding_mode

    ).get_data_loader()

    # TODO: use other loss functions and optimizer

    scaler = torch.cuda.amp.GradScaler(enabled=config.use_fp16)

    criterion = nn.CrossEntropyLoss()

    optimizer = Adam(model.parameters(), lr=config.lr)

    model = model.to(config.device)

    best_valid_loss = evaluate(model, valid_loader, criterion, config.device, config.use_fp16) # TODO: track model performance with other metrics

    valid_losses = [best_valid_loss]

    (exp_path / "valid_losses.json").write_text(json.dumps(valid_losses, indent=4))

    for epoch in range(config.nepochs):

        model.train()

        pbar = tqdm(train_loader)

        pbar.set_description(f"Epoch {epoch + 1}")

        for i, (inputs, labels) in enumerate(train_loader):

            inputs = inputs.to(config.device)

            labels = labels.to(config.device)

            with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=config.use_fp16):

                preds = model(inputs)

                loss = criterion(preds, labels)

            scaler.scale(loss).backward()

            scaler.step(optimizer)

            scaler.update()

            optimizer.zero_grad(set_to_none=True)

            pbar.update()

            if (i != 0 and i % config.valid_steps == 0) or (i == len(train_loader) - 1):

                total_valid_loss = evaluate(model, valid_loader, criterion, config.device, config.use_fp16)

                valid_losses.append(total_valid_loss)

                print(f"Valid loss: {total_valid_loss:.4f}")

                # save model if validation loss is improved

                if total_valid_loss < best_valid_loss:

                    best_valid_loss = total_valid_loss

                    torch.save(model.state_dict(), exp_path / "model.pth")

                    (exp_path / "valid_losses.json").write_text(json.dumps(valid_losses, indent=4))

                    print(f"Model saved at {config.save_path}")

if __name__ == "__main__":

    config = load_yaml("config.yml")

    train(config)