import json
import random
import sys
import time

import numpy as np
import torch
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from tqdm import tqdm
from transformers import (
    AdamW,
    GPT2Config,
    GPT2LMHeadModel,
    GPT2Tokenizer,
    get_linear_schedule_with_warmup,
)

import wandb
from crazyneuraluser.user_model_code.argument import get_args

# from interact import interact
from crazyneuraluser.user_model_code.dataset import SGD_Dataset
from crazyneuraluser.user_model_code.utils_generation import decode_e2e
from crazyneuraluser.user_model_code.utils_sgd import get_special_tokens


def print_loss(epoch, data_type, LOSS, t0):
    print(
        "Epoch: {} | {} loss: {:.3f} | time: {:.1f}".format(
            epoch, data_type, LOSS, time.time() - t0
        )
    )


def print_score(epoch, data_type, res, t0):
    print(
        "Epoch: {} | {}: joint_acc: {:.2f}%, slot_acc: {:.2f}% | time: {:.1f}".format(
            epoch,
            data_type,
            res["avg_joint_acc"],
            res["avg_slot_acc"],
            time.time() - t0,
        )
    )


def run_one_epoch(data_type, dataloader, trainer, epoch, run_type, collector=None):
    t0 = time.time()
    assert data_type in ["dev", "test"]
    assert run_type in ["teacher_force", "generation"]
    model, optimizer, scheduler, tokenizer = trainer

    LOSS = 0
    # result = {"slot_acc": [], "joint_acc": []}
    # mention_match = 0
    # coref_lines = []
    iterator = enumerate(
        tqdm(
            dataloader,
            desc="Epoch {} {}".format(epoch, run_type),
            disable=args.disable_display,
        )
    )
    for step, batch in iterator:
        if run_type == "teacher_force":
            loss, logits, _ = model(
                input_ids=batch["input_ids"],
                attention_mask=batch["attention_mask"],
                token_type_ids=batch["token_type_ids"],
                labels=batch["label_ids"],
            ).values()
            LOSS += loss
        else:
            decode_e2e(args, batch, model, tokenizer, collector=collector)

    # print log
    if run_type == "teacher_force":
        LOSS /= step + 1
        print_loss(epoch, data_type, LOSS, t0)
        return LOSS
    else:  # generation
        # TODO: add evaluation code here
        return None


def set_dataloader(args, tokenizer, data_type, run_type, data_size=-1):
    dataset = SGD_Dataset(
        args, tokenizer, data_type, run_type == "generation", data_size
    )
    # 	sys.exit(1)
    if data_type == "train":
        sampler = RandomSampler(
            dataset
        )  # if args.local_rank == -1 else DistributedSampler(train_dataset)
    else:
        sampler = SequentialSampler(dataset)

    dataloader = DataLoader(
        dataset,
        sampler=sampler,
        batch_size=args.train_batch_size
        if data_type == "train"
        else args.eval_batch_size,
        collate_fn=dataset.collate_fn,
    )
    return dataloader


def train(args, tokenizer, model):

    wandb.init(
        # Set the project where this run will be logged
        project="E2E User Simulator (Alistair)",
        entity="byrne-lab",
        # We pass a run name (otherwise it’ll be randomly assigned, like sunshine-lollypop-10)
        name=args.wandb_train_run_name,
        # Track hyperparameters and run metadata
        config={
            "data_dir": args.data_dir,
            "model_name": args.model_name,
            "learning_rate": args.learning_rate,
            "gradient_accumulation_steps": args.gradient_accumulation_steps,
            "train_batch_size": args.train_batch_size,
            "eval_batch_size": args.eval_batch_size,
        },
    )

    # load data
    train_dataloader = set_dataloader(
        args, tokenizer, "train", "teacher_force", data_size=args.train_size
    )
    dev_dataloader = set_dataloader(
        args, tokenizer, "dev", "teacher_force", data_size=args.eval_size
    )

    optimizer = AdamW(model.parameters(), lr=args.learning_rate, eps=args.adam_epsilon)
    if args.use_scheduler:
        t_total = (
            len(train_dataloader) // args.gradient_accumulation_steps * args.max_epoch
        )
        scheduler = get_linear_schedule_with_warmup(
            optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
        )
    else:
        scheduler = None
    trainer = (model, optimizer, scheduler, tokenizer)

    print("Do evaluation before training!")
    model.eval()
    with torch.no_grad():
        _ = run_one_epoch("dev", dev_dataloader, trainer, -1, "teacher_force")

    print("Start training!\n{}".format("***" * 30))
    eval_step = args.eval_interval // args.train_batch_size
    best_score = -100
    global_step = 0
    no_improve_count = 0
    for epoch in range(args.max_epoch):
        # initialize for each epoch training
        t0 = time.time()
        model.train()
        model.zero_grad()
        LOSS = 0
        iterator = enumerate(
            tqdm(
                train_dataloader,
                desc="Epoch {}".format(epoch),
                disable=args.disable_display,
            )
        )
        for local_step, batch in iterator:
            loss, logits, _ = model(
                input_ids=batch["input_ids"],
                attention_mask=batch["attention_mask"],
                token_type_ids=batch["token_type_ids"],
                labels=batch["label_ids"],
            ).values()
            LOSS += loss
            global_step += 1

            wandb.log({"loss": loss})

            # update model
            if loss != 0:
                loss = loss / args.gradient_accumulation_steps
                loss.backward()

            if global_step % args.gradient_accumulation_steps == 0:
                # norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                optimizer.step()
                if args.use_scheduler:
                    scheduler.step()
                optimizer.zero_grad()

            # evaluate model
            if global_step % eval_step == 0:
                model.eval()
                with torch.no_grad():
                    loss = run_one_epoch(
                        "dev", dev_dataloader, trainer, epoch, "teacher_force"
                    )
                    score = -loss  # dev loss as criterion for early training
                    wandb.log({"dev_loss": loss})
                model.train()

                save_checkpoint(
                    args, tokenizer, model, global_step * args.train_batch_size
                )
                if score > best_score:
                    best_score = score
                    print("Best score: {:.2f}".format(best_score))
                    no_improve_count = 0
                else:
                    no_improve_count += 1

                # early stop
                if no_improve_count == args.no_improve_max:
                    print("Early stop!")
                    return

        LOSS /= local_step + 1
        print_loss(epoch, "train", LOSS, t0)
        print("***" * 30)

        wandb.log({"epoch": epoch, "epoch_loss": LOSS})

    # Mark the run as finished on wandb
    wandb.finish()


