Updated drivetrain and main code