Add good State class and small changes