import * as anchor from "@project-serum/anchor";
import { Program, BN } from "@project-serum/anchor";
import { assert } from "chai";
import { Keypair, PublicKey } from "@solana/web3.js";
import { PsyStake } from "../target/types/psy_stake";
import {
  deriveStakePool,
  deriveStakingTokenAccount,
  deriveRewardPoolPdas,
  deriveStakingRecord,
} from "./pdas";
const { setupTestAccounts, MIN_EPOCH_DURATION } = require("./setup");
const {
  getAccountData,
  defaultVec,
  waitForStakePoolEpoch,
  calculateLockupDuration,
  assertParsedError,
} = require("./utils");
import {
  stakeTokenRpc,
  claimRewardRpc,
  createStakingRecordRpc,
  createRewardPoolRpc,
  createRewardRecordRpc,
  createStakePoolRpc,
  updateRewardPoolRpc,
} from "../packages/psystake/src/instructions";
import {
  DistributionType,
  calculateRewardUnits,
} from "../packages/psystake/src";

describe("Test stake token instructions.", () => {
  const program = anchor.workspace.PsyStake as Program<PsyStake>;
  const provider = anchor.Provider.env();
  anchor.setProvider(provider);

  let userUsdcTokenAccount: PublicKey,
    userBtcTokenAccount: PublicKey,
    network,
    userWalletKeypair: Keypair,
    userWalletKeypair2: Keypair,
    userUsdcTokenAccount2: PublicKey,
    stakePoolAuthority: Keypair,
    rewardPoolAuthority: Keypair,
    mintAuthority: PublicKey;

  const uniqueSeed = 1;

  let stakePool: PublicKey,
    stakingTokenAccount: PublicKey,
    stakingTokenMint: PublicKey,
    rewardTokenAccount: PublicKey,
    stakingRecord: PublicKey,
    stakingRecordBump: number;

  const startingEpoch = 1;
  let poolId = 1;
  const epochDuration = new BN(MIN_EPOCH_DURATION);

  const distributionType = DistributionType.CONSTANT;
  const rewardPerEpoch = new BN(1000000);
  const epochRewardDecimals = 8;

  const createStakingRecordHelper = async (
    recordOwner = userWalletKeypair.publicKey,
    signer = userWalletKeypair
  ) => {
    await createStakingRecordRpc(program, recordOwner, stakePool, [signer]);
  };

  const stakeTokenHelper = async (
    stakeAmount,
    lockUpPeriod,
    ownerTokenAccount,
    recordOwner = userWalletKeypair.publicKey,
    signer = userWalletKeypair
  ) => {
    await stakeTokenRpc(
      program,
      recordOwner,
      stakePool,
      ownerTokenAccount,
      stakeAmount,
      lockUpPeriod,
      [signer]
    );
  };

  const createRewardPoolHelper = async (startingEpoch, poolId) => {
    await createRewardPoolRpc(
      program,
      stakePoolAuthority.publicKey,
      rewardPoolAuthority.publicKey,
      stakePool,
      network.btcToken.publicKey,
      poolId,
      startingEpoch,
      distributionType,
      rewardPerEpoch,
      epochRewardDecimals,
      [stakePoolAuthority]
    );
  };

  const updateRewardPoolHelper = async (isActive, poolId) => {
    await updateRewardPoolRpc(
      program,
      rewardPoolAuthority.publicKey,
      stakePool,
      poolId,
      distributionType,
      rewardPerEpoch,
      isActive,
      [rewardPoolAuthority]
    );
  };

  const claimRewardHelper = async (
    startRewardRecordNumber,
    endRewardRecordNumber
  ) => {
    await claimRewardRpc(
      program,
      userWalletKeypair.publicKey,
      stakePool,
      userBtcTokenAccount,
      startRewardRecordNumber,
      endRewardRecordNumber,
      poolId,
      [userWalletKeypair]
    );
  };

  const createRewardRecordHelper = async (epochToCreate) => {
    await createRewardRecordRpc(
      program,
      userWalletKeypair.publicKey,
      stakePool,
      poolId,
      epochToCreate,
      [userWalletKeypair]
    );
  };

  before(async () => {
    ({
      userWalletKeypair,
      userUsdcTokenAccount,
      userBtcTokenAccount,
      userWalletKeypair2,
      userUsdcTokenAccount2,
      mintAuthority,
      network,
      stakePoolAuthority,
      rewardPoolAuthority,
    } = await setupTestAccounts(program, provider));
    stakingTokenMint = network.usdcToken.publicKey;

    [stakePool] = await deriveStakePool(
      program,
      network.usdcToken.publicKey,
      new BN(uniqueSeed)
    );
    [stakingTokenAccount] = await deriveStakingTokenAccount(program, stakePool);

    const nextEpochStartTime = new BN(new Date().getTime() / 1000 + 3);

    // Create a StakePool.
    await createStakePoolRpc(
      program,
      stakePoolAuthority.publicKey,
      stakePool,
      stakingTokenMint,
      uniqueSeed,
      nextEpochStartTime,
      epochDuration,
      [stakePoolAuthority]
    );

    // Create a RewardPool.
    const poolAccounts = await deriveRewardPoolPdas(
      program,
      stakePool,
      poolId,
      startingEpoch
    );
    ({ rewardTokenAccount } = poolAccounts);
    await createRewardPoolHelper(startingEpoch, poolId);

    // Create a StakingRecord.
    [stakingRecord, stakingRecordBump] = await deriveStakingRecord(
      program,
      userWalletKeypair.publicKey,
      stakePool
    );
    await createStakingRecordHelper();

    // Mint tokens to reward account.
    await network.btcToken.mintTo(
      rewardTokenAccount,
      mintAuthority,
      [],
      10 ** 6
    );
  });

  it("Stake token fails when owner token account is from different mint.", async () => {
    const stakeAmount = new BN(500);
    const lockUpPeriod = 0;

    try {
      await stakeTokenHelper(stakeAmount, lockUpPeriod, userBtcTokenAccount);
    } catch (err) {
      assertParsedError(err, "Owner token account is invalid.");
    }
  });

  it("Stake token fails when amount in owner's account is insufficient.", async () => {
    const stakeAmount = new BN(10 ** 12);
    const lockUpPeriod = 0;

    try {
      await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);
    } catch (err) {
      assertParsedError(err, "Insufficient amount for staking or unstaking.");
    }
  });

  it("Stake token successfully with no lockup.", async () => {
    const stakeAmount = new BN(1000);
    const lockUpPeriod = 0;

    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountData = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordData = accountData.stakingRecord;

    const expectedRewardUnits = calculateRewardUnits(stakeAmount, lockUpPeriod);

    // Verify reward units in StakePool.
    assert(accountData.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordData.rewardUnits.eq(expectedRewardUnits));
    assert(stakingRecordData.stakedAmount.eq(stakeAmount));
    assert(accountData.stakingTokenAccountAmount.eq(stakeAmount));

    // Other fields should remain unchanged.
    assert(stakingRecordData.stakePool.equals(stakePool));
    assert(stakingRecordData.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordData.recordBump, stakingRecordBump);

    // Last epoch claimed still remains the same, at current epoch of 0.
    assert.deepEqual(stakingRecordData.lastEpochClaimedVec, defaultVec());

    // Verify that lock up should not be in the future.
    const currentTime = new Date().getTime() / 1000;
    assert(stakingRecordData.lockUpExpiry.toNumber() <= currentTime);
  });

  it("Stake additonal tokens successfully for a longer lockUpPeriod 1.", async () => {
    const stakeAmount = new BN(2000);
    const lockUpPeriod = 1;

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    // Last epoch claimed still remains the same, at current epoch of 0.
    assert.deepEqual(stakingRecordAfter.lastEpochClaimedVec, defaultVec());

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake additonal tokens successfully for a longer lockUpPeriod 2.", async () => {
    const stakeAmount = new BN(2000);
    const lockUpPeriod = 2;

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    // Last epoch claimed still remains the same, at current epoch of 0.
    assert.deepEqual(stakingRecordAfter.lastEpochClaimedVec, defaultVec());

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake additonal tokens successfully for a longer lockUpPeriod 3.", async () => {
    const stakeAmount = new BN(1232);
    const lockUpPeriod = 3;

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    // Last epoch claimed still remains the same, at current epoch of 0.
    assert.deepEqual(stakingRecordAfter.lastEpochClaimedVec, defaultVec());

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake additonal tokens successfully for a longer lockUpPeriod 4.", async () => {
    const stakeAmount = new BN(2222);
    const lockUpPeriod = 4;

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    // Last epoch claimed still remains the same, at current epoch of 0.
    assert.deepEqual(stakingRecordAfter.lastEpochClaimedVec, defaultVec());

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake token fails when existing lockup is longer than chosen lockup.", async () => {
    const stakeAmount = new BN(1000);
    const lockUpPeriod = 2;

    try {
      await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);
    } catch (err) {
      assertParsedError(err, "Selected lockup period is too short.");
    }
  });

  it("Stake token fails when there are some rewards unclaimed.", async () => {
    // Wait for Epoch 3.
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      3
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    const stakeAmount = new BN(1000);
    const lockUpPeriod = 4;

    try {
      await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);
    } catch (err) {
      assertParsedError(err, "Some RewardPool is unclaimed.");
    }
  });

  it("Stake token succeeds after rewards are claimed.", async () => {
    // Create RewardRecords for epoch 2 and 3.
    await createRewardRecordHelper(2);
    await createRewardRecordHelper(3);

    // Claim rewards for epoch 1 and 2.
    await claimRewardHelper(1, 2);

    const stakeAmount = new BN(2222);
    const lockUpPeriod = 4;

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    // Expect first RewardPool to be claimed up to epoch 3, when the stake instruction was called.
    const expectedLastClaimedVec = defaultVec();
    expectedLastClaimedVec[0] = 3;
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      expectedLastClaimedVec
    );

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake token succeeds even when there is a RewardPool that has not started yet.", async () => {
    // Create a new RewardPool that starts in epoch 5.
    await createRewardPoolHelper(5, 2);

    // Wait for Epoch 4.
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      4
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    // Create RewardRecord for first RewardPool for epoch 4.
    await createRewardRecordHelper(4);

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    const stakeAmount = new BN(323);
    const lockUpPeriod = 4;
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    // Expect first and second RewardPool to be claimed up to epoch 4, when the stake instruction was called.
    const expectedLastClaimedVec = defaultVec();
    expectedLastClaimedVec[0] = 4;
    expectedLastClaimedVec[1] = 4;
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      expectedLastClaimedVec
    );

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake token when RewardPool is turned inactive during the epoch, with last claimed set at the same epoch.", async () => {
    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    // Set first RewardPool as inactive during epoch 4.
    await updateRewardPoolHelper(false, 1);

    // Expect first and second RewardPool to be claimed up to epoch 4, when the stake instruction was called.
    const expectedLastClaimedVec = defaultVec();
    expectedLastClaimedVec[0] = 4;
    expectedLastClaimedVec[1] = 4;
    assert.deepEqual(
      stakingRecordBefore.lastEpochClaimedVec,
      expectedLastClaimedVec
    );

    const stakeAmount = new BN(24);
    const lockUpPeriod = 4;

    // Stake token during epoch 4.
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      expectedLastClaimedVec
    );

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake token after RewardPool has been inactive for more than one epoch.", async () => {
    // Wait for Epoch 5.
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      5
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    const stakeAmount = new BN(222);
    const lockUpPeriod = 4;
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    const expectedLastClaimedVec = defaultVec();
    // Expect first RewardPool to be claimed up to epoch 4, when it was last active.
    expectedLastClaimedVec[0] = 4;
    // Expect second RewardPool to be claimed up to epoch 5, when the stake instruction was called.
    expectedLastClaimedVec[1] = 5;
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      expectedLastClaimedVec
    );

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake token when RewardPool is turned inactive during the epoch, with last claimed set at a previous epoch.", async () => {
    // Wait for Epoch 6.
    let isTargetEpochReached = await waitForStakePoolEpoch(
      program,
      stakePool,
      6
    );
    assert(isTargetEpochReached);
    console.log("Target Epoch Reached");

    const accountDataBefore = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordBefore = accountDataBefore.stakingRecord;

    // Set second RewardPool as inactive in epoch 6.
    await updateRewardPoolHelper(false, 2);

    const expectedLastClaimedVec1 = defaultVec();
    // Expect first RewardPool to be claimed up to epoch 4, when it was last active.
    expectedLastClaimedVec1[0] = 4;
    // Expect second RewardPool to be claimed up to epoch 5.
    expectedLastClaimedVec1[1] = 5;
    assert.deepEqual(
      stakingRecordBefore.lastEpochClaimedVec,
      expectedLastClaimedVec1
    );

    const stakeAmount = new BN(222);
    const lockUpPeriod = 4;

    // Try staking token during epoch 6, before second RewardPool finishes the inactive epoch.
    await stakeTokenHelper(stakeAmount, lockUpPeriod, userUsdcTokenAccount);

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord,
      stakingTokenAccount,
    });
    const stakingRecordAfter = accountDataAfter.stakingRecord;

    // Reward units is calculated with entire staked amount using new lockup period.
    const expectedRewardUnits = calculateRewardUnits(
      stakeAmount.add(stakingRecordBefore.stakedAmount),
      lockUpPeriod
    );

    // Verify reward units in StakePool.
    assert(accountDataAfter.stakePool.totalRewardUnits.eq(expectedRewardUnits));

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordAfter.rewardUnits.eq(expectedRewardUnits));
    assert(
      stakingRecordAfter.stakedAmount
        .sub(stakingRecordBefore.stakedAmount)
        .eq(stakeAmount)
    );
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    // Other fields should remain unchanged.
    assert(stakingRecordAfter.stakePool.equals(stakePool));
    assert(stakingRecordAfter.recordOwner.equals(userWalletKeypair.publicKey));
    assert.equal(stakingRecordAfter.recordBump, stakingRecordBump);

    const expectedLastClaimedVec2 = defaultVec();
    // Expect first RewardPool to be claimed up to epoch 4, when it was last active.
    expectedLastClaimedVec2[0] = 4;
    // Expect second RewardPool to be claimed up to epoch 6, the current epoch.
    expectedLastClaimedVec2[1] = 6;
    assert.deepEqual(
      stakingRecordAfter.lastEpochClaimedVec,
      expectedLastClaimedVec2
    );

    // Verify that lock up period is within range of expected expiry time.
    const expectedExpiryTime =
      new Date().getTime() / 1000 + calculateLockupDuration(lockUpPeriod);
    const actualExpiryTime = stakingRecordAfter.lockUpExpiry.toNumber();
    assert(actualExpiryTime > expectedExpiryTime - 3);
    assert(actualExpiryTime < expectedExpiryTime + 3);
  });

  it("Stake token should increment total reward units in StakePool", async () => {
    // Create a StakingRecord owned by second user.
    const [stakingRecord2, stakingRecordBump2] = await deriveStakingRecord(
      program,
      userWalletKeypair2.publicKey,
      stakePool
    );
    await createStakingRecordHelper(
      userWalletKeypair2.publicKey,
      userWalletKeypair2
    );

    const accountDataBefore = await getAccountData(program, {
      stakingTokenAccount,
      stakePool,
    });

    const stakeAmount = new BN(1000);
    const lockUpPeriod = 1;

    // Expect stake token to succeed even when all RewardPools are inactive.
    await stakeTokenHelper(
      stakeAmount,
      lockUpPeriod,
      userUsdcTokenAccount2,
      userWalletKeypair2.publicKey,
      userWalletKeypair2
    );

    const accountDataAfter = await getAccountData(program, {
      stakePool,
      stakingRecord: stakingRecord2,
      stakingTokenAccount,
    });
    const stakingRecordData = accountDataAfter.stakingRecord;
    const expectedRewardUnits = calculateRewardUnits(stakeAmount, lockUpPeriod);

    // Verify reward units in StakePool should increment by staked amount of user 2.
    assert(
      accountDataAfter.stakePool.totalRewardUnits
        .sub(accountDataBefore.stakePool.totalRewardUnits)
        .eq(expectedRewardUnits)
    );

    // Verify staked amount and reward units in StakingRecord changed.
    assert(stakingRecordData.rewardUnits.eq(expectedRewardUnits));
    assert(stakingRecordData.stakedAmount.eq(stakeAmount));
    assert(
      accountDataAfter.stakingTokenAccountAmount
        .sub(accountDataBefore.stakingTokenAccountAmount)
        .eq(stakeAmount)
    );

    assert(stakingRecordData.stakePool.equals(stakePool));
    assert(stakingRecordData.recordOwner.equals(userWalletKeypair2.publicKey));
    assert.equal(stakingRecordData.recordBump, stakingRecordBump2);

    const expectedLastClaimedVec = defaultVec();
    // Expect first RewardPool to be claimed up to epoch 4, when it was last active.
    expectedLastClaimedVec[0] = 4;
    // Expect second RewardPool to be claimed up to epoch 6, when it was last active.
    expectedLastClaimedVec[1] = 6;
    assert.deepEqual(
      stakingRecordData.lastEpochClaimedVec,
      expectedLastClaimedVec
    );
  });
});