fixed issues, improved code