# -*- coding: utf-8 -*-
"""
Created on Wed Apr 10 17:34:48 2019

@author: Arnold Yu
"""
# import libraries
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# calculate accuracy
def calculateAcc(matrix):
    denom = matrix[0][0] + matrix[0][1] + matrix[1][0] + matrix[1][1]
    num = matrix[0][0] + matrix[1][1]
    return float(num)/float(denom)

# calculate percision
def calculatePer( matrix):
    denom = matrix[0][0]+ matrix[1][0]
    return float(matrix[0][0])/float(denom)

# calculate recall
def calculateRec( matrix):
    denom = matrix[0][0]+ matrix[0][1]
    return float(matrix[0][0])/float(denom)

# calculate F1
def calculateF1(percision, recall):
    return float(2.0 * percision * recall)/ float(percision + recall)

# calculate false positive rate
def FPR(matrix):
    denom = matrix[0][1]+ matrix[1][1]
    return float(matrix[0,1])/float(denom)

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# import dataset
train = pd.read_csv('data/train_replaceMissing.csv')
test = pd.read_csv('data/test_replaceMissing.csv')

from sklearn.preprocessing import StandardScaler, LabelEncoder

# turn categorical data into numerical
le = LabelEncoder()
columns_cat = ['Workclass', 'Education', 'Marital', 'Occupation', 'Relationship', 'Race', 'Sex', 'Country', 'Income']
for col in columns_cat:
    if train[col].dtype == object:
        train[col] = le.fit_transform(train[col])
        test[col] = le.fit_transform(test[col])

# standardlize
sc = StandardScaler()
all_data = pd.concat([train, test])
all_X = all_data.iloc[:, :-1]
all_y = all_data.iloc[:, -1]
all_X = sc.fit_transform(all_X)

# import all classifiers
from sklearn.naive_bayes import GaussianNB
from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import AdaBoostClassifier
from sklearn.svm import SVC

models = {'Naïve Bayesian': GaussianNB(),
          'Decision tree classifier': DecisionTreeClassifier(criterion = 'entropy', random_state = 0),
          'k-NN': KNeighborsClassifier(n_neighbors = 5, metric = 'minkowski', p = 2),
          'Random Forest': RandomForestClassifier(n_estimators = 10, criterion = 'entropy', random_state = 0),
          'Logistic regression': LogisticRegression( solver='lbfgs', random_state = 0),
          'AdaBoost': AdaBoostClassifier(n_estimators = 50, learning_rate = 1. ,  random_state = 0),
          'SVM linear': SVC(kernel = 'linear', random_state = 0, probability = True)
          }
results = {'Naïve Bayesian': {},
          'Decision tree classifier':  {},
          'k-NN': {},
          'Random Forest':  {},
          'Logistic regression':  {},
          'AdaBoost':  {},
          'SVM linear':  {}
          }
from sklearn.model_selection import cross_val_predict
from sklearn.metrics import confusion_matrix
from sklearn.metrics import classification_report
from sklearn.metrics import roc_curve
from sklearn.metrics import roc_auc_score
methodP = ['predict', 'predict_proba']
target_names = ['Class 0', 'Class 1']

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# Naïve Bayesian classifier

y_prob_nb = cross_val_predict(estimator = models['Naïve Bayesian'],
                           X = all_X,
                           y= all_y,
                           groups=None,
                           cv= 10,
                           n_jobs= -1,
                           verbose=0,
                           fit_params=None,
                           pre_dispatch='2*n_jobs',
                           method= methodP[1])

y_pred_nb = np.zeros(len(y_prob_nb))
for i in range(0, len(y_prob_nb)):
    if y_prob_nb[i][0] > y_prob_nb[i][1]:
        y_pred_nb[i] = 0
    else:
        y_pred_nb[i] = 1
#import sklearn
#print('The scikit-learn version is {}.'.format(sklearn.__version__))
cm_nb = confusion_matrix(all_y, y_pred_nb)
dict_nb = classification_report(all_y, y_pred_nb, target_names = target_names,  output_dict= True)

results['Naïve Bayesian']['Confusion_matrix'] = cm_nb
results['Naïve Bayesian']['Classification_report'] = dict_nb
results['Naïve Bayesian']['Accuracy'] = calculateAcc(cm_nb)
results['Naïve Bayesian']['Percision'] = calculatePer(cm_nb)
results['Naïve Bayesian']['Recall'] = calculateRec(cm_nb)
results['Naïve Bayesian']['F1'] = calculateF1(results['Naïve Bayesian']['Percision'],results['Naïve Bayesian']['Recall'])

fpr, tpr, thresholds = roc_curve(all_y, y_prob_nb[:, 1])
auc = roc_auc_score(all_y, y_prob_nb[:, 1])
plt.plot(fpr, tpr, label='ROC curve:(area = %0.3f)' % (auc))
plt.plot([0, 1], [0, 1], 'k--')  # random predictions curve
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.0])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate or (Recall)')
plt.title('ROC : %s' % 'Naïve Bayesian')
plt.legend(loc="lower right")

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# Decision tree classifier
y_prob_dt = cross_val_predict(estimator = models['Decision tree classifier'],
                           X = all_X,
                           y= all_y,
                           groups=None,
                           cv= 10,
                           n_jobs=-1,
                           verbose=0,
                           fit_params=None,
                           pre_dispatch='2*n_jobs',
                           method= methodP[1])

y_pred_dt = np.zeros(len(y_prob_dt))
for i in range(0, len(y_prob_dt)):
    if y_prob_dt[i][0] > y_prob_dt[i][1]:
        y_pred_dt[i] = 0
    else:
        y_pred_dt[i] = 1

cm_dt = confusion_matrix(all_y, y_pred_dt)
dict_dt = classification_report(all_y, y_pred_dt, target_names = target_names,  output_dict= True)

results['Decision tree classifier']['Confusion_matrix'] = cm_dt
results['Decision tree classifier']['Classification_report'] = dict_dt
results['Decision tree classifier']['Accuracy'] = calculateAcc(cm_dt)
results['Decision tree classifier']['Percision'] = calculatePer(cm_dt)
results['Decision tree classifier']['Recall'] = calculateRec(cm_dt)
results['Decision tree classifier']['F1'] = calculateF1(results['Decision tree classifier']['Percision'],results['Decision tree classifier']['Recall'])

fpr, tpr, thresholds = roc_curve(all_y, y_prob_dt[:, 1])
auc = roc_auc_score(all_y, y_prob_dt[:, 1])
plt.plot(fpr, tpr, label='ROC curve:(area = %0.3f)' % (auc))
plt.plot([0, 1], [0, 1], 'k--')  # random predictions curve
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.0])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate or (Recall)')
plt.title('ROC : %s' % 'Decision tree classifier')
plt.legend(loc="lower right")

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# k-NN classifier
y_prob_knn = cross_val_predict(estimator = models['k-NN'],
                           X = all_X,
                           y= all_y,
                           groups=None,
                           cv= 10,
                           n_jobs=-1,
                           verbose=0,
                           fit_params=None,
                           pre_dispatch='2*n_jobs',
                           method= methodP[1])

y_pred_knn = np.zeros(len(y_prob_knn))
for i in range(0, len(y_prob_knn)):
    if y_prob_knn[i][0] > y_prob_knn[i][1]:
        y_pred_knn[i] = 0
    else:
        y_pred_knn[i] = 1

cm_knn = confusion_matrix(all_y, y_pred_knn)
dict_knn = classification_report(all_y, y_pred_knn, target_names = target_names,  output_dict= True)

results['k-NN']['Confusion_matrix'] = cm_knn
results['k-NN']['Classification_report'] = dict_knn
results['k-NN']['Accuracy'] = calculateAcc(cm_knn)
results['k-NN']['Percision'] = calculatePer(cm_knn)
results['k-NN']['Recall'] = calculateRec(cm_knn)
results['k-NN']['F1'] = calculateF1(results['k-NN']['Percision'],results['k-NN']['Recall'])

fpr, tpr, thresholds = roc_curve(all_y, y_prob_knn[:, 1])
auc = roc_auc_score(all_y, y_prob_knn[:, 1])
plt.plot(fpr, tpr, label='ROC curve:(area = %0.3f)' % (auc))
plt.plot([0, 1], [0, 1], 'k--')  # random predictions curve
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.0])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate or (Recall)')
plt.title('ROC : %s' % 'k-NN')
plt.legend(loc="lower right")

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# Random Forest classifier
y_prob_rf= cross_val_predict(estimator = models['Random Forest'],
                           X = all_X,
                           y= all_y,
                           groups=None,
                           cv= 10,
                           n_jobs= -1,
                           verbose=0,
                           fit_params=None,
                           pre_dispatch='2*n_jobs',
                           method= methodP[1])

y_pred_rf = np.zeros(len(y_prob_rf))
for i in range(0, len(y_prob_rf)):
    if y_prob_rf[i][0] > y_prob_rf[i][1]:
        y_pred_rf[i] = 0
    else:
        y_pred_rf[i] = 1

cm_rf = confusion_matrix(all_y, y_pred_rf)
dict_rf = classification_report(all_y, y_pred_rf, target_names = target_names,  output_dict= True)

results['Random Forest']['Confusion_matrix'] = cm_rf
results['Random Forest']['Classification_report'] = dict_rf
results['Random Forest']['Accuracy'] = calculateAcc(cm_rf)
results['Random Forest']['Percision'] = calculatePer(cm_rf)
results['Random Forest']['Recall'] = calculateRec(cm_rf)
results['Random Forest']['F1'] = calculateF1(results['Random Forest']['Percision'],results['Random Forest']['Recall'])

fpr, tpr, thresholds = roc_curve(all_y, y_prob_rf[:, 1])
auc = roc_auc_score(all_y, y_prob_rf[:, 1])
plt.plot(fpr, tpr, label='ROC curve:(area = %0.3f)' % (auc))
plt.plot([0, 1], [0, 1], 'k--')  # random predictions curve
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.0])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate or (Recall)')
plt.title('ROC : %s' % 'Random Forest')
plt.legend(loc="lower right")

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# Logistic regression classifier
y_prob_logis = cross_val_predict(estimator = models['Logistic regression'],
                           X = all_X,
                           y= all_y,
                           groups=None,
                           cv= 10,
                           n_jobs= -1,
                           verbose=0,
                           fit_params=None,
                           pre_dispatch='2*n_jobs',
                           method= methodP[1])

