Trio/FreeAPS/Sources/APS/Storage/TDDStorage.swift

import Foundation
import Swinject

protocol TDDStorage {
    func calculateTDD(pumpHistory: [PumpHistoryEvent], basalProfile: [BasalProfileEntry], basalIncrement: Decimal) async
        -> TDDResult
}

/// Structure containing the results of TDD calculations
struct TDDResult {
    let total: Decimal
    let bolus: Decimal
    let tempBasal: Decimal
    let scheduledBasal: Decimal
    let weightedAverage: Decimal?
    let hoursOfData: Double
}

/// Implementation of the TDD Calculator
final class BaseTDDStorage: TDDStorage, Injectable {
    init(resolver: Resolver) {
        injectServices(resolver)
    }

    /// Main function to calculate TDD from pump history and basal profile
    /// - Parameters:
    ///   - pumpHistory: Array of pump history events
    ///   - basalProfile: Array of basal profile entries
    /// - Returns: TDDResult containing all calculated values
    func calculateTDD(
        pumpHistory: [PumpHistoryEvent],
        basalProfile: [BasalProfileEntry],
        basalIncrement: Decimal
    ) async -> TDDResult {
        debug(.apsManager, "Starting TDD calculation with \(pumpHistory.count) pump events")

        var bolusInsulin: Decimal = 0
        var tempInsulin: Decimal = 0
        var scheduledBasalInsulin: Decimal = 0

        let pumpData = calculatePumpDataHours(pumpHistory)
        debug(.apsManager, "Hours of pump data available: \(pumpData)")

        let bolusEvents = pumpHistory.filter({ $0.type == .bolus })
        let tempBasalEvents = pumpHistory.filter({ $0.type == .tempBasal })

        let gaps = findBasalGaps(in: tempBasalEvents)

        if !gaps.isEmpty {
            scheduledBasalInsulin = calculateScheduledBasalInsulin(gaps: gaps, profile: basalProfile).rounded(toPlaces: 2)
        }

        bolusInsulin = calculateBolusInsulin(bolusEvents).rounded(toPlaces: 2)
        debug(.apsManager, "Total bolus insulin: \(bolusInsulin)U")

        tempInsulin = calculateTempBasalInsulin(tempBasalEvents, basalIncrement: basalIncrement).rounded(toPlaces: 2)
        debug(.apsManager, "Total temp basal insulin: \(tempInsulin)U")

        let total = bolusInsulin + tempInsulin + scheduledBasalInsulin
        let weightedAverage = calculateWeightedAverage()

        debug(.apsManager, """
        TDD Summary:
        - Total: \(total)U
        - Bolus: \(bolusInsulin)U (\((bolusInsulin / total * 100).rounded(toPlaces: 1))%)
        - Temp Basal: \(tempInsulin)U (\((tempInsulin / total * 100).rounded(toPlaces: 1))%)
        - Scheduled Basal: \(scheduledBasalInsulin)U (\((scheduledBasalInsulin / total * 100).rounded(toPlaces: 1))%)
        - WeightedAverage: \(weightedAverage ?? 0)
        - Hours of Data: \(pumpData)
        """)

        return TDDResult(
            total: total,
            bolus: bolusInsulin,
            tempBasal: tempInsulin,
            scheduledBasal: scheduledBasalInsulin,
            weightedAverage: weightedAverage,
            hoursOfData: pumpData
        )
    }
    
    /// Finds gaps between tempBasal events where scheduled basal ran
    /// - Parameter tempBasalEvents: Array of pump history events of type tempBasal
    /// - Returns: Array of gaps, where each gap has a start and end time
    private func findBasalGaps(in tempBasalEvents: [PumpHistoryEvent]) -> [(start: Date, end: Date)] {
        guard !tempBasalEvents.isEmpty else {
            // No events = full day gap
            let startOfDay = Calendar.current.startOfDay(for: Date())
            let endOfDay = startOfDay.addingTimeInterval(24 * 60 * 60 - 1)
            return [(start: startOfDay, end: endOfDay)]
        }

        // Sort events by timestamp
        let sortedEvents = tempBasalEvents.sorted { $0.timestamp < $1.timestamp }

        var gaps: [(start: Date, end: Date)] = []

        // Track the end time of the last temp basal event
        var lastEndTime: Date?

        for (index, event) in sortedEvents.enumerated() {
            // Calculate the actual end time for the current event
            guard let duration = event.duration else { continue }
            var currentEndTime = event.timestamp.addingTimeInterval(TimeInterval(duration * 60))

            // Check for a cancellation
            if index < sortedEvents.count - 1 {
                let nextEvent = sortedEvents[index + 1]
                if nextEvent.timestamp < currentEndTime {
                    // The next event cancels this one, adjust the end time
                    currentEndTime = nextEvent.timestamp
                }
            }

            // If there’s a gap between the last event's end time and the current event's start time, record it
            if let lastEnd = lastEndTime, event.timestamp > lastEnd {
                gaps.append((start: lastEnd, end: event.timestamp))
            }

            // Update the last end time to the current event's (possibly adjusted) end time
            lastEndTime = currentEndTime
        }

        // Handle gap at the end of the dataset (if needed)
        if let lastEnd = lastEndTime {
            let endOfDay = Calendar.current.startOfDay(for: sortedEvents.first!.timestamp)
                .addingTimeInterval(24 * 60 * 60 - 1)
            if lastEnd < endOfDay {
                gaps.append((start: lastEnd, end: endOfDay))
            }
        }

        return gaps
    }

