InvestmentTrackerApp/PortfolioJournal/Services/PredictionEngine.swift

import Foundation

class PredictionEngine {
    static let shared = PredictionEngine()

    private let context = CoreDataStack.shared.viewContext

    // MARK: - Performance: Cached Calendar reference
    private static let calendar = Calendar.current

    private init() {}

    // MARK: - Main Prediction Interface

    func predict(
        snapshots: [Snapshot],
        monthsAhead: Int = 12,
        algorithm: PredictionAlgorithm? = nil
    ) -> PredictionResult {
        guard snapshots.count >= 3 else {
            return PredictionResult(
                predictions: [],
                algorithm: .linear,
                accuracy: 0,
                volatility: 0
            )
        }

        // Sort snapshots by date
        let sortedSnapshots = snapshots.sorted { $0.date < $1.date }

        // Calculate volatility for algorithm selection
        let volatility = calculateVolatility(snapshots: sortedSnapshots)

        // Select algorithm if not specified
        let selectedAlgorithm = algorithm ?? selectBestAlgorithm(volatility: volatility)

        // Generate predictions
        let predictions: [Prediction]
        let accuracy: Double

        switch selectedAlgorithm {
        case .linear:
            predictions = predictLinear(snapshots: sortedSnapshots, monthsAhead: monthsAhead)
            accuracy = calculateLinearAccuracy(snapshots: sortedSnapshots)
        case .exponentialSmoothing:
            predictions = predictExponentialSmoothing(snapshots: sortedSnapshots, monthsAhead: monthsAhead)
            accuracy = calculateESAccuracy(snapshots: sortedSnapshots)
        case .movingAverage:
            predictions = predictMovingAverage(snapshots: sortedSnapshots, monthsAhead: monthsAhead)
            accuracy = calculateMAAccuracy(snapshots: sortedSnapshots)
        case .holtTrend:
            predictions = predictHoltTrend(snapshots: sortedSnapshots, monthsAhead: monthsAhead)
            accuracy = calculateHoltAccuracy(snapshots: sortedSnapshots)
        }

        return PredictionResult(
            predictions: predictions,
            algorithm: selectedAlgorithm,
            accuracy: accuracy,
            volatility: volatility
        )
    }

    // MARK: - Algorithm Selection

    private func selectBestAlgorithm(volatility: Double) -> PredictionAlgorithm {
        switch volatility {
        case 0..<8:
            return .holtTrend
        case 8..<20:
            return .exponentialSmoothing
        default:
            return .movingAverage
        }
    }

    // MARK: - Linear Regression

    func predictLinear(snapshots: [Snapshot], monthsAhead: Int = 12) -> [Prediction] {
        guard snapshots.count >= 3 else { return [] }

        guard let firstDate = snapshots.first?.date else { return [] }

        let dataPoints: [(x: Double, y: Double)] = snapshots.map { snapshot in
            let daysSinceStart = snapshot.date.timeIntervalSince(firstDate) / 86400
            return (x: daysSinceStart, y: snapshot.decimalValue.doubleValue)
        }

        let (slope, intercept) = calculateLinearRegression(dataPoints: dataPoints)
        let residualStdDev = calculateResidualStdDev(dataPoints: dataPoints, slope: slope, intercept: intercept)

        var predictions: [Prediction] = []
        let lastDate = snapshots.last!.date

        for month in 1...monthsAhead {
            guard let futureDate = Self.calendar.date(
                byAdding: .month,
                value: month,
                to: lastDate
            ) else { continue }

            let daysFromStart = futureDate.timeIntervalSince(firstDate) / 86400
            let predictedValue = max(0, slope * daysFromStart + intercept)

            // Widen confidence interval for further predictions
            let confidenceMultiplier = 1.0 + (Double(month) * 0.02)
            let intervalWidth = residualStdDev * 1.96 * confidenceMultiplier

            predictions.append(Prediction(
                date: futureDate,
                predictedValue: Decimal(predictedValue),
                algorithm: .linear,
                confidenceInterval: Prediction.ConfidenceInterval(
                    lower: Decimal(max(0, predictedValue - intervalWidth)),
                    upper: Decimal(predictedValue + intervalWidth)
                )
            ))
        }

        return predictions
    }

