from sentence_transformers import SentenceTransformer
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.manifold import TSNE
import csv
from tqdm import tqdm
from pathlib import Path
import json

def main_cluster_programs(args):
    cluster_model = args.cluster_model
    output_dir = Path(args.output_dir)
    input_file = args.input_file
    program_name_column = args.program_name_column
    asr_column = args.asr_column # column containing the ASR estimate

    output_dir.mkdir(parents=True, exist_ok=True)
    programs = []
    asrs = []

    with open(input_file, 'r') as f:
        input_reader = csv.DictReader(f)
        for row in tqdm(list(input_reader)):
            program = row[program_name_column]
            asr = row[asr_column]
            programs.append(program)
            asrs.append(float(asr))

    embedder = SentenceTransformer(cluster_model)
    embeddings = embedder.encode(programs, convert_to_numpy=True, batch_size=1024, show_progress_bar=True)

    tsne = TSNE(n_components=2, learning_rate='auto',
                  init='random', perplexity=5)
    tsne_results = tsne.fit_transform(np.array(embeddings))

    df = pd.DataFrame()
    df['tsne-2d-one'] = tsne_results[:,0]
    df['tsne-2d-two'] = tsne_results[:,1]
    df['asr'] = asrs

    plt.figure(figsize=(16,10))
    scatter_plot = sns.scatterplot(
        x="tsne-2d-one", y="tsne-2d-two", hue='asr',
        palette=sns.color_palette("hls"),
        data=df,
        legend="auto",
        alpha=0.3
    ).set_title('Clustering of CodeBERT embeddings of Jailbreak Attacks')
    fig = scatter_plot.get_figure()
    fig.savefig(output_dir / 'cluster.png')
    plt.close()

    # length analysis
    # do longer compositions perform better?

    length_asr_dict = {}
    for idx, p in enumerate(programs):
        if ".then(" not in p:
            length = 1
        else: 
            length = len(p.split(".then("))

        if length not in length_asr_dict:
            length_asr_dict[length] = []

        length_asr_dict[length].append(asrs[idx])

    # dump this to a file
    json.dump(length_asr_dict, open(output_dir / 'lengths.json', 'w'))

    avg_asrs = []
    lengths = []
    std_dev_asrs = []
    for _ in length_asr_dict:
        lengths.append(_)
        avg_asrs.append(np.mean(length_asr_dict[_]))
        std_dev_asrs.append(np.var(length_asr_dict[_]))


    plt.bar(lengths, avg_asrs, yerr=std_dev_asrs, capsize=5, align='center')
    plt.xlabel("Number of Primitives in Attack")
    plt.ylabel("Average ASR")
    plt.title("ASR vs Number of Primitives in Attack")
    plt.grid(axis='y')
    plt.savefig(output_dir / 'asr_vs_length.png')
    plt.close()

    # high performing attacks analysis
    # only keep attacks with asrs >= 0.8
    programs = np.array(programs)
    asrs = np.array(asrs)

    programs = programs[asrs >= 0.8]
    asrs = asrs[asrs >= 0.8]
    
    embeddings = embedder.encode(programs, convert_to_numpy=True, batch_size=1024, show_progress_bar=True)

    tsne = TSNE(n_components=2, learning_rate='auto',
                  init='random', perplexity=5)
    tsne_results = tsne.fit_transform(np.array(embeddings))

    df = pd.DataFrame()
    df['tsne-2d-one'] = tsne_results[:,0]
    df['tsne-2d-two'] = tsne_results[:,1]
    df['asr'] = asrs

    plt.figure(figsize=(16,10))
    scatter_plot = sns.scatterplot(
        x="tsne-2d-one", y="tsne-2d-two", hue='asr',
        palette=sns.color_palette("hls"),
        data=df,
        legend="auto",
        alpha=0.3
    ).set_title('Clustering of CodeBERT embeddings of Jailbreak Attacks (ASR >= 0.8)')
    fig = scatter_plot.get_figure()
    fig.savefig(output_dir / 'cluster_0.8.png')
    plt.close()

    length_asr_dict = {}
    for idx, p in enumerate(programs):
        if ".then(" not in p:
            length = 1
        else: 
            length = len(p.split(".then("))

        if length not in length_asr_dict:
            length_asr_dict[length] = []

        length_asr_dict[length].append(asrs[idx])

    for _ in length_asr_dict:
        print(f'Number of attacks of length {_} = {len(length_asr_dict[_])}')

    primitives_dict = {}
    for p in programs:
        primitives = p.split(').then(')
        processed_primitives = [_[:_.find('(')] for _ in primitives]

        for primitive in processed_primitives:
            if primitive not in primitives_dict:
                primitives_dict[primitive] = 0

            primitives_dict[primitive] += 1

    for _ in primitives_dict:
        print(f'Number of attacks with primitive {_} = {primitives_dict[_]}')