Merge parts of Stage 4 implementation

Fixes bugs #38, #40 and #51