#!/usr/bin/python
import os
from datetime import datetime
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# Perform operation on CPU to avoid loading of libraries for GPU (no training required)
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
from os import path
from sys import argv
from Model import MCModel
from Logger import debug, error, info
import numpy as np

class MalwareClassificator:
    def __init__(self, family_n, image_size, family_labels=None):
        self.family_n = family_n
        self.family_labels = family_labels
        self.image_size = image_size
        self.channels = 3 # Even if it is a GS image, they needs to be three
        self.weights_loaded = False
        self.model = MCModel(family_n, image_size, 3)

    def load_weights(self, path="./weights_result"):
        if not os.path.exists('{}.index'.format(path)):
            raise FileNotFoundError('weights file {} does not exist or is unaccessible'.format(path))
        else:
            try:
                self.model.load_weights(path)
                self.model.compile(
                    optimizer=tf.keras.optimizers.Adam(),
                    loss=tf.keras.losses.CategoricalCrossentropy(),
                    metrics=['accuracy']
                )
            except Exception as ex:
                raise Exception('error during model loading: {}'.format(ex))
            self.weights_loaded = True 
    
    def predict(self, path):
        if self.weights_loaded:
            image_content = tf.io.read_file(argv[1])
            image = tf.image.decode_image(image_content, channels=3)
            image = tf.image.resize(image, [self.image_size, self.image_size])
            image_array = np.array([image])
            return self.model.predict(image_array)
        else:
            raise Exception('model weights have not been loaded')

    def get_label_from_prediction(self, prediction):
        if not self.family_labels is None:
            index = np.argmax(prediction)
            return self.family_labels[index]
        else:
            raise Exception('family labels not set')

def help(name):
    error("usage: {} <image_path> [<model_weights>]".format(name))
    exit()

family_labels = [
    'Airpush_variety1',
    'Dowgin_variety1',
    'FakeInst_variety3',
    'Fusob_variety2',  
    'Mecor_variety1',
    'trusted'
]

if __name__ == '__main__':
    help(argv[0]) if len(argv) < 2 else None

    weigths_path = argv[2] if argv[2] else './weights_result'

    classificator = MalwareClassificator(6, 200, family_labels)

    try:
        if len(argv) >= 3:
            classificator.load_weights(argv[2])
        else:
            classificator.load_weights()
    except IOError as e:
        error(str(e))
        exit()
    
    if not '-q' in argv:
            debug('model correctly loaded')
            debug('loading image {}'.format(argv[1]))
    
    try:
        prediction = classificator.predict(argv[1])
    except Exception as e:
        error(str(e))
        exit()
    
    classificator.model.model().summary()
    exit()

    info('prediction for the image: {}'.format(classificator.get_label_from_prediction(prediction)))
