import XCTest
@testable import MarkBase

final class MultilangEmbeddingTest: XCTestCase {
    var engine: MarkBaseEngine!
    var embeddingModel: TextEmbeddingModel!
    let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"

    override func setUp() {
        super.setUp()
        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else { return }
        engine = try? MarkBaseEngine(autoCompile: true)
        embeddingModel = try? TextEmbeddingModel(modelDir: modelDir, engine: engine, config: TextEmbeddingConfig())
    }

    func testMultilangEmbeddings() throws {
        try XCTSkipIf(embeddingModel == nil, "Model not found")

        let texts: [String: String] = [
            "en": "The weather is beautiful today",
            "zh": "今天天氣很好",
            "ja": "今日は天気がいいです",
            "ko": "오늘 날씨가 좋습니다",
            "es": "El clima está hermoso hoy",
            "fr": "Il fait beau aujourd'hui",
            "de": "Das Wetter ist heute schön",
            "ru": "Сегодня прекрасная погода",
            "ar": "الطقس جميل اليوم",
            "hi": "आज मौसम बहुत सुंदर है"
        ]

        var embeddings: [String: [Float]] = [:]
        for (lang, text) in texts {
            let emb = try embeddingModel.embed(text: text)
            XCTAssertEqual(emb.count, 2560, "\(lang) embedding dimension")
            let norm = sqrt(emb.reduce(0) { $0 + $1 * $1 })
            XCTAssertEqual(norm, 1.0, accuracy: 0.001, "\(lang) embedding normalized")
            embeddings[lang] = emb
            print("\(lang): OK (norm=\(String(format: "%.4f", norm)))")
        }

        // Cross-language similarity (en-zh should be higher than en-ru for weather context)
        let enEmb = embeddings["en"]!
        let zhEmb = embeddings["zh"]!
        let jaEmb = embeddings["ja"]!

        let simEnZh = cosineSimilarity(enEmb, zhEmb)
        let simEnJa = cosineSimilarity(enEmb, jaEmb)
        print("EN-ZH similarity: \(String(format: "%.4f", simEnZh))")
        print("EN-JA similarity: \(String(format: "%.4f", simEnJa))")
    }

    func testCrossLingualSemanticSimilarity() throws {
        try XCTSkipIf(embeddingModel == nil, "Model not found")

        // Same meaning, different languages
        let enCat = try embeddingModel.embed(text: "The cat is sleeping")
        let zhCat = try embeddingModel.embed(text: "貓在睡覺")
        let enDog = try embeddingModel.embed(text: "The dog is running")

        let simSame = cosineSimilarity(enCat, zhCat)
        let simDiff = cosineSimilarity(enCat, enDog)
        print("EN(cat) - ZH(cat): \(String(format: "%.4f", simSame))")
        print("EN(cat) - EN(dog): \(String(format: "%.4f", simDiff))")
        print("Cross-lingual > Same-lingual different topic: \(simSame > simDiff)")
    }

    private func cosineSimilarity(_ 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 0..<a.count {
            dot += a[i] * b[i]
            normA += a[i] * a[i]
            normB += b[i] * b[i]
        }
        return dot / (sqrt(normA) * sqrt(normB))
    }
}