src/utils/metricAcwr.ts

import type { Activity, MetricACWRData } from '@/types';

/**
 * Format date to YYYY-MM-DD in local timezone (not UTC)
 */
function formatDateLocal(date: Date): string {
    const year = date.getFullYear();
    const month = String(date.getMonth() + 1).padStart(2, '0');
    const day = String(date.getDate()).padStart(2, '0');
    return `${year}-${month}-${day}`;
}

/**
 * Calculate optimal activity values for the next 7 days to minimize MSE between ACWR and target
 * Uses optimization to find values that maintain ACWR close to target across all predicted days
 * @param allDates - Array of all historical dates
 * @param dailyValues - Map of date strings to activity values
 * @param targetACWR - Target ACWR to achieve
 * @param numFutureDays - Number of future days to predict (default 7)
 * @param startOffset - Day offset to start predictions (0 for today, 1 for tomorrow)
 */
function calculateOptimalFutureDays(
    allDates: Date[],
    dailyValues: Map<string, number>,
    targetACWR: number,
    numFutureDays: number = 7,
    startOffset: number = 1
): {
    futureDates: string[];
    futureValues: number[];
    futureAverage7d: number[];
    futureAverage28d: number[];
    futureAcwr: number[];
    allFutureDates: string[];
    allFutureValues: number[];
    allFutureAverage7d: number[];
    allFutureAverage28d: number[];
    allFutureAcwr: number[];
} {
    // Variance weight parameter (0-1): controls trade-off between ACWR accuracy and value consistency
    // Higher weight = more preference for consistent daily values
    const varianceWeight = 0.5;

    // Helper function to calculate variance of predicted values
    const calculateVariance = (values: number[]): number => {
        if (values.length === 0) return 0;
        const mean = values.reduce((sum, v) => sum + v, 0) / values.length;
        const variance = values.reduce((sum, v) => sum + Math.pow(v - mean, 2), 0) / values.length;
        return variance;
    };

    // Helper function to simulate predictions and calculate error (MSE + variance penalty)
    const simulateAndCalculateError = (values: number[]): { error: number; acwrValues: number[]; variance: number } => {
        const simulatedDailyValues = new Map(dailyValues);
        const simulatedAllDates = [...allDates];
        const lastDate = new Date(allDates[allDates.length - 1]);
        const acwrValues: number[] = [];

        for (let dayOffset = 0; dayOffset < numFutureDays; dayOffset++) {
            const futureDate = new Date(lastDate);
            futureDate.setDate(lastDate.getDate() + 1 + startOffset + dayOffset);
            const futureDateStr = formatDateLocal(futureDate);

            simulatedDailyValues.set(futureDateStr, values[dayOffset]);
            simulatedAllDates.push(futureDate);

            const currentIndex = simulatedAllDates.length - 1;

            // Calculate 7-day average
            let acuteSum = 0;
            for (let i = Math.max(0, currentIndex - 6); i <= currentIndex; i++) {
                const checkDateStr = formatDateLocal(simulatedAllDates[i]);
                acuteSum += simulatedDailyValues.get(checkDateStr) || 0;
            }
            const acuteAvg = acuteSum / 7;

            // Calculate 28-day average
            let chronicSum = 0;
            for (let i = Math.max(0, currentIndex - 27); i <= currentIndex; i++) {
                const checkDateStr = formatDateLocal(simulatedAllDates[i]);
                chronicSum += simulatedDailyValues.get(checkDateStr) || 0;
            }
            const chronicAvg = chronicSum / 28;

            const acwr = chronicAvg > 0 ? acuteAvg / chronicAvg : 0;
            acwrValues.push(acwr);
        }

        // Calculate MSE for ACWR
        let mse = 0;
        for (const acwr of acwrValues) {
            const diff = acwr - targetACWR;
            mse += diff * diff;
        }
        mse /= numFutureDays;

        // Calculate variance of predicted values
        const variance = calculateVariance(values);

        // Normalize variance by mean to make it scale-independent
        const mean = values.reduce((sum, v) => sum + v, 0) / values.length;
        const normalizedVariance = mean > 0 ? variance / (mean * mean) : 0;

        // Combined error: MSE + weighted variance penalty
        const error = mse + varianceWeight * normalizedVariance;

        return { error, acwrValues, variance };
    };

    // Initialize with greedy solution as starting point
    const initialValues: number[] = [];
    const simulatedDailyValues = new Map(dailyValues);
    const simulatedAllDates = [...allDates];
    const lastDate = new Date(allDates[allDates.length - 1]);

    for (let dayOffset = 0; dayOffset < numFutureDays; dayOffset++) {
        const futureDate = new Date(lastDate);
        futureDate.setDate(lastDate.getDate() + 1 + startOffset + dayOffset);
        const futureDateStr = formatDateLocal(futureDate);

        simulatedAllDates.push(futureDate);
        const currentIndex = simulatedAllDates.length - 1;

        // Calculate sums for greedy solution
        let acuteSum = 0;
        for (let i = Math.max(0, currentIndex - 6); i < currentIndex; i++) {
            const checkDateStr = formatDateLocal(simulatedAllDates[i]);
            acuteSum += simulatedDailyValues.get(checkDateStr) || 0;
        }

        let chronicSum = 0;
        for (let i = Math.max(0, currentIndex - 27); i < currentIndex; i++) {
            const checkDateStr = formatDateLocal(simulatedAllDates[i]);
            chronicSum += simulatedDailyValues.get(checkDateStr) || 0;
        }

        // Greedy optimal value
        const numerator = 4 * acuteSum - targetACWR * chronicSum;
        const denominator = targetACWR - 4;
        let optimalValue = 0;

        if (Math.abs(denominator) > 0.001) {
            optimalValue = numerator / denominator;
            if (optimalValue < 0) optimalValue = 0;
        } else {
            optimalValue = acuteSum / 7;
        }

        initialValues.push(optimalValue);
        simulatedDailyValues.set(futureDateStr, optimalValue);
    }

    // Iterative refinement to minimize error (MSE + variance penalty)
    let bestValues = [...initialValues];
    let bestError = simulateAndCalculateError(bestValues).error;

    // Simple gradient descent-like optimization
    const iterations = 50;
    const learningRate = 0.3;

    for (let iter = 0; iter < iterations; iter++) {
        const currentResult = simulateAndCalculateError(bestValues);

        // Try adjusting each day's value
        for (let dayIdx = 0; dayIdx < numFutureDays; dayIdx++) {
            const originalValue = bestValues[dayIdx];

            // Calculate gradient by finite difference
            const delta = Math.max(1, originalValue * 0.1);
            const testValues = [...bestValues];
            testValues[dayIdx] = originalValue + delta;
            const upResult = simulateAndCalculateError(testValues);

            const gradient = (upResult.error - currentResult.error) / delta;

            // Update value using gradient
            let newValue = originalValue - learningRate * gradient * originalValue;
            newValue = Math.max(0, newValue); // Ensure non-negative

            // Test if this improves error
            testValues[dayIdx] = newValue;
            const newResult = simulateAndCalculateError(testValues);

            if (newResult.error < bestError) {
                bestValues[dayIdx] = newValue;
                bestError = newResult.error;
            }
        }
    }

    // Build final results with optimized values
    // Filter out insignificant values (< 1) - likely rest days
    const MIN_THRESHOLD = 1;

    const futureDates: string[] = [];
    const futureValues: number[] = [];
    const futureAverage7d: number[] = [];
    const futureAverage28d: number[] = [];
    const futureAcwr: number[] = [];

    const finalSimulatedDailyValues = new Map(dailyValues);
    const finalSimulatedAllDates = [...allDates];
    
    // We need to track ALL dates and their ACWR for curve continuity
    const allFutureDates: string[] = [];
    const allFutureValues: number[] = [];
    const allFutureAverage7d: number[] = [];
    const allFutureAverage28d: number[] = [];
    const allFutureAcwr: number[] = [];

    for (let dayOffset = 0; dayOffset < numFutureDays; dayOffset++) {
        const futureDate = new Date(lastDate);
        futureDate.setDate(lastDate.getDate() + 1 + startOffset + dayOffset);
        const futureDateStr = formatDateLocal(futureDate);
        
        // Store ALL dates for ACWR calculation
        allFutureDates.push(futureDateStr);
        allFutureValues.push(bestValues[dayOffset]);

        // Only include days with significant values (>= MIN_THRESHOLD) in the filtered arrays
        if (bestValues[dayOffset] >= MIN_THRESHOLD) {
            futureDates.push(futureDateStr);
            futureValues.push(bestValues[dayOffset]);
        }

        finalSimulatedDailyValues.set(futureDateStr, bestValues[dayOffset]);
        finalSimulatedAllDates.push(futureDate);

        const currentIndex = finalSimulatedAllDates.length - 1;

        // Calculate averages and ACWR for ALL days to maintain curve continuity
        // Calculate 7-day average
        let newAcuteSum = 0;
        for (let i = Math.max(0, currentIndex - 6); i <= currentIndex; i++) {
            const checkDateStr = formatDateLocal(finalSimulatedAllDates[i]);
            newAcuteSum += finalSimulatedDailyValues.get(checkDateStr) || 0;
        }
        const newAcuteAvg = newAcuteSum / 7;
        allFutureAverage7d.push(newAcuteAvg);

        // Calculate 28-day average
        let newChronicSum = 0;
        for (let i = Math.max(0, currentIndex - 27); i <= currentIndex; i++) {
            const checkDateStr = formatDateLocal(finalSimulatedAllDates[i]);
            newChronicSum += finalSimulatedDailyValues.get(checkDateStr) || 0;
        }
        const newChronicAvg = newChronicSum / 28;
        allFutureAverage28d.push(newChronicAvg);

        // Calculate ACWR
        const newACWR = newChronicAvg > 0 ? newAcuteAvg / newChronicAvg : 0;
        allFutureAcwr.push(newACWR);
    }

    return {
        futureDates,
        futureValues,
        futureAverage7d,
        futureAverage28d,
        futureAcwr,
        // Include all dates and ACWR for curve continuity
        allFutureDates,
        allFutureValues,
        allFutureAverage7d,
        allFutureAverage28d,
        allFutureAcwr,
    };
}

