The solution is quite simple. There is an app by Cordless Dog called Stay. When you run it, OS X asks you to change your security settings to allow Stay to control your computer. Once that is done, it's ready to go.
Stay can be configured to launch on startup, which is not enabled by default, but that's the option you would really want. The other option, which is on by default is to restore windows as displays are connected and disconnected, that is the option that fixes the issue.
I tested it by turning my display on and off and it works perfectly.
The app costs $15, which is money well spent.