InvestmentTrackerApp/InvestmentTracker/Services/CalculationService.swift

import Foundation

class CalculationService {
    static let shared = CalculationService()

    private init() {}

    // MARK: - Portfolio Summary

    func calculatePortfolioSummary(
        from sources: [InvestmentSource],
        snapshots: [Snapshot]
    ) -> PortfolioSummary {
        let totalValue = sources.reduce(Decimal.zero) { $0 + $1.latestValue }
        let totalContributions = sources.reduce(Decimal.zero) { $0 + $1.totalContributions }

        // Calculate period changes
        let now = Date()
        let dayAgo = Calendar.current.date(byAdding: .day, value: -1, to: now) ?? now
        let weekAgo = Calendar.current.date(byAdding: .weekOfYear, value: -1, to: now) ?? now
        let monthAgo = Calendar.current.date(byAdding: .month, value: -1, to: now) ?? now
        let yearAgo = Calendar.current.date(byAdding: .year, value: -1, to: now) ?? now

        let dayChange = calculatePeriodChange(sources: sources, from: dayAgo)
        let weekChange = calculatePeriodChange(sources: sources, from: weekAgo)
        let monthChange = calculatePeriodChange(sources: sources, from: monthAgo)
        let yearChange = calculatePeriodChange(sources: sources, from: yearAgo)

        let allTimeReturn = totalValue - totalContributions
        let allTimeReturnPercentage = totalContributions > 0
            ? NSDecimalNumber(decimal: allTimeReturn / totalContributions).doubleValue * 100
            : 0

        let lastUpdated = snapshots.map { $0.date }.max()

        return PortfolioSummary(
            totalValue: totalValue,
            totalContributions: totalContributions,
            dayChange: dayChange.absolute,
            dayChangePercentage: dayChange.percentage,
            weekChange: weekChange.absolute,
            weekChangePercentage: weekChange.percentage,
            monthChange: monthChange.absolute,
            monthChangePercentage: monthChange.percentage,
            yearChange: yearChange.absolute,
            yearChangePercentage: yearChange.percentage,
            allTimeReturn: allTimeReturn,
            allTimeReturnPercentage: allTimeReturnPercentage,
            sourceCount: sources.count,
            lastUpdated: lastUpdated
        )
    }

    private func calculatePeriodChange(
        sources: [InvestmentSource],
        from startDate: Date
    ) -> (absolute: Decimal, percentage: Double) {
        var previousTotal: Decimal = 0
        var currentTotal: Decimal = 0

        for source in sources {
            currentTotal += source.latestValue

            // Find snapshot closest to start date
            let snapshots = source.sortedSnapshotsByDateAscending
            let previousSnapshot = snapshots.last { $0.date <= startDate } ?? snapshots.first
            previousTotal += previousSnapshot?.decimalValue ?? 0
        }

        let absolute = currentTotal - previousTotal
        let percentage = previousTotal > 0
            ? NSDecimalNumber(decimal: absolute / previousTotal).doubleValue * 100
            : 0

        return (absolute, percentage)
    }

    // MARK: - Investment Metrics

    func calculateMetrics(for snapshots: [Snapshot]) -> InvestmentMetrics {
        guard !snapshots.isEmpty else { return .empty }

        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }
        let values = sortedSnapshots.map { $0.decimalValue }

        guard let firstValue = values.first,
              let lastValue = values.last,
              firstValue != 0 else {
            return .empty
        }

        let totalValue = lastValue
        let totalContributions = sortedSnapshots.reduce(Decimal.zero) { $0 + $1.decimalContribution }
        let absoluteReturn = lastValue - firstValue
        let percentageReturn = (absoluteReturn / firstValue) * 100

        // Calculate monthly returns for advanced metrics
        let monthlyReturns = calculateMonthlyReturns(from: sortedSnapshots)

        let cagr = calculateCAGR(
            startValue: firstValue,
            endValue: lastValue,
            startDate: sortedSnapshots.first?.date ?? Date(),
            endDate: sortedSnapshots.last?.date ?? Date()
        )

        let twr = calculateTWR(snapshots: sortedSnapshots)
        let volatility = calculateVolatility(monthlyReturns: monthlyReturns)
        let maxDrawdown = calculateMaxDrawdown(values: values)
        let sharpeRatio = calculateSharpeRatio(
            averageReturn: monthlyReturns.map { $0.returnPercentage }.average(),
            volatility: volatility
        )

        let winRate = calculateWinRate(monthlyReturns: monthlyReturns)
        let averageMonthlyReturn = monthlyReturns.map { $0.returnPercentage }.average()

