modules/retirement_planning.py

import base64
import csv
from io import StringIO
import matplotlib.pyplot as plt
import plotly.graph_objects as go
import io
from modules.utils import load_css, format_value

import plotly.graph_objects as go
import base64

def create_retirement_graph(post_retirement_investments):
    ages = [investment[0] for investment in post_retirement_investments]
    income_required = [investment[2] for investment in post_retirement_investments]
    dividend_income = [investment[3] for investment in post_retirement_investments]
    
    # 절대값의 필요 자금 계산
    cash_requirements = [max(0, -investment[6]) for investment in post_retirement_investments]  # difference의 절대값

    fig = go.Figure()
    
    # Dividends 그래프
    fig.add_trace(go.Bar(x=ages, y=dividend_income, name='Dividends', marker=dict(color='#89c9e6'), showlegend=False))
    
    # Needs (Shortfall) 그래프
    fig.add_trace(go.Bar(x=ages, y=cash_requirements, name='Shortfall', marker=dict(color='#ffca28'), showlegend=False))
    
    # Income Required 선 그래프 추가
    fig.add_trace(go.Scatter(x=ages, y=income_required, mode='lines', name='Needs', 
                             line=dict(color='red', dash='dash'), showlegend=False))

    fig.update_layout(
        xaxis_title='Age',
        barmode='stack',  # 누적 막대 그래프 설정
        xaxis=dict(
            showline=True,
            showgrid=True,
            gridcolor='LightGray',
            gridwidth=1
        ),
        yaxis=dict(
            showline=True,
            showgrid=True,
            gridcolor='LightGray',
            gridwidth=1,
            tickformat='.2s'
        ),
        hovermode='x unified',        
        margin=dict(l=0, r=0, t=0, b=0),
        autosize=True
    )

    graph_html = fig.to_html(full_html=False, include_plotlyjs='cdn')
    return base64.b64encode(graph_html.encode('utf-8')).decode('utf-8')

def create_retirement_table(post_retirement_investments, monthly_income_required):

    result_html = ""
    additional_cash_total = 0
    income_required_total = 0
    dividend_income_total = 0
    for age, investment, annual_income_required, annual_dividend_income, monthly_dividend_income, income_required, difference in post_retirement_investments:
        additional_cash_required = f'{format_value(abs(difference))}' if difference < 0 else ''
        additional_cash_total += difference if difference < 0 else 0
        income_required_total += annual_income_required
        dividend_income_total += annual_dividend_income

        result_html += f"""
                <tr>
                    <td style='font-weight: bold'>{age}</td>
                    <td>{investment:,.0f}</td>
                    <td>{annual_income_required:,.0f}</td>
                    <td>{annual_dividend_income:,.0f}</td>
                    <td>{additional_cash_required}</td>
                </tr>
        """
    result_html = f"""
    <div class='table-container'>
        <table>
            <thead>         
                <tr>
                    <th>Age</th>
                    <th>Returns</th>
                    <th>Needs</th>                        
                    <th>Dividends</th>
                    <th>Shortfall</th>

                </tr>
            </thead>
            <tbody>
                {result_html}
            </tbody>
        </table>
    </div>
    """
    return result_html

def create_csv_file(post_retirement_investments, monthly_income_required):
    csv_output = StringIO()
    csv_writer = csv.writer(csv_output)
    csv_writer.writerow([
        'Age', 
        'Returns', 
        'Needs', 
        'Dividends', 
        'Shortfall'
    ])
    for age, investment, annual_income_required, annual_dividend_income, monthly_dividend_income, income_required, difference in post_retirement_investments:
        additional_cash_required = f'{abs(difference):,.0f}' if difference < 0 else ''
        csv_writer.writerow([
            age, 
            f'{investment:,.0f}', 
            f'{annual_income_required:,.0f}', 
            f'{annual_dividend_income:,.0f}', 
            additional_cash_required
        ])
    csv_data = csv_output.getvalue()
    csv_base64 = base64.b64encode(csv_data.encode('utf-8')).decode('utf-8')

    return csv_base64

