Cleaned up code and tweaked colour mapping further