package review

import (
    "math"
    "time"

    "gonum.org/v1/gonum/optimize"
)

// PowerLawFunc models weight as a function of reps: w = a * reps^b
func PowerLawFunc(a, b, reps float64) float64 {
    return a * math.Pow(reps, b)
}

// WeightedResiduals computes weighted residuals for curve fitting
func WeightedResiduals(params []float64, weight, reps []float64, dates []time.Time, now time.Time, halfLifeDays float64) float64 {
    a, b := params[0], params[1]
    var sum float64
    for i := range weight {
        // Exponential time decay weighting
        daysAgo := now.Sub(dates[i]).Hours() / 24
        weightDecay := math.Exp(-math.Ln2 * daysAgo / halfLifeDays) // Half-life decay
        predicted := PowerLawFunc(a, b, reps[i])
        residual := weight[i] - predicted
        sum += weightDecay * residual * residual
    }
    return sum
}

// FitPowerLaw fits the power law curve with time weighting
func FitPowerLaw(weight, reps []float64, dates []time.Time, halfLifeDays float64) (a, b float64) {
    now := time.Now()
    // Initial guess: a = max(weight), b = -0.1
    params := []float64{max(weight), -0.1}
    problem := optimize.Problem{
        Func: func(x []float64) float64 {
            return WeightedResiduals(x, weight, reps, dates, now, halfLifeDays)
        },
    }
    result, err := optimize.Minimize(problem, params, nil, nil)
    if err != nil {
        panic(err)
    }
    return result.X[0], result.X[1]
}

// max returns the maximum value in a slice
func max(slice []float64) float64 {
    m := slice[0]
    for _, v := range slice {
        if v > m {
            m = v
        }
    }
    return m
}

// Estimate1RM estimates the current 1RM (reps=1) using fitted parameters
func Estimate1RM(a, b float64) float64 {
    return PowerLawFunc(a, b, 1)
}

// EstimateReps returns the predicted number of reps at a given weight
func EstimateReps(a, b, targetWeight float64) float64 {
    // Avoid division by zero or negative exponent issues
    if a == 0 || b == 0 {
        return 0
    }
    return math.Pow(targetWeight/a, 1/b)
}

// EstimateMaxWeight returns the predicted max weight for a given number of reps
func EstimateMaxWeight(a, b, nReps float64) float64 {
    return a * math.Pow(nReps, b)
}