    private func calculateLinearRegression(
        dataPoints: [(x: Double, y: Double)]
    ) -> (slope: Double, intercept: Double) {
        let n = Double(dataPoints.count)
        let sumX = dataPoints.reduce(0) { $0 + $1.x }
        let sumY = dataPoints.reduce(0) { $0 + $1.y }
        let sumXY = dataPoints.reduce(0) { $0 + ($1.x * $1.y) }
        let sumX2 = dataPoints.reduce(0) { $0 + ($1.x * $1.x) }

        let denominator = n * sumX2 - sumX * sumX
        guard denominator != 0 else { return (0, sumY / n) }

        let slope = (n * sumXY - sumX * sumY) / denominator
        let intercept = (sumY - slope * sumX) / n

        return (slope, intercept)
    }

    private func calculateResidualStdDev(
        dataPoints: [(x: Double, y: Double)],
        slope: Double,
        intercept: Double
    ) -> Double {
        guard dataPoints.count > 2 else { return 0 }

        let residuals = dataPoints.map { point in
            let predicted = slope * point.x + intercept
            return pow(point.y - predicted, 2)
        }

        let meanSquaredError = residuals.reduce(0, +) / Double(dataPoints.count - 2)
        return sqrt(meanSquaredError)
    }

    private func calculateLinearAccuracy(snapshots: [Snapshot]) -> Double {
        guard snapshots.count >= 5 else { return 0.5 }

        // Use last 20% of data for validation
        let splitIndex = Int(Double(snapshots.count) * 0.8)
        let trainingData = Array(snapshots.prefix(splitIndex))
        let validationData = Array(snapshots.suffix(from: splitIndex))

        guard let firstDate = trainingData.first?.date else { return 0.5 }

        let trainPoints = trainingData.map { snapshot in
            (x: snapshot.date.timeIntervalSince(firstDate) / 86400, y: snapshot.decimalValue.doubleValue)
        }

        let (slope, intercept) = calculateLinearRegression(dataPoints: trainPoints)

        // Calculate R-squared on validation data
        let validationValues = validationData.map { $0.decimalValue.doubleValue }
        let meanValidation = validationValues.reduce(0, +) / Double(validationValues.count)

        var ssRes: Double = 0
        var ssTot: Double = 0

        for snapshot in validationData {
            let x = snapshot.date.timeIntervalSince(firstDate) / 86400
            let actual = snapshot.decimalValue.doubleValue
            let predicted = slope * x + intercept

            ssRes += pow(actual - predicted, 2)
            ssTot += pow(actual - meanValidation, 2)
        }

        guard ssTot != 0 else { return 0.5 }
        let rSquared = max(0, 1 - (ssRes / ssTot))

        return min(1.0, rSquared)
    }

    // MARK: - Exponential Smoothing

