import os
import pandas as pd

FAMILIES = {
    'bamital.csv': 1,
    'bedep.csv': 2,
    'blackhole.csv': 3,
    'ccleaner.csv': 4,
    'chinad.csv': 5,
    'chir.csv': 6,
    'conficker.csv': 7,
    'corebot.csv': 8,
    'cryptolocker.csv': 9,
    'diamondfox.csv': 10,
    'dircrypt.csv': 11,
    'dmsniff.csv': 12,
    'dyre.csv': 13,
    'ebury.csv': 14,
    'ekforward.csv': 15,
    'emotet.csv': 16,
    'feodo.csv': 17,
    'fobber.csv': 18,
    'gameover.csv': 19,
    'goznym.csv': 20,
    'gspy.csv': 21,
    'hesperbot.csv': 22,
    'infy.csv': 23,
    'locky.csv': 24,
    'madmax.csv': 25,
    'makloader.csv': 26,
    'mirai.csv': 27,
    'modpack.csv': 28,
    'monerominer.csv': 29,
    'murofet.csv': 30,
    'murofetweekly.csv': 31,
    'mydoom.csv': 32,
    'necurs.csv': 33,
    'nymaim.csv': 34,
    'oderoor.csv': 35,
    'omexo.csv': 36,
    'padcrypt.csv': 37,
    'pandabanker.csv': 38,
    'proslikefan.csv': 39,
    'pushdo.csv': 40,
    'pushdotid.csv': 41,
    'pykspa.csv': 42,
    'pykspa2.csv': 43,
    'pykspa2s.csv': 44,
    'qadars.csv': 45,
    'qakbot.csv': 46,
    'ramdo.csv': 47,
    'ramnit.csv': 48,
    'ranbyus.csv': 49,
    'rovnix.csv': 50,
    'shifu.csv': 51,
    'simda.csv': 52,
    'sisron.csv': 53,
    'sphinx.csv': 54,
    'sutra.csv': 55,
    'szribi.csv': 56,
    'tempedreve.csv': 57,
    'tempedrevetdd.csv': 58,
    'tinba.csv': 59,
    'tinynuke.csv': 60,
    'tofsee.csv': 61,
    'torpig.csv': 62,
    'ud2.csv': 63,
    'ud3.csv': 64,
    'ud4.csv': 65,
    'urlzone.csv': 66,
    'vawtrak.csv': 67,
    'vidro.csv': 68,
    'vidrotid.csv': 69,
    'virut.csv': 70,
    'wd.csv': 71,
    'xshellghost.csv': 72,
    'xxhex.csv': 73
}

CHUNK_SIZE = 1000000
NUMBER_OF_FAMILIES = 73
MIXED_CHUNK_SIZE = 30000
NUMBER_OF_MIXED_SETS = 1


# -----------------------------------------------------------------------
def generate_mixed_sets():
    """
    """
    for file in sorted(os.listdir(".")):
        if file.endswith(".csv"):
            data = pd.read_csv(file, header=0, names=["domain"])
            
            for i in range(NUMBER_OF_MIXED_SETS):
                filename = os.path.join("..", "datasets_mixed", "dataset_mixed_1m_test_{0:02d}.csv".format(i))
                
                if data.size >= MIXED_CHUNK_SIZE:
                    data_sample = data.sample(MIXED_CHUNK_SIZE).reset_index(drop=True)
                    data_sample["group"] = FAMILIES.get(file)
                    data_sample.to_csv(filename, mode="a", header=None, index=False)
                
                else:
                    data["group"] = FAMILIES.get(file)
                    data.to_csv(filename, mode="a", header=None, index=False)
    
    for dataset in os.listdir("..\\datasets_mixed"):
        filename = os.path.join("..", "datasets_mixed", dataset)
        
        with open(filename, "r") as fr:
            data = fr.read()
        
        with open(filename, "w") as fw:
            fw.write("raw_domain,family\n" + data)



########################################################################
if __name__ == "__main__":
    generate_ratio_sets()