Fixed logic error and removed unused import