#import

from pyimagesearch.transform import four_point_transform
from skimage.filters import threshold_local
import numpy as numpy
import argparse
import cv2
import imutils

# construct the arg parser
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required = True,
	help = "Path to the image to be scanned")
args = vars(ap.parse_args())
# load the image and compute the ratio of the old height
# to the new height, clone it, and resize it
image = cv2.imread(args["image"])
ratio = image.shape[0] / 500.0
orig = image.copy()
image = imutils.resize(image, height = 500)

refPt = []
cropping = False
def click_and_crop(event, x, y, flags, param):
	# grab references to the global variables
	global refPt, cropping
	# if the left mouse button was clicked, record the starting
	# (x, y) coordinates and indicate that cropping is being
	# performed
	if event == cv2.EVENT_LBUTTONDOWN:
		refPt = [(x, y)]
		cropping = True
	# check to see if the left mouse button was released
	elif event == cv2.EVENT_LBUTTONUP:
		# record the ending (x, y) coordinates and indicate that
		# the cropping operation is finished
		refPt.append((x, y))
		cropping = False
		# draw a rectangle around the region of interest
		cv2.rectangle(image, refPt[0], refPt[1], (0, 255, 0), 2)
		cv2.imshow("image", image)

# convert the image to grayscale, blur it, and find edges
# in the image
#print("image shape", image.shape[1])

y_end = image.shape[0]
x_end = image.shape[1]

gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
gray = cv2.GaussianBlur(gray, (5, 5), 0)
edged = cv2.Canny(gray, 75, 200)

# show the original image and the edge detected image
print("STEP 1: Edge Detection")
#cv2.imshow("Image", image)
#cv2.imshow("Edged", edged)
#cv2.waitKey(0)
#cv2.destroyAllWindows()