y_pred_logis = np.zeros(len(y_prob_logis))
for i in range(0, len(y_prob_logis)):
    if y_prob_logis[i][0] > y_prob_logis[i][1]:
        y_pred_logis[i] = 0
    else:
        y_pred_logis[i] = 1

cm_logis = confusion_matrix(all_y, y_pred_logis)
dict_logis = classification_report(all_y, y_pred_logis, target_names = target_names,  output_dict= True)

results['Logistic regression']['Confusion_matrix'] = cm_logis
results['Logistic regression']['Classification_report'] = dict_logis
results['Logistic regression']['Accuracy'] = calculateAcc(cm_logis)
results['Logistic regression']['Percision'] = calculatePer(cm_logis)
results['Logistic regression']['Recall'] = calculateRec(cm_logis)
results['Logistic regression']['F1'] = calculateF1(results['Logistic regression']['Percision'],results['Logistic regression']['Recall'])

fpr, tpr, thresholds = roc_curve(all_y, y_prob_logis[:, 1])
auc = roc_auc_score(all_y, y_prob_logis[:, 1])
plt.plot(fpr, tpr, label='ROC curve:(area = %0.3f)' % (auc))
plt.plot([0, 1], [0, 1], 'k--')  # random predictions curve
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.0])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate or (Recall)')
plt.title('ROC : %s' % 'Logistic regression')
plt.legend(loc="lower right")

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# AdaBoost classifier
y_prob_ada = cross_val_predict(estimator = models['AdaBoost'],
                           X = all_X,
                           y= all_y,
                           groups=None,
                           cv= 10,
                           n_jobs= -1,
                           verbose=0,
                           fit_params=None,
                           pre_dispatch='2*n_jobs',
                           method= methodP[1])

y_pred_ada = np.zeros(len(y_prob_ada))
for i in range(0, len(y_prob_ada)):
    if y_prob_ada[i][0] > y_prob_ada[i][1]:
        y_pred_ada[i] = 0
    else:
        y_pred_ada[i] = 1

cm_ada = confusion_matrix(all_y, y_pred_ada)
dict_ada = classification_report(all_y, y_pred_ada, target_names = target_names,  output_dict= True)

results['AdaBoost']['Confusion_matrix'] = cm_ada
results['AdaBoost']['Classification_report'] = dict_ada
results['AdaBoost']['Accuracy'] = calculateAcc(cm_ada)
results['AdaBoost']['Percision'] = calculatePer(cm_ada)
results['AdaBoost']['Recall'] = calculateRec(cm_ada)
results['AdaBoost']['F1'] = calculateF1(results['AdaBoost']['Percision'],results['AdaBoost']['Recall'])

fpr, tpr, thresholds = roc_curve(all_y, y_prob_ada[:, 1])
auc = roc_auc_score(all_y, y_prob_ada[:, 1])
plt.plot(fpr, tpr, label='ROC curve:(area = %0.3f)' % (auc))
plt.plot([0, 1], [0, 1], 'k--')  # random predictions curve
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.0])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate or (Recall)')
plt.title('ROC : %s' % 'AdaBoost')
plt.legend(loc="lower right")

#----------------------------------------------------------------------------------------------------------------------------------------------------------

# svm linear classifier
y_prob_svm = cross_val_predict(estimator = models['SVM linear'],
                           X = all_X,
                           y= all_y,
                           groups=None,
                           cv= 10,
                           n_jobs= -1,
                           verbose=0,
                           fit_params=None,
                           pre_dispatch='2*n_jobs',
                           method= methodP[1])

y_pred_svm = np.zeros(len(y_prob_svm))
for i in range(0, len(y_prob_svm)):
    if y_prob_svm[i][0] > y_prob_svm[i][1]:
        y_pred_svm[i] = 0
    else:
        y_pred_svm[i] = 1

cm_svm= confusion_matrix(all_y, y_pred_svm)
dict_svm = classification_report(all_y, y_pred_svm, target_names = target_names,  output_dict= True)

results['SVM linear']['Confusion_matrix'] = cm_svm
results['SVM linear']['Classification_report'] = dict_svm
results['SVM linear']['Accuracy'] = calculateAcc(cm_svm)
results['SVM linear']['Percision'] = calculatePer(cm_svm)
results['SVM linear']['Recall'] = calculateRec(cm_svm)
results['SVM linear']['F1'] = calculateF1(results['SVM linear']['Percision'],results['SVM linear']['Recall'])

fpr, tpr, thresholds = roc_curve(all_y, y_prob_svm[:, 1])
auc = roc_auc_score(all_y, y_prob_svm[:, 1])
plt.plot(fpr, tpr, label='ROC curve:(area = %0.3f)' % (auc))
plt.plot([0, 1], [0, 1], 'k--')  # random predictions curve
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.0])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate or (Recall)')
plt.title('ROC : %s' % 'SVM linear')
plt.legend(loc="lower right")

#----------------------------------------------------------------------------------------------------------------------------------------------------------
# graph learning curves of all classifiers

from sklearn.model_selection import learning_curve

train_sizes, train_scores, test_scores = learning_curve(estimator = models['Naïve Bayesian'],
                                                        X = all_X,
                                                        y = all_y,
                                                        # Number of folds in cross-validation
                                                        cv = 10,
                                                        # Evaluation metric
                                                        scoring='neg_log_loss', # this scoring can change to accuracy
                                                        # Use all computer cores
                                                        n_jobs = -1,
                                                        # 10 different sizes of the training set
                                                        train_sizes=np.linspace(0.01, 1.0, 10))

# Create means and standard deviations of training set scores
train_mean = np.mean(train_scores, axis=1)
train_std = np.std(train_scores, axis=1)

# Create means and standard deviations of test set scores
test_mean = np.mean(test_scores, axis=1)
test_std = np.std(test_scores, axis=1)
plt.grid()
# Draw lines
plt.plot(train_sizes, train_mean, '--', color="#111111",  label="Training score")
plt.plot(train_sizes, test_mean, color="#111111", label="Cross-validation score")

# Draw bands
plt.fill_between(train_sizes, train_mean - train_std, train_mean + train_std, color="#DDDDDD")
plt.fill_between(train_sizes, test_mean - test_std, test_mean + test_std, color="#DDDDDD")

# Create plot
plt.title("Learning Curve : Naïve Bayesian")
plt.xlabel("Training Set Size"), plt.ylabel("Neg Log Loss"), plt.legend(loc="best")
plt.tight_layout()
plt.show()

#----------------------------------------------------------------------------------------------------------------------------------------------------------

train_sizes, train_scores, test_scores = learning_curve(estimator = models['Decision tree classifier'],
                                                        X = all_X,
                                                        y = all_y,
                                                        # Number of folds in cross-validation
                                                        cv = 10,
                                                        # Evaluation metric
                                                        scoring='neg_log_loss', # this scoring can change to accuracy
                                                        # Use all computer cores
                                                        n_jobs = -1,
                                                        # 10 different sizes of the training set
                                                        train_sizes=np.linspace(0.01, 1.0, 10))

# Create means and standard deviations of training set scores
train_mean = np.mean(train_scores, axis=1)
train_std = np.std(train_scores, axis=1)

# Create means and standard deviations of test set scores
test_mean = np.mean(test_scores, axis=1)
test_std = np.std(test_scores, axis=1)
plt.grid()
# Draw lines
plt.plot(train_sizes, train_mean, '--', color="#111111",  label="Training score")
plt.plot(train_sizes, test_mean, color="#111111", label="Cross-validation score")

# Draw bands
plt.fill_between(train_sizes, train_mean - train_std, train_mean + train_std, color="#DDDDDD")
plt.fill_between(train_sizes, test_mean - test_std, test_mean + test_std, color="#DDDDDD")

# Create plot
plt.title("Learning Curve : Decision tree classifier")
plt.xlabel("Training Set Size"), plt.ylabel("Neg Log Loss"), plt.legend(loc="best")
plt.tight_layout()
plt.show()

#----------------------------------------------------------------------------------------------------------------------------------------------------------

train_sizes, train_scores, test_scores = learning_curve(estimator = models['k-NN'],
                                                        X = all_X,
                                                        y = all_y,
                                                        # Number of folds in cross-validation
                                                        cv = 10,
                                                        # Evaluation metric
                                                        scoring='neg_log_loss', # this scoring can change to accuracy
                                                        # Use all computer cores
                                                        n_jobs = -1,
                                                        # 10 different sizes of the training set
                                                        train_sizes=np.linspace(0.01, 1.0, 10))

# Create means and standard deviations of training set scores
train_mean = np.mean(train_scores, axis=1)
train_std = np.std(train_scores, axis=1)

# Create means and standard deviations of test set scores
test_mean = np.mean(test_scores, axis=1)
test_std = np.std(test_scores, axis=1)
plt.grid()
# Draw lines
plt.plot(train_sizes, train_mean, '--', color="#111111",  label="Training score")
plt.plot(train_sizes, test_mean, color="#111111", label="Cross-validation score")

# Draw bands
plt.fill_between(train_sizes, train_mean - train_std, train_mean + train_std, color="#DDDDDD")
plt.fill_between(train_sizes, test_mean - test_std, test_mean + test_std, color="#DDDDDD")

# Create plot
plt.title("Learning Curve : k-NN")
plt.xlabel("Training Set Size"), plt.ylabel("Neg Log Loss"), plt.legend(loc="best")
plt.tight_layout()
plt.show()

#----------------------------------------------------------------------------------------------------------------------------------------------------------

train_sizes, train_scores, test_scores = learning_curve(estimator = models['Random Forest'],
                                                        X = all_X,
                                                        y = all_y,
                                                        # Number of folds in cross-validation
                                                        cv = 10,
                                                        # Evaluation metric
                                                        scoring='neg_log_loss', # this scoring can change to accuracy
                                                        # Use all computer cores
                                                        n_jobs = -1,
                                                        # 10 different sizes of the training set
                                                        train_sizes=np.linspace(0.01, 1.0, 10))

