import numpy as np
import random 
import nltk
nltk.download('stopwords')
nltk.download('punkt')
nltk.download('wordnet')
from nltk.corpus import stopwords, wordnet
from collections import defaultdict, OrderedDict
from src.models import create_model
import sys
import re

### AutoDAN ###
def autodan_sample_control(control_suffixs, score_list, num_elites, batch_size, crossover=0.5,
                           num_points=5, mutation=0.01, API_key=None, reference=None, if_softmax=True, if_api=True, device=None):
    score_list = [-x for x in score_list]
    # Step 1: Sort the score_list and get corresponding control_suffixs
    sorted_indices = sorted(range(len(score_list)), key=lambda k: score_list[k], reverse=True)
    sorted_control_suffixs = [control_suffixs[i] for i in sorted_indices]

    # Step 2: Select the elites
    elites = sorted_control_suffixs[:num_elites]

    # Step 3: Use roulette wheel selection for the remaining positions
    parents_list = roulette_wheel_selection(control_suffixs, score_list, batch_size - num_elites, if_softmax)

    # Step 4: Apply crossover and mutation to the selected parents
    offspring = apply_crossover_and_mutation(parents_list, crossover_probability=crossover,
                                                     num_points=num_points,
                                                     mutation_rate=mutation, API_key=API_key, reference=reference,
                                                     if_api=if_api, device=device)

    # Combine elites with the mutated offspring
    next_generation = elites + offspring[:batch_size-num_elites]

    assert len(next_generation) == batch_size
    return next_generation

def roulette_wheel_selection(data_list, score_list, num_selected, if_softmax=True):
    if if_softmax:
        selection_probs = np.exp(score_list - np.max(score_list))
        selection_probs = selection_probs / selection_probs.sum()
    else:
        total_score = sum(score_list)
        selection_probs = [score / total_score for score in score_list]

    selected_indices = np.random.choice(len(data_list), size=num_selected, p=selection_probs, replace=True)

    selected_data = [data_list[i] for i in selected_indices]
    return selected_data

def apply_crossover_and_mutation(selected_data, crossover_probability=0.5, num_points=3, mutation_rate=0.01,
                                 API_key=None,
                                 reference=None, if_api=True, device=None):
    offspring = []

    for i in range(0, len(selected_data), 2):
        parent1 = selected_data[i]
        parent2 = selected_data[i + 1] if (i + 1) < len(selected_data) else selected_data[0]

        if random.random() < crossover_probability:
            child1, child2 = crossover(parent1, parent2, num_points)
            offspring.append(child1)
            offspring.append(child2)
        else:
            offspring.append(parent1)
            offspring.append(parent2)

    mutated_offspring = apply_gpt_mutation(offspring, mutation_rate, API_key, reference, if_api, device)

    return mutated_offspring

def crossover(str1, str2, num_points):
    # Function to split text into paragraphs and then into sentences
    def split_into_paragraphs_and_sentences(text):
        paragraphs = text.split('\n\n')
        return [re.split('(?<=[,.!?])\s+', paragraph) for paragraph in paragraphs]

    paragraphs1 = split_into_paragraphs_and_sentences(str1)
    paragraphs2 = split_into_paragraphs_and_sentences(str2)

    new_paragraphs1, new_paragraphs2 = [], []

    for para1, para2 in zip(paragraphs1, paragraphs2):
        max_swaps = min(len(para1), len(para2)) - 1
        num_swaps = min(num_points, max_swaps)

        swap_indices = sorted(random.sample(range(1, max_swaps + 1), num_swaps))

        new_para1, new_para2 = [], []
        last_swap = 0
        for swap in swap_indices:
            if random.choice([True, False]):
                new_para1.extend(para1[last_swap:swap])
                new_para2.extend(para2[last_swap:swap])
            else:
                new_para1.extend(para2[last_swap:swap])
                new_para2.extend(para1[last_swap:swap])
            last_swap = swap

        if random.choice([True, False]):
            new_para1.extend(para1[last_swap:])
            new_para2.extend(para2[last_swap:])
        else:
            new_para1.extend(para2[last_swap:])
            new_para2.extend(para1[last_swap:])

        new_paragraphs1.append(' '.join(new_para1))
        new_paragraphs2.append(' '.join(new_para2))

    return '\n\n'.join(new_paragraphs1), '\n\n'.join(new_paragraphs2)

