Updated main code