text:
import math
import numpy as np
import pandas as pd
import csv
import sys
import os
from datetime import datetime

from sklearn import metrics
from sklearn.model_selection import cross_validate
from sklearn.linear_model import BayesianRidge, Lasso, Ridge, LinearRegression, HuberRegressor
from sklearn.neural_network import MLPRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.kernel_ridge import KernelRidge
from sklearn.svm import SVR
from sklearn.metrics import mean_squared_error, r2_score

#custom class import

try:
    data_input = sys.argv[1]
except:
    print ("Please include input filename.")
    print ("USAGE: python3 train_models.py input_data.csv")
    exit()

if not os.path.isfile(data_input):
    print ("File " + data_input + " not found.")
    print ("USAGE: python3 train_models.py input_data.csv")
    exit()

from bestModel import BestModelObject

#some global variables here
num_of_cross_val_folds = 5

data = pd.read_csv(data_input).values #create numpy array from csv sheet
#create training and test set (80/20 split)
total_samples = data.shape[0]
train_test_index = math.floor( 0.8 * total_samples ) # index to divide 80/20 train/test data split (must be integer)
train_data = data[:train_test_index]
test_data = data[train_test_index:]

#randomize training data (do this before before we make cross validation folds)
np.random.shuffle(train_data)

#create numpy arrays of training and test data
X = train_data[:,:3]#first 20 rows and first 3 columns(remember zero based indexing) are X
y = train_data[:,3]# first 20 rows and 4th column are y
#we'll use these to test our final model.
X_test = test_data[:,:3]
y_test = test_data[:,3]

n, d = X.shape
t = X_test.shape[0]

print ("Training inputs shape: " + str(X.shape))
print ("Training labels shape: " + str(y.shape))
print ("Test inputs shape:     " + str(X_test.shape))
print ("Test labels shape:     " + str(y_test.shape))

#this function returns a list of mean scores
def evaluate_model(model, name, ):

    #let's try doing 5 fold vs 4 fold
    #lets try leave on out? might overfit?

    scores = cross_validate(model, X, y, cv=num_of_cross_val_folds, return_train_score=True, scoring=('r2','neg_mean_squared_error','explained_variance'))

    #note
    #scores['train_neg_mean_squared_error'] will return an array of values, each element will correspond to the train MSE for a given fold. so 4 folds will get you 4 elements in this array

    mse_train = np.mean(-scores['train_neg_mean_squared_error'])
    mse_test  = np.mean(-scores['test_neg_mean_squared_error'])#this is more the validation score than test score (test on validation)
    #explantory power of our model. if 0.6 = 60% of the variation in y is explained by the model. the greater proportion that is explained means the greater explanotry power the model has
    r2_train  = np.mean(scores['train_r2'])# if model has small r2, then it has low explantory power, if r2 is high, above 70% this has really good explanatory power.
    var_test  = np.mean(scores['test_explained_variance'])

    #add the above values to an list and return the list
    final_scores_to_return = [mse_train, mse_test, r2_train, var_test]

    return final_scores_to_return

#each model we want to test can have a set of hyperparameters we want to test as well. so we encapsulate each model with its hyperparameters in their own functions
def tune_linear_regression_hyperparameters():
    #hyperparameters --> none

    myModel = LinearRegression()#should we set the 'normalize" parameter to true?, default is false

    #we get back a list of the scores
    myScores = evaluate_model(myModel, 'Linear Regression') # Least squares loss with L1 reg.

    #trackers for best model and its scores
    best_model = myModel
    best_score = myScores
    best_hyperparameters =[]

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

def tune_lasso_regression_hyperparameters():
    #hyperparameters
    #depending on what get's picked as the best alpha, maybe we should consider another pass with a smaller window around the best value.
    alpha = [1e-15, 1e-10, 1e-8, 1e-5, 1e-4, 1e-3, 1e-2, 1, 5, 10]#range from 0 to infinity, don't use 0 - use plain linear regression if this is the case
    #alpha2 = []#if make a new set of values centered around the best from above?
    tol = [0.01, .001, 0.0001, .00001]
    positive = [True, False]

    #trackers for best model and its scores
    best_model = None
    best_model_scores = None
    best_hyperparameters = []

    run_once_flag = False
    for alpha_element in alpha:
        for tol_element in tol:
            for positive_element in positive:
                myModel = Lasso(alpha=alpha_element, tol=tol_element, positive=positive_element)#should we set the 'normalize" parameter to true?, default is false

                #we get back a list of the scores
                myScores = evaluate_model(myModel, 'Lasso Regression') # Least squares loss with L1 reg.

                #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                if run_once_flag == False:
                    best_model = myModel
                    best_score = myScores
                    best_hyperparameters.append(alpha_element)
                    best_hyperparameters.append(tol_element)
                    best_hyperparameters.append(positive_element)
                    run_once_flag = True

                #check if we have a better model based on validaiton MSE score, and update if we do
                if myScores[1] < best_score[1]: #we want to check the validation MSE
                    best_model = myModel
                    best_score = myScores
                    best_hyperparameters =[]#clear any old ones
                    best_hyperparameters.append(alpha_element)
                    best_hyperparameters.append(tol_element)
                    best_hyperparameters.append(positive_element)

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

def tune_ridge_regression_hyperparameters():
    #hyperparameters
    alpha = [1e-15, 1e-10, 1e-8, 1e-4, 1e-3, 1e-2, 1, 5, 10, 20]
    tol = [0.01, .001, 0.0001, .00001]
    solver = ['auto', 'svd', 'cholesky', 'lsqr', 'sparse_cg', 'sag', 'saga']

    #trackers for best model and its scores
    best_model = None
    best_model_scores = None
    best_hyperparameters = []

    run_once_flag = False
    for alpha_element in alpha:
        for tol_element in tol:
            for solver_element in solver:

                myModel = Ridge(alpha=alpha_element, tol=tol_element, solver=solver_element)#should we set the 'normalize" parameter to true?, default is false

                #we get back a list of the scores
                myScores = evaluate_model(myModel, 'Ridge Regression') # Least squares loss with L2 reg.

                #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                if run_once_flag == False:
                    best_model = myModel
                    best_score = myScores
                    best_hyperparameters.append(alpha_element)
                    best_hyperparameters.append(tol_element)
                    best_hyperparameters.append(solver_element)
                    run_once_flag = True

                #check if we have a better model based on validaiton MSE score, and update if we do
                if myScores[1] < best_score[1]: #we want the validation MSE
                    best_model = myModel
                    best_score = myScores
                    best_hyperparameters =[]#clear any old ones
                    best_hyperparameters.append(alpha_element)
                    best_hyperparameters.append(tol_element)
                    best_hyperparameters.append(solver_element)

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

def tune_huber_regression_hyperparameters():
    #hyperparameters
    alpha = [1e-15, 1e-10, 1e-8, 1e-4, 1e-3,1e-2, 1, 5, 10, 20]
    #eplsion: The parameter epsilon controls the number of samples that should be classified as outliers. The smaller the epsilon, the more robust it is to outliers
    #epsilon = list(range(1.0,6.0,0.5))#float, greater than 1.0, default 1.35 (do 1 to 5)
    epsilon = np.append(np.linspace(1,5,9),[1.35])
    tol = [0.01, .001, 0.0001, .00001]

    #trackers for best model and its scores
    best_model = None
    best_model_scores = None
    best_hyperparameters = []

    run_once_flag = False
    for alpha_element in alpha:
        for epsilon_element in epsilon:
            for tol_element in tol:

                myModel = HuberRegressor(alpha=alpha_element, epsilon=epsilon_element, tol=tol_element)#should we set the 'normalize" parameter to true?, default is false

                #we get back a list of the scores
                myScores = evaluate_model(myModel, 'Huber Regression') # Least squares loss with L2 reg.

                #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                if run_once_flag == False:
                    best_model = myModel
                    best_score = myScores
                    best_hyperparameters.append(alpha_element)
                    best_hyperparameters.append(epsilon_element)
                    best_hyperparameters.append(tol_element)
                    run_once_flag = True

                #check if we have a better model based on validaiton MSE score, and update if we do
                if myScores[1] < best_score[1]: #we want the validation MSE
                    best_model = myModel
                    best_score = myScores
                    best_hyperparameters =[]#clear any old ones
                    best_hyperparameters.append(alpha_element)
                    best_hyperparameters.append(epsilon_element)
                    best_hyperparameters.append(tol_element)

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object, bestModelObjec is just a custom class that acts as a container
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

