import * as math from 'npm:mathjs'

/**
 * Performs non-linear least squares regression of data to an arbitrary function using math.js.
 *
 * @param {number[]} xData - Array of x-coordinates.
 * @param {number[]} yData - Array of y-coordinates.
 * @param {(x: number, params: number[]) => number} func - The function to fit to the data. Takes x and an array of parameters as input.
 * @param {number[]} initialParams - Initial guesses for the parameters.
 * @param {number} maxIterations - Maximum number of iterations.
 * @param {number} tolerance - Tolerance for convergence.
 * @returns {{ params: number[], residuals: number[], rSquared: number, iterations: number, converged: boolean }} - An object containing the fitted parameters and other information.
 */
export function curveFit<T extends number[]>(
	xData: number[],
	yData: number[],
	func: (x: number, params: T) => number,
	initialParams: T,
	maxIterations: number = 1000,
	tolerance: number = 1e-8,
): {
	params: T
	residuals: number[]
	rSquared: number
	iterations: number
	converged: boolean
} {
	const numPoints: number = xData.length
	const numParams: number = initialParams.length

	if (numPoints !== yData.length) {
		throw new Error('xData and yData must have the same length.')
	}

	const params: number[] = [...initialParams]
	let iterations: number = 0
	let converged: boolean = false

	while (iterations < maxIterations) {
		const jacobian: number[][] = []
		const residuals: number[] = []

		for (let i: number = 0; i < numPoints; i++) {
			const x: number = xData[i]
			const y: number = yData[i]
			const predicted: number = func(x, params as T)
			const residual: number = y - predicted
			residuals.push(residual)

			const row: number[] = []
			for (let j: number = 0; j < numParams; j++) {
				const delta: number = 1e-6
				const paramsPlusDelta = [...params] as T
				paramsPlusDelta[j] += delta
				const predictedPlusDelta: number = func(x, paramsPlusDelta)
				const derivative: number = (predictedPlusDelta - predicted) / delta
				row.push(derivative)
			}
			jacobian.push(row)
		}

		const jacobianTranspose: number[][] = math.transpose(jacobian)
		const hessian: math.Matrix = math.matrix(
			math.multiply(jacobianTranspose, jacobian),
		)
		const gradient: math.Matrix = math.matrix(
			math.multiply(jacobianTranspose, residuals),
		)

		try {
			const deltaParams: math.Matrix = math.usolve(
				hessian,
				gradient,
			) as math.Matrix

			let maxDelta: number = 0
			for (let i: number = 0; i < numParams; i++) {
				params[i] += deltaParams.get([i, 0])
				maxDelta = Math.max(maxDelta, Math.abs(deltaParams.get([i, 0])))
			}

			if (maxDelta < tolerance) {
				converged = true
				break
			}
		} catch (_error) {
			// console.error('Matrix solve failed:', error)
			return {
				params: params as T,
				residuals: residuals,
				rSquared: NaN,
				iterations: iterations,
				converged: false,
			}
		}

		iterations++
	}

	const calculatedResiduals: number[] = xData.map((x: number, i: number) =>
		yData[i] - func(x, params as T)
	)
	const ssRes: number = calculatedResiduals.reduce(
		(sum: number, r: number) => sum + r * r,
		0,
	)
	const yMean: number = yData.reduce((sum: number, y: number) => sum + y, 0) /
		numPoints
	const ssTot: number = yData.reduce(
		(sum: number, y: number) => sum + (y - yMean) * (y - yMean),
		0,
	)
	const rSquared: number = 1 - ssRes / ssTot

	return {
		params: params as T,
		residuals: calculatedResiduals,
		rSquared: rSquared,
		iterations: iterations,
		converged: converged,
	}
}