# Create means and standard deviations of training set scores
train_mean = np.mean(train_scores, axis=1)
train_std = np.std(train_scores, axis=1)

# Create means and standard deviations of test set scores
test_mean = np.mean(test_scores, axis=1)
test_std = np.std(test_scores, axis=1)
plt.grid()
# Draw lines
plt.plot(train_sizes, train_mean, '--', color="#111111",  label="Training score")
plt.plot(train_sizes, test_mean, color="#111111", label="Cross-validation score")

# Draw bands
plt.fill_between(train_sizes, train_mean - train_std, train_mean + train_std, color="#DDDDDD")
plt.fill_between(train_sizes, test_mean - test_std, test_mean + test_std, color="#DDDDDD")

# Create plot
plt.title("Learning Curve : Random Forest")
plt.xlabel("Training Set Size"), plt.ylabel("Neg Log Loss"), plt.legend(loc="best")
plt.tight_layout()
plt.show()

#----------------------------------------------------------------------------------------------------------------------------------------------------------

train_sizes, train_scores, test_scores = learning_curve(estimator = models['Logistic regression'],
                                                        X = all_X,
                                                        y = all_y,
                                                        # Number of folds in cross-validation
                                                        cv = 10,
                                                        # Evaluation metric
                                                        scoring='neg_log_loss', # this scoring can change to accuracy
                                                        # Use all computer cores
                                                        n_jobs=-1,
                                                        # 10 different sizes of the training set
                                                        train_sizes=np.linspace(0.01, 1.0, 10))

# Create means and standard deviations of training set scores
train_mean = np.mean(train_scores, axis=1)
train_std = np.std(train_scores, axis=1)

# Create means and standard deviations of test set scores
test_mean = np.mean(test_scores, axis=1)
test_std = np.std(test_scores, axis=1)
plt.grid()
# Draw lines
plt.plot(train_sizes, train_mean, '--', color="#111111",  label="Training score")
plt.plot(train_sizes, test_mean, color="#111111", label="Cross-validation score")

# Draw bands
plt.fill_between(train_sizes, train_mean - train_std, train_mean + train_std, color="#DDDDDD")
plt.fill_between(train_sizes, test_mean - test_std, test_mean + test_std, color="#DDDDDD")

# Create plot
plt.title("Learning Curve : Logistic regression")
plt.xlabel("Training Set Size"), plt.ylabel("Neg Log Loss"), plt.legend(loc="best")
plt.tight_layout()
plt.show()

#----------------------------------------------------------------------------------------------------------------------------------------------------------

train_sizes, train_scores, test_scores = learning_curve(estimator = models['AdaBoost'],
                                                        X = all_X,
                                                        y = all_y,
                                                        # Number of folds in cross-validation
                                                        cv = 10,
                                                        # Evaluation metric
                                                        scoring='neg_log_loss', # this scoring can change to accuracy
                                                        # Use all computer cores
                                                        n_jobs=-1,
                                                        # 10 different sizes of the training set
                                                        train_sizes=np.linspace(0.01, 1.0, 10))

# Create means and standard deviations of training set scores
train_mean = np.mean(train_scores, axis=1)
train_std = np.std(train_scores, axis=1)

# Create means and standard deviations of test set scores
test_mean = np.mean(test_scores, axis=1)
test_std = np.std(test_scores, axis=1)
plt.grid()
# Draw lines
plt.plot(train_sizes, train_mean, '--', color="#111111",  label="Training score")
plt.plot(train_sizes, test_mean, color="#111111", label="Cross-validation score")

# Draw bands
plt.fill_between(train_sizes, train_mean - train_std, train_mean + train_std, color="#DDDDDD")
plt.fill_between(train_sizes, test_mean - test_std, test_mean + test_std, color="#DDDDDD")

# Create plot
plt.title("Learning Curve : AdaBoost")
plt.xlabel("Training Set Size"), plt.ylabel("Neg Log Loss"), plt.legend(loc="best")
plt.tight_layout()
plt.show()

#----------------------------------------------------------------------------------------------------------------------------------------------------------

train_sizes, train_scores, test_scores = learning_curve(estimator = models['SVM linear'],
                                                        X = all_X,
                                                        y = all_y,
                                                        # Number of folds in cross-validation
                                                        cv = 10,
                                                        # Evaluation metric
                                                        scoring='neg_log_loss', # this scoring can change to accuracy
                                                        # Use all computer cores
                                                        n_jobs = -1,
                                                        # 10 different sizes of the training set
                                                        train_sizes=np.linspace(0.01, 1.0, 10))

# Create means and standard deviations of training set scores
train_mean = np.mean(train_scores, axis=1)
train_std = np.std(train_scores, axis=1)

# Create means and standard deviations of test set scores
test_mean = np.mean(test_scores, axis=1)
test_std = np.std(test_scores, axis=1)
plt.grid()
# Draw lines
plt.plot(train_sizes, train_mean, '--', color="#111111",  label="Training score")
plt.plot(train_sizes, test_mean, color="#111111", label="Cross-validation score")

# Draw bands
plt.fill_between(train_sizes, train_mean - train_std, train_mean + train_std, color="#DDDDDD")
plt.fill_between(train_sizes, test_mean - test_std, test_mean + test_std, color="#DDDDDD")

# Create plot
plt.title("Learning Curve : SVM linear")
plt.xlabel("Training Set Size"), plt.ylabel("Neg Log Loss"), plt.legend(loc="best")
plt.tight_layout()
plt.show()


from decimal import Decimal
from django.conf import settings
from shop.models import Product

class Cart(object):
    def __init__(self, request):
        #Ініціалізація обєкта корзини
        self.session = request.session
        cart = self.session.get(settings.CART_SESSION_ID)
        if not cart:
            #Зберігаємо в сесії пусту корзину
            cart = self.session[settings.CART_SESSION_ID] = {}
        self.cart = cart

    def add(self, product, quantity=1, update_quantity=False):
        #Додавання товару в корзину або поновлення його кількості
        product_id = str(product.id)
        if product_id not in self.cart:
            self.cart[product_id] = {'quantity': 0, 'price': str(product.price)}
        if update_quantity:
            self.cart[product_id]['quantity'] = quantity
        else:
            self.cart[product_id]['quantity'] += quantity
        self.save()

    def save(self):
        # Визначаємо сесію, як така, що була змінена
        self.session.modified = True

    def remove(self, product):
        # Видалення товарів із корзини
        product_id = str(product.id)
        if product_id in self.cart:
            del self.cart[product_id]
            self.save()

    def __iter__(self):
        #Проходимо по товарам корзини і отримуємо відповідні обєкти Product
        product_ids = self.cart.keys()
        # Отримуємо обєкти моделі Product  і передаємо їх в корзину
        products = Product.objects.filter(id__in=product_ids)
        print(products)
        cart = self.cart.copy()
        for product in products:
            cart[str(product.id)]['product'] = product
        for item in cart.values():
            # print(item['product'])
            item['price'] = Decimal(item['price'])
            item['total_price'] = item['price'] * item['quantity']
            yield item

    def __len__(self):
        # Повертаємо загальну кількість товарів в корзині
        return sum(item['quantity'] for item in self.cart.values())

    def get_total_price(self):
        return sum(
            Decimal(item['price']) * item['quantity']
            for item in self.cart.values()
        )

    def clear(self):
        # Очищаємо корзину
        del self.session[settings.CART_SESSION_ID]
        self.save()

    def __str__(self):
        return str(self.cart)


import multiprocessing as mp
import random as python_random
import sys
import logging
import regressors
logger = logging.getLogger('MLP_training')

max_depth = float('inf')
max_layer = float('inf')
overdetermined_level = 1
training_seqs = 0
threads = 1

def builder(generations,num_per_generation,depth,level,n_input,train,layer,par,seed):
	"""Build MLP topologies based on maximum width, depth, overdetermined level"""
	global max_depth
	max_depth = depth
	global overdetermined_level
	overdetermined_level = level
	global training_seqs
	training_seqs = train
	global max_layer
	max_layer = layer
	global threads
	threads = par
	python_random.seed(seed)

	logger.info('The theoretical maximum number of MLP topologies to build: '+str(float((2*(max_layer-1))**max_depth)))
	logger.info('The maximum number of MLP topologies to train over all generations: '+str(generations*num_per_generation))

	NNs =[]
	if ((2*(max_layer-1))**max_depth<generations*num_per_generation and (2*(max_layer-1))**max_depth<536870912):
		#build brute force if theoretical number is less than max to be sampled by genetic algorithm AND less than the maximum possible list size in python
		logger.info('Generating all possible MLPs')
		NNs = [[2]]
		while not(NNs[-1][-1]==max_layer and len(NNs[-1])==max_depth):
			new_NN = [x for x in NNs[-1]]
			if new_NN[-1]<max_layer:
				new_NN[-1] = new_NN[-1]+1
			else:
				new_NN = new_NN+[2]
			NNs.append(new_NN)
		NNs = [NN for NN in NNs if training_seqs/(overdetermined_level*regressors.connections(to_add))>=1.0]
		return (True,NNs)
	else:
		#build random subset
		logger.info('Generating a random subset of MLPs')
		if (training_seqs/(overdetermined_level*regressors.connections([2]))<1.0):
			return(True,[])
		max_layer = int(((training_seqs/overdetermined_level)-n_input*2-2)/3) #assuming the max is a two-layer network with 2 as the first layer
		logger.info('The maximum number of nodes in any layer and still within the overdetermined level is: '+str(max_layer))
		while len(NNs)<num_per_generation:
			num_layers = python_random.randint(1,max_depth)
			to_add =[]
			for _ in range(0,num_layers,1):
				to_add.append(python_random.randint(2,max_layer))
			if training_seqs/(overdetermined_level*regressors.connections(to_add))>=1.0:
				if not(to_add in NNs):
					NNs.append(to_add)
		return (False,NNs)

#implement a genetic algorithm to search for an optimal topology
def recombiner(input):
	"""Randomly recombine two topologies"""
	(net1,net2) = input
	joint1 = python_random.randint(0,len(net1)-1)
	joint2 = python_random.randint(0,len(net2)-1)
	new_network = net1[:joint1]+net2[joint2:]
	if training_seqs/(overdetermined_level*regressors.connections(new_network))>=1.0:
		if max_depth >= len(new_network):
			return (True,new_network)
	return (False,None)