def tune_bayesian_ridge_regression_hyperparameters():
    #hyperparameters
    alpha1 = [1e-08, 1e-07,1e-06, 1e-05, 1e-04]#for to -8 to -4 in steps of 1
    lambda1 = [1e-08, 1e-07,1e-06, 1e-05, 1e-04]#for to -8 to -4 in steps of 1
    alpha2 = [1e-08, 1e-07,1e-06, 1e-05, 1e-04]#for to -8 to -4 in steps of 1
    lambda2 = [1e-08, 1e-07,1e-06, 1e-05, 1e-04]
    n_iter = list(range(200,401,50))#n_iter go from 200 to 400 in 50 size steps

    #trackers for best model and its scores
    best_model = None
    best_model_scores = None
    best_hyperparameters = []

    run_once_flag = False
    for alpha1_element in alpha1:
        for lambda1_element in lambda1:
            for alpha2_element in alpha2:
                for lambda2_element in lambda2:
                    for n_iter_element in n_iter:

                        myModel = BayesianRidge(alpha_1=alpha1_element, lambda_1=lambda1_element, alpha_2=alpha2_element, lambda_2=lambda2_element, n_iter=n_iter_element)#should we set the 'normalize" parameter to true?, default is false

                        #we get back a list of the scores
                        myScores = evaluate_model(myModel, 'Bayesian Ridge Regression') # Least squares loss with L2 reg.

                        #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                        if run_once_flag == False:
                            best_model = myModel
                            best_score = myScores
                            best_hyperparameters.append(alpha1_element)
                            best_hyperparameters.append(lambda1_element)
                            best_hyperparameters.append(alpha2_element)
                            best_hyperparameters.append(lambda2_element)
                            best_hyperparameters.append(n_iter_element)
                            run_once_flag = True

                        #check if we have a better model based on validaiton MSE score, and update if we do
                        if myScores[1] < best_score[1]: #we want the validation MSE
                            best_model = myModel
                            best_score = myScores
                            best_hyperparameters =[]#clear any old ones
                            best_hyperparameters.append(alpha1_element)
                            best_hyperparameters.append(lambda1_element)
                            best_hyperparameters.append(alpha2_element)
                            best_hyperparameters.append(lambda2_element)
                            best_hyperparameters.append(n_iter_element)

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

def tune_neural_network_regression_hyperparameters():
    #hyperparameters
    hidden_layer_sizes = []
    for i in range(1,2,1):#we want 1 to 3 layers
        for j in range(10,101,10):#layers from size 10 to 100 in steps of 10
            myList = [j] * i

            # if i == 1:
            #     myList = [s + ',' for s in map(str, myList)]

            #myString = ','.join(map(str, myList))
            myTuple = tuple(myList)
            hidden_layer_sizes.append(myTuple)

    activation = ['identity', 'logistic', 'tanh', 'relu']
    sovler = ['lbfgs']
    alpha = [1e-15, 1e-10, 1e-8, 1e-4, 1e-3, 1e-2, 1, 5, 10, 20]
    learning_rate = ['constant', 'invscaling', 'adaptive']
    #learning_rate_init = [0.001] dont use this unless we use sdg as our solver, but i don't think were doing that for such a small dataset
    tol = [0.01, .001, 0.0001, .00001]

    #trackers for best model and its scores
    best_model = None
    best_model_scores = None
    best_hyperparameters = []

    run_once_flag = False
    for activation_element in activation:
        for solver_element in sovler:
            for alpha_element in alpha:
                for learning_rate_element in learning_rate:
                    for tol_element in tol:
                        for hidden_element in hidden_layer_sizes:
                            #default is 1 hidden layer
                            myModel = MLPRegressor(activation=activation_element, solver=solver_element, alpha=alpha_element, learning_rate=learning_rate_element, tol=tol_element, hidden_layer_sizes=hidden_element)

                            #we get back a list of the scores
                            myScores = evaluate_model(myModel, 'Neural Network Regression') # Least squares loss with L2 reg.

                            #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                            if run_once_flag == False:
                                best_model = myModel
                                best_score = myScores
                                best_hyperparameters.append(activation_element)
                                best_hyperparameters.append(solver_element)
                                best_hyperparameters.append(alpha_element)
                                best_hyperparameters.append(learning_rate_element)
                                best_hyperparameters.append(tol_element)
                                best_hyperparameters.append(hidden_element)
                                run_once_flag = True

                            #check if we have a better model based on validaiton MSE score, and update if we do
                            if myScores[1] < best_score[1]: #we want the validation MSE
                                best_model = myModel
                                best_score = myScores
                                best_hyperparameters =[]#clear any old ones
                                best_hyperparameters.append(activation_element)
                                best_hyperparameters.append(solver_element)
                                best_hyperparameters.append(alpha_element)
                                best_hyperparameters.append(learning_rate_element)
                                best_hyperparameters.append(tol_element)
                                best_hyperparameters.append(hidden_element)

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

def tune_random_forest_regression_hyperparameters():
    #hyperparameters
    #this has bootstrp sampling enabled by default, each tree gets its own sample
    n_estimators = list(range(1,26,1))#The number of trees in the forest.
    max_depth = list(range(5,26,1))#Max depth of the tree
    max_features = list(range(1,d+1,1))

    #trackers for best model and its scores
    best_model = None
    best_model_scores = None
    best_hyperparameters = []

    run_once_flag = False
    for n_estimators_element in n_estimators:
        #for max_depth_element in max_depth:
        for max_features_element in max_features:

            myModel = RandomForestRegressor(n_estimators=n_estimators_element, max_features=max_features_element)#should we set the 'normalize" parameter to true?, default is false

            #we get back a list of the scores
            myScores = evaluate_model(myModel, 'Random Forest Regression') # Least squares loss with L2 reg.

            #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
            if run_once_flag == False:
                best_model = myModel
                best_score = myScores
                best_hyperparameters.append(n_estimators_element)
                #best_hyperparameters.append(max_depth_element)
                best_hyperparameters.append(max_features_element)
                run_once_flag = True

            #check if we have a better model based on validaiton MSE score, and update if we do
            if myScores[1] < best_score[1]: #we want the validation MSE
                best_model = myModel
                best_score = myScores
                best_hyperparameters =[]#clear any old ones
                best_hyperparameters.append(n_estimators_element)
                #best_hyperparameters.append(max_depth_element)
                best_hyperparameters.append(max_features_element)

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

def tune_kernel_ridge_regression_hyperparameters():
    #hyperparameters
    alpha = [1e-15, 1e-10, 1e-8, 1e-4, 1e-3, 1e-2, 1, 5, 10, 20]
    kernel = ['linear', 'poly', 'sigmoid', 'rbf', 'laplacian']
    degree = [2, 3, 4]#this will only be used if kernel is poly

    #trackers for best model and its scores
    best_model = None
    best_model_scores = None
    best_hyperparameters = []

    run_once_flag = False
    for alpha_element in alpha:
        for kernel_element in kernel:
            if kernel_element == 'poly':#use the degree
                for degree_element in degree:
                    myModel = KernelRidge(alpha=alpha_element, kernel=kernel_element, degree=degree_element)#should we set the 'normalize" parameter to true?, default is false

                    #we get back a list of the scores
                    myScores = evaluate_model(myModel, 'Kernel Ridge Regression') # Least squares loss with L2 reg.

                    #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                    if run_once_flag == False:
                        best_model = myModel
                        best_score = myScores
                        best_hyperparameters.append(alpha_element)
                        best_hyperparameters.append(kernel_element)
                        best_hyperparameters.append(degree_element)
                        run_once_flag = True

                    #check if we have a better model based on validaiton MSE score, and update if we do
                    if myScores[1] < best_score[1]: #we want the validation MSE
                        best_model = myModel
                        best_score = myScores
                        best_hyperparameters =[]#clear any old ones
                        best_hyperparameters.append(alpha_element)
                        best_hyperparameters.append(kernel_element)
                        best_hyperparameters.append(degree_element)
            else:
                myModel = KernelRidge(alpha=alpha_element, kernel=kernel_element)#should we set the 'normalize" parameter to true?, default is false

                #we get back a list of the scores
                myScores = evaluate_model(myModel, 'Kernel Ridge Regression') # Least squares loss with L2 reg.

                #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                if run_once_flag == False:
                    best_model = myModel
                    best_score = myScores
                    best_hyperparameters.append(alpha_element)
                    best_hyperparameters.append(kernel_element)
                    run_once_flag = True

                #check if we have a better model based on validaiton MSE score, and update if we do
                if myScores[1] < best_score[1]: #we want the validation MSE
                    best_model = myModel
                    best_score = myScores
                    best_hyperparameters =[]#clear any old ones
                    best_hyperparameters.append(alpha_element)
                    best_hyperparameters.append(kernel_element)

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