    func predictExponentialSmoothing(
        snapshots: [Snapshot],
        monthsAhead: Int = 12,
        alpha: Double = 0.3
    ) -> [Prediction] {
        guard snapshots.count >= 3 else { return [] }

        let values = snapshots.map { $0.decimalValue.doubleValue }

        // Calculate smoothed values
        var smoothed = values[0]
        for i in 1..<values.count {
            smoothed = alpha * values[i] + (1 - alpha) * smoothed
        }

        // Calculate trend
        var trend: Double = 0
        if values.count >= 2 {
            let recentChange = values.suffix(3).reduce(0) { $0 + $1 } / 3.0 -
                              values.prefix(3).reduce(0) { $0 + $1 } / 3.0
            trend = recentChange / Double(values.count)
        }

        // Calculate standard deviation for confidence interval
        let stdDev = calculateStdDev(values: values)

        var predictions: [Prediction] = []
        let lastDate = snapshots.last!.date

        for month in 1...monthsAhead {
            guard let futureDate = Self.calendar.date(
                byAdding: .month,
                value: month,
                to: lastDate
            ) else { continue }

            let predictedValue = max(0, smoothed + trend * Double(month))
            let intervalWidth = stdDev * 1.96 * (1.0 + Double(month) * 0.05)

            predictions.append(Prediction(
                date: futureDate,
                predictedValue: Decimal(predictedValue),
                algorithm: .exponentialSmoothing,
                confidenceInterval: Prediction.ConfidenceInterval(
                    lower: Decimal(max(0, predictedValue - intervalWidth)),
                    upper: Decimal(predictedValue + intervalWidth)
                )
            ))
        }

        return predictions
    }

    private func calculateESAccuracy(snapshots: [Snapshot]) -> Double {
        guard snapshots.count >= 5 else { return 0.5 }

        let values = snapshots.map { $0.decimalValue.doubleValue }
        let splitIndex = Int(Double(values.count) * 0.8)

        var smoothed = values[0]
        for i in 1..<splitIndex {
            smoothed = 0.3 * values[i] + 0.7 * smoothed
        }

        let validationValues = Array(values.suffix(from: splitIndex))
        let meanValidation = validationValues.reduce(0, +) / Double(validationValues.count)

        var ssRes: Double = 0
        var ssTot: Double = 0

        for (i, actual) in validationValues.enumerated() {
            let predicted = smoothed + (smoothed - values[splitIndex - 1]) * Double(i + 1) / Double(splitIndex)
            ssRes += pow(actual - predicted, 2)
            ssTot += pow(actual - meanValidation, 2)
        }

        guard ssTot != 0 else { return 0.5 }
        return max(0, min(1.0, 1 - (ssRes / ssTot)))
    }

    // MARK: - Moving Average

    func predictMovingAverage(
        snapshots: [Snapshot],
        monthsAhead: Int = 12,
        windowSize: Int = 3
    ) -> [Prediction] {
        guard snapshots.count >= windowSize else { return [] }

        let values = snapshots.map { $0.decimalValue.doubleValue }

        // Calculate moving average of last window
        let recentValues = Array(values.suffix(windowSize))
        let movingAverage = recentValues.reduce(0, +) / Double(windowSize)

        // Calculate average monthly change
        var changes: [Double] = []
        for i in 1..<values.count {
            changes.append(values[i] - values[i - 1])
        }
        let avgChange = changes.isEmpty ? 0 : changes.reduce(0, +) / Double(changes.count)

        let stdDev = calculateStdDev(values: values)

        var predictions: [Prediction] = []
        let lastDate = snapshots.last!.date

        for month in 1...monthsAhead {
            guard let futureDate = Self.calendar.date(
                byAdding: .month,
                value: month,
                to: lastDate
            ) else { continue }

            let predictedValue = max(0, movingAverage + avgChange * Double(month))
            let intervalWidth = stdDev * 1.96 * (1.0 + Double(month) * 0.03)

            predictions.append(Prediction(
                date: futureDate,
                predictedValue: Decimal(predictedValue),
                algorithm: .movingAverage,
                confidenceInterval: Prediction.ConfidenceInterval(
                    lower: Decimal(max(0, predictedValue - intervalWidth)),
                    upper: Decimal(predictedValue + intervalWidth)
                )
            ))
        }

        return predictions
    }