def recombine(possible):
	"""Recombine topologies"""
	pairs = [(possible[x],possible[y]) for x in range(0,len(possible),1) for y in range(0,len(possible),1) if not(x == y)]
	p = mp.Pool(threads)
	results = p.map(recombiner,pairs,1)
	p.close()
	p.join()
	results = [x[1] for x in results if x[0]]
	return results

def mutater(network):
	"""Randomly mutate a topology"""
	new_network =[x for x in network]
	layer = python_random.randint(0,len(new_network)-1)
	new_network[layer] = python_random.randint(2,max_layer)
	if training_seqs/(overdetermined_level*regressors.connections(new_network))>=1.0:
		return (True,network)
	return (False,None)

def mutate(possible):
	"""Mutate topologies"""
	p = mp.Pool(threads)
	results = p.map(mutater,possible,1)
	p.close()
	p.join()
	results = [result[1] for result in results if result[0]]
	return results


import torch
from torch import Tensor, nn
from torch.utils.data import DataLoader

from tqdm import tqdm

class ZeroShotPredictor():

    def __init__(self,
        model: nn.Module, data_loader: DataLoader, device: str = 'cpu'
    ) -> None:
        self.model = model.to(device)
        self.data_loader = data_loader
        self.device = device

    def predict(self) -> dict:
        """
        Predict on the dataset
        """
        self.model.eval()
        pbar = tqdm(range(len(self.data_loader)), desc='Predicting')

        all_outputs, all_labels, all_metadata = [], [], []

        for i, data in enumerate(self.data_loader):
            with torch.no_grad():
                pbar.update()
                inputs, metadata = data['inputs'], data['metadata']
                inputs = {k: v.to(self.device) for k,v in inputs.items()}
                labels = inputs.pop('label')

                outputs = self.model(**inputs)
                logits = outputs['logits']

                all_outputs.append(logits.detach().cpu())
                all_labels.append(labels.cpu())
                all_metadata.append(metadata)

        metadata_dict = {k: sum([x[k] for x in all_metadata], []) for k in all_metadata[0].keys()}
        outputs_dict = dict(outputs=all_outputs, labels=all_labels)
        outputs_dict = {k: torch.cat(v) for k,v in outputs_dict.items()}

        res_dict = {**outputs_dict, **metadata_dict}

        metrics = self.model.compute_metrics(res_dict)

        return metrics


import zipfile
import PIL

from PIL import Image
import pytesseract
import cv2 as cv
import numpy as np
from PIL import Image
from PIL import ImageEnhance
from PIL import ImageFont, ImageDraw

filename = 'readonly/small_img.zip'
#filename = 'readonly/images.zip'

with zipfile.ZipFile(filename,'r') as myzip:
    myzip.printdir()
    myzip.extractall()
    print("Done")

# loading the face detection classifier
def extractor(extracted_file,keyword_search):
    filename = extracted_file
    keyword_search = keyword_search
    face_cascade = cv.CascadeClassifier('readonly/haarcascade_frontalface_default.xml')
    text = pytesseract.image_to_string(filename)
    if(keyword_search in text):
        image = Image.open(filename)
        cv_image = cv.imread(filename)
        faces = face_cascade.detectMultiScale(cv_image,1.72)
        drawing=ImageDraw.Draw(image)
        image_width = 500
        image_height = 200
        drawingboard = PIL.Image.new(image.mode,(image_width,image_height))
        row = 0
        col = 0
        i=1
        for x,y,w,h in faces:
            drawing.rectangle((x,y,x+w,y+h), outline="white")
            cropped = image.crop((x,y,x+w,y+h))

            if(row>=500):
                row = 0
                col = 100
                i=0
            drawingboard.paste(cropped,(row,col))
            row  = 100*i
            res = 1
            i+=1
        if(res > 0):
            print('Results found in {} '.format(filename))
            display(drawingboard)
        else:
            print('Result found in {} but no faces were found'.format(filename))
    return

extracted_files = []

for i in range(15):
    zip_filename = 'a-{}.png'.format(i)
    try:
        temp = Image.open(zip_filename)
        extracted_files.append(zip_filename)
    except:
        pass
keyword = input('Enter the keyword :')

for extracted_file in extracted_files:
    result = extractor(extracted_file,keyword)

print('Process completed , Have a nice day.')


import torchvision
import numpy as np
import os
import pickle

class Fmnist(object):
    def __init__(self, data_path: str, save_path: str, n_agents: int = 100, max_num_samples=500):
        data_path = data_path
        n_agents = n_agents

        data = torchvision.datasets.FashionMNIST(root=os.path.join(data_path, "FashionMNIST"),
                                                 train=True, download=True)
        x, y = data.data.numpy().reshape(-1, 1, 28, 28) / 255, np.array(data.targets)

        x = x.reshape(x.shape[0], x.shape[-1] * x.shape[-2])
        random_ids = np.random.randint(0, n_agents, x.shape[0])
        assert len(np.unique(random_ids)) == n_agents

        save_path += str(n_agents) + '_n' + str(max_num_samples).lower() + '/'
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        for agent in range(n_agents):
            agent_ids = np.where(random_ids == agent)[0]
            if max_num_samples is not None:
                assert max_num_samples <= len(agent_ids)
                x_agent = x[agent_ids, :][0:max_num_samples, :]
                y_agent = y[agent_ids][0:max_num_samples]
            else:
                x_agent = x[agent_ids, :]
                y_agent = y[agent_ids]
            agent_dict = {'dataset': {'X': x_agent,
                                      'Y': y_agent},
                          'id': str(agent)}
            with open(save_path + str(agent) + '.pickle', 'wb') as handle:
                pickle.dump(agent_dict, handle)
        self.save_path = save_path

if __name__ == '__main__':
    DATA_PATH = '/Users/etagjal/Air/data/cfl_boosting/'
    Fmnist(data_path=DATA_PATH, save_path=DATA_PATH + 'FashionMNIST/data/', n_agents=10, max_num_samples=5000)


import requests
import os
from datetime import datetime
import time
import shodan
import json
import helpers
import geome
from docker import write_docker_file
from kubernetes import write_kube_file
from schedule import shodan_query_manager

def query_shodan(api_key, query_file=None, keyword=None):
    api = shodan.Shodan(api_key)
    total = 0

    if keyword is not None:
        try:
            # Search Shodan
            results = api.search('{}'.format(keyword))
            time.sleep(3)
            # Show the results
            total += results['total']
            for result in results['matches']:
                ip = result['ip_str']
                data = result['data']
                geome.write_to_geo_file(result, ip)
                # geome.plot_map(result,ip)
                if 'docker' in result.keys():
                    docker = result['docker']
                else:
                    docker = None
                try:
                    create_report(ip, data, docker)
                except:
                    print("Report failed")
                    exit()
        except shodan.APIError as e:
            print('Error: {} \n {}'.format(e, line))
            pass
        print ("Total Results: {}".format(total))
        return 1
    if query_file is not None:
        with open(query_file, 'r', encoding="utf-8") as query_f:
            for line in query_f:
                try:
                    # Search Shodan
                    results = api.search('{}'.format(line))
                    time.sleep(3)
                    # Show the results
                    total += results['total']
                    for result in results['matches']:
                        ip = result['ip_str']
                        data = result['data']
                        if 'docker' in result.keys():
                            docker = result['docker']
                        else:
                            docker = None
                        create_report(ip, data, docker)
                except shodan.APIError as e:
                    print('Error: {} \n {}'.format(e, line))

        print ("Total Results: {}".format(total))
        return 1

def create_report(ip, data, docker=None, kubernetes=None):
    # ex. directory 2019_1_1/
    date_string = helpers.get_date_string()
    helpers.report_dir_check(date_string)

    report = "reports/" + date_string +'/'+ ip +'.txt'
    try:
        helpers.write_to_ip_file(report, ip, data)

        if docker is not None:
            write_docker_file(docker=docker,ip=ip, date_string=date_string)

        if kubernetes is not None:
            write_kube_file(ip=ip, date_string=date_string, kubernetes=kubernetes)

        with open("reports/" + date_string + "/all_ips.txt", 'a+') as ip_file:
            ip_file.write("{}\n".format(ip))
            ip_file.close()

        geome.get_location_data_from_file(data, ip)

    except Exception as e:
        pass

def all_results(api_key, keyword=None, limit_results=None):
    api = shodan.Shodan(api_key)

    if limit_results is None:
        limit = helpers.get_total_results(api, keyword=keyword)
    else:
        limit = limit_results
    date_string = helpers.get_date_string()
    helpers.report_dir_check(date_string)

    if keyword is not None:
        output_file = "reports/" + date_string + "/-export" + keyword.replace(" ", "_").replace(":", "_").replace("\"","_").replace(";","").replace(".","").replace("-","").replace("&","")
        if os.path.exists(output_file+ ".json") is False:

            os.system("shodan download --limit {} {} {}".format(limit,output_file, keyword))
            os.system("gzip -d {}".format(output_file + ".json.gz"))
        else:
            print("That query has already been ran...exiting\n\n")
            time.sleep(3)
            exit()
        with open(output_file + ".json", "r") as out_file:
            docker=None
            kubernetes=None
            for line in out_file.readlines():
                result = json.loads(line)
                ip = result['ip_str']
                data = result['data']
                if 'opts' in result.keys():

                    if 'docker' in result['opts'].keys():
                        docker = result['opts']['docker']
                    else:
                        docker = None

                if result['product'] == "Kubernetes":
                    kubernetes= result
                create_report(ip, data, docker=docker, kubernetes=kubernetes)

        out_file.close()


from apscheduler.schedulers.background import BackgroundScheduler
from schedule.constants import SCHEDULER_TIME
from service import operation_service, OperationService
import logging

class Schedule(object):
    def __init__(self):
        self.schedule = None

    def start_schedule(self) -> None:
        logging.info("------------ starting the schedule -----------")
        self.schedule = BackgroundScheduler(daemon=True)
        self.schedule.add_job(id='send_next_job', func=self._send_next_job_to_iot, trigger='interval',
                              kwargs={'service': operation_service}, seconds=SCHEDULER_TIME)
        self.schedule.start()

    def stop_schedule(self) -> None:
        if self.schedule:
            self.schedule.shutdown()
            logging.info("-------------- stopping the schedule ------------")

    @staticmethod
    def _send_next_job_to_iot(service: OperationService) -> None:
        logging.info("--------------- scheduler sending next job -------------")
        service.get_next_operation()


