InvestmentTrackerApp/PortfolioJournal/Services/PredictionEngine.swift

import Foundation

class PredictionEngine {
    static let shared = PredictionEngine()

    private static let calendar = Calendar.current

    private init() {}

    // MARK: - Public Interface

    func predict(snapshots: [Snapshot], monthsAhead: Int = 12, algorithm: PredictionAlgorithm? = nil) -> PredictionResult {
        guard snapshots.count >= 3 else { return emptyResult() }
        let sorted = snapshots.sorted { $0.date < $1.date }
        return predictFromValues(
            sorted.map { $0.decimalValue.doubleValue },
            dates: sorted.map { $0.date },
            monthsAhead: monthsAhead,
            algorithm: algorithm
        )
    }

    func predict(series: [(date: Date, value: Decimal)], monthsAhead: Int = 12, algorithm: PredictionAlgorithm? = nil) -> PredictionResult {
        guard series.count >= 3 else { return emptyResult() }
        let sorted = series.sorted { $0.date < $1.date }
        return predictFromValues(
            sorted.map { NSDecimalNumber(decimal: $0.value).doubleValue },
            dates: sorted.map { $0.date },
            monthsAhead: monthsAhead,
            algorithm: algorithm
        )
    }

    // MARK: - Private Core

    private func emptyResult() -> PredictionResult {
        PredictionResult(predictions: [], algorithm: .linear, accuracy: 0, volatility: 0)
    }

    private func predictFromValues(_ values: [Double], dates: [Date], monthsAhead: Int, algorithm: PredictionAlgorithm?) -> PredictionResult {
        let volatility = calculateVolatility(values: values)
        let selectedAlgorithm = algorithm ?? selectBestAlgorithm(volatility: volatility)
        let lastDate = dates.last!
        let firstDate = dates.first!

        let predictions: [Prediction]
        let accuracy: Double

        switch selectedAlgorithm {
        case .linear:
            predictions = linearPredictions(values: values, firstDate: firstDate, lastDate: lastDate, monthsAhead: monthsAhead)
            accuracy = linearAccuracy(values: values, firstDate: firstDate)
        case .exponentialSmoothing:
            predictions = esPredictions(values: values, lastDate: lastDate, monthsAhead: monthsAhead)
            accuracy = esAccuracy(values: values)
        case .movingAverage:
            predictions = maPredictions(values: values, lastDate: lastDate, monthsAhead: monthsAhead)
            accuracy = maAccuracy(values: values)
        case .holtTrend:
            predictions = holtPredictions(values: values, lastDate: lastDate, monthsAhead: monthsAhead)
            accuracy = holtAccuracy(values: values)
        }

        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

    private func linearPredictions(values: [Double], firstDate: Date, lastDate: Date, monthsAhead: Int) -> [Prediction] {
        let dataPoints = values.enumerated().map { (x: Double($0.offset), y: $0.element) }
        let (slope, intercept) = calculateLinearRegression(dataPoints: dataPoints)
        let residualStdDev = calculateResidualStdDev(dataPoints: dataPoints, slope: slope, intercept: intercept)
        let n = Double(values.count)

        return (1...monthsAhead).compactMap { month in
            guard let futureDate = Self.calendar.date(byAdding: .month, value: month, to: lastDate) else { return nil }
            let x = n - 1 + Double(month)
            let predicted = max(0, slope * x + intercept)
            let width = residualStdDev * 1.96 * (1.0 + Double(month) * 0.02)
            return Prediction(date: futureDate, predictedValue: Decimal(predicted), algorithm: .linear,
                              confidenceInterval: .init(lower: Decimal(max(0, predicted - width)), upper: Decimal(predicted + width)))
        }
    }

    private func linearAccuracy(values: [Double], firstDate: Date) -> Double {
        guard values.count >= 5 else { return 0.5 }
        let splitIndex = Int(Double(values.count) * 0.8)
        let trainPoints = Array(values.prefix(splitIndex)).enumerated().map { (x: Double($0.offset), y: $0.element) }
        let (slope, intercept) = calculateLinearRegression(dataPoints: trainPoints)
        let validation = Array(values.suffix(from: splitIndex))
        let mean = validation.reduce(0, +) / Double(validation.count)
        var ssRes = 0.0, ssTot = 0.0
        for (i, actual) in validation.enumerated() {
            let x = Double(splitIndex + i)
            ssRes += pow(actual - (slope * x + intercept), 2)
            ssTot += pow(actual - mean, 2)
        }
        guard ssTot != 0 else { return 0.5 }
        return min(1.0, max(0, 1 - ssRes / ssTot))
    }

    // MARK: - Exponential Smoothing

    private func esPredictions(values: [Double], lastDate: Date, monthsAhead: Int, alpha: Double = 0.3) -> [Prediction] {
        var smoothed = values[0]
        for i in 1..<values.count { smoothed = alpha * values[i] + (1 - alpha) * smoothed }

        var trend = 0.0
        if values.count >= 2 {
            let window = Double(min(3, values.count))
            let recentAvg: Double = values.suffix(3).reduce(0, +) / window
            let earlyAvg: Double = values.prefix(3).reduce(0, +) / window
            trend = (recentAvg - earlyAvg) / Double(values.count)
        }
        let stdDev = calculateStdDev(values: values)

        return (1...monthsAhead).compactMap { month in
            guard let futureDate = Self.calendar.date(byAdding: .month, value: month, to: lastDate) else { return nil }
            let predicted = max(0, smoothed + trend * Double(month))
            let width = stdDev * 1.96 * (1.0 + Double(month) * 0.05)
            return Prediction(date: futureDate, predictedValue: Decimal(predicted), algorithm: .exponentialSmoothing,
                              confidenceInterval: .init(lower: Decimal(max(0, predicted - width)), upper: Decimal(predicted + width)))
        }
    }

    private func esAccuracy(values: [Double]) -> Double {
        guard values.count >= 5 else { return 0.5 }
        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 validation = Array(values.suffix(from: splitIndex))
        let mean = validation.reduce(0, +) / Double(validation.count)
        var ssRes = 0.0, ssTot = 0.0
        for (i, actual) in validation.enumerated() {
            let predicted = smoothed + (smoothed - values[splitIndex - 1]) * Double(i + 1) / Double(splitIndex)
            ssRes += pow(actual - predicted, 2)
            ssTot += pow(actual - mean, 2)
        }
        guard ssTot != 0 else { return 0.5 }
        return max(0, min(1.0, 1 - ssRes / ssTot))
    }

    // MARK: - Moving Average

    private func maPredictions(values: [Double], lastDate: Date, monthsAhead: Int, windowSize: Int = 3) -> [Prediction] {
        guard values.count >= windowSize else { return [] }
        let recent = Array(values.suffix(windowSize))
        let movingAvg = recent.reduce(0, +) / Double(windowSize)
        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)

        return (1...monthsAhead).compactMap { month in
            guard let futureDate = Self.calendar.date(byAdding: .month, value: month, to: lastDate) else { return nil }
            let predicted = max(0, movingAvg + avgChange * Double(month))
            let width = stdDev * 1.96 * (1.0 + Double(month) * 0.03)
            return Prediction(date: futureDate, predictedValue: Decimal(predicted), algorithm: .movingAverage,
                              confidenceInterval: .init(lower: Decimal(max(0, predicted - width)), upper: Decimal(predicted + width)))
        }
    }

