import pandas as pd
import numpy as np
from domain import psl
from sklearn import preprocessing
from sklearn.cluster import DBSCAN

min_samples = 2
eps = 0.6
topMap = {}
df1 = pd.read_csv('cluster_domain_1.csv', encoding='utf-8')
df2 = pd.read_csv('cluster_domain_2.csv', encoding='utf-8')
df = pd.concat([df1,df2])
topset = list(df.top)
topLen = len(topset)

for i in range(topLen):
    topMap[topset[i]] = [0] * topLen
    topMap[topset[i]][i] = 1

def saveFeatureDF(df):
    def countNum(x):
        numCnt = 0
        for i in x:
            if i >= '0' and i <= '9':
                numCnt += 1
        return numCnt

    def meanFeatures(data):
        return np.mean([len(x) for x in data])

    def preNum(data):
        return np.max([countNum(x.split('.')[0]) for x in data])

    def meanBlock(data):
        return np.mean([len(x.split('.')) for x in data])

    def featureFromFilename(filename):
        if filename is None or type(filename) == float:
            filename = '/'
        fileLen = len(filename)
        fileBlock = len(filename.split('/'))
        filename = filename.strip('/')
        fileChar = 0
        fileNum = 0
        fileSpecial = 0
        hasFileNum = 0
        for char in filename:
            if char >= '0' and char <= '9':
                fileNum += 1
                hasFileNum = 1
            elif char >= 'a' and char <= 'z':
                fileChar += 1
            else:
                fileSpecial += 1
        return [fileBlock, fileLen, fileChar, fileNum, fileSpecial, hasFileNum]

    t = df.groupby('host', as_index=False).agg({'domain': meanFeatures})
    t.columns = ['host', 'mean_len']
    df = pd.merge(df, t, on='host')

    t = df.groupby('host', as_index=False).agg({'domain': preNum})
    t.columns = ['host', 'numCnt']
    df = pd.merge(df, t, on='host')

    t = df.groupby('host', as_index=False).agg({'domain': meanBlock})
    t.columns = ['host', 'mean_block']
    df = pd.merge(df, t, on='host')
    min_max_scaler = preprocessing.MinMaxScaler()
    subCount = preprocessing.scale(df['subCount'], axis=0, with_mean=True, with_std=True, copy=True)
    subCount = subCount.reshape(-1, 1)
    df['sub_count'] = min_max_scaler.fit_transform(subCount)
    meanLen = preprocessing.scale(df['mean_len'], axis=0, with_mean=True, with_std=True,
                                  copy=True)  # np.array(df['mean_len'])
    meanLen = meanLen.reshape(-1, 1)
    df['mean_len_normal'] = min_max_scaler.fit_transform(meanLen)

    df['fileBlock'], df['fileLen'], df['fileChar'], df['fileNum'], df['fileSpecial'], df['hasFileNum'] = zip(
        *df['file'].map(featureFromFilename))
    t = df.groupby('host', as_index=False).agg(
        {'fileBlock': 'mean', 'fileLen': 'mean', 'fileChar': 'mean', 'fileNum': 'max', 'fileSpecial': 'mean',
         'hasFileNum': 'max'})
    df = pd.merge(df, t, on='host')
    df.to_csv('features.csv',index=False)


def firstFeature(data):
    domain = data['domain']
    top = data['top']
    vals = domain.lower().split('.')
    block = len(vals)
    charNum =0
    numberNum = 0
    allLen = 0
    for char in domain:
        if char == '.':
            continue
        allLen+=1
        if char>='0' and char<='9':
            numberNum+=1
        elif char>='a' and char<='z':
            charNum+=1
    mean_len = data['mean_len']
    blockLens = [len(x) for x in vals]
    blockDis = (max(blockLens)- min(blockLens))/mean_len
    std = np.std(blockLens)
    return [block,charNum/(allLen+0.0),numberNum/2.0,mean_len/10,blockDis,std] \
    +[data['fileBlock_y'],data['fileLen_y']/3,data['fileChar_y']/3,data['fileNum_y']/2,data['fileSpecial_y']/2,data['hasFileNum_y']]\
    + topMap[top]




def secondFeature(data):
    host = data['host']
    top = data['top']
    vals = host.lower().split('.')[0]
    block = data['mean_block']
    charNum =0
    numberNum = 0
    allLen = 0
    for char in vals[0]:
        if char == '.':
            continue
        allLen+=1
        if char>='0' and char<='9':
            numberNum+=1
        elif char>='a' and char<='z':
            charNum+=1
    subDomainNum = data['sub_count']
    mean_len = data['mean_len_normal']
    pre_len = data['mean_len']-allLen
    hasNumber = 1 if numberNum>0 else 0
    return [block,data['numCnt'],charNum/(allLen+0.0),numberNum/2,subDomainNum,mean_len,hasNumber,allLen,pre_len] + topMap[top]+\
    [data['fileBlock_y'],data['fileLen_y']/2,data['fileChar_y'],data['fileNum_y'],data['fileSpecial_y'],data['hasFileNum_y']]

def clusterFirst(df):
    first = df[df.subCount==1]
    print(first.shape)
    train = []
    def addToTrain(data):
        train.append(firstFeature(data))
    first.apply(addToTrain,axis=1)
    X = list(first.domain)
    return cluster(X,train)

def clusterSecond(df):
    second = df[df.subCount>1]
    del(second['domain'])
    del(second['file'])
    second = second.drop_duplicates()
    print(second.shape)
    print(second.columns)
    train = []
    def addToTrain(data):  
        train.append(secondFeature(data))
    second.apply(addToTrain,axis=1)
    X = list(second.host)
    return cluster(X,train)



def cluster(X,train):
    y_pred = DBSCAN(min_samples=min_samples,eps=eps).fit_predict(train)
    result = {}
    for i in range(len(X)):
        key = y_pred[i]
        if key not in result:
            result[key] = set()
        result[key].add(X[i])
    for key in result:
        print(key)
        print(result[key])
    print(len(result.keys()))
    return result


min_samples_list = [5,3,4,2,8,10]
eps_list  = [0.5,0.3,0.4,0.2,0.6,0.8]
#
# for sample in min_samples_list:
#     for ep in eps_list:
#         min_samples = sample
#         eps = ep
#         total = 0
#         for data in [clusterSecond(),clusterFirst()]:
#             for key,val in data.items():
#                 if key<0-1:
#                     print('not clustered %d'%len(val))
#                     # for domain in val:
#                     #     outf.write(domain+'\n')
#                     continue
#                 total += len(val)
#         print(min_samples,eps,total)

if __name__=='__main__':
    saveFeatureDF(df)
    temp = pd.read_csv('features.csv')
    total = 0
    with open('dga.txt','w+') as outf:
        # for data in [clusterFirst(df)]:
        for data in [clusterSecond(temp),clusterFirst(temp)]:
            for key,val in data.items():
                if key<0:
                    print('not clustered %d'%len(val))
                    # for domain in val:
                    #     outf.write(domain+'\n')
                    continue
                line = ' '.join(list(val))+'\n'
                total += len(val)
                outf.write(line)
    print(total)