def apply_gpt_mutation(offspring, mutation_rate=0.01, API_key=None, reference=None, if_api=True, device=None):
    for i in range(len(offspring)):
        if random.random() < mutation_rate:
            offspring[i] = replace_with_synonyms(offspring[i])    
    return offspring

def gpt_mutate(sentence, API_key=None, device=None):
    azureGPT = create_model("model_configs/azuregpt_config.json", device)
    revised_sentence = azureGPT.mutate_query(sentence)

    if len(revised_sentence) < len(sentence)/2:
        revised_sentence = (revised_sentence + ' ')*2
    
    return revised_sentence

def replace_with_synonyms(sentence, num=10):
    T = {"llama2", "meta", "vicuna", "lmsys", "guanaco", "theblokeai", "wizardlm", "mpt-chat",
         "mosaicml", "mpt-instruct", "falcon", "tii", "chatgpt", "modelkeeper", "prompt"}
    stop_words = set(stopwords.words('english'))
    words = nltk.word_tokenize(sentence)
    uncommon_words = [word for word in words if word.lower() not in stop_words and word.lower() not in T]
    selected_words = random.sample(uncommon_words, min(num, len(uncommon_words)))
    for word in selected_words:
        synonyms = wordnet.synsets(word)
        if synonyms and synonyms[0].lemmas():
            synonym = synonyms[0].lemmas()[0].name()
            sentence = sentence.replace(word, synonym, 1)
    return sentence

### HGA ###
def autodan_sample_control_hga(word_dict, control_suffixs, score_list, num_elites, batch_size, crossover=0.5,
                               mutation=0.01, API_key=None, reference=None, if_api=True, device=None):
    score_list = [-x for x in score_list]
    # Step 1: Sort the score_list and get corresponding control_suffixs
    sorted_indices = sorted(range(len(score_list)), key=lambda k: score_list[k], reverse=True)
    sorted_control_suffixs = [control_suffixs[i] for i in sorted_indices]

    # Step 2: Select the elites
    elites = sorted_control_suffixs[:num_elites]
    parents_list = sorted_control_suffixs[num_elites:]

    # Step 3: Construct word list
    word_dict = construct_momentum_word_dict(word_dict, control_suffixs, score_list)
    print(f"Length of current word dictionary: {len(word_dict)}")

    # check the length of parents
    parents_list = [x for x in parents_list if len(x) > 0]
    if len(parents_list) < batch_size - num_elites:
        print("Not enough parents, using reference instead.")
        parents_list += random.choices(reference[batch_size:], k = batch_size - num_elites - len(parents_list))
        
    # Step 4: Apply word replacement with roulette wheel selection
    offspring = apply_word_replacement(word_dict, parents_list, crossover)
    offspring = apply_gpt_mutation(offspring, mutation, API_key, reference, if_api, device)

    # Combine elites with the mutated offspring
    next_generation = elites + offspring[:batch_size-num_elites]

    assert len(next_generation) == batch_size
    return next_generation, word_dict

def construct_momentum_word_dict(word_dict, control_suffixs, score_list, topk=-1):
    T = {"llama2", "meta", "vicuna", "lmsys", "guanaco", "theblokeai", "wizardlm", "mpt-chat",
         "mosaicml", "mpt-instruct", "falcon", "tii", "chatgpt", "modelkeeper", "prompt"}
    stop_words = set(stopwords.words('english'))
    if len(control_suffixs) != len(score_list):
        raise ValueError("control_suffixs and score_list must have the same length.")

    word_scores = defaultdict(list)

    for prefix, score in zip(control_suffixs, score_list):
        words = set(
            [word for word in nltk.word_tokenize(prefix) if word.lower() not in stop_words and word.lower() not in T])
        for word in words:
            word_scores[word].append(score)

    for word, scores in word_scores.items():
        avg_score = sum(scores) / len(scores)
        if word in word_dict:
            word_dict[word] = (word_dict[word] + avg_score) / 2
        else:
            word_dict[word] = avg_score

    sorted_word_dict = OrderedDict(sorted(word_dict.items(), key=lambda x: x[1], reverse=True))
    if topk == -1:
        topk_word_dict = dict(list(sorted_word_dict.items()))
    else:
        topk_word_dict = dict(list(sorted_word_dict.items())[:topk])
    return topk_word_dict

