import { Program } from "@coral-xyz/anchor";
import { PublicKey } from "@solana/web3.js";

import {
    SortBy,
    Rule,
    CreateBasketParams,
    NUM_OF_DAYS_IN_DATABASE,
    NUM_OF_TOKENS_IN_ASSET_POOL,
    BPS_DIVIDER,
    EXPO_DIVIDER,
    WEIGHT_MULTIPLIER,
    COMBINED_TOKENS_IN_A_BASKET,
} from "./config";
import { BasketsIDL } from "./basketsIDL";

export interface DataPoint {
    price: number,
    circulatingSupply: number,
    volume: number,
}

export interface BasketState {
    currentCompToken: number[],
    currentCompAmount: number[],
    currentCompWeight: number[],
    targetWeight: number[],
    numOfTokens: number,
    weightSum: number,
    basketWorth: number,
    lastRefilterTime: number,
    lastReweightTime: number,
    lastRebalanceTime: number,
    singleRuleAssets: number[][],
    ruleAssets: number[],
    ruleWeights: number[],
    numRuleTokens: number,
}

export interface TokenStats {
    stats: Stats[][],
}

export interface Stats {
    days: number,
    performance: number,
    volume: number,
    mcap: number,
}

export async function fetchDatabase(
    program: Program<BasketsIDL>,
    database: PublicKey,
) {
    let data: any = await program.account.database.fetch(database, "confirmed");
    
    let processedData: DataPoint[][] = [];
    for(let i=0; i<NUM_OF_TOKENS_IN_ASSET_POOL; i++){
        let index = data.data[i].index.toNumber();
        let currentTokenData: DataPoint[] = [];
        while(true) {
            let dataPoint: DataPoint = {
                price: data.data[i].price[index].toNumber(),
                circulatingSupply: data.data[i].circulatingSupply[index].toNumber(),
                volume: data.data[i].volume[index].toNumber(),
            };
            currentTokenData.push(dataPoint);
            index += 1;
            if(index == NUM_OF_DAYS_IN_DATABASE) index = 0;
            if(index == data.data[i].index.toNumber())break;
        }
        processedData.push(currentTokenData);
    }
    return processedData;
}

export function defaultBasketState() {
    let array: number[] = [1000000];
    let token: number[] = [0];
    let amount: number[] = [100000000];
    for(let i=1; i<20; i++){
        amount.push(0);
        array.push(0);
        token.push(0);
    }
    return {
        currentCompToken: token,
        currentCompAmount: amount,
        currentCompWeight: Object.assign([], array),
        targetWeight: Object.assign([], array),
        numOfTokens: 1,
        weightSum: 1000000,
        basketWorth: 100000000,
        lastRefilterTime: 0,
        lastReweightTime: 0,
        lastRebalanceTime: 0,
        singleRuleAssets: new Array(),
        ruleWeights: Object.assign([], array),
        ruleAssets: Object.assign([], token),
        numRuleTokens: 1,
    };
}

export function calculateStats(
    data: DataPoint[],
    days: number,
    index: number,
) {
    let result = [0, 0, 0];
    let startingDay = index - days + 1;
    let firstPrice = 0, lastPrice = 0;
    if(startingDay < 0) startingDay = 0;
    for(let i=startingDay; i<=index; i++){
        if(data[i].price == 0)continue;
        if(firstPrice == 0){
            firstPrice = data[i].price;
        }
        lastPrice = data[i].price;
        result[1] += data[i].volume;
        result[2] += data[i].circulatingSupply * data[i].price / 1000000;
    }

    result[0] = lastPrice * WEIGHT_MULTIPLIER / firstPrice;

    return result;
}

export function updateTokenStats(
    data: DataPoint[][],
    index: number,
) {
    let tokenStats: TokenStats = {
        stats: [],
    };
    let daysArray = [1, 7, 30, 91, 182, 365];
    for(let tokenId=0; tokenId<NUM_OF_TOKENS_IN_ASSET_POOL; tokenId++){
        let statsForToken = [];
        for(let i=0; i<6; i++) {
            let result = calculateStats(
                data[tokenId],
                daysArray[i],
                index,
            );
            statsForToken.push({
                days: daysArray[i],
                performance: result[0],
                volume: result[1],
                mcap: result[2],
            });
        }
        tokenStats.stats.push(statsForToken);
    }
    return tokenStats;
}

