Added figure environment for code blocks
Some changes to styling