// UQO Éval — modélisation hédonique : estimateurs (MCO, Huber, LAD, ridge, LASSO, SAR-2SLS) et covariances.
import { cholesky, cholInverse, cholSolve, gram, hcat, matVec, solveSPD, takeCols, zeros, type Mat } from "./linalg";
import { lag, KnnIndex } from "./spatial";
import type { SeType } from "./spec";

export interface Fit {
  beta: Float64Array;
  fitted: Float64Array;
  resid: Float64Array;
  /** (X'WX)^-1 — « pain » du sandwich ; null pour les méthodes sans inférence. */
  bread: Float64Array | null;
  /** Matrice de design effectivement utilisée pour l'inférence (X ou X̂ en 2SLS). */
  Xinf: Mat | null;
  weights: Float64Array | null;
  sigma2: number;
  df: number;
  k: number;
  extra: Record<string, number>;
}

function residuals(X: Mat, y: Float64Array, beta: Float64Array) {
  const fitted = matVec(X, beta);
  const resid = new Float64Array(X.n);
  for (let i = 0; i < X.n; i++) resid[i] = y[i] - fitted[i];
  return { fitted, resid };
}

/** MCO (ou MCP si poids). */
export function ols(X: Mat, y: Float64Array, weights?: Float64Array): Fit {
  const { G, c } = gram(X, y, weights);
  const { x: beta, inv } = solveSPD(G, X.k, c);
  const { fitted, resid } = residuals(X, y, beta);
  let ss = 0;
  for (let i = 0; i < X.n; i++) ss += (weights ? weights[i] : 1) * resid[i] * resid[i];
  const df = X.n - X.k;
  return { beta, fitted, resid, bread: inv, Xinf: X, weights: weights ?? null, sigma2: ss / Math.max(1, df), df, k: X.k, extra: {} };
}

/** Covariance des coefficients : classique, HC1 ou grappes. */
export function covariance(fit: Fit, type: SeType, clusters?: Int32Array, nClusters = 0): Float64Array | null {
  if (!fit.bread || !fit.Xinf) return null;
  const X = fit.Xinf;
  const { n, k, a } = X;
  const B = fit.bread;
  const w = fit.weights;
  if (type === "classic") {
    const V = new Float64Array(k * k);
    for (let i = 0; i < k * k; i++) V[i] = fit.sigma2 * B[i];
    return V;
  }
  // viande : Σ_g s_g s_g' avec s_g = Σ_{i∈g} w_i e_i x_i (HC1 : chaque obs = une grappe)
  const M = new Float64Array(k * k);
  const useCl = type === "cluster" && clusters && nClusters > 1;
  const nG = useCl ? nClusters : n;
  const S = useCl ? new Float64Array(nG * k) : null;
  for (let i = 0; i < n; i++) {
    const s = (w ? w[i] : 1) * fit.resid[i];
    const off = i * k;
    if (useCl && S) {
      const g = clusters![i] * k;
      for (let j = 0; j < k; j++) S[g + j] += s * a[off + j];
    } else {
      for (let j = 0; j < k; j++) {
        const sj = s * a[off + j];
        if (sj === 0) continue;
        for (let l = 0; l < k; l++) M[j * k + l] += sj * s * a[off + l];
      }
    }
  }
  if (useCl && S) for (let g = 0; g < nG; g++) for (let j = 0; j < k; j++) for (let l = 0; l < k; l++) M[j * k + l] += S[g * k + j] * S[g * k + l];
  // V = B M B × correction petits échantillons
  const BM = new Float64Array(k * k);
  for (let i = 0; i < k; i++) for (let j = 0; j < k; j++) {
    let s = 0;
    for (let l = 0; l < k; l++) s += B[i * k + l] * M[l * k + j];
    BM[i * k + j] = s;
  }
  const V = new Float64Array(k * k);
  const corr = useCl ? (nG / (nG - 1)) * ((n - 1) / (n - k)) : n / (n - k);
  for (let i = 0; i < k; i++) for (let j = 0; j < k; j++) {
    let s = 0;
    for (let l = 0; l < k; l++) s += BM[i * k + l] * B[l * k + j];
    V[i * k + j] = s * corr;
  }
  return V;
}

const mad = (e: Float64Array) => {
  const s = Float64Array.from(e, Math.abs).sort();
  return s[Math.floor(s.length / 2)] / 0.6745;
};

