//
//  SentenceEmbedder.swift
//  Prisme
//
//  Author: Simon-Pierre Boucher <contact@spboucher.ai>
//

import Foundation
import NaturalLanguage

/// On-device sentence embeddings (NaturalLanguage framework — no Foundation
/// Models involved, works on every device, offline). Vectors from different
/// language models live in different spaces: the store tags each vector
/// with its language and only compares like with like.
@MainActor
final class SentenceEmbedder {
    private var cache: [NLLanguage: NLEmbedding] = [:]

    func languageCode(for text: String) -> String? {
        let recognizer = NLLanguageRecognizer()
        recognizer.processString(String(text.prefix(400)))
        return recognizer.dominantLanguage?.rawValue
    }

    func embed(_ text: String, languageCode: String?) -> [Float]? {
        guard let languageCode else { return nil }
        let language = NLLanguage(rawValue: languageCode)
        let embedding: NLEmbedding?
        if let cached = cache[language] {
            embedding = cached
        } else {
            embedding = NLEmbedding.sentenceEmbedding(for: language)
            if let embedding { cache[language] = embedding }
        }
        guard let vector = embedding?.vector(for: String(text.prefix(600))) else {
            return nil
        }
        return vector.map(Float.init)
    }

    static func cosine(_ a: [Float], _ b: [Float]) -> Float {
        guard a.count == b.count, !a.isEmpty else { return 0 }
        var dot: Float = 0, normA: Float = 0, normB: Float = 0
        for i in a.indices {
            dot += a[i] * b[i]
            normA += a[i] * a[i]
            normB += b[i] * b[i]
        }
        let denominator = (normA.squareRoot() * normB.squareRoot())
        return denominator > 0 ? dot / denominator : 0
    }
}