import pandas as pd
import matplotlib.pyplot as plt

def simple_interest(balances, rate=0.05, terms=3):
    if terms <= 0:
        print("Number of terms must be >0!")
        return
    new_bal = []
    for balance in balances:
        b = balance
        for i in range(0,terms):
            b = b + (b * rate)
        new_bal.append(b)
    return new_bal

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, dividend=10, terms=3):
    new_bal = []
    for balance in balances:
        b = balance
        for i in range(0,terms):
            b = b + dividend
        new_bal.append(b)
    return new_bal

def get_balances_over_time(names, balances, func, terms=4):
    if func is None:
        print("Need a function!")
        return
    if len(names) != len(balances):
        print("Every person needs a balance.")
        return
    data = {
        'name' : names,
        'initial' : balances
    }
    df = pd.DataFrame(data)
    for t in range(1,terms):
        key = "term%i" % t 
        frame = func(balances, terms=t)
        if frame is not None:
            df[key] = frame
    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, simple_interest
    )
    print(df.to_html())

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

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

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

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

# illustrate_share_of_wealth()

def visualize_ubi(terms=25):

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

    df = get_balances_over_time (
        participants, balances, ubi,
        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)
    print(sow.keys())
    x = [i for i in range(1,terms)]
    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='--'
    )

    plt.title("Change in Wealth Distribution Under UBI")
    plt.legend()
    plt.xlabel("Terms")
    plt.ylabel("Share of Wealth")
    plt.show()
    return

visualize_ubi(terms=50)