from sklearn.svm import LinearSVC
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from skimage.feature import hog
from cv2 import imread
from sklearn.externals import joblib

import random as rand
import numpy as np 
import cv2
import glob
import time

import matplotlib.pyplot as plt

from helpers import convert, show_images
from featuresourcer import FeatureSourcer

sourcer_params = {
  'color_model': 'yuv',                # hls, hsv, yuv, ycrcb
  'bounding_box_size': 64,             #
  'number_of_orientations': 11,        # 6 - 12
  'pixels_per_cell': 16,               # 8, 16
  'cells_per_block': 2,                # 1, 2
  'do_transform_sqrt': True
}

start_frame = imread("/home/h/Pulpit/exit/cropP_20191121_003105_vHDR_Auto.jpg")
sourcer = FeatureSourcer(sourcer_params, start_frame)

print("Loading images to memory...")
t_start = time.time()

vehicle_imgs, nonvehicle_imgs = [], []
vehicle_paths = glob.glob('/home/h/Pulpit/exit/*.jpg')
nonvehicle_paths = glob.glob('/home/h/Pulpit/non-vehicles/*/*.png')

def exit_reader(path):
    img = cv2.imread(path)
    return cv2.resize(img, (64,64))

for path in vehicle_paths: vehicle_imgs.append(exit_reader(path))
for path in nonvehicle_paths: nonvehicle_imgs.append(imread(path))

vehicle_imgs, nonvehicle_imgs = np.asarray(vehicle_imgs), np.asarray(nonvehicle_imgs)
total_vehicles, total_nonvehicles = vehicle_imgs.shape[0], nonvehicle_imgs.shape[0]

print("... Done")
print("Time Taken:", np.round(time.time() - t_start, 2))
print("Vehicle images shape: ", vehicle_imgs.shape)
print("Non-vehicle images shape: ", nonvehicle_imgs.shape)

print("Extracting features... This might take a while...")
t_start = time.time()

vehicles_features, nonvehicles_features = [], []

print("Vehicles...")
for img in vehicle_imgs:
  vehicles_features.append(sourcer.features(img))
  print('█', end = '')

print()
print("Non-Vehicles...")
for img in nonvehicle_imgs:
  nonvehicles_features.append(sourcer.features(img))
  print('█', end = '')
                         
vehicles_features = np.asarray(vehicles_features)
nonvehicles_features = np.asarray(nonvehicles_features)

print()
print("...Done")
print("Time Taken:", np.round(time.time() - t_start, 2))
print("Vehicles features shape: ", vehicles_features.shape)
print("Non-vehicles features shape: ", nonvehicles_features.shape)

print("Scaling features...")
t_start = time.time()

unscaled_x = np.vstack((vehicles_features, nonvehicles_features)).astype(np.float64)
scaler = StandardScaler().fit(unscaled_x)
x = scaler.transform(unscaled_x)
y = np.hstack((np.ones(total_vehicles), np.zeros(total_nonvehicles)))

print("...Done")
print("Time Taken:", np.round(time.time() - t_start, 2))
print(" x shape: ", x.shape, " y shape: ", y.shape)

print("Training classifier...")
t_start = time.time()

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size = 0.2,
                                                    random_state = rand.randint(1, 100))
svc = LinearSVC()
svc.fit(x_train, y_train)
accuracy = svc.score(x_test, y_test)

print("...Done")
print("Time Taken:", np.round(time.time() - t_start, 2))
print("Accuracy: ", np.round(accuracy, 4))

sourcer = FeatureSourcer(sourcer_params, start_frame)

f1 = sourcer.features(nonvehicle_imgs[rand.randint(0, total_nonvehicles)])
f2 = sourcer.features(vehicle_imgs[rand.randint(0, total_vehicles)])
f3 = sourcer.features(vehicle_imgs[rand.randint(0, total_vehicles)])
f4 = sourcer.features(nonvehicle_imgs[rand.randint(0, total_nonvehicles)])

f = scaler.transform([f1, f2, f3, f4])
print(svc.predict(f))

print ("Saving models...")

joblib.dump(svc, 'svc2.pkl')
joblib.dump(scaler, 'scaler2.pkl')

print("...Done")

print ("Loading models...")

svc = joblib.load('svc2.pkl')
scaler = joblib.load('scaler2.pkl')

print("...Done")

sourcer = FeatureSourcer(sourcer_params, start_frame)

f1 = sourcer.features(nonvehicle_imgs[rand.randint(0, total_nonvehicles)])
f2 = sourcer.features(vehicle_imgs[rand.randint(0, total_vehicles)])
f3 = sourcer.features(vehicle_imgs[rand.randint(0, total_vehicles)])
f4 = sourcer.features(nonvehicle_imgs[rand.randint(0, total_nonvehicles)])

f = scaler.transform([f1, f2, f3, f4])
print(svc.predict(f))