import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.ticker as tkr

def stuloan(gross_salary, tax_year, plan):
    if tax_year == '20/21':
        plan_1_threshold = 19380
        plan_2_threshold = 26568
    elif tax_year == '19/20':
        plan_1_threshold = 18935
        plan_2_threshold = 25725
    if plan == 1:
        if gross_salary > plan_1_threshold:
            return (gross_salary - plan_1_threshold) * 0.09
        else:
            return 0
    elif plan == 2:
        if gross_salary > plan_2_threshold:
            return (gross_salary - plan_2_threshold) * 0.09
        else:
            return 0
    else:
        return 0
 
    
def inctax(gross_salary, tax_year):
    if tax_year == '20/21':
        base_personal_allowance = 12500
    elif tax_year == '19/20':
        base_personal_allowance = 12500
    if gross_salary < 100000:
        personal_allowance = base_personal_allowance
    elif gross_salary <= 125000:
        personal_allowance = base_personal_allowance - ((gross_salary - 100000) / 2)
    else:
        personal_allowance = 0
    higher_rate_threshold = personal_allowance + 37500
    if gross_salary <= personal_allowance:
        return 0
    else:
        basic_rate_paid = min((gross_salary - personal_allowance) * 0.2, (higher_rate_threshold - personal_allowance) * 0.2)
        if gross_salary <= higher_rate_threshold:
            higher_rate_paid = 0
        else:
            higher_rate_paid = min((gross_salary - higher_rate_threshold) * 0.4, (150000 - higher_rate_threshold) * 0.4)
        if gross_salary <= 150000:
            additional_rate_paid = 0
        else:
            additional_rate_paid = (gross_salary - 150000) * 0.45
        return basic_rate_paid + higher_rate_paid + additional_rate_paid

def natins(gross_salary, tax_year):
    if tax_year == '20/21':
        primary_threshold = 9500
        upper_earnings_limit = 50000
    elif tax_year == '19/20':
        primary_threshold = 8632
        upper_earnings_limit = 50000
    if gross_salary <= primary_threshold:
        return 0
    else:
        primary_rate_paid = min((gross_salary - primary_threshold) * 0.12, (upper_earnings_limit - primary_threshold) * 0.12)
        if gross_salary <= upper_earnings_limit:
            upper_rate_paid = 0
        else:
            upper_rate_paid = (gross_salary - upper_earnings_limit) * 0.02
        return primary_rate_paid + upper_rate_paid
  
taxyear = '20/21'
      
noloan_x = []
noloan_y = [] 
plan1_x = []
plan1_y = []
plan2_x = []
plan2_y = []
salary = 0

while salary <= 100000:
    noloan_x.append(salary)
    noloan_y.append(salary - inctax(salary, taxyear) - natins(salary, taxyear))
    plan1_x.append(salary)
    plan1_y.append(salary - inctax(salary, taxyear) - natins(salary, taxyear) - stuloan(salary, taxyear, 1))
    plan2_x.append(salary)
    plan2_y.append(salary - inctax(salary, taxyear) - natins(salary, taxyear) - stuloan(salary, taxyear, 2))
    salary += 1000

fig, ax = plt.subplots()
plt.plot(plan1_x, plan1_y, color='blue', linewidth=1.2)
plt.plot(plan2_x, plan2_y, color='orange', linewidth=1.2)
plt.plot(noloan_x, noloan_y, color='grey', linewidth=1.2)
plt.xticks(fontsize=7)
plt.yticks(fontsize=7)
plt.xlabel("Gross annual pay", fontsize=10)
plt.ylabel("Take home annual pay", fontsize=10)
plt.title("Gross vs take home pay by student loan plan (20/21 tax year)", fontsize=10)
noloan_label = mlines.Line2D([], [], color='grey', label='No student loan', linewidth=1.2)
plan1_label = mlines.Line2D([], [], color='blue', label='Plan 1 loan', linewidth=1.2)
plan2_label = mlines.Line2D([], [], color='orange', label='Plan 2 loan', linewidth=1.2)
plt.legend(handles=[noloan_label, plan1_label, plan2_label], fontsize=7)
plt.grid(linewidth=0.3)
ax.yaxis.set_major_formatter(
        tkr.FuncFormatter(lambda y,  p: format(int(y), ',')))
ax.xaxis.set_major_formatter(
        tkr.FuncFormatter(lambda y,  p: format(int(y), ',')))
plt.savefig('stuloan.pdf')