def tune_SVR_hyperparameters():
    #hyperparameters
    tol = [0.01, .001, 0.0001, .00001]
    C = [0.001, 0.01, 0.1, 0.5, 1, 1.5, 2.0, 2.5, 3]#C 5 to 3, in steps of .5
    kernel = ['linear', 'poly', 'sigmoid', 'rbf']
    degree = [2, 3, 4]#this will only be used if kernel is poly

    #trackers for best model and its scores
    best_model = None
    best_model_scores = None
    best_hyperparameters = []

    run_once_flag = False
    for tol_element in tol:
        for c_element in C:
            for kernel_element in kernel:
                if kernel_element == 'poly':#use the degree
                    for degree_element in degree:

                        myModel = SVR(tol=tol_element, C=c_element ,kernel=kernel_element, degree=degree_element)#should we set the 'normalize" parameter to true?, default is false

                        #we get back a list of the scores
                        myScores = evaluate_model(myModel, 'Support Vector Regression') # Least squares loss with L2 reg.

                        #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                        if run_once_flag == False:
                            best_model = myModel
                            best_score = myScores
                            best_hyperparameters.append(tol_element)
                            best_hyperparameters.append(c_element)
                            best_hyperparameters.append(kernel_element)
                            best_hyperparameters.append(degree_element)
                            run_once_flag = True

                        #check if we have a better model based on validaiton MSE score, and update if we do
                        if myScores[1] < best_score[1]: #we want the validation MSE
                            best_model = myModel
                            best_score = myScores
                            best_hyperparameters =[]#clear any old ones
                            best_hyperparameters.append(tol_element)
                            best_hyperparameters.append(c_element)
                            best_hyperparameters.append(kernel_element)
                            best_hyperparameters.append(degree_element)
                else:
                    myModel = SVR(tol=tol_element, C=c_element ,kernel=kernel_element)#should we set the 'normalize" parameter to true?, default is false

                    #we get back a list of the scores
                    myScores = evaluate_model(myModel, 'Support Vector Regression') # Least squares loss with L2 reg.

                    #if index is 0, this is teh first iteration of this loop, so just set best model and score b.c. otherwise we'd have nothing to compare against
                    if run_once_flag == False:
                        best_model = myModel
                        best_score = myScores
                        best_hyperparameters.append(tol_element)
                        best_hyperparameters.append(c_element)
                        best_hyperparameters.append(kernel_element)
                        run_once_flag = True

                    #check if we have a better model based on validaiton MSE score, and update if we do
                    if myScores[1] < best_score[1]: #we want the validation MSE
                        best_model = myModel
                        best_score = myScores
                        best_hyperparameters =[]#clear any old ones
                        best_hyperparameters.append(tol_element)
                        best_hyperparameters.append(c_element)
                        best_hyperparameters.append(kernel_element)

    #now that we've gon through all combinations of hyperparameters store everything in a bestModelObject and return the object
    return BestModelObject(best_model, best_score, best_hyperparameters)#return best model with best hyperparameters

from sklearn import datasets
from sklearn.externals import joblib

#section to call the evaluation for a given model and it's hyperparamters
print ("Training linear regression model")
best_LR_model = tune_linear_regression_hyperparameters()#get the model with the best hyperparameters
best_LR_model.model.fit(X,y)#retrain the model on the whole dataset
y_hat_train = best_LR_model.model.predict(X)
LR_MSE_train = mean_squared_error(y, y_hat_train)
LR_r2_train = r2_score(y, y_hat_train)
y_hat = best_LR_model.model.predict(X_test)#make predictions
LR_MSE_test = mean_squared_error(y_test, y_hat)#compute MSE of predictions
LR_r2_test = r2_score(y_test, y_hat)
#print("LR hyperparameters and test error")
#print(best_LR_model.hyperparameters)
#print(LR_MSE_test)
#print(LR_r2_test)

# save model using joblib
joblib.dump(best_LR_model.model, "LR_model.joblib")

print ("Training lasso regression model")
best_lasso_model = tune_lasso_regression_hyperparameters()#get the model with the best hyperparameters
best_lasso_model.model.fit(X,y)#retrain the model on the whole dataset
y_hat_train = best_lasso_model.model.predict(X)
lasso_MSE_train = mean_squared_error(y, y_hat_train)
lasso_r2_train = r2_score(y, y_hat_train)
y_hat = best_lasso_model.model.predict(X_test)#make predictions
lasso_MSE = mean_squared_error(y_test, y_hat)#compute MSE of predictions
lasso_r2 = r2_score(y_test, y_hat)
#print("Lasso hyperparameters and test error")
#print(best_lasso_model.hyperparameters)
#print(lasso_MSE)
#print(lasso_r2)

# save model using joblib
joblib.dump(best_lasso_model.model, "lasso_model.joblib")

print ("Training ridge regression model")
best_ridge_model = tune_ridge_regression_hyperparameters()#get the model with the best hyperparameters
best_ridge_model.model.fit(np.ascontiguousarray(X),np.ascontiguousarray(y))#retrain the model on the whole dataset
y_hat_train = best_ridge_model.model.predict(X)
ridge_MSE_train = mean_squared_error(y, y_hat_train)
ridge_r2_train = r2_score(y, y_hat_train)
y_hat = best_ridge_model.model.predict(X_test)#make predictions
ridge_MSE = mean_squared_error(y_test, y_hat)#compute MSE of predictions
ridge_r2 = r2_score(y_test, y_hat)
#print("Ridge hyperparameters and test error")
#print(best_ridge_model.hyperparameters)
#print(ridge_MSE)
#print(ridge_r2)

# save model using joblib
joblib.dump(best_ridge_model.model, "ridge_model.joblib")

print ("Training Huber regression model")
best_huber_model = tune_huber_regression_hyperparameters()#get the model with the best hyperparameters
best_huber_model.model.fit(X,y)#retrain the model on the whole dataset
y_hat_train = best_huber_model.model.predict(X)
huber_MSE_train = mean_squared_error(y, y_hat_train)
huber_r2_train = r2_score(y, y_hat_train)
y_hat = best_huber_model.model.predict(X_test)#make predictions
huber_MSE = mean_squared_error(y_test, y_hat)#compute MSE of predictions
huber_r2 = r2_score(y_test, y_hat)
#print("Huber hyperparameters and test error")
#print(best_huber_model.hyperparameters)
#print(huber_MSE)
#print(huber_r2)

# save model using joblib
joblib.dump(best_huber_model.model, "huber_model.joblib")

print ("Training Bayesian ridge regression model")
best_bayesian_ridge_model = tune_bayesian_ridge_regression_hyperparameters()#get the model with the best hyperparameters
best_bayesian_ridge_model.model.fit(X,y)#retrain the model on the whole dataset
y_hat_train = best_bayesian_ridge_model.model.predict(X)
bayesRidge_MSE_train = mean_squared_error(y, y_hat_train)
bayesRidge_r2_train = r2_score(y, y_hat_train)
y_hat = best_bayesian_ridge_model.model.predict(X_test)#make predictions
bayesian_ridge_MSE = mean_squared_error(y_test, y_hat)#compute MSE of predictions
bayesian_ridge_r2 = r2_score(y_test, y_hat)
#print("Bayesian Ridge hyperparameters and test error")
#print(best_bayesian_ridge_model.hyperparameters)
#print(bayesian_ridge_MSE)
#print(bayesian_ridge_r2)

# save model using joblib
joblib.dump(best_bayesian_ridge_model.model, "bayesian_ridge_model.joblib")

print ("Training neural net model")
best_neural_network_model = tune_neural_network_regression_hyperparameters()#get the model with the best hyperparameters
best_neural_network_model.model.fit(X,y)#retrain the model on the whole dataset
y_hat_train = best_neural_network_model.model.predict(X)
NN_MSE_train = mean_squared_error(y, y_hat_train)
NN_r2_train = r2_score(y, y_hat_train)
y_hat = best_neural_network_model.model.predict(X_test)#make predictions
neural_network_MSE = mean_squared_error(y_test, y_hat)#compute MSE of predictions
neural_network_r2 = r2_score(y_test, y_hat)
#print("Neural Network hyperparameters and test error")
#print(best_neural_network_model.hyperparameters)
#print(neural_network_MSE)
#print(neural_network_r2)

# save model using joblib
joblib.dump(best_neural_network_model.model, "NN_model.joblib")

print ("Training random forest regression model")
best_random_forest_model = tune_random_forest_regression_hyperparameters()#get the model with the best hyperparameters
best_random_forest_model.model.fit(X,y)#retrain the model on the whole dataset
y_hat_train = best_random_forest_model.model.predict(X)
randomForest_MSE_train = mean_squared_error(y, y_hat_train)
randomForest_r2_train = r2_score(y, y_hat_train)
y_hat = best_random_forest_model.model.predict(X_test)#make predictions
random_forest_MSE = mean_squared_error(y_test, y_hat)#compute MSE of predictions
random_forest_r2 = r2_score(y_test, y_hat)
#print("Random Forest Regression hyperparameters and test error")
#print(best_random_forest_model.hyperparameters)
#print(random_forest_MSE)
#print(random_forest_r2)

# save model using joblib
joblib.dump(best_random_forest_model.model, "randomforest_model.joblib")

print ("Training kernel ridge regression model")
best_kernel_ridge_model = tune_kernel_ridge_regression_hyperparameters()#get the model with the best hyperparameters
best_kernel_ridge_model.model.fit(X,y)#retrain the model on the whole dataset
y_hat_train = best_kernel_ridge_model.model.predict(X)
kernel_MSE_train = mean_squared_error(y, y_hat_train)
kernel_r2_train = r2_score(y, y_hat_train)
y_hat = best_kernel_ridge_model.model.predict(X_test)#make predictions
kernel_ridge_MSE = mean_squared_error(y_test, y_hat)#compute MSE of predictions
kernel_ridge_r2 = r2_score(y_test, y_hat)
#print("Kernel Ridge Regression hyperparameters and test error")
#print(best_kernel_ridge_model.hyperparameters)
#print(kernel_ridge_MSE)
#print(kernel_ridge_r2)

# save model using joblib
joblib.dump(best_kernel_ridge_model.model, "kernel_ridge_model.joblib")

