import os
import json

import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler
from tqdm import tqdm 

from utils.helper import get_lr, evaluate
from utils import constants
from pytorchtools import EarlyStopping
from devkit.core.dist_utils import average_gradients, broadcast_params


def fit_one_cycle(epochs, start_epoch, max_lr, model, train_loader, val_loader, train_sampler, val_sampler, 
                  pathwrapper, weight_decay=0, grad_clip=None, opt_func=torch.optim.SGD, logfunc=None, 
                  ddpmodule=False, savepath=None, g_writer=None, last_best_loss=float('inf')):
    if logfunc is None:
        logfunc = print
    
    torch.cuda.empty_cache()
    scaler = GradScaler()

    history = []
    best_val_acc = 0
    best_epoch = start_epoch
    best_loss = last_best_loss

    # patience = 20
    # early_stopping = EarlyStopping(patience=patience, verbose=True)
    # setup custom optimizer with weight decay
    optimizer = opt_func(model.parameters(), max_lr, weight_decay=weight_decay)
    # setup one-cycle learning rate scheduler
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr, epochs=epochs, steps_per_epoch=len(train_loader))

    for epoch in range(start_epoch, epochs):
        # Training Phase
        model.train()

        ##################################
        if constants.USE_DDP:
            train_sampler.set_epoch(epoch)
            val_sampler.set_epoch(epoch)
        ##################################

        train_losses = []
        train_accs = []
        lrs = []

        for batch in tqdm(train_loader):
            optimizer.zero_grad()
            
            if constants.MIX_TRAINING:
                ############################
                # 建立autocast的上下文语句
                with autocast():  
                    if ddpmodule:
                        loss, acc = model.module.training_step(batch)
                    else:
                        loss, acc = model.training_step(batch)

                last_scale = scaler.get_scale()
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
                skip_lr_sched = (last_scale != scaler.get_scale())
                ##########################

                train_losses.append(loss)
                train_accs.append(acc)
                
                ##########################
                if constants.USE_DDP:
                    average_gradients(model)
                ##########################

                # Gradient clipping
                if grad_clip:
                    nn.utils.clip_grad_value_(model.parameters(), grad_clip)

                # optimizer.step()
                # Record and update learning rate
                ##########################
                if not skip_lr_sched:
                    lrs.append(get_lr(optimizer))
                    scheduler.step()
                ##########################

            else:
                if ddpmodule:
                    loss, acc = model.module.training_step(batch)
                else:
                    loss, acc = model.training_step(batch)
                train_losses.append(loss)
                train_accs.append(acc)
                
                ##########################
                loss.backward()
                ##########################

                ##########################
                if constants.USE_DDP:
                    average_gradients(model)
                ##########################

                # Gradient clipping
                if grad_clip:
                    nn.utils.clip_grad_value_(model.parameters(), grad_clip)

                optimizer.step()
                # Record and update learning rate
                lrs.append(get_lr(optimizer))
                scheduler.step()

        current_lr = get_lr(optimizer)

        # Validation phase
        if ddpmodule:
            result = evaluate(model.module, val_loader)
        else:
            result = evaluate(model, val_loader)
        train_losses_mean = torch.stack(train_losses).mean().item()
        train_accs_mean = torch.stack(train_accs).mean().item()

        result['train_loss'] = train_losses_mean
        result['train_acc'] = train_accs_mean
        result['lrs'] = lrs
        
        if ddpmodule:
            model.module.epoch_end(epoch, result)
        else:
            model.epoch_end(epoch, result)

        # EarlyStopping
        # early_stopping(result['val_loss'], model)
        # if early_stopping.early_stop:
        #     print("Early stopping")
        #     break
      
        if constants.G_RANK == 0:
            # ＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝
            # 计算整体的性能指标
            # 通过之前获取的writer反馈给平台
            # 这些指标的名称、作为key、是需要与平台端对齐
            if g_writer is not None:
                g_writer.append_one_line([epoch, train_accs_mean, current_lr, train_losses_mean])

            else:
                logfunc(fr'input g_writer is None ! write nothing')
            # ＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝

            # finding out the model with best val_acc
            # is_best = result['val_acc'] > best_val_acc
            is_best = result['val_loss'] < best_loss
            # model_path = constants.CHECKPOINT_PATH if savepath is None else fr'{savepath}/model.pth'
            model_path = constants.CHECKPOINT_PATH if savepath is None else fr'{savepath}/model.pth'
            if is_best:
                best_epoch = epoch
                logfunc(fr'saving best model @ Loss: {result["val_loss"]}, ACC: {best_val_acc} -> {result["val_acc"]}, epoch: {epoch}')
            else:
                logfunc(fr'saving model @ Loss: {result["val_loss"]}, ACC: {result["val_acc"]}, epoch: {epoch}')

            best_val_acc = max(best_val_acc, result['val_acc'])
            best_loss = min(best_loss, result['val_loss'])

            state = {
                'epoch': epoch,
                'state_dict': model.module.state_dict() if ddpmodule else model.state_dict() ,
                'optimizer_state_dict': optimizer.state_dict(),
                'best_loss': result['val_loss'],
                'best_acc': result['val_acc'],
            }
            # save_checkpoint(state=state, is_best=is_best, filename=model_path)

            # ＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝
            # 存储训练的权重
            logfunc(fr'writing ckpt to {pathwrapper.get_output_train_latestmodel_pth_filepath()}')
            torch.save(state, pathwrapper.get_output_train_latestmodel_pth_filepath())
                
            if is_best:
                logfunc(fr'writing best ckpt to {pathwrapper.get_output_train_bestmodel_pth_filepath()}')
                torch.save(state, pathwrapper.get_output_train_bestmodel_pth_filepath())
            # ＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝

            history.append(result)

        # ＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝
        # wait for all process before check cancle flag
        torch.distributed.barrier() # change
        if pathwrapper.get_input_cancle_flag():
            logfunc(fr'epoch end, cancle flag detected, stop training epochs')
            break
        # ＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝＝

    return history, best_epoch, best_loss, best_val_acc


def save_checkpoint(state, is_best, filename):
    import shutil
    import os

    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, os.path.join(os.path.dirname(filename), 'best_loss.pth'))
