import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import random
import sys
import tabulate

# it = 0

def compounding_interest(balances, rate=0.05, terms=3):
    if terms <= 0:
        print("Number of terms must be >0!")
        return
    global it 
    balances_over_time = []
    current_balances = balances
    for i in range(0, terms):
        new_bal = []
        for balance in current_balances:
            b = balance + (balance * rate)
            new_bal.append(b)
            # it += 1
        assert(len(new_bal) == len(balances))
        balances_over_time.append(new_bal)
        current_balances = new_bal
    return balances_over_time

def calc_share_of_wealth(balances):
    total_money = sum(balances)
    shares = []
    for balance in balances:
        shares.append(balance/total_money)
    return shares

def ubi(balances, rate=10, terms=3):
    dividend = rate
    balances_over_time = []
    current_balances = balances
    for i in range(0, terms):
        new_bal = []
        for balance in current_balances:
            b = balance + dividend
            new_bal.append(b)
        assert(len(new_bal) == len(balances))
        balances_over_time.append(new_bal)
        current_balances = new_bal
    return balances_over_time

def mining(balances, halving_frequency=1, reward=10, terms=3):
    balances_over_time = []
    current_balances = balances
    for i in range(0, terms):
        new_bal = []
        for balance in current_balances:
            b = balance + reward
            new_bal.append(b)
        assert(len(new_bal) == len(balances))
        if i % halving_frequency == 0:
            reward = reward/2
        balances_over_time.append(new_bal)
        current_balances = new_bal
    return balances_over_time

def get_balances_over_time(names, balances, data):
    if len(names) != len(balances):
        print("Every person needs a balance.")
        return
    df = pd.DataFrame({
        'name' : names,
        'initial' : balances
    })
    for i in range(0, len(data)):
        key = str(i+1)
        df[key] = data[i]
    return df

# -------------

def illustrate_share_of_wealth():

    participants = ["Alice", "Bob", "Charlie"]
    balances = [100,40,20]

    print("In this demo, we have three participants.\n", participants)
    print("They respectively have ", balances)

    print("\nThe initial distribution of wealth is ")
    print(calc_share_of_wealth(balances))
    print("\n")

    df = get_balances_over_time (
        participants, balances, compounding_interest(balances)
    )
    # print(it)
    df["share"] = [val * 100 for val in calc_share_of_wealth(df["3"]) ]
    
    print(df.to_html())

    terms = 222
    print(f"How much money exists after {terms} terms?")
    #### In 222 terms, even C becomes a millionaire.
    nb = compounding_interest(balances, terms=terms)
    print(nb[-1])
    print("\n")

    df = get_balances_over_time (
        participants, balances, ubi(balances)
    )
    print(df.to_html())

    print("How much money exists after 999 terms?")
    u = ubi(balances, terms=999)
    print(u[-1])

    print("What is the share of wealth in a UBI economy?")
    print(calc_share_of_wealth(u[-1]))



def visualize_ubi(terms=25):

    participants = ["Alice", "Bob", "Charlie"]
    balances = [100,40,20]

    df = get_balances_over_time (
        participants, balances, ubi(balances, terms=terms)
    )
    # print(df.keys)
    shares = []
    for key, data in df.items():
        if key == "name":
            continue
        if key == "initial":
            continue
        shares.append(calc_share_of_wealth(data))
    sow = pd.DataFrame(shares)
    x = [i for i in range(0,terms)]
    try:
        assert(len(x) == len(shares))
    except AssertionError:
        print(len(x), len(shares))
        exit()

    plt.style.use('dark_background')
    plt.plot(
        x, sow[0], color="red", label=participants[0]
    )
    plt.plot(
        x, sow[1], color="lightgreen", label=participants[1]
    )
    plt.plot(
        x, sow[2], color="cyan", label=participants[2]
    )
    plt.axhline(
        y=0.33, color='violet', linestyle='--', label="0.33"
    )

    plt.title("Change in Wealth Distribution With a Constant Dividend")
    plt.legend()
    plt.xlabel("Terms")
    plt.ylabel("Share of Wealth")
    plt.savefig("ubi-wealth-distribution.png")
    plt.close()
    # plt.show()
    return

def calc_total_supply(df):
    total = []
    for key, data in df.items():
        if key == "name":
            continue
        total.append(sum(data))
    return total