print ("Training SVR model")
best_SVR_model = tune_SVR_hyperparameters()#get the model with the best hyperparameters
best_SVR_model.model.fit(X,y)#retrain the model on the whole dataset
y_hat_train = best_SVR_model.model.predict(X)
svr_MSE_train = mean_squared_error(y, y_hat_train)
svr_r2_train = r2_score(y, y_hat_train)
y_hat = best_SVR_model.model.predict(X_test)#make predictions
svr_MSE = mean_squared_error(y_test, y_hat)#compute MSE of predictions
svr_r2 = r2_score(y_test, y_hat)
#print("SVR Regression hyperparameters and test error")
#print(best_SVR_model.hyperparameters)
#print(svr_MSE)
#print(svr_r2)

# save model using joblib
joblib.dump(best_SVR_model.model, "SVR_model.joblib")

#prep csv file to output to
filename = "results_" + str(datetime.now().strftime('%Y-%m-%d_%H:%M:%S')) + ".csv"
with open(filename, 'w', newline='') as csvfile:
    spamwriter = csv.writer(csvfile, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL)
    # validation MSE from best cros validation model,
    # training MSE from retraining model on whole training set
    # training r^2 from retraning model on whole training set
    # test MSE
    # test r^2
    spamwriter.writerow(['Model'] + ['Validation MSE'] + ['Training MSE'] + ['Training R2'] + ['Test MSE'] + ['Test R2'] + ['Hyper-parameters'])
    spamwriter.writerow(['LR Model'] + [best_LR_model.scores[1]] + [LR_MSE_train] + [LR_r2_train] + [LR_MSE_test] + [LR_r2_test] + best_LR_model.hyperparameters)
    spamwriter.writerow(['Lasso Model'] + [best_lasso_model.scores[1]] + [lasso_MSE_train] + [lasso_r2_train] + [lasso_MSE] + [lasso_r2] + best_lasso_model.hyperparameters)
    spamwriter.writerow(['Ridge Model'] + [best_ridge_model.scores[1]] + [ridge_MSE_train] + [ridge_r2_train] + [ridge_MSE] + [ridge_r2] + best_ridge_model.hyperparameters)
    spamwriter.writerow(['Huber Model'] + [best_huber_model.scores[1]] + [huber_MSE_train] + [huber_r2_train] + [huber_MSE] + [huber_r2] + best_huber_model.hyperparameters)
    spamwriter.writerow(['Bayesian Ridge Model'] + [best_bayesian_ridge_model.scores[1]] + [bayesRidge_MSE_train] + [bayesRidge_r2_train] + [bayesian_ridge_MSE] + [bayesian_ridge_r2] + best_bayesian_ridge_model.hyperparameters)
    spamwriter.writerow(['Neural Net Model'] + [best_neural_network_model.scores[1]] + [NN_MSE_train] + [NN_r2_train] + [neural_network_MSE] + [neural_network_r2] + best_neural_network_model.hyperparameters)
    spamwriter.writerow(['Random Forest Model'] + [best_random_forest_model.scores[1]] + [randomForest_MSE_train] + [randomForest_r2_train] + [random_forest_MSE] + [random_forest_r2] + best_random_forest_model.hyperparameters)
    spamwriter.writerow(['Kernel Model'] + [best_kernel_ridge_model.scores[1]] + [kernel_MSE_train] + [kernel_r2_train] + [kernel_ridge_MSE] + [kernel_ridge_r2] + best_kernel_ridge_model.hyperparameters)
    spamwriter.writerow(['SVR Model'] + [best_SVR_model.scores[1]] + [svr_MSE_train] + [svr_r2_train] + [svr_MSE] + [svr_r2] + best_SVR_model.hyperparameters)

max_stars_repo_name: lpalazzi/LLFI-CrashRatePredictor

max_issues_repo_name: lpalazzi/LLFI-CrashRatePredictor

max_forks_repo_name: lpalazzi/LLFI-CrashRatePredictor


text:
'''Reference: https://github.com/AI4Finance-LLC/FinRL'''

import pandas as pd
import yfinance as yf
import numpy as np
from stockstats import StockDataFrame as Sdf
from finrl.config import config

class YahooFinanceEngineer():
    """Provides methods for retrieving daily stock data from
    Yahoo Finance API
    Attributes
    ----------
        start_date : str
            start date of the data (modified from config.py)
        end_date : str
            end date of the data (modified from config.py)
        ticker_list : list
            a list of stock tickers (modified from config.py)
    Methods
    -------
    fetch_data()
        Fetches data from yahoo API
    """

    def __init__(self):
        pass

    def data_fetch(self, start_date: str, end_date: str, ticker_list: list) -> pd.DataFrame:
        """Fetches data from Yahoo API
        Parameters
        ----------
        Returns
        -------
        `pd.DataFrame`
            7 columns: A date, open, high, low, close, volume and tick symbol
            for the specified stock ticker
        """
        # Download and save the data in a pandas DataFrame:
        data_df = pd.DataFrame()
        for tic in ticker_list:
            temp_df = yf.download(tic, start=start_date, end=end_date)
            temp_df["tic"] = tic
            data_df = data_df.append(temp_df)
        # reset the index, we want to use numbers as index instead of dates
        data_df = data_df.reset_index()
        try:
            # convert the column names to standardized names
            data_df.columns = [
                "date",
                "open",
                "high",
                "low",
                "close",
                "adjcp",
                "volume",
                "tic",
            ]
        except NotImplementedError:
            print("the features are not supported currently")
        # create day of the week column (monday = 0)
        data_df["day"] = data_df["date"].dt.dayofweek
        # convert date to standard string format, easy to filter
        data_df["date"] = data_df.date.apply(lambda x: x.strftime("%Y-%m-%d"))
        # drop missing data
        data_df = data_df.dropna()
        data_df = data_df.reset_index(drop=True)
        print("Shape of DataFrame: ", data_df.shape)
        # print("Display DataFrame: ", data_df.head())

        data_df = data_df.sort_values(by=['date','tic']).reset_index(drop=True)

        return data_df

    def select_equal_rows_stock(self, df):
        df_check = df.tic.value_counts()
        df_check = pd.DataFrame(df_check).reset_index()
        df_check.columns = ["tic", "counts"]
        mean_df = df_check.counts.mean()
        equal_list = list(df.tic.value_counts() >= mean_df)
        names = df.tic.value_counts().index
        select_stocks_list = list(names[equal_list])
        df = df[df.tic.isin(select_stocks_list)]
        return df

    def add_technical_indicator(self, data, tech_indicator_list):
        """
        calculate technical indicators
        use stockstats package to add technical inidactors
        :param data: (df) pandas dataframe
        :return: (df) pandas dataframe
        """
        df = data.copy()
        df = df.sort_values(by=['tic','date'])
        stock = Sdf.retype(df.copy())
        unique_ticker = stock.tic.unique()

        for indicator in tech_indicator_list:
            indicator_df = pd.DataFrame()
            for i in range(len(unique_ticker)):
                try:
                    temp_indicator = stock[stock.tic == unique_ticker[i]][indicator]
                    temp_indicator = pd.DataFrame(temp_indicator)
                    temp_indicator['tic'] = unique_ticker[i]
                    temp_indicator['date'] = df[df.tic == unique_ticker[i]]['date'].to_list()
                    indicator_df = indicator_df.append(
                        temp_indicator, ignore_index=True
                    )
                except Exception as e:
                    print(e)
            df = df.merge(indicator_df[['tic','date',indicator]],on=['tic','date'],how='left')
        df = df.sort_values(by=['date','tic'])
        return df

    def add_user_defined_feature(self, data):
        """
         add user defined features
        :param data: (df) pandas dataframe
        :return: (df) pandas dataframe
        """
        df = data.copy()
        df["daily_return"] = df.close.pct_change(1)
        # df['return_lag_1']=df.close.pct_change(2)
        # df['return_lag_2']=df.close.pct_change(3)
        # df['return_lag_3']=df.close.pct_change(4)
        # df['return_lag_4']=df.close.pct_change(5)
        return df

    def add_turbulence(self, data):
        """
        add turbulence index from a precalcualted dataframe
        :param data: (df) pandas dataframe
        :return: (df) pandas dataframe
        """
        df = data.copy()
        turbulence_index = self.calculate_turbulence(df)
        df = df.merge(turbulence_index, on="date")
        df = df.sort_values(["date", "tic"]).reset_index(drop=True)
        return df

    def calculate_turbulence(self, data, time_period = 252):
        """calculate turbulence index based on dow 30"""
        # can add other market assets
        df = data.copy()
        df_price_pivot = df.pivot(index="date", columns="tic", values="close")
        # use returns to calculate turbulence
        df_price_pivot = df_price_pivot.pct_change()

        unique_date = df.date.unique()
        # start after a year
        start = time_period
        turbulence_index = [0] * start
        # turbulence_index = [0]
        count = 0
        for i in range(start, len(unique_date)):
            current_price = df_price_pivot[df_price_pivot.index == unique_date[i]]
            # use one year rolling window to calcualte covariance
            hist_price = df_price_pivot[
                (df_price_pivot.index < unique_date[i])
                & (df_price_pivot.index >= unique_date[i - time_period])
            ]
            # Drop tickers which has number missing values more than the "oldest" ticker
            filtered_hist_price = hist_price.iloc[hist_price.isna().sum().min():].dropna(axis=1)

            cov_temp = filtered_hist_price.cov()
            current_temp = current_price[[x for x in filtered_hist_price]] - np.mean(filtered_hist_price, axis=0)
            #cov_temp = hist_price.cov()
            #current_temp=(current_price - np.mean(hist_price,axis=0))

            temp = current_temp.values.dot(np.linalg.pinv(cov_temp)).dot(
                current_temp.values.T
            )
            if temp > 0:
                count += 1
                if count > 2:
                    turbulence_temp = temp[0][0]
                else:
                    # avoid large outlier because of the calculation just begins
                    turbulence_temp = 0
            else:
                turbulence_temp = 0
            turbulence_index.append(turbulence_temp)

        turbulence_index = pd.DataFrame(
            {"date": df_price_pivot.index, "turbulence": turbulence_index}
        )
        return turbulence_index

    def df_to_ary(self,df,tech_indicator_list):
        """transform final df to numpy arrays"""
        unique_ticker = df.tic.unique()
        print(unique_ticker)
        if_first_time = True
        for tic in unique_ticker:
            if if_first_time:
                price_ary = df[df.tic==tic][['adjcp']].values
                #price_ary = df[df.tic==tic]['close'].values
                tech_ary = df[df.tic==tic][tech_indicator_list].values
                turbulence_ary = df[df.tic==tic]['turbulence'].values
                if_first_time = False
            else:
                price_ary = np.hstack([price_ary, df[df.tic==tic][['close']].values])
                tech_ary = np.hstack([tech_ary, df[df.tic==tic][tech_indicator_list].values])
        assert price_ary.shape[0] == tech_ary.shape[0]
        assert tech_ary.shape[0] == turbulence_ary.shape[0]
        print('Successfully transformed into array')
        return price_ary,tech_ary,turbulence_ary

