import argparse
from argparse import Namespace
import os
from pathlib import Path

import subprocess

from submit import submit_local, submit_job, format_kwargs_fairseq, get_name, check_not_known

import sys
currentdir = os.path.dirname(os.path.realpath(__file__))
parentdir = os.path.dirname(currentdir)
sys.path.append(parentdir)


from checkpoints import get_checkpoints, load_checkpoint

import logging

from pretrain_v2.get_port import get_port

logging.basicConfig(level=logging.INFO)

# def format_args_fairseq(args: Namespace):
#     def lower(s):
#         return "--" + s.strip("-").replace("_", "-")
#     arg_list = []
#     for key in args._get_args():
#         print(key)
#         input()
#         arg_list.append(lower(key))
#     return arg_list


def format_args(exp_args: Namespace, args: Namespace):
    kwargs_list = format_kwargs_fairseq(args)
    args_list = [
        "--layernorm-embedding",
        "--fp16",
        "--no-epoch-checkpoints",
        "--skip-invalid-size-inputs-valid-test"
    ]
    # args_list.extend(format_args_fairseq(args))
    return kwargs_list, args_list

def get_slurm_str(args, exp_args: Namespace):
    if exp_args.sbatch:
        Path(f"slurm_output/{args.wandb_name}").mkdir(exist_ok=True, parents=True)
        o_dir = f"{exp_args.save_dir}/slurm-%j.out"
        e_dir = f"{exp_args.save_dir}/slurm-%j.err"
        print(f"output: {exp_args.save_dir}/slurm-")
        slurm_str = f"sbatch -c {exp_args.cpus} --gpus={exp_args.gpus} --nodes={exp_args.nodes} -n 1 -t {exp_args.ndays}-00:00:00 -o {o_dir} -e {e_dir} --constraint=\"type_e\""  # 
    else:
        slurm_str = ""
    return slurm_str

def fairseq_train(exp_args: Namespace, args: Namespace) -> str:
    slurm_str = get_slurm_str(args, exp_args)
    fairseq_path, data_dir = exp_args.fairseq_path, exp_args.data_dir
    kwargs_list, args_list = format_args(exp_args, args)
    args = [data_dir] + kwargs_list + args_list

    args = " ".join([slurm_str] + [fairseq_path] + list(map(str,args)))
    if exp_args.sbatch:
        job_id = submit_job(args)
    else:
        job_id = submit_local(args)
    return job_id

def add_multilingual_denoising_arguments(parser):
    # Tasks args
    parser.add_argument("--multilang_sampling_alpha", default=0.3, type=float)
    parser.add_argument("--insert", default=0, type=str)
    parser.add_argument("--permute_sentences", default=0, type=float)
    parser.add_argument("--permute_sentences_delim", default='\"_;\"', type=str)
    parser.add_argument("--poisson_lambda", default=3.5, type=float)
    parser.add_argument("--mask", default=0.3, type=float, 
                        help='fraction of words/subwords that will be masked')
    parser.add_argument("--mask_length", default='span-poisson', type=str)
    parser.add_argument("--replace_length", default=1, type=float)
    parser.add_argument("--rotate", default=0, type=float)
    parser.add_argument("--mask_random", default=0.1, type=float, \
                        help='instead of using [MASK], use random token this often')
    parser.add_argument("--tokens_per_sample", default=512, type=int)
    parser.add_argument("--add-lang-token", action="store_false")

    parser.add_argument("--varmisuse", default=0.0, type=float)

def add_translation_from_pretrained_bart_arguments(parser):
    parser.add_argument('--prepend_bos', action='store_true',
        help='prepend bos token to each sentence, which matches '
                'mBART pretraining')
    parser.add_argument('-s', '--source-lang', default="src", metavar='SRC',
                            help='source language')
    parser.add_argument('-t', '--target-lang', default="tgt", metavar='TARGET',
                        help='target language')



def get_task():
    task_parser = argparse.ArgumentParser("task")
    task_parser.add_argument("--task", default="multilingual_denoising", 
        choices=["multilingual_denoising","translation_from_pretrained_bart"], type=str)
    args, _ = task_parser.parse_known_args()
    return args.task

def add_taskspecific_args(parser):
    task = get_task()
    if task == "multilingual_denoising":
        add_multilingual_denoising_arguments(parser)
    elif task == "translation_from_pretrained_bart":
        add_translation_from_pretrained_bart_arguments(parser)
    else:
        raise ValueError(f"Unknown task {task_parser.task}")



