import click
import torch
from asm2vec.utils import (
    TraceData,
    train,
    save_model,
    cosine_similarities,
)
from asm2vec.datatype import Tokens
import json


def length_heuristic(l0, l1, debug=False):
    """
    Function length heuristic (since asm2vec is terrible at handling mismatched lengths)

    Returns a similarity metric in the range [0, 1]
    """
    length_diff = abs(l0 - l1)
    # Weight mismatched lengths considerably lower, but clip factor to 0
    length_factor = max(1 - 4 * (length_diff / (l0 + l1)), 0)
    if debug:
        print("Length factor", l0, l1, length_factor)
    return length_factor


def full_similarity_matrix(
    cosine_similarity_matrix,
    old_trace_data: TraceData,
    new_trace_data: TraceData,
):
    """
    Generates a matrix comparing all old opcodes to all new opcodes.

    Returns (old opcodes, new_opcodes, similarity_matrix)
    """
    old_fns = old_trace_data.traces
    new_fns = new_trace_data.traces

    old_opcodes = []
    new_opcodes = []
    # Full similarity matrix mapping old_opcodes => new_opcodes
    similarity_matrix = []

    for old_data in old_trace_data.opcodes.values():
        similarities = []
        for new_data in new_trace_data.opcodes.values():
            old_idx = old_data["fn_idx"]
            new_idx = new_data["fn_idx"]
            # Use the length of the instructions for the length heuristic
            l0 = len(old_fns[old_idx].insts)
            l1 = len(new_fns[new_idx].insts)
            length_factor = length_heuristic(l0, l1)

            h0 = old_data["packet_size_hint"]
            h1 = new_data["packet_size_hint"]

            packet_size_factor = 0
            if h0 == h1 and h0 != 0:
                packet_size_factor = 0.5
            else:
                packet_size_factor = -0.5
            # Since cosine similarity is in the range (-1, 1), add 1 to push it
            # into the range (0, 2).
            cs = cosine_similarity_matrix[old_idx, new_idx] + 1

            # Multiply the length factor and cosine similarity together to
            # yield some value in the range (0, 2), then subtract 1 to get a
            # score from range (-1, 1)
            score = length_factor * cs - 1

            # Add or subtract score depending on the packet size matching
            # Also clamp value to between (-1, 1)
            score = max(min(score + packet_size_factor, 1.0), -1.0)

            # Now we copy this similarity value for all opcodes in the new
            # switch case
            for op in new_data["opcodes"]:
                similarities.append(float(score))
        # Now we copy this similarity mapping for all opcodes in the old
        # switch case
        for op in old_data["opcodes"]:
            similarity_matrix.append(similarities)

    for old_data in old_trace_data.opcodes.values():
        for op in old_data["opcodes"]:
            old_opcodes.append(op)

    for new_data in new_trace_data.opcodes.values():
        for op in new_data["opcodes"]:
            new_opcodes.append(op)

    return (old_opcodes, new_opcodes, similarity_matrix)


def print_banner(text):
    print("")
    print(f"======= {text} =======")
    print("")


def write_matrix_to_file(output_file, old_opcodes, new_opcodes, similarity_matrix):
    with open(output_file, "w+") as f:
        json.dump(
            {
                "old_opcodes": old_opcodes,
                "new_opcodes": new_opcodes,
                "matrix": similarity_matrix,
            },
            f,
            indent=4,
        )
        print_banner(f"Output written to {output_file}")


@click.command()
@click.argument(
    "old_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
)
@click.argument(
    "new_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
)
@click.argument("output_file", type=click.Path(dir_okay=False, resolve_path=True))
def generate_similarity_matrix(old_traces, new_traces, output_file):
    """
    Compares the OLD_TRACES and NEW_TRACES directories generated by the
    `generate_deep_traces.py` script.

    Creates a JSON OUTPUT_FILE containing a pairwise similarity matrix of all
    opcodes found.

    \b
    {
        "old_opcodes": (list of old opcodes indexing dimension 0),
        "new_opcodes": (list of new opcodes indexing dimenision 1),
        "matrix": (m by n array of floats: [[]]),
    }
    Example:

    python generate_similarity_matrix.py old-traces/ new-traces/ similarity.json
    """
    tokens = Tokens()
    old_trace_data = TraceData.load_data(old_traces, tokens)
    new_trace_data = TraceData.load_data(new_traces, tokens)

    opath = "model.pt"

    def training_callback(context):
        progress = f'{context["epoch"]} | time = {context["time"]:.2f}, loss = {context["loss"]:.4f}'
        if context["accuracy"]:
            progress += f', accuracy = {context["accuracy"]:.4f}'
        print(progress)
        save_model(opath, context["model"], context["tokens"])

    training_params = {
        "embedding_size": 100,
        "batch_size": 1024,
        "epochs": 20,
        "neg_sample_num": 25,
        "calc_acc": True,
        "device": "cuda" if torch.cuda.is_available() else "cpu",
        "callback": training_callback,
        "learning_rate": 0.02,
    }

    print_banner("Training embeddings from scratch on old trace data")
    model = train(old_trace_data, **training_params)

    # Prepare the model for new trace data and freeze all training from old trace data
    model.init_estimation_mode(len(new_trace_data.traces))

    print_banner("Calculating embeddings for new trace data")
    model = train(new_trace_data, model=model, mode="test", **training_params)

    print_banner("Calculating cosine similarities")
    csm = cosine_similarities(model)

    print_banner("Computing full similarity matrix")
    old_opcodes, new_opcodes, similarity_matrix = full_similarity_matrix(
        csm, old_trace_data, new_trace_data
    )

    write_matrix_to_file(output_file, old_opcodes, new_opcodes, similarity_matrix)


if __name__ == "__main__":
    generate_similarity_matrix()