import argparse
import importlib
import numpy as np
import os
import random
import sys
import textattack
import time
import torch
import pickle
import copy

from .attack_args import *

def set_seed(random_seed):
    random.seed(random_seed)
    np.random.seed(random_seed)
    torch.manual_seed(random_seed)

def get_args():
    # Parser for regular arguments
    parser = argparse.ArgumentParser(
        description='A commandline parser for TextAttack', 
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    transformation_names = set(BLACK_BOX_TRANSFORMATION_CLASS_NAMES.keys()) | set(WHITE_BOX_TRANSFORMATION_CLASS_NAMES.keys())
    parser.add_argument('--transformation', type=str, required=False,
        default='word-swap-embedding', choices=transformation_names,
        help='The transformation to apply. Usage: "--transformation {transformation}:{arg_1}={value_1},{arg_3}={value_3}. Choices: ' + str(transformation_names))
    
    model_group = parser.add_mutually_exclusive_group()
    
    model_names = list(TEXTATTACK_MODEL_CLASS_NAMES.keys()) + list(HUGGINGFACE_DATASET_BY_MODEL.keys()) + list(TEXTATTACK_DATASET_BY_MODEL.keys())
    model_group.add_argument('--model', type=str, required=False, default='bert-base-uncased-yelp-sentiment',
        choices=model_names, help='The pre-trained model to attack.')
    model_group.add_argument('--model-from-file', type=str, required=False,
        help='File of model and tokenizer to import.')     
    model_group.add_argument('--model-from-huggingface', type=str, required=False,
        help='huggingface.co ID of pre-trained model to load')
        
    dataset_group = parser.add_mutually_exclusive_group()
    dataset_group.add_argument('--dataset-from-nlp', type=str, required=False, default=None,
        help='Dataset to load from `nlp` repository.')
    dataset_group.add_argument('--dataset-from-file', type=str, required=False, default=None,
        help='Dataset to load from a file.')
    
    parser.add_argument('--constraints', type=str, required=False, nargs='*',
        default=['repeat', 'stopword'],
        help='Constraints to add to the attack. Usage: "--constraints {constraint}:{arg_1}={value_1},{arg_3}={value_3}". Choices: ' + str(CONSTRAINT_CLASS_NAMES.keys()))
    
    parser.add_argument('--out-dir', type=str, required=False, default=None,
        help='A directory to output results to.')
    
    parser.add_argument('--enable-visdom', action='store_true',
        help='Enable logging to visdom.')
    
    parser.add_argument('--enable-wandb', action='store_true',
        help='Enable logging to Weights & Biases.')
    
    parser.add_argument('--disable-stdout', action='store_true',
        help='Disable logging to stdout')
   
    parser.add_argument('--enable-csv', nargs='?', default=None, const='fancy', type=str,
        help='Enable logging to csv. Use --enable-csv plain to remove [[]] around words.')

    parser.add_argument('--num-examples', '-n', type=int, required=False, 
        default='5', help='The number of examples to process.')
    
    parser.add_argument('--num-examples-offset', '-o', type=int, required=False, 
        default=0, help='The offset to start at in the dataset.')

    parser.add_argument('--shuffle', action='store_true', required=False, 
        default=False, help='Randomly shuffle the data before attacking')
    
    parser.add_argument('--interactive', action='store_true', default=False,
        help='Whether to run attacks interactively.')
    
    parser.add_argument('--attack-n', action='store_true', default=False,
        help='Whether to run attack until `n` examples have been attacked (not skipped).')
    
    parser.add_argument('--parallel', action='store_true', default=False,
        help='Run attack using multiple GPUs.')

    goal_function_choices = ', '.join(GOAL_FUNCTION_CLASS_NAMES.keys())
    parser.add_argument('--goal-function', '-g', default='untargeted-classification',
        help=f'The goal function to use. choices: {goal_function_choices}')
   
    def str_to_int(s): return sum((ord(c) for c in s))
    parser.add_argument('--random-seed', default=str_to_int('TEXTATTACK'))

    parser.add_argument('--checkpoint-dir', required=False, type=str, default=default_checkpoint_dir(),
        help='The directory to save checkpoint files.')

    parser.add_argument('--checkpoint-interval', required=False, type=int, 
        help='If set, checkpoint will be saved after attacking every N examples. If not set, no checkpoints will be saved.')

    parser.add_argument('--query-budget', '-q', type=int, default=float('inf'),
        help='The maximum number of model queries allowed per example attacked.')

    attack_group = parser.add_mutually_exclusive_group(required=False)
    search_choices = ', '.join(SEARCH_CLASS_NAMES.keys())
    attack_group.add_argument('--search', '--search-method', '-s', type=str, 
        required=False, default='greedy-word-wir', 
        help=f'The search method to use. choices: {search_choices}')
    attack_group.add_argument('--recipe', '--attack-recipe', '-r', type=str, required=False, default=None,
        help='full attack recipe (overrides provided goal function, transformation & constraints)',
        choices=RECIPE_NAMES.keys())
    attack_group.add_argument('--attack-from-file', type=str, required=False, default=None,
        help='attack to load from file (overrides provided goal function, transformation & constraints)',
        )

    # Parser for parsing args for resume
    resume_parser = argparse.ArgumentParser(
        description='A commandline parser for TextAttack', 
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    resume_parser.add_argument('--checkpoint-file', '-f', type=str, required=True, 
        help='Path of checkpoint file to resume attack from. If "latest" (or "{directory path}/latest") is entered,'\
        'recover latest checkpoint from either current path or specified directory.')

    resume_parser.add_argument('--checkpoint-dir', '-d', required=False, type=str, default=None,
        help='The directory to save checkpoint files. If not set, use directory from recovered arguments.')

    resume_parser.add_argument('--checkpoint-interval', '-i', required=False, type=int, 
        help='If set, checkpoint will be saved after attacking every N examples. If not set, no checkpoints will be saved.')

    resume_parser.add_argument('--parallel', action='store_true', default=False,
        help='Run attack using multiple GPUs.')
    
    # Resume attack from checkpoint.
    if sys.argv[1:] and sys.argv[1].lower() == 'resume':
        args = resume_parser.parse_args(sys.argv[2:])
        setattr(args, 'checkpoint_resume', True)
    else:
        command_line_args = None if sys.argv[1:] else ['-h'] # Default to help with empty arguments.
        args = parser.parse_args(command_line_args)
        setattr(args, 'checkpoint_resume', False)

        if args.checkpoint_interval and args.shuffle:
            # Not allowed b/c we cannot recover order of shuffled data
            raise ValueError('Cannot use `--checkpoint-interval` with `--shuffle=True`')
        
        set_seed(args.random_seed)
    
    # Shortcuts for huggingface models using --model.
    if args.model in HUGGINGFACE_DATASET_BY_MODEL:
        args.model_from_huggingface, args.dataset_from_nlp = HUGGINGFACE_DATASET_BY_MODEL[args.model]
        args.model = None
    
    return args

def parse_transformation_from_args(args, model):
    # Transformations
    transformation_name = args.transformation
    if ':' in transformation_name:
        transformation_name, params = transformation_name.split(':')
        
        if transformation_name in WHITE_BOX_TRANSFORMATION_CLASS_NAMES:
            transformation = eval(f'{WHITE_BOX_TRANSFORMATION_CLASS_NAMES[transformation_name]}(model, {params})')
        elif transformation_name in BLACK_BOX_TRANSFORMATION_CLASS_NAMES:
            transformation = eval(f'{BLACK_BOX_TRANSFORMATION_CLASS_NAMES[transformation_name]}({params})')
        else:
            raise ValueError(f'Error: unsupported transformation {transformation_name}')
    else:
        if transformation_name in WHITE_BOX_TRANSFORMATION_CLASS_NAMES:
            transformation = eval(f'{WHITE_BOX_TRANSFORMATION_CLASS_NAMES[transformation_name]}(model)')
        elif transformation_name in BLACK_BOX_TRANSFORMATION_CLASS_NAMES:
            transformation = eval(f'{BLACK_BOX_TRANSFORMATION_CLASS_NAMES[transformation_name]}()')
        else:
            raise ValueError(f'Error: unsupported transformation {transformation_name}')
    return transformation

def parse_goal_function_from_args(args, model):
    # Goal Functions
    goal_function = args.goal_function
    if ':' in goal_function:
        goal_function_name, params = goal_function.split(':')
        if goal_function_name not in GOAL_FUNCTION_CLASS_NAMES:
            raise ValueError(f'Error: unsupported goal_function {goal_function_name}')
        goal_function = eval(f'{GOAL_FUNCTION_CLASS_NAMES[goal_function_name]}(model, {params})')
    elif goal_function in GOAL_FUNCTION_CLASS_NAMES:
        goal_function = eval(f'{GOAL_FUNCTION_CLASS_NAMES[goal_function]}(model)')
    else:
        raise ValueError(f'Error: unsupported goal_function {goal_function}')
    goal_function.query_budget = args.query_budget
    return goal_function

def parse_constraints_from_args(args):
    # Constraints
    if not args.constraints: 
        return []
    
    _constraints = []
    for constraint in args.constraints:
        if ':' in constraint:
            constraint_name, params = constraint.split(':')
            if constraint_name not in CONSTRAINT_CLASS_NAMES:
                raise ValueError(f'Error: unsupported constraint {constraint_name}')
            _constraints.append(eval(f'{CONSTRAINT_CLASS_NAMES[constraint_name]}({params})'))
        elif constraint in CONSTRAINT_CLASS_NAMES:
            _constraints.append(eval(f'{CONSTRAINT_CLASS_NAMES[constraint]}()'))
        else:
            raise ValueError(f'Error: unsupported constraint {constraint}')
    
    return _constraints

def parse_attack_from_args(args):
    model = parse_model_from_args(args)
    if args.recipe:
        if ':' in args.recipe:
            recipe_name, params = args.recipe.split(':')
            if recipe_name not in RECIPE_NAMES:
                raise ValueError(f'Error: unsupported recipe {recipe_name}')
            recipe = eval(f'{RECIPE_NAMES[recipe_name]}(model, {params})')
        elif args.recipe in RECIPE_NAMES:
            recipe = eval(f'{RECIPE_NAMES[args.recipe]}(model)')
        else:
            raise ValueError(f'Invalid recipe {args.recipe}')
        recipe.goal_function.query_budget = args.query_budget
        return recipe
    elif args.attack_from_file:
        if ':' in args.attack_from_file:
            attack_file, attack_name = args.attack_from_file.split(':')
        else:
            attack_file, attack_name = args.attack_from_file, 'attack'
        attack_file = attack_file.replace('.py', '').replace('/', '.')
        attack_module = importlib.import_module(attack_file)
        attack_func = getattr(attack_module, attack_name)
        return attack_func(model)
    else:
        goal_function = parse_goal_function_from_args(args, model)
        transformation = parse_transformation_from_args(args, model)
        constraints = parse_constraints_from_args(args)
        if ':' in args.search:
            search_name, params = args.search.split(':')
            if search_name not in SEARCH_CLASS_NAMES:
                raise ValueError(f'Error: unsupported search {search_name}')
            search_method = eval(f'{SEARCH_CLASS_NAMES[search_name]}({params})')
        elif args.search in SEARCH_CLASS_NAMES:
            search_method = eval(f'{SEARCH_CLASS_NAMES[args.search]}()')
        else:
            raise ValueError(f'Error: unsupported attack {args.search}')
    return textattack.shared.Attack(goal_function, constraints, transformation, search_method)

def parse_model_from_args(args):
    if args.model_from_file:
        colored_model_name = textattack.shared.utils.color_text(args.model_from_file, color='blue', method='ansi')
        textattack.shared.logger.info(f'Loading model and tokenizer from file: {colored_model_name}')
        if ':' in args.model_from_file:
            model_file, model_name, tokenizer_name = args.model_from_file.split(':')
        else:
            model_file, model_name, tokenizer_name = args.model_from_file, 'model', 'tokenizer'
        try:
            model_file = args.model_from_file.replace('.py', '').replace('/', '.')
            model_module = importlib.import_module(model_file)
        except:
            raise ValueError(f'Failed to import model or tokenizer from file {args.model_from_file}')
        try:
            model = getattr(model_module, model_name)
        except AttributeError:
            raise AttributeError(f'``{model_name}`` not found in module {args.model_from_file}')
        try:
            tokenizer = getattr(model_module, tokenizer_name)
        except AttributeError:
            raise AttributeError(f'``{tokenizer_name}`` not found in module {args.model_from_file}')
        model = model.to(textattack.shared.utils.device)
        setattr(model, 'tokenizer', tokenizer)
    elif args.model_from_huggingface:
        import transformers
        if ':' in args.model_from_huggingface:
            model_class, model_name = args.model_from_huggingface.split(':')
            model_class = eval(f'transformers.{model_class}')
        else:
            model_class, model_name = transformers.AutoModelForSequenceClassification, args.model_from_huggingface
        colored_model_name = textattack.shared.utils.color_text(model_name, color='blue', method='ansi')
        textattack.shared.logger.info(f'Loading pre-trained model from HuggingFace model repository: {colored_model_name}')
        model = model_class.from_pretrained(model_name)
        model = model.to(textattack.shared.utils.device)
        try:
            tokenizer = textattack.tokenizers.AutoTokenizer(args.model_from_huggingface)
        except OSError:
            textattack.shared.logger.warn(f'AutoTokenizer {args.model_from_huggingface} not found. Defaulting to `bert-base-uncased`')
            tokenizer = textattack.tokenizers.AutoTokenizer('bert-base-uncased')
        setattr(model, 'tokenizer', tokenizer)
    else:
        if ':' in args.model:
            model_name, params = args.model.split(':')
            colored_model_name = textattack.shared.utils.color_text(model_name, color='blue', method='ansi')
            textattack.shared.logger.info(f'Loading pre-trained TextAttack model: {colored_model_name}')
            if model_name not in TEXTATTACK_MODEL_CLASS_NAMES:
                raise ValueError(f'Error: unsupported model {model_name}')
            model = eval(f'{TEXTATTACK_MODEL_CLASS_NAMES[model_name]}({params})')
        elif args.model in TEXTATTACK_MODEL_CLASS_NAMES:
            colored_model_name = textattack.shared.utils.color_text(args.model, color='blue', method='ansi')
            textattack.shared.logger.info(f'Loading pre-trained TextAttack model: {colored_model_name}')
            model = eval(f'{TEXTATTACK_MODEL_CLASS_NAMES[args.model]}()')
        elif args.model in TEXTATTACK_DATASET_BY_MODEL:
            colored_model_name = textattack.shared.utils.color_text(args.model, color='blue', method='ansi')
            model_path, args.dataset_from_nlp = TEXTATTACK_DATASET_BY_MODEL[args.model]
            if args.model.startswith('lstm'):
                textattack.shared.logger.info(f'Loading pre-trained TextAttack LSTM: {colored_model_name}')
                model = textattack.models.helpers.LSTMForClassification(model_path=model_path)
            elif args.model.startswith('cnn'):
                textattack.shared.logger.info(f'Loading pre-trained TextAttack CNN: {colored_model_name}')
                model = textattack.models.helpers.WordCNNForClassification(model_path=model_path)
            else:
                raise ValueError(f'Unknown TextAttack pretrained model {args.model}')
        else: 
            raise ValueError(f'Error: unsupported model {args.model}')
    return model

def parse_dataset_from_args(args):
    args.dataset_from_nlp = ('glue', 'sst2', 'validation')
    if args.dataset_from_file:
        textattack.shared.logger.info(f'Loading model and tokenizer from file: {args.model_from_file}')
        if ':' in args.dataset_from_file:
            dataset_file, dataset_name = args.dataset_from_file.split(':')
        else:
            dataset_file, dataset_name = args.dataset_from_file, 'dataset'
        try:
            dataset_file = dataset_file.replace('.py', '').replace('/', '.')
            dataset_module = importlib.import_module(dataset_file)
        except:
            raise ValueError(f'Failed to import dataset from file {args.dataset_from_file}')
        try:
            dataset = getattr(dataset_module, dataset_name)
        except AttributeError:
            raise AttributeError(f'``dataset`` not found in module {args.dataset_from_file}')
    elif args.dataset_from_nlp:
        dataset_args = args.dataset_from_nlp
        if ':' in dataset_args:
            dataset_args = dataset_args.split(':')
        dataset = textattack.datasets.HuggingFaceNLPDataset(*dataset_args, shuffle=args.shuffle)
    else:
        if not args.model:
            raise ValueError('Must supply pretrained model or dataset')
        elif args.model in DATASET_BY_MODEL:
            dataset = DATASET_BY_MODEL[args.model](offset=args.num_examples_offset)
        else:
            raise ValueError(f'Error: unsupported model {args.model}')
    return dataset

def parse_logger_from_args(args):
    # Create logger
    attack_log_manager = textattack.loggers.AttackLogManager()
    # Set default output directory to `textattack/outputs`.
    if not args.out_dir:
        current_dir = os.path.dirname(os.path.realpath(__file__))
        outputs_dir = os.path.join(current_dir, os.pardir, os.pardir, os.pardir, 'outputs')
        args.out_dir = os.path.normpath(outputs_dir)
        
    # Output file.
    out_time = int(time.time()*1000) # Output file
    outfile_name = 'attack-{}.txt'.format(out_time)
    attack_log_manager.add_output_file(os.path.join(args.out_dir, outfile_name))
        
    # CSV
    if args.enable_csv:
        outfile_name = 'attack-{}.csv'.format(out_time)
        color_method = None if args.enable_csv == 'plain' else 'file'
        csv_path = os.path.join(args.out_dir, outfile_name)
        attack_log_manager.add_output_csv(csv_path, color_method)
        print('Logging to CSV at path {}.'.format(csv_path))

    # Visdom
    if args.enable_visdom:
        attack_log_manager.enable_visdom()
        
    # Weights & Biases
    if args.enable_wandb:
        attack_log_manager.enable_wandb()

    # Stdout
    if not args.disable_stdout:
        attack_log_manager.enable_stdout()
    return attack_log_manager

def parse_checkpoint_from_args(args):
    file_name = os.path.basename(args.checkpoint_file)   
    if file_name.lower() == 'latest':
        dir_path = os.path.dirname(args.checkpoint_file)
        chkpt_file_names = [f for f in os.listdir(dir_path) if f.endswith('.ta.chkpt')]
        assert chkpt_file_names, "Checkpoint directory is empty"
        timestamps = [int(f.replace('.ta.chkpt', '')) for f in chkpt_file_names]
        latest_file = str(max(timestamps)) + '.ta.chkpt'
        checkpoint_path = os.path.join(dir_path, latest_file)
    else:
        checkpoint_path = args.checkpoint_file
    
    checkpoint = textattack.shared.Checkpoint.load(checkpoint_path)
    set_seed(checkpoint.args.random_seed)

    return checkpoint

def default_checkpoint_dir():
    current_dir = os.path.dirname(os.path.realpath(__file__))
    checkpoints_dir = os.path.join(current_dir, os.pardir, os.pardir, os.pardir, 'checkpoints')
    return os.path.normpath(checkpoints_dir)

def merge_checkpoint_args(saved_args, cmdline_args):
    """ Merge previously saved arguments for checkpoint and newly entered arguments """
    args = copy.deepcopy(saved_args)
    # Newly entered arguments take precedence
    args.checkpoint_resume = cmdline_args.checkpoint_resume
    args.parallel =  cmdline_args.parallel
    # If set, replace
    if cmdline_args.checkpoint_dir:
        args.checkpoint_dir = cmdline_args.checkpoint_dir
    if cmdline_args.checkpoint_interval:
        args.checkpoint_interval = cmdline_args.checkpoint_interval
    
    return args