import numpy as np
import matplotlib.pyplot as plt

numStudents = 500
numClasses = 30
numClassesPerStudent = 5

studentGoodness = np.random.uniform(0, 4, numStudents)
classDifficulty = np.random.uniform(0, 4, numClasses)

def noise():
	return np.random.normal(0, 1)

X, Y, naiveGpa = [], [], []
for i in range(numStudents):
	I = np.arange(numClasses)
	# Randomly choose class for students.  In this case we have strong students (above
	# 75th percentile) prefer harder classes and weaker students prefer easier classes.
	if studentGoodness[i] > 3.0:
		I = np.argsort(-classDifficulty + np.random.normal(0, 1, numClasses))
	else:
		I = np.argsort( classDifficulty + np.random.normal(0, 1, numClasses))
	for j in I[:numClassesPerStudent]:
		X.append([0] * (numStudents + numClasses))
		X[-1][i] = 1
		X[-1][numStudents + j] = 1
		Y.append(studentGoodness[i] - classDifficulty[j] + noise())
	naiveGpa.append(sum(Y[-10:]) / 10)

X = np.array(X)
Y = np.array(Y)
naiveGpa = np.array(naiveGpa)
w = np.linalg.lstsq(X, Y, rcond=0.0001)[0]
gpa = w[:numStudents]
classDifficulty = w[numStudents:]

w -= w.mean(); w /= w.std()
studentGoodness -= studentGoodness.mean(); studentGoodness /= studentGoodness.std()
naiveGpa -= naiveGpa.mean(); naiveGpa /= naiveGpa.std()

plt.scatter(studentGoodness, naiveGpa, label='naive gpa')
plt.scatter(studentGoodness, w[:numStudents], label='corrected gpa')
plt.legend()
plt.show()