/** Régression robuste de Huber (IRLS, c = 1,345). */
export function huber(X: Mat, y: Float64Array, iters = 30): Fit {
  let fit = ols(X, y);
  let w = new Float64Array(X.n).fill(1);
  let downweighted = 0;
  for (let it = 0; it < iters; it++) {
    const s = Math.max(1e-9, mad(fit.resid));
    const c = 1.345 * s;
    const w2 = new Float64Array(X.n);
    downweighted = 0;
    for (let i = 0; i < X.n; i++) {
      const ae = Math.abs(fit.resid[i]);
      w2[i] = ae <= c ? 1 : c / ae;
      if (w2[i] < 1) downweighted++;
    }
    const next = ols(X, y, w2);
    let delta = 0;
    for (let j = 0; j < X.k; j++) delta = Math.max(delta, Math.abs(next.beta[j] - fit.beta[j]) / (1e-9 + Math.abs(fit.beta[j])));
    fit = next;
    w = w2;
    if (delta < 1e-6) break;
  }
  fit.extra = { scale: mad(fit.resid), downweighted, downweightedPct: (100 * downweighted) / X.n };
  fit.weights = w;
  return fit;
}

/** Régression médiane (LAD) par IRLS — coefficients seulement. */
export function lad(X: Mat, y: Float64Array, iters = 50): Fit {
  let fit = ols(X, y);
  let sy = 0;
  for (let i = 0; i < X.n; i++) sy += Math.abs(y[i]);
  const eps = 1e-6 * (sy / X.n);
  for (let it = 0; it < iters; it++) {
    const w = new Float64Array(X.n);
    for (let i = 0; i < X.n; i++) w[i] = 1 / Math.max(Math.abs(fit.resid[i]), eps);
    const next = ols(X, y, w);
    let delta = 0;
    for (let j = 0; j < X.k; j++) delta = Math.max(delta, Math.abs(next.beta[j] - fit.beta[j]) / (1e-9 + Math.abs(fit.beta[j])));
    fit = next;
    if (delta < 1e-7) break;
  }
  let sae = 0;
  for (let i = 0; i < X.n; i++) sae += Math.abs(fit.resid[i]);
  return { ...fit, bread: null, Xinf: null, weights: null, extra: { mae: sae / X.n } };
}

/* ------------------------------ pénalisation ------------------------------ */
interface Std {
  mean: Float64Array;
  sd: Float64Array;
  ymean: number;
}
function standardize(X: Mat, y: Float64Array): { Z: Mat; yc: Float64Array; std: Std } {
  const { n, k, a } = X;
  const mean = new Float64Array(k);
  const sd = new Float64Array(k);
  for (let j = 1; j < k; j++) {
    let s = 0;
    for (let i = 0; i < n; i++) s += a[i * k + j];
    mean[j] = s / n;
    let v = 0;
    for (let i = 0; i < n; i++) v += (a[i * k + j] - mean[j]) ** 2;
    sd[j] = Math.sqrt(v / n) || 1;
  }
  const Z = zeros(n, k - 1);
  for (let i = 0; i < n; i++) for (let j = 1; j < k; j++) Z.a[i * (k - 1) + j - 1] = (a[i * k + j] - mean[j]) / sd[j];
  let ym = 0;
  for (let i = 0; i < n; i++) ym += y[i];
  ym /= n;
  const yc = new Float64Array(n);
  for (let i = 0; i < n; i++) yc[i] = y[i] - ym;
  return { Z, yc, std: { mean, sd, ymean: ym } };
}

function destandardize(b: Float64Array, std: Std, k: number): Float64Array {
  const beta = new Float64Array(k);
  let b0 = std.ymean;
  for (let j = 1; j < k; j++) {
    beta[j] = b[j - 1] / std.sd[j];
    b0 -= beta[j] * std.mean[j];
  }
  beta[0] = b0;
  return beta;
}

/** Ridge sur données standardisées : (G/n + λI)^-1 c/n. */
function ridgeStd(G: Float64Array, c: Float64Array, p: number, n: number, lambda: number): Float64Array {
  const A = new Float64Array(p * p);
  const cc = new Float64Array(p);
  for (let i = 0; i < p * p; i++) A[i] = G[i] / n;
  for (let j = 0; j < p; j++) {
    A[j * p + j] += lambda;
    cc[j] = c[j] / n;
  }
  const L = cholesky(A, p, 1e-10);
  if (!L) throw new Error("ridge : matrice singulière");
  return cholSolve(L, p, cc);
}

