import assert from 'assert';
import { BigRational } from 'big-rational-ts';
import * as types from './types.js';

const ZERO = new BigRational(0n, 1n);

function maxRational(a: BigRational, b: BigRational) {
  return a.gt(b) ? a : b;
}

function minRational(a: BigRational, b: BigRational) {
  return a.lt(b) ? a : b;
}

/**
 * Computes the next state of the portfolio after an interaction.
 * @param prices The prices of the assets in ADA
 * @param targetWeights The target weights of the assets in the portfolio
 * @param state The current state of the portfolio
 * @param adaInput The amount of ADA to be added (if negative removed) to the portfolio
 *
 * @returns The new state of the portfolio
 */
export function computeInteraction(
  prices: types.Prices,
  targetWeights: types.Weights,
  state: types.PortfolioState<bigint, bigint>,
  adaInput: BigRational,
): types.PortfolioState<BigRational, BigRational> {
  /** How much per asset in ADA in the portfolio before the deposit */
  const adaPerToken: types.RationalDict = {};

  const tvl = Object.entries(state.assets).reduce(
    (accTvl, [asset, assetAmount]) => {
      adaPerToken[asset] = new BigRational(assetAmount, 1n).mul(prices[asset]);
      return accTvl.add(adaPerToken[asset]);
    },
    ZERO,
  );

  /**
   * First condition: At least for one token in the collection,
   * there is a state for which this asset has the same amount of tokens
   * and the portfolio is balanced without removing tokens of a
   * different asset (i.e., it only needs to add tokens to the rest)
   *
   * Notice that tokens that satisfy this condition,
   * maximize the worth of tokens needed to reach balance.
   * If we find the asset with the max worth, we can know the minimum
   * amount of ADA needed to balance the portfolio.
   */

  const [maybeBalancedTokens, adaLeft] = getTokensToBalance(
    prices,
    targetWeights,
    adaPerToken,
    tvl,
    adaInput,
  );

  /**
   * Only enters here when the assets are actually balanced and there is more to
   * change in the portfolio
   */
  const adaLeftToModify = adaInput.gte(ZERO)
    ? adaLeft.gt(ZERO)
    : adaLeft.lt(ZERO);

  if (adaLeftToModify) {
    for (const [asset, amount] of Object.entries(maybeBalancedTokens)) {
      maybeBalancedTokens[asset] = adaLeft
        .mul(targetWeights[asset])
        .div(prices[asset])
        .add(amount)
        .reduce();
    }
  }

  const mtkPrice = tvl.div(new BigRational(state.mtkSupply, 1n));
  const newMtkSupply = tvl.add(adaInput).div(mtkPrice).reduce();
  return { assets: maybeBalancedTokens, mtkSupply: newMtkSupply };
}

/**
 * This function returns the amount of tokens that would end up in the portfolio
 * if the input is used to balance the portfolio.
 * @param prices The prices of the assets in ADA
 * @param targetWeights The target weights of the assets in the portfolio
 * @param adaPerToken The amount of ADA per token in the portfolio
 * @param tvl The total value locked in the portfolio
 * @param adaAvailable The amount of ADA available to balance the portfolio
 *
 * @returns The amount of tokens that would end up in the portfolio and how much
 * ADA would be left after the interaction
 *
 */