ext: py

max_stars_repo_name: ZerounNet/FinRL

max_issues_repo_name: ZerounNet/FinRL

max_forks_repo_name: ZerounNet/FinRL

reason: import numpy


text:
import numba

import numpy as np
from ..spectra.source_function import Bv_T, PopuSource_AB
from ..spectra.abscoeff import basic

@numba.jit
def FOSC(tau, Sb, Sm, Ib):
    """ First Order Short Characteristics

    See "Theory of stellar atmospheres: An introduction to astrophysical
    non-equilibrium quantitative spectroscopic analysis" by Ivan Hubeny
    and Dimitri Mihalas, ISBN 978-0-691-16328-4

    Parameters:
        tau: optical depth between two layers
        Sb: Source function at adjacent grid points
        Sm: Source function at target point
        Ib: Intensity adjacent grid points

    Returns:
        Inwards-directed intensity, lambda used in equation
    """
    yd = tau - 1. + np.exp(-tau)  # (12.120)

    dev_cond = tau * 1.
    dev_cond[tau == 0] = np.nan

    lambda_m = yd/dev_cond  # (12.117)
    lambda_b = - (yd / dev_cond) + 1. - np.exp(-tau)  # (12.118, 116)
    Im = Ib * np.exp(-tau) + lambda_m * Sm + lambda_b * Sb  # (12.114)
    Im[tau == 0] = Ib[tau == 0]
    return Im, lambda_m

def SOSC(tau1, tau3, S1, S2, S3, I1):
    """
    Calculate intensity by second order short characteristic method
    written in Kunasz and Auer, 1988, J. Quanr Spectrosc. Radiar Transfer

    Parameters:
        tau1: optical depth at grid 1
        tau3: optical depth at grid 3
        S1: Source function at grid 1
        S2: Source function at grid 2
        S3: Source function at grid 3
        I1: Intensity at grid 1

    Returns:
        intensity at grid 2, lambda used in equation
    """
    lambda_1 = np.empty_like(tau1)
    lambda_2 = np.empty_like(tau1)
    lambda_3 = np.empty_like(tau1)

    tau1[tau1 != tau1] = 0
    tau3[tau3 != tau3] = 0

    linear = np.logical_or(tau1 < 1.e-4,
                           tau3 < 1.e-10)  # Computational error region
    exp_tau = np.exp(-tau1)
    w0 = 1 - exp_tau

    lambda_1[linear] = w0[linear]
    lambda_2[linear] = 0
    lambda_3[linear] = 0

    linear = linear == False  # Second order calculation for not too small tau

    w1 = tau1[linear] - w0[linear]
    w2 = tau1[linear]**2 - 2 * w1  # w2 < 1.e-16 is 0
    lambda_1[linear] = w0[linear] + (w2 - (tau3[linear] + 2 * tau1[linear])
                                     * w1) / (tau1[linear] *
                                              (tau1[linear] + tau3[linear]))
    lambda_2[linear] = (w1 * (tau1[linear] + tau3[linear])
                        - w2) / (tau1[linear] * tau3[linear])
    lambda_3[linear] = (w2 - w1 * tau1[linear]) / \
                       (tau3[linear] * (tau1[linear] + tau3[linear]))

    source_function = lambda_3 * S3 + lambda_2 * S2 + lambda_1 * S1

    I2 = I1 * exp_tau + source_function

    # return I2, lambda_2+I1 * exp_tau
    return I2, lambda_2 + lambda_1 * exp_tau

def SOSCdamy(tau, tau3, Sb, Sm, S3, Ib):
    """
    Calculate intensity by second order short characteristic method
    written in Kunasz and Auer, 1988, J. Quanr Spectrosc. Radiar Transfer

    Parameters:
        tau1: optical depth at grid 1
        tau3: optical depth at grid 3
        S1: Source function at grid 1
        S2: Source function at grid 2
        S3: Source function at grid 3
        I1: Intensity at grid 1

    Returns:
        intensity at grid 2, lambda used in equation
    """
    yd = tau - 1. + np.exp(-tau)  # (12.120)

    dev_cond = tau * 1.

    lambda_m = yd/dev_cond  # (12.117)
    lambda_b = - (yd / dev_cond) + 1. - np.exp(-tau)  # (12.118, 116)
    Im = Ib * np.exp(-tau) + lambda_m * Sm + lambda_b * Sb  # (12.114)
    return Im, lambda_m

max_stars_repo_name: tmieslinger/typhon

max_issues_repo_name: tmieslinger/typhon

max_forks_repo_name: tmieslinger/typhon

reason: import numpy,import numba


text:
'''
Created on 11 Oct 2016

@author: Federico Ressi
'''

# pylint: disable=missing-docstring

import argparse
import collections
import csv
import json
import logging
import os
import re
import shutil
import sys
import zipfile
import zlib

import jinja2
import msgpack
import numpy
import pandas
import yaml

import pubtransit

LOG = logging.getLogger(__name__)

OUT_STREAM = sys.stdout

TEMPLATE_MANAGER = jinja2.Environment(
    loader=jinja2.PackageLoader(pubtransit.__name__, ''))

TARGET_METHODS = {}

DEFAULT_STOPS_PER_TILE = 128

def main():
    logging.basicConfig(
        level=logging.WARNING, format="%(asctime)-15s | %(message)s")

    parser = argparse.ArgumentParser(
        description='Departures transit feed compiler.')

    parser.add_argument(
        '--target', type=str, nargs=1, choices=TARGET_METHODS,
        default='all',
        help='One between: {}'.format(', '.join(TARGET_METHODS)))

    parser.add_argument(
        '--build-dir', type=str, default='build',
        help='Folder where to put produced data.')

    parser.add_argument(
        '--quiet', dest='logging_level', action='store_const',
        const=logging.ERROR, default=None, help='Show only error messages.')

    parser.add_argument(
        '--logging-level', dest='logging_level', default=logging.WARNING,
        type=int,
        help='Set logging level (from {min} to {max}.'.format(
            min=logging.DEBUG, max=logging.FATAL))

    parser.add_argument(
        '--max-stops', dest='max_stops', default=DEFAULT_STOPS_PER_TILE,
        type=int, help='Set maximum number of stops-per-tile.')

    parser.add_argument(
        '--verbose', dest='logging_level', action='store_const',
        const=logging.INFO, default=None, help='Show verbose messages.')

    parser.add_argument(
        '--debug', dest='logging_level', action='store_const',
        const=logging.DEBUG, default=None, help='Show debug messages.')

    parser.add_argument(
        'files', type=str, default=['site.yml'], nargs='*',
        help='Feed file to extract feed rules from.')

    parser.add_argument(
        '--dest', type=str, help='Destination feed file.')

    args = parser.parse_args()

    if args.logging_level:
        # Raise logging level
        logging.getLogger().setLevel(args.logging_level)

    method = TARGET_METHODS[args.target[0]]

    try:
        for inpute_file in args.files or [None]:
            method(args, inpute_file)

    except Exception as error:  # pylint: disable=broad-except
        if args.logging_level is logging.DEBUG:
            LOG.fatal("Unhandled exception.", exc_info=1)
        else:
            LOG.fatal(str(error) or str(type(error)))
        exit(1)

    except BaseException:
        logging.warning('interrupted', exc_info=1)
        raise

    else:
        logging.debug('SUCCESS')

def target_method(name):

    def decorator(func):
        TARGET_METHODS[name] = func
        return func

    return decorator

MethodParameters = collections.namedtuple(
    'MethodParameters', ['site', 'feed', 'target_path'])

@target_method("version")
def print_version(args):
    # pylint: disable=unused-argument
    OUT_STREAM.write(pubtransit.__version__ + '\n')