# find the contours in the edged image, keeping only the
# largest ones, and initialize the screen contour
cnts = cv2.findContours(edged.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
cnts = imutils.grab_contours(cnts)
cnts = sorted(cnts, key = cv2.contourArea, reverse = True)[:5]

flag = False
# loop over the contours
for c in cnts:
	# approximate the contour
	peri = cv2.arcLength(c, True)
	approx = cv2.approxPolyDP(c, 0.02 * peri, True)
	# if our approximated contour has four points, then we
	# can assume that we have found our screen
	if len(approx) == 4:
		screenCnt = approx
		print("flag yes")
		flag = True
		break

if flag == False :
	screenCnt = numpy.array([[[0,0]] ,[[x_end,0]] ,[[x_end, y_end]] ,[[0 ,y_end]]])
# show the contour (outline) of the piece of paper
print("STEP 2: Find contours of paper")
cv2.drawContours(image, [screenCnt], -1, (0, 255, 0), 2)
#cv2.imshow("Outline", image)
#cv2.waitKey(0)
#cv2.destroyAllWindows()

check_if_cropped = False

clone = image.copy()
cv2.namedWindow("image")
cv2.setMouseCallback("image", click_and_crop)
# keep looping until the 'q' key is pressed
while True:
	# display the image and wait for a keypress
	cv2.imshow("image", image)
	key = cv2.waitKey(1) & 0xFF
	# if the 'r' key is pressed, reset the cropping region
	if key == ord("r"):
		image = clone.copy()
	# if the 'c' key is pressed, break from the loop
	elif key == ord("c"):
		check_if_cropped = True
		if len(refPt) == 2:
			roi = clone[refPt[0][1]:refPt[1][1], refPt[0][0]:refPt[1][0]]
			cv2.imshow("ROI", roi)
			y_end_cr = roi.shape[0]
			x_end_cr = roi.shape[1]
			cv2.waitKey(0)
		break
	elif key == ord("s"):
		roi = clone.copy()
		break
# if there are two reference points, then crop the region of interest
# from teh image and display it

# close all open windows
cv2.destroyAllWindows()

if check_if_cropped == True :
	screenCnt = numpy.array([[[0,0]] ,[[x_end_cr,0]] ,[[x_end_cr, y_end_cr]] ,[[0 ,y_end_cr]]])

# apply the four point transform to obtain a top-downs
# view of the original image

#warped = four_point_transform(imutils.resize(roi, height = 500, width = 500), screenCnt.reshape(4, 2))
warped = four_point_transform(roi, screenCnt.reshape(4, 2))

# convert the warped image to grayscale, then threshold it
# to give it that 'black and white' paper effect

    # Threshold the HSV image to get only blue colors
warped = cv2.cvtColor(warped,cv2.COLOR_RGBA2GRAY)

T = threshold_local(warped, 11, offset = 5, method = "gaussian")
warped = (warped > T).astype("uint8") * 255

# show the original and scanned images
print("STEP 3: Apply perspective transform")
cv2.imshow("Original", imutils.resize(orig, height = 500))
#cv2.imshow("Scanned", imutils.resize(warped, height = 500, width = 500))
cv2.imshow("Scanned", warped)

cv2.waitKey(0)


# Module imports
from flask import request, abort, current_app
from sqlalchemy import exc
from marshmallow import ValidationError

class masqlapi():
    """ RestfulAPI SQLA/MarshMallow/jSend class """

    def __init__(self, session, resource, roschema, rwschema):
        """ Constructor """
        self.session = session
        self.resource = resource
        self.roschema = roschema
        self.rwschema = rwschema

    def get(self, **kwargs):
        """ Get record for a given resource """
        try:
            current_app.logger.debug(f'kwargs:{kwargs}')
            obj = self.resource.query.filter_by(**kwargs).first()
        except ValueError as err:
            current_app.logger.error(err)
            abort(400, err.args)
        return obj

    def get_many(self, **kwargs):
        """ Gets zero or more records for a given resource """
        try:
            obj = self.resource.query.filter_by(**kwargs)
        except ValueError as err:
            current_app.logger.error(err)
            abort(400, err.args)
        return obj

    def getMethod(self, obj, many=False):
        """ GET one """
        # Dump object
        try:
            data = self.roschema().dump(obj, many=many)
        except ValidationError as error:
            errMsg = {
                "errorType":"Schema Dump ValidationError",
                "errMsg":error.messages
            }
            abort(400, errMsg)
        except :
            abort(500, "schema dump error with unknown reason")

        # Return entity
        return {'message': None, 'status_code': 200, 'status': 'success', 'data': data}
    def add(self,obj):
        """ add the object from manual create"""
        # Merge and Commit
        try:
            self.session.add(obj)
            self.session.commit()
            data = self.roschema().dump(obj)
            return {'message': 'OK - entity add successfully', 'status_code': 200, 'status': 'success', 'data': data}
        except exc.IntegrityError as err:
            self.session.rollback()
            return {'message': 'Conflict', 'status_code': 409, 'status': 'failure', 'data': err.orig.__str__().strip()}
        except Exception as err:
            self.session.rollback()
            raise err

    def update(self,obj):
        """ just update the object from manual patch"""
        # Merge and Commit
        try:
            current_app.logger.debug(f"obj before patch:{obj.__dict__}")
            self.session.add(obj)
            self.session.commit()
            data = self.roschema().dump(obj)
            return {'message': 'OK - entity updated successfully', 'status_code': 200, 'status': 'success', 'data': data}
        except exc.IntegrityError as err:
            self.session.rollback()
            return {'message': 'Conflict', 'status_code': 409, 'status': 'failure', 'data': err.orig.__str__().strip()}
        except Exception as err:
            self.session.rollback()
            raise err

    def patch(self, obj):
        """ PATCH """
        # Grab request data
        json = request.get_json()
        if not json:
            return {'message': 'Bad request', 'status_code': 400, 'status': 'failure', 'data': 'JSON input object is missing or cannot be parsed'}

        try:
            current_app.logger.debug(f"isObj shown:{obj}, json:{json}")
            data = self.rwschema().load(json, instance=obj, partial=True)
            # data = self.rwschema().dump(obj)
            # data_id = data.pop('id',None)
            # data = {k: v for k, v in data.items() if v is not None}
            # data = self.rwschema().load(data)

            # current_app.logger.debug(f"data before patching:{data}")
            # patch_data = self.rwschema().load(json, partial=True)
            # data.update(patch_data)
            # if data_id:
            #     data['id'] = data_id
            # current_app.logger.debug(f"data for patching:{data}")
        except ValidationError as error:
            errMsg = {
                "errorType":"Schema Load ValidationError",
                "errMsg":error.messages
            }
            abort(400, errMsg)
        except :
            abort(500, "schema load error with unknown reason")

        # Merge and Commit
        try:
            for key, value in data.items():
                setattr(obj,key, value)
            current_app.logger.debug(f"obj before patch:{obj.__dict__}")
            self.session.add(obj)
            self.session.commit()
            data = self.roschema().dump(obj)
            return {'message': 'OK - entity updated successfully', 'status_code': 200, 'status': 'success', 'data': data}
        except exc.IntegrityError as err:
            self.session.rollback()
            return {'message': 'Conflict', 'status_code': 409, 'status': 'failure', 'data': err.orig.__str__().strip()}
        except Exception as err:
            self.session.rollback()
            raise err

    def delete(self, obj):
        """ DELETE """
        # Delete object
        try:
            self.session.delete(obj)
            self.session.commit()
            return {'message': 'No data', 'status_code': 200, 'status': 'success', 'data': 'Entity deleted successfully'}
        except exc.IntegrityError as err:
            self.session.rollback()
            return {'message': 'Conflict', 'status_code': 409, 'status': 'failure', 'data': err.orig.__str__().strip()}, 409
        except Exception as err:
            self.session.rollback()
            raise err

    def post(self,uniqueField):
        """ POST """

        # if uniqueField is not a list we make it into one
        if not isinstance(uniqueField, list):
            uniqueField_list = [uniqueField]
        else:
            uniqueField_list = uniqueField
        # Grab request data
        json = request.get_json()
        if not json:
            return {'message': 'Bad request', 'status_code': 400, 'status': 'failure', 'data': 'JSON input object is missing or cannot be parsed'}, 400
        # print("json:{0}, type:{1}".format(json,type(json)))
        # Validate and deserialize input
        try:
            data = self.rwschema().load(json)
        except ValidationError as error:
            errMsg = {
                "errorType":"Schema Load ValidationError",
                "errMsg":error.messages
            }
            abort(400, errMsg)
        except :
            abort(500, "schema load error with unknown reason")

        query = {}
        for keys in uniqueField_list:
            if keys in data:
                query[keys] = data[keys]
        # Check for existing row based on passed JSON
        current_app.logger.info(f"query data in post function:{query}")
        existing = self.resource.query.filter_by(**query).first()

        if existing:
            obj=self.roschema().dump(existing)
            return {'message': 'OK - entity exists', 'status_code': 200, 'status': 'success', 'data': obj}

        # Add and Commit
        try:
            obj = self.resource(**data)
            self.session.add(obj)
            self.session.commit()
            data = self.roschema().dump(obj)
            return {'message': 'Created - entity created successfully', 'status_code': 201, 'status': 'success', 'data': data}, 201
        except exc.IntegrityError as err:
            self.session.rollback()
            return {'message': 'Conflict', 'status_code': 409, 'status': 'failure', 'data': err.orig.__str__().strip()},409
        except Exception as err:
            self.session.rollback()
            raise err


#%%
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import numpy as np
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import confusion_matrix, plot_confusion_matrix
from sklearn.metrics import accuracy_score, mean_squared_error, mean_absolute_error, auc, roc_curve, plot_roc_curve
from sklearn.metrics import  precision_score, recall_score, f1_score, classification_report, roc_auc_score
from sklearn.model_selection import StratifiedKFold

from pre_pro import  X_train, X_test, y_train, y_test#, X_val, y_val #sem val por causa do kfold
#from pre_pro import PCA_X_test, PCA_X_train,

import tensorflow as tf
import os #disable logs
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

from tensorflow.keras.optimizers import SGD, Adam
from keras.callbacks import EarlyStopping
from keras.wrappers.scikit_learn import KerasClassifier
from keras.regularizers import l2
from keras.backend import clear_session
from ann_visualizer.visualize import ann_viz;
#%% K FOLD CROSS VALIDATION
kfold = StratifiedKFold(n_splits=2, shuffle=True, random_state=1)
cvscores = []

for train, test in kfold.split(X_train, y_train):
    # create model
    clear_session()
    ann = tf.keras.models.Sequential(name='MLPFullyConnected')
    ann.add(tf.keras.layers.InputLayer(input_shape=(44,)))
    ann.add(tf.keras.layers.Dense(units= 6, activation='relu'))
    ann.add(tf.keras.layers.Dense(units= 1, activation='sigmoid'))

    # Compile model
    optimizer = Adam()

    #optimizer = SGD(learning_rate=0.01, decay=1e-5)# , momentum=.9,

    ann.compile(optimizer=optimizer,
                loss= 'mean_squared_error',
                metrics= ['accuracy']) #adam

    early_stop = EarlyStopping(monitor='val_loss',
                            min_delta=0,
                            patience=5,
                            verbose=0,
                            mode='auto')  #testar patiance

    ann_history = ann.fit(X_train.iloc[train],y_train[train],
                        shuffle=True,
                        use_multiprocessing=True, #validation_split=0.1,
                        batch_size= 128,
                        epochs= 200,
                        validation_data = (X_train.iloc[test],y_train[test]),
                        callbacks=[early_stop])

    scores = ann.evaluate(X_test, y_test, verbose=0)
    print("%s: %.2f%%" % (ann.metrics_names[1], scores[1]*100))
    cvscores.append(scores[1] * 100)
#%%
print(cvscores)
print("%.2f%% (+/- %.2f%%)" % (np.mean(cvscores), np.std(cvscores)))

#%% KERAS ANN MLP
#maybe add dropout?
clear_session()

ann = tf.keras.models.Sequential(name='MLPFullyConnected')
ann.add(tf.keras.layers.InputLayer(input_shape=(44,)))
ann.add(tf.keras.layers.Dense(units= 6, activation='sigmoid')) #testar tgh
#ann.add(tf.keras.layers.Dense(units= 10, activation='relu')) #testar tgh
ann.add(tf.keras.layers.Dense(units= 1, activation='sigmoid')) #sigmoid

optimizer = Adam()

#testar sgd ainda.
#optimizer = SGD(learning_rate=0.01, momentum=.9, decay=1e-5,)#

ann.compile(optimizer=optimizer,
            loss= 'mean_squared_error',
            metrics= ['accuracy']) #adam

early_stop = EarlyStopping(monitor='val_loss',
                           min_delta=0,
                           patience=5,
                           verbose=0,
                           mode='auto')  #testar patiance

ann_history = ann.fit(X_train,y_train,
                      shuffle=True,
                      use_multiprocessing=True, #validation_split=0.1,
                      batch_size= 128,
                      epochs= 200,
                      validation_data = (X_val, y_val),
                      callbacks=[early_stop])

scores = ann.evaluate(X_test, y_test, batch_size=128)
#%%
#Loss
loss_train = ann_history.history['loss']
loss_val = ann_history.history['val_loss']
epochs = range(1,len(loss_train)+1)
plt.plot(epochs, loss_train, 'g', label='Training loss')
plt.plot(epochs, loss_val, 'b', label='validation loss')
plt.title('Training and Validation loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()
#%%
#Acuracia
print("%s: %.2f%%" % (ann.metrics_names[1], scores[1]*100))

ann.summary()
ann.get_config()
loss_train = ann_history.history['accuracy']
loss_val = ann_history.history['val_accuracy']
epochs = range(1,len(loss_train)+1)
plt.plot(epochs, loss_train, 'g', label='Training accuracy')
plt.plot(epochs, loss_val, 'b', label='validation accuracy')
plt.title('Training and Validation accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

#%% model visualization
from keras.utils.vis_utils import plot_model
plot_model(ann, to_file='model.png')
#%%confusion matrix
y_pred = ann.predict_classes(X_test)
#print(y_pred)
cma = confusion_matrix(y_test, y_pred) #labels='1,2'
#print(cma)

group_names = ['True Neg','False Pos','False Neg','True Pos']
group_counts = ['{0:0.0f}'.format(value) for value in
                cma.flatten()]
group_percentages = ['{0:.2%}'.format(value) for value in
                     cma.flatten()/np.sum(cma)]
labels = [f'{v1}\n{v2}\n{v3}' for v1, v2, v3 in
          zip(group_names,group_counts,group_percentages)]
labels = np.asarray(labels).reshape(2,2)
sns.heatmap(cma, annot=labels, fmt='', cmap='Blues')

#%%
print(classification_report(y_test, y_pred))
print(f'ROC AUC score: {roc_auc_score(y_test, y_pred)}')
print('Accuracy Score: ',accuracy_score(y_test, y_pred))
#%%
precision = precision_score(y_test, y_pred)  #Precision Score
recall = recall_score(y_test, y_pred)  #Recall Score
f1 = f1_score(y_test, y_pred)  #F1 Score

print(precision, recall, f1)
#%%
y_test.shape
#%%roc
false_positive_rate, true_positive_rate, thresholds = roc_curve(y_test, y_pred)
roc_auc = auc(false_positive_rate, true_positive_rate)
sns.set_theme(style = 'white')
plt.figure(figsize = (8, 8))
plt.plot(false_positive_rate,true_positive_rate, color = '#b01717', label = 'AUC = %0.3f' % roc_auc)
plt.legend(loc = 'lower right')
plt.plot([0, 1], [0, 1], linestyle = '--', color = '#174ab0')
plt.axis('tight')
plt.ylabel('True Positive Rate')
plt.xlabel('False Positive Rate')
plt.legend()
plt.show()

#%%error histogram for train, test, validation

sns.histplot(data=ann_history.history['loss'], kde=True, log_scale=True, element='step', fill=False)

# %%


"""Collective motion"""
import numba
import numpy as np
from cell_lists import add_to_cells, neighboring_cells, iter_nearest_neighbors
from numba import f8, i8
from numba.typing.typeof import typeof

from crowddynamics.core.sensory_region import is_obstacle_between_points
from crowddynamics.core.structures import obstacle_type_linear
from crowddynamics.core.vector2D import length, normalize, weighted_average, dot
from crowddynamics.simulation.agents import NO_LEADER, NO_TARGET

MISSING_NEIGHBOR = -1

@numba.jit([f8[:, :](f8[:, :])], nopython=True, nogil=True, cache=True)
def normalize_nx2(v):
    out = np.zeros_like(v)
    for i in range(len(v)):
        out[i, :] = normalize(v[i, :])
    return out

@numba.jit([(f8[:], f8[:], f8[:], f8[:], f8)],
           nopython=True, nogil=True, cache=True)
def herding_relationship(x1, x2, v1, v2, phi=np.pi / 2):
    """Determine leader-follower relationship

    Args:
        x1:
        x2:
        v1:
        v2:
        phi: Angle between [0, pi]

    Returns:
        (bool ,bool):
    """
    if length(v1) == 0 or length(v2) == 0:
        return False, False

    e_rel = normalize(x2 - x1)

    c_i = dot(e_rel, normalize(v1))
    c_j = -dot(e_rel, normalize(v2))

    # 1.0 = cos(0)
    cos_phi = np.cos(phi)
    if cos_phi < c_i < 1.0:
        if cos_phi < c_j < 1.0:
            return False, False
        else:
            return True, False
    else:
        if cos_phi < c_j < 1.0:
            return False, True
        else:
            return True, True

@numba.jit(nopython=True, nogil=True)
def set_neighbor(i, j, l, neighbors, distances, distances_max):
    argmax = np.argmax(distances[i, :])
    neighbors[i, argmax] = j
    distances[i, argmax] = l
    distances_max[i] = np.max(distances[i, :])

@numba.jit([(f8[:, :], f8, i8, i8[:], i8[:], i8[:],
             i8[:], i8[:], typeof(obstacle_type_linear)[:])],
           nopython=True, nogil=True, cache=True)
def find_nearest_neighbors(
        position, sight, size_nearest_other,
        cell_indices, neigh_cells, points_indices, cells_count,
        cells_offset, obstacles):
    size = len(position)

    neighbors = np.full(
        (size, size_nearest_other),
        fill_value=MISSING_NEIGHBOR, dtype=np.int64)
    '''Current nearest neighbours.'''

    distances = np.full(
        (size, size_nearest_other),
        fill_value=sight, dtype=np.float64)
    '''Distance to current nearest neighbours.'''

    distances_max = np.full(
        size, fill_value=sight, dtype=np.float64)
    '''Distance to furthest neighbor.'''

    for i, j in iter_nearest_neighbors(
            cell_indices, neigh_cells, points_indices, cells_count,
            cells_offset):
        # Test if line of sight is obstructed by an obstacle
        if is_obstacle_between_points(position[i], position[j], obstacles):
            continue

        l = length(position[i] - position[j])

        if l < distances_max[i]:
            set_neighbor(i, j, l, neighbors, distances, distances_max)

        if l < distances_max[j]:
            set_neighbor(j, i, l, neighbors, distances, distances_max)

    return neighbors

@numba.jit(nopython=True, nogil=True, cache=True)
def herding_interaction(is_herding, position, velocity, neighbors,
                        weight_position, phi):
    r"""Herding effect.

        .. math::
           \mathbf{\hat{e}_{herding}} = \mathcal{N}
           \left(\sum_{j \in Neigh} \mathbf{\hat{e}}_j\right)

    """
    # [(boolean[:], f8[:, :], f8[:, :], i8[:, :], f8, f8, i8[:])]
    new_direction = np.zeros_like(position)
    has_new_direction = np.zeros(position.shape[0], dtype=np.bool_)

    mean_position = np.zeros(position.shape[1])
    mean_velocity = np.zeros(velocity.shape[1])

    for i in range(position.shape[0]):
        if not is_herding[i]:
            continue

        num_neigh = 0

        for j in neighbors[i, :]:
            if j == MISSING_NEIGHBOR:
                continue

            is_heading_away, _ = herding_relationship(
                position[i], position[j], velocity[i], velocity[j], phi)

            if is_heading_away:
                mean_position[:] += position[j, :]
                mean_velocity[:] += velocity[j, :]
                num_neigh += 1

        if num_neigh > 0:
            new_direction[i, :] = normalize(weighted_average(
                normalize(mean_position[:] / num_neigh - position[i, :]),
                normalize(mean_velocity[:]),
                weight_position))
            has_new_direction[i] = True

        mean_position[:] = 0
        mean_velocity[:] = 0

    return new_direction, has_new_direction

@numba.jit(nopython=True, nogil=True, cache=True)
def leader_follower_interaction_brute(
        is_follower, is_leader, position, velocity, weight_position, phi,
        target, index_leader, obstacles, sight):
    has_strategy = np.zeros(len(position), dtype=np.bool_)
    new_direction = np.zeros_like(position)
    # new_target = np.zeros_like(target)
    # new_index_leader = np.zeros_like(index_leader)

    indices = np.arange(len(position))
    leaders = indices[is_leader]
    distances = np.zeros(leaders.shape, dtype=np.float64)

    # Iterate over followers
    for i in indices[is_follower]:
        behind_obstacle = 0
        heading_away = 0

        # Find distances to the leaders
        for k, j in enumerate(leaders):
            distances[k] = length(position[i] - position[j])

        # Iterate over leaders in orders from closest to furthest leader
        # and choose a strategy.
        for k in np.argsort(distances):
            if distances[k] > sight:
                continue

            j = leaders[k]

            if is_obstacle_between_points(position[i], position[j], obstacles):
                # We are not seeing this leader.
                leader = index_leader[i]
                # Check if we were following this leader before.
                if leader != NO_LEADER and leader == j:
                    behind_obstacle += 1
                    # Use navigation of this leader.
                    target[i] = target[leader]
                    has_strategy[i] = True
                    break
                else:
                    continue
            else:
                is_heading_away, _ = herding_relationship(
                    position[i], position[j], velocity[i], velocity[j], phi)

                if is_heading_away:
                    heading_away += 1
                    # We see this leader. Remember it and follow.
                    index_leader[i] = j
                    target[i] = NO_TARGET
                    new_direction[i, :] = normalize(weighted_average(
                        normalize(position[j, :] - position[i, :]),
                        normalize(velocity[j, :]), weight_position))
                    has_strategy[i] = True
                    break

        if behind_obstacle == 0 and heading_away == 0:
            leader = index_leader[i]
            if leader != NO_LEADER:
                target[i] = target[leader]
                has_strategy[i] = True

    return new_direction, has_strategy

def leader_follower_interaction(
        agents, obstacles, sight, phi=0.45 * np.pi,
        weight_position_leader=0.40):
    # Follow the leader
    direction_leader, has_strategy = leader_follower_interaction_brute(
        agents['is_follower'], agents['is_leader'], agents['position'],
        agents['velocity'], weight_position_leader, phi,
        agents['target'], agents['index_leader'], obstacles, sight)

    # Use familiar exits
    is_lost = ~has_strategy & agents['is_follower']
    agents['target'][is_lost] = agents['familiar_exit'][is_lost]

    return direction_leader

def leader_follower_with_herding_interaction(
        agents, obstacles, sight,
        size_nearest_other,
        phi=0.45 * np.pi,
        weight_position_herding=0.15,
        weight_position_leader=0.40,
        weight_direction_leader=0.65):
    position = agents['position']
    velocity = agents['velocity']
    is_leader = agents['is_leader']
    is_follower = agents['is_follower']
    sight_leader = 20.0

    cell_size = sight
    points_indices, cells_count, cells_offset, grid_shape = add_to_cells(
        agents['position'], cell_size)
    cell_indices = np.arange(len(cells_count))
    neigh_cells = neighboring_cells(grid_shape)

    neighbors = find_nearest_neighbors(
        position, sight,
        size_nearest_other, cell_indices, neigh_cells, points_indices,
        cells_count, cells_offset, obstacles)

    direction, has_direction = herding_interaction(
        is_follower, position, velocity, neighbors, weight_position_herding,
        phi)

    agents['target'][has_direction] = NO_TARGET

    direction_leader, has_strategy = leader_follower_interaction_brute(
        is_follower, is_leader, position, velocity, weight_position_leader,
        phi, agents['target'], agents['index_leader'], obstacles, sight_leader)

    # Use familiar exits
    is_lost = ~(has_direction | has_strategy) & is_follower
    agents['target'][is_lost] = agents['familiar_exit'][is_lost]

    return normalize_nx2(weighted_average(
        direction_leader, direction, weight_direction_leader))


import torch
import torch.nn as nn
import numpy as np

class FeatureExtractor(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_size,hidden_size),
            nn.LeakyReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.LeakyReLU(),
            nn.Linear(hidden_size, output_size),
            nn.LeakyReLU(),
        )

    def forward(self, input):
        # input : [B, N, C, W, H]
        # flatten to [B, N, D]
        input = input.view(input.size(0), input.size(1), -1)
        return self.net(input)

##
# Scores
class DotScore(nn.Module):
    def __init__(self, H):
        super(DotScore, self).__init__()
        self.H = H

    def forward(self, s, context):
        """
        s : [B, N, D](B, T, H) shape
        context: [B, D]
        score : [B, N]
        """
        # bmm([B, N, D], [B, D, 1]) -> [B, N]
        score = torch.bmm(s, context.unsqueeze(-1)) / np.sqrt(self.H)
        return score


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from matplotlib import pyplot
import numpy as np
import os
import shutil
import operator
import caffe2.python.predictor.predictor_exporter as pe
import argparse
from caffe2.proto import caffe2_pb2

from caffe2.python import (
    brew,
    core,
    model_helper,
    net_drawer,
    optimizer,
    visualize,
    workspace,
)

# If you would like to see some really detailed initializations,
# you can change --caffe2_log_level=0 to --caffe2_log_level=-1
core.GlobalInit(['caffe2', '--caffe2_log_level=1'])
print("Necessities imported!")

class TrainCaffe(object):
    USE_LENET_MODEL = True
    data_folder = None
    root_folder = None
    train_model = None
    test_model = None
    deploy_model = None

    def __init__(self):
        print('Init train class')

    # This section preps your image and test set in a lmdb database
    def DownloadResource(self, url, path):
        '''Downloads resources from s3 by url and unzips them to the provided path'''
        import requests, zipfile, StringIO
        print("Downloading... {} to {}".format(url, path))
        r = requests.get(url, stream=True)
        z = zipfile.ZipFile(StringIO.StringIO(r.content))
        z.extractall(path)
        print("Completed download and extraction.")

    def PrepareData(self):
        # Setup the paths for the necessary directories
        current_folder = os.path.join(os.path.expanduser('~'), 'data')
        self.data_folder = os.path.join(current_folder, 'tutorial_data', 'breakhis')
        self.root_folder = os.path.join(current_folder, 'tutorial_files', 'tutorial_breakhis')
        db_missing = False

        # Check if the data folder already exists
        if not os.path.exists(self.data_folder):
            os.makedirs(self.data_folder)
            print("Your data folder was not found!! This was generated: {}".format(self.data_folder))

        # Check if the training lmdb exists in the data folder
        if os.path.exists(os.path.join(self.data_folder, "breakhis40x-train-nchw-lmdb")):
            print("lmdb train db found!")
        else:
            db_missing = True

        # Check if the testing lmdb exists in the data folder
        if os.path.exists(os.path.join(self.data_folder, "breakhis40x-test-nchw-lmdb")):
            print("lmdb test db found!")
        else:
            db_missing = True

        # Attempt the download of the db if either was missing
        if db_missing:
            print("one or both of the MNIST lmbd dbs not found!!")
            db_url = "http://download.caffe2.ai/databases/mnist-lmdb.zip"
            try:
                self.DownloadResource(db_url, self.data_folder)
            except Exception as ex:
                print("Failed to download dataset. Please download it manually from {}".format(db_url))
                print("Unzip it and place the two database folders here: {}".format(self.data_folder))
                raise ex

        # Clean up statistics from any old runs
        if os.path.exists(self.root_folder):
            print("Looks like you ran this before, so we need to cleanup those old files...")
            shutil.rmtree(self.root_folder)

        os.makedirs(self.root_folder)
        workspace.ResetWorkspace(self.root_folder)

        print("training data folder:" + self.data_folder)
        print("workspace root folder:" + self.root_folder)

    def AddInput(self, model, batch_size, db, db_type):
        ### load the data from db - Method 1 using brew
        #data_uint8, label = brew.db_input(
        #    model,
        #    blobs_out=["data_uint8", "label"],
        #    batch_size=batch_size,
        #    db=db,
        #    db_type=db_type,
        #)
        ### load the data from db - Method 2 using TensorProtosDB
        data_uint8, label = model.TensorProtosDBInput(
            [], ["data_uint8", "label"], batch_size=batch_size,
            db=db, db_type=db_type)
        # cast the data to float
        data = model.Cast(data_uint8, "data", to=core.DataType.FLOAT)
        # scale data from [0,255] down to [0,1]
        data = model.Scale(data, data, scale=float(1./256))
        # don't need the gradient for the backward pass
        data = model.StopGradient(data, data)
        return data, label

    # Function to construct a MLP neural network
    # The input 'model' is a model helper and 'data' is the input data blob's name
    def AddMLPModel(self, model, data):
        size = 28 * 28 * 1
        sizes = [size, size * 2, size * 2, 10]
        layer = data
        for i in range(len(sizes) - 1):
            layer = brew.fc(model, layer, 'dense_{}'.format(i), dim_in=sizes[i], dim_out=sizes[i + 1])
            layer = brew.relu(model, layer, 'relu_{}'.format(i))
        softmax = brew.softmax(model, layer, 'softmax')
        return softmax

    def AddLeNetModel(self, model, data):
        '''
        This part is the standard LeNet model: from data to the softmax prediction.

        For each convolutional layer we specify dim_in - number of input channels
        and dim_out - number or output channels. Also each Conv and MaxPool layer changes the
        image size. For example, kernel of size 5 reduces each side of an image by 4.

        While when we have kernel and stride sizes equal 2 in a MaxPool layer, it divides
        each side in half.
        '''
        # Image size: 28 x 28 -> 24 x 24
        conv1 = brew.conv(model, data, 'conv1', dim_in=1, dim_out=20, kernel=5)
        # Image size: 24 x 24 -> 12 x 12
        pool1 = brew.max_pool(model, conv1, 'pool1', kernel=2, stride=2)
        # Image size: 12 x 12 -> 8 x 8
        conv2 = brew.conv(model, pool1, 'conv2', dim_in=20, dim_out=50, kernel=5)
        # Image size: 8 x 8 -> 4 x 4
        pool2 = brew.max_pool(model, conv2, 'pool2', kernel=2, stride=2)
        # 50 * 4 * 4 stands for dim_out from previous layer multiplied by the image size
        # Here, the data is flattened from a tensor of dimension 50x4x4 to a vector of length 50*4*4
        fc3 = brew.fc(model, pool2, 'fc3', dim_in=50 * 4 * 4, dim_out=500)
        relu3 = brew.relu(model, fc3, 'relu3')
        # Last FC Layer
        pred = brew.fc(model, relu3, 'pred', dim_in=500, dim_out=10)
        # Softmax Layer
        softmax = brew.softmax(model, pred, 'softmax')

        return softmax

    def AddLeNetModelBreakHis(self, model, data):
        '''
        This part is the standard LeNet model: from data to the softmax prediction.

        For each convolutional layer we specify dim_in - number of input channels
        and dim_out - number or output channels. Also each Conv and MaxPool layer changes the
        image size. For example, kernel of size 5 reduces each side of an image by 4.

        While when we have kernel and stride sizes equal 2 in a MaxPool layer, it divides
        each side in half.
        '''
        # Image size: 28 x 28 -> 24 x 24
        # Image size: 224 x 224-> 220 x 220
        conv1 = brew.conv(model, data, 'conv1', dim_in=3, dim_out=20, kernel=5)
        # Image size: 24 x 24 -> 12 x 12
        # Image size: 220 x 220-> 110 x 110
        pool1 = brew.max_pool(model, conv1, 'pool1', kernel=2, stride=2)
        # Image size: 12 x 12 -> 8 x 8
        # Image size: 110 x 110 -> 106 x 106
        conv2 = brew.conv(model, pool1, 'conv2', dim_in=20, dim_out=50, kernel=5)
        # Image size: 8 x 8 -> 4 x 4
        # Image size: 106 x 106 -> 53 x 53
        pool2 = brew.max_pool(model, conv2, 'pool2', kernel=2, stride=2)
        # 50 * 4 * 4 stands for dim_out from previous layer multiplied by the image size
        # Here, the data is flattened from a tensor of dimension 50x4x4 to a vector of length 50*4*4
        fc3 = brew.fc(model, pool2, 'fc3', dim_in=50 * 53 * 53, dim_out=500)
        relu3 = brew.relu(model, fc3, 'relu3')
        # Last FC Layer
        pred = brew.fc(model, relu3, 'pred', dim_in=500, dim_out=2)
        # Softmax Layer
        softmax = brew.softmax(model, pred, 'softmax')

        return softmax

    def AddModel(self,model, data):
        if self.USE_LENET_MODEL:
            return self.AddLeNetModelBreakHis(model, data)
        else:
            return self.AddMLPModel(model, data)

    def AddAccuracy(self, model, softmax, label):
        """Adds an accuracy op to the model"""
        accuracy = brew.accuracy(model, [softmax, label], "accuracy")
        return accuracy

    def AddTrainingOperators(self, model, softmax, label):
        """Adds training operators to the model."""
        # Compute cross entropy between softmax scores and labels
        xent = model.LabelCrossEntropy([softmax, label], 'xent')
        # Compute the expected loss
        loss = model.AveragedLoss(xent, "loss")
        # Track the accuracy of the model
        self.AddAccuracy(model, softmax, label)
        # Use the average loss we just computed to add gradient operators to the model
        model.AddGradientOperators([loss])
        # Specify the optimization algorithm
        optimizer.build_sgd(
            model,
            base_learning_rate=0.1,
            policy="step",
            stepsize=1,
            gamma=0.999,
        )

    def AddBookkeepingOperators(self, model):
        """This adds a few bookkeeping operators that we can inspect later.

        These operators do not affect the training procedure: they only collect
        statistics and prints them to file or to logs.
        """
        # Print basically prints out the content of the blob. to_file=1 routes the
        # printed output to a file. The file is going to be stored under
        #     root_folder/[blob name]
        model.Print('accuracy', [], to_file=1)
        model.Print('loss', [], to_file=1)
        # Summarizes the parameters. Different from Print, Summarize gives some
        # statistics of the parameter, such as mean, std, min and max.
        for param in model.params:
            model.Summarize(param, [], to_file=1)
            model.Summarize(model.param_to_grad[param], [], to_file=1)
        # Now, if we really want to be verbose, we can summarize EVERY blob
        # that the model produces; it is probably not a good idea, because that
        # is going to take time - summarization do not come for free. For this
        # demo, we will only show how to summarize the parameters and their
        # gradients.

    def SetDeviceOption(self, model, device_option):
        # Clear op-specific device options and set global device option.
        for net in ("net", "param_init_net"):
            net_def = getattr(model, net).Proto()
            net_def.device_option.CopyFrom(device_option)
            #for op in net_def.op:
            #    # Some operators are CPU-only.
            #    if str(op.output[0]) not in ("optimizer_iteration", "iteration_mutex", "data", "data_uint8"):
            #        op.ClearField("device_option")
            #        op.ClearField("engine")
            setattr(model, net, core.Net(net_def))
    '''
    def SetRunGPU(self, model):
        device_option = caffe2_pb2.DeviceOption()
        device_option.device_type = caffe2_pb2.CUDA
        device_option.cuda_gpu_id = 0
        self.SetDeviceOption(model, device_option)
        return device_option
    '''

    def ConstructModel(self):
        #### Train Model
        # Specify the data will be input in NCHW order
        #  (i.e. [batch_size, num_channels, height, width])
        arg_scope = {"order": "NCHW"}
        # Create the model helper for the train model
        self.train_model = model_helper.ModelHelper(name="mnist_train", arg_scope=arg_scope)
        #self.train_model.param_init_net.RunAllOnGPU()
        #self.train_model.net.RunAllOnGPU()
        # Specify the input is from the train lmdb
        data, label = self.AddInput(
            self.train_model, batch_size=64,
            db=os.path.join(self.data_folder, 'breakhis40x-train-nchw-lmdb'),
            db_type='lmdb')
        # Add the model definition (fc layers, conv layers, softmax, etc.)
        softmax = self.AddModel(self.train_model, data)
        # Add training operators, specify loss function and optimization algorithm
        self.AddTrainingOperators(self.train_model, softmax, label)
        # Add bookkeeping operators to save stats from training
        self.AddBookkeepingOperators(self.train_model)

        #### Testing model.
        # We will set the batch size to 100, so that the testing
        #   pass is 100 iterations (10,000 images in total).
        #   For the testing model, we need the data input part, the main AddModel
        #   part, and an accuracy part. Note that init_params is set False because
        #   we will be using the parameters obtained from the train model which will
        #   already be in the workspace.
        self.test_model = model_helper.ModelHelper(
            name="mnist_test", arg_scope=arg_scope, init_params=False)
        #self.test_model.param_init_net.RunAllOnGPU()
        #self.test_model.net.RunAllOnGPU()

        data, label = self.AddInput(
            self.test_model, batch_size=100,
            db=os.path.join(self.data_folder, 'breakhis40x-test-nchw-lmdb'),
            db_type='lmdb')
        softmax = self.AddModel(self.test_model, data)
        self.AddAccuracy(self.test_model, softmax, label)

        #### Deployment model.
        # We simply need the main AddModel part.
        self.deploy_model = model_helper.ModelHelper(
            name="mnist_deploy", arg_scope=arg_scope, init_params=False)
        #self.deploy_model.param_init_net.RunAllOnGPU()
        #self.deploy_model.net.RunAllOnGPU()

        self.AddModel(self.deploy_model, "data")

        print("*******train_model.net.Proto()*******\n")
        print(str(self.train_model.net.Proto())[:400] + '\n...')
        print("\n*******train_model.param_init_net.Proto()*******\n")
        print(str(self.train_model.param_init_net.Proto())[:400] + '\n...')

    def RunTraining(self):
        # The parameter initialization network only needs to be run once.
        # Now all the parameter blobs are initialized in the workspace.
        workspace.RunNetOnce(self.train_model.param_init_net)
        #for name in workspace.C.blobs():
        #    blob = workspace.FetchBlob(name)
        #    workspace.FeedBlob(name, blob, device_option=core.DeviceOption(caffe2_pb2.CUDA, 0))

        # Creating an actual network as a C++ object in memory.
        #   We need this as the object is going to be used a lot
        #   so we avoid creating an object every single time it is used.
        # overwrite=True allows you to run this cell several times and avoid errors
        workspace.CreateNet(self.train_model.net, overwrite=True)
        #device_option = self.SetRunGPU(self.train_model)

        # Set the iterations number and track the accuracy & loss
        total_iters = 200
        accuracy = np.zeros(total_iters)
        loss = np.zeros(total_iters)

        # MAIN TRAINING LOOP!
        # Now, we will manually run the network for 200 iterations.
        for i in range(total_iters):
            workspace.RunNet(self.train_model.net)
            accuracy[i] = workspace.blobs['accuracy']
            loss[i] = workspace.blobs['loss']
            # Check the accuracy and loss every so often
            if i % 25 == 0:
                print("Iter: {}, Loss: {}, Accuracy: {}".format(i, loss[i], accuracy[i]))

        # After the execution is done, let's plot the values.
        pyplot.plot(loss, 'b')
        pyplot.plot(accuracy, 'r')
        pyplot.title("Summary of Training Run")
        pyplot.xlabel("Iteration")
        pyplot.legend(('Loss', 'Accuracy'), loc='upper right')

    def VisualizeResults(self):
        ### Let's look at some of the training data
        pyplot.figure()
        pyplot.title("Training Data Sample")
        # Grab the most recent data blob (i.e. batch) from the workspace
        data = workspace.FetchBlob('data')
        # Use visualize module to show the examples from the last batch that was fed to the model
        _ = visualize.NCHW.ShowMultiple(data)

        ### Let's visualize the softmax result
        pyplot.figure()
        pyplot.title('Softmax Prediction for the first image above')
        pyplot.ylabel('Confidence')
        pyplot.xlabel('Label')
        # Grab and visualize the softmax blob for the batch we just visualized. Since batch size
        #  is 64, the softmax blob contains 64 vectors, one for each image in the batch. To grab
        #  the vector for the first image, we can simply index the fetched softmax blob at zero.
        softmax = workspace.FetchBlob('softmax')
        _ = pyplot.plot(softmax[0], 'ro')

    def VisualizeLearnedFilters(self):
        if self.USE_LENET_MODEL:
            pyplot.figure()
            pyplot.title("Conv1 Output Feature Maps for Most Recent Mini-batch")
            # Grab the output feature maps of conv1. Change this to conv2 in order to look into the second one.
            #  Remember, early convolutional layers tend to learn human-interpretable features but later conv
            #  layers work with highly-abstract representations. For this reason, it may be harder to understand
            #  features of the later conv layers.
            conv = workspace.FetchBlob('conv1')

            # We can look into any channel. Think of it as a feature model learned.
            # In this case we look into the 5th channel. Play with other channels to see other features
            conv1 = conv[:, [1], :, :]

            _ = visualize.NCHW.ShowMultiple(conv1)

            pyplot.figure()
            pyplot.title("Conv1 Output Feature Maps for Most Recent Mini-batch")
            conv2 = conv[:, [5], :, :]

            e = visualize.NCHW.ShowMultiple(conv2)

    def TestModel(self):
        # param_init_net here will only create a data reader
        # Other parameters won't be re-created because we selected init_params=False before
        workspace.RunNetOnce(self.test_model.param_init_net)
        workspace.CreateNet(self.test_model.net, overwrite=True)

        # Testing Loop
        test_accuracy = np.zeros(100)
        for i in range(100):
            # Run a forward pass of the net on the current batch
            workspace.RunNet(self.test_model.net)
            # Collect the batch accuracy from the workspace
            test_accuracy[i] = workspace.FetchBlob('accuracy')

        # After the execution is done, let's plot the accuracy values.
        pyplot.figure()
        pyplot.plot(test_accuracy, 'r')
        pyplot.title('Accuracy over test batches.')
        print('test_accuracy: %f' % test_accuracy.mean())

    def DeployModel(self):
        # construct the model to be exported
        # the inputs/outputs of the model are manually specified.
        pe_meta = pe.PredictorExportMeta(
            predict_net=self.deploy_model.net.Proto(),
            parameters=[str(b) for b in self.deploy_model.params],
            inputs=["data"],
            outputs=["softmax"],
        )

        # save the model to a file. Use minidb as the file format
        pe.save_to_db("minidb", os.path.join(self.root_folder, "mnist_model.minidb"), pe_meta)
        print("Deploy model saved to: " + self.root_folder + "/mnist_model.minidb")

    def LoadModel(self):
        # Grab and display the last data batch used before we scratch the workspace. This purely for our convenience...
        blob = workspace.FetchBlob("data")
        pyplot.figure()
        pyplot.title("Batch of Testing Data")
        _ = visualize.NCHW.ShowMultiple(blob)

        # reset the workspace, to make sure the model is actually loaded
        workspace.ResetWorkspace(self.root_folder)

        # verify that all blobs from training are destroyed.
        print("The blobs in the workspace after reset: {}".format(workspace.Blobs()))

        # load the predict net
        predict_net = pe.prepare_prediction_net(os.path.join(self.root_folder, "mnist_model.minidb"), "minidb")

        # verify that blobs are loaded back
        print("The blobs in the workspace after loading the model: {}".format(workspace.Blobs()))

        # feed the previously saved data to the loaded model
        workspace.FeedBlob("data", blob)

        # predict
        workspace.RunNetOnce(predict_net)
        softmax = workspace.FetchBlob("softmax")

        print("Shape of softmax: ", softmax.shape)

        # Quick way to get the top-1 prediction result
        # Squeeze out the unnecessary axis. This returns a 1-D array of length 10
        # Get the prediction and the confidence by finding the maximum value and index of maximum value in preds array
        curr_pred, curr_conf = max(enumerate(softmax[0]), key=operator.itemgetter(1))
        print("Prediction: ", curr_pred)
        print("Confidence: ", curr_conf)

        # the first letter should be predicted correctly
        pyplot.figure()
        pyplot.title('Prediction for the first image')
        pyplot.ylabel('Confidence')
        pyplot.xlabel('Label')
        _ = pyplot.plot(softmax[0], 'ro')

def main():
    parser = argparse.ArgumentParser(
        description="BreakHis: breast cancer classification"
    )
    parser.add_argument("--predict_net", type=str, default='models/squeezenet/predict_net.pb',
                        help="Predict net file (net with operators")
    args = parser.parse_args()

    print(args)
    tr = TrainCaffe()
    tr.PrepareData()
    tr.ConstructModel()
    tr.RunTraining()
    tr.VisualizeResults()
    tr.VisualizeLearnedFilters()
    tr.TestModel()
    tr.DeployModel()
    tr.LoadModel()

if __name__ == '__main__':
    workspace.GlobalInit(['caffe2', '--caffe2_log_level=2'])
    main()


from sklearn import svm
from MFCCStuff import *

class SVMStuff:
    # 22 according to paper (https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.148.5904&rep=rep1&type=pdf)
    order = 22

    def __init__(self, target_x, impostor_x):
        self.target_num_rows, self.target_num_cols = target_x.shape
        self.impostor_num_rows, self.impostor_num_cols = impostor_x.shape

        self.target_x = target_x
        self.target_y = [1] * self.target_num_rows

        self.impostor_x = impostor_x
        self.impostor_y = [-1] * self.impostor_num_rows

        self.svm = None  # Support vector machine
        self.fit_model()

    def fit_model(self):
        self.svm = svm.SVC(kernel="rbf", gamma='scale', decision_function_shape='ovr')
        x = np.concatenate((self.target_x, self.impostor_x), axis=0)
        y = self.target_y + self.impostor_y
        self.svm.fit(x, y)

    def predict(self, impostor_x):
        return self.svm.predict(impostor_x)


# -*- coding: utf-8 -*-
"""
Created on Thu Jul  6 09:09:54 2017

@author: user
"""
import operator
from math import log
import matplotlib.pyplot as plt

def calcShannonEnt(dataSet):
    numEntries = len(dataSet)
    labelCounts = {}
    for featVec in dataSet:
        currentLabel = featVec[-1]
        if currentLabel not in labelCounts.keys():
            labelCounts[currentLabel] = 0
        labelCounts[currentLabel] += 1
        shannonEnt = 0.0
    for key in labelCounts:
        prob = float(labelCounts[key])/numEntries
        shannonEnt -= prob * log(prob,2)
    return shannonEnt

def createDataSet():
    dataSet = [[1, 1, 'yes'],
    [1, 1, 'yes'],
    [1, 0, 'no'],
    [0, 1, 'no'],
    [0, 1, 'no']]
    labels = ['no surfacing','flippers']
    return dataSet, labels
"""
myDat,labels=createDataSet()
print(myDat)

print(calcShannonEnt(myDat))
myDat[0][-1]='maybe'
print(myDat)
print(calcShannonEnt(myDat))
"""
def splitDataSet(dataSet, axis, value):
    retDataSet = []
    for featVec in dataSet:
        if featVec[axis] == value:
            reducedFeatVec = featVec[:axis]
            reducedFeatVec.extend(featVec[axis+1:])
            retDataSet.append(reducedFeatVec)
    return retDataSet
"""
myDat,labels = createDataSet()

print(myDat)
print(splitDataSet(myDat,0,1))
print(splitDataSet(myDat,0,0))
"""

def chooseBestFeatureToSplit(dataSet):
    numFeatures = len(dataSet[0]) - 1
    baseEntropy = calcShannonEnt(dataSet)
    bestInfoGain = 0.0; bestFeature = -1
    for i in range(numFeatures):
        featList = [example[i] for example in dataSet]
        uniqueVals = set(featList)
        newEntropy = 0.0
        for value in uniqueVals:
            subDataSet = splitDataSet(dataSet, i, value)
            prob = len(subDataSet)/float(len(dataSet))
            newEntropy += prob * calcShannonEnt(subDataSet)
        infoGain = baseEntropy - newEntropy
        if (infoGain > bestInfoGain):
            bestInfoGain = infoGain
            bestFeature = i
    return bestFeature

myDat,labels = createDataSet()
chooseBestFeatureToSplit(myDat)
# print(myDat)

def majorityCnt(classList):
    classCount={}
    for vote in classList:
        if vote not in classCount.keys(): classCount[vote] = 0
        classCount[vote] += 1
    sortedClassCount = sorted(classCount.iteritems(),
    key=operator.itemgetter(1), reverse=True)
    return sortedClassCount[0][0]

def createTree(dataSet,labels):
    classList = [example[-1] for example in dataSet]
    if classList.count(classList[0]) == len(classList):
        return classList[0]#stop splitting when all of the classes are equal
    if len(dataSet[0]) == 1: #stop splitting when there are no more features in dataSet
        return majorityCnt(classList)
    bestFeat = chooseBestFeatureToSplit(dataSet)
    bestFeatLabel = labels[bestFeat]
    myTree = {bestFeatLabel:{}}
    del(labels[bestFeat])
    featValues = [example[bestFeat] for example in dataSet]
    uniqueVals = set(featValues)
    for value in uniqueVals:
        subLabels = labels[:]       #copy all of labels, so trees don't mess up existing labels
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value),subLabels)
    return myTree

