added main class
fixed import
made a minor change