def apply_word_replacement(word_dict, parents_list, crossover=0.5):
    return [replace_with_best_synonym(sentence, word_dict, crossover) for sentence in parents_list]

def replace_with_best_synonym(sentence, word_dict, crossover_probability):
    stop_words = set(stopwords.words('english'))
    T = {"llama2", "meta", "vicuna", "lmsys", "guanaco", "theblokeai", "wizardlm", "mpt-chat",
         "mosaicml", "mpt-instruct", "falcon", "tii", "chatgpt", "modelkeeper", "prompt"}
    paragraphs = sentence.split('\n\n')
    modified_paragraphs = []
    min_value = min(word_dict.values())

    for paragraph in paragraphs:
        words = replace_quotes(nltk.word_tokenize(paragraph))
        count = 0
        for i, word in enumerate(words):
            if random.random() < crossover_probability:
                if word.lower() not in stop_words and word.lower() not in T:
                    synonyms = get_synonyms(word.lower())
                    word_scores = {syn: word_dict.get(syn, min_value) for syn in synonyms}
                    best_synonym = word_roulette_wheel_selection(word, word_scores)
                    if best_synonym:
                        words[i] = best_synonym
                        count += 1
                        if count >= 5:
                            break
            else:
                if word.lower() not in stop_words and word.lower() not in T:
                    synonyms = get_synonyms(word.lower())
                    word_scores = {syn: word_dict.get(syn, 0) for syn in synonyms}
                    best_synonym = word_roulette_wheel_selection(word, word_scores)
                    if best_synonym:
                        words[i] = best_synonym
                        count += 1
                        if count >= 5:
                            break
        modified_paragraphs.append(join_words_with_punctuation(words))
    return '\n\n'.join(modified_paragraphs)

def replace_with_synonyms(sentence, num=10):
    T = {"llama2", "meta", "vicuna", "lmsys", "guanaco", "theblokeai", "wizardlm", "mpt-chat",
         "mosaicml", "mpt-instruct", "falcon", "tii", "chatgpt", "modelkeeper", "prompt"}
    stop_words = set(stopwords.words('english'))
    words = nltk.word_tokenize(sentence)
    uncommon_words = [word for word in words if word.lower() not in stop_words and word.lower() not in T]
    selected_words = random.sample(uncommon_words, min(num, len(uncommon_words)))
    for word in selected_words:
        synonyms = wordnet.synsets(word)
        if synonyms and synonyms[0].lemmas():
            synonym = synonyms[0].lemmas()[0].name()
            sentence = sentence.replace(word, synonym, 1)
    # print(f'revised: {sentence}')
    return sentence

def replace_quotes(words):
    new_words = []
    quote_flag = True

    for word in words:
        if word in ["``", "''"]:
            if quote_flag:
                new_words.append('“')
                quote_flag = False
            else:
                new_words.append('”')
                quote_flag = True
        else:
            new_words.append(word)
    return new_words

def get_synonyms(word):
    synonyms = set()
    for syn in wordnet.synsets(word):
        for lemma in syn.lemmas():
            synonyms.add(lemma.name())
    return list(synonyms)

def word_roulette_wheel_selection(word, word_scores):
    if not word_scores:
        return word
    min_score = min(word_scores.values())
    adjusted_scores = {k: v - min_score for k, v in word_scores.items()}
    total_score = sum(adjusted_scores.values())
    pick = random.uniform(0, total_score)
    current_score = 0
    for synonym, score in adjusted_scores.items():
        current_score += score
        if current_score > pick:
            if word.istitle():
                return synonym.title()
            else:
                return synonym

def join_words_with_punctuation(words):
    sentence = words[0]
    previous_word = words[0]
    flag = 1
    for word in words[1:]:
        if word in [",", ".", "!", "?", ":", ";", ")", "]", "}", '”']:
            sentence += word
        else:
            if previous_word in ["[", "(", "'", '"', '“']:
                if previous_word in ["'", '"'] and flag == 1:
                    sentence += " " + word
                else:
                    sentence += word
            else:
                if word in ["'", '"'] and flag == 1:
                    flag = 1 - flag
                    sentence += " " + word
                elif word in ["'", '"'] and flag == 0:
                    flag = 1 - flag
                    sentence += word
                else:
                    if "'" in word and re.search('[a-zA-Z]', word):
                        sentence += word
                    else:
                        sentence += " " + word
        previous_word = word
    return sentence