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 reach and maintain target ACWR
 * Uses a greedy algorithm that optimizes each day to get closer to target ACWR
 * @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[];
} {
    const futureDates: string[] = [];
    const futureValues: number[] = [];
    const futureAverage7d: number[] = [];
    const futureAverage28d: number[] = [];
    const futureAcwr: number[] = [];

    // Create a mutable copy of daily values to simulate future
    const simulatedDailyValues = new Map(dailyValues);
    const simulatedAllDates = [...allDates];

    // Get the last date in the data
    const lastDate = new Date(allDates[allDates.length - 1]);

    for (let dayOffset = 0; dayOffset < numFutureDays; dayOffset++) {
        const futureDate = new Date(lastDate);
        // Add 1 + startOffset + dayOffset to go beyond lastDate
        // startOffset=0: start from day after lastDate (tomorrow if lastDate is today)
        // startOffset=1: skip one day, start from 2 days after lastDate
        futureDate.setDate(lastDate.getDate() + 1 + startOffset + dayOffset);
        const futureDateStr = formatDateLocal(futureDate);

        futureDates.push(futureDateStr);
        simulatedAllDates.push(futureDate);

        const currentIndex = simulatedAllDates.length - 1;

        // Calculate what the 7-day sum would be (last 6 days + today)
        let acuteSum = 0;
        for (let i = Math.max(0, currentIndex - 6); i < currentIndex; i++) {
            const checkDateStr = formatDateLocal(simulatedAllDates[i]);
            acuteSum += simulatedDailyValues.get(checkDateStr) || 0;
        }

        // Calculate what the 28-day sum would be (last 27 days + today)
        let chronicSum = 0;
        for (let i = Math.max(0, currentIndex - 27); i < currentIndex; i++) {
            const checkDateStr = formatDateLocal(simulatedAllDates[i]);
            chronicSum += simulatedDailyValues.get(checkDateStr) || 0;
        }

        // Calculate optimal value for today using the formula:
        // targetACWR = [(acuteSum + X) / 7] / [(chronicSum + X) / 28]
        // Solving for X:
        // X = (4 * acuteSum - targetACWR * chronicSum) / (targetACWR - 4)

        const numerator = 4 * acuteSum - targetACWR * chronicSum;
        const denominator = targetACWR - 4;

        let optimalValue = 0;

        if (Math.abs(denominator) > 0.001) {
            optimalValue = numerator / denominator;

            // Ensure non-negative values (can't have negative activity)
            if (optimalValue < 0) {
                optimalValue = 0;
            }
        } else {
            // Special case: when targetACWR ≈ 4, we need a different approach
            // In this case, set to average of recent values to maintain stability
            const recentSum = acuteSum;
            optimalValue = recentSum / 7;
        }

        // Store the optimal value for this future day
        futureValues.push(optimalValue);
        simulatedDailyValues.set(futureDateStr, optimalValue);

        // Calculate the resulting 7-day average (acute load)
        let newAcuteSum = 0;
        for (let i = Math.max(0, currentIndex - 6); i <= currentIndex; i++) {
            const checkDateStr = formatDateLocal(simulatedAllDates[i]);
            newAcuteSum += simulatedDailyValues.get(checkDateStr) || 0;
        }
        const newAcuteAvg = newAcuteSum / 7;
        futureAverage7d.push(newAcuteAvg);

        // Calculate the resulting 28-day average (chronic load)
        let newChronicSum = 0;
        for (let i = Math.max(0, currentIndex - 27); i <= currentIndex; i++) {
            const checkDateStr = formatDateLocal(simulatedAllDates[i]);
            newChronicSum += simulatedDailyValues.get(checkDateStr) || 0;
        }
        const newChronicAvg = newChronicSum / 28;
        futureAverage28d.push(newChronicAvg);

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

    return {
        futureDates,
        futureValues,
        futureAverage7d,
        futureAverage28d,
        futureAcwr,
    };
}

/**
 * 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
 * @param predictToday - If true, include today in predictions; if false, predictions start tomorrow
 */
export function calculateMetricACWR(
    activities: Activity[],
    metricExtractor: (activity: Activity) => number | undefined,
    dateRange?: { start: Date; end: Date },
    targetACWR: number = 1.3,
    predictToday: boolean = false
): 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 (need at least 28 days of data)
        if (index < 27) {
            acwr.push(null);
            return;
        }

        // For ACWR, use the full 28-day sum (not average)
        let acwrChronicSum = 0;
        for (let i = index - 27; i <= index; i++) {
            const checkDateStr = formatDateLocal(allDates[i]);
            const val = dailyValues.get(checkDateStr);
            if (val !== undefined) {
                acwrChronicSum += val;
            }
        }
        const acwrChronicAvg = acwrChronicSum / 28;

        // Calculate ACWR
        if (acwrChronicAvg > 0 && acuteAvg !== null) {
            acwr.push(acuteAvg / acwrChronicAvg);
        } 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 and target ACWR < 1.0, set to 0 (rest day recommended)
            // If negative and target ACWR >= 1.0, set to null (target unreachable)
            if (targetTomorrowValue < 0) {
                if (targetACWR < 1.0) {
                    targetTomorrowValue = 0;
                } else {
                    targetTomorrowValue = null;
                }
            }
        } 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
        // Determine start offset based on predictToday and whether today has activities
        ...(allDates.length >= 28 ? (() => {
            const today = new Date();
            const todayStr = formatDateLocal(today);
            const lastDateStr = formatDateLocal(allDates[allDates.length - 1]);
            
            // Check if today is the last date in our data and if it has any activities
            const todayHasActivities = todayStr === lastDateStr && dailyValues.has(todayStr);
            
            // startOffset determines which day to start predictions from relative to lastDate
            // startOffset=0: predictions start from lastDate + 1 (next day after last data)
            // startOffset=1: predictions start from lastDate + 2 (skip one day)
            // predictToday unchecked: start from next day (offset 0)
            // predictToday checked and today has activities: start from next day (offset 0) 
            // predictToday checked and today has no activities: start from next day (offset 0), but lastDate was yesterday
            const startOffset = 0;
            
            return calculateOptimalFutureDays(allDates, dailyValues, targetACWR, 7, startOffset);
        })() : {}),
    };
}