        return InvestmentMetrics(
            totalValue: totalValue,
            totalContributions: totalContributions,
            absoluteReturn: absoluteReturn,
            percentageReturn: percentageReturn,
            cagr: cagr,
            twr: twr,
            volatility: volatility,
            maxDrawdown: maxDrawdown,
            sharpeRatio: sharpeRatio,
            bestMonth: monthlyReturns.max(by: { $0.returnPercentage < $1.returnPercentage }),
            worstMonth: monthlyReturns.min(by: { $0.returnPercentage < $1.returnPercentage }),
            winRate: winRate,
            averageMonthlyReturn: averageMonthlyReturn,
            startDate: sortedSnapshots.first?.date,
            endDate: sortedSnapshots.last?.date,
            totalMonths: monthlyReturns.count
        )
    }

    // MARK: - CAGR (Compound Annual Growth Rate)

    func calculateCAGR(
        startValue: Decimal,
        endValue: Decimal,
        startDate: Date,
        endDate: Date
    ) -> Double {
        guard startValue > 0 else { return 0 }

        let years = Calendar.current.dateComponents(
            [.day],
            from: startDate,
            to: endDate
        ).day.map { Double($0) / 365.25 } ?? 0

        guard years > 0 else { return 0 }

        let ratio = NSDecimalNumber(decimal: endValue / startValue).doubleValue
        let cagr = pow(ratio, 1 / years) - 1

        return cagr * 100
    }

    // MARK: - TWR (Time-Weighted Return)

    func calculateTWR(snapshots: [Snapshot]) -> Double {
        guard snapshots.count >= 2 else { return 0 }

        var twr: Double = 1.0

        for i in 1..<snapshots.count {
            let previousValue = snapshots[i-1].decimalValue
            let currentValue = snapshots[i].decimalValue
            let contribution = snapshots[i].decimalContribution

            guard previousValue > 0 else { continue }

            // Adjust for contributions
            let adjustedPreviousValue = previousValue + contribution
            let periodReturn = NSDecimalNumber(
                decimal: currentValue / adjustedPreviousValue
            ).doubleValue

            twr *= periodReturn
        }

        return (twr - 1) * 100
    }

    // MARK: - Volatility (Annualized Standard Deviation)

    func calculateVolatility(monthlyReturns: [InvestmentMetrics.MonthlyReturn]) -> Double {
        let returns = monthlyReturns.map { $0.returnPercentage }
        guard returns.count >= 2 else { return 0 }

        let mean = returns.average()
        let squaredDifferences = returns.map { pow($0 - mean, 2) }
        let variance = squaredDifferences.reduce(0, +) / Double(returns.count - 1)
        let stdDev = sqrt(variance)

        // Annualize (multiply by sqrt(12) for monthly data)
        return stdDev * sqrt(12)
    }

    // MARK: - Max Drawdown

    func calculateMaxDrawdown(values: [Decimal]) -> Double {
        guard !values.isEmpty else { return 0 }

        var maxDrawdown: Double = 0
        var peak = values[0]

        for value in values {
            if value > peak {
                peak = value
            }

            guard peak > 0 else { continue }

            let drawdown = NSDecimalNumber(
                decimal: (peak - value) / peak
            ).doubleValue * 100

            maxDrawdown = max(maxDrawdown, drawdown)
        }

        return maxDrawdown
    }

    // MARK: - Sharpe Ratio

    func calculateSharpeRatio(
        averageReturn: Double,
        volatility: Double,
        riskFreeRate: Double = 2.0 // Assume 2% annual risk-free rate
    ) -> Double {
        guard volatility > 0 else { return 0 }

        // Convert to monthly risk-free rate
        let monthlyRiskFree = riskFreeRate / 12

        return (averageReturn - monthlyRiskFree) / volatility * sqrt(12)
    }

    // MARK: - Win Rate

    func calculateWinRate(monthlyReturns: [InvestmentMetrics.MonthlyReturn]) -> Double {
        guard !monthlyReturns.isEmpty else { return 0 }

        let positiveMonths = monthlyReturns.filter { $0.returnPercentage > 0 }.count
        return Double(positiveMonths) / Double(monthlyReturns.count) * 100
    }

    // MARK: - Monthly Returns

    func calculateMonthlyReturns(from snapshots: [Snapshot]) -> [InvestmentMetrics.MonthlyReturn] {
        guard snapshots.count >= 2 else { return [] }

        var monthlyReturns: [InvestmentMetrics.MonthlyReturn] = []
        let calendar = Calendar.current

        // Group snapshots by month
        var monthlySnapshots: [String: [Snapshot]] = [:]
        let formatter = DateFormatter()
        formatter.dateFormat = "yyyy-MM"

        for snapshot in snapshots {
            let key = formatter.string(from: snapshot.date)
            monthlySnapshots[key, default: []].append(snapshot)
        }

        // Sort months
        let sortedMonths = monthlySnapshots.keys.sorted()

        for i in 1..<sortedMonths.count {
            let previousMonth = sortedMonths[i-1]
            let currentMonth = sortedMonths[i]

            guard let previousSnapshots = monthlySnapshots[previousMonth],
                  let currentSnapshots = monthlySnapshots[currentMonth],
                  let previousValue = previousSnapshots.last?.decimalValue,
                  let currentValue = currentSnapshots.last?.decimalValue,
                  previousValue > 0 else {
                continue
            }

            let returnPercentage = NSDecimalNumber(
                decimal: (currentValue - previousValue) / previousValue
            ).doubleValue * 100

            if let date = formatter.date(from: currentMonth) {
                monthlyReturns.append(InvestmentMetrics.MonthlyReturn(
                    date: date,
                    returnPercentage: returnPercentage
                ))
            }
        }

        return monthlyReturns
    }

    // MARK: - Category Metrics

    func calculateCategoryMetrics(
        for categories: [Category],
        totalPortfolioValue: Decimal
    ) -> [CategoryMetrics] {
        categories.map { category in
            let allSnapshots = category.sourcesArray.flatMap { $0.snapshotsArray }
            let metrics = calculateMetrics(for: allSnapshots)

            let percentage = totalPortfolioValue > 0
                ? NSDecimalNumber(decimal: category.totalValue / totalPortfolioValue).doubleValue * 100
                : 0

            return CategoryMetrics(
                id: category.id,
                categoryName: category.name,
                colorHex: category.colorHex,
                icon: category.icon,
                totalValue: category.totalValue,
                percentageOfPortfolio: percentage,
                metrics: metrics
            )
        }
    }
}

// MARK: - Array Extension

extension Array where Element == Double {
    func average() -> Double {
        guard !isEmpty else { return 0 }
        return reduce(0, +) / Double(count)
    }
}