/** LASSO par descente de coordonnées sur la matrice de Gram (données standardisées). */
function lassoStd(G: Float64Array, c: Float64Array, p: number, n: number, lambda: number, b0?: Float64Array, iters = 200): Float64Array {
  const b = b0 ? Float64Array.from(b0) : new Float64Array(p);
  const soft = (z: number, l: number) => (z > l ? z - l : z < -l ? z + l : 0);
  for (let it = 0; it < iters; it++) {
    let maxd = 0;
    for (let j = 0; j < p; j++) {
      let z = c[j] / n;
      const row = j * p;
      for (let l = 0; l < p; l++) if (l !== j && b[l] !== 0) z -= (G[row + l] / n) * b[l];
      const nb = soft(z, lambda) / (G[row + j] / n || 1);
      maxd = Math.max(maxd, Math.abs(nb - b[j]));
      b[j] = nb;
    }
    if (maxd < 1e-7) break;
  }
  return b;
}

export interface PenalizedResult {
  fit: Fit;
  lambda: number;
  grid: { lambda: number; cvErr: number }[];
  selected: number[]; // indices de colonnes (X) non nulles (LASSO)
  postFit?: Fit; // MCO post-LASSO sur les colonnes retenues
}

export function penalized(X: Mat, y: Float64Array, method: "ridge" | "lasso", folds = 5): PenalizedResult {
  const { Z, yc, std } = standardize(X, y);
  const p = Z.k;
  const n = Z.n;
  const full = gram(Z, yc);
  // grille λ
  let lmax = 0;
  for (let j = 0; j < p; j++) lmax = Math.max(lmax, Math.abs(full.c[j]) / n);
  const grid: number[] = [];
  const nl = 24;
  if (method === "lasso") for (let i = 0; i < nl; i++) grid.push(lmax * Math.pow(1e-3, i / (nl - 1)));
  else for (let i = 0; i < nl; i++) grid.push(10 * Math.pow(1e-5, i / (nl - 1)));
  // validation croisée par blocs (l'échantillon est déjà mélangé)
  const cvErr = new Float64Array(grid.length);
  const foldOf = (i: number) => i % folds;
  for (let f = 0; f < folds; f++) {
    const tr: number[] = [];
    const te: number[] = [];
    for (let i = 0; i < n; i++) (foldOf(i) === f ? te : tr).push(i);
    const Ztr = zeros(tr.length, p);
    const ytr = new Float64Array(tr.length);
    for (let r = 0; r < tr.length; r++) {
      Ztr.a.set(Z.a.subarray(tr[r] * p, tr[r] * p + p), r * p);
      ytr[r] = yc[tr[r]];
    }
    const g = gram(Ztr, ytr);
    let warm: Float64Array | undefined;
    for (let li = 0; li < grid.length; li++) {
      const b = method === "lasso" ? lassoStd(g.G, g.c, p, tr.length, grid[li], warm) : ridgeStd(g.G, g.c, p, tr.length, grid[li]);
      warm = b;
      let se = 0;
      for (const i of te) {
        let pred = 0;
        for (let j = 0; j < p; j++) pred += Z.a[i * p + j] * b[j];
        se += (yc[i] - pred) ** 2;
      }
      cvErr[li] += se / n;
    }
  }
  let best = 0;
  for (let li = 1; li < grid.length; li++) if (cvErr[li] < cvErr[best]) best = li;
  const lambda = grid[best];
  const bStd = method === "lasso" ? lassoStd(full.G, full.c, p, n, lambda, undefined, 500) : ridgeStd(full.G, full.c, p, n, lambda);
  const beta = destandardize(bStd, std, X.k);
  const { fitted, resid } = residuals(X, y, beta);
  let ss = 0;
  for (let i = 0; i < n; i++) ss += resid[i] * resid[i];
  const selected = [0];
  for (let j = 1; j < X.k; j++) if (Math.abs(beta[j]) > 1e-12) selected.push(j);
  const dfEff = method === "lasso" ? selected.length : X.k;
  const fit: Fit = { beta, fitted, resid, bread: null, Xinf: null, weights: null, sigma2: ss / Math.max(1, n - dfEff), df: n - dfEff, k: X.k, extra: { lambda, nonzero: selected.length - 1 } };
  const res: PenalizedResult = { fit, lambda, grid: grid.map((l, i) => ({ lambda: l, cvErr: cvErr[i] })), selected };
  if (method === "lasso" && selected.length > 1) res.postFit = ols(takeCols(X, selected), y);
  return res;
}

