# Copyright © 2014 Zack Weinberg
# Copying and distribution of this program, with or without modification,
# are permitted in any medium without royalty provided the copyright
# notice and this notice are preserved.  This program is offered as-is,
# without any warranty.

library(plyr)
library(reshape2)
library(ggplot2)
library(scales)

# from http://www.cookbook-r.com/Graphs/Multiple_graphs_on_one_page_%28ggplot2%29/
multiplot <- function(..., plotlist=NULL, file, cols=1, layout=NULL) {
  require(grid)

  # Make a list from the ... arguments and plotlist
  plots <- c(list(...), plotlist)

  numPlots = length(plots)

  # If layout is NULL, then use 'cols' to determine layout
  if (is.null(layout)) {
    # Make the panel
    # ncol: Number of columns of plots
    # nrow: Number of rows needed, calculated from # of cols
    layout <- matrix(seq(1, cols * ceiling(numPlots/cols)),
                    ncol = cols, nrow = ceiling(numPlots/cols))
  }

 if (numPlots==1) {
    print(plots[[1]])

  } else {
    # Set up the page
    grid.newpage()
    pushViewport(viewport(layout = grid.layout(nrow(layout), ncol(layout))))

    # Make each plot, in the correct location
    for (i in 1:numPlots) {
      # Get the i,j matrix positions of the regions that contain this subplot
      matchidx <- as.data.frame(which(layout == i, arr.ind = TRUE))

      print(plots[[i]], vp = viewport(layout.pos.row = matchidx$row,
                                      layout.pos.col = matchidx$col))
    }
  }
}

##
## Tax rate comparisons.
##

# USA 2013 tax rates for comparison, assuming a single taxpayer, under
# the age of 65, all income from wages, taking the standard deduction,
# one exemption, and no other adjustments or credits.
# Source: IRS Schedule X/2013, Form 1040/2013.
Musa <- stepfun(c(0, 8925, 36250, 87850, 183250, 398350, 400000)+10000,
                c(0, .1, .15, .25, .28, .33, .35, .396), right=TRUE)

M <- function(i, B, r, s) ifelse(i <= B, 0, (r*(B-i))/(B-s+r*(s-i)))
A <- function(I, B, r, s, P)
    ifelse(I <= B + P,
           I,
           B + P +
           ((r*B + (1-r)*s)/r) * (log(r*I+(1-r)*s)-log(r*(B+P)+(1-r)*s)))

# Income range of interest; we go by 1000s up to 100,000, then by 10,000s
# up to 1,000,000, then by 100,000s up to 10,000,000.
Income = unique(c(seq(0, 1e5, 1e3),
                  seq(1e5, 1e6, 1e4),
                  seq(1e6, 1e7, 1e5)))

# For simplicity, we fix B = 24,000 and s = 500,000.
B. = 24000
s. = 500000

tax.p <- expand.grid(Income = Income, r = c(0.2, 0.4, 0.6, 0.8))
tax.p$marginal <- M(tax.p$Income, B., tax.p$r, s.)
tax.p$takehome <- A(tax.p$Income, B., tax.p$r, s., 0)
tax.p$r <- factor(tax.p$r, labels=c(
                               "r = 20%, s = 500,000",
                               "r = 40%, \"   \"",
                               "r = 60%, \"   \"",
                               "r = 80%, \"   \""))

tax.usa <- data.frame(Income = Income,
                      r = factor(c("USA 2013")),
                      marginal = Musa(Income))
tax.usa$takehome <- cumsum(c(0,diff(Income))*(1 - tax.usa$marginal))

tax <- rbind(tax.p, tax.usa)

tax.marginal.plot <-
ggplot(tax, aes(x=Income, y=marginal, colour=r, group=r)) + geom_line() +
    scale_x_log10(name="Pre-tax income (US$)", expand=c(0,0),
                  limits=c(7500,1e7), labels=comma) +
    scale_y_continuous(name="Marginal tax rate", expand=c(0,0),
                       limits=c(0,1), labels=percent) +
    scale_colour_manual(name="Rate function",
                        values=c("#e69f00", "#56b4e9", "#009e73",
                                 "#cc79a7", "#999999")) +
    theme(legend.position=c(0.15,0.85), axis.text.x=element_text(hjust=0.9))

tax.takehome.plot <-
ggplot(tax, aes(x=Income, y=takehome, colour=r, group=r)) + geom_line() +
    scale_x_log10(name="Pre-tax income (US$)", expand=c(0,0),
                  limits=c(7500,1e7), labels=comma) +
    scale_y_log10(name="After-tax income", expand=c(0,0),
                  limits=c(7500,1e7), labels=comma) +
    scale_colour_manual(name="Rate function",
                        values=c("#e69f00", "#56b4e9", "#009e73",
                                 "#cc79a7", "#999999"),
                        guide=FALSE) +
    theme(axis.text.x=element_text(hjust=0.9),
          axis.text.y=element_text(hjust=0.9, angle=90))

