Fixed diagonal correction. Sort of works now.