Cleaned up the proof of concept a bit