myDat,labels = createDataSet()
myTree = createTree(myDat,labels)

# print(myTree)

decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")

def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    createPlot.ax1.annotate(nodeTxt, xy=parentPt,
    xycoords='axes fraction',
    xytext=centerPt, textcoords='axes fraction',
    va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)

def createPlot():
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    createPlot.ax1 = plt.subplot(111, frameon=False)
    plotNode('a decision node', (0.5, 0.1), (0.1, 0.5), decisionNode)
    plotNode('a leaf node', (0.8, 0.1), (0.3, 0.8), leafNode)
    plt.show()

# createPlot()

def getNumLeafs(myTree):
    numLeafs = 0
    firstStr = myTree.keys()[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]).__name__=='dict':
            numLeafs += getNumLeafs(secondDict[key])
        else:
            numLeafs +=1
    return numLeafs

def getTreeDepth(myTree):
    maxDepth = 0
    firstStr = myTree.keys()[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]).__name__=='dict':
            thisDepth = 1 + getTreeDepth(secondDict[key])
        else:thisDepth = 1
        if thisDepth > maxDepth: maxDepth = thisDepth
    return maxDepth

def retrieveTree(i):
    listOfTrees =[{'no surfacing': {0: 'no', 1: {'flippers': \
    {0: 'no', 1: 'yes'}}}},
    {'no surfacing': {0: 'no', 1: {'flippers': \
    {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}
    ]
    return listOfTrees[i]

print(retrieveTree (1))

def plotMidText(cntrPt, parentPt, txtString):
    xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0]
    yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]
    createPlot.ax1.text(xMid, yMid, txtString)

def plotTree(myTree, parentPt, nodeTxt):
    numLeafs = getNumLeafs(myTree)
    getTreeDepth(myTree)
    firstStr = myTree.keys()[0]
    cntrPt = (plotTree.xOff + (1.0 + float(numLeafs))/2.0/plotTree.totalW,\
    plotTree.yOff)
    plotMidText(cntrPt, parentPt, nodeTxt)
    plotNode(firstStr, cntrPt, parentPt, decisionNode)
    secondDict = myTree[firstStr]
    plotTree.yOff = plotTree.yOff - 1.0/plotTree.totalD
    for key in secondDict.keys():
        if type(secondDict[key]).__name__=='dict':
            plotTree(secondDict[key],cntrPt,str(key))
        else:
            plotTree.xOff = plotTree.xOff + 1.0/plotTree.totalW
            plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff),
            cntrPt, leafNode)
            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))
    plotTree.yOff = plotTree.yOff + 1.0/plotTree.totalD

