/**
 * earth-now.co
 * Author:  Simon-Pierre Boucher
 * Contact: contact@spboucher.ai
 * File:    packages/models/src/fourier.ts
 * Purpose: Least-squares fit of trend + Fourier harmonics in the VALUE domain, converted to rate harmonics for CounterModel
 */

import {
  type Harmonic,
  DAY_SECONDS,
  SEASONAL_EPOCH_MS,
  WEEK_SECONDS,
  YEAR_SECONDS,
  parseIsoUtc,
} from "@earth-now/counter";
import { leastSquares } from "./linalg.js";

export interface FourierFitInput {
  /** Observations: ISO time → value. Needs >= 2·harmonics + 2 points. */
  observations: Array<{ time: string; value: number }>;
  /** Harmonic spec to fit, e.g. [{period:"year", order:1}, {period:"year", order:2}]. */
  harmonics: Array<{ period: "year" | "week" | "day"; order: number }>;
}

export interface FourierFitResult {
  /** Value at the seasonal epoch (intercept). */
  intercept: number;
  /** Linear trend of the value, per second. Becomes rateFn.base. */
  trendPerSecond: number;
  /** Rate-domain harmonics ready for a `seasonal` RateFunction. */
  rateHarmonics: Harmonic[];
  /** Evaluate the fitted VALUE curve at t (ms). */
  valueAt: (tMs: number) => number;
  /** Root-mean-square residual of the fit. */
  rmse: number;
}

function omegaOf(period: "year" | "week" | "day", order: number): number {
  const seconds =
    period === "year" ? YEAR_SECONDS : period === "week" ? WEEK_SECONDS : DAY_SECONDS;
  return (2 * Math.PI * order) / seconds;
}

/**
 * Fit value(τ) = a + b·τ + Σ_h [ c_h·cos(ω_h τ) + s_h·sin(ω_h τ) ], τ in seconds
 * since the seasonal epoch. The derivative of each fitted value harmonic is an
 * exact rate harmonic, so the returned Harmonics reproduce the fitted curve when
 * integrated by packages/counter — the client animates the very curve we fitted.
 */
export function fitFourier(input: FourierFitInput): FourierFitResult {
  const { observations, harmonics } = input;
  const k = 2 + 2 * harmonics.length;
  if (observations.length < k) {
    throw new Error(
      `fitFourier: need at least ${k} observations for ${harmonics.length} harmonics, got ${observations.length}`,
    );
  }

  const taus = observations.map((o) => (parseIsoUtc(o.time) - SEASONAL_EPOCH_MS) / 1000);
  // Center τ for conditioning; un-center the intercept afterwards.
  const tauMean = taus.reduce((s, t) => s + t, 0) / taus.length;
  const omegas = harmonics.map((h) => omegaOf(h.period, h.order));

  const X = taus.map((tau) => {
    const row = [1, tau - tauMean];
    for (const omega of omegas) row.push(Math.cos(omega * tau), Math.sin(omega * tau));
    return row;
  });
  const y = observations.map((o) => o.value);
  const beta = leastSquares(X, y);

  const b = beta[1]!;
  const a = beta[0]! - b * tauMean;
  const coefs = harmonics.map((h, i) => ({
    ...h,
    c: beta[2 + 2 * i]!,
    s: beta[3 + 2 * i]!,
    omega: omegas[i]!,
  }));

  const valueAt = (tMs: number): number => {
    const tau = (tMs - SEASONAL_EPOCH_MS) / 1000;
    let v = a + b * tau;
    for (const { c, s, omega } of coefs) v += c * Math.cos(omega * tau) + s * Math.sin(omega * tau);
    return v;
  };

  let sse = 0;
  for (let i = 0; i < taus.length; i++) {
    const r = valueAt(parseIsoUtc(observations[i]!.time)) - y[i]!;
    sse += r * r;
  }

  // d/dτ [c·cos(ωτ) + s·sin(ωτ)] = (s·ω)·cos(ωτ) − (c·ω)·sin(ωτ)
  //                              = A·ω·cos(ωτ + φ) with A = √(c²+s²), φ = atan2(c, s)... derived below.
  const rateHarmonics: Harmonic[] = coefs.map(({ period, order, c, s, omega }) => {
    const amplitude = Math.hypot(c, s) * omega;
    // Write derivative as R·cos(ωτ + φ): R·cosφ = s·ω, R·sinφ = c·ω ⇒ φ = atan2(c, s).
    const phase = Math.atan2(c, s);
    return { period, order, amplitude, phase };
  });

  return {
    intercept: a,
    trendPerSecond: b,
    rateHarmonics,
    valueAt,
    rmse: Math.sqrt(sse / taus.length),
  };
}