Cleaned up code
Cleaned up code
Added missing features