import numpy as np
import pandas as pd
import lightgbm as lgb
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns

import gc
import pickle
import cPickle
import datetime
import os
import copy
import sys
import time
from multiprocessing import Process, Pool, Manager 
from distutils.version import StrictVersion, LooseVersion

import preprocessing
import common

train_path = "data/trainset.csv"
test_path = "data/testset.csv"

def dataset_boxplot(trainset, testset):
    for col in trainset.columns:
        fname = "gen/%s_boxplot.png" % col
        print("Processing %s ......" % fname)

        fig, axes = plt.subplots(1, 2, figsize=(20, 10))
        axes = axes.flatten()

        sns.boxplot(trainset[col], ax=axes[0])
        sns.boxplot(testset[col], ax=axes[1])

        fig.savefig(fname, bbox_inches="tight")

def dist_distribution_plot(args):
    col, train_col, test_col = args\

    fname = u"gen/%s_distribution.png" % col
    print("Plotting %s ......" % fname)

    fig, axes = plt.subplots(1, 2, figsize=(20, 10))
    axes = axes.flatten()

    sns.distplot(train_col, ax=axes[0])
    sns.distplot(test_col, ax=axes[1])
    fig.savefig(fname, bbox_inches="tight")
    plt.close(fig)
    print("%s generated ! ......" % fname)

    train_col = None
    test_col = None
    gc.collect()

def bar_distribution_plot(args):
    col, uniques, train_col, test_col = args

    fname = u"gen/%s_distribution.png" % col
    print("Plotting %s ......" % fname)

    tr_count = [(train_col == val).sum() for val in uniques]
    te_count = [(test_col == val).sum() for val in uniques]

    fig, axes = plt.subplots(1, 2, figsize=(20, 10))
    axes = axes.flatten()

    sns.barplot(x=uniques, y=tr_count, ax=axes[0])
    sns.barplot(x=uniques, y=te_count, ax=axes[1])
    fig.savefig(fname, bbox_inches="tight")
    plt.close(fig)
    print("%s generated ! ......" % fname)

    train_col = None 
    test_col = None
    uniques = None
    gc.collect()

def dataset_distribution_plot(merged, trainset, testset):
    barplot_dtypes = ["category", "int8", "int16", "int32"]
    distplot_dtypes = ["float16", "float32", "float64"]

    barplot_data = trainset.select_dtypes(include=barplot_dtypes)
    distplot_data = trainset.select_dtypes(include=distplot_dtypes)

    bar_sub_args = []
    for col in barplot_data.columns:
        if(col == "HasDetections"):
            continue

        uniques = pd.unique(merged[col])
        bar_sub_args.append([col, uniques, trainset[col], testset[col]])
    
    dist_sub_args = []
    for col in distplot_data.columns:
        dist_sub_args.append([col, trainset[col], testset[col]])

    sub_proc = []
    for sub_arg in bar_sub_args:
        sub_proc.append(Process(target=bar_distribution_plot, args=[sub_arg]))
    for p in sub_proc:
        p.start()
    for p in sub_proc:
        p.join()

    sub_proc = []
    for sub_arg in dist_sub_args:
        sub_proc.append(Process(target=dist_distribution_plot, args=[sub_arg]))
    for p in sub_proc:
        p.start()
    for p in sub_proc:
        p.join()

def scatter_plot(args):
    x, y, fname = args

    print("Plotting %s ......" % fname)
    fig, ax = plt.subplots(figsize=(10, 10))
    sns.scatterplot(x=x, y=y, ax=ax)
    fig.savefig(fname, bbox_inches="tight")
    print("%s generated ! ......" % fname)

def col_pair_plot(data):
    names = data.columns

    sub_args = []
    for i in range(len(names)):
        for j in range(i):
            fname = "gen/scatter_%s-%s.png" % (names[i], names[j])
            sub_args.append([data[ names[i] ], data[ names[j] ], fname])

    task = []
    for sub_arg in sub_args:
        task.append(Process(target=scatter_plot, args=[sub_arg]))

    for p in task:
        p.start()
    for p in task:
        p.join()

def col_hist_plot(data):
    for col in data.columns:
        if("float" not in str(data[col].dtypes)):
            continue

        fname = "gen/hist_%s.png" % col
        print("Plotting %s ......" % fname)
        fig, ax = plt.subplots(figsize=(10, 10))
        sns.distplot(data[col], ax=ax)
        fig.savefig(fname, bbox_inches="tight")
        print("%s generated ! ......" % fname)

def category_cardinality_plot(data):
    fname = "gen/category_cardinality.png"

    fig, ax = plt.subplots(figsize=(10, 40))
    category_data = data.select_dtypes(include="category")
    length = [category_data[col].nunique() for col in category_data.columns]
    sns.barplot(x=length, y=category_data.columns, ax=ax)
    
    fig.savefig(fname, bbox_inches="tight")
    print("%s generated ! ......" % fname)

def load_test_data():
    sys.stderr.write("Loading test data ......\n")
    test_data = pd.read_csv(test_path, dtype=common.dtypes)
    machine_id = copy.copy(test_data["MachineIdentifier"])
    del test_data["MachineIdentifier"]

    return test_data, machine_id

def load_train_data():
    sys.stderr.write("Loading train data ......\n")

    data = pd.read_csv(train_path, dtype=common.dtypes)
    del data["MachineIdentifier"]
    
    return data

def main():
    tr_data = load_train_data()
    te_data, machine_id = load_test_data()
    offset = tr_data.shape[0]

    X = pd.concat([tr_data, te_data], sort=False)     
    del tr_data, te_data
    gc.collect()
    
    d = {}
    for col in common.category_col:
        X[col] = X[col].astype("category")
        X[col].cat.add_categories([""], inplace=True)
        d[col] = ""
    for col in common.non_category_col:
        d[col] = X[col].max() + 1
    X = X.fillna(d)

    ret = preprocessing.col_factorize(X)
    ret = preprocessing.reduce_mem_usage(ret)

    ret_train_data = ret[:offset]
    ret_test_data = ret[offset:]
    del ret, X
    gc.collect()

    dataset_boxplot(ret_train_data, ret_test_data)

main()
