from feather import read_dataframe as read_feather
import numpy as np
import pandas as pd
import pdb

import pydens
from zpylib import data_path as dp


def select_group(cols, n):
    return np.random.choice(cols, size=n, replace=False)

def cols2density(df, cols):
    cade = pydens.cade.Cade(
        sim_size=1000000,
        verbose=True)
    diagnostics = cade.train(
        df=df[cols].copy(),
        diagnostics=True
    )
    print('classifier in-sample AUC = ' + str(diagnostics['auc']))
    return cade

def densify(infile, outfile, groups, models=None):
    print("### Loading " + infile)
    df = read_feather(infile)
    if models is None:
        models = [cols2density(df, cols) for cols in groups]
    print("Scoring on all groups ...")
    scores_df = pd.DataFrame({
        'pydens_' + str(k): models[k].density(df[groups[k]]) for k in range(len(groups))
    })
    print("Concatenating ...")
    new_df = pd.concat([df, scores_df], axis=1)
    print("Feathering ...")
    new_df.to_feather(outfile)
    return models


# Constants
np.random.seed(0)
ALLCOLS = pd.read_csv(dp('raw/train.csv'), nrows=2).columns.tolist()
DENSIFIABLE_COLS = [a for a in ALLCOLS
                    if ((a != 'HasDetections') and (a != 'MachineIdentifier'))]
N = 25 # number of features to densify per run
N_GROUPS = 10
GROUPS = [select_group(DENSIFIABLE_COLS, N).tolist() for k in range(N_GROUPS)]

# Main
models = densify(
    infile=dp("refactored/train.feather"),
    outfile=dp("refactored/densified_train.feather"),
    groups=GROUPS
)
_ = densify(
    infile=dp("refactored/test.feather"),
    outfile=dp("refactored/densified_test.feather"),
    groups=GROUPS,
    models=models
)