InvestmentTrackerApp/PortfolioJournal/Services/CalculationService.swift

import Foundation

class CalculationService {
    static let shared = CalculationService()

    // MARK: - Performance: Caching
    private var cachedMonthlyReturns: [ObjectIdentifier: [InvestmentMetrics.MonthlyReturn]] = [:]
    private var cacheVersion: Int = 0

    private init() {}

    /// Call this when underlying data changes to invalidate caches
    func invalidateCache() {
        cachedMonthlyReturns.removeAll()
        cacheVersion += 1
    }

    // MARK: - Shared DateFormatter (avoid repeated allocations)
    private static let monthYearFormatter: DateFormatter = {
        let formatter = DateFormatter()
        formatter.dateFormat = "yyyy-MM"
        return formatter
    }()

    // 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: - Monthly Summary

    func calculateMonthlySummary(
        sources: [InvestmentSource],
        snapshots: [Snapshot],
        range: DateRange = .thisMonth
    ) -> MonthlySummary {
        let startDate = range.start
        let endDate = range.end

        var startingValue: Decimal = 0
        var endingValue: Decimal = 0
        var contributions: Decimal = 0

        let snapshotsBySource = Dictionary(grouping: snapshots) { $0.source?.id }

        for source in sources {
            let sourceSnapshots = snapshotsBySource[source.id] ?? []
            let sorted = sourceSnapshots.sorted { $0.date < $1.date }

            let startSnapshot = sorted.last { $0.date <= startDate }
            let endSnapshot = sorted.last { $0.date <= endDate }

            startingValue += startSnapshot?.decimalValue ?? 0
            endingValue += endSnapshot?.decimalValue ?? 0

            let rangeContributions = sorted
                .filter { $0.date >= startDate && $0.date <= endDate }
                .reduce(Decimal.zero) { $0 + $1.decimalContribution }
            contributions += rangeContributions
        }

        let netPerformance = endingValue - startingValue - contributions

        return MonthlySummary(
            periodLabel: "This Month",
            startDate: startDate,
            endDate: endDate,
            startingValue: startingValue,
            endingValue: endingValue,
            contributions: contributions,
            netPerformance: netPerformance
        )
    }

    // 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 [] }

        // Performance: Use shared formatter instead of creating new one each call
        let formatter = Self.monthYearFormatter

        // Group snapshots by month - pre-allocate capacity
        var monthlySnapshots: [String: [Snapshot]] = [:]
        monthlySnapshots.reserveCapacity(min(snapshots.count, 60)) // Reasonable max months

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

        // Sort months
        let sortedMonths = monthlySnapshots.keys.sorted()
        guard sortedMonths.count >= 2 else { return [] }

        // Pre-allocate result array
        var monthlyReturns: [InvestmentMetrics.MonthlyReturn] = []
        monthlyReturns.reserveCapacity(sortedMonths.count - 1)

        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],
        sources: [InvestmentSource],
        totalPortfolioValue: Decimal
    ) -> [CategoryMetrics] {
        categories.map { category in
            let categorySources = sources.filter { $0.category?.id == category.id }
            let allSnapshots = categorySources.flatMap { $0.snapshotsArray }
            let metrics = calculateCategoryMetrics(from: allSnapshots)
            let categoryValue = categorySources.reduce(Decimal.zero) { $0 + $1.latestValue }

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

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

    private struct SeriesPoint {
        let date: Date
        let value: Decimal
        let contribution: Decimal
    }

    private func calculateCategoryMetrics(from snapshots: [Snapshot]) -> InvestmentMetrics {
        let series = buildCategorySeries(from: snapshots)
        guard !series.isEmpty else { return .empty }

        let sortedSeries = series.sorted { $0.date < $1.date }
        let values = sortedSeries.map { $0.value }

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

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

        let monthlyReturns = calculateMonthlyReturns(from: sortedSeries)

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

        let twr = calculateTWR(series: sortedSeries)
        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: sortedSeries.first?.date,
            endDate: sortedSeries.last?.date,
            totalMonths: monthlyReturns.count
        )
    }

    private func buildCategorySeries(from snapshots: [Snapshot]) -> [SeriesPoint] {
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }
        let uniqueDates = Array(Set(sortedSnapshots.map { Calendar.current.startOfDay(for: $0.date) }))
            .sorted()
        guard !uniqueDates.isEmpty else { return [] }

        var snapshotsBySource: [UUID: [(date: Date, value: Decimal)]] = [:]
        var contributionsByDate: [Date: Decimal] = [:]

        for snapshot in sortedSnapshots {
            guard let sourceId = snapshot.source?.id else { continue }
            snapshotsBySource[sourceId, default: []].append(
                (date: snapshot.date, value: snapshot.decimalValue)
            )
            let day = Calendar.current.startOfDay(for: snapshot.date)
            contributionsByDate[day, default: 0] += snapshot.decimalContribution
        }

        var indices: [UUID: Int] = [:]
        var series: [SeriesPoint] = []

        for (index, date) in uniqueDates.enumerated() {
            let nextDate = index + 1 < uniqueDates.count
                ? uniqueDates[index + 1]
                : Date.distantFuture
            var total: Decimal = 0

            for (sourceId, sourceSnapshots) in snapshotsBySource {
                var currentIndex = indices[sourceId] ?? 0
                var latest: (date: Date, value: Decimal)?

                while currentIndex < sourceSnapshots.count && sourceSnapshots[currentIndex].date < nextDate {
                    latest = sourceSnapshots[currentIndex]
                    currentIndex += 1
                }

                indices[sourceId] = currentIndex

                if let latest {
                    total += latest.value
                }
            }

            series.append(
                SeriesPoint(
                    date: date,
                    value: total,
                    contribution: contributionsByDate[date] ?? 0
                )
            )
        }

        return series
    }

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

        // Performance: Use shared formatter
        let formatter = Self.monthYearFormatter

        var monthlySeries: [String: [SeriesPoint]] = [:]
        monthlySeries.reserveCapacity(min(series.count, 60))

        for point in series {
            let key = formatter.string(from: point.date)
            monthlySeries[key, default: []].append(point)
        }

        let sortedMonths = monthlySeries.keys.sorted()
        guard sortedMonths.count >= 2 else { return [] }

        var monthlyReturns: [InvestmentMetrics.MonthlyReturn] = []
        monthlyReturns.reserveCapacity(sortedMonths.count - 1)

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

            guard let previousPoints = monthlySeries[previousMonth],
                  let currentPoints = monthlySeries[currentMonth],
                  let previousValue = previousPoints.last?.value,
                  let currentValue = currentPoints.last?.value,
                  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
    }

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

        var twr: Double = 1.0
        for i in 1..<series.count {
            let previousValue = series[i - 1].value
            let currentValue = series[i].value
            let contribution = series[i].contribution

            guard previousValue > 0 else { continue }

            let periodReturn = (currentValue - contribution) / previousValue
            twr *= NSDecimalNumber(decimal: periodReturn).doubleValue
        }

        return (twr - 1) * 100
    }
}

// MARK: - Array Extension

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