/* --------------------------------- SAR 2SLS --------------------------------- */
export interface SarResult {
  fit: Fit; // beta = [ρ, β...] ; Xinf = X̂ = [Ŵy, X]
  rho: number;
  Wy: Float64Array;
  nb: Int32Array[];
  index: KnnIndex;
}

/** y = ρWy + Xβ + ε, instruments Z = [X, WX, W²X] (Kelejian & Prucha, 1998). */
export function sar(X: Mat, y: Float64Array, lat: Float64Array, lng: Float64Array, k = 8): SarResult {
  const index = new KnnIndex(lat, lng);
  const nb = index.selfNeighbours(k);
  const Wy = lag(nb, y);
  const { n, k: p } = X;
  // WX et W²X (hors constante, hors colonnes quasi-constantes)
  const cols: number[] = [];
  for (let j = 1; j < p; j++) cols.push(j);
  const Xnc = takeCols(X, cols);
  const WX = zeros(n, Xnc.k);
  const W2X = zeros(n, Xnc.k);
  for (let j = 0; j < Xnc.k; j++) {
    const col = new Float64Array(n);
    for (let i = 0; i < n; i++) col[i] = Xnc.a[i * Xnc.k + j];
    const w1 = lag(nb, col);
    const w2 = lag(nb, w1);
    for (let i = 0; i < n; i++) {
      WX.a[i * Xnc.k + j] = w1[i];
      W2X.a[i * Xnc.k + j] = w2[i];
    }
  }
  const Zm = hcat(hcat(X, WX), W2X);
  // 1re étape : Ŵy = Z (Z'Z)^-1 Z' Wy  (ridge infime pour la colinéarité des instruments)
  const gz = gram(Zm, Wy);
  let tr = 0;
  for (let j = 0; j < Zm.k; j++) tr += gz.G[j * Zm.k + j];
  const { x: gamma } = solveSPD(gz.G, Zm.k, gz.c, (tr / Zm.k) * 1e-8);
  const WyHat = matVec(Zm, gamma);
  // 2e étape : régresser y sur X̂ = [Ŵy, X]
  const WyM: Mat = { n, k: 1, a: WyHat };
  const Xhat = hcat(WyM, X);
  const g2 = gram(Xhat, y);
  const { x: beta, inv } = solveSPD(g2.G, Xhat.k, g2.c);
  // résidus structurels avec le Wy observé
  const Xfull = hcat({ n, k: 1, a: Wy }, X);
  const { fitted, resid } = residuals(Xfull, y, beta);
  let ss = 0;
  for (let i = 0; i < n; i++) ss += resid[i] * resid[i];
  const df = n - Xhat.k;
  const fit: Fit = { beta, fitted, resid, bread: inv, Xinf: Xhat, weights: null, sigma2: ss / Math.max(1, df), df, k: Xhat.k, extra: { rho: beta[0], knn: k } };
  return { fit, rho: beta[0], Wy, nb, index };
}

/** Prédiction SAR hors échantillon : Wy des k voisins d'entraînement. */
export function sarPredict(res: SarResult, Xte: Mat, latTe: Float64Array, lngTe: Float64Array, yTrain: Float64Array, beta: Float64Array): Float64Array {
  const out = new Float64Array(Xte.n);
  for (let i = 0; i < Xte.n; i++) {
    const q = res.index.queryLatLng(latTe[i], lngTe[i], res.nb[0]?.length ?? 8);
    let wy = 0;
    for (let j = 0; j < q.idx.length; j++) wy += yTrain[q.idx[j]];
    wy = q.idx.length ? wy / q.idx.length : 0;
    let s = beta[0] * wy;
    for (let j = 0; j < Xte.k; j++) s += Xte.a[i * Xte.k + j] * beta[j + 1];
    out[i] = s;
  }
  return out;
}

export { cholInverse };