Using Pandas to explore the power of compounding

import pandas as pd
import plotly.graph_objects as go
import plotly.express as px

if __name__ == "__main__":
    def returns(age_start: int, age_stop: int, s_age: int, early_investor: bool) -> list:
        initial_investment_age_stop = s_age
        annual_contribution = 2000
        annual_return_rate = 0.08

        # Initialize variables
        total_amount = 0
        returns_list = []

        # Loop through each year, simulating the investment growth
        for age in range(age_start, age_stop + 1):
            if early_investor:
                if age <= initial_investment_age_stop:
                    # Add principal for ages 19 to 27
                    total_amount += annual_contribution
            else:
                if age >= initial_investment_age_stop:
                    # Add principal for ages 28 to 64
                    total_amount += annual_contribution
            # Calculate the growth for the year
            total_amount *= (1 + annual_return_rate)
            # Add the total amount for the year to the list
            returns_list.append(total_amount)

        # Adjust the list to show the total amount at the end of each year, rounded to 2 decimal places
        returns_list = [round(amount, 2) for amount in returns_list]

        return returns_list

    # Create a dataframe according to the following:
    # 1. Create a datafrom with an age column from age 19 to 65
    # 2. Create column #2 called 'jack_invests' with value 2000 from age 19 to 27, and then 0 from age 28 to 64
    # 3. Create column #4 called 'jack_returns' with initial 2000 compounded at 6% from age 19 to 64, and also adding
    #   2000 from age 19 to 27
    # 4. Create column #3 called 'jill_invests' with value 0 from age 19 to 27, and then 2000 from age 28 to 64
    # 5. Create column #5 called 'jill_returns' with initial 2000 compounded at 8% from age 19 to 64, and also adding
    #   2000 from age 28 to 64
    start_age = 25
    end_age = 65
    s_age = 35
    df = pd.DataFrame({'age': range(start_age, end_age + 1),
                       'jack_invests': [2000 if age < s_age else 0 for age in range(start_age, end_age + 1)],
                       'jack_returns': returns(start_age, end_age, s_age, True),
                       'jill_invests': [0 if age < s_age else 2000 for age in range(start_age, end_age + 1)],
                       'jill_returns': returns(start_age, end_age, s_age, False)
                       })

    # Plot the returns of Jack and Jill using plotly as a bar plot
    fig = px.bar(df, x='age', y=['jack_returns', 'jill_returns'], title='Returns of Jack and Jill', barmode='group')
    fig.update_layout(xaxis_title='Age', yaxis_title='Return')
    fig.show()

    # Plot total investment of Jack and Jill using plotly. Also plot the total returns on investment as a percentage
    # of the total investment at age 65
    df['jack_total_investment'] = df['jack_invests'].cumsum()
    df['jill_total_investment'] = df['jill_invests'].cumsum()

    jack_total_investment = df['jack_total_investment'].iloc[-1]
    jill_total_investment = df['jill_total_investment'].iloc[-1]
    jack_final_return = df['jack_returns'].iloc[-1]
    jill_final_return = df['jill_returns'].iloc[-1]

    jack_total_return = (jack_final_return - jack_total_investment) / jack_total_investment * 100
    jill_total_return = (jill_final_return - jill_total_investment) / jill_total_investment * 100

    fig = go.Figure()
    fig.add_trace(go.Bar(x=['Jack Total Contribution'], y=[jack_total_investment], name='Jack Total Contributions', marker_color='rgb(55, 83, 109)'))
    fig.add_trace(go.Bar(x=['Jack Total Return'], y=[jack_final_return], name='Jack Total Return', marker_color='rgb(26, 118, 255)'))
    fig.add_trace(go.Bar(x=['Jill Total Contribution'], y=[jill_total_investment], name='Jill Total Contributions', marker_color='rgb(255, 153, 51)'))
    fig.add_trace(go.Bar(x=['Jill Total Return'], y=[jill_final_return], name='Jill Total Return', marker_color='rgb(255, 99, 71)'))
    fig.show()