@target_method("makefile")
def make_makefiles(args, site_file=None):
    feeds_conf = read_yaml_file(site_file or 'site.yml')
    for site in feeds_conf['feed']:
        for feed in site["feeds"]:
            target_path = os.path.join(
                args.build_dir, site["name"], feed["name"])

            if not os.path.isdir(target_path):
                os.makedirs(target_path)

            url = feed.get("url") or (site["url"] + '/' + feed["path"])

            # pylint: disable=unused-argument,no-member
            OUT_STREAM.write(target_path + ".mk ")
            target_template = TEMPLATE_MANAGER.get_template("feed_item.mk")
            target_make = target_template.render(
                install_dir=os.path.join('$(INSTALL_DIR)', 'feed'),
                build_dir=args.build_dir,
                target=os.path.join(site["name"], feed["name"]),
                url=url,
                make_flags="--logging-level " + str(args.logging_level),
                make_me='python -m pubtransit ' + ' '.join(
                    repr(arg) for arg in sys.argv[1:]),
                script_name="pubtransit")
            with open(target_path + ".mk", 'wt') as target_stream:
                target_stream.write(target_make)

@target_method("datastore")
def generate_datastores(args, feed_file):
    dest_dir = os.path.splitext(args.dest or args.source)[0]
    if os.path.isdir(dest_dir):
        shutil.rmtree(dest_dir)
    os.makedirs(dest_dir)

    with zipfile.ZipFile(feed_file) as zip_file:
        routes = read_routes(zip_file)
        generate_routes(dest_dir=dest_dir, routes=routes)

        trips = read_trips(zip_file)
        generate_trips(dest_dir=dest_dir, trips=trips, route_id=routes.id)

        stops = read_stops(zip_file)
        tiles = generate_tiled_stops(
            dest_dir=dest_dir, stops=stops, max_rows=args.max_stops)

        stop_times = read_stop_times(zip_file)
        generate_tiled_stop_times(
            dest_dir=dest_dir, stop_times=stop_times, trip_id=trips.id,
            tiles=tiles)

    feed_info = dict(
        west=stops.west, east=stops.east, south=stops.south, north=stops.north)
    store_object(feed_info, dest_dir, 'feed')

@target_method("index")
def make_index(args, feed_file=None):
    feeds_conf = read_yaml_file(feed_file or 'feeds.yaml')
    paths = []
    west = []
    east = []
    south = []
    north = []
    for site in feeds_conf['feed']:
        for feed in site["feeds"]:
            target_path = os.path.join(
                args.build_dir, site["name"], feed["name"])

            if not os.path.isdir(target_path):
                LOG.error('Target feed dir not found: %r', target_path)
                os.makedirs(target_path)

            with open(os.path.join(target_path, 'feed.gz')) as in_stream:
                zipped = in_stream.read()
                packed = zlib.decompress(zipped)
                feed_info = msgpack.unpackb(packed)

            paths.append(os.path.join(site["name"], feed["name"]))
            west.append(feed_info['west'])
            east.append(feed_info['east'])
            south.append(feed_info['south'])
            north.append(feed_info['north'])

    store_column(paths, args.build_dir, 'index', 'path')
    store_column(west, args.build_dir, 'index', 'west', float)
    store_column(east, args.build_dir, 'index', 'east', float)
    store_column(south, args.build_dir, 'index', 'south', float)
    store_column(north, args.build_dir, 'index', 'north', float)

def generate_routes(dest_dir, routes):
    store_column(routes.name, dest_dir, 'routes', 'name')

def generate_trips(dest_dir, trips, route_id):
    trip_route_id = numpy.searchsorted(route_id, trips.route_id)
    store_column(trip_route_id, dest_dir, 'trips', 'route_id', int)
    store_column(trips.name, dest_dir, 'trips', 'name')

def generate_tiled_stops(dest_dir, stops, max_rows=None):
    if not max_rows:
        max_rows = len(stops.id)
    max_rows = max(max_rows, 4)

    tiles_tree, tiles = create_tree(
        stops, index_columns=['lon', 'lat'], max_rows=max_rows)

    tiles_num = len(tiles)
    tiles_shape = tiles_num,
    tiles_id_format = '0' + str(len(str(tiles_num)))
    tiles_west = numpy.zeros(tiles_shape, dtype=float)
    tiles_east = numpy.zeros(tiles_shape, dtype=float)
    tiles_south = numpy.zeros(tiles_shape, dtype=float)
    tiles_north = numpy.zeros(tiles_shape, dtype=float)
    for i, tile in enumerate(tiles):
        tile_dir = os.path.join(dest_dir, format(i, tiles_id_format))
        store_column(tile.name, tile_dir, 'stops', 'name')
        store_column(tile.lon, tile_dir, 'stops', 'lon', float)
        store_column(tile.lat, tile_dir, 'stops', 'lat', float)
        tiles_west[i] = tile.west
        tiles_east[i] = tile.east
        tiles_south[i] = tile.south
        tiles_north[i] = tile.north

    store_object(tiles_tree, os.path.join(dest_dir, 'tiles'), 'tree')
    store_column(tiles_west, dest_dir, 'tiles', 'west')
    store_column(tiles_east, dest_dir, 'tiles', 'east')
    store_column(tiles_south, dest_dir, 'tiles', 'south')
    store_column(tiles_north, dest_dir, 'tiles', 'north')
    return tiles

def generate_tiled_stop_times(dest_dir, stop_times, trip_id, tiles):
    # pylint: disable=too-many-locals
    tiles_num = len(tiles)
    tiles_id_format = '0' + str(len(str(tiles_num)))

    stop_times = stop_times.sort_by('stop_id')
    trip_id_sorter = numpy.argsort(trip_id)

    for tile_id, stops in enumerate(tiles):
        stop_id = stops.id
        stop_id_sorter = numpy.argsort(stop_id)

        tile_start = numpy.searchsorted(
            stop_id, stop_times.stop_id, side='left', sorter=stop_id_sorter)
        tile_stop = numpy.searchsorted(
            stop_id, stop_times.stop_id, side='right', sorter=stop_id_sorter)
        tiled_stop_times = stop_times.select(tile_start != tile_stop)

        tiled_stop_times_stop_id = stop_id_sorter[numpy.searchsorted(
            stop_id, tiled_stop_times.stop_id, sorter=stop_id_sorter)]

        tiled_stop_times_trip_id = trip_id_sorter[numpy.searchsorted(
            trip_id, tiled_stop_times.trip_id,
            sorter=trip_id_sorter)]

        tile_dir = os.path.join(dest_dir, format(tile_id, tiles_id_format))
        store_column(
            tiled_stop_times_stop_id, tile_dir, 'stop_times', 'stop_id')
        store_column(
            tiled_stop_times_trip_id, tile_dir, 'stop_times', 'trip_id')
        store_column(
            tiled_stop_times.departure_minutes, tile_dir, 'stop_times',
            'departure_minutes')

        stop_empty = numpy.ones_like(stop_id, dtype=int)
        stop_empty[tiled_stop_times_stop_id] = 0
        store_column(stop_empty, tile_dir, 'stops', 'empty')

def timestamp_to_minutes(timestamp):
    timestamp = numpy.asarray(timestamp, dtype='S8')
    timestamp = numpy.core.defchararray.rjust(timestamp, 8, '0')

    hours = array_from_data(  # pylint: disable=no-member
        data=timestamp.__array_interface__['data'],
        shape=timestamp.shape, typestr='S2', strides=(8,)
        ).astype(int) % 24

    minutes = array_from_data(  # pylint: disable=no-member
        data=timestamp.__array_interface__['data'],
        shape=timestamp.shape, typestr='S2', strides=(8,),
        offset=3).astype(int) % 60

    timestamp = (hours * 60) + minutes
    assert timestamp.max() < 24 * 60
    return timestamp

def named_tuple(*fields):

    def decorator(cls):
        _named_tuple_type = collections.namedtuple(cls.__name__, fields)
        return type(cls.__name__, (cls, _named_tuple_type), {})

    return decorator

class BaseTable(tuple):

    def sort_by(self, index_name):
        # sort all columns by given index
        return self.sort_by_array(getattr(self, index_name))

    def sort_by_array(self, index_array, sort_index_array=False, sorter=None):
        # sort all columns by given index
        if sorter is None:
            sorter = numpy.argsort(index_array)
        if sort_index_array:
            index_array[:] = index_array[sorter]
        return self.select(sorter)

    def select(self, item):
        values = [
            column[item] for column in self if column is not None]
        return type(self)(*values)

def create_tree(table, index_columns, max_rows=128):
    # pylint: disable=too-many-locals
    table_class = type(table)
    ndim = len(index_columns)
    tree = {}
    stack = [(tree, table, 0)]
    leaves = []

    while stack:
        node, table, level = stack.pop()
        column_id = index_columns[level % ndim]
        column = getattr(table, column_id)
        if len(column) <= max_rows:
            node['leaf'] = len(leaves)
            leaves.append(table)

        else:
            table = table.sort_by(column_id)
            column = getattr(table, column_id)  # this is another table!!
            mid = int(len(column) / 2)
            node['col'] = column_id
            node['min'] = column[0]
            node['mid'] = column[mid]
            node['max'] = column[-1]
            node['left'] = {}
            node['right'] = {}
            assert node['min'] <= node['max']
            assert node['min'] <= node['mid']
            assert node['mid'] <= node['max']

            left = table_class(*[col[:mid] for col in table])
            right = table_class(*[col[mid:] for col in table])
            stack.append((node['right'], right, level + 1))
            stack.append((node['left'], left, level + 1))

    LOG.debug('Tree generated:\n%s',
              json.dumps(tree, indent=4, sort_keys=True))
    return tree, leaves

