Updated train logic