def createPlot(inTree):
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
    plotTree.totalW = float(getNumLeafs(inTree))
    plotTree.totalD = float(getTreeDepth(inTree))
    plotTree.xOff = -0.5/plotTree.totalW; plotTree.yOff = 1.0;
    plotTree(inTree, (0.5,1.0), '')
    plt.show()


import requests as req
import pandas as pd

from bs4 import BeautifulSoup
import re
import psycopg2
import datetime as dt
from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook

def postgre_read():
    # Obtner data de la base de datos postgres e imprimirlo
    hook = PostgresHook(postgres_conn_id="Canasta")
    conn = hook.get_conn()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM precios;")

    print(cursor.description)

    cursor.close()
    conn.close()

with DAG(
    dag_id="postres_prueba",
    description="DAG para scrapping de canasta familiar",
    schedule_interval="15 08 * * *",
    default_args={
        "owner": "airflow",
        "retries": 1,
        "retry_delay": timedelta(minutes=20),
        "start_date": datetime(2023, 3, 16),
        "email": ["ismaelpiovani@gmail.com"],
        "email_on_success": True,
        "email_on_failure": True,
        "email_on_retry": True,
    },
    catchup=False,
) as f:
    task1 = PythonOperator(task_id="prueba_postgres", python_callable=postgre_read)