def retirement_planning(
    current_age=None, 
    retirement_age=None, 
    life_expectancy=None, 
    monthly_income_required=None, 
    inflation_rate=None, 
    initial_investment=None, 
    monthly_contribution=None, 
    annual_increase_in_monthly_contribution=None, 
    reinvest_dividends=None, 
    pre_retirement_dividend_growth=None, 
    post_retirement_dividend_growth=None, 
    pre_retirement_dividend_yield=None, 
    post_retirement_dividend_yield=None
):
    # NoneType일 때 0으로 처리
    current_age = current_age if current_age is not None else 0
    retirement_age = retirement_age if retirement_age is not None else 0
    initial_investment = initial_investment if initial_investment is not None else 0
    monthly_contribution = monthly_contribution if monthly_contribution is not None else 0
    annual_increase_in_monthly_contribution = annual_increase_in_monthly_contribution if annual_increase_in_monthly_contribution is not None else 0
    pre_retirement_dividend_growth = pre_retirement_dividend_growth if pre_retirement_dividend_growth is not None else 0
    post_retirement_dividend_growth = post_retirement_dividend_growth if post_retirement_dividend_growth is not None else 0
    pre_retirement_dividend_yield = pre_retirement_dividend_yield if pre_retirement_dividend_yield is not None else 0
    post_retirement_dividend_yield = post_retirement_dividend_yield if post_retirement_dividend_yield is not None else 0
    life_expectancy = life_expectancy if life_expectancy is not None else 0
    monthly_income_required = monthly_income_required if monthly_income_required is not None else 0
    inflation_rate = inflation_rate if inflation_rate is not None else 0

    # 은퇴 전후의 년 수 계산
    if retirement_age > life_expectancy:
        return "<p style='color: red;'>Error: Retirement age cannot be greater than life expectancy.</p>"
    
    if retirement_age < current_age:
        return "<p style='color: red;'>Error: Retirement age cannot be less than current age.</p>"
    
    years_to_retirement = retirement_age - current_age
    post_retirement_years = life_expectancy - retirement_age

    # 현재 투자액으로 초기 투자 설정
    total_investment = initial_investment

    # 은퇴 전 월간 이자율 계산
    monthly_return_pre = (1 + pre_retirement_dividend_growth / 100) ** (1 / 12) - 1

    # 은퇴 시점의 투자 계산
    for year in range(years_to_retirement):
        for month in range(12):
            # 월간 투자액과 이자율을 적용하여 총 투자액 갱신
            total_investment = (total_investment + monthly_contribution) * (1 + monthly_return_pre)
            # 배당금을 재투자할 경우 배당금 추가
            if reinvest_dividends:
                total_investment += total_investment * (pre_retirement_dividend_yield / 100 / 12)
        # 연간 증가액을 월 투자액에 추가
        monthly_contribution += annual_increase_in_monthly_contribution

    # 은퇴 시작 시점의 총 투자액과 연간 배당 수익 저장
    investment_at_retirement = total_investment
    annual_dividend_at_retirement = investment_at_retirement * (pre_retirement_dividend_yield / 100)
    monthly_dividend_at_retirement = annual_dividend_at_retirement / 12

    # 은퇴 후 월간 이자율 계산
    monthly_return_post = (1 + post_retirement_dividend_growth / 100) ** (1 / 12) - 1

    # 연간 물가상승률을 반영한 월 생활비 계산
    monthly_income_required_inflated = monthly_income_required
    monthly_income_required_over_time = []
    for age in range(current_age, life_expectancy + 1):
        if age >= retirement_age:
            monthly_income_required_over_time.append((age, monthly_income_required_inflated))
        monthly_income_required_inflated *= (1 + inflation_rate / 100)

    annual_income_required_at_retirement = monthly_income_required_over_time[0][1] * 12
    monthly_income_required_at_retirement = monthly_income_required_over_time[0][1]

    # 은퇴 후 투자 목록 초기화
    post_retirement_investments = [(retirement_age, investment_at_retirement, annual_income_required_at_retirement, annual_dividend_at_retirement, monthly_dividend_at_retirement, monthly_income_required_at_retirement, annual_dividend_at_retirement - annual_income_required_at_retirement)]

    # 은퇴 후 각 년도의 투자 및 배당 수익 계산
    for year in range(1, post_retirement_years + 1):
        # 은퇴 후 수익률을 적용하여 총 투자액 갱신
        total_investment *= (1 + post_retirement_dividend_growth / 100)
        # 연간 배당 수익 계산
        annual_dividend_income = total_investment * (post_retirement_dividend_yield / 100)
        # 월간 배당 수익 계산
        monthly_dividend_income = annual_dividend_income / 12
        # 연도별 물가상승률 반영한 월 생활비 갱신
        inflated_income_required = monthly_income_required_over_time[year][1] if year < len(monthly_income_required_over_time) else monthly_income_required_over_time[-1][1]
        # 각 연도별 투자와 배당 수익 및 월 생활비를 리스트에 추가
        difference = annual_dividend_income - inflated_income_required * 12
        post_retirement_investments.append((retirement_age + year, total_investment, inflated_income_required * 12, annual_dividend_income, monthly_dividend_income, inflated_income_required, difference))

    # 마이너스 값의 difference 합계 계산
    negative_differences_sum = sum(diff for _, _, _, _, _, _, diff in post_retirement_investments if diff < 0)

    # CSV 파일 생성
    csv_base64 = create_csv_file(post_retirement_investments, monthly_income_required)

    # CSS 스타일 로딩
    css = load_css()

    # 그래프 생성
    svg_base64 = create_retirement_graph(post_retirement_investments)
    
    # 표 생성
    table_html = create_retirement_table(post_retirement_investments, monthly_income_required)

    # HTML 결과 생성
    result_html = css + f"""
    <div class="wrap-text">
        <div style="margin-bottom: 1.5rem;">
            <iframe src="data:text/html;base64,{svg_base64}" style="width: 100%; height: 400px; border: none;"></iframe>
            <div style="display: flex; align-items: center; justify-content: flex-end;">💡Total Shortfall&nbsp;<span style="font-weight: bold; color: #1678fb">{abs(negative_differences_sum):,.0f}</span></div>
        </div>
    </div>
    
    <div style="margin-bottom: 2rem;"></div>
    
    <div style="margin-bottom: 2rem;"></div>
    <div style="display: flex; align-items: center; justify-content: flex-end;">
    <a href="data:text/csv;base64,{csv_base64}" download="retirement_planning.csv" style="padding: 10px 20px; border: 1px solid; border-radius: 5px; background-color: #1678fb; color: white; text-decoration: none;">Download CSV</a></div>

    """ + table_html + """
    <p style="padding: 10px; border: 1px solid; border-radius: 5px; text-align: center; max-width: 400px; margin: 20px auto;">
        <strong>Note:</strong> No additional investments or reinvestment of dividends after retirement.
    </p>
    """
    
    return result_html