import re
import torch
import numpy as np

from PIL import Image
from gensim.models.doc2vec import Doc2Vec
from torch.nn.functional import softmax
from bytes.create_data.bytes2img import hex2mat


def get_opcode_sequence(file):
    """
    Get opcode sequence from .asm file
    :param file: The name of .asm file
    :return: The opcode sequence
    """
    # Save opcode sequences in string arrays
    opcode_seq = []
    # Use regular expression to get all opcodes and combine opcodes into sequences
    p = re.compile(r'\s([a-fA-F0-9]{2}\s)+\s*([a-z]+)')
    with open(file, errors='ignore') as f:
        for line in f:
            # If a line starts with .text, it is a assembly instruction line
            if line.startswith(".text"):
                m = re.findall(p, line)
                if m:
                    opc = m[0][1]
                    if opc != "align":
                        opcode_seq.append(opc)
    return opcode_seq


def vectorize_asm(file, model_path):
    """
    Vectorize uploaded .asm file.
    :param file: the path of uploaded .asm file
    :param model_path: the path of Doc2Vec model
    :return: 1*256 vector
    """
    model = Doc2Vec.load(model_path)
    vector = model.infer_vector(get_opcode_sequence(file))

    return vector


def predict_file(file, file_type):
    """
    Predict uploaded file.
    :param file: the path of uploaded file
    :param file_type: the type of uploaded file (asm/bytes/bmp)
    :return: probability of each category
    """
    probability = []
    if file_type == "asm":
        # get Doc2Vec vector
        vector = vectorize_asm(file, "./display/model/doc2vec.model")
        vector = np.reshape(vector, (1, 1, 256))

        x = torch.from_numpy(vector).cuda()

        model = torch.load("./display/model/lstm_dpcnn.pth").cuda()

        model.eval()
        output = model(x).cpu()
        probability = softmax(output, dim=1)
        probability = probability.detach().numpy()
        probability = np.reshape(probability, 9).tolist()

        for i in range(9):
            if probability[i] < 1e-6:
                probability[i] = 0

    elif file_type == "bytes":
        # get n * width array
        mat = hex2mat(file)

        img = Image.fromarray(mat)
        # convert image to bit map
        img = img.convert("L")
        img = img.resize((224, 224))
        mat = np.asarray(img, dtype=np.float32)
        mat = np.reshape(mat, (1, 1, 224, 224))
        mat = np.repeat(mat, 3, axis=1)
        # get 1*3*256*256 tensor
        x = torch.from_numpy(mat).cuda()

        model = torch.load("./display/model/resnet.pth").cuda()

        model.eval()
        output = model(x).cpu()
        probability = softmax(output, dim=1)
        probability = probability.detach().numpy()
        probability = np.reshape(probability, 9).tolist()

        for i in range(9):
            if probability[i] < 1e-6:
                probability[i] = 0

    elif file_type == "bmp":
        image = Image.open(file)
        image = image.resize((224, 224))
        mat = np.asarray(image, dtype=np.float32)
        mat = np.reshape(mat, (1, 1, 224, 224))
        mat = np.repeat(mat, 3, axis=1)
        # get 1*3*256*256 tensor
        x = torch.from_numpy(mat).cuda()

        model = torch.load("./display/model/resnet.pth").cuda()

        model.eval()
        output = model(x).cpu()
        probability = softmax(output, dim=1)
        probability = probability.detach().numpy().tolist()[0]

        for i in range(9):
            if probability[i] < 1e-6:
                probability[i] = 0

    return probability
