TO DO lists are wonderful things.
I'm not talking about // #TODO: blah blah comments. I mean get an honest to God notebook.
You never know when you'll remember something important to do. A notebook will quietly sit there and let you think without complaining about how your handwriting wont compile. Some of my best ideas happen in the bathroom (yes I do own a water proof notebook but you don't have to go that far).
You can get pocket sized ones that are sewn (not glued) so they don't fall apart in your pocket. Didn't manage to get a fancy one with a built in book mark? Tape, scissors, ribbon and no one will ever know.
When an idea hits just jot it down. Draw little boxes next to each idea and you can easily mark it as done. Put a box at the top of the page and you know when the page is done.
What sequential access isn't good enough for you? Yeah they make pocket binders as well. This all might seem like a bit much but it's better than drowning in post it notes or trying to capture everything in Jira.
Don't leave things half implemented
Keep your improvements small and achievable. Don't start anything that can't be finished in one sitting. If it's to big for that then break it down into smaller steps. Always leave code that compiles and passes it's tests. Oh and don't leave passing tests you've never seen fail. Making a test both pass and fail is how you test the test.
Stop thinking you need the whole design on paper
What you need to do is capture your evolving plan. You don't know how things are going to look when you're done so stop pretending you do. Capture what you have figured out as well as you can. Use a napkin and crayon if you have to. Few people understand 90% of UML anyway. Use whatever way you can to show what you need to show. I focus on showing my interfaces and what knows about what.
Write notes when you stop coding
The moment you take your fingers off the keys is the last time you will understand what you've done (and what you have planned) as well as you do now. Capture that understanding as best you can in some notes. If all you have is comments then you're still tied to the computer and likely to leave a puddle in the chair. Again, having a notebook is an awesome thing.
This way you can land your brain gracefully, save your bladder, and take off again later without resorting to caffeine and teeth gritting.