'''
    Reference : https://github.com/rlcode/per
    Date: 2022.11.30
    Authored: Dala zhu
'''

import random
import math
import sys
import gymnasium as gym
import numpy as np
import numpy
import pylab
import torch
import torch.nn as nn

EPSILON_START = 1.0
EPSILON_FINAL = 0.01
EPSILON_DECAY = 1000
epsilon_by_frame = lambda i : EPSILON_FINAL + (EPSILON_START - EPSILON_FINAL) * math.exp(-1 * i / EPSILON_DECAY)

# SumTree
# a binary tree data structure where the parent’s value is the sum of its children
class SumTree:
    write = 0
    def __init__(self, capacity):
        self.capacity = capacity
        self.tree = numpy.zeros(2 * capacity - 1)  #start from 0
        self.data = numpy.zeros(capacity, dtype=object)
        self.n_entries = 0  #tree的leaf中被占用的个数，也即data中不为空的个数

    # update to the root node
    def _propagate(self, idx, change):
        parent = (idx - 1) // 2
        self.tree[parent] += change
        if parent != 0:
            self._propagate(parent, change)

    # find sample on leaf node
    def _retrieve(self, idx, s):
        left = 2 * idx + 1
        right = left + 1

        if left >= len(self.tree):
            return idx

        if s <= self.tree[left]:
            return self._retrieve(left, s)
        else:
            return self._retrieve(right, s - self.tree[left])

    def total(self):
        return self.tree[0]

    # store priority and sample
    def add(self, p, data):
        idx = self.write + self.capacity - 1

        self.data[self.write] = data
        self.update(idx, p)

        self.write += 1
        if self.write >= self.capacity:
            self.write = 0

        if self.n_entries < self.capacity:
            self.n_entries += 1

    # update priority
    def update(self, idx, p):
        change = p - self.tree[idx]
        self.tree[idx] = p
        self._propagate(idx, change)

    # get priority and sample
    def get(self, s):
        idx = self._retrieve(0, s)
        dataIdx = idx - self.capacity + 1

        return (idx, self.tree[idx], self.data[dataIdx])