def get_fairseq_arguments():
    parser = argparse.ArgumentParser("fairseq-train")
    parser.add_argument("--task", default="multilingual_denoising", type=str)
    parser.add_argument("--wandb_project", type=str, default="PLBART")
    parser.add_argument("--wandb_name", default="test", type=str)
    parser.add_argument("--wandb_group", default="default", type=str)

    # Distributed args
    parser.add_argument("--ddp_backend", choices=["c10d", "no_c10d"], default="no_c10d", type=str)
    # parser.add_argument("--distributed_port", default=None, type=int)
    # parser.add_argument("--distributed_world_size", default=None, type=int)
    parser.add_argument("--num_workers", default=4, type=int)


    parser.add_argument("--langs", default="java,python,en_XX", type=str)
    parser.add_argument("--dataset_impl", default='mmap', type=str)
    parser.add_argument("--bpe", default='sentencepiece', type=str)
    parser.add_argument("--sentencepiece_model", default="sentencepiece/sentencepiece.bpe.model", type=str)
    parser.add_argument("--arch", default="mbart_small", type=str)
    parser.add_argument("--user-dir", default="../fairseq_plugins", type=str)
    parser.add_argument("--max_sentences", default=32, type=int)
    parser.add_argument("--max_tokens", default=2048 * 4, type=int)
    parser.add_argument("--update_freq", default=12, type=str)
    

    parser.add_argument("--train_subset", default="train", type=str)
    parser.add_argument("--valid_subset", default="valid", type=str)
    parser.add_argument("--required_batch_size_multiple", default=8, type=int)

    # Optimizer args
    parser.add_argument("--criterion", default="cross_entropy", type=str)
    parser.add_argument("--dropout", default=0.1, type=float)
    parser.add_argument("--attention_dropout", default=0.1, type=float)
    parser.add_argument("--relu_dropout", default=0.0, type=float)
    parser.add_argument("--weight_decay", default=0.01, type=float)
    parser.add_argument("--optimizer", default="adam", type=str)
    parser.add_argument("--adam_eps", default=1e-06, type=float)
    parser.add_argument("--clip_norm", default=0.1, type=float)
    parser.add_argument("--lr", default=3e-4, type=float)
    parser.add_argument("--lr_scheduler", default="polynomial_decay", type=str)
    parser.add_argument("--warmup_updates", default=2000, type=str)
    parser.add_argument("--total_num_update", default=100000, type=str,
                        help='is set to max_update')
    parser.add_argument("--max_update", default=100000, type=str)
    
    parser.add_argument("--save_interval_updates", default=2000, type=int)
    parser.add_argument("--keep_interval_updates", default=30, type=int)
    parser.add_argument("--log_format", default="json", type=str)
    parser.add_argument("--log_interval", default=20, type=int)
    parser.add_argument("--seed", default=1234, type=int)

    parser.add_argument("--restore_file", type=str, default=None, help="do not touch")

    add_taskspecific_args(parser)

    args, _ = parser.parse_known_args()

    args.total_num_update = args.max_update
    return args, _

def get_exp_arguments():
    parser = argparse.ArgumentParser("exp")

    parser.add_argument("--fairseq_path", default="fairseq-train", type=str)
    parser.add_argument("--exp_dir", default="pretrain_v2/saved", type=str,
        help='where to save models')
    parser.add_argument("--data_home", default="data_processed", type=str)
    parser.add_argument("--data_dir", default="binarized_data", type=str)
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--pretrain_dir", type=str, default="saved")
    parser.add_argument("--restore_dir", type=str, default=None,
                        help='example default-mask-random-0.5/exp_20211004-204923/ or None')

    parser.add_argument("--sbatch", action="store_true")
    parser.add_argument("--ndays", type=int, default=3)
    parser.add_argument("--gpus", type=int, default=1)
    parser.add_argument("--cpus", type=int, default=4)
    parser.add_argument("--nodes", type=int, default=1)
    
    exp_args, unknown = parser.parse_known_args()
    if exp_args.debug:
        pass
        exp_args.data_dir = "binarized_data_smaller"

    pretrain_tasks = get_pretrain_info(exp_args)

    home_dir = Path(os.path.abspath(str(Path(os.path.realpath("."), ".."))))

    data_dir = Path(home_dir / exp_args.data_home, exp_args.data_dir)
    assert data_dir.exists(), f"data_dir {data_dir} does not exist"

    data_dir = data_dir / get_pretrain_info(exp_args)
    
    print("data dir", data_dir)
    if get_task() == "translation_from_pretrained_bart":
        # TODO
        data_dir = ":".join(sorted([str(path / "java") for path in data_dir.glob("shard*")]))
    else:
        data_dir = ":".join(sorted([str(path) for path in data_dir.glob("shard*")]))

    # assert data_dir.exists(), f"data_dir {data_dir} does not exist"
    data_dir = str(data_dir)
    logging.info(f"data dir: {data_dir}")

    wandb_name = fairseq_args.wandb_name if hasattr(fairseq_args, "wandb_name") else "TEST"
    save_dir = home_dir / exp_args.exp_dir /  wandb_name / get_name()
    save_dir.mkdir(parents=True)

    logging.info(f"Exp args, {str(exp_args)}")
    logging.info("--"*80)

    exp_args.home_dir = str(home_dir)
    exp_args.data_dir = str(data_dir)
    exp_args.save_dir = str(save_dir)
    return exp_args, unknown

ALL_TASKS = ["mlm", "puzzle", "varm"]
def get_pretrain_info(args):
    task = get_task()
    name = []
    if task == "multilingual_denoising":
        name += ["mlm"]
    elif task == "translation_from_pretrained_bart":
        name += ["varm"]
    else:
        raise ValueError(f"Unknown task {task_parser.task}")
    assert len(name) == 1, "you should choose one pretrain task!"
    return ":".join(name)



if __name__ == "__main__":
    # Experiment arguments

    fairseq_args, not_known = get_fairseq_arguments()
    exp_args, _ = get_exp_arguments()
    check_not_known(not_known, exp_args)

    if hasattr(fairseq_args, "distributed_port") and fairseq_args.distributed_port is None:
        fairseq_args.distributed_port = get_port()
        fairseq_args.distributed_world_size = exp_args.gpus

    if exp_args.debug:
        fairseq_args.wandb_project = "debug"
        fairseq_args.arch = "mbart_small"
    fairseq_args.sentencepiece_model = f"{exp_args.home_dir}/{fairseq_args.sentencepiece_model}"

    fairseq_args.save_dir = exp_args.save_dir

    if exp_args.restore_dir is None:
        # Train from scratch
        fairseq_args.restore_file = None
    else:
        fairseq_args.restore_file = load_checkpoint(fairseq_args, exp_args)
    
    logging.info(f"Fairseq args, {str(fairseq_args)}")
    logging.info("--"*80)

    format_kwargs_fairseq(fairseq_args)
    print(fairseq_train(exp_args, fairseq_args))
