%load_ext watermark
%watermark -u -d -v -p numpy,matplotlib,scipy,pandas,sklearn,mlxtend

"""
Explanation: <a id='top'> </a>
Author: James Bourbeau
End of explanation
"""


import sys
sys.path.append('/home/jbourbeau/cr-composition')
print('Added to PYTHONPATH')

%matplotlib inline
from __future__ import division, print_function
from collections import defaultdict
import itertools
import numpy as np
from scipy import interp
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

from sklearn.metrics import accuracy_score, confusion_matrix, roc_curve, auc
from sklearn.model_selection import cross_val_score, ShuffleSplit, KFold, StratifiedKFold
from sklearn.cluster import KMeans
from mlxtend.feature_selection import SequentialFeatureSelector as SFS

import composition as comp
import composition.analysis.plotting as plotting
    
color_dict = {'light': 'C0', 'heavy': 'C1', 'total': 'C2',
             'P': 'C0', 'He': 'C1', 'O': 'C3', 'Fe':'C4'}

"""
Explanation: Cosmic-ray composition clustering
Table of contents

Define analysis free parameters
Data preprocessing
Fitting random forest
Fraction correctly identified
Spectrum
Unfolding
Feature importance
End of explanation
"""


comp_class = True
comp_list = ['light', 'heavy'] if comp_class else ['P', 'He', 'O', 'Fe']

"""
Explanation: Define analysis free parameters
[ back to top ]
Whether or not to train on 'light' and 'heavy' composition classes, or the individual compositions
End of explanation
"""


pipeline_str = 'xgboost'
pipeline = comp.get_pipeline(pipeline_str)

"""
Explanation: Get composition classifier pipeline
End of explanation
"""


energybins = comp.analysis.get_energybins()

"""
Explanation: Define energy binning for this analysis
End of explanation
"""


sim_train, sim_test = comp.preprocess_sim(comp_class=comp_class, return_energy=True)

data = comp.preprocess_data(comp_class=comp_class, return_energy=True)

"""
Explanation: Data preprocessing
[ back to top ]
1. Load simulation/data dataframe and apply specified quality cuts
2. Extract desired features from dataframe
3. Get separate testing and training datasets
4. Feature transformation
End of explanation
"""


clf_name = pipeline.named_steps['classifier'].__class__.__name__
print('=' * 30)
print(clf_name)
pipeline.fit(sim_train.X, sim_train.y)
train_pred = pipeline.predict(sim_train.X)
train_acc = accuracy_score(sim_train.y, train_pred)
print('Training accuracy = {:.2%}'.format(train_acc))
test_pred = pipeline.predict(sim_test.X)
test_acc = accuracy_score(sim_test.y, test_pred)
print('Testing accuracy = {:.2%}'.format(test_acc))
# scores = cross_val_score(
#     estimator=pipeline, X=sim_train.X, y=sim_train.y, cv=3, n_jobs=10)
# print('CV score: {:.2%} (+/- {:.2%})'.format(scores.mean(), scores.std()))
print('=' * 30)

splitter = ShuffleSplit(n_splits=1, test_size=.5, random_state=2)
for set1_index, set2_index in splitter.split(sim_train.X):
    sim_train1 = sim_train[set1_index]
    sim_train2 = sim_train[set2_index]

kmeans = KMeans(n_clusters=4)

pred = kmeans.fit_predict(sim_train.X)

MC_comp_mask = {}
for composition in comp_list:
    MC_comp_mask[composition] = sim_train.le.inverse_transform(sim_train.y) == composition
MC_comp_mask

light_0 = np.sum(pred[MC_comp_mask['light']] == 0)/np.sum(MC_comp_mask['light'])
light_1 = np.sum(pred[MC_comp_mask['light']] == 1)/np.sum(MC_comp_mask['light'])
print('percent light cluster in 0 = {}'.format(light_0))
print('percent light cluster in 1 = {}'.format(light_1))

heavy_0 = np.sum(pred[MC_comp_mask['heavy']] == 0)/np.sum(MC_comp_mask['heavy'])
heavy_1 = np.sum(pred[MC_comp_mask['heavy']] == 1)/np.sum(MC_comp_mask['heavy'])
print('percent heavy cluster in 0 = {}'.format(heavy_0))
print('percent heavy cluster in 1 = {}'.format(heavy_1))

"""
Explanation: Run classifier over training and testing sets to get an idea of the degree of overfitting
End of explanation
"""