function getTokensToBalance(
  prices: types.Prices,
  targetWeights: types.Weights,
  adaPerToken: types.RationalDict,
  tvl: BigRational,
  adaAvailable: BigRational,
): [types.RationalDict, BigRational] {
  /** The actual token ratio of the portfolio */
  const weights: types.RationalDict = {};
  for (const [asset, adaEqAmount] of Object.entries(adaPerToken)) {
    weights[asset] = adaEqAmount.div(tvl).reduce();
  }

  // Already balanced => No changes needed
  if (
    Object.entries(targetWeights).every(([asset, targetWeight]) =>
      targetWeight.eq(weights[asset]),
    )
  ) {
    return [
      Object.fromEntries(
        Object.entries(targetWeights).map(([asset, weight]) => [
          asset,
          tvl.mul(weight).div(prices[asset]).reduce(),
        ]),
      ),
      adaAvailable,
    ];
  }

  /**
   * Base case for minimum tvl to balance can't be a token with target 0 as it may not be balanceable
   */
  const maybeBaseCase = Object.entries(adaPerToken).find(([asset]) => {
    const targetWeight = targetWeights[asset];
    assert(targetWeight, new Error('Invalid portfolio weights'));
    return !targetWeight.eq(ZERO);
  });

  assert(maybeBaseCase, new Error('Invalid portfolio assets'));
  const [baseCaseAsset, baseCaseValue] = maybeBaseCase;
  const baseCase = baseCaseValue.div(targetWeights[baseCaseAsset]);

  /**
   * The minimum balanced TVL if no tokens are removed in a deposit (or inserted in a withdraw) is determined by
   * the token for which, if nothing is added to it, the portfolio
   * needs the most (hence maxRational) ADA to complete the rest of
   * the ratios.
   *
   */
  const minTvlIfBalancedAndNoTokensRemovedOnDepositNorAddedOnWithdraw =
    Object.entries(adaPerToken).reduce((acc, [asset, assetValue]) => {
      // skip zero target weight
      if (targetWeights[asset].eq(ZERO)) {
        return acc;
      }
      /** Check if this is the token that needs the most (least in case of withdrawal) ADA to complete the rest */
      const tvlIfThisTokenAmountDoesntChange = assetValue.div(
        targetWeights[asset],
      );

      return adaAvailable.gte(ZERO)
        ? maxRational(tvlIfThisTokenAmountDoesntChange, acc)
        : minRational(tvlIfThisTokenAmountDoesntChange, acc);
    }, baseCase);

  /**
   * How much the input should be to have a a perfect balance after
   * the interaction (restricting to the tokens with non-zero weights
   */
  const tvlOfTokensWithZeroWeight = Object.entries(targetWeights)
    .filter(([_asset, weight]) => weight.eq(ZERO))
    .reduce((acc, [asset]) => {
      return acc.add(adaPerToken[asset]).reduce();
    }, ZERO);

  const tvlOfNonNullWeights = tvl
    .add(tvlOfTokensWithZeroWeight.negate())
    .reduce();

  const adaDiffToBalance = adaAvailable.gte(ZERO)
    ? minTvlIfBalancedAndNoTokensRemovedOnDepositNorAddedOnWithdraw
        .add(tvlOfNonNullWeights.negate())
        .reduce()
    : minTvlIfBalancedAndNoTokensRemovedOnDepositNorAddedOnWithdraw.add(
        tvl.negate(),
      );
  // If there is enough input to get to a balanced state
  const enoughAdaToBalance = adaAvailable.gte(ZERO)
    ? /*
       * If the input is positive, then the portfolio can be balanced with this
       * this interaction if the adaNecessaryToBalance is less than the input
       * Enough ADA (equivalent) coming in to add to the tokens that need it
       */
      adaDiffToBalance.lte(adaAvailable)
    : /*
       * If the input is negative, then the portfolio can be balanced with this
       * this interaction if the adaNecessaryToBalance is greater than the input
       * Removing enough ADA (equivalent) to the tokens that have too much
       */
      adaDiffToBalance.gte(adaAvailable);

  if (enoughAdaToBalance) {
    const balancedTokens: types.RationalDict = {};
    for (const [asset, targetWeight] of Object.entries(targetWeights)) {
      if (targetWeight.eq(ZERO) && adaAvailable.gte(ZERO)) {
        balancedTokens[asset] = adaPerToken[asset].div(prices[asset]).reduce();
      } else {
        balancedTokens[asset] =
          minTvlIfBalancedAndNoTokensRemovedOnDepositNorAddedOnWithdraw
            .mul(targetWeight)
            .div(prices[asset])
            .reduce();
      }
    }
    const adaLeft = adaAvailable.add(adaDiffToBalance.negate()).reduce();
    return [balancedTokens, adaLeft];
  }
  // there is not enough to balance it out
  else {
    const approximatedBalance: types.RationalDict = {};
    for (const [asset, assetAda] of Object.entries(adaPerToken)) {
      if (targetWeights[asset].eq(ZERO) && adaAvailable.gte(ZERO)) {
        approximatedBalance[asset] = assetAda.div(prices[asset]).reduce();
      } else {
        /**
         * How much of a given token is missing (or too much of) to reach
         * perfect balance
         */
        const tokenAdaDiffToBalance =
          minTvlIfBalancedAndNoTokensRemovedOnDepositNorAddedOnWithdraw
            .mul(targetWeights[asset])
            .add(assetAda.negate());

        approximatedBalance[asset] = adaAvailable
          .mul(tokenAdaDiffToBalance)
          .div(adaDiffToBalance)
          .div(prices[asset])
          .add(assetAda.div(prices[asset]))
          .reduce();
      }
    }
    const adaLeft = ZERO;
    return [approximatedBalance, adaLeft];
  }
}

/**
 * Computes how much of each fee in the smart contract this interaction contains
 */
export function computeFees({
  portfolioState,
  amount,
  type,
}: types.FeeComputation) {
  const batcherFee = portfolioState.batcherFee;

  const platformFee = amount
    .multipliedBy(portfolioState.platformFee)
    .div(10_000);

  const microMTKsInOrder = amount.div(portfolioState.microMtkPrice);

  // divided by 10_000 to turn the fee into a multiplying factor
  const userFee =
    type === 'mint'
      ? microMTKsInOrder.multipliedBy(portfolioState.entryFee).div(10_000)
      : microMTKsInOrder.multipliedBy(portfolioState.exitFee).div(10_000);

  return {
    /**
     * Batcher fee in lovelace
     */
    batcherFee: batcherFee.abs(),
    /**
     * Platform fee in lovelace
     */
    platformFee: platformFee.abs(),
    /**
     * User fees in microMTKs
     */
    userFee: userFee.abs(),
  };
}