@named_tuple('id', 'name')
class RouteTable(BaseTable):

    @classmethod
    def from_zip_file(cls, zip_file):
        columns = read_table(
            zip_file=zip_file, table_name='routes',
            columns=['route_id', 'route_short_name'])
        return cls(*columns).sort_by('id')

read_routes = RouteTable.from_zip_file  # pylint: disable=invalid-name

@named_tuple('id', 'route_id', 'name')
class TripTable(BaseTable):

    @classmethod
    def from_zip_file(cls, zip_file):
        columns = read_table(
            zip_file=zip_file, table_name='trips',
            columns=['trip_id', 'route_id', 'trip_headsign'])
        return cls(*columns).sort_by('id')

read_trips = TripTable.from_zip_file  # pylint: disable=invalid-name

@named_tuple('lon', 'lat', 'id', 'name', 'indexes')
class StopTable(BaseTable):
    # pylint: disable=no-member

    @classmethod
    def from_zip_file(cls, zip_file):
        columns = read_table(
            zip_file=zip_file, table_name='stops',
            columns=['stop_lon', 'stop_lat', 'stop_id', 'stop_name'],
            dtypes={'stop_lon': float, 'stop_lat': float})
        columns += (numpy.arange(len(columns[0])),)
        return cls(*columns).sort_by('id')

    _west = None

    @property
    def west(self):
        west = self._west
        if west is None:
            self._west = west = numpy.amin(self.lon)
        return west

    _east = None

    @property
    def east(self):
        east = self._east
        if east is None:
            self._east = east = numpy.amax(self.lon)
        return east

    _south = None

    @property
    def south(self):
        south = self._south
        if south is None:
            self._south = south = numpy.amin(self.lat)
        return south

    _north = None

    @property
    def north(self):
        north = self._north
        if north is None:
            self._north = north = numpy.amax(self.lat)
        return north

read_stops = StopTable.from_zip_file  # pylint: disable=invalid-name

@named_tuple('stop_id', 'trip_id', 'departure_minutes')
class StopTimeTable(BaseTable):

    @classmethod
    def from_zip_file(cls, zip_file):
        stop_id, trip_id, departure_time = read_table(
            zip_file=zip_file, table_name='stop_times',
            columns=['stop_id', 'trip_id', 'departure_time'])
        return cls(
            stop_id=stop_id, trip_id=trip_id,
            departure_minutes=timestamp_to_minutes(departure_time))

    _left = None

read_stop_times = StopTimeTable.from_zip_file  # pylint: disable=invalid-name

GET_COLUMN_NAME_REGEX = re.compile(b'[^\\w]')

def read_table(zip_file, table_name, columns, dtypes=None):
    if not dtypes:
        dtypes = {}

    with zip_file.open(table_name + '.txt', 'r') as csv_stream:
        hearer = csv_stream.readline().strip()

        names = [
            GET_COLUMN_NAME_REGEX.sub(b'', name).decode('ascii')
            for name in hearer.split(b',')]

        table = pandas.read_csv(
            csv_stream, names=names, quotechar='"', quoting=csv.QUOTE_ALL,
            usecols=[col for col in columns])
        table = [
            numpy.asarray(remove_nans(table[column]), dtype=dtypes.get(column))
            for column in columns]
        return table

def remove_nans(series):
    if series.dtype == object:
        series = series.fillna('')
    return series

def store_column(array, dest_dir, table, column, dtype=None):
    if dtype:
        array = numpy.asarray(array, dtype=dtype)

    if isinstance(array, list):
        obj = array
    else:
        obj = list(array)
    store_object(obj, os.path.join(dest_dir, table), column)

def store_object(obj, dest_dir, name):
    if not os.path.isdir(dest_dir):
        os.makedirs(dest_dir)
    packed = msgpack.packb(obj)
    zipped = zlib.compress(packed, 9)
    dest_file = os.path.join(dest_dir, name + '.gz')
    with open(dest_file, 'wb') as dest_stream:
        dest_stream.write(zipped)
    LOG.info('Object stored:\n'
             '    path: %r\n'
             '    packet size: %d bytes\n'
             '    zipped size: %d bytes.\n',
             dest_file, len(packed), len(zipped))

def read_yaml_file(feed_file):
    with open(feed_file, 'rt') as feed_file_stream:
        return yaml.load(feed_file_stream.read())

def array_from_data(data, shape, typestr, strides, offset=0):
    if offset:
        data = data[0] + offset, data[1],

    # pylint: disable=no-member
    return numpy.array(ArrayInterface(
        shape=shape, typestr=typestr, data=data, strides=strides))

@named_tuple('shape', 'typestr', 'data', 'strides')
class ArrayInterface(object):
    # pylint: disable=too-few-public-methods,no-member

    @property
    def __array_interface__(self):
        return {
            'shape': self.shape,
            'typestr': self.typestr,
            'data': self.data,
            'strides': self.strides,
            'version': 3}

max_stars_repo_name: pubtransit/transit

max_issues_repo_name: pubtransit/pubtransit

max_forks_repo_name: pubtransit/pubtransit


text:
import numpy as np
import os.path
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
import json
import re
import collections
import cv2

def parse_labelmap(path):
  # item {
  #   id: 1
  #   name: 'rema'
  # }
  #
  # to
  #
  # {1: {'id': 1, 'name': u'rema'}}
  with open(path) as f:
    labels = json.loads("[%s]" % re.sub(" *([^ :]*):", '"\\1":', f.read()).replace("'", '"').replace("item {\n", "{").replace("\n}", "}").strip().replace("\n", ","))
    return {label['id']: label for label in labels}

PATH_TO_CKPT = os.path.join(os.path.dirname(__file__), 'my_faster_rcnn_resnet101_coco.pb/frozen_inference_graph.pb')
#PATH_TO_CKPT = os.path.join(os.path.dirname(__file__), 'my_pipeline.pb/frozen_inference_graph.pb')
#PATH_TO_CKPT = os.path.join(os.path.dirname(__file__), 'my_ssd_mobilenet_v1_coco.pb/frozen_inference_graph.pb')
PATH_TO_LABELS = os.path.join(os.path.dirname(__file__), 'label_map.pbtxt')

def merge_boxes(boxes,
                classes,
                scores,
                category_index,
                instance_masks=None,
                keypoints=None,
                max_boxes_to_draw=100,
                min_score_thresh=.7,
                line_thickness=4):
  """Extracted from tensorflow.research.object_detection.utils.visualization_utils.visualize_boxes_and_labels_on_image_array"""
  box_to_class_map = collections.defaultdict(str)
  box_to_class_name_map = collections.defaultdict(str)
  box_to_score_map = collections.defaultdict(float)
  box_to_instance_masks_map = {}
  box_to_keypoints_map = collections.defaultdict(list)
  if not max_boxes_to_draw:
    max_boxes_to_draw = boxes.shape[0]
  for i in range(min(max_boxes_to_draw, boxes.shape[0])):
    if scores is None or scores[i] > min_score_thresh:
      box = tuple(boxes[i].tolist())
      if instance_masks is not None:
        box_to_instance_masks_map[box] = instance_masks[i]
      if keypoints is not None:
        box_to_keypoints_map[box].extend(keypoints[i])
      if scores is None:
        box_to_class_map[box] = -1
      else:
        if classes[i] in category_index.keys():
          class_name = category_index[classes[i]]['name']
        else:
          class_name = 'N/A'
        box_to_class_map[box] = classes[i]
        box_to_class_name_map[box] = class_name
        box_to_score_map[box] = 100*scores[i]

  return {box: {"class": cls,
                "class_name": box_to_class_name_map.get(box),
                "score": box_to_score_map.get(box),
                "mask": box_to_instance_masks_map.get(box),
                "keypoints": keypoints and box_to_keypoints_map.get(box)}
          for box, cls in box_to_class_map.iteritems()}

detection_graph = tf.Graph()
with detection_graph.as_default():
  od_graph_def = tf.GraphDef()
  with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
    serialized_graph = fid.read()
    od_graph_def.ParseFromString(serialized_graph)
    tf.import_graph_def(od_graph_def, name='')

category_index = parse_labelmap(PATH_TO_LABELS)

def load_image_into_numpy_array(image):
  (im_width, im_height) = image.size
  return np.array(image.getdata()).reshape(
      (im_height, im_width, 3)).astype(np.uint8)

sess = None
def find_labels(image):
  global sess
  if sess is None:
    sess = tf.Session(graph=detection_graph)
  with detection_graph.as_default():
      image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
      detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
      detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
      detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
      num_detections = detection_graph.get_tensor_by_name('num_detections:0')

      # This is how the model was trained, so we have to use the same byte order
      image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
      image_np_expanded = np.expand_dims(image_np, axis=0)
      #image_np_expanded = image
      (boxes, scores, classes, num) = sess.run(
          [detection_boxes, detection_scores, detection_classes, num_detections],
          feed_dict={image_tensor: image_np_expanded})

      h, w = image.shape[:2]

      def fixbbox(b):
        return int(b[1]*w), int(b[0]*h), int((b[3]-b[1])*w), int((b[2]-b[0])*h)

      return {fixbbox(bbox): attrs
              for bbox, attrs in
              merge_boxes(
                np.squeeze(boxes),
                np.squeeze(classes).astype(np.int32),
                np.squeeze(scores),
                category_index).iteritems()}

