import pandas as pd
from collections import defaultdict
from SQLiteInterface import SQLiteInterface

# 定义一个类，用于计算财务指标
class FinancialMetricsCalculator:
    def __init__(self, db_interface):
        # 初始化时传入一个数据库接口对象
        self.db_interface = db_interface
        self.data = {}  # 存储加载的数据

    def load_data(self):
        # 加载数据的方法
        try:
            # 从数据库中读取支付宝和微信支付的现金流量表
            transactions_alipay = self.db_interface.read("SELECT * FROM AlipayCashFlowStatement")
            transactions_weichatpay = self.db_interface.read("SELECT * FROM WeichatpayCashFlowStatement")

            # 合并两个表的交易记录
            transactions = transactions_alipay + transactions_weichatpay

            # 获取表的列信息
            columns_info = self.db_interface.read("PRAGMA table_info(AlipayCashFlowStatement)")
            column_names = [info[1] for info in columns_info]

            # 创建一个包含所有交易的DataFrame
            df_transactions = pd.DataFrame(transactions, columns=column_names)
            # 将日期列转换为datetime类型
            df_transactions['date'] = pd.to_datetime(df_transactions['date'].apply(lambda x: x.split(' ')[0]))

            # 从数据库中读取活动类别及其方向（收入或支出）
            categories = self.db_interface.read("SELECT activity_name, direction FROM ActivityCategories")
            category_directions = {name: direction for name, direction in categories}

            # 将数据存储到实例变量中
            self.data['df_transactions'] = df_transactions
            self.data['category_directions'] = category_directions
            return True
        except Exception as e:
            print(f"Error loading data: {e}")
            self.data = {}
            return False

    def calculate_monthly_inflows_outflows(self):
        # 计算每月的收入和支出
        monthly_inflows = defaultdict(float)
        monthly_outflows = defaultdict(float)

        df_transactions = self.data.get('df_transactions')
        category_directions = self.data.get('category_directions')

        if df_transactions is None or category_directions is None:
            print("Data not loaded properly.")
            return monthly_inflows, monthly_outflows

        # 遍历每一行交易记录，计算每个月的收入和支出
        for _, row in df_transactions.iterrows():
            date, amount, label = row['date'], row['amount'], row['label']
            month = date.strftime("%Y-%m")
            direction = category_directions.get(label, None)
            if direction == 'Inflow':
                monthly_inflows[month] += amount
            elif direction == 'Outflow':
                monthly_outflows[month] += amount

        return monthly_inflows, monthly_outflows

    def calculate_savings_rate(self, monthly_inflows, monthly_outflows):
        # 计算储蓄率
        savings_data = []
        for month in sorted(monthly_inflows.keys()):
            total_income = monthly_inflows[month]
            total_expense = monthly_outflows.get(month, 0)
            net_savings = total_income - total_expense
            savings_rate = net_savings / total_income if total_income > 0 else 0
            savings_data.append((month, total_income, total_expense, net_savings, savings_rate))

        return savings_data

    def calculate_years_to_fire(self, savings_rate=0.5, annual_return=0.05, current_savings=0, annual_income=100000):
        # 计算达到财务独立所需的年数
        fire_target_multiplier = 25  # 财务独立的目标是年支出的25倍
        r = annual_return  # 年回报率

        # 计算年储蓄和年支出
        annual_savings = savings_rate * annual_income
        annual_spending = (1 - savings_rate) * annual_income

        # 计算财务独立目标金额
        fire_target = fire_target_multiplier * annual_spending

        years = 0
        total_savings = current_savings

        # 计算达到目标所需的年数
        while total_savings < fire_target:
            total_savings += annual_savings
            total_savings *= (1 + r)
            years += 1

        return years, fire_target

# 主程序
db_interface = SQLiteInterface()  
calculator = FinancialMetricsCalculator(db_interface)  

if calculator.load_data():
    # 计算每月收入和支出
    monthly_inflows, monthly_outflows = calculator.calculate_monthly_inflows_outflows()
    # 计算储蓄率数据
    savings_data = calculator.calculate_savings_rate(monthly_inflows, monthly_outflows)

    # 打印每个月的储蓄率数据
    for month, total_income, total_expense, net_savings, savings_rate in savings_data:
        print(f"{month}: 总收入: {total_income:.2f}, 总支出: {total_expense:.2f}, 净储蓄: {net_savings:.2f}, 储蓄率: {savings_rate:.2%}")
else:
    print("Failed to load data.")

# 计算达到财务独立所需的年数
savings_rate_input = input("\nEnter your savings rate (e.g., 0.5 for 50%) [default 0.5]: ")
annual_return_input = input("Enter your annual investment return (e.g., 0.05 for 5%) [default 0.05]: ")
current_savings_input = input("Enter your current savings [default 0]: ")
annual_income_input = input("Enter your annual income [default 100000]: ")

savings_rate = float(savings_rate_input) if savings_rate_input else 0.5
annual_return = float(annual_return_input) if annual_return_input else 0.05
current_savings = float(current_savings_input) if current_savings_input else 0
annual_income = float(annual_income_input) if annual_income_input else 100000

# 打印结果
years, fire_target = calculator.calculate_years_to_fire(savings_rate, annual_return, current_savings, annual_income)
    
print(f"\n\nWith a savings rate of {savings_rate*100:.0f}% and an annual return of {annual_return*100:.0f}%,")
print(f"and current savings of ${current_savings:,.0f},")
print(f"your savings target is ${fire_target:,.0f}, and you will need approximately {years:.1f} years to achieve FIRE.\n\n")