Wrote som dirty code, cleaned it up