/**
 * Calculate ACWR for a specific metric (distance, duration, or TSS)
 * @param activities - Array of activities
 * @param metricExtractor - Function to extract the metric value from an activity
 * @param dateRange - Optional date range to maintain consistency
 * @param targetACWR - Target ACWR value for predictions
 */
export function calculateMetricACWR(
    activities: Activity[],
    metricExtractor: (activity: Activity) => number | undefined,
    dateRange?: { start: Date; end: Date },
    targetACWR: number = 1.3
): MetricACWRData {
    if (activities.length === 0 && !dateRange) {
        return {
            dates: [],
            values: [],
            average7d: [],
            average28d: [],
            acwr: [],
        };
    }

    // Sort activities by date
    const sortedActivities = [...activities].sort((a, b) => a.date.getTime() - b.date.getTime());

    // Get date range
    const startDate = dateRange?.start || (sortedActivities.length > 0 ? new Date(sortedActivities[0].date) : new Date());
    const endDate = dateRange?.end || (sortedActivities.length > 0 ? new Date(sortedActivities[sortedActivities.length - 1].date) : new Date());

    // Create a map of date -> daily sum
    const dailyValues = new Map<string, number>();
    // Create a map of date -> activities
    const activitiesByDate = new Map<string, Activity[]>();

    sortedActivities.forEach(activity => {
        const dateStr = formatDateLocal(activity.date);
        const value = metricExtractor(activity);
        if (value !== undefined) {
            dailyValues.set(dateStr, (dailyValues.get(dateStr) || 0) + value);
        }

        // Group activities by date
        if (!activitiesByDate.has(dateStr)) {
            activitiesByDate.set(dateStr, []);
        }
        activitiesByDate.get(dateStr)!.push(activity);
    });

    const dates: string[] = [];
    const values: (number | null)[] = [];
    const average7d: (number | null)[] = [];
    const average28d: (number | null)[] = [];
    const acwr: (number | null)[] = [];

    // Generate all dates in range
    const allDates: Date[] = [];
    for (let d = new Date(startDate); d <= endDate; d.setDate(d.getDate() + 1)) {
        allDates.push(new Date(d));
    }

    // Calculate for each date
    allDates.forEach((date, index) => {
        const dateStr = formatDateLocal(date);
        dates.push(dateStr);

        // Get daily value
        const dailyValue = dailyValues.get(dateStr) || null;
        values.push(dailyValue);

        // Calculate 7-day rolling average (acute load)
        let acuteSum = 0;
        let acuteCount = 0;
        for (let i = Math.max(0, index - 6); i <= index; i++) {
            const checkDateStr = formatDateLocal(allDates[i]);
            const val = dailyValues.get(checkDateStr);
            if (val !== undefined) {
                acuteSum += val;
                acuteCount++;
            }
        }
        const acuteAvg = acuteCount > 0 ? acuteSum / 7 : null;
        average7d.push(acuteAvg);

        // Calculate 28-day rolling average
        let chronicSum = 0;
        let chronicCount = 0;
        for (let i = Math.max(0, index - 27); i <= index; i++) {
            const checkDateStr = formatDateLocal(allDates[i]);
            const val = dailyValues.get(checkDateStr);
            if (val !== undefined) {
                chronicSum += val;
                chronicCount++;
            }
        }
        const chronicAvg = chronicCount > 0 ? chronicSum / 28 : null;
        average28d.push(chronicAvg);

        // Calculate ACWR only after we have at least 28 days of data range
        // This ensures accuracy by having a full chronic period (28 days)
        if (index >= 27 && acuteAvg !== null && chronicAvg !== null && chronicAvg > 0) {
            acwr.push(acuteAvg / chronicAvg);
        } else {
            acwr.push(null);
        }
    });    // Calculate tomorrow's required value to reach ACWR
    let targetTomorrowValue: number | null = null;

    if (allDates.length >= 28) {

        // Calculate tomorrow's 7-day acute sum (last 6 days including zeros)
        let tomorrowAcuteSum = 0;
        for (let i = allDates.length - 6; i < allDates.length; i++) {
            const checkDateStr = formatDateLocal(allDates[i]);
            const val = dailyValues.get(checkDateStr) || 0; // Include 0 for rest days
            tomorrowAcuteSum += val;
        }

        // Calculate tomorrow's 28-day chronic sum (last 27 days including zeros)
        let tomorrowChronicSum = 0;
        for (let i = allDates.length - 27; i < allDates.length; i++) {
            const checkDateStr = formatDateLocal(allDates[i]);
            const val = dailyValues.get(checkDateStr) || 0; // Include 0 for rest days
            tomorrowChronicSum += val;
        }

        // Solve for tomorrow's value (X):
        // Tomorrow's ACWR = [(tomorrowAcuteSum + X) / 7] / [(tomorrowChronicSum + X) / 28]
        // Simplifies to: targetACWR = (tomorrowAcuteSum + X) / (tomorrowChronicSum + X) * 28 / 7
        // targetACWR = (tomorrowAcuteSum + X) / (tomorrowChronicSum + X) * 4
        // targetACWR * (tomorrowChronicSum + X) = 4 * (tomorrowAcuteSum + X)
        // targetACWR * tomorrowChronicSum + targetACWR * X = 4 * tomorrowAcuteSum + 4 * X
        // targetACWR * X - 4 * X = 4 * tomorrowAcuteSum - targetACWR * tomorrowChronicSum
        // X * (targetACWR - 4) = 4 * tomorrowAcuteSum - targetACWR * tomorrowChronicSum
        // X = (4 * tomorrowAcuteSum - targetACWR * tomorrowChronicSum) / (targetACWR - 4)

        const numerator = 4 * tomorrowAcuteSum - targetACWR * tomorrowChronicSum;
        const denominator = targetACWR - 4;

        if (Math.abs(denominator) > 0.001) {
            targetTomorrowValue = numerator / denominator;
            // If negative, set to 0 (rest day recommended)
            if (targetTomorrowValue < 0) {
                targetTomorrowValue = 0;
            }
        } else {
            // Special case: targetACWR ≈ 4, need different approach
            // If ACWR = 4, then acute = 4 * chronic, which means you need massive increase
            targetTomorrowValue = null;
        }
    }

    // Calculate what ACWR would be with a rest day tomorrow
    let restTomorrowACWR: number | null = null;

    if (allDates.length >= 28) {
        // Calculate tomorrow's 7-day acute sum with rest (last 6 days + 0)
        let tomorrowAcuteSum = 0;
        for (let i = allDates.length - 6; i < allDates.length; i++) {
            const checkDateStr = formatDateLocal(allDates[i]);
            const val = dailyValues.get(checkDateStr) || 0;
            tomorrowAcuteSum += val;
        }
        // Add 0 for rest day (no need to add)

        // Calculate tomorrow's 28-day chronic sum with rest (last 27 days + 0)
        let tomorrowChronicSum = 0;
        for (let i = allDates.length - 27; i < allDates.length; i++) {
            const checkDateStr = formatDateLocal(allDates[i]);
            const val = dailyValues.get(checkDateStr) || 0;
            tomorrowChronicSum += val;
        }
        // Add 0 for rest day (no need to add)

        // Calculate ACWR with rest day
        const tomorrowAcuteAvg = tomorrowAcuteSum / 7;
        const tomorrowChronicAvg = tomorrowChronicSum / 28;

        if (tomorrowChronicAvg > 0) {
            restTomorrowACWR = tomorrowAcuteAvg / tomorrowChronicAvg;
        }
    }

    return {
        dates,
        values,
        average7d,
        average28d,
        acwr,
        targetTomorrowValue,
        targetACWR: allDates.length >= 28 ? targetACWR : undefined,
        restTomorrowACWR: allDates.length >= 28 ? restTomorrowACWR : undefined,
        todayValue: undefined,
        activitiesByDate,
        // Add future predictions if we have enough data
        // Predictions always start from the day after lastDate
        ...(allDates.length >= 28 ? calculateOptimalFutureDays(allDates, dailyValues, targetACWR, 7, 0) : {}),
    };
}