import argparse
import random
import torch

from engine import Engine
from dataloader import autodetect_device_type, build_model


RED = "\033[0;31m"
GREEN = "\033[0;32m"
YELLOW = "\033[0;33m"
RESET = "\033[0;37m"

CATEGORIES = ["general", "artist", "character", "copyright"]

RATING_TAGS = ["general", "sensitive", "questionable", "explicit"]
SCORE_TAGS = ["score_0", "score_1", "score_2", "score_3"]
LENGTH_TAGS = ["len_0", "len_1", "len_2", "len_3"]
YEAR_TAGS = ["year_16", "year_21", "year_23", "year_26"]


parser = argparse.ArgumentParser()
parser.add_argument('-p', '--positive-tags', type=str, default=None)
parser.add_argument('-n', '--negative-tags', type=str, default=None)
parser.add_argument('-c', '--banned-categories', type=str, default=None)
parser.add_argument('-t', '--temperature',   type=float, default=1)
parser.add_argument('-k', '--top-k',         type=int, default=None)
parser.add_argument('-d', '--device',        type=str, default='', choices=['cuda', 'cpu', 'mps'], help='Device type for evaluation: cuda|cpu|mps. empty => autodetect')
args = parser.parse_args()


device = autodetect_device_type() if args.device == "" else args.device
model, tokenizer = build_model(device)

engine = Engine(model, tokenizer)
rng = torch.Generator(device=device)


def prepare_tags(tags: str):
    tags = tags.lower().strip().replace('\\', '')
    
    comma_separated = ',' in tags
    space_separated = ' ' in tags and not comma_separated
    underscores = '_' in tags

    if comma_separated:
        tags_ = [tag.strip() for tag in tags.split(',')]
            
    elif space_separated:
        tags_ = tags.split()
        
    if not comma_separated and not space_separated:
        comma_separated = True
        tags_ = [tags]
        
    if not underscores:
        tags_ = [tag.replace(' ', '_') for tag in tags_]
    
    return comma_separated, space_separated, tags_


while True:
    print('\n' + GREEN + '='*64 + RESET, end='')
    
    if args.positive_tags is not None:
        positive_tags = args.positive_tags
        negative_tags = args.negative_tags
        banned_categories = args.banned_categories
    else:
        # Get the prompt interactively from the console
        try:
            print(GREEN + '\nPositive tags:' + RESET)
            positive_tags = input()
            print(RED + '\nNegative tags:' + RESET)
            negative_tags = input()
            print(RED + '\nBanned categories:' + RESET)
            banned_categories = input()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break
        
    if not positive_tags:
        r_rating, r_score, r_length, r_year = random.choice(RATING_TAGS), random.choice(SCORE_TAGS), random.choice(LENGTH_TAGS), random.choice(YEAR_TAGS)
        positive_tags = ', '.join([r_year, r_rating, r_score, r_length])
        print(GREEN + '\nPositive tags:' + RESET)
        print(positive_tags)
    
    comma_separated, space_separated, positive_tags = prepare_tags(positive_tags)
    
    if negative_tags:
        _, _, negative_tags = prepare_tags(negative_tags)
        
    if banned_categories:
        _, _, banned_categories = prepare_tags(banned_categories)
        
        banned_categories = set(banned_categories)
        
        difference = banned_categories - set(CATEGORIES)
        
        if difference:
            print(RED + '\nValueError: ' + f'"{" ".join(difference)}" Category doesn\'t exist' + RESET)
            continue
    
    
    ntags_by_category = []
    if banned_categories:
        for data in tokenizer.tags:
            if data[2] in banned_categories:
                ntags_by_category.append(data[0])

    if ntags_by_category:
        if isinstance(negative_tags, str):
            negative_tags = []
            
        negative_tags.extend(ntags_by_category)
        
    
    try:
        positive_tokens = tokenizer.encode(positive_tags)
        negative_tokens = tokenizer.encode(negative_tags)
    except ValueError as e:
        print(RED + '\nValueError: ' + str(e) + RESET)
        continue

    generate_kwargs = {
        "negative_tokens": [ntoken[0] for ntoken in negative_tokens],
        "num_samples": 1,
        "max_tokens": 100,
        "temperature": args.temperature,
        "top_k": args.top_k,
        "seed": rng.seed()
    }
    
    result_tags = []
    for token_column, _ in engine.generate([ptoken[0] for ptoken in positive_tokens], **generate_kwargs):
        token = token_column[0]
        token = tokenizer.decode([token])[0]
        tag = token[0]
        
        if tag == 'EOS':
            break
        
        result_tags.append(token)
        
    print('\nResult:')
        
    for tag_data in result_tags:
        tag, tag_category, tag_count = tag_data
        
        if tag_count < 63406 and tag_count >= 4900:
            print(YELLOW, end='')
        elif tag_count <= 4900:
            print(RED, end='')
        elif tag_count >= 63406:
            print(GREEN, end='')

        if space_separated:
            print(tag, end=" ")
            
        elif comma_separated:
            print(tag.replace('_', ' ')
                     .replace('(', r'\(')
                     .replace(')', r'\)'),
            end=RESET + ", ")
            
    print(RESET)
    
    all_tags = positive_tags + result_tags
    duplicates = [i for i in set(all_tags) if all_tags.count(i) > 1]
    
    if duplicates:
        print(RED + "\nDuplicates: " + f'"{" ".join(d for d in duplicates)}"' + RESET)
    
    if args.positive_tags is not None:
        break