import torch
import time
import logging

from graphgym.config import cfg
from graphgym.loss import compute_loss
from graphgym.utils.epoch import is_eval_epoch, is_ckpt_epoch
from graphgym.checkpoint import load_ckpt, save_ckpt, clean_ckpt


def add_custom_stats_from_batch(batch):
    custom_stats = {}

    keys = list(batch.__dict__.keys())
    if len(keys) == 1 and keys[0] == '_store':
       my_dict = batch.__dict__['_store']
    else:
        my_dict =  batch.__dict__



    for attr_name in my_dict.keys():
        if 'alpha_' in attr_name:
            custom_stats[attr_name] =  round(getattr(batch, attr_name).item(), 6)
        if 'lmbda' in attr_name:
            custom_stats[attr_name] = round(getattr(batch, attr_name), 6)
        if 'score_' in attr_name:
            custom_stats[attr_name] = round(getattr(batch, attr_name).item(), 6)
    return custom_stats

def train_epoch(logger, loader, model, optimizer, scheduler):
    model.train()
    time_start = time.time()

    i = 0
    num_max = len(loader)
    if cfg.dataset.use_subset < 1.0:
        num_max = int(num_max*cfg.dataset.use_subset)
    else:
        num_max = len(loader) + 4
    for batch in loader:
        i+=1
        if i == num_max: break
        optimizer.zero_grad()
        batch.to(torch.device(cfg.device))
        pred, true = model(batch)
        loss, pred_score = compute_loss(pred, true)
        loss.backward()
        if cfg.train.clip >0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.train.clip)

        optimizer.step()

        custom_stats = add_custom_stats_from_batch(batch)
        logger.update_stats(true=true.detach().cpu(),
                            pred=pred_score.detach().cpu(),
                            loss=loss.item(),
                            lr=scheduler.get_last_lr()[0],
                            time_used=time.time() - time_start,
                            params=cfg.params,
                            **custom_stats)
        time_start = time.time()
    scheduler.step()


@torch.no_grad()
def eval_epoch(logger, loader, model):
    model.eval()
    time_start = time.time()
    for batch in loader:
        batch.to(torch.device(cfg.device))
        pred, true = model(batch)
        loss, pred_score = compute_loss(pred, true)
        custom_stats = add_custom_stats_from_batch(batch)
        logger.update_stats(true=true.detach().cpu(),
                            pred=pred_score.detach().cpu(),
                            loss=loss.item(),
                            lr=0,
                            time_used=time.time() - time_start,
                            params=cfg.params,
                            **custom_stats)
        time_start = time.time()


def train(loggers, loaders, model, optimizer, scheduler):
    start_epoch = 0
    if cfg.train.auto_resume:
        start_epoch = load_ckpt(model, optimizer, scheduler)
    if start_epoch == cfg.optim.max_epoch:
        logging.info('Checkpoint found, Task already done')
    else:
        logging.info('Start from epoch {}'.format(start_epoch))

    num_splits = len(loggers)
    for cur_epoch in range(start_epoch, cfg.optim.max_epoch):
        train_epoch(loggers[0], loaders[0], model, optimizer, scheduler)
        loggers[0].write_epoch(cur_epoch)
        if is_eval_epoch(cur_epoch):
            for i in range(1, num_splits):
                eval_epoch(loggers[i], loaders[i], model)
                loggers[i].write_epoch(cur_epoch)
        if is_ckpt_epoch(cur_epoch):
            save_ckpt(model, optimizer, scheduler, cur_epoch)
    for logger in loggers:
        logger.close()
    if cfg.train.ckpt_clean:
        clean_ckpt()

    logging.info('Task done, results saved in {}'.format(cfg.out_dir))