    private func calculateMAAccuracy(snapshots: [Snapshot]) -> Double {
        guard snapshots.count >= 5 else { return 0.5 }

        let values = snapshots.map { $0.decimalValue.doubleValue }
        let windowSize = 3
        let splitIndex = Int(Double(values.count) * 0.8)

        guard splitIndex > windowSize else { return 0.5 }

        let recentWindow = Array(values[(splitIndex - windowSize)..<splitIndex])
        let movingAvg = recentWindow.reduce(0, +) / Double(windowSize)

        let validationValues = Array(values.suffix(from: splitIndex))
        let meanValidation = validationValues.reduce(0, +) / Double(validationValues.count)

        var ssRes: Double = 0
        var ssTot: Double = 0

        for actual in validationValues {
            ssRes += pow(actual - movingAvg, 2)
            ssTot += pow(actual - meanValidation, 2)
        }

        guard ssTot != 0 else { return 0.5 }
        return max(0, min(1.0, 1 - (ssRes / ssTot)))
    }

    // MARK: - Holt Trend (Double Exponential Smoothing)

    func predictHoltTrend(
        snapshots: [Snapshot],
        monthsAhead: Int = 12,
        alpha: Double = 0.4,
        beta: Double = 0.3
    ) -> [Prediction] {
        guard snapshots.count >= 3 else { return [] }

        let values = snapshots.map { $0.decimalValue.doubleValue }
        var level = values[0]
        var trend = values[1] - values[0]

        var fitted: [Double] = []
        for value in values {
            let lastLevel = level
            level = alpha * value + (1 - alpha) * (level + trend)
            trend = beta * (level - lastLevel) + (1 - beta) * trend
            fitted.append(level + trend)
        }

        let residuals = zip(values, fitted).map { $0 - $1 }
        let stdDev = calculateStdDev(values: residuals)

        var predictions: [Prediction] = []
        let lastDate = snapshots.last!.date

        for month in 1...monthsAhead {
            guard let futureDate = Self.calendar.date(
                byAdding: .month,
                value: month,
                to: lastDate
            ) else { continue }

            let predictedValue = max(0, level + Double(month) * trend)
            let intervalWidth = stdDev * 1.96 * (1.0 + Double(month) * 0.04)

            predictions.append(Prediction(
                date: futureDate,
                predictedValue: Decimal(predictedValue),
                algorithm: .holtTrend,
                confidenceInterval: Prediction.ConfidenceInterval(
                    lower: Decimal(max(0, predictedValue - intervalWidth)),
                    upper: Decimal(predictedValue + intervalWidth)
                )
            ))
        }

        return predictions
    }

    private func calculateHoltAccuracy(snapshots: [Snapshot]) -> Double {
        guard snapshots.count >= 5 else { return 0.5 }

        let values = snapshots.map { $0.decimalValue.doubleValue }
        let splitIndex = Int(Double(values.count) * 0.8)
        guard splitIndex >= 2 else { return 0.5 }

        var level = values[0]
        var trend = values[1] - values[0]

        for value in values.prefix(splitIndex) {
            let lastLevel = level
            level = 0.4 * value + 0.6 * (level + trend)
            trend = 0.3 * (level - lastLevel) + 0.7 * trend
        }

        let validationValues = Array(values.suffix(from: splitIndex))
        let meanValidation = validationValues.reduce(0, +) / Double(validationValues.count)

        var ssRes: Double = 0
        var ssTot: Double = 0

        for (i, actual) in validationValues.enumerated() {
            let predicted = level + Double(i + 1) * trend
            ssRes += pow(actual - predicted, 2)
            ssTot += pow(actual - meanValidation, 2)
        }

        guard ssTot != 0 else { return 0.5 }
        return max(0, min(1.0, 1 - (ssRes / ssTot)))
    }

    // MARK: - Helpers

    private func calculateVolatility(snapshots: [Snapshot]) -> Double {
        let values = snapshots.map { $0.decimalValue.doubleValue }
        guard values.count >= 2 else { return 0 }

        var returns: [Double] = []
        for i in 1..<values.count {
            guard values[i - 1] != 0 else { continue }
            let periodReturn = (values[i] - values[i - 1]) / values[i - 1] * 100
            returns.append(periodReturn)
        }

        return calculateStdDev(values: returns)
    }

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

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

        return sqrt(variance)
    }
}