    private func maAccuracy(values: [Double], windowSize: Int = 3) -> Double {
        guard values.count >= 5 else { return 0.5 }
        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 validation = Array(values.suffix(from: splitIndex))
        let mean = validation.reduce(0, +) / Double(validation.count)
        var ssRes = 0.0, ssTot = 0.0
        for actual in validation {
            ssRes += pow(actual - movingAvg, 2)
            ssTot += pow(actual - mean, 2)
        }
        guard ssTot != 0 else { return 0.5 }
        return max(0, min(1.0, 1 - ssRes / ssTot))
    }

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

    private func holtPredictions(values: [Double], lastDate: Date, monthsAhead: Int, alpha: Double = 0.4, beta: Double = 0.3) -> [Prediction] {
        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 stdDev = calculateStdDev(values: zip(values, fitted).map { $0 - $1 })

        return (1...monthsAhead).compactMap { month in
            guard let futureDate = Self.calendar.date(byAdding: .month, value: month, to: lastDate) else { return nil }
            let predicted = max(0, level + Double(month) * trend)
            let width = stdDev * 1.96 * (1.0 + Double(month) * 0.04)
            return Prediction(date: futureDate, predictedValue: Decimal(predicted), algorithm: .holtTrend,
                              confidenceInterval: .init(lower: Decimal(max(0, predicted - width)), upper: Decimal(predicted + width)))
        }
    }

    private func holtAccuracy(values: [Double]) -> Double {
        guard values.count >= 5 else { return 0.5 }
        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 validation = Array(values.suffix(from: splitIndex))
        let mean = validation.reduce(0, +) / Double(validation.count)
        var ssRes = 0.0, ssTot = 0.0
        for (i, actual) in validation.enumerated() {
            ssRes += pow(actual - (level + Double(i + 1) * trend), 2)
            ssTot += pow(actual - mean, 2)
        }
        guard ssTot != 0 else { return 0.5 }
        return max(0, min(1.0, 1 - ssRes / ssTot))
    }

    // MARK: - Helpers

    private func calculateVolatility(values: [Double]) -> Double {
        guard values.count >= 2 else { return 0 }
        var returns: [Double] = []
        for i in 1..<values.count {
            guard values[i - 1] != 0 else { continue }
            returns.append((values[i] - values[i - 1]) / values[i - 1] * 100)
        }
        return calculateStdDev(values: returns)
    }

    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 denom = n * sumX2 - sumX * sumX
        guard denom != 0 else { return (0, sumY / n) }
        let slope = (n * sumXY - sumX * sumY) / denom
        return (slope, (sumY - slope * sumX) / n)
    }

    private func calculateResidualStdDev(dataPoints: [(x: Double, y: Double)], slope: Double, intercept: Double) -> Double {
        guard dataPoints.count > 2 else { return 0 }
        let mse = dataPoints.map { pow($0.y - (slope * $0.x + intercept), 2) }.reduce(0, +) / Double(dataPoints.count - 2)
        return sqrt(mse)
    }

    private func calculateStdDev(values: [Double]) -> Double {
        guard values.count >= 2 else { return 0 }
        let mean = values.reduce(0, +) / Double(values.count)
        let variance = values.map { pow($0 - mean, 2) }.reduce(0, +) / Double(values.count - 1)
        return sqrt(variance)
    }

    // MARK: - Public compatibility (kept for external callers)

    func predictLinear(snapshots: [Snapshot], monthsAhead: Int = 12) -> [Prediction] {
        guard snapshots.count >= 3 else { return [] }
        let sorted = snapshots.sorted { $0.date < $1.date }
        return linearPredictions(
            values: sorted.map { $0.decimalValue.doubleValue },
            firstDate: sorted.first!.date,
            lastDate: sorted.last!.date,
            monthsAhead: monthsAhead
        )
    }
}