class Memory:  # stored as ( s, a, r, s_, done) in SumTree
    e = 0.01  #ebusilon
    a = 0.6 #alpha
    beta = 0.4
    beta_increment_per_sampling = 0.001

    def __init__(self, capacity):
        self.tree = SumTree(capacity)
        self.capacity = capacity

    def _get_priority(self, error):
        return (np.abs(error) + self.e) ** self.a

    def add(self, error, sample):   #sample is (s, a, r, s_, done)
        p = self._get_priority(error)
        self.tree.add(p, sample)

    def sample(self, n):
        batch = []
        idxs = []

        segment = self.tree.total() / n  #抽取时为何会进行这样化等分抽取？
        priorities = []

        self.beta = np.min([1., self.beta + self.beta_increment_per_sampling])  #每次采样，beta都会增加，为何？

        for i in range(n):
            a = segment * i
            b = segment * (i + 1)

            s = random.uniform(a, b)   # sample from [a,b]
            (idx, p, data) = self.tree.get(s)
            priorities.append(p)
            batch.append(data)
            idxs.append(idx)

        sampling_probabilities = np.array(priorities) / self.tree.total()
        is_weight = np.power(self.tree.n_entries * sampling_probabilities, - self.beta)
        is_weight /= is_weight.max()

        return batch, idxs, is_weight #idxs is the index in tree.tree, not in tree.data

    def update(self, idx, error):
        p = self._get_priority(error)
        self.tree.update(idx, p)

class DuelingDQN(nn.Module):
    def __init__(self, n_state, n_action):

        super(DuelingDQN, self).__init__()
        self.n_state = n_state
        self.n_action = n_action

        self.feature = nn.Sequential(
            nn.Linear(self.n_state, 128),
            nn.ReLU()
        )

        self.advantage = nn.Sequential(
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, self.n_action)
        )

        self.value = nn.Sequential(
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )

    def forward(self, x):
        x = self.feature(x)
        advantage = self.advantage(x)
        value = self.value(x)
        return value + advantage - advantage.mean()  # V(s) + A(s,a) - mean(A(s,a)),这里是将max(A(s,a))换成了mean(A(s,a))

class Agent:
    def __init__(self, num_state, num_action):

        self.num_state = num_state
        self.num_action = num_action
        self.learning_rate = 1e-3
        self.gamma = 0.99
        self.batch_size = 64  # 如果太大训练不好训练
        self.load_model = False #超参，是否记载模型

        self.online_net = DuelingDQN(num_state, num_action)

        if self.load_model:
            self.online_net = torch.load('./save_model/cartpole_dqn')

        self.target_net = DuelingDQN(num_state, num_action)
        self.update_target()

        self.memo = Memory(capacity = 1000)

        self.optimizer = torch.optim.Adam(self.online_net.parameters(), lr= self.learning_rate)

    def act(self, state, epsilon):
        if random.random() > epsilon:
            state = torch.as_tensor(state).unsqueeze(0)
            q_value = self.online_net(state)
            # action = q_value.max(1)[1].item()
            max_q_idx = torch.argmax(input=q_value)
            action = max_q_idx.detach().item()
        else:
            action = random.randrange(self.num_action)
        return action

    def push(self, state, action, reward, next_state, done):
        # error最大，这个需要看论文确定
        error = max(self.memo.tree.tree[self.memo.capacity - 1: self.memo.tree.n_entries + self.memo.tree.capacity])
        self.memo.add(error = error, sample = (state, action, reward, next_state, done))

    def compute_td_loss(self):

        batch, idxs, is_weight = self.memo.sample(self.batch_size)
        batch_state, batch_action, batch_reward, batch_next_state, batch_done = zip(*batch)
        batch_state_tensor = torch.as_tensor(np.asarray(batch_state), dtype=torch.float32)
        batch_action_tensor = torch.as_tensor(np.asarray(batch_action), dtype=torch.int64).unsqueeze(-1)
        batch_reward_tensor = torch.as_tensor(np.asarray(batch_reward), dtype=torch.float32).unsqueeze(-1)  # unsqueeze 进行升维
        batch_done_tensor = torch.as_tensor(np.asarray(batch_done), dtype=torch.float32).unsqueeze(-1)
        batch_next_state_tensor = torch.as_tensor(np.asarray(batch_next_state), dtype=torch.float32)

        # compute double DQN , y_i
        next_q_values_online = self.online_net(batch_next_state_tensor)
        next_q_values_target = self.target_net(batch_next_state_tensor)

        next_q_value = next_q_values_target.gather(dim=1, index=torch.max(next_q_values_online, dim=1, keepdim=True)[1])

        pred_q_value = batch_reward_tensor + self.gamma * next_q_value * (1 - batch_done_tensor)

        # compute Q(s,a)
        q_values = self.online_net(batch_state_tensor)
        true_q_value = q_values.gather(1, batch_action_tensor)

        # update priority
        errors = torch.abs(pred_q_value - true_q_value).squeeze(dim = 1).detach().numpy()
        for i in range(self.batch_size):
            idx = idxs[i]
            self.memo.update(idx, errors[i])

        # compute loss
        loss = (true_q_value - pred_q_value.data).pow(2).mean()
        # loss = nn.functional.smooth_l1_loss(expected_q_value, q_value)

        # gradient descent
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss

    # update target net parameters
    def update_target(self):
        self.target_net.load_state_dict(self.online_net.state_dict())

class Environment:
    def __init__(self, env_id, render_mode):
        self.env = gym.make(env_id, render_mode = render_mode)

        self.num_action = self.env.action_space.n
        self.num_state = self.env.observation_space.shape[0]
        self.agent = Agent(self.num_state, self.num_action)
        self.updata_target_frame = 500
        self.n_episode = 500

    def train(self):
        eposide_rewards = []
        eposide_ids = []
        frame_losses = []
        frames = 0 # 记录所有的帧数
        for i_eposide in range(self.n_episode):
            done = False
            eposide_reward = 0
            # 初始化状态
            state, _ = self.env.reset()
            while not done:
                frames += 1
                epsilon = epsilon_by_frame(frames)
                action = self.agent.act(state, epsilon)
                next_state, reward, done, _, _ = self.env.step(action)

                # if an action make the episode end, then gives penalty of -100
                reward = reward if not done or eposide_reward == 499 else -10

                # add transitions to memory
                self.agent.push(state, action, reward, next_state, done)

                #compute td loss
                loss = self.agent.compute_td_loss()
                frame_losses.append(loss.item())

                # pylab.figure(1)
                # pylab.plot(frame_losses, 'r')
                # pylab.xlabel('step')
                # pylab.ylabel('loss')
                # pylab.savefig("./save_graph/cartpole_dqn_loss.png")

                eposide_reward += reward
                state = next_state

                if done:
                    eposide_reward = eposide_reward if eposide_reward == 500 else eposide_reward + 10
                    eposide_rewards.append(eposide_reward)
                    eposide_ids.append(i_eposide)

                    # plot the play time
                    # pylab.figure(2)
                    # pylab.plot(eposide_ids, eposide_rewards, 'b')
                    # pylab.xlabel('eposide')
                    # pylab.ylabel('eposide_reward')
                    # pylab.savefig("./save_graph/cartpole_dqn_reward.png")

                    print("episode:", i_eposide, "  eposids_reward:", eposide_reward, "  memory length:",
                          self.agent.memo.tree.n_entries, "  epsilon:", epsilon)

                    # if the mean of scores of last 10 episode is bigger than 490
                    # stop training
                    if np.mean(eposide_rewards[-min(10, len(eposide_rewards)):]) > 490:
                        torch.save(self.agent.online_net, "./save_model/cartpole_dqn")
                        sys.exit()

                if frames % self.updata_target_frame == 0:
                    self.agent.update_target()

if __name__ == '__main__':
    env_id = 'CartPole-v1'
    env = Environment(env_id, 'human')
    env.train()


import contextlib
import glob
import json
import nodetool
import os
import pytest
import subprocess
import tempfile
import random
import util

# To run the Scylla tools, we need to run Scylla executable itself, so we
# need to find the path of the executable that was used to run Scylla for
# this test. We do this by trying to find a local process which is listening
# to the address and port to which our our CQL connection is connected.
# If such a process exists, we verify that it is Scylla, and return the
# executable's path. If we can't find the Scylla executable we use
# pytest.skip() to skip tests relying on this executable.
@pytest.fixture(scope="module")
def scylla_path(cql):
    pid = util.local_process_id(cql)
    if not pid:
        pytest.skip("Can't find local Scylla process")
    # Now that we know the process id, use /proc to find the executable.
    try:
        path = os.readlink(f'/proc/{pid}/exe')
    except:
        pytest.skip("Can't find local Scylla executable")
    # Confirm that this executable is a real tool-providing Scylla by trying
    # to run it with the "--list-tools" option
    try:
        subprocess.check_output([path, '--list-tools'])
    except:
        pytest.skip("Local server isn't Scylla")
    return path

# A fixture for finding Scylla's data directory. We get it using the CQL
# interface to Scylla's configuration. Note that if the server is remote,
# the directory retrieved this way may be irrelevant, whether or not it
# exists on the local machine... However, if the same test that uses this
# fixture also uses the scylla_path fixture, the test will anyway be skipped
# if the running Scylla is not on the local machine local.
@pytest.fixture(scope="module")
def scylla_data_dir(cql):
    try:
        dir = json.loads(cql.execute("SELECT value FROM system.config WHERE name = 'data_file_directories'").one().value)[0]
        return dir
    except:
        pytest.skip("Can't find Scylla sstable directory")