export function selectTokens(
    createBasketParams: CreateBasketParams,
    rule: Rule,
    tokenStats: TokenStats,
) {
    let ruleTokens = new Array(20);
    if(rule.filterBy == 0) {
        ruleTokens[0] = rule.fixedAsset;
        return ruleTokens;
    }
    let allTokens: {
        value: number,
        id: number,
    }[] = [];
    for(let i=0; i<NUM_OF_TOKENS_IN_ASSET_POOL; i++){
        if(createBasketParams.assetPool.indexOf(i) == -1)continue;
        if(rule.excludeAssets.indexOf(i) != -1)continue;
        if(tokenStats.stats[i][0].volume == 0)continue;
        let tokenStatsInfo = tokenStats.stats[i][rule.filterDays];
        let x;
        switch(rule.filterBy){
            case 1:
                x = tokenStatsInfo.mcap;
                break;
            case 2:
                x = tokenStatsInfo.volume;
                break;
            default:
                x = tokenStatsInfo.performance;
        }
        allTokens.push({
            value: x,
            id: i,
        });
    }
    allTokens.sort((a, b) => 0 - (a.value > b.value ? -1 : 1));

    for(let i=0; i<rule.numAssets; i++){
        switch(rule.sortBy) {
            case SortBy.AscendingOrder:
                ruleTokens[i] = allTokens[i].id;
                break;
            default:
                ruleTokens[i] = allTokens[allTokens.length - i - 1].id;
        }
    }
    return ruleTokens;
}

export function refilter(
    createBasketParams: CreateBasketParams,
    tokenStats: TokenStats,
    basketState: BasketState,
) {
    basketState.singleRuleAssets = [];
    for(let i=0; i<createBasketParams.rules.length; i++){
        let rule = createBasketParams.rules[i];
        let ruleTokens = selectTokens(createBasketParams, rule, tokenStats);
        basketState.singleRuleAssets.push(ruleTokens);
    }
}

export function calculateWeight(
    rawWeight: number,
    expo: number,
    weightBy: number,
): number {
    let multiplier = 1;
    if(weightBy == 3)multiplier = WEIGHT_MULTIPLIER;
    return Math.floor(Math.pow(rawWeight/multiplier , expo / EXPO_DIVIDER) * multiplier);
}

export function generateWeights(
    tokenStats: TokenStats,
    rule: Rule,
    ruleAssets: any,
) {
    let ruleWeights = [];
    for(let i=0; i<rule.numAssets; i++){
        let token = ruleAssets[i];
        let tokenStatsInfo = tokenStats.stats[token][rule.weightDays];
        let rawWeight = 1000;
        if(rule.filterBy != 0){
            switch(rule.weightBy) {
                case 1: 
                    rawWeight = tokenStatsInfo.mcap;break;
                case 2:
                    rawWeight = tokenStatsInfo.volume;break;
                default:
                    rawWeight = tokenStatsInfo.performance;
            }
        }
        let assetWeight = calculateWeight(
            rawWeight,
            rule.weightExpo,
            rule.weightBy,
        );
        ruleWeights.push({
            token,
            assetWeight
        });
    }
    return ruleWeights;
}

export function combineRules(
    basketState: BasketState,
    ruleAssets: any,
    createBasketParams: CreateBasketParams,
) {
    basketState.ruleAssets = new Array<number>(20).fill(0);
    basketState.ruleWeights = new Array<number>(20).fill(0);
    basketState.numRuleTokens = 1;
    for(let i=0; i<createBasketParams.rules.length; i++){
        let sumOfWeights = 0;
        for(let j=0; j<ruleAssets[i].length; j++)sumOfWeights += ruleAssets[i][j].assetWeight;
        for(let j=0; j<ruleAssets[i].length; j++){
            let token = ruleAssets[i][j].token;
            let optionIndex = basketState.ruleAssets.indexOf(token);
            let index = 0;
            switch(optionIndex) {
                case -1:
                    basketState.ruleAssets[basketState.numRuleTokens] = token;
                    basketState.numRuleTokens += 1;
                    index = basketState.numRuleTokens -1;
                    break;
                default:
                    index = optionIndex;
            }

            let x = BigInt(WEIGHT_MULTIPLIER) * BigInt(ruleAssets[i][j].assetWeight)
                    * BigInt(createBasketParams.rules[i].totalWeight) / BigInt(sumOfWeights);
            let weight = Math.floor(Number(x));
            basketState.ruleWeights[index] += weight;
        }
    }
}