import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from torch.hub import load_state_dict_from_url
except ImportError:
    from torch.utils.model_zoo import load_url as load_state_dict_from_url

def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)

def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out

class CifarResNet(nn.Module):

    def __init__(self, block, layers, num_classes=10):
        super(CifarResNet, self).__init__()
        self.inplanes = 16
        self.conv1 = conv3x3(3, 16)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu = nn.ReLU(inplace=True)

        self.layer1 = self._make_layer(block, 16, layers[0])
        self.layer2 = self._make_layer(block, 32, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 64, layers[2], stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(64 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x, return_hidden=False):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        if return_hidden:
            return x, self.fc(x)
        x = self.fc(x)
        return x

def ResNet20(num_classes=10):
    model = CifarResNet(BasicBlock, [3, 3, 3], num_classes=num_classes)
    return model

def train_ResNet(model, train_loader, optimizer, criterion, device):
    model.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model.forward(inputs).squeeze()
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        predicted = torch.round(torch.sigmoid(outputs))
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    return train_loss / len(train_loader), correct / total

def test_ResNet(model, test_loader, criterion, device):
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(test_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model.forward(inputs).squeeze()
            loss = criterion(outputs, targets)

            test_loss += loss.item()
            predicted = torch.round(torch.sigmoid(outputs))
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

    return test_loss / len(test_loader), correct / total

def calculate_intermediate_output(model, data_iterator, device):
    model.eval()
    with torch.no_grad():
        intermediate = []
        pred = []
        labels = []
        outputs = []
        for batch_idx, (inputs, targets) in enumerate(data_iterator):
            inputs = inputs.to(device)
            targets = targets.to(device)
            hidden, output = model(inputs, return_hidden=True)
            outputs.append(torch.nn.functional.softmax(output, dim=1))
            intermediate.append(hidden)
            pred.append(torch.sigmoid(output))
            labels.append(targets)
        intermediate = torch.cat(intermediate).data.cpu().numpy()
        outputs = torch.cat(outputs).data.cpu().numpy()
        pred = torch.cat(pred).data.cpu().numpy()
        labels = torch.cat(labels).data.cpu().numpy()
        return intermediate, outputs, pred, labels

def test():
    model = ResNet20()
    print(sum(p.numel() for p in model.parameters() if p.requires_grad))


# https://morvanzhou.github.io/tutorials/data-manipulation/plt/2-4-axis2/
# https://github.com/MorvanZhou/tutorials/blob/master/matplotlibTUT/plt6_ax_setting2.py
import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(-3, 3, 50)
y1 = 2*x+1
y2 = x**2

plt.figure(num=3, figsize=(8, 5)) # 以下是這張figure的數據
plt.plot(x, y2)
plt.plot(x, y1, color='red', linewidth=1.0, linestyle='--') # 設定個別參數

plt.xlim((-1, 2)) # x show limit range
plt.ylim((-2, 3)) # y show limit range
plt.xlabel('I am x') # show x name
plt.ylabel('I am y') # show y name

# show new ticks
new_ticks = np.linspace(-1, 2, 5)
plt.xticks(new_ticks) # 替換x座標的範圍文字
plt.yticks([-2, -1.8, -1, 1.22, 3], [r'$really\ good$', r'$bad\ \alpha$', r'$normal$', r'$good$', r'$really\ good$']) # 替換y座標的範圍文字, r是正則表達式, $是變好看的字體, \alpha是數學符號

# gca = 'get current axis' 把現在的軸拿出來
ax = plt.gca()
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')

ax.xaxis.set_ticks_position('bottom') # 用哪個軸代替X軸
# ACCEPTS: [ 'top' | 'bottom' | 'both' | 'default' | 'none' ]
ax.spines['bottom'].set_position(('data', 0))
# the 1st is in 'outward' | 'axes' | 'data'
# axes: percentage of y axis
# data: depend on y data

ax.yaxis.set_ticks_position('left') # 用哪個軸代替Y軸
# ACCEPTS: [ 'left' | 'right' | 'both' | 'default' | 'none' ]
ax.spines['left'].set_position(('data',0))

plt.show()


import requests
from bs4 import BeautifulSoup
import json
import sys
from langdetect import detect
from flask import Flask, request, jsonify, render_template, url_for, redirect, flash, session
import openai

app = Flask(__name__)
app.secret_key = 'your_secret_key_here'

# OpenAI API Key
openai_apikey = "YOUR_OPENAI_API_KEY"
openai.api_key = openai_apikey

def read_webpage(url):
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3',
        'accpet': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
        # 'accept-encoding': 'gzip, deflate, br',
        'referer': 'https://www.google.com/',
        'accept-language': 'ko,en-US;q=0.9,en;q=0.8'
    }

    response = requests.get(url, headers=headers)

    if response.status_code == 200:
        html = response.text
        soup = BeautifulSoup(html, 'html.parser')

        # get the title
        title = soup.select_one('title').get_text().strip()

        # get the language
        meta_tag = soup.find('meta', {'charset': True})
        charset = meta_tag['charset']

        # get the body
        body = soup.select_one('body')

        # list the first level children of the body
        children = body.findChildren()
        content = []
        text_content = ''
        for child in children:
            if child.name in ['h1', 'h2', 'h3', 'h4', 'p']:
                content.append({
                    'type': child.name,
                    'text': child.get_text().strip()
                })
                text_content += child.get_text().strip() + '\n'

        language = detect(text_content)

        print("Charset: ", charset)
        print("Language: ", language)

        return {
            'result': 'success',
            'language': language,
            'title': title,
            'content': content
        }
    else:
        return {
            'result': 'error',
            'status_code': response.status_code,
            'err_msg': response.text
        }

def complete(messages):
    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=messages,
            temperature=0.1,
            max_tokens=1024,
            top_p=1,
            frequency_penalty=0.5,
            presence_penalty=0.5,
            stop=None
        )

        return response['choices'][0]['message']['content']
    except Exception as e:
        print(e)
        return 'Error'

def count_sentences(text):
    length = 0
    for line in text.split('\n'):
        if line.strip() != '':
            for str in line.split('.'):
                if str.strip() != '':
                    length += 1
    return length

