#!/usr/bin/env python

'''
This is a simple program that plots the results of compounding when there are
monthly deposits.  It takes into account starting points so that information
from Social Security can be plotted and compared.  The idea is to "what if"
various scenarios.  When is it best to start taking Soc Sec if I expect to
spend every penny of it?  When is it best to start if I invest some percentage
at a give return rate?

This does nothing more than plot the equation at Wikipedia's page on Compound
Interest in the section titled Monthly Deposits.

Do not make important financial decisions base on this program!  For that
matter, do not make -any- decisions based on this program.  This is not
a detailed simulation of your financial future, just a very rough cut from a
single equation of compounding.

Mike Markowski
May 2026
'''

import matplotlib.pyplot as plt
import numpy as np
import os, sys

def cli(argv):
    '''Get user's command line arguments.  See usage() for detail.
    '''

    prog = os.path.basename(argv[0])
    argv = argv[1:]
    inv = 0 # Invest nothing by default.
    ret = 0.000001 # Avoid div by 0 in compound().

    age = []
    pay = []
    i = 0
    while i < len(argv):
        arg = argv[i]
        if arg == '-i': # 0 to 100%, monthly amount invested.
            i += 1
            inv = float(argv[i])/100
        elif arg == '-r': # 0 to 100%, long term return on investment.
            i += 1
            ret = float(argv[i])/100
        elif arg[0] == '-' and arg[1:].isnumeric(): # Age/$ soc sec pair.
            age.append(int(arg[1:]))
            i += 1
            pay.append(float(argv[i]))
        else:
            usage(prog)
        i += 1
    if len(age) == 0:
        usage(prog)
    return age, pay, inv, ret

def compound(P, r, M, t):
    '''
    Compound principal while monthly deposits occur.  Taken from
    https://en.wikipedia.org/wiki/Compound_interest
    "Monthly despoits" section, closed form formula.

    Inputa:
      P (float): dollars, principal amount.
      r (float): rate of return, monthly.
      M (float): dollars, monthly deposit.
      t (float): time in months.
    Output:
      (float): dollars, new principal.
    '''
    return M*((1 + r)**t - 1)/r + P*(1 + r)**t

def intersect(y1, y2):
    '''Given y values of two curves sharing the same x axis, find their
    points of intersection.

    Input:
      y1 (numbers[]): y values of first curve.
      y2 (numbers[]): y values of second curve.
    Output:
      (int[]): list of 0 or more indices of points of intersection.
    '''

    # return np.where(np.diff(np.sign(y1 - y2)))[0] # Indices of intersection.
    dy = y1 - y2          # Vertical distances between curves.
    signs = np.sign(dy)   # Array of -1s, 0s and 1s.
    twos = np.diff(signs) # Array of -2s, many 0s and 2s.
    pts = np.where(twos)  # Indices of non-zeros are intersections.
    return pts[0]

def plot(title, t, curve, ages):
    '''Plot the results!

    Inputs:
      title (string): Plot title.
      t (float[]): array of ages.
      curve (float[][]): for each age in ages[], the compounding curve.
      ages (int[]): ages at which compounding was calculated.
    '''

    plt.rcParams['axes.grid'] = True
    f, ax = plt.subplots(1, 1, figsize=(7,7))

    ax.set_title(title)
    ax.set_xlabel('Age')
    ax.set_ylabel('Total Earned ($)')

    for age,earned_usd in zip(ages, curve): # Plot each compounding curve.
        lbl = 'Start age %d' % age
        ax.plot(t, earned_usd, lw=0.75, label=lbl)

    # Put a dot at each pt of intersection and draw a vertical line.
    for i in range(len(curve)):
        for j in range(i+1, len(curve)):
            k = intersect(curve[i], curve[j])
            if len(k) == 0:
                continue
            k = k[-1]
            if k != -1 and curve[i][k] > 0:
                ax.scatter(x=[t[k]], y=[curve[i][k]], s=10) # Dot.
                ax.vlines(x=t[k], ymin=0, ymax=curve[i][k],
                    linestyle='--', color='blue', lw=0.75)  # Dashed line.

    ax.legend()
    plt.tight_layout()
    plt.show()

def usage(prog):
    print('Usage: %s [-i invest%%] [-r rate%%] -age earn$ ...\n' % prog)
    print('Example: %s -i 70 -r 6 -62 1400 -67 2000 -70 2250' % prog)
    print('         compares investing 70% of Soc Sec with a return of 6%.')
    sys.exit(1)

def main(argv):

    ages, pays, inv, ret = cli(argv) # Get cmd line args.
    ageMax_yr = 100 # Max age of calculations.

    P = 0      # Principal, empty bank account!
    r = ret/12 # Annual rate to monthly rate.
    curve = [] # Each curve will be a value vs age compounding plot.

    # Create the set of compounding curves.
    for age,pay in zip(ages, pays):
        M = inv*pay                         # Fraction of Soc Sec invested.
        dly_mo = 12*(age - ages[0])         # Months offset from 1st plot.
        t = np.arange(12*(ageMax_yr - age)) # Months to perform compounding.
        earned_usd = compound(P, r, M, t)   # Compounding curve.
        earned_usd = np.concatenate((np.zeros(dly_mo), earned_usd)) # Offset.
        curve.append(earned_usd)

    # Make a plot title.
    title = 'Age vs Earning\n'
    if ret > 0.0001:
        title += '%.1f%% return on %d%% investment of\n' % (100*ret, 100*inv)
    title += 'monthly SSA income at %d' % ages[0]
    for a in ages[1:]:
        title += ',%d' % a
    title += ' of %d' % pays[0]
    for p in pays[1:]:
        title += ',%d' % p
    # Convert compounding months to human ages in years.
    t = ages[0] + np.arange(12*(ageMax_yr - ages[0]))/12
    # Share the hard work!
    plot(title, t, curve, ages)

if __name__ == '__main__':
    main(sys.argv)