import os
import json
import math
import string
import numpy as np

from zlib import crc32

from . import basic_data
from . import modify_data

VOWELS = set("aeiouy")
CONSONANTS = set("bcdfghjklmnpqrstvwxz")
HEX_CHARS = string.hexdigits[:16]


#######################################################################
class FeatureExtraction(basic_data.BasicData):
    """
    Extract features from domain names.
    All features are calculated from the second-level domain string unless stated otherwise.
    """
    
    #----------------------------------------------------------------------
    def __init__(self, logger, config, type, filepath, filename, is_dga):
        """
        Constructor
        """
        super().__init__(logger, config, type, filepath, filename)
        self.is_dga = is_dga
    
    
    #----------------------------------------------------------------------
    def extract_all_features(self):
        self.get_domain()
        self.get_tld()
        
        for f in self.features_sld_apply:
            # f[0] - name of the function to apply
            # f[1] - name of the new column
            self.data[f[1]] = self.data["domain"].apply(f[0])
        
        for f in self.features_sld:
            f(self.data["domain"])
            
        for f in self.features_sld_with_length:
            f(self.data["domain"], self.data["domain_length"])
            
        for f in self.features_tld_apply:
            # f[0] - name of the function to apply
            # f[1] - name of the new column
            self.data[f[1]] = self.data["tld"].apply(f[0])
            
        for f in self.features_tld:
            f(self.data["tld"])
            
        for n in range(self.config.getint("features_ngrams_min"), self.config.getint("features_ngrams_max")):
            with open("{}grams.txt".format(n), "r") as ng:
                self.ngrams = json.load(ng)
                
                for f in self.features_ngrams:
                    # f[0] - name of the function to apply
                    # f[1] - name of the new column
                    # f[2] - number of repetitions
                    counts = self.data["domain"].apply(self.ngrams_counts, args=(n,), repetition_number=f[2])
                    self.data[f[1].format(n)] = counts.apply(f[0])
    
    
    #----------------------------------------------------------------------
    def get_tld(self):
        """
        Get top level domain.
        """
        self.data["tld"] = self.data.raw_domain.str.rsplit(".").str.get(-1)
    
    
    #----------------------------------------------------------------------
    def get_domain(self):
        """
        Get domain name.
        """
        self.data["domain"] = self.data.raw_domain.str.rsplit(".").str[:-1].str.join(".")
        
        
    #----------------------------------------------------------------------
    def domain_length(self, domain):
        """
        Calculate the length of the domain.
        """
        self.data["domain_length"] = domain.str.len()
        

    #----------------------------------------------------------------------
    def shannon_entropy(self, domain):
        """
        Calculate Shannon entropy.
        """
        domain = list(domain)
        n_labels = len(domain)
    
        value, counts = np.unique(domain, return_counts=True)
        probs = counts / n_labels
        n_classes = np.count_nonzero(probs)
    
        if n_classes <= 1:
            return 0
    
        ent = 0.
        for p in probs:
            ent -= p * math.log(p, 2)
    
        return ent


    # ----------------------------------------------------------------------
    def gini_coefficient(self, domain):
        """
        Calculate the Gini coefficient/index of the characters of the domain.
        """
        domain = list(domain)
        n_labels = len(domain)

        value, counts = np.unique(domain, return_counts=True)
        probs = counts / n_labels
        n_classes = np.count_nonzero(probs)

        if n_classes <= 1:
            return 0

        gini = 0.
        for p in probs:
            gini += math.pow(p, 2)

        return 1 - gini

    # ----------------------------------------------------------------------
    def classification_error_of_characters(self, domain):
        """
        Calculate the classification error of characters of the domain.
        """
        domain = list(domain)
        n_labels = len(domain)

        value, counts = np.unique(domain, return_counts=True)
        probs = counts / n_labels
        n_classes = np.count_nonzero(probs)

        if n_classes <= 1:
            return 0

        return 1 - np.max(probs)
    
    
    #----------------------------------------------------------------------
    def vowel_ratio(self, domain, length):
        """
        Calculate vowel character ratio of the domain.
        """
        self.data["vowel_ratio"] = domain.str.count("[aeiouy]") / length
        
        
    #----------------------------------------------------------------------
    def consonant_ratio(self, domain, length):
        """
        Calculate consonant character ratio of the domain.
        """
        self.data["consonant_ratio"] = domain.str.count("[bcdfghjklmnpqrstvwxz]") / length
        
        
    #----------------------------------------------------------------------
    def hex_character_ratio(self, domain, length):
        """
        Calculate hex character ratio of the domain.
        """
        self.data["hex_character_ratio"] = domain.str.count("[0-9a-f]") / length
        
        
    #----------------------------------------------------------------------
    def digit_ratio(self, domain, length):
        """
        Calculate digit ratio of the domain.
        """
        self.data["digit_ratio"] = domain.str.count("\d") / length
        
        
    #----------------------------------------------------------------------
    def digit_to_letter_ratio(self, domain):
        """
        Calculate digit to letter ratio of the domain.
        """
        self.data["digit_to_letter_ratio"] = np.where(domain.str.count("[a-z]") == 0,
                                                      domain.str.count("\d"),
                                                      domain.str.count("\d") / domain.str.count("[a-z]"))
        
        
    #----------------------------------------------------------------------
    def is_first_character_digit(self, domain):
        """
        Check if the first character of the domain is a digit.
        """
        self.data["is_first_character_digit"] = np.where(domain.str.get(0).str.isdigit(), 1, 0)
    
    
    #----------------------------------------------------------------------
    def number_of_digits(self, domain):
        """
        Calculate the number of digits in the domain.
        """
        self.data["number_of_digits"] = domain.str.count("\d")


    # ----------------------------------------------------------------------
    def number_of_unique_characters(self, domain):
        """
        Calculate the number of unique characters in the domain.
        """
        return len(set(domain))


    #----------------------------------------------------------------------
    def longest_consonant_sequence(self, domain):
        """
        Calculate the length of the longest consonant sequence in the domain.
        """
        con_seq_len = 0
        con_seq_len_max = 0
        
        for char in domain:
            if char in CONSONANTS:
                con_seq_len += 1
            else:
                if con_seq_len > con_seq_len_max:
                    con_seq_len_max = con_seq_len
                con_seq_len = 0
                
        return con_seq_len_max
    
    
    #----------------------------------------------------------------------
    def longest_vowel_sequence(self, domain):
        """
        Calculate the length of the longest vowel sequence in the domain.
        """
        vow_seq_len = 0
        vow_seq_len_max = 0
        
        for char in domain:
            if char in VOWELS:
                vow_seq_len += 1
            else:
                if vow_seq_len > vow_seq_len_max:
                    vow_seq_len_max = vow_seq_len
                vow_seq_len = 0
                
        return vow_seq_len_max
    
    
    #----------------------------------------------------------------------
    def longest_digit_sequence(self, domain):
        """
        Calculate the length of the longest digit sequence in the domain.
        """
        dig_seq_len = 0
        dig_seq_len_max = 0
        
        for char in domain:
            if char.isdigit():
                dig_seq_len += 1
            else:
                if dig_seq_len > dig_seq_len_max:
                    dig_seq_len_max = dig_seq_len
                dig_seq_len = 0
                
        return dig_seq_len_max
    
    
    #----------------------------------------------------------------------
    def is_md5_like(self, domain):
        """
        Check if the domain looks like MD5 hash.
        """
        self.data["is_md5_like"] = np.where(domain.str.len() == 32, np.where(domain.str.count("[0-9a-f]") == 32, 1, 0), 0)
        
    
    #----------------------------------------------------------------------
    def more_vowels(self, domain):
        """
        Check if the domain contains more (or same) vowels than consonants.
        """
        self.data["more_vowels"] = np.where(domain.str.count("[aeiouy]") >= domain.str.count("[bcdfghjklmnpqrstvwxz]"), 1, 0)
    
    
    #----------------------------------------------------------------------
    def tld_length(self, tld):
        """
        Calculate the length of the top-level domain.
        """
        self.data["tld_length"] = tld.str.len()
        
        
    #----------------------------------------------------------------------
    def tld_hash(self, tld):
        """
        Calculate the CRC32 hash of the top-level domain
        and normalize it to a value between 0 and 1.
        """
        return float(crc32(tld.encode()) & 0xffffffff) / 2**32
    
    
    #----------------------------------------------------------------------
    def ngrams_counts(self, domain, n, repetition_number=None):
        """
        Calculate frequency of ngrams of the domain.
        """
        if repetition_number:
            domain = domain * repetition_number
        
        domain_ngrams = []
        domain_length = len(domain)
        _range = domain_length - n + 1
        
        if domain_length > n:
            for i in range(_range):
                ngram = domain[i : i+n]
                domain_ngrams.append(self.ngrams.get(ngram, 0))
                
            return domain_ngrams
        
        else:
            return [0]
        
        
    #----------------------------------------------------------------------
    def ngram_average(self, counts):
        """
        Calculate the average of ngram counts of the domain.
        """
        avg = np.average(counts)
        return np.log10(avg) if avg > 1 else avg
    
    
    #----------------------------------------------------------------------
    def ngram_median(self, counts):
        """
        Calculate the median of ngram counts of the domain.
        """
        med = np.median(counts)
        return np.log10(med) if med > 1 else med
    
    
    #----------------------------------------------------------------------
    def ngram_std(self, counts):
        """
        Calculate the standard deviation of ngram counts of the domain.
        """
        std = np.std(counts)
        return np.log10(std) if std > 1 else std
    
    
    #----------------------------------------------------------------------
    def dga_flag(self, domain):
        """
        Set flag if the domain is generated by DGA.
        """
        self.data["dga_flag"] = self.is_dga
    
    
    #----------------------------------------------------------------------
    def main(self, subdir=""):
        """
        Self-explaining.
        """
        self.features_sld_apply = []
        for feature in self.config.get("features_sld_apply").split(";"):
            self.features_sld_apply.append((getattr(self, feature), feature))

        self.features_sld = []
        for feature in self.config.get("features_sld").split(";"):
            self.features_sld.append(getattr(self, feature))

        self.features_sld_with_length = []
        for feature in self.config.get("features_sld_with_length").split(";"):
            self.features_sld_with_length.append(getattr(self, feature))

        self.features_tld = []
        for feature in self.config.get("features_tld").split(";"):
            self.features_tld.append(getattr(self, feature))

        self.features_tld_apply = []
        for feature in self.config.get("features_tld_apply").split(";"):
            self.features_tld_apply.append((getattr(self, feature), feature))
        
        self.features_ngrams = (
            (self.ngram_average, "{}gram_avg", None),
            (self.ngram_median, "{}gram_med", None),
            (self.ngram_std, "{}gram_std", None),
        )
        
        self.extract_all_features()
        
        self.logger.debug("Reordering columns...")
        self.data = modify_data.reorder_columns(self.data, self.config.get("features_labels").split(";"))
        
        features_filepath = os.path.join(self.config.get("features_dir"), subdir, self.filename)
        os.makedirs(os.path.dirname(features_filepath), exist_ok=True)
        self.save_data_to_disk(features_filepath)



if __name__ == '__main__':
    pass