def summarize_content(title, content, language='en'):
    MIN_SECTIONS = 5
    SUM_RATIO_KOR = 5
    SUM_RATIO_ENG = 5
    SUM_RATIO_SUMMARY = 3
    MAX_CHUNK_SIZE = 500

    strcuture = {
        'titles': [{'title': title}],
        'sections': []
    }

    section_list = {
        'h1': [],
        'h2': [],
        'h3': [],
        'h4': []
    }
    for item in content:
        if item['type'] in ['h1', 'h2', 'h3', 'h4']:
            section_list[item['type']].append(item.copy())

    section_tag = None
    if len(section_list['h1']) >= MIN_SECTIONS:
        strcuture['sections'] = section_list['h1']
        section_tag = 'h1'
    elif len(section_list['h2']) >= MIN_SECTIONS:
        strcuture['sections'] = section_list['h2']
        strcuture['titles'].extend(section_list['h1'])
        section_tag = 'h2'
    elif len(section_list['h3']) >= MIN_SECTIONS:
        strcuture['sections'] = section_list['h3']
        strcuture['titles'].extend(section_list['h1'])
        strcuture['titles'].extend(section_list['h2'])
        section_tag = 'h3'
    elif len(section_list['h4']) >= MIN_SECTIONS:
        strcuture['sections'] = section_list['h4']
        strcuture['titles'].extend(section_list['h1'])
        strcuture['titles'].extend(section_list['h2'])
        strcuture['titles'].extend(section_list['h3'])
        section_tag = 'h4'
    else:
        if len(section_list['h1']) > 1:
            strcuture['sections'] = section_list['h1']
            section_tag = 'h1'
        elif len(section_list['h2']) > 1:
            strcuture['sections'] = section_list['h2']
            strcuture['titles'].extend(section_list['h1'])
            section_tag = 'h2'
        elif len(section_list['h3']) > 1:
            strcuture['sections'] = section_list['h3']
            strcuture['titles'].extend(section_list['h1'])
            strcuture['titles'].extend(section_list['h2'])
            section_tag = 'h3'
        elif len(section_list['h4']) > 1:
            strcuture['sections'] = section_list['h4']
            strcuture['titles'].extend(section_list['h1'])
            strcuture['titles'].extend(section_list['h2'])
            strcuture['titles'].extend(section_list['h3'])
            section_tag = 'h4'

    messages = [
        {'role': 'system', 'content': 'You are a helpful assistant'}
    ]

    if language == 'ko':
        message = '글의 구조는 title > h1 > h2 > h3 > p 의 순서이다. 하나의 문장으로 제목을 추출하라:\n'
    else:
        message = 'While strcuture is title > h1 > h2 > h3 > h4 > p, please get a title as a sentence from:\n'
    for title in strcuture['titles']:
        if 'title' in title:
            message += "title: {}\n".format(title['title'])
        else:
            message += "{}: {}\n".format(title['type'], title['text'])

    messages.append({'role': 'user', 'content': message})
    strcuture['messages'] = messages
    title = complete(messages)
    # if title contains a quotation, extract it
    if title.find('"') != -1:
        # Extract a sentence quoted by "
        import re
        title = re.findall(r'\"(.+?)\"', title)[0]
        strcuture['title'] = title
    else:
        strcuture['title'] = title.split('\n')[0]

    # return {'structure': json.dumps(strcuture, indent=4)}

    print("Found title: {}".format(strcuture['title']))

    if section_tag is not None:
        current_section = -1
        for item in content:
            for section in strcuture['sections']:
                if item['type'] == section['type']:
                    if item['text'] == section['text']:
                        current_section += 1
                        strcuture['sections'][current_section]['content'] = []
                        break

            if current_section == -1:
                strcuture['sections'].insert(0, {
                    'type': 'preface',
                    'text': '',
                    'content': []
                })
                current_section = 0

            strcuture['sections'][current_section]['content'].append(item.copy())

        print("Found {} sections".format(len(strcuture['sections'])))

        # Firstly, try to summarize all the sections at once
        messages = [
            {'role': 'system', 'content': 'You are a helpful assistant'}
        ]

        if language == 'ko':
            message = '글의 구조는 title > h1 > h2 > h3 > p 의 순서이다\n\n'
        else:
            message = "Text strcuture is title > h1 > h2 > h3 > h4 > p.\n\n"
        message += "title: {}\n".format(strcuture['title'])
        chunk = ''
        for section in strcuture['sections']:
            if section['type'] == 'preface':
                for item in section['content']:
                    if item['type'] == 'p':
                        chunk += "{}: {}\n".format(item['type'], item['text'])
            else:
                chunk += "{}: {}\n".format(section['type'], section['text'])
                for item in section['content']:
                    chunk += "{}: {}\n".format(item['type'], item['text'])
        # Count the number of sentences of message
        length = count_sentences(chunk)

        if language == 'ko':
            length = int(length / SUM_RATIO_KOR / SUM_RATIO_SUMMARY) + 1
            command = "\nPlease summarize the text as a {} sentence long paragraph in Korean:\n".format(length)
        else:
            length = int(length / SUM_RATIO_ENG / SUM_RATIO_SUMMARY) + 1
            command = "\nPlease summarize the text as a {} sentence long paragraph:\n".format(length)

        print("Trying to summarize all sections as {} sentences".format(length))

        messages.append({'role': 'user', 'content': "{}{}{}".format(message, chunk, command)})
        strcuture['messages'] = messages

        # Write structure as json text into a file
        with open('structure.json', 'w') as f:
            f.write(json.dumps(strcuture, indent=4))

        print("Before: {}\n".format(messages[1]['content']))
        summary = complete(messages)
        if summary == 'Error':
            # Secondly, try to summarize each section separately
            for section in strcuture['sections']:
                messages = [
                    {'role': 'system', 'content': 'You are a helpful assistant'}
                ]
                if language == 'ko':
                    message = '글의 구조는 title > h1 > h2 > h3 > p 의 순서이다.\n\n'
                else:
                    message = "Text strcuture is title > h1 > h2 > h3 > h4 > p.\n\n"
                # message += "please summarize the content as maximum {} sentences paragraph from:\n".format(MAX_LINES)
                message += "title: {}\n".format(strcuture['title'])

                chunk = ''
                if section['type'] == 'preface':
                    chunk += "{}: {}\n".format(section_tag, section['text'])
                    for item in section['content']:
                        if item['type'] == 'p':
                            chunk += "{}: {}\n".format(item['type'], item['text'])
                else:
                    chunk += "{}: {}\n".format(section['type'], section['text'])
                    for item in section['content']:
                        chunk += "{}: {}\n".format(item['type'], item['text'])

                # Count the number of sentences of message
                length = count_sentences(chunk)

                if language == 'ko':
                    length = int(length / SUM_RATIO_KOR) + 1
                    command = "\nPlease summarize the content as maximum {} bullets in Korean:\n\n".format(length)
                else:
                    length = int(length / SUM_RATIO_ENG) + 1
                    command = "\nPlease summarize the content as maximum {} bullets:\n\n".format(length)

                print("Trying to summarize section as {} sentences".format(length))
                print('section: {}'.format(section['text']))

                messages.append({'role': 'user', 'content': "{}{}{}".format(message, chunk, command)})
                section['messages'] = messages

                summary = complete(messages)
                if summary == 'Error':
                    # Thirdly, try to summarize each chunk of section separately
                    chunk_array = chunk.split('\n')
                    chunk_index = 0
                    chunk_size = 0
                    chunk_segment = ''
                    previous_summary = ''
                    for chunk_line in chunk_array:
                        chunk_size += len(chunk_line.split(' '))
                        if chunk_size > MAX_CHUNK_SIZE:
                            chunk_size = 0
                            chunk_index += 1

                            messages = [
                                {'role': 'system', 'content': 'You are a helpful assistant'},
                                {'role': 'user', 'content': "Previous summary is:\n{}\n".format(previous_summary)}
                            ]

                            # Count the number of sentences of message
                            length = count_sentences("{}\n{}".format(previous_summary, chunk_segment))

                            if language == 'ko':
                                length = int(length / SUM_RATIO_KOR) + 1
                                command = "\nPlease summarize the content with previous summary as maximum {} bullets in Korean:\n\n".format(length)
                            else:
                                length = int(length / SUM_RATIO_ENG) + 1
                                command = "\nPlease summarize the content with previous summary as maximum {} bullets:\n\n".format(length)

                            print("Trying to summarize each chunk of section as {} sentences".format(length))
                            print('section: {}'.format(section['text']))

                            messages.append({'role': 'user', 'content': "{}{}{}".format(message, chunk_segment, command)})
                            summary = complete(messages)
                            print("previous summary: {}".format(previous_summary))
                            print("chunk summary: {}".format(summary))
                            previous_summary = summary
                            chunk_segment = ''

                        chunk_segment += chunk_line + '\n'

                    if chunk_segment != '':
                        messages = [
                            {'role': 'system', 'content': 'You are a helpful assistant'},
                            {'role': 'user', 'content': "Previous summary is:\n{}\n".format(previous_summary)}
                        ]

                        if language == 'ko':
                            command = "\nPlease summarize the content with previous summary as maximum {} bullets in Korean:\n\n".format(length)
                        else:
                            command = "\nPlease summarize the content with previous summary as maximum {} bullets:\n\n".format(length)

                        print("Trying to summarize each chunk of section as {} sentences".format(length))
                        print('section: {}'.format(section['text']))

                        messages.append({'role': 'user', 'content': "{}{}{}".format(message, chunk_segment, command)})
                        summary = complete(messages)
                        print("previous summary: {}".format(previous_summary))
                        print("chunk summary: {}".format(summary))

                section['summary'] = summary

                print('section: {}'.format(section['text']))
                print("summary: {}".format(summary))

            # Finally summarize all the sections
            messages = [
                {'role': 'system', 'content': 'You are a helpful assistant'}
            ]
            if language == 'ko':
                message = 'These are title and sections.\n\n'
            else:
                message = "These are title and sections.\n\n"
            message += "Title: {}\n".format(strcuture['title'])
            chunk = ''
            for section in strcuture['sections']:
                if section['type'] == 'preface':
                    chunk += "Section: {}\n".format(section['text'])
                    chunk += "{}\n".format(section['summary'])
                else:
                    chunk += "Section: {}\n".format(section['text'])
                    chunk += "{}\n".format(section['summary'])
            # Count the number of sentences of message
            length = count_sentences(chunk)

            if language == 'ko':
                length = int(length / SUM_RATIO_SUMMARY) + 1        # Korean summarization ratio is too high, so apply English ratio
                command = "\nPlease summarize all sections as a {} sentence long paragraph in Korean:\n".format(length)
            else:
                length = int(length / SUM_RATIO_SUMMARY) + 1
                command = "\nPlease summarize all sections as a {} sentence long paragraph:\n".format(length)

            print("Trying to summarize section as {} sentences".format(length))

            messages.append({'role': 'user', 'content': "{}{}{}".format(message, chunk, command)})
            strcuture['messages'] = messages

            print("Before: {}\n".format(messages[1]['content']))
            summary = complete(messages)

        print("Final summary: {}".format(summary))
        strcuture['summary'] = summary

    else:
        messages = [
            {'role': 'system', 'content': 'You are a helpful assistant'}
        ]
        if language == 'ko':
            message = '글의 구조는 title > h1 > h2 > h3 > p 의 순서이다.\n\n'
        else:
            message = "Text strcuture is title > h1 > h2 > h3 > h4 > p.\n\n"
        # message += "please summarize the content as maximum {} sentences paragraph from:\n".format(MAX_LINES)
        message += "title: {}\n".format(strcuture['title'])

        chunk = ''
        for item in content:
            chunk += "{}: {}\n".format(item['type'], item['text'])

        # Count the number of sentences of message
        length = count_sentences(chunk)

        if language == 'ko':
            length = int(length / SUM_RATIO_KOR / SUM_RATIO_SUMMARY) + 1
            command = "\nPlease summarize the text as a {} sentence long paragraph in Korean:\n".format(length)
        else:
            length = int(length / SUM_RATIO_ENG / SUM_RATIO_SUMMARY) + 1
            command = "\nPlease summarize the text as a {} sentence long paragraph:\n".format(length)

        print("Trying to summarize all as {} sentences without section".format(length))

        messages.append({'role': 'user', 'content': "{}{}{}".format(message, chunk, command)})
        strcuture['messages'] = messages

        # Write structure as json text into a file
        with open('structure.json', 'w') as f:
            f.write(json.dumps(strcuture, indent=4))

        print("Before: {}\n".format(messages[1]['content']))
        summary = complete(messages)
        print("Final summary: {}".format(summary))
        strcuture['summary'] = summary

    return {
        'title': strcuture['title'],
        'summary': strcuture['summary']
    }

    return {'structure': json.dumps(strcuture, indent=4)}

@app.route('/')
def index():
    if 'url' not in session:
        session['url'] = 'https://edition.cnn.com/2023/04/15/americas/darien-gap-migrants-colombia-panama-whole-story-cmd-intl/index.html'
    return render_template('index.html', url=session.get('url'))

@app.route('/summarize', methods=['POST'])
def summarize():
    url = request.form['url']
    session['url'] = url
    print(url)
    result = read_webpage(url)
    if result['result'] == 'success':
        summary = summarize_content(result['title'], result['content'], result['language'])
        if 'title' in summary:
            return render_template('index.html', url=session.get('url'), title=summary['title'], content=summary['summary'])
        else:
            return render_template('index.html', url=session.get('url'), structure=summary['structure'])
        # return render_template('index.html', title=result['title'], content=json.dumps(result['content'], indent=4))
    else:
        return render_template('index.html', url=session.get('url'), title=result['status_code'], content=result['err_msg'])

if len(sys.argv) > 1:
    # get the url from the command line
    url = sys.argv[1]
    print(url)

    # read the webpage
    result = read_webpage(url)
    if result['result'] == 'success':
        print(result['title'])
        print(json.dumps(result['content'], indent=4))
    else:
        print(result['status_code'])
        print(result['content'])

    sys.exit(0)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)