    /// Calculates the number of hours of available pump history data
    /// - Parameter pumpHistory: Array of pump history events
    /// - Returns: Number of hours of available data
    private func calculatePumpDataHours(_ pumpHistory: [PumpHistoryEvent]) -> Double {
        guard let firstEvent = pumpHistory.last, // we are fetching in a descending order
              let lastEvent = pumpHistory.first
        else {
            return 0
        }

        let startDate = firstEvent.timestamp
        var endDate = lastEvent.timestamp

        // If last event is a temp basal, use current time
        if lastEvent.type == .tempBasalDuration {
            endDate = Date()
        }

        return Double(endDate.timeIntervalSince(startDate)) / 3600.0
    }

    /// Calculates total bolus insulin from pump history
    /// - Parameter bolusEvents: Array of pump history events of type bolus
    /// - Returns: Total bolus insulin
    private func calculateBolusInsulin(_ bolusEvents: [PumpHistoryEvent]) -> Decimal {
        bolusEvents
            .reduce(Decimal(0)) { totalBolusInsulin, event in
                totalBolusInsulin + (event.amount ?? 0)
            }
    }

    /// Calculates insulin delivered via temporary basal rates, accounting for interruptions
    /// - Parameters:
    ///   - tempBasalEvents: Array of pump history events of type tempBasal
    ///   - basalIncrement: The smallest increment for basal rates
    /// - Returns: Total temporary basal insulin
    private func calculateTempBasalInsulin(_ tempBasalEvents: [PumpHistoryEvent], basalIncrement: Decimal) -> Decimal {
        guard !tempBasalEvents.isEmpty else { return Decimal(0) }

        // Sort events by timestamp to ensure proper order
        let sortedEvents = tempBasalEvents.sorted { $0.timestamp < $1.timestamp }

        return sortedEvents.enumerated().reduce(Decimal(0)) { totalInsulin, currentEvent in
            let (index, event) = currentEvent

            // Ensure the event is of type tempBasal and has valid data
            guard let rate = event.amount, // Rate in U/hr
                  let durationMinutes = event.duration, // Duration in minutes
                  rate > 0, durationMinutes > 0 else { return totalInsulin }

            // Calculate the actual duration in minutes the temp basal ran
            let actualDurationMinutes: Int
            if index < sortedEvents.count - 1 {
                // Next event exists; calculate if it interrupts the current event
                let nextEvent = sortedEvents[index + 1]
                let currentEndTime = event.timestamp.addingTimeInterval(TimeInterval(durationMinutes * 60))

                if nextEvent.timestamp < currentEndTime {
                    // Interrupted; calculate the actual duration
                    let interruptedDuration = nextEvent.timestamp.timeIntervalSince(event.timestamp) / 60
                    actualDurationMinutes = Int(interruptedDuration)
                } else {
                    // Not interrupted; use full duration
                    actualDurationMinutes = durationMinutes
                }
            } else {
                // Last event in the list; use its full duration
                actualDurationMinutes = durationMinutes
            }

            // Convert the duration to hours and calculate insulin
            let durationHours = Decimal(actualDurationMinutes) / 60
            let insulin = accountForIncrements(rate * durationHours, basalIncrement: basalIncrement)

            debug(
                .apsManager,
                "Temp basal: \(rate)U/hr for \(Decimal(actualDurationMinutes) / 60)hr = \(insulin)U (adjusted for interruptions if needed)"
            )

            // Add the calculated insulin to the total
            return totalInsulin + insulin
        }
    }