def draw_apy_inflation(terms=50, linear_label="? (linear)"):
    participants = ["Alice", "Bob", "Charlie"]
    balances = [100,40,20]

    df = get_balances_over_time (
        participants, balances, ubi(balances, terms=terms)
    )
    total_supply_ubi = calc_total_supply(df)

    df_si = get_balances_over_time (
        participants, balances, 
        compounding_interest(balances, terms=terms)
    )
    total_supply_si = calc_total_supply(df_si)
    # + 1 because the initial frame is included this time
    x = [i for i in range(0,terms+1)]
    try:
        assert(len(x) == len(total_supply_si))
    except AssertionError:
        print(len(x), len(total_supply_si))
        print(df_si)
        exit()

    plt.style.use('dark_background')
    plt.plot(
        x, total_supply_ubi, color="cyan", label="dividend = 10"
    )
    plt.plot(
        x, total_supply_si, color="red", label="apy = 0.05"
    )
    plt.plot(
        x,
        calc_total_supply(get_balances_over_time(
            participants, 
            balances, 
            compounding_interest(balances, terms=terms, rate=0.04)
        )),
        color="orange", label="apy = 0.04"
    )
    plt.plot(
        x,
        calc_total_supply(get_balances_over_time(
            participants, 
            balances, 
            compounding_interest(balances, terms=terms, rate=0.03)
        )),
        color="yellow", label="apy = 0.03"
    )
    plt.title("Supply of Money Over Time")
    plt.legend()
    plt.xlabel("Terms")
    plt.ylabel("Total Currency")
    plt.savefig("inflation-ubi-vs-5apy.png")
    plt.close()
    # plt.show()

def draw_apy_inflation2(terms=50, linear_label="? (linear)"):
    participants = ["Alice", "Bob", "Charlie"]
    balances = [100,40,20]

    df = get_balances_over_time (
        participants, balances, ubi(balances, terms=terms)
    )
    total_supply_ubi = calc_total_supply(df)

    df_si = get_balances_over_time (
        participants, balances, 
        compounding_interest(balances, terms=terms)
    )
    total_supply_si = calc_total_supply(df_si)
    # + 1 because the initial frame is included this time
    x = [i for i in range(1913,1913+terms+1)]
    try:
        assert(len(x) == len(total_supply_si))
    except AssertionError:
        print("Assert Error:")
        print(len(x), len(total_supply_si))
        print(df_si)
        exit()

    plt.style.use('dark_background')
    plt.plot(
        x, total_supply_ubi, color="cyan", label="dividend = 10"
    )
    plt.plot(
        x, total_supply_si, color="red", label="apy = 0.05"
    )
    plt.axvline(x=1960, color='yellow', linestyle='--')
    plt.title("Supply of Money Over Time")
    plt.legend()
    plt.xlabel("Year")
    plt.ylabel("Total Currency")
    plt.savefig("inflation-ubi-vs-5apy2.png")
    plt.close()
    # plt.show()

def draw_apy_inflation3(terms=50, linear_label="? (linear)"):
    participants = ["Alice", "Bob", "Charlie"]
    balances = [100,40,20]

    df = get_balances_over_time (
        participants, balances, ubi(balances, terms=terms)
    )
    total_supply_ubi = calc_total_supply(df)

    df_si = get_balances_over_time (
        participants, balances, 
        compounding_interest(balances, terms=terms)
    )
    total_supply_si = calc_total_supply(df_si)
    # + 1 because the initial frame is included this time
    x = [i for i in range(1913,1913+terms+1)]
    try:
        assert(len(x) == len(total_supply_si))
    except AssertionError:
        print("Assert Error:")
        print(len(x), len(total_supply_si))
        print(df_si)
        exit()
    dy1 = np.gradient(total_supply_ubi, x)
    dy2 = np.gradient(total_supply_si, x)

    plt.style.use('dark_background')
    plt.plot(
        x, dy1, color="cyan", label="deriv. const.div"
    )
    plt.plot(
        x, dy2, color="red", label="deriv. apy"
    )
    plt.axvline(x=1960, color='yellow', linestyle='--')
    plt.axvline(x=1940, color='orange', linestyle='--')
    plt.title("Change in Supply of Money")
    plt.legend()
    plt.xlabel("Year")
    plt.ylabel("Rate of Inflation")
    plt.savefig("inflation-ubi-vs-5apy3.png")
    plt.close()
    # plt.show()