# -*- coding: utf-8 -*-

# Run this app with `python app.py` and
# visit http://127.0.0.1:8050/ in your web browser.

import dash
import dash_core_components as dcc
import dash_html_components as html
import plotly.express as px
import pandas as pd
import numpy as np
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.graph_objects as go
from dash.dependencies import Input, Output
from jupyter_dash import JupyterDash

def plotConfiguration(data):
        data["particle"]=data["particle"].astype(str)
        graphs=[]
        for particle,particleData in data.groupby("particle"):
            graphs.append(go.Scatter3d( x=particleData['x'], y=particleData['y'], z=particleData['z'] ,name=str(particle)))
        fig=go.Figure(data=graphs)
        return fig

class configurationsApp:

    def __init__(self,filename):
        external_stylesheets = ['https://codepen.io/chriddyp/pen/bWLwgP.css']

        self.app = JupyterDash(__name__, external_stylesheets=external_stylesheets)
        fig=go.Figure(data=[])

        try:
            conf=pd.read_csv(filename,sep=" ",index_col=False)
            fig=plotConfiguration(conf)

        except FileNotFoundError:
            print ("File not found")

        self.app.layout = html.Div(children=[
            html.H1(children='Visualization of PIMC simulation'),

            html.H2(children='''
                Configurations
            '''),

            dcc.Graph(
                id='last-configuration',
                figure=fig
            )
            ])


import numpy as np
import scipy.spatial.distance as dist

from sklearn.neighbors import KDTree

# TODO
# - Vary alpha in epsilon in the qlearner to encourage more exploration early
#
# Notes
# - The state should at least include alignment.
# - The state may also include a neighbor density function
# - Without the neighbor density function, the flock tends
#   to oscillate more;  they appear to not have a memory
#   of density.  IOW there is not positively-rewarded state
#   of "nice density" so they are constantly searching for it.

class Flock(object):
    """Object representing a flock of birds in two dimensions."""

    num_turns = 8
    num_bins_alignment = 8
    num_bins_neighbors = 4

    def __init__(self, n_birds, speed, time_step,
                 bounding_box=None,
                 perceptual_radius=2,
                 birth_box=None,
                 random_seed=42):
        self.n_birds = n_birds
        self.speed = speed
        self.delta_t = time_step
        self.bbox = bounding_box
        self.birth_box = self.bbox if birth_box is None else birth_box
        self.perceptual_radius = perceptual_radius
        self.seed = random_seed

        self.turns = np.linspace(-np.pi/6, np.pi/6, self.num_turns)

        self.TOL = 1e-6

        self.position, self.direction = self.initialize()
        self.neighbors_in_range = None

        self.n_states = self.num_bins_alignment*self.num_bins_neighbors
        self.state_bin_boundaries_neigh = np.linspace(-1-self.TOL, 1+self.TOL, self.num_bins_neighbors)
        self.state_bin_boundaries_align = np.linspace(-1-self.TOL, 1+self.TOL, self.num_bins_alignment)

        self.ideal_neighbors = 20

    def initialize(self):
        """Randomly initialize positions and directions of the flock."""
        np.random.seed(self.seed)

        position_x = np.random.uniform(self.birth_box[0], self.birth_box[1] ,self.n_birds)
        position_y = np.random.uniform(self.birth_box[0], self.birth_box[1] ,self.n_birds)
        direction_theta = np.random.uniform(0.0, 2*np.pi, self.n_birds)

        position = np.vstack((position_x, position_y)).T
        direction = np.vstack((np.cos(direction_theta), np.sin(direction_theta))).T

        return position, direction

    def get_neighbors(self):
        kdt = KDTree(self.position)
        neighbors, distances = kdt.query_radius(self.position,
                                                r=self.perceptual_radius,
                                                return_distance=True,
                                                count_only=False)

        self.neighbors_in_range = []
        self.neighbor_distances_in_range = []
        for b in range(self.n_birds):
            nbrs, dists = neighbors[b], distances[b]

            rel_pos = self.position[nbrs,:]-self.position[b,:]
            in_perception = np.dot(rel_pos, self.direction[b,:])
            inds = np.where((in_perception>0)&(dists>0))[0]
            self.neighbors_in_range.append(nbrs[inds])
            self.neighbor_distances_in_range.append(dists[inds])

        return self.neighbors_in_range

    def compute_center_of_mass(self):
        com = np.zeros_like(self.position)
        for b in range(self.n_birds):
            nbr_inds = self.neighbors_in_range[b]
            if nbr_inds.shape[0]>0:
                com[b,:] = np.mean(self.position[nbr_inds,:], axis=0)
            else:
                com[b,:] = np.inf

        return com

    def compute_average_velocity(self):
        avg_vel = np.zeros_like(self.direction)
        for b in range(self.n_birds):
            nbr_inds = self.neighbors_in_range[b]
            if nbr_inds.shape[0]>0:
                avg_vel[b,:] = np.mean(self.direction[nbr_inds,:], axis=0)
            else:
                avg_vel[b,:] = np.inf

        return avg_vel

    def update(self, turn_angles=None, delta_v=None, apply_bcs=True):
        """Update the positions and directions of birds based on a set of turn angles."""
        if turn_angles is not None:
            direction_new_x = self.direction[:,0]*np.cos(turn_angles) - self.direction[:,1]*np.sin(turn_angles)
            direction_new_y = self.direction[:,0]*np.sin(turn_angles) + self.direction[:,1]*np.cos(turn_angles)
            self.direction = np.vstack((direction_new_x, direction_new_y)).T
        else:
            self.direction += delta_v
            norm = np.linalg.norm(self.direction, axis=1)
            for i in range(self.direction.shape[1]):
                self.direction[:,i] /= norm

        self.position += self.speed*self.direction*self.delta_t
        if apply_bcs:
            self._apply_boundary_conditions()

    def _apply_boundary_conditions(self):
        """Apply periodic boundary conditions on the flock."""
        boundary_conds = {
            'left':{
                'dim':0,
                'periodic_boundary':1,
                'dist_to_boundary':self.position[:,0]-self.bbox[0],
                'on_boundary':(self.position[:,0]-self.bbox[0])<self.TOL,
                'exiting':np.dot(self.direction, np.array([-1,0])) > self.TOL
            },
            'right':{
                'dim':0,
                'periodic_boundary':0,
                'dist_to_boundary':self.bbox[1]-self.position[:,0],
                'on_boundary':(self.bbox[1]-self.position[:,0])<self.TOL,
                'exiting':np.dot(self.direction, np.array([1,0])) > self.TOL
            },
            'up':{
                'dim':1,
                'periodic_boundary':0,
                'dist_to_boundary':self.bbox[1]-self.position[:,1],
                'on_boundary':(self.bbox[1]-self.position[:,1])<self.TOL,
                'exiting':np.dot(self.direction, np.array([0,1])) > self.TOL
            },
            'down':{
                'dim':1,
                'periodic_boundary':1,
                'dist_to_boundary':self.position[:,1]-self.bbox[0],
                'on_boundary':(self.position[:,1]-self.bbox[0])<self.TOL,
                'exiting':np.dot(self.direction, np.array([0,-1])) > self.TOL
            }
        }
        for _, condition in boundary_conds.items():
            idx = np.where(condition['on_boundary'] & condition['exiting'])[0]
            multi = 1 if condition['periodic_boundary']==1 else -1
            self.position[idx,condition['dim']] = self.bbox[condition['periodic_boundary']] + multi*condition['dist_to_boundary'][idx]

    def get_state(self):
        """Return the state of each bird in the flock.

        This is the state used by the qlearner."""
        self.get_neighbors()

        cohesion = np.zeros(self.n_birds)
        alignment = np.zeros(self.n_birds)
        n_neighbors = np.zeros(self.n_birds, dtype=int)

        for i in range(self.n_birds):
            nbrs = self.neighbors_in_range[i]
            n_neighbors[i] = nbrs.shape[0]
            if nbrs.shape[0]==0:
                cohesion[i] = 0.0
                continue

            avg_direction = self.direction[nbrs,:].mean(axis=0)
            avg_direction = avg_direction/np.linalg.norm(avg_direction)

            cohesion[i] = np.dot(avg_direction, self.direction[i,:])
            alignment[i] = np.cross(avg_direction, self.direction[i,:])

        state_n = 1/(1 + np.exp(self.ideal_neighbors-n_neighbors))
        state_n = np.digitize(state_n, self.state_bin_boundaries_neigh)-1
        state_a = np.digitize(alignment, self.state_bin_boundaries_align)-1
        state = state_a + self.num_bins_alignment*state_n

        assert np.all(state_n<self.num_bins_neighbors)
        assert np.all(state_a<self.num_bins_alignment)

        return state, cohesion, alignment, n_neighbors

    def get_reward(self, cohesion, n_neighbors):
        r_n = 1 - np.abs(self.ideal_neighbors-n_neighbors)/(n_neighbors+1)
        r_c = cohesion

        return (r_n+r_c)/2

    def plot(self, figsize=(5,5)):
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=figsize)
        ax.quiver(self.position[:,0], self.position[:,1], self.direction[:,0], self.direction[:,1])
        ax.axis('equal')
        ax.set_xlim(self.bbox)


import asyncio
import datetime
import json
from functools import partial
from typing import Any, Dict, List, Optional, Union

import pandas
import pytz

from lyra.core.config import cfg
from lyra.core.errors import HydstraIOError
from lyra.core.io import load_file
from lyra.core.utils import local_path
from lyra.src.hydstra import helper
from lyra.src.mnwd.helper import get_timeseries_from_dt_metrics
from lyra.src.timeseries import utils

AGG_REMAP = {
    # hydstra : pandas
    "tot": "sum",
    "cum": "cumsum",
    "mean": "mean",
    "max": "max",
    "min": "min",
}

INTERVAL_REMAP = {
    "year": "YS",
    "month": "MS",
    "day": "D",
    "hour": "H",
}

SWN_SITES_PATH = local_path("data/mount/swn/hydstra").resolve() / "swn_sites.json"

