//
//  Correlate.swift
//  Metrika
//
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//

/// Pearson correlation matrix on listwise-complete data (the engine
/// applies listwise deletion before calling, matching Stata `correlate`).
public enum ZQCorrelate {
    public static func matrix(columns: [[Double]]) -> [[Double]] {
        let k = columns.count
        guard k > 0 else { return [] }
        let n = columns[0].count

        let means = columns.map { $0.reduce(0, +) / Double(n) }
        var centered = columns
        for j in 0..<k {
            for i in 0..<n { centered[j][i] -= means[j] }
        }
        let norms = centered.map { column in
            column.reduce(0) { $0 + $1 * $1 }.squareRoot()
        }

        var result = [[Double]](repeating: [Double](repeating: 1, count: k), count: k)
        for a in 0..<k {
            for b in (a + 1)..<k {
                var dot = 0.0
                for i in 0..<n { dot += centered[a][i] * centered[b][i] }
                let denominator = norms[a] * norms[b]
                let r = denominator > 0 ? dot / denominator : .nan
                result[a][b] = r
                result[b][a] = r
            }
        }
        return result
    }
}