def draw_3ubi_3apy(terms=50):
    participants = ["Alice", "Bob", "Charlie"]
    balances = [100,40,20]

    df = get_balances_over_time (
        participants, balances, ubi(balances, terms=terms)
    )
    total_supply_ubi = calc_total_supply(df)

    df_si = get_balances_over_time (
        participants, balances, 
        compounding_interest(balances, terms=terms)
    )
    total_supply_si = calc_total_supply(df_si)
    # + 1 because the initial frame is included this time
    x = [i for i in range(0,terms+1)]
    assert(len(x) == len(total_supply_si))

    plt.style.use('dark_background')
    plt.plot(
        x, total_supply_ubi, color="cyan", label="dividend = 10"
    )
    plt.plot(
        x,
        calc_total_supply(get_balances_over_time(
            participants, 
            balances, 
            ubi(balances, rate=5, terms=terms)
        )),
        color="violet", label="dividend = 5"
    )
    plt.plot(
        x,
        calc_total_supply(get_balances_over_time(
            participants, 
            balances, 
            ubi(balances, rate=15, terms=terms)
        )),
        color="lightgreen", label="dividend = 15"
    )
    plt.plot(
        x, total_supply_si, color="red", label="apy = 0.05"
    )
    plt.plot(
        x,
        calc_total_supply(get_balances_over_time(
            participants, 
            balances, 
            compounding_interest(balances, rate=0.04, terms=terms)
        )),
        color="orange", label="apy = 0.04"
    )
    plt.plot(
        x,
        calc_total_supply(get_balances_over_time(
            participants, 
            balances, 
            compounding_interest(balances, rate=0.03, terms=terms),
        )),
        color="yellow", label="apy = 0.03"
    )
    plt.title("Supply of Money Over Time")
    plt.legend()
    plt.xlabel("Terms")
    plt.ylabel("Total Currency")
    plt.savefig("inflation-ubi-vs-apy2.png")
    plt.close()
    # plt.show()

def visualize_inflation(df, title="Inflation", filename="inflation.png"):
    supply_over_time = calc_total_supply(df)
    time_span = len(supply_over_time)
    x = [i for i in range(time_span)]
    assert(len(x) == time_span)

    plt.style.use('dark_background')
    plt.plot(
        x, supply_over_time, color="red")
    plt.title(title)
    plt.xlabel("Terms")
    plt.ylabel("Total Currency")
    plt.savefig(filename)
    plt.close()

def visualize_wealth_dist(df, title="Wealth Distribution", filename="wealth-distribution.png"):
    participants = df["name"]
    shares = []
    for key, data in df.items():
        if key == "name":
            continue
        if key == "initial":
            continue
        shares.append(calc_share_of_wealth(data))
    sow = pd.DataFrame(shares)
    # print(sow.keys())
    terms = sow.shape[0] + 1
    x = [i for i in range(1,terms)]
    print(len(x))
    print(len(shares))
    assert(len(x) == len(shares))

    plt.style.use('dark_background')
    plt.plot(
        x, sow[0], color="red", label=participants[0]
    )
    plt.plot(
        x, sow[1], color="lightgreen", label=participants[1]
    )
    plt.plot(
        x, sow[2], color="cyan", label=participants[2]
    )
    plt.axhline(
        y=0.33, color='violet', linestyle='--', label="0.33"
    )

    plt.title(title)
    plt.legend()
    plt.xlabel("Terms")
    plt.ylabel("Share of Wealth")
    plt.savefig(filename)
    plt.close()

def draw_simple_mining_demo():
    participants = ["Alice", "Bob", "Charlie"]
    balances = [100,40,20]
    m = get_balances_over_time(
        participants, balances, 
        mining(balances, halving_frequency=1, terms=10)
    )
    print(m)
    print(m.shape[1])
    visualize_inflation(m, title="Inflation Given block_reward = 10 and Halving Every Term", filename="inflation-pow.png")
    # print(m.iloc[:,:5].to_markdown())
    visualize_wealth_dist(m, filename="wealth-distribution-pow.png")