class Timeseries(object):
    def __init__(
        self,
        site: str,
        variable: str,
        aggregation_method: str,
        interval: str = None,
        start_date: str = None,  # "yyyy-mm-dd"
        end_date: str = None,
        trace_upstream: bool = None,
        weather_condition: str = None,
        hydstra_kwargs: Optional[Dict] = None,
        warnings: Optional[List[Any]] = None,
        nearest_rainfall_station: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        self.site = site
        self.variable = variable
        self.start_date = start_date or "2020-01-01"
        self.end_date = (
            end_date
            or datetime.datetime.now(pytz.timezone("US/Pacific")).date().isoformat()
        )
        self.interval = interval or "month"
        self.aggregation_method = aggregation_method
        self.trace_upstream = True if trace_upstream is None else trace_upstream
        self.weather_condition = weather_condition or "both"

        self.cfg = cfg
        self.variable_info = self.cfg["variables"].get(self.variable, {})
        allowed_aggs = self.variable_info["allowed_aggregations"]
        self.aggregation_method = (
            aggregation_method
            if aggregation_method in allowed_aggs
            else allowed_aggs[0]
        )
        self.warnings = warnings or []

        # load_file will cache this so it doesn't happen for frequent requests
        self.all_props: List[Any] = [
            f["properties"]
            for f in json.loads(load_file(self.cfg["site_path"]))["features"]
        ]

        self.site_props: Dict[str, Any] = next(
            (x for x in self.all_props if x["station"] == site), {}
        )

        self.nearest_rainfall_station = nearest_rainfall_station or self.site_props.get(
            "nearest_rainfall_station", "not_set"
        )

        self.hydstra_kwargs: Dict[str, Any] = hydstra_kwargs or {}
        self.kwargs: Dict[str, Any] = kwargs

        # properties
        self._timeseries = None
        self._timeseries_src = None
        self._nearest_rainfall_station_props: Optional[Dict[str, Any]] = None
        self._nearest_rainfall_ts: Optional[Timeseries] = None
        self._weather_condition_series = None
        self.label = self.__repr__()

    def __repr__(self):

        is_dt_metric = self.variable_info["source"] == "dt_metrics"
        _var_name = self.variable_info["name"]
        _var_units = self.variable_info["units"]
        _us = "Upstream from" if self.trace_upstream and is_dt_metric else "from"
        _method = self.aggregation_method.title() if self.aggregation_method else ""
        _weather_condition = (
            f"{self.weather_condition.title()} Weather "
            if self.weather_condition != "both"
            else ""
        )

        repr = (
            f"{_method} 1 {self.interval} {_weather_condition}{_var_name} ({_var_units}) {_us} "
            f"{self.site} from {self.start_date} to {self.end_date}"
        )

        return repr

    @property
    def timeseries(self) -> pandas.Series:
        if self._timeseries is None:  # pragma: no branch
            asyncio.run(self.init_ts())
        return self._timeseries

    @timeseries.setter
    def timeseries(self, timeseries):  # pragma: no cover
        self._timeseries = timeseries

    @property
    def timeseries_src(self) -> pandas.DataFrame:
        if self._timeseries_src is None:  # pragma: no branch
            self._timeseries_src = self.timeseries.assign(site=self.site).assign(
                variable=self.variable
            )
        return self._timeseries_src

    @property
    def nearest_rainfall_station_props(self):  # pragma: no cover
        if self._nearest_rainfall_station_props is None:
            self._nearest_rainfall_station_props = next(
                (
                    x
                    for x in self.all_props
                    if x["station"] == self.nearest_rainfall_station
                ),
                {},
            )
        return self._nearest_rainfall_station_props

    def get_nearest_rainfall_ts(self):
        if self._nearest_rainfall_ts is None:
            _ = asyncio.run(self.get_nearest_rainfall_ts_async())

        return self._nearest_rainfall_ts

    async def get_nearest_rainfall_ts_async(self):
        if self._nearest_rainfall_ts is None:

            rainfall_inputs = dict(
                site=self.nearest_rainfall_station,
                variable="rainfall",
                start_date=self.start_date,
                end_date=self.end_date,
                interval="hour",
                aggregation_method="tot",
            )

            self._nearest_rainfall_ts = Timeseries(**rainfall_inputs)  # type: ignore
            await self._nearest_rainfall_ts.init_ts()

        return self._nearest_rainfall_ts

    @property
    def weather_condition_series(self):
        if all(
            [
                self.variable == "rainfall",
                self.interval == "hour",
                self._weather_condition_series is None,
            ]
        ):
            self._weather_condition_series = utils.identify_dry_weather(
                self.timeseries.value
            )

        return self._weather_condition_series

    async def _init_hydstra(self) -> pandas.Series:
        hyd_variable_info: Dict[str, Any] = self.site_props.get(
            self.variable + "_info", {}
        )

        interval = self.interval

        if self.weather_condition != "both":
            interval = "hour"

        inputs = dict(
            site=self.site,
            varfrom=hyd_variable_info["varfrom"],
            varto=hyd_variable_info["varto"],
            start_date=self.start_date,
            end_date=self.end_date,
            interval=interval,
            agg_method=self.aggregation_method,
            **self.hydstra_kwargs,
        )

        async def process_hydstra_errors(**inputs: Dict) -> Any:
            """recursive helper function to try to get hydstra data and
            handle errors when they occur

            uses globals: self

            """
            timeseries_details = await helper.get_site_variable_as_trace(**inputs)

            if timeseries_details is None:
                self.warnings.append(
                    "Hydstra Database is temporarily unavailable. See "
                    "http://hydstra.ocpublicworks.com/web.htm for most recent data available."
                )
                return "empty"

            if "error_msg" in timeseries_details.keys():
                num = timeseries_details["error_num"]
                if num == 220:
                    raise HydstraIOError(
                        f"{timeseries_details['error_msg']}", data=inputs
                    )

                if num == 124:
                    self.warnings.append(
                        f'Warning: No conversion or rating table for "{hyd_variable_info["varfrom"]}" '
                        f'to "{hyd_variable_info["varfrom"]}".'
                    )
                    return "empty"

                if num == 126:
                    self.warnings.append(
                        f'Warning: No data within requested period for variable "{self.variable}" '
                        f'at site "{self.site}".'
                    )
                    return "empty"

                if num == 125:

                    inputs["varfrom"] = hyd_variable_info["varfrom_fallback"]
                    self.warnings.append(
                        f"Warning: Variable '{hyd_variable_info['varfrom']}' not available. "
                        f"Falling back to '{hyd_variable_info['varfrom_fallback']}'."
                    )

                    # retry
                    return await process_hydstra_errors(**inputs)

            return timeseries_details

        timeseries_details = await process_hydstra_errors(**inputs)

        if timeseries_details == "empty":
            return pandas.DataFrame([])

        trace = timeseries_details.get("trace")

        if not trace:  # pragma: no cover
            raise ValueError(f"inputs failed: {inputs}, {timeseries_details}")

        hydstra_result = helper.hydstra_trace_to_series(trace)

        quality = self.cfg["hydstra"]["max_quality_flag"]
        questionable_data = hydstra_result.query("q > @quality")
        if not questionable_data.empty:
            ds, de = (
                questionable_data.index.min().date(),
                questionable_data.index.max().date(),
            )
            flags = questionable_data.q.unique()

            self.warnings.append(
                f"{len(questionable_data)} data points have Hydstra data quality "
                f"flags for {self.variable} at {self.site} between {ds} and {de}. "
                f"Quality flags include: {', '.join((str(i) for i in flags))}."
            )

        # this will block the async await... but so it goes.
        return await self.process_weather_condition(hydstra_result)

    async def process_weather_condition(self, df):
        if self.weather_condition == "both":
            return df

        if self.weather_condition == "dry":
            q = "is_dry"
        else:
            q = "~is_dry"

        nearest_rain_ts = await self.get_nearest_rainfall_ts_async()

        rain_start_date = nearest_rain_ts.timeseries.index.min()
        rain_end_date = nearest_rain_ts.timeseries.index.max()

        if rain_start_date > df.index.min() or rain_end_date > df.index.max():
            df = df.loc[(df.index >= rain_start_date) & (df.index <= rain_end_date)]

            self.warnings.append(
                f"Warning: Setting date range to start_date={df.index.min().isoformat()} and end_date={df.index.max().isoformat()} "
                f'at site "{self.site}" to support determination of weather_condition = {self.weather_condition}.'
            )

        result = (
            df.join(nearest_rain_ts.weather_condition_series["is_dry"], how="left")
            .fillna({"is_dry": True})
            .query(q)[["value"]]
            .resample(INTERVAL_REMAP[self.interval])
            .aggregate(AGG_REMAP[self.aggregation_method])
        )

        return result

    async def _init_dt_metrics(self) -> pandas.Series:
        variable = self.variable_info["varfrom"]
        catchidn = self.site_props["CatchIDN"]
        self.aggregation_method = "tot"
        self.interval = (
            "month" if self.interval not in ["year", "month"] else self.interval
        )

        inputs = dict(
            variable=variable,
            site=catchidn,
            start_date=self.start_date,
            end_date=self.end_date,
            agg_method=AGG_REMAP[self.aggregation_method],
            trace_upstream=self.trace_upstream,
            interval=self.interval,
            **self.kwargs,
        )

        loop = asyncio.get_running_loop()
        f = partial(get_timeseries_from_dt_metrics, **inputs)

        timeseries = await loop.run_in_executor(None, f)

        return timeseries

    async def init_ts(self, delay: Optional[float] = None) -> pandas.Series:
        source = self.variable_info.get("source")

        if source == "hydstra":
            if delay:
                # need to rate limit hydstra to 1 request per second, unfortunately.
                # this issue was reported to Hydstra support on 2021-07-12.
                # per communication 20201-07-19 this issue has been resolved.
                await asyncio.sleep(delay)

            self._timeseries = await self._init_hydstra()

        elif source == "dt_metrics":
            self._timeseries = await self._init_dt_metrics()

        else:  # pragma: no cover
            raise NotImplementedError(
                f"no method for loading {self.variable} from source: {source}. See config file."
            )

async def gather_timeseries(ts):

    await asyncio.gather(*(t.init_ts(delay=0) for i, t in enumerate(ts)))