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

import ZQData
import ZQParser

/// Reference native plugin (CLAUDE.md §4): `zscore varname` generates
/// `z_varname`, the standardized version of a numeric variable. Declares
/// `mutates: true` and returns the updated dataset via
/// `ZQResult.replacementFrame` — the pattern SPM-compiled-in user
/// plugins follow.
public struct ZScorePlugin: ZQCommandPlugin {
    public static let verb = "zscore"
    public static let syntax = ZQSyntaxSpec(
        acceptsVarlist: true,
        acceptsCondition: false,
        acceptsRange: false,
        acceptsWeight: false,
        options: [],
        mutates: true
    )

    public init() {}

    public func execute(_ ctx: ZQContext) async throws -> ZQResult {
        let names = ctx.command.varlist.flatMap(\.referencedNames)
        guard names.count == 1, let name = names.first else {
            throw ZQDataError("zscore: syntax is 'zscore varname'")
        }
        let (values, missing) = try ctx.frame.requireNumeric(name)

        var count = 0
        var sum = 0.0
        for i in 0..<values.count where !missing[i] {
            sum += values[i]
            count += 1
        }
        guard count > 1 else { throw ZQDataError("zscore: insufficient observations") }
        let mean = sum / Double(count)
        var squared = 0.0
        for i in 0..<values.count where !missing[i] {
            squared += (values[i] - mean) * (values[i] - mean)
        }
        let sd = (squared / Double(count - 1)).squareRoot()
        guard sd > 0 else { throw ZQDataError("zscore: '\(name)' is constant") }

        var scores = [Double](repeating: .nan, count: values.count)
        for i in 0..<values.count where !missing[i] {
            scores[i] = (values[i] - mean) / sd
        }

        var frame = ctx.frame
        try frame.addColumn(ZQColumn(
            name: "z_\(name)",
            data: .float64(values: scores, missing: missing)
        ))
        return ZQResult(
            text: "(variable z_\(name) generated: mean \(mean), sd \(sd))",
            scalars: ["mean": mean, "sd": sd],
            replacementFrame: frame
        )
    }
}