import numpy as np 
import matplotlib.pyplot as plt 
import pickle

from storm.sketch import *
from storm.hashes.srp_hash import SRPHash
from storm.hashes.asymmetric_simple_hash import AsymmetricSimpleHash
from storm.optimization.iterative_optimizer import IterativeOptimizer
from storm.optimization.utils import *
from storm.datasets import *
from storm.baselines import *


from sklearn import linear_model

import sys
import random
import pickle
import os 
import time 

testfile = 'data/autos.csv'
datafile = 'data/autos.csv'


reps = 200
p = 8

sigma = 0.1
n_iters = 1000
eta = 0.3
beta = 0.5
n_components = 8
use_intercept = True

x,y,x_test,y_test = load(datafile,testfile)

dataset = format_dataset(x,y,intercept = use_intercept)
N = dataset.shape[0]
d = dataset.shape[1]

# Sanity check baselines
baseline_model = optimal_linregress(x,y,intercept = use_intercept)
optimal_loss = linear_regression_loss(baseline_model, dataset) * 1.0/N
trivial_loss = np.sum(y**2)*1.0/N
mean_loss = np.sum( (y - np.mean(y))**2 )*1.0/N
print("Optimal Loss: ",optimal_loss)
print("Trivial Loss: ",trivial_loss)
print("Mean Loss: ",mean_loss)

for i in range(10):
	seed = i+1
	LSH = SRPHash(N = reps, d = d, p = p, seed = seed)

	S = STORM(reps, 2**p, LSH)

	S.addMulti(dataset)
	S.addMulti(-dataset)

	def constraint(theta):
		theta[-1] = -1
		return theta

	gradient = BallGradApprox(S, sigma, n_components)
	opt = IterativeOptimizer(S, gradient, constraint)

	theta = np.zeros(d)
	theta[-1] = -1
	theta, losses = opt.optimize(theta, eta, beta, n_iters, compute_function = True)

	storm_loss = linear_regression_loss(theta[:-1], dataset) * 1.0/N
	print("STORM Model Loss: ", storm_loss)

	plt.plot(losses)
	sys.stdout.flush()

plt.show()