def pos2bbox(p):
    p = list(p)
    p[2] -= p[0]
    p[3] -= p[1]
    return tuple(p)

def bbox2pos(b):
    b = list(b)
    b[2] += b[0]
    b[3] += b[1]
    return b

if __name__ == "__main__":
  img = cv2.imread(sys.argv[1])
#  w,h,bd = img.shape
#  img = cv2.resize(img, (800, int(h*800./w)))
#  img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
#  img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,401,0)
  # img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
#  img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)

  res = find_labels(img)
  out = img.copy()
  for bbox, info in res.iteritems():
    p = bbox2pos(bbox)
    cv2.rectangle(out, (p[0],p[1]), (p[2],p[3]), (0,0,255),2)
    cv2.putText(out, ("%(class_name)s (%(score).2f" % info).encode("utf-8"), (p[0], p[1]), cv2.FONT_HERSHEY_COMPLEX, 0.6, (255, 0, 0), 2)
  cv2.imwrite(os.path.splitext(sys.argv[1])[0] + ".out.jpg", out)

max_stars_repo_name: innovationgarage/ChalkMail

max_issues_repo_name: innovationgarage/ChalkMail

max_forks_repo_name: innovationgarage/ChalkMail


text:
# -*- coding: utf8 -*-

from pathlib import Path
import numpy as np

from sumagg.sumagg import BaseSummary, SumAgg

class MeanSummary(BaseSummary):

    def __init__(self, name):
        self.name = name
        self.reset()

    def add(self, data):
        self.all_data.append(data)

    def write_metrics(self):
        sumagg = SumAgg()

        if len(self.all_data) > 0:
            log_value = np.mean(self.all_data)
        else:
            log_value = np.nan
        sumagg.tensorboard_writer.add_scalar(f"{sumagg.phase}/{self.name}", log_value, sumagg.epoch)

        sumagg.epoch_metrics[self.name] = log_value

    def reset(self):
        self.all_data = []

class LastSummary(BaseSummary):

    def __init__(self, name):
        self.name = name
        self.reset()

    def add(self, data):
        if data is not None:
            self.log_value = data

    def write_metrics(self):
        sumagg = SumAgg()

        if self.log_value is not None:
            sumagg.tensorboard_writer.add_scalar(f"{sumagg.phase}/{self.name}", self.log_value, sumagg.epoch)
            sumagg.epoch_metrics[self.name] = self.log_value

    def reset(self):
        self.log_value = None

class HistAppendSummary(BaseSummary):

    def __init__(self, name):
        self.name = name
        self.reset()

    def add(self, data):
        self.all_data.append(data)

    def write_metrics(self):
        sumagg = SumAgg()

        if len(self.all_data) > 0:
            all_data = np.array(self.all_data)
            sumagg.tensorboard_writer.add_histogram(f"{sumagg.phase}/{self.name}", all_data, sumagg.epoch)

            sumagg.epoch_metrics[self.name] = all_data

    def reset(self):
        self.all_data = []

class NoRestrictionImageSummary(BaseSummary):

    def __init__(self, name):
        self.name = name
        self.reset()

    def add(self, data):
        if data is not None: # Unless None, Overwrite
            self.log_img = data

    def write_metrics(self):
        sumagg = SumAgg()

        if self.log_img is not None:
            if len(self.log_img.shape) == 2:
                self.log_img = np.expand_dims(self.log_img, 0)

            sumagg.tensorboard_writer.add_image(f"{sumagg.phase}/{self.name}", self.log_img, sumagg.epoch)
            sumagg.epoch_metrics[self.name] = self.log_img

    def reset(self):
        self.log_img = None


text:
from collections import deque
from dataclasses import dataclass

import numpy as np

from day import Day
from utils import ITER_END_MARKER

def iter_with_terminator(iterable, end_marker=ITER_END_MARKER, include_end_marker=True, transform=None, predicate=None):
    if predicate is None:
        predicate = lambda x: True

    if transform is None:
        transform = lambda x: x

    # iterators return themselves if you call iter() on it, this is to ensure it's a iterator
    it = iter(iterable)

    while (value := next(it, end_marker)) != end_marker:
        if predicate(value):
            yield transform(value)

    if include_end_marker:
        yield end_marker

class LineType:
    BOARD_LINE = 'BOARD_LINE'
    HEADER = 'HEADER'
    SEPARATOR = 'SEPARATOR'

@dataclass
class Board:
    grid: np.ndarray
    number_locations: dict[int, tuple[int, int]] = None
    flattened: np.ndarray = None

    def __post_init__(self):
        self.number_locations = {value: index for index, value in np.ndenumerate(self.grid)}
        self.flattened = np.ravel(self.grid)

    def active_indexes_from_numbers(self, active_numbers):
        return [index for number, index in self.number_locations.items() if number in active_numbers]

class Day4Part2(Day):
    day = 4
    part = 2

    def get_sample_input(self):
        return ('7,4,9,5,11,17,23,2,0,14,21,24,10,16,13,6,15,25,12,22,18,20,8,19,3,26,1\n'
                '\n'
                '22 13 17 11  0\n'
                ' 8  2 23  4 24\n'
                '21  9 14 16  7\n'
                ' 6 10  3 18  5\n'
                ' 1 12 20 15 19\n'
                '\n'
                ' 3 15  0  2 22\n'
                ' 9 18 13 17  5\n'
                '19  8  7 25 23\n'
                '20 11 10 24  4\n'
                '14 21 16 12  6\n'
                '\n'
                '14 21 17 24  4\n'
                '10 16 15  9 19\n'
                '18  8 23 26 20\n'
                '22 11 13  6  5\n'
                ' 2  0 12  3  7')

    def parse_input(self):
        drawing_numbers = deque()
        buffer = []
        boards = []

        def _trans(x):
            if ',' in x:
                return x, LineType.HEADER
            if x:
                return x, LineType.BOARD_LINE
            return LineType.SEPARATOR

        for pair in iter_with_terminator(self.input_text_lines, transform=_trans, end_marker=LineType.SEPARATOR):
            match pair:
                case (x, LineType.HEADER):
                    drawing_numbers.extend(map(int, x.split(',')))
                case (x, LineType.BOARD_LINE):
                    buffer.append([int(x) for x in x.split()])
                case LineType.SEPARATOR if buffer:
                    boards.append(buffer)
                    buffer = []

        return drawing_numbers, [Board(np.array(b)) for b in boards]

    def check_win(self, board: Board, drawn_numbers: set[int]):
        for y, x in board.active_indexes_from_numbers(drawn_numbers):
            if all(val in drawn_numbers for val in board.grid[y, :]) or all(val in drawn_numbers for val in board.grid[:, x]):
                return True
        return False

    def solve(self):
        numbers, boards = self.parse_input()
        drawn = set()
        last_draw = None
        last_board = None

        while boards:
            last_draw = numbers.popleft()
            drawn.add(last_draw)
            boards = [b for b in boards if not self.check_win(b, drawn)]
            if boards:
                last_board = boards[-1]

        print(sum(set(last_board.flattened) - drawn) * last_draw)


text:
import numpy as np
import cv2
import matplotlib.pyplot as plt
from functools import reduce
from funcoes import *

img = cv2.imread('img/moire.tif')
img = img[:,:,0] # Retiro somente um canal, pois já está em escala de cinza
cv2.imshow('Original', cv2.resize(img, None, fx=2,fy=2))

h, w = img.shape
padded = np.zeros((2*h, 2*w), dtype=np.uint8)
padded[int(h/2):int(3*h/2),int(w/2):int(3*w/2)] = img

transform = fourier_transform(padded)
cv2.imshow('Transformada', exibicao(getMagDB(transform)))
# plt.matshow(getMagDB(transform))
# plt.show()

# Posição retirada da visualização da imagem
# notch_pos = (114, 160)
# Gero posições espelhadas nos 4 quadrantes
# notches = mirrorAroundCenter(notch_pos, transform)

# Gero posições dos notches
notches = [
  (108,170),
  (222,160),
  (112,328),
  (228,320),
]

# Crio filtros notch nas 4 posições
filtros = [getButterworth(notch, transform, D0=30, n=4, invert=True) for notch in notches]
# Junto os filtros
filtro = reduce(lambda a, b: a*b, filtros)
cv2.imshow('Filtro', exibicao(255*getMag(filtro)))

# Multiplico pela máscara
filtrado = np.multiply(transform, filtro)
cv2.imshow('Filtrado', exibicao(getMagDB(filtrado)))

# Faço a transformada inversa
back = inverse_transform(filtrado)[int(h/2):int(3*h/2),int(w/2):int(3*w/2)]
cv2.imshow('Resultado', cv2.resize(back, None, fx=2, fy=2))

# cv2.imwrite('img/moire.png', img)
# cv2.imwrite('img/resultado/problema2_filtrada.png', back)

cv2.waitKey(0)
cv2.destroyAllWindows()