import numpy as np
from sklearn import preprocessing
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import scipy.ndimage.interpolation as compress
from sklearn import linear_model
import pandas as pd
from sklearn.feature_selection import SelectKBest
from sklearn.feature_selection import chi2

X_cropped = np.load('/home/evrova/MLTask1/ml-project/data/X_cropped.npy')
y_train = np.genfromtxt('/home/evrova/MLTask1/ml-project/data/y_1.csv')

print("Loading done")

NO_SAMPLES = X_cropped.shape[0]
#SLICES_AXE_SIZE = 176

X_train_3D = np.reshape(X_cropped, (-1, 136, 170, 144))

#We need an array of indexes of the non-zero pixels to see which one is the maximum
# non_zero_indexes = np.empty
# maximal_values = np.empty((1, 1))


#print(type(X_cropped))
# min_x = 176
# max_x = 0
# min_y = 208
# max_y = 0
# min_z = 176
# max_z = 0
#
#
# for i in range(0, NO_SAMPLES):
#     for z in range(0, SLICES_AXE_SIZE):
#         for x in range(0, 176):
#             non_zero_indexes_y = np.argwhere(X_train_3D[i,x,:,z])
#             if non_zero_indexes_y.size > 0:
#                 min_y = min(min_y, np.amin(non_zero_indexes_y))
#                 max_y = max(max_y, np.amax(non_zero_indexes_y))
#
#
# for i in range(0, NO_SAMPLES):
#     for z in range(0, SLICES_AXE_SIZE):
#         for y in range(0, 208):
#             non_zero_indexes_x = np.argwhere(X_train_3D[i,:,y,z])
#             if non_zero_indexes_x.size > 0:
#                 min_x = min(min_x, np.amin(non_zero_indexes_x))
#                 max_x = max(max_x, np.amax(non_zero_indexes_x))
#
# for i in range(0, NO_SAMPLES):
#     for x in range(0, 176):
#         for y in range(0, 208):
#             non_zero_indexes_z = np.argwhere(X_train_3D[i,x,y,:])
#             if non_zero_indexes_z.size > 0:
#                 min_z = min(min_z, np.amin(non_zero_indexes_z))
#                max_z = max(max_z, np.amax(non_zero_indexes_z))
NO_PATCHES = 10000
PATCH_SIZE = 5

X_test = np.load('/home/evrova/MLTask1/ml-project/data/X_test.npy')
X_test_3D = np.reshape(X_test, (-1, 176, 208, 176))
X_test_cropped = X_test_3D[:, 18:154, 18:188, 7:151]

averages = np.zeros((NO_SAMPLES, NO_PATCHES))
averages_test = np.zeros((X_test_cropped.shape[0], NO_PATCHES))
k = 0
temp_average = 1

for r in range(0, NO_SAMPLES):
    np.random.seed(10)  # TO CHECK
    k = 0
    while(k<NO_PATCHES):
        distr_x = np.random.randint(low=5, high=X_cropped.shape[1]-5)
        distr_y = np.random.randint(low=5, high=X_cropped.shape[2]-5)
        distr_z = np.random.randint(low=5, high=X_cropped.shape[3]-5)


        temp_average = np.average(X_cropped[r,distr_x:distr_x+PATCH_SIZE-1,distr_y:distr_y+PATCH_SIZE-1,distr_z:distr_z+PATCH_SIZE-1])
        if r < X_test_cropped.shape[0]:
            temp_average_test = np.average(X_test_cropped[r, distr_x:distr_x + PATCH_SIZE - 1, distr_y:distr_y + PATCH_SIZE - 1, distr_z:distr_z + PATCH_SIZE - 1])

        if temp_average > 1:
            averages[r, k] = temp_average
            if r < X_test_cropped.shape[0]:
                averages_test[r,k] = temp_average_test
            k = k + 1

slct = SelectKBest(k=5000)
X_new = slct.fit_transform(averages, y_train)
X_test_new = slct.transform(averages_test)
#
# NO_BINS = np.sqrt(X_new.shape[1])
# hist, bins = np.histogram(X_new, NO_BINS)
#
# width = 0.7 * (bins[1] - bins[0])
# center = (bins[:-1] + bins[1:]) / 2
# plt.bar(center, hist, align='center', width=width)
# plt.show()


print("Before fitting")

model = linear_model.LassoLarsCV(cv=5)
model.fit(X_new, y_train)

print("After fitting")


err = model.score(X_new, y_train)

err1 = model.predict(X_new)

error = np.sum((y_train - err1)**2)


print("Score: ", err)
print("Our score: ", error)



y_test = model.predict(X_test_new)

y_test = pd.DataFrame(y_test)
y_test.to_csv('test.csv', sep=',')


print("Done!")