import logging
import os
import random
import sys
import time

import numpy as np
import torch
from data import create_dataloader
from networks.trainer import Trainer
from options.train_options import TrainOptions
from tensorboardX import SummaryWriter
from validate import validate


def seed_torch(seed=1029):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = False


def get_val_opt():
    val_opt = TrainOptions().parse(print_options=False)
    val_opt.isTrain = False
    val_opt.no_resize = False
    val_opt.no_crop = False
    val_opt.serial_batches = True
    return val_opt


if __name__ == '__main__':
    opt = TrainOptions().parse()
    seed_torch(100)
    opt.dataroot = f'{opt.dataroot}/{opt.train_split}/'
    print('  '.join(list(sys.argv)))
    val_opt = get_val_opt()
    val_opt.dataroot = f'{val_opt.dataroot}/{val_opt.val_split}/'

    data_loader = create_dataloader(opt)
    val_loader = create_dataloader(val_opt)

    train_writer = SummaryWriter(os.path.join(opt.checkpoints_dir, opt.name, 'train'))
    val_writer = SummaryWriter(os.path.join(opt.checkpoints_dir, opt.name, 'val'))

    model = Trainer(opt)
    model.train()

    # early_stopping = EarlyStopping(patience=opt.earlystop_epoch, delta=-0.001, verbose=True)

    # Configure logger
    logging.basicConfig(
        level=logging.INFO,  # Set logging level
        format='%(asctime)s - %(levelname)s - %(message)s',  # Log format
        datefmt='%Y-%m-%d %H:%M:%S',  # Date format without milliseconds
        handlers=[
            logging.FileHandler('log.log', mode='w'),  # Log file output
            logging.StreamHandler(),
        ],
    )  # Console output

    print(f'cwd: {os.getcwd()}')
    print(
        f'{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} Length of data loader: {len(data_loader)}'
    )

    for epoch in range(opt.niter):
        for i, data in enumerate(data_loader):
            model.total_steps += 1

            # Move data to cuda
            model.set_input(data)
            # Train
            model.optimize_parameters()

            if model.total_steps % opt.loss_freq == 0:
                print(
                    f'{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} Train loss: {model.loss} at step: {model.total_steps} lr {model.lr}'
                )
                train_writer.add_scalar('loss', model.loss, model.total_steps)

        if epoch % opt.save_epoch_freq == 0 and epoch != 0:
            # if epoch == 0:
            # Optionally save model at each epoch
            print(f'saving the model at the end of epoch {epoch}')
            model.save_networks(epoch)

        if epoch % 10 == 0 and epoch != 0:
            print(
                f'{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} changing lr at the end of epoch {epoch}, iters {model.total_steps}'
            )
            model.adjust_learning_rate()

        # Validation
        model.eval()
        acc, ap, r_acc, f_acc = validate(model.model, val_loader)
        val_writer.add_scalar('accuracy', acc, model.total_steps)
        val_writer.add_scalar('ap', ap, model.total_steps)
        print(
            f'{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())} (Val @ epoch {epoch}) acc: {acc}; ap: {ap} r_acc: {r_acc}; f_acc: {f_acc}'
        )

        model.train()

    model.eval()
    model.save_networks('last')