
# loanit/core.py

import csv

class Loan:
    def __init__(self, principal, annual_rate, years, payments_per_year=1, loan_type="term"):
        """
        Initialize a loan object.
        :param principal: Total loan amount
        :param annual_rate: Annual interest rate (e.g., 0.1 for 10%)
        :param years: Loan duration in years
        :param payments_per_year: Number of payments per year (default = 1 for annual)
        :param loan_type: Type of loan ('term', 'equal_amortization', 'interest_only', 'balloon', 'bullet')
        """
        self.principal = principal
        self.annual_rate = annual_rate
        self.years = years
        self.payments_per_year = payments_per_year
        self.loan_type = loan_type.lower()

        self.periods = self.years * self.payments_per_year
        self.rate_per_period = self.annual_rate / self.payments_per_year

    def amortization_schedule(self):
        """To be implemented for each loan type."""
        if self.loan_type == "term":
            return self._term_loan_schedule()
        elif self.loan_type == "equal_amortization":
            return self._equal_amortization_schedule()
        elif self.loan_type == "interest_only":
            return self._interest_only_schedule()
        elif self.loan_type == "balloon":
            return self._balloon_schedule()
        elif self.loan_type == "bullet":
            return self._bullet_loan_schedule()

        else:
            raise ValueError(f"Loan type '{self.loan_type}' not supported yet.")


    def _term_loan_schedule(self):
        """Calculates the amortization schedule for a term (equal payment) loan."""
        r = self.rate_per_period
        n = self.periods
        P = self.principal

        if r == 0:
            payment = P / n
        else:
            payment = P * (r * (1 + r)**n) / ((1 + r)**n - 1)

        schedule = []
        balance = P
        total_balance = payment * n  # total payments remaining

        for period in range(1, n + 1):
            interest = balance * r
            principal_paid = payment - interest
            balance -= principal_paid
            total_balance -= payment
            if balance < 0:
                balance = 0
            if total_balance < 0:
                total_balance = 0

            schedule.append({
                "Year": period,
                "Payment": round(payment, 2),
                "Principal Paid": round(principal_paid, 2),
                "Interest Paid": round(interest, 2),
                "Principal Balance": round(balance, 2),
                "Total Balance": round(total_balance, 2)
            })

        return schedule

    def _equal_amortization_schedule(self):
        """Calculates the schedule for Equal Amortization Term Loan.
        Fixed principal repayment each period, decreasing interest.
        """
        P = self.principal
        r = self.rate_per_period
        n = self.periods

        principal_payment = P / n
        balance = P
        schedule = []

        # Pre-calculate all payments for total balance tracking
        total_payments = []

        for period in range(1, n + 1):
            interest = balance * r
            payment = principal_payment + interest
            total_payments.append(payment)
            balance -= principal_payment
            if balance < 0:
                balance = 0

        # Now generate schedule with running total balance
        balance = P
        remaining_payments = sum(total_payments)

        for period in range(1, n + 1):
            interest = balance * r
            payment = principal_payment + interest
            balance -= principal_payment
            remaining_payments -= payment
            if balance < 0:
                balance = 0
            if remaining_payments < 0:
                remaining_payments = 0

            schedule.append({
                "Year": period,
                "Payment": round(payment, 2),
                "Principal Paid": round(principal_payment, 2),
                "Interest Paid": round(interest, 2),
                "Principal Balance": round(balance, 2),
                "Total Balance": round(remaining_payments, 2)
            })

        return schedule

    def _interest_only_schedule(self):
        """Calculates the schedule for an Interest-Only Loan.
        Interest paid every year; principal repaid at the end.
        """
        P = self.principal
        r = self.rate_per_period
        n = self.periods
        schedule = []

        interest_payment = P * r  # same each year

        for period in range(1, n + 1):
            if period < n:
                payment = interest_payment
                principal_paid = 0
                principal_balance = P
            else:
                # Last year — pay interest + full principal
                payment = interest_payment + P
                principal_paid = P
                principal_balance = 0

            total_balance = principal_balance + (interest_payment if period < n else 0)

            schedule.append({
                "Year": period,
                "Payment": round(payment, 2),
                "Principal Paid": round(principal_paid, 2),
                "Interest Paid": round(interest_payment, 2),
                "Principal Balance": round(principal_balance, 2),
                "Total Balance": round(total_balance, 2)
            })

        return schedule

    def _balloon_schedule(self):
        """Calculates the amortization schedule for a Balloon Loan.
        Regular payments with a large balloon (remaining principal) payment at the end.
        """
        P = self.principal
        r = self.rate_per_period
        n = self.periods
        balloon_pct = 0.3  # You can adjust this default (e.g., 30% balloon)
        balloon_amount = P * balloon_pct

        # Remaining principal (non-balloon portion) amortized over n-1 periods
        amortized_principal = P - balloon_amount

        if r == 0:
            payment = amortized_principal / n
        else:
            payment = amortized_principal * (r * (1 + r)**n) / ((1 + r)**n - 1)

        balance = P
        schedule = []
        total_balance = payment * n + balloon_amount  # total future payments

        for period in range(1, n + 1):
            interest = balance * r

            if period < n:
                principal_paid = payment - interest
                payment_total = payment
            else:
                # Final period — add balloon amount
                principal_paid = balance
                payment_total = interest + balance

            balance -= principal_paid
            total_balance -= payment_total
            if balance < 0:
                balance = 0
            if total_balance < 0:
                total_balance = 0

            schedule.append({
                "Year": period,
                "Payment": round(payment_total, 2),
                "Principal Paid": round(principal_paid, 2),
                "Interest Paid": round(interest, 2),
                "Principal Balance": round(balance, 2),
                "Total Balance": round(total_balance, 2)
            })

        return schedule

    def _bullet_loan_schedule(self):
        """Calculates the schedule for a Bullet Loan.
        No payments in periods 1..n-1 (interest is capitalized each period).
        Entire principal + accrued interest is paid in the last period.
        """
        P = self.principal
        r = self.rate_per_period
        n = self.periods
        schedule = []

        # Start with initial principal
        balance = P

        for period in range(1, n + 1):
            # interest for this period (on current balance)
            interest = balance * r

            if period < n:
                # No payment this period; interest is capitalized (added to balance)
                payment = 0.0
                principal_paid = 0.0
                # Capitalize interest onto balance
                balance = round(balance + interest, 10)  # keep higher precision here
                # Remaining payments: only one payment at the end, which will be balance*(1+r)^(remaining_periods)
                remaining_periods = n - period
                # final payment still to come (after compounding remaining_periods)
                total_balance = round(balance * ((1 + r) ** remaining_periods), 2)
            else:
                # Final period: pay the accumulated balance plus this period's interest
                # Note: balance currently equals principal + interest capitalized up to period n-1
                final_interest = interest  # interest on the accumulated balance for the final period
                payment = round(balance + final_interest, 2)
                principal_paid = round(balance, 2)  # principal portion repaid now
                # After final payment, no balance remains
                balance = 0.0
                total_balance = 0.0

            schedule.append({
                "Year": period,
                "Payment": round(payment, 2),
                "Principal Paid": round(principal_paid, 2),
                "Interest Paid": round(round(interest, 2), 2),
                "Principal Balance": round(balance, 2),
                "Total Balance": round(total_balance, 2)
            })

        return schedule


    def total_interest(self):
        """Returns the total interest paid over the full loan period."""
        schedule = self.amortization_schedule()
        return round(sum(row["Interest Paid"] for row in schedule), 2)

    def total_payment(self):
        """Returns the total amount paid (principal + interest)."""
        schedule = self.amortization_schedule()
        return round(sum(row["Payment"] for row in schedule), 2)

    def breakdown_of_year(self, year):
        """Returns the breakdown (payment, interest, principal, balance) for a specific year."""
        schedule = self.amortization_schedule()
        for row in schedule:
            if row["Year"] == year:
                return row
        raise ValueError("Invalid year number.")

    def outstanding_principal_after(self, year):
        """Returns the remaining principal after a given year."""
        schedule = self.amortization_schedule()
        if year > len(schedule):
            return 0.0
        # After completing `year`, we look at the same year's record
        return schedule[year - 1]["Principal Balance"]

    def balance_after_year(self, year):
        """Returns the total outstanding balance (principal + remaining interest) after a given year."""
        schedule = self.amortization_schedule()
        if year > len(schedule):
            return 0.0
        return schedule[year - 1]["Total Balance"]

    def outstanding_interest_after(self, year):
        """Total interest yet to be paid after a given year."""
        schedule = self.amortization_schedule()
        total_interest = self.total_interest()
        interest_paid = sum(row["Interest Paid"] for row in schedule if row["Year"] <= year)
        return round(total_interest - interest_paid, 2)

    def export_to_csv(self, filename="amortization_schedule.csv"):
        """Exports the amortization schedule to a CSV file."""
        schedule = self.amortization_schedule()
        if not schedule:
            raise ValueError("No schedule found. Make sure to define a valid loan type.")

        fieldnames = schedule[0].keys()

        with open(filename, mode="w", newline="", encoding="utf-8") as file:
            writer = csv.DictWriter(file, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(schedule)

        print(f"Amortization schedule exported successfully to '{filename}'.")