export function addTargetAssets(
    basketState: BasketState
) {
    let i = 0;
    while(i < basketState.numRuleTokens){
        let optionIndex = basketState.currentCompToken.indexOf(basketState.ruleAssets[i]);
        if(optionIndex != -1){
            basketState.targetWeight[optionIndex] = basketState.ruleWeights[i];
        }
        else {
            if (basketState.numOfTokens < COMBINED_TOKENS_IN_A_BASKET) {
                basketState.currentCompAmount[basketState.numOfTokens] = 0;
                basketState.currentCompToken[basketState.numOfTokens] = basketState.ruleAssets[i];
                basketState.targetWeight[basketState.numOfTokens] = basketState.ruleWeights[i];
                basketState.numOfTokens += 1;
            }
        }
        i++;
    }
}

export function removeTargetAssets(
    basketState: BasketState
) {
    let i = 1;
    while(i < basketState.numOfTokens){
        let optionIndex = basketState.ruleAssets.indexOf(basketState.currentCompToken[i]);
        if(optionIndex == -1){
            if(basketState.currentCompAmount[i] == 0){
                basketState.numOfTokens -= 1;
                let lastTokenIndex = basketState.numOfTokens;
                // we write lastToken data on this index
                basketState.currentCompAmount[i] = basketState.currentCompAmount[lastTokenIndex];
                basketState.currentCompToken[i] = basketState.currentCompToken[lastTokenIndex];
                basketState.targetWeight[i] = basketState.targetWeight[lastTokenIndex];
                // after we moved data of lastToken on index i, we can make everything 0 on lastTokenIndex
                basketState.targetWeight[lastTokenIndex] = 0;
                basketState.currentCompToken[lastTokenIndex] = 0;
                basketState.currentCompAmount[lastTokenIndex] = 0;
                // since now we have the lastToken data on this index, we shouldn't
                // increase i in the while loop. because lastToken wasn't processed yet
                i = i - 1;
            }
            else {
                basketState.targetWeight[i] = 0;
            }
        }
        i++;
    }
}

export function updateTargetAssets(
    basketState: BasketState
) {
    removeTargetAssets(basketState);
    addTargetAssets(basketState);
    basketState.weightSum = 0;
    basketState.ruleWeights.forEach((element) => {
        basketState.weightSum += element;
    });
}

export function reweight(
    createBasketParams: CreateBasketParams,
    tokenStats: TokenStats,
    basketState: BasketState,
) {
    basketState.ruleAssets = new Array(20);
    basketState.ruleWeights = new Array(20);
    basketState.weightSum = 0;
    let ruleAssets = [];
    for(let i=0; i<createBasketParams.rules.length; i++){
        ruleAssets.push(
            generateWeights(
                tokenStats,
                createBasketParams.rules[i],
                basketState.singleRuleAssets[i]
            )
        );
    }
    combineRules(
        basketState,
        ruleAssets,
        createBasketParams,
    );
    updateTargetAssets(basketState);
}

export function rebalance(
    createBasketParams: CreateBasketParams,
    basketState: BasketState,
    data: DataPoint[][],
    day: number,
    tokenList: any,
) {
    let sell = [];
    let buy = [];
    for(let i=0; i<basketState.numOfTokens; i++){
        let currentPercentage = basketState.currentCompWeight[i];
        let targetPercentage = Math.floor(
            basketState.targetWeight[i] * WEIGHT_MULTIPLIER / basketState.weightSum
        );
        // checking rebalance threshold
        if(currentPercentage <= Math.floor(
                targetPercentage * 
                (BPS_DIVIDER + createBasketParams.rebalanceThreshold) /
                BPS_DIVIDER
            ) &&
            currentPercentage >= Math.floor(
                targetPercentage * 
                (BPS_DIVIDER - createBasketParams.rebalanceThreshold) /
                BPS_DIVIDER
            )
            )continue;
        
        // swapping
        if(currentPercentage > targetPercentage){
            if(i == 0)continue;
            sell.push(i);
        }
        else {
            buy.push(i);
        }
    }
    for(let ii=0; ii<sell.length; ii++){
        let i = sell[ii];
        let currentPercentage = basketState.currentCompWeight[i];
        let targetPercentage = Math.floor(
            basketState.targetWeight[i] * (WEIGHT_MULTIPLIER / basketState.weightSum)
        );
        let amountToSell = Math.min(Math.floor(
            basketState.currentCompAmount[i] * 
            (currentPercentage - targetPercentage) / 
            currentPercentage
        ), basketState.currentCompAmount[i]);
        let toUsdcAmount = Math.floor(
            amountToSell / 1000 * data[basketState.currentCompToken[i]][day].price * 950 / Math.pow(10, tokenList.list[basketState.currentCompToken[i]].decimals)
        );
        basketState.currentCompAmount[i] -= amountToSell;
        basketState.currentCompAmount[0] += toUsdcAmount;
    }

    for(let ii=0; ii<buy.length; ii++){
        let i = buy[ii];
        let currentPercentage = basketState.currentCompWeight[i];
        let targetPercentage = Math.floor(
            basketState.targetWeight[i] * (WEIGHT_MULTIPLIER / basketState.weightSum)
        );
        let amountToSpend = Math.min(Math.floor(
            basketState.basketWorth *
            ((targetPercentage - currentPercentage) /
            WEIGHT_MULTIPLIER)
        ), basketState.currentCompAmount[0]);
        let toTokenAmount = Math.floor(
            amountToSpend * 950 * Math.pow(10, tokenList[basketState.currentCompToken[i]].decimals) / data[basketState.currentCompToken[i]][day].price / 1000
        );
        basketState.currentCompAmount[i] += toTokenAmount;
        basketState.currentCompAmount[0] -= amountToSpend;
    }
}