def test(args, tokenizer, model):
    # load data
    test_gen_dataloader = set_dataloader(args, tokenizer, "test", "generation")

    trainer = (model, None, None, tokenizer)
    model.eval()
    collector = {"decode-dev": {}, "decode-test": {}}
    with torch.no_grad():
        # 		# evaluate on dev
        # 		_ = run_one_epoch('dev', dev_dataloader, trainer, 'Eval', 'teacher_force')

        # 		# generate on dev
        # 		res_dev = run_one_epoch('dev', dev_gen_dataloader, trainer, 'Dev', 'generation',
        #       collector=collector['decode-dev'])
        # 		collector['result-dev'] = res_dev
        # 		print_qr_result(res_dev['qr'], 'dev')

        # generate on test
        res_test = run_one_epoch(
            "test",
            test_gen_dataloader,
            trainer,
            "Test",
            "generation",
            collector=collector["decode-test"],
        )
        collector["result-test"] = res_test

    out_file = args.decode_file
    with open(out_file, "w") as f:
        json.dump(collector, f, indent=4, sort_keys=True)
    print("Decode file is saved at {}".format(out_file))
    print("Done decoding!")


def save_checkpoint(args, tokenizer, model, step):
    save_path = args.checkpoint + "_step" + str(step)
    print("Save model in {}!".format(save_path))
    tokenizer.save_pretrained(save_path)
    model.save_pretrained(save_path)


def load_checkpoint(args):
    save_path = args.checkpoint  # + '_step' + str(args.step)
    print("Load model, tokenizer from {}".format(save_path))
    tokenizer = GPT2Tokenizer.from_pretrained(save_path)
    model = GPT2LMHeadModel.from_pretrained(save_path)
    model.to(args.device)
    return tokenizer, model


def load_pretrained_model(args):
    save_path = args.pre_checkpoint
    print("Load model, tokenizer from {}".format(save_path))
    tokenizer = GPT2Tokenizer.from_pretrained(save_path)
    model = GPT2LMHeadModel.from_pretrained(save_path)
    model.to(args.device)
    return tokenizer, model


def set_model(args, SPECIAL_TOKENS):
    """initiate config, tokenizer and model"""
    # add special tokens into tokenizer
    config = GPT2Config.from_pretrained(args.model_name_or_path)
    tokenizer = GPT2Tokenizer.from_pretrained(args.model_name_or_path)
    tokenizer.add_special_tokens(SPECIAL_TOKENS)
    model = GPT2LMHeadModel.from_pretrained(
        args.model_name_or_path, config=config
    )  # GPT2LMHeadModel
    model.resize_token_embeddings(len(tokenizer))
    model.to(args.device)
    print("Done setting model")
    return config, tokenizer, model


def set_seed(args):
    """for reproduction"""
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = False
    torch.backends.cudnn.benchmark = False


if __name__ == "__main__":
    # Load arguments
    args = get_args()

    # Set seed, device
    set_seed(args)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.device = device

    # Load special tokens
    SPECIAL_TOKENS = get_special_tokens()

    if args.mode == "training":
        config, tokenizer, model = set_model(args, SPECIAL_TOKENS)
        train(args, tokenizer, model)

    elif args.mode == "finetune":
        tokenizer, model = load_pretrained_model(args)
        train(args, tokenizer, model)

    elif args.mode == "testing":
        tokenizer, model = load_checkpoint(args)
        test(args, tokenizer, model)

    # 	elif args.mode == 'interact':
    # 		tokenizer, model = load_checkpoint(args)
    # 		interact(args, tokenizer, model)

    else:
        sys.exit(1)