png("taxrates.png", width=960, height=480)
multiplot(tax.marginal.plot, tax.takehome.plot, cols=2)
dev.off()

##
## Pension hypotheticals.
##

age <- seq(18,89)

# Alice lives on the BLS till age 22 (when she has completed her
# masters' degree) and then gets a white-collar job that initially
# pays $50,000 a year, with $3000 raises every other year to a maximum
# of $101,000.  She retires at age 65.

alice <- c(0,0,0,0, # age 22
    50000, 50000, 53000, 53000, # age 25
    56000, 56000, 59000, 59000, # age 29
    62000, 62000, 65000, 65000, # age 33
    68000, 68000, 71000, 71000, # age 37
    74000, 74000, 77000, 77000, # age 41
    80000, 80000, 83000, 83000, # age 45
    86000, 86000, 89000, 89000, # age 49
    92000, 92000, 95000, 95000, # age 53
    98000, 98000, 101000, 101000, # age 57
    101000, 101000, 101000, 101000, # age 61
    101000, 101000, 101000, 101000, # age 65
    0,0,0,0, 0,0,0,0, 0,0,0,0, # age 77
    0,0,0,0, 0,0,0,0, 0,0,0,0) # age 89

# Bob starts a rock'n'roll band at the tender age of 18.  It records
# four albums over the course of 12 years.  Each album earns Bob (and
# his bandmates) $500,000 in the first year of its release,
# dropping to 75% of that amount each year thereafter, and going out
# of print after eight years.  Bob has no further income after that.

album <- cumprod(c(5e5, rep(.75, 7)))

bob <- rep(0, length(age))
bob[1:length(album)] <- bob[1:length(album)] + album
bob[3:(2+length(album))] <- bob[3:(2+length(album))] + album
bob[6:(5+length(album))] <- bob[6:(5+length(album))] + album
bob[9:(8+length(album))] <- bob[9:(8+length(album))] + album

# Carol works for a peace advocacy organization that can only afford
# to pay her $6,000 a year.  But at age 25 she starts writing novels;
# she can only write one every two years, but they are solidly
# received and earn her $50,000 each (at first, but decaying toward
# zero later on) Unfortunately, that's all advance.  She retires
# from the peace advocacy organization at age 40 and continues to
# write novels for the rest of her life.

carol <- rep(0, length(age))
carol[18:40 - 18] <- 6000
carol[seq(25,89,by=2) - 18] <- carol[seq(25,89,by=2) - 18] +
    cumprod(c(50000, rep(.9,32)))

pension = melt(data.frame(age=age, alice=alice, bob=bob, carol=carol),
    id.vars=c('age'))
colnames(pension) <- c('age', 'person', 'income')


# Compute pension and take home pay for each hypothetical, with these
# parameters:

r.  <- 0.4
p0. <- 0.5
pr. <- 0.05
d.  <- 0.99

P <- function(T, p0, pr, s)
    max(0, ((p0 * pr * (s-1))/(p0 - pr)) * (
        log(pr*s - p0 + (p0-pr)*T) -
        log(pr*(s-1))))

# Unfortunately, this calculation is now history-dependent and I don't
# see an easy way to vectorize it.

takehome.history <- function (case) {
    n <- nrow(case)
    pension <- rep(0, n)
    takehome <- rep(0, n)
    prev.P <- 0
    cumu.T <- 0
    for (i in 1:nrow(case)) {
        full.I <- case$income[i] + prev.P + B.
        takehome[i] <- A(full.I, B., r., s., prev.P)
        cumu.T <- cumu.T * d. + (full.I - takehome[i])
        pension[i] <- prev.P <- P(cumu.T, p0., pr., s.)
    }
    case$pension <- pension
    case$takehome <- takehome
    case
}

pension <- ddply(pension, .(person), takehome.history)

pension.t <- melt(pension, id.vars=c('age','person'))
png("pensions.png", width=960, height=480)
ggplot(pension.t, aes(x=age, y=value, colour=person, group=person)) +
    facet_wrap(nrow=1, ~variable) +
    geom_step() +
    scale_y_log10(limits=c(500,1e6),
        breaks=c(1000,2000,5000,10000,20000,50000,
            100000, 200000, 500000, 1e6), labels=comma) +
    theme(axis.title.y=element_blank(),
          legend.position=c(0.95,0.9))
dev.off()