export function updateBasketState(
    createBasketParams: CreateBasketParams,
    basketState: BasketState,
    tokenStats: TokenStats,
    database: DataPoint[][],
    day: number,
    tokenList: any,
) {
    if(basketState.lastRefilterTime + createBasketParams.refilterInterval <= day * 24 * 60 * 60){
        refilter(createBasketParams, tokenStats, basketState);
        basketState.lastRefilterTime = day * 24 * 60 * 60;
    }
    if(basketState.lastReweightTime + createBasketParams.reweightInterval <= day * 24 * 60 * 60){
        reweight(createBasketParams, tokenStats, basketState);
        basketState.lastReweightTime = day * 24 * 60 * 60;
    }
    if(basketState.lastRebalanceTime + createBasketParams.rebalanceInterval <= day * 24 * 60 * 60){
        updateCurrentWeights(database, day, basketState, tokenList);
        rebalance(createBasketParams, basketState, database, day, tokenList);
        basketState.lastRebalanceTime = day * 24 * 60 * 60;
    }
}

export function updateCurrentWeights(
    data: DataPoint[][],
    day: number,
    basketState: BasketState,
    tokenList: any,
) {
    basketState.basketWorth = 0;
    for(let i=0; i<COMBINED_TOKENS_IN_A_BASKET; i++){
        basketState.basketWorth += Math.floor(basketState.currentCompAmount[i] * data[basketState.currentCompToken[i]][day].price
                 / Math.pow(10, tokenList[basketState.currentCompToken[i]].decimals));
    }
    for(let i=0; i<COMBINED_TOKENS_IN_A_BASKET; i++){
        let x = BigInt(WEIGHT_MULTIPLIER) * BigInt(basketState.currentCompAmount[i])
                    * BigInt(data[basketState.currentCompToken[i]][day].price) / BigInt(basketState.basketWorth)
                    / BigInt(Math.pow(10, tokenList[basketState.currentCompToken[i]].decimals));
        basketState.currentCompWeight[i] = Math.floor(Number(x));
    }
}

export async function simulate(
    program: Program<BasketsIDL>,
    database: PublicKey,
    tokenList: PublicKey,
    createBasketParams: CreateBasketParams,
    simulationDays: number,
) {
    if(simulationDays > 365) simulationDays = 365;
    let data = await fetchDatabase(program, database);
    let tokenListData = await program.account.tokenList.fetch(tokenList, "confirmed");
    let basketState: BasketState = await defaultBasketState();
    let simulationData: {
        price: number,
        currentCompAmount: number[],
        currentCompToken: number[],
    }[] = [];
    for(let i = NUM_OF_DAYS_IN_DATABASE - simulationDays; i < NUM_OF_DAYS_IN_DATABASE; i++){
        let tokenStats = await updateTokenStats(data, i);
        await updateCurrentWeights(data, i, basketState, tokenListData);
        await updateBasketState(createBasketParams, basketState, tokenStats, data, i, tokenListData);
        simulationData.push({
            price: basketState.basketWorth,
            currentCompAmount: Object.assign([], basketState.currentCompAmount),
            currentCompToken: Object.assign([], basketState.currentCompToken),
        });
    }

    return simulationData;
}