#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Predict the Malware categories
"""

import os

import numpy as np
import pandas as pd

from keras.models import load_model
from keras.preprocessing import sequence

from .utils import (url2domain, get_app_file_path, download_file, find_ngrams,
                    MODELS_BASE_URL)

NGRAMS = 2
FEATURE_LEN = 128


def load_model_data(year, latest=False):
    model_fn = 'malware_cat_lstm_{0:d}.h5'.format(year)
    model_path = get_app_file_path('pydomains', model_fn)
    if not os.path.exists(model_path) or latest:
        print("Downloading Malware model data from the server ({0!s})..."
              .format(model_fn))
        if not download_file(MODELS_BASE_URL + model_fn, model_path):
            print("ERROR: Cannot download Malware model data file")
            return None, None
    else:
        print("Using cached Malware model data from local ({0!s})...".format(model_path))
    vocab_fn = 'malware_cat_vocab_{0:d}.csv'.format(year)
    vocab_path = get_app_file_path('pydomains', vocab_fn)
    if not os.path.exists(vocab_path) or latest:
        print("Downloading Malware vocab data from the server ({0!s})..."
              .format(vocab_fn))
        if not download_file(MODELS_BASE_URL + vocab_fn, vocab_path):
            print("ERROR: Cannot download Malware vocab data file")
            return None, None
    else:
        print("Using cached Malware vocab data from local ({0!s})...".format(vocab_path))
    print("Loading Malware model and vocab data file...")
    #  sort n-gram by freq (highest -> lowest)
    vdf = pd.read_csv(vocab_path)
    vocab = vdf.vocab.tolist()
    model = load_model(model_path)

    return (model, vocab)


def pred_malware(df, domain_names="domain_names", year=2017, latest=False):
    """Predict the probability of a domain being implicated in
       distributing malware using the Malware model.

    Args:
        df (:obj:`DataFrame`): Pandas DataFrame. No default value.
        domain_names (str): Column name of the domain in DataFrame. 
            Default in `domain_names`.
        year (int): Malware model year. Only 2017 is available.
            Default is `2017`.
        latest (Boolean): Whether or not to download latest 
            model available from GitHub. Default is `False`.

    Returns:
        DataFrame: Pandas DataFrame with 3 additional columns:
            - `pred_malware_year_domain`: the domain name
            - `pred_malware_year_lab`: most probable category
            - `pred_malware_year_prob`: probability of the domain being
              implicated in distributing malware.
    """

    if domain_names not in df.columns:
        print("No column `{0!s}` in the DataFrame".format(column))
        return None

    model, vocab = load_model_data(year, latest)

    if model is None or vocab is None:
        print("ERROR: Couldn't load model data.")
        return None

    col_domain = 'pred_malware_{0:04d}_domain'.format(year)
    col_lab =  'pred_malware_{0:04d}_lab'.format(year)
    col_prob =  'pred_malware_{0:04d}_prob'.format(year)
    df[col_domain] = df[domain_names].apply(lambda c: url2domain(c, exclude_subdomains=['www']))

    # build X from index of n-gram sequence
    X = np.array(df[col_domain].apply(lambda c: find_ngrams(vocab, c, NGRAMS)))
    X = sequence.pad_sequences(X, maxlen=FEATURE_LEN)

    df[col_lab] = model.predict_classes(X, verbose=2)
    df[col_prob] = model.predict_proba(X, verbose=2)

    return df


if __name__ == "__main__":
    pass
