contract airdrop is Ownable {
    using SafeMath for uint256;
    using SafeERC20 for IERC20;

    uint256 public stakingEnd; // What block does staking end?
    uint256 public _totalSupply;
    mapping(address => uint256) private _balances;

    IERC20 public stakeToken; 
    IERC20 public rewardToken;


    function totalSupply() public view returns (uint256) {
        return _totalSupply;
    }

    function balanceOf(address account) public view returns (uint256) {
        return _balances[account];
    }

    function rewardRemaining() public view returns (uint256) {
        return rewardToken.balanceOf(address(this));
    }

    function stake(uint256 amount) public {
        require(stakeToken != IERC20(0x0), "address not set");
        require(block.number < stakingEnd, "can't stake after staking period is over"); 
        _totalSupply = _totalSupply.add(amount);
        _balances[msg.sender] = _balances[msg.sender].add(amount);
        stakeToken.safeTransferFrom(msg.sender, address(this), amount);
    }

    function withdraw() internal {
        require(stakeToken != IERC20(0x0), "address not set");
        require(block.number > stakingEnd, "can't unstake until after staking period is over"); 
        uint256 amount = _balances[msg.sender];
        _balances[msg.sender] = 0;
        _totalSupply = _totalSupply.sub(amount);
        stakeToken.safeTransfer(msg.sender, amount);
    }
 
    // set the staking token and the reward token
    function setStake(address contractAddress) public onlyOwner {
          stakeToken = IERC20(contractAddress);
    }
    function setReward(address contractAddress) public onlyOwner {
          rewardToken = IERC20(contractAddress);
    }

    function viewAllo(address staker) public view returns(uint256) {
          return _balances[staker].mul(1e18).div(_totalSupply); // divided by 1e18 is the % staked
    }

    function viewReward(address staker) public view returns(uint256) {
          return viewAllo(staker).mul(rewardRemaining()).div(1e18);
    }

    function claimRewards() public returns(bool) { // withdraw as well to prevent exploits
          require(block.number > stakingEnd);
          uint256 reward = viewReward(msg.sender);
          withdraw();
          rewardToken.safeTransfer(msg.sender, reward);
          return true;
    }

    function setEnd(uint256 blocknumber) public onlyOwner {
          require(block.number < blocknumber);
          stakingEnd = blocknumber;      
    }

}
function setStake(address contractAddress) public onlyOwner {
      stakeToken = IERC20(contractAddress);
}

// set the staking token and the reward token