Added additional logic