Change constant into if statement

I have create ERC721-B which requires the wallet to hold ERC721-A in order to be able to mint. This works well (as per intention).

Now I would like to "force" the holder of ERC721-A to keep the NFT and not sell it by offering a higher reward rate, by checking the balance of ERC721-A.

ERC721-B will now be staked in my contract for 1 per block, but if the wallet holds ERC721-A when staking, the reward rate will be increased by 1 or more depending on the amount of ERC721-A the wallet has.

Now the REWARDS_PER_BLOCK is a constant:

uint256 constant REWARDS_PER_BLOCK = 1;

and I check the balance like this:

uint256 balance = nft721.balanceOf(msg.sender);
require(balance>0,Error no nft1);

Is there any way I can make a statement like below and have a variable reward based on NFT's held in wallet?:

if(balance==1) {
		return 1;
	} else if (balance==2) {
		return 2;
	} else if (balance<=5) {
		return 4;
	} else (balance>5) {
		return 8;
	}
}

So I tried to create a function like this:

function REWARDS_PER_BLOCK2(uint256 balance) public view returns (uint256) {
        uint256 balance = nft721.balanceOf(msg.sender);
        if(balance == 0) {
            return 1;
        } else if (balance >= 1) {
            return 10;
        }
    }

And this checks the balance of the ERC721 in the wallet just fine, and returns the correct number according to the function, but HOW do include the return in another function? I was hoping that I could "import" the return from the function above this text into the funcion below?

    function calculateStakingRewards(uint256 tokenID) public view returns (uint256) {
        require(lastClaimedBlockForToken[tokenID] != 0, "token not staked");
        uint256 toBlock = rewardsEnd < block.number ? rewardsEnd : block.number;
        return REWARDS_PER_BLOCK2 * (toBlock - lastClaimedBlockForToken[tokenID]);
    }

Hi @Kion_Larsen, This is a pretty basic programming question so I won't provide a lot of guidance, but REWARDS_PER_BLOCK2 is a function that you will need to call like this: REWARDS_PER_BLOCK2(...) where ... are the function arguments.