    /// Calculates total scheduled basal insulin within gaps
    /// - Parameters:
    ///   - tempBasalEvents: Array of pump history events of type tempBasal
    ///   - profile: Array of basal profile entries
    /// - Returns: Total scheduled basal insulin
    private func calculateScheduledBasalInsulin(
        gaps: [(start: Date, end: Date)],
        profile: [BasalProfileEntry]
    ) -> Decimal {
        var totalInsulin: Decimal = 0

        for gap in gaps {
            var currentTime = gap.start

            while currentTime < gap.end {
                guard let rate = findBasalRate(for: getTimeString(from: currentTime), in: profile) else {
                    debug(.apsManager, "No basal rate found for time \(currentTime)")
                    break
                }

                // Determine the next switch time in the basal profile or the end of the gap
                let nextSwitchTime = getNextBasalRateSwitch(after: currentTime, in: profile) ?? gap.end
                let endTime = min(nextSwitchTime, gap.end)

                // Calculate duration in hours and insulin delivered
                let duration = Decimal(endTime.timeIntervalSince(currentTime)) / 3600
                let insulin = rate * duration
                totalInsulin += insulin

                debug(.apsManager, "Scheduled basal: \(rate)U/hr from \(currentTime) to \(endTime) = \(insulin)U")

                // Move to the next time block
                currentTime = endTime
            }
        }

        return totalInsulin
    }

    /// Finds the next basal profile switch after a given time
    /// - Parameters:
    ///   - time: Current time
    ///   - profile: Array of basal profile entries
    /// - Returns: The time of the next switch, if any
    private func getNextBasalRateSwitch(after time: Date, in profile: [BasalProfileEntry]) -> Date? {
        let calendar = Calendar.current
        let timeMinutes = calendar.component(.hour, from: time) * 60 + calendar.component(.minute, from: time)

        // Find the next entry in the profile after the current time
        for entry in profile {
            if entry.minutes > timeMinutes {
                let nextSwitchTime = calendar.startOfDay(for: time).addingTimeInterval(TimeInterval(entry.minutes * 60))
                return nextSwitchTime
            }
        }

        return nil // No further switches; end of day
    }

    /// Converts a Date to a time string in "HH:mm:ss" format
    private func getTimeString(from date: Date) -> String {
        let formatter = DateFormatter()
        formatter.dateFormat = "HH:mm:ss"
        return formatter.string(from: date)
    }

    /// Rounds insulin amounts according to pump increment constraints
    /// - Parameter insulin: Raw insulin amount
    /// - Returns: Rounded insulin amount
    private func accountForIncrements(_ insulin: Decimal, basalIncrement: Decimal) -> Decimal {
        let incrementsRaw = insulin / basalIncrement

        if incrementsRaw >= 1 {
            // Convert to NSDecimalNumber to use its rounding capabilities
            let nsIncrements = NSDecimalNumber(decimal: incrementsRaw)
            let roundedIncrements = nsIncrements.rounding(
                accordingToBehavior:
                NSDecimalNumberHandler(
                    roundingMode: .down,
                    scale: 0,
                    raiseOnExactness: false,
                    raiseOnOverflow: false,
                    raiseOnUnderflow: false,
                    raiseOnDivideByZero: false
                )
            )
            return (roundedIncrements.decimalValue * basalIncrement).rounded(toPlaces: 3)
        }
        return 0
    }

    /// Finds the basal rate for a specific time in the profile, considering closest increments or wide coverage.
    /// - Parameters:
    ///   - timeString: Time string in "HH:mm:ss" format
    ///   - profile: Array of basal profile entries
    /// - Returns: Basal rate if found
    private func findBasalRate(for timeString: String, in profile: [BasalProfileEntry]) -> Decimal? {
        // Convert the timeString to minutes since midnight
        let timeComponents = timeString.split(separator: ":").compactMap { Int($0) }
        guard timeComponents.count == 3 else { return nil }
        let totalMinutes = timeComponents[0] * 60 + timeComponents[1]

        // If only one entry exists, return its rate (covers full 24 hours)
        guard profile.count > 1 else {
            return profile.first?.rate
        }

        // Find the closest matching basal entry
        for (index, entry) in profile.enumerated() {
            // Check if the time falls within the range of the current entry
            let startMinutes = entry.minutes
            let endMinutes = (index + 1 < profile.count) ? profile[index + 1].minutes : 1440 // End of the day

            if totalMinutes >= startMinutes, totalMinutes < endMinutes {
                return entry.rate
            }
        }

        // Default to nil if no match found
        return nil
    }

    /// Calculates weighted average of TDD from historical data
    /// - Returns: Weighted average if available
    private func calculateWeightedAverage() -> Decimal? {
        // Implementation of weighted average calculation
        // Would use historical TDD data from Core Data
        nil
    }
}

/// Extension for rounding Decimal numbers
extension Decimal {
    /// Rounds a decimal to specified number of places
    /// - Parameter places: Number of decimal places
    /// - Returns: Rounded decimal
    func rounded(toPlaces places: Int) -> Decimal {
        var value = self
        var result = Decimal()
        NSDecimalRound(&result, &value, places, .plain)
        return result
    }
}