def simple_no_clustering_table(cql, keyspace):
    table = util.unique_name()
    schema = f"CREATE TABLE {keyspace}.{table} (pk int PRIMARY KEY, v int) WITH compaction = {{'class': 'NullCompactionStrategy'}}"

    cql.execute(schema)

    for pk in range(0, 10):
        x = random.randrange(0, 4)
        if x == 0:
            # partition tombstone
            cql.execute(f"DELETE FROM {keyspace}.{table} WHERE pk = {pk}")
        else:
            # live row
            cql.execute(f"INSERT INTO {keyspace}.{table} (pk, v) VALUES ({pk}, 0)")

        if pk == 5:
            nodetool.flush(cql, f"{keyspace}.{table}")

    nodetool.flush(cql, f"{keyspace}.{table}")

    return table, schema

def simple_clustering_table(cql, keyspace):
    table = util.unique_name()
    schema = f"CREATE TABLE {keyspace}.{table} (pk1 int, pk2 int, ck1 int, ck2 int, v int, s int STATIC, PRIMARY KEY ((pk1, pk2), ck1, ck2)) WITH compaction = {{'class': 'NullCompactionStrategy'}}"

    cql.execute(schema)

    for pk in range(0, 10):
        for ck in range(0, 10):
            x = random.randrange(0, 8)
            if x == 0:
                # ttl
                cql.execute(f"INSERT INTO {keyspace}.{table} (pk1, pk2, ck1, ck2, v) VALUES ({pk}, {pk}, {ck}, {ck}, 0) USING TTL 6000")
            elif x == 1:
                # row tombstone
                cql.execute(f"DELETE FROM {keyspace}.{table} WHERE pk1 = {pk} AND pk2 = {pk} AND ck1 = {ck} AND ck2 = {ck}")
            elif x == 2:
                # cell tombstone
                cql.execute(f"DELETE v FROM {keyspace}.{table} WHERE pk1 = {pk} AND pk2 = {pk} AND ck1 = {ck} AND ck2 = {ck}")
            elif x == 3:
                # range tombstone
                l = ck * 10
                u = ck * 11
                cql.execute(f"DELETE FROM {keyspace}.{table} WHERE pk1 = {pk} AND pk2 = {pk} AND ck1 > {l} AND ck1 < {u}")
            else:
                # live row
                cql.execute(f"INSERT INTO {keyspace}.{table} (pk1, pk2, ck1, ck2, v) VALUES ({pk}, {pk}, {ck}, {ck}, 0)")

        if pk == 5:
            cql.execute(f"UPDATE {keyspace}.{table} SET s = 10 WHERE pk1 = {pk} AND pk2 = {pk}")
            nodetool.flush(cql, f"{keyspace}.{table}")

    nodetool.flush(cql, f"{keyspace}.{table}")

    return table, schema

def clustering_table_with_collection(cql, keyspace):
    table = util.unique_name()
    schema = f"CREATE TABLE {keyspace}.{table} (pk int, ck int, v1 map<int, text>, v2 set<int>, v3 list<int>, PRIMARY KEY (pk, ck)) WITH compaction = {{'class': 'NullCompactionStrategy'}}"

    cql.execute(schema)

    for pk in range(0, 10):
        for ck in range(0, 10):
            map_vals = {f"{p}: '{c}'" for p in range(0, pk) for c in range(0, ck)}
            map_str = ", ".join(map_vals)
            set_list_vals = list(range(0, pk))
            set_list_str = ", ".join(map(str, set_list_vals))
            cql.execute(f"INSERT INTO {keyspace}.{table} (pk, ck, v1, v2, v3) VALUES ({pk}, {ck}, {{{map_str}}}, {{{set_list_str}}}, [{set_list_str}])")
        if pk == 5:
            nodetool.flush(cql, f"{keyspace}.{table}")

    nodetool.flush(cql, f"{keyspace}.{table}")

    return table, schema

def clustering_table_with_udt(cql, keyspace):
    table = util.unique_name()
    create_type_schema = f"CREATE TYPE IF NOT EXISTS {keyspace}.type1 (f1 int, f2 text)"
    create_table_schema = f" CREATE TABLE {keyspace}.{table} (pk int, ck int, v type1, PRIMARY KEY (pk, ck)) WITH compaction = {{'class': 'NullCompactionStrategy'}}"

    cql.execute(create_type_schema)
    cql.execute(create_table_schema)

    for pk in range(0, 10):
        for ck in range(0, 10):
            cql.execute(f"INSERT INTO {keyspace}.{table} (pk, ck, v) VALUES ({pk}, {ck}, {{f1: 100, f2: 'asd'}})")
        if pk == 5:
            nodetool.flush(cql, f"{keyspace}.{table}")

    nodetool.flush(cql, f"{keyspace}.{table}")

    return table, "; ".join((create_type_schema, create_table_schema))

def table_with_counters(cql, keyspace):
    table = util.unique_name()
    schema = f"CREATE TABLE {keyspace}.{table} (pk int PRIMARY KEY, v counter) WITH compaction = {{'class': 'NullCompactionStrategy'}}"

    cql.execute(schema)

    for pk in range(0, 10):
        for c in range(0, 4):
            cql.execute(f"UPDATE {keyspace}.{table} SET v = v + 1 WHERE pk = {pk};")
        if pk == 5:
            nodetool.flush(cql, f"{keyspace}.{table}")

    nodetool.flush(cql, f"{keyspace}.{table}")

    return table, schema

@contextlib.contextmanager
def scylla_sstable(table_factory, cql, ks, data_dir):
    table, schema = table_factory(cql, ks)

    schema_file = os.path.join(data_dir, "..", "test_tools_schema.cql")
    with open(schema_file, "w") as f:
        f.write(schema)

    sstables = glob.glob(os.path.join(data_dir, ks, table + '-*', '*-Data.db'))

    try:
        yield (schema_file, sstables)
    finally:
        cql.execute(f"DROP TABLE {ks}.{table}")
        os.unlink(schema_file)

def one_sstable(sstables):
    assert len(sstables) > 1
    return [sstables[0]]

def all_sstables(sstables):
    assert len(sstables) > 1
    return sstables

@pytest.mark.parametrize("what", ["index", "compression-info", "summary", "statistics", "scylla-metadata"])
@pytest.mark.parametrize("which_sstables", [one_sstable, all_sstables])
def test_scylla_sstable_dump_component(cql, test_keyspace, scylla_path, scylla_data_dir, what, which_sstables):
    with scylla_sstable(simple_clustering_table, cql, test_keyspace, scylla_data_dir) as (schema_file, sstables):
        out = subprocess.check_output([scylla_path, "sstable", f"dump-{what}", "--schema-file", schema_file] + which_sstables(sstables))

    print(out)

    assert out
    assert json.loads(out)

@pytest.mark.parametrize("table_factory", [
        simple_no_clustering_table,
        simple_clustering_table,
        clustering_table_with_collection,
        clustering_table_with_udt,
        table_with_counters,
])
@pytest.mark.parametrize("merge", [True, False])
@pytest.mark.parametrize("output_format", ["text", "json"])
def test_scylla_sstable_dump_data(cql, test_keyspace, scylla_path, scylla_data_dir, table_factory, merge, output_format):
    with scylla_sstable(simple_clustering_table, cql, test_keyspace, scylla_data_dir) as (schema_file, sstables):
        args = [scylla_path, "sstable", "dump-data", "--schema-file", schema_file, "--output-format", output_format]
        if merge:
            args.append("--merge")
        out = subprocess.check_output(args + sstables)

    print(out)

    assert out
    if output_format == "json":
        assert json.loads(out)

@pytest.mark.parametrize("table_factory", [
        simple_no_clustering_table,
        simple_clustering_table,
])
def test_scylla_sstable_write(cql, test_keyspace, scylla_path, scylla_data_dir, table_factory):
    with scylla_sstable(table_factory, cql, test_keyspace, scylla_data_dir) as (schema_file, sstables):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dump_common_args = [scylla_path, "sstable", "dump-data", "--schema-file", schema_file, "--output-format", "json", "--merge"]
            generation = util.unique_key_int()

            original_out = subprocess.check_output(dump_common_args + sstables)
            original_json = json.loads(original_out)["sstables"]["anonymous"]

            input_file = os.path.join(tmp_dir, 'input.json')

            with open(input_file, 'w') as f:
                json.dump(original_json, f)

            subprocess.check_call([scylla_path, "sstable", "write", "--schema-file", schema_file, "--input-file", input_file, "--output-dir", tmp_dir, "--generation", str(generation), '--logger-log-level', 'scylla-sstable=trace'])

            sstable_file = os.path.join(tmp_dir, f"me-{generation}-big-Data.db")

            actual_out = subprocess.check_output(dump_common_args + [sstable_file])
            actual_json = json.loads(actual_out)["sstables"]["anonymous"]

            assert actual_json == original_json

def script_consume_test_table_factory(cql, keyspace):
    table = util.unique_name()
    schema = f"CREATE TABLE {keyspace}.{table} (pk int, ck int, v int, s int STATIC, PRIMARY KEY (pk, ck)) WITH compaction = {{'class': 'NullCompactionStrategy'}}"

    cql.execute(schema)

    partitions = 4

    for sst in range(0, 2):
        for pk in range(sst * partitions, (sst + 1) * partitions):
            # static row
            cql.execute(f"UPDATE {keyspace}.{table} SET s = 10 WHERE pk = {pk}")
            # range tombstone
            cql.execute(f"DELETE FROM {keyspace}.{table} WHERE pk = {pk} AND ck >= 0 AND ck <= 4")
            # 2 rows
            for ck in range(0, 4):
                cql.execute(f"INSERT INTO {keyspace}.{table} (pk, ck, v) VALUES ({pk}, {ck}, 0)")

        nodetool.flush(cql, f"{keyspace}.{table}")

    return table, schema

def test_scylla_sstable_script_consume_sstable(cql, test_keyspace, scylla_path, scylla_data_dir):
    script_file = os.path.join(scylla_data_dir, "..", "test_scylla_sstable_script_consume_sstable.lua")

    script = """
wr = Scylla.new_json_writer()
i = 0

function arg(args, arg)
    ret = nil
    wr:key(arg)
    if args[arg] then
        ret = tonumber(args[arg])
        wr:int(ret)
    else
        wr:null()
    end
    return ret
end

function basename(path)
    s, e = string.find(string.reverse(path), '/', 1, true)
    return string.sub(path, #path - s + 2)
end

function consume_stream_start(args)
    wr:start_object()
    start_sst = arg(args, "start_sst")
    end_sst = arg(args, "end_sst")
    wr:key("content")
    wr:start_array()
end

function consume_sstable_start(sst)
    wrote_ps = false
    i = i + 1
    if i == start_sst then
        return false
    end
    wr:string(basename(sst.filename))
end

function consume_partition_start(ps)
    if not wrote_ps then
        wr:string("ps")
        wrote_ps = true
    end
end

function consume_sstable_end()
    if i == end_sst then
        return false
    end
end

function consume_stream_end()
    wr:end_array()
    wr:end_object()
end
"""
    with open(script_file, 'w') as f:
        f.write(script)

    with scylla_sstable(script_consume_test_table_factory, cql, test_keyspace, scylla_data_dir) as (schema_file, sstables):
        sst1 = os.path.basename(sstables[0])
        sst2 = os.path.basename(sstables[1])
        def run_scenario(script_args, expected):
            print(f"Scenario: '{script_args}'\n")
            if script_args:
                script_args = ["--script-arg", script_args]
            else:
                script_args = []
            script_args = [scylla_path, "sstable", "script", "--schema-file", schema_file, "--script-file", script_file] + script_args + sstables[0:2]
            res = json.loads(subprocess.check_output(script_args))
            assert res == expected

        run_scenario("", {'start_sst': None, 'end_sst': None, 'content': [sst1, "ps", sst2, "ps"]})
        run_scenario("start_sst=1", {'start_sst': 1, 'end_sst': None, 'content': [sst2, "ps"]})
        run_scenario("start_sst=2", {'start_sst': 2, 'end_sst': None, 'content': [sst1, "ps"]})
        run_scenario("start_sst=1:end_sst=1", {'start_sst': 1, 'end_sst': 1, 'content': []})
        run_scenario("start_sst=2:end_sst=2", {'start_sst': 2, 'end_sst': 2, 'content': [sst1, "ps"]})
        run_scenario("end_sst=1", {'start_sst': None, 'end_sst': 1, 'content': [sst1, "ps"]})
        run_scenario("end_sst=2", {'start_sst': None, 'end_sst': 2, 'content': [sst1, "ps", sst2, "ps"]})

