#!/usr/bin/env python3

import sys
import logging
import yaml
import pandas as pd
import numpy as np
from collections import defaultdict

from sklearn.model_selection import train_test_split
from sklearn.ensemble import IsolationForest
from sklearn.impute import SimpleImputer

from anoflows.hpo import find_best_flows

from data_loading import load_data

logging.getLogger().setLevel(logging.INFO)

if len(sys.argv) == 1:
    logging.error("YAML data specification missing from the command line arguments")
    exit(1)

spec_file = sys.argv[1]
df, spec = load_data(spec_file)
max_rows = min(len(df), spec.get("max_rows", 40000))
novelty_detection = spec.get("novelty", True)
normal_classes = spec["normal_classes"]

precision = defaultdict(list)

for rounds in range(spec.get("rounds", 1)):
    # random sampling
    df = df.sample(n=max_rows, replace=False)
    label_col = spec["label_column"]
    y = df[label_col].values
    other = df.drop(label_col, inplace=False, axis=1)
    X = other.values

    # imputing
    X = SimpleImputer(copy=False).fit_transform(X)

    # train/test split
    X_train, X_test, y_train, y_test = \
        train_test_split(X, y, shuffle=False, test_size=0.5)
    if novelty_detection:
        keep = np.where(np.isin(y_train, normal_classes))[0]
        X_train = X_train[keep, :]
        y_train = y_train[keep]

    # training
    #flows, loss = find_best_flows(X_train, device='cpu', n_trials=1)
    from anoflows.anoflow_bagging import AnoFlowBagging
    flows = AnoFlowBagging()
    flows.fit(X_train)
    iforest = IsolationForest().fit(X_train)

    # prediction
    pred = {
        "anoflows": flows.likelihood(X_test),
        "iforest": iforest.decision_function(X_test)
    }

    # evaluation
    y_true = np.where(np.isin(y_test, spec["anomaly_classes"]))[0]
    ref = np.zeros(len(y_test))
    ref[y_true] = 1
    k = len(y_true)
    for name, y_pred in pred.items():
        anomaly_indices = y_pred.argsort()[:k]
        prec = ref[anomaly_indices].sum() / k
        logging.info("%s: %.1f%% (%d anomalies / %d rows)" % (name, 100*prec, k, len(y_test)))
        precision[name].append(prec)

logging.info("* SUMMARY %s", spec_file)
for name, prec in precision.items():
    prec = 100 * np.array(prec)
    mean = np.mean(prec)
    std = np.std(prec)
    logging.info("%s; mean=%.1f%% std=%.1f%%" % (name, mean, std))

[((25, 11, 25, 31), 'data_loading.load_data', 'load_data', ({(25, 21, 25, 30): 'spec_file'}, {}), '(spec_file)', False, 'from data_loading import load_data\n'), ((30, 12, 30, 29), 'collections.defaultdict', 'defaultdict', ({(30, 24, 30, 28): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((75, 0, 75, 39), 'logging.info', 'logging.info', ({(75, 13, 75, 27): '"""* SUMMARY %s"""', (75, 29, 75, 38): 'spec_file'}, {}), "('* SUMMARY %s', spec_file)", False, 'import logging\n'), ((21, 4, 21, 84), 'logging.error', 'logging.error', ({(21, 18, 21, 83): '"""YAML data specification missing from the command line arguments"""'}, {}), "('YAML data specification missing from the command line arguments'\n    )", False, 'import logging\n'), ((45, 8, 45, 60), 'sklearn.model_selection.train_test_split', 'train_test_split', (), '', False, 'from sklearn.model_selection import train_test_split\n'), ((54, 12, 54, 28), 'anoflows.anoflow_bagging.AnoFlowBagging', 'AnoFlowBagging', ({}, {}), '()', False, 'from anoflows.anoflow_bagging import AnoFlowBagging\n'), ((78, 11, 78, 24), 'numpy.mean', 'np.mean', ({(78, 19, 78, 23): 'prec'}, {}), '(prec)', True, 'import numpy as np\n'), ((79, 10, 79, 22), 'numpy.std', 'np.std', ({(79, 17, 79, 21): 'prec'}, {}), '(prec)', True, 'import numpy as np\n'), ((80, 4, 80, 66), 'logging.info', 'logging.info', ({(80, 17, 80, 65): "('%s; mean=%.1f%% std=%.1f%%' % (name, mean, std))"}, {}), "('%s; mean=%.1f%% std=%.1f%%' % (name, mean, std))", False, 'import logging\n'), ((18, 0, 18, 19), 'logging.getLogger', 'logging.getLogger', ({}, {}), '()', False, 'import logging\n'), ((77, 17, 77, 31), 'numpy.array', 'np.array', ({(77, 26, 77, 30): 'prec'}, {}), '(prec)', True, 'import numpy as np\n'), ((41, 8, 41, 33), 'sklearn.impute.SimpleImputer', 'SimpleImputer', (), '', False, 'from sklearn.impute import SimpleImputer\n'), ((56, 14, 56, 31), 'sklearn.ensemble.IsolationForest', 'IsolationForest', ({}, {}), '()', False, 'from sklearn.ensemble import IsolationForest\n'), ((65, 22, 65, 62), 'numpy.isin', 'np.isin', ({(65, 30, 65, 36): 'y_test', (65, 38, 65, 61): "spec['anomaly_classes']"}, {}), "(y_test, spec['anomaly_classes'])", True, 'import numpy as np\n'), ((47, 24, 47, 56), 'numpy.isin', 'np.isin', ({(47, 32, 47, 39): 'y_train', (47, 41, 47, 55): 'normal_classes'}, {}), '(y_train, normal_classes)', True, 'import numpy as np\n')]