#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Example analysis of Scanette logs to generate missing tests.

@author: utting@usc.edu.au
"""

from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import AdaBoostClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.gaussian_process import GaussianProcessClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import LinearSVC
from sklearn.dummy import DummyClassifier

from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import cross_val_score
from sklearn.pipeline import Pipeline
from sklearn.tree import export_graphviz, export_text
import matplotlib.pyplot as plt
from pathlib import Path
import random
import timeit
import sys

import agilkia

# %%

DISPLAY_TIME = 3
TEST_MODE = "--test" in sys.argv
if TEST_MODE:
    print(f"Running in test mode - all figures will close after {DISPLAY_TIME} seconds.")

# %%

INPUT = Path("log_split.json")
cust = agilkia.TraceSet.load_from_json(INPUT)
print(f"Loaded {len(cust)} customer traces from {INPUT}")
print(cust.get_event_chars())
cust.meta_data["dataset"] = "Scanette customer traces"


# %% Get some data about each trace, to use for clustering traces.

data = cust.get_trace_data(method="action_counts")  # or add: method="action_status_counts"
print(data.sum().sort_values())

# %% Now cluster the traces (default MeanShift)

num_clusters = cust.create_clusters(data)
print(num_clusters, "clusters found")
for i in range(num_clusters):
    print(f"Cluster {i}:")
    n = 0
    for tr in cust.get_cluster(i):
        n += 1
        if n > 5:
            print("    ...")
            break
        print(f"    {tr}")


# %% Visualise clusters (using TSNE)

vis = PCA(n_components=2)
cust.visualize_clusters(algorithm=vis, block=(not TEST_MODE))
if TEST_MODE:
    plt.pause(3)
    plt.close()

# %%

print(cust.meta_data)

# %% The event signatures for Scanette

signature = {
    "abandon": {"input":{}, "output":{"Status":"int"}},
    "ajouter": {"input":{}, "output":{"Status":"int"}},
    "debloquer": {"input":{}, "output":{"Status":"int"}},
    "fermerSession": {"input":{}, "output":{"Status":"int"}},
    "ouvrirSession": {"input":{}, "output":{"Status":"int"}},
    "payer": {"input":{}, "output":{"Status":"int"}},
    "scanner": {"input":{}, "output":{"Status":"int"}},
    #"supprimer": {"input":{}, "output":{"Status":"int"}},
    "transmission": {"input":{}, "output":{"Status":"int"}}
    }
    
# %% Some largish clusters that are missing system tests.

missing = [3, 4, 5]

# %% Build models for those clusters and evaluate the models.
for n in missing:
    cluster = agilkia.TraceSet(cust.get_cluster(n))
    print(f"========== cluster {n} has {len(cluster)} traces ============")

    # Learn a test-generation model for this cluster.
    ex = agilkia.TracePrefixExtractor()
    X = ex.fit_transform(cluster)
    y = ex.get_labels()
    #print(X.head())
    #print(f"y: {y[0:20]}...")

    # Evaluate various classifiers
    for name,alg in [
                ("Tree", DecisionTreeClassifier()),   # fast, 0.951
                ("GBC", GradientBoostingClassifier()),  # slower, 0.951
                ("RandForest", RandomForestClassifier(n_estimators=100)),  # med 0.951
                ("AdaBoost", AdaBoostClassifier()),  # 0.421, some labels have no predictions
                #("Gaussian", GaussianProcessClassifier(max_iter_predict=10)),   # VERY slow, 0.886
                ("NeuralNet", MLPClassifier(solver='lbfgs')),  # adam solver doesn't converge. 0.924
                ("KNeighbors", KNeighborsClassifier()),  # fast, 0.948
                ("NaiveBayes", GaussianNB()),  # fast, F1 undef. 0.839
                ("LinearSVC", LinearSVC()),  # fast, 0.886
                ("Dummy", DummyClassifier()),  # fast, 0.130
                ("LogReg", LogisticRegression(solver='lbfgs', max_iter=200, multi_class='auto'))  # med 0.89
            ]:
        pipe = Pipeline([
            ("Normalize", MinMaxScaler()),
            (name, alg)
            ])
        scores = cross_val_score(pipe, X, y, cv=10, scoring='f1_macro')
        # print(scores)
        print(f"{name:20s} & {scores.mean():0.3f} (+/- {scores.std() * 2:0.3f})\\\\")


# %% Now use Decision Tree model to generate some tests.
        
def gen_tests_for(traceset, name, traces=5):
    print(f"========== {name} has {len(traceset)} traces ============")

    # Learn a test-generation model for this cluster.
    ex = agilkia.TracePrefixExtractor()
    X = ex.fit_transform(traceset)
    y = ex.get_labels()
    print(f"  it has {len(X)} trace prefixes")
    # Train a decision tree model on this cluster
    model = Pipeline([
        ("Extractor", ex),
        ("Normalize", MinMaxScaler()),
        ("Tree", DecisionTreeClassifier())   # fast, 0.951
        ])
    model.fit(traceset, y)
    
    rand = random.Random(1234)
    smart = agilkia.SmartSequenceGenerator([], method_signatures=signature, rand=rand)
    smart.trace_set.set_event_chars(traceset.get_event_chars())
    # generate some tests
    for i in range(traces):
        smart.generate_trace_with_model(model, length=100)
    for tr in smart.trace_set:
        print(f"    {tr}")
    return model

# %%

for n in missing:
    cluster = agilkia.TraceSet(cust.get_cluster(n))
    gen_tests_for(cluster, f"cluster {n}")

# %%
    
model = gen_tests_for(cust, "all customer traces", traces=30)

# %% Get frequency of a generated trace (maybe partial)

max_length = 35
chars = cust.get_event_chars()


# %%

def freq(trace:agilkia.Trace) -> float:
    return trace.meta_data["freq"]
    
def gen_all():
    smart = agilkia.SmartSequenceGenerator(urls=[], method_signatures=signature, 
                                           verbose=False, action_chars=chars)
    seqs = smart.generate_all_traces(model, length=max_length, action_prob=0.01, path_prob=0.01,
                                     partial=True)
    return seqs

# %% Time the test suite generation
    
duration = timeit.timeit(gen_all, number=10) / 10.0
print(f"duration = {duration} seconds")

# %%

seqs = gen_all()
sorted_seqs = sorted(seqs, key=freq, reverse=True)
print(f"generated {len(seqs)} sequences")
total = 0.0
for tr in sorted_seqs:
    etc = "" if len(tr) < max_length else "???"
    percent = freq(tr) * 100
    total += percent
    print(f"  {percent:5.2f}% {tr.to_string(to_char=chars)}{etc}")
print(f" {total:6.2f}% total")

# %%

ex = model[0]
tree = model[-1]
print(tree)

# %% 

# Export as dot file
export_graphviz(tree, out_file='tree.dot', 
                feature_names = ex.get_feature_names(),
                class_names = model.classes_,
                rounded = True, proportion = False, 
                precision = 2, filled = True)

# %% print as text
tree_str = export_text(tree, feature_names=ex.get_feature_names())
print(tree_str)

# %%

if not TEST_MODE:
    # Convert to png using system command (requires Graphviz)
    from subprocess import call
    call(['dot', '-Tpng', 'tree.dot', '-o', 'tree.png', '-Gdpi=300'])

    # Display in jupyter notebook (does not work in Spyder)
    from IPython.display import Image
    Image(filename = 'tree.png')
