import os
import json
import argparse
import logging
from pathlib import Path

import torch


logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
log = logging.getLogger(__name__)


def get_args():
    parser = argparse.ArgumentParser(description='ESPER')


    parser.add_argument(
        '--init-model', type=str, default='gpt2', help='language model used for policy.')
    parser.add_argument(
        '--label_path', type=str, default='./data/esper_demo/labels_all.json', help='style label info file path')
    parser.add_argument(
        '--checkpoint', type=str, default='./data/esper_demo/ckpt/gpt2_style', help='checkpoint file path')

    parser.add_argument(
        '--prefix_length', type=int, default=10, help='prefix length for the visual mapper')
    parser.add_argument(
        '--clipcap_num_layers', type=int, default=1, help='num_layers for the visual mapper')
    parser.add_argument(
        '--use_transformer_mapper', action='store_true', default=False, help='use transformer mapper instead of mlp')
    parser.add_argument(
        '--use_label_prefix', action='store_true', default=False, help='label as prefixes')
    parser.add_argument(
        '--clip_model_type', type=str, default='ViT-B/32', help='clip backbone type')

    parser.add_argument(
        '--infer_no_repeat_size', type=int, default=2, help="no repeat ngram size for inference")
    parser.add_argument(
        '--response-length', type=int, default=20, help='number of tokens to generate for each prompt.')
    parser.add_argument(
        '--num-gpus', type=int, default=None, help='number of gpus. use all available if none')
    parser.add_argument(
        '--port', type=int, default=None, help="port for the demo server")

    args = parser.parse_args()
    args.cuda = torch.cuda.is_available()

    if args.use_label_prefix:
        log.info(f'using label prefix')
    num_gpus = torch.cuda.device_count()
    if args.num_gpus is None:
        args.num_gpus = num_gpus
    else:
        args.num_gpus = min(num_gpus, args.num_gpus)

    if args.checkpoint is not None:
        args.checkpoint = str(Path(args.checkpoint).resolve())
    return args