def trade(players, balances, spend_limit=0.10):
    assert(len(players) == len(balances))
    PRICE_FLOOR = 1
    pl_idx = len(players) - 1
    #  randomly pick two players that will trade
    buyer = 0
    seller = 0
    while buyer == seller:
    # reroll until the two arent the same
        buyer = random.choice([i for i in range(0, pl_idx)])
        seller = random.choice([i for i in range(0, pl_idx)])
    price = random.uniform(
        PRICE_FLOOR, spend_limit * balances[buyer]
    )
    balances[buyer] -= price
    balances[seller] += price
    # print(balances[buyer], balances[seller])
    return

class Sim:
    def __init__(self, players, balances, terms=50, starting_amount=0):
        self.players = players.copy()
        self.balances = balances.copy()
        self.terms = terms

        self.ADD_NEW_PPL_EVERY = 10
        self.NEW_PPL_START_WITH_BALANCE = starting_amount
        self.IS_TRADING = True

        self.IS_HALVING = False
        self.HALVING_EVERY = 5

        self.TITLE = str()
        self.events = []

    def add_player(self, 
                   balances, name="Player %i", starting_amount=0):
        name = name % (len(self.players) + 1)
        self.players.append(name)
        balances.append(starting_amount)
        try:
            assert(len(self.players) == len(balances))
        except AssertionError as e:
            print(e)
            print(len(self.players))
            print(len(balances))
            sys.exit()

    def run(self, update):
        ADD_NEW_PPL_EVERY = self.ADD_NEW_PPL_EVERY
        INIT_AMT = self.NEW_PPL_START_WITH_BALANCE
        current_balances = self.balances
        result = list()
        
        dividend = 10            

        i = 0
        while i < self.terms:
            # add new player?
            if i > 0 \
            and ADD_NEW_PPL_EVERY > 0 \
            and i % ADD_NEW_PPL_EVERY == 0:
                self.add_player(
                    current_balances, starting_amount=INIT_AMT)
            new_bal = []
            # make players trade if an
            if self.IS_TRADING:
                trade(self.players, current_balances)
            for balance in current_balances:
                if not self.IS_HALVING \
                or update == mini_apy:
                    b = update(balance)
                else:
                    b = mini_dividend(balance, dividend=dividend)
                new_bal.append(b)
            
            result.append(new_bal)
            current_balances = new_bal
            if i > 0 and i % self.HALVING_EVERY == 0:
                dividend = dividend/2
            i += 1
        return result

    def visualize(self, data, filename=str()):
        data.insert(0, self.balances)
        data2 = [calc_share_of_wealth(row) for row in data]
        df = pd.DataFrame(data2)

        # +1 to offset 0 index
        x = [i for i in range(1,df.shape[0] + 1)]
        assert(len(x) == df.shape[0])

        plt.style.use('dark_background')        
        # each column is the balance over time of an individual
        for column in df.columns:
            plt.plot(x, df[column], label=self.players[column] )
        plt.xlabel("Terms")
        plt.ylabel("Share of Wealth")
        plt.title(self.TITLE)
        # plt.legend()
        if filename != str():
            plt.savefig(filename)
        else:
            plt.show()
        plt.close()

def mini_apy(balance, rate=0.05):
    return balance + (balance * rate)

def mini_dividend(balance, dividend = 10):
    return balance + dividend

if __name__ == "__main__":

    # illustrate_share_of_wealth()
    # visualize_ubi(terms=50)
    # draw_apy_inflation(terms=75)
    # draw_apy_inflation2(terms=75, linear_label="dividend = 10")
    # draw_apy_inflation3(terms=75, linear_label="dividend = 10")
    # draw_3ubi_3apy(terms=75)
    # 
    # draw_simple_mining_demo()

    participants = ["Alice", "Bob", "Charlie"]
    balances = [100,40,20]
    s = Sim(participants, balances)
    result = s.run(update=mini_dividend)
    s.TITLE = "Constant dividend with new players joining and trading"
    s.visualize(result, filename="const-new-players-trade.png")

    t = Sim(participants, balances)
    result2 = t.run(update=mini_apy)
    t.TITLE = "APY with new people joining and trading"
    t.visualize(result2, filename="apy-new-playrs-trade.png")

    alt_starting_balance = [1,1,1]
    u = Sim(participants, balances)
    u.IS_HALVING = True
    result3 = u.run(update=mini_dividend)
    # print(pd.DataFrame(result3))
    u.TITLE = "Halving dividend when new users mine and trade"
    u.visualize(result3, filename="halving-dividend-new-users-trade.png")