def test_scylla_sstable_script_slice(cql, test_keyspace, scylla_path, scylla_data_dir):
    class bound:
        @staticmethod
        def unpack_value(value):
            if isinstance(value, tuple):
                return value
            else:
                return None, value

        def __init__(self, value, weight):
            self.token, self.value = self.unpack_value(value)
            self.weight = weight

        def tri_cmp(self, value):
            if self.token is None and self.value is None:
                assert(self.weight)
                return -self.weight
            token, value = self.unpack_value(value)
            if token is None:
                res = 0
            else:
                res = int(token) - int(self.token)
            if res == 0 and not value is None and not self.value is None :
                res = int(value) - int(self.value)
            return res if res else -self.weight

        def get_value(self, lookup_table, is_start):
            if self.token is None and self.value is None:
                return '-inf' if is_start else '+inf'
            if self.value is None:
                return "t{}".format(int(self.token))
            return lookup_table[self.value]

        @staticmethod
        def before(value):
            return bound(value, -1)

        @staticmethod
        def at(value):
            return bound(value, 0)

        @staticmethod
        def after(value):
            return bound(value, 1)

    class interval:
        def __init__(self, start_bound, end_bound):
            self.start = start_bound
            self.end = end_bound

        def contains(self, value):
            return self.start.tri_cmp(value) >= 0 and self.end.tri_cmp(value) <= 0

    def summarize_dump(dump):
        summary = []
        for partition in list(dump["sstables"].items())[0][1]:
            partition_summary = {"pk": partition["key"]["value"], "token": partition["key"]["token"], "frags": []}
            if "static_row" in partition:
                partition_summary["frags"].append(("sr", None))
            for clustering_fragment in partition.get("clustering_elements", []):
                type_str = "cr" if clustering_fragment["type"] == "clustering-row" else "rtc"
                partition_summary["frags"].append((type_str, clustering_fragment["key"]["value"]))
            summary.append(partition_summary)
        return summary

    def filter_summary(summary, partition_ranges, clustering_ranges):
        if not partition_ranges:
            return summary
        filtered_summary = []
        for partition in summary:
            if any(map(lambda x: interval.contains(x, (partition["token"], partition["pk"])), partition_ranges)):
                filtered_summary.append({"pk": partition["pk"], "token": partition["token"], "frags": []})
                for (t, k) in partition["frags"]:
                    if t == "rtc" or k is None or not clustering_ranges or any(map(lambda x: interval.contains(x, k), clustering_ranges)):
                        filtered_summary[-1]["frags"].append((t, k))

        return filtered_summary

    def serialize_ranges(prefix, ranges, lookup_table):
        serialized_ranges = []
        i = 0
        for r in ranges:
            s = r.start.get_value(lookup_table, True)
            e = r.end.get_value(lookup_table, False)
            serialized_ranges.append("{}{}={}{},{}{}".format(
                prefix,
                i,
                "(" if s == '-inf' or r.start.weight > 0 else "[",
                s,
                e,
                ")" if e == '+inf' or r.end.weight < 0 else "]"))
            i = i + 1
        return serialized_ranges

    scripts_path = os.path.realpath(os.path.join(__file__, '../../../tools/scylla-sstable-scripts'))
    script_file = os.path.join(scripts_path, 'slice.lua')

    with scylla_sstable(script_consume_test_table_factory, cql, test_keyspace, scylla_data_dir) as (schema_file, sstables):
        reference_summary = summarize_dump(json.loads(subprocess.check_output([scylla_path, "sstable", "dump-data", "--schema-file", schema_file, "--merge"] + sstables)))

        # same order as in dump
        pks = [(p["token"], p["pk"]) for p in reference_summary]
        cks = set()
        for p in reference_summary:
            for t, ck in p["frags"]:
                if not ck is None:
                    cks.add(ck)
        cks = sorted(list(cks))
        serialized_pk_lookup = {pk: subprocess.check_output([scylla_path, "types", "serialize", "--full-compound", "-t", "Int32Type", "--", pk]).strip().decode() for t, pk in pks}
        serialized_ck_lookup = {ck: subprocess.check_output([scylla_path, "types", "serialize", "--prefix-compound", "-t", "Int32Type", "--", ck]).strip().decode() for ck in cks}

        script_common_args = [scylla_path, "sstable", "script", "--schema-file", schema_file, "--merge", "--script-file", script_file]

        def run_scenario(scenario, partition_ranges, clustering_ranges):
            print(f"running scenario {scenario}")
            script_args = serialize_ranges("pr", partition_ranges, serialized_pk_lookup) + serialize_ranges("cr", clustering_ranges, serialized_ck_lookup)
            if script_args:
                script_args = ["--script-arg"] + [":".join(script_args)]
            print(f"script_args={script_args}")
            expected = filter_summary(reference_summary, partition_ranges, clustering_ranges)
            out = subprocess.check_output(script_common_args + script_args + sstables)
            summary = summarize_dump(json.loads(out))
            assert summary == expected

        run_scenario("no args", [], [])
        run_scenario("full range", [interval(bound.before(None), bound.after(None))], [])
        run_scenario("(pks[0], +inf)", [interval(bound.after(pks[0]), bound.after(None))], [])
        run_scenario("(-inf, pks[-3]]", [interval(bound.before(None), bound.after(pks[-3]))], [])
        run_scenario("[pks[2], pks[-2]]", [interval(bound.before(pks[2]), bound.after(pks[-2]))], [])
        run_scenario("[pks[0], pks[1]], [pks[2], pks[3]]", [interval(bound.before(pks[1]), bound.after(pks[2])), interval(bound.before(pks[3]), bound.after(pks[4]))], [])
        run_scenario("[t:pks[2], t:pks[-2]]", [interval(bound.before((pks[2][0], None)), bound.after((pks[-2][0], None)))], [])
        run_scenario("full pk range | [-inf, cks[2]]", [interval(bound.before(None), bound.after(None))], [interval(bound.before(None), bound.after(cks[2]))])
        run_scenario("[pks[0], pks[1]] | (cks[0], cks[1]], (cks[2], +inf)", [interval(bound.before(pks[1]), bound.after(pks[2]))],
                     [interval(bound.after(cks[0]), bound.after(cks[1])), interval(bound.after(cks[2]), bound.after(None))])

@pytest.mark.parametrize("table_factory", [
        simple_no_clustering_table,
        simple_clustering_table,
        clustering_table_with_collection,
        clustering_table_with_udt,
        table_with_counters,
])
def test_scylla_sstable_script(cql, test_keyspace, scylla_path, scylla_data_dir, table_factory):
    scripts_path = os.path.realpath(os.path.join(__file__, '../../../tools/scylla-sstable-scripts'))
    slice_script_path = os.path.join(scripts_path, 'slice.lua')
    dump_script_path = os.path.join(scripts_path, 'dump.lua')
    with scylla_sstable(table_factory, cql, test_keyspace, scylla_data_dir) as (schema_file, sstables):
        dump_common_args = [scylla_path, "sstable", "dump-data", "--schema-file", schema_file, "--output-format", "json"]
        script_common_args = [scylla_path, "sstable", "script", "--schema-file", schema_file]

        # without --merge
        cxx_json = json.loads(subprocess.check_output(dump_common_args + sstables))
        dump_lua_json = json.loads(subprocess.check_output(script_common_args + ["--script-file", dump_script_path] + sstables))
        slice_lua_json = json.loads(subprocess.check_output(script_common_args + ["--script-file", slice_script_path] + sstables))

        assert dump_lua_json == cxx_json
        assert slice_lua_json == cxx_json

        # with --merge
        cxx_json = json.loads(subprocess.check_output(dump_common_args + ["--merge"] + sstables))
        dump_lua_json = json.loads(subprocess.check_output(script_common_args + ["--merge", "--script-file", dump_script_path] + sstables))
        slice_lua_json = json.loads(subprocess.check_output(script_common_args + ["--merge", "--script-file", slice_script_path] + sstables))

        assert dump_lua_json == cxx_json
        assert slice_lua_json == cxx_json


#!/usr/env python3
# -*- coding: utf-8 -*- #

## source: http://nafiulis.me/making-a-static-blog-with-pelican.html#writing-content
import sys
from datetime import datetime

TEMPLATE = """
{title}
{hashes}

:date: {year}-{month}-{day} {hour}:{minute:02d}
:tags:
:category:
:slug: {slug}
:summary:
:status: draft

"""

def make_entry(title):
    today = datetime.today()
    slug = title.lower().strip().replace(' ', '-')

    f_create = "content/{}_{:0>2}_{:0>2}_{}.rst".format(
        today.year, today.month, today.day, slug)
    t = TEMPLATE.strip().format(title=title,
                                hashes='#' * len(title),
                                year=today.year,
                                month=today.month,
                                day=today.day,
                                hour=today.hour,
                                minute=today.minute,
                                slug=slug)
    with open(f_create, 'w') as w:
        w.write(t)
    print("File created -> " + f_create)

if __name__ == '__main__':

    if len(sys.argv) > 1:
        make_entry(sys.argv[1])
    else:
        print("No title given")


import yaml
import re
import sys
import os

if len(sys.argv) != 4:
    sys.stderr.write("Usage: "+sys.argv[0]+" <controllerfile> <computefile> <outfile>\n")
    sys.exit(1)

controller = sys.argv[1]
if not os.path.exists(controller):
    sys.stderr.write("ERROR: The file "+controller+" could not be opened\n")
    sys.exit(1)

compute = sys.argv[2]
if not os.path.exists(compute):
    sys.stderr.write("ERROR: The file "+compute+" could not be opened\n")
    sys.exit(1)

outfile = sys.argv[3]

f_controller = open(controller, 'r')
doc_controller = yaml.load(f_controller)
f_controller.close()

f_compute = open(compute, 'r')
doc_compute = yaml.load(f_compute)
f_compute.close()

out = {}
out["opnfv"] = {}
out["opnfv"]["controller"] = doc_controller["opnfv"]
out["opnfv"]["compute"] = doc_compute["opnfv"]

f2 = open(outfile, 'a')
f2.write(yaml.dump(out, default_flow_style=False))
f2.close()