main.swift

import Accelerate
import AVFoundation
import CoreML
import Foundation

struct Segment: Hashable {
    let start: Double
    let end: Double
}

struct SlidingWindow {
    var start: Double
    var duration: Double
    var step: Double
    
    func time(forFrame index: Int) -> Double {
        return start + Double(index) * step
    }
    
    func segment(forFrame index: Int) -> Segment {
        let s = time(forFrame: index)
        return Segment(start: s, end: s + duration)
    }
}

struct SlidingWindowFeature {
    var data: [[[Float]]] // (1, 589, 3)
    var slidingWindow: SlidingWindow
}

var speakerDB: [String: [Float]] = [:]  // Global speaker database
let threshold: Float = 0.7              // Distance threshold

func cosineDistance(_ x: [Float], _ y: [Float]) -> Float {
    precondition(x.count == y.count, "Vectors must be same size")
    let dot = zip(x, y).map(*).reduce(0, +)
    let normX = sqrt(x.map { $0 * $0 }.reduce(0, +))
    let normY = sqrt(y.map { $0 * $0 }.reduce(0, +))
    return 1.0 - (dot / (normX * normY + 1e-6))
}

func updateSpeakerDB(_ speaker: String, _ newEmbedding: [Float], alpha: Float = 0.9) {
    guard var oldEmbedding = speakerDB[speaker] else { return }
    for i in 0..<oldEmbedding.count {
        oldEmbedding[i] = alpha * oldEmbedding[i] + (1 - alpha) * newEmbedding[i]
    }
    speakerDB[speaker] = oldEmbedding
}

func assignSpeaker(embedding: [Float], threshold: Float = 0.7) -> String {
    if speakerDB.isEmpty {
        let speaker = "Speaker 1"
        speakerDB[speaker] = embedding
        return speaker
    }

    var minDistance: Float = Float.greatestFiniteMagnitude
    var identifiedSpeaker: String? = nil

    for (speaker, refEmbedding) in speakerDB {
        let distance = cosineDistance(embedding, refEmbedding)
        if distance < minDistance {
            minDistance = distance
            identifiedSpeaker = speaker
        }
    }

    if let bestSpeaker = identifiedSpeaker {
        if minDistance > threshold {
            let newSpeaker = "Speaker \(speakerDB.count + 1)"
            speakerDB[newSpeaker] = embedding
            return newSpeaker
        } else {
            updateSpeakerDB(bestSpeaker, embedding)
            return bestSpeaker
        }
    }

    return "Unknown"
}

func getAnnotation(annotation: inout [Segment: String],
                   speakerMapping: [Int: Int],
                   binarizedSegments: [[[Float]]],
                   slidingWindow: SlidingWindow) {
    
    let segmentation = binarizedSegments[0] // shape: [589][3]
    let numFrames = segmentation.count

    // Step 1: argmax to get dominant speaker per frame
    var frameSpeakers: [Int] = []
    for frame in segmentation {
        if let maxIdx = frame.indices.max(by: { frame[$0] < frame[$1] }) {
            frameSpeakers.append(maxIdx)
        } else {
            frameSpeakers.append(0) // fallback
        }
    }

    // Step 2: group contiguous same-speaker segments
    var currentSpeaker = frameSpeakers[0]
    var startFrame = 0

    for i in 1..<numFrames {
        if frameSpeakers[i] != currentSpeaker {
            let startTime = slidingWindow.time(forFrame: startFrame)
            let endTime = slidingWindow.time(forFrame: i)

            let segment = Segment(start: startTime, end: endTime)
            if let mappedSpeaker = speakerMapping[currentSpeaker] {
                annotation[segment] = "Speaker \(mappedSpeaker)"
            }
            currentSpeaker = frameSpeakers[i]
            startFrame = i
        }
    }

    // Final segment
    let finalStart = slidingWindow.time(forFrame: startFrame)
    let finalEnd = slidingWindow.segment(forFrame: numFrames - 1).end
    let finalSegment = Segment(start: finalStart, end: finalEnd)
    if let mappedSpeaker = speakerMapping[currentSpeaker] {
        annotation[finalSegment] = "Speaker \(mappedSpeaker)"
    }
}


func getEmbedding(audioChunk: [Float],
                  binarizedSegments _: [[[Float]]],
                  slidingWindowSegments: SlidingWindowFeature,
                  chunkSize: Int = 10 * 16000,
                  embeddingModel: MLModel) -> MLMultiArray?
{
    // 1. Create audio_tensor of shape (1, 1, chunkSize)
    let audioTensor = audioChunk

    let batchSize = slidingWindowSegments.data.count
    let numFrames = slidingWindowSegments.data[0].count
    let numSpeakers = slidingWindowSegments.data[0][0].count

    // 2. Compute clean_frames = 1.0 where active speakers < 2
    var cleanFrames = Array(repeating: Array(repeating: 0.0 as Float, count: 1), count: numFrames)

    for f in 0 ..< numFrames {
        let frame = slidingWindowSegments.data[0][f]
        let speakerSum = frame.reduce(0, +)
        cleanFrames[f][0] = (speakerSum < 2.0) ? 1.0 : 0.0
    }

    // 3. Multiply slidingWindowSegments.data by cleanFrames
    var cleanSegmentData = Array(
        repeating: Array(repeating: Array(repeating: 0.0 as Float, count: numSpeakers), count: numFrames),
        count: 1
    )

    for f in 0 ..< numFrames {
        for s in 0 ..< numSpeakers {
            cleanSegmentData[0][f][s] = slidingWindowSegments.data[0][f][s] * cleanFrames[f][0]
        }
    }

    // 4. Flatten audio tensor to shape (3, 160000)
    var audioBatch: [[Float]] = []
    for _ in 0 ..< 3 {
        audioBatch.append(audioTensor)
    }

    // 5. Transpose mask shape to (3, 589)
    var cleanMasks: [[Float]] = Array(repeating: Array(repeating: 0.0, count: numFrames), count: numSpeakers)

    for s in 0 ..< numSpeakers {
        for f in 0 ..< numFrames {
            cleanMasks[s][f] = cleanSegmentData[0][f][s]
        }
    }

    // 6. Prepare MLMultiArray inputs
    guard let waveformArray = try? MLMultiArray(shape: [3, chunkSize] as [NSNumber], dataType: .float32),
          let maskArray = try? MLMultiArray(shape: [3, numFrames] as [NSNumber], dataType: .float32)
    else {
        print("Failed to allocate MLMultiArray")
        return nil
    }

    // Fill waveform
    for s in 0 ..< 3 {
        for i in 0 ..< chunkSize {
            waveformArray[s * chunkSize + i] = NSNumber(value: audioBatch[s][i])
        }
    }

    // Fill mask
    for s in 0 ..< 3 {
        for f in 0 ..< numFrames {
            maskArray[s * numFrames + f] = NSNumber(value: cleanMasks[s][f])
        }
    }

    // 7. Run model
    let inputs: [String: Any] = [
        "waveform": waveformArray,
        "mask": maskArray,
    ]

    guard let output = try? embeddingModel.prediction(from: MLDictionaryFeatureProvider(dictionary: inputs)) else {
        print("Embedding model prediction failed")
        return nil
    }

    return output.featureValue(for: "embedding")?.multiArrayValue
}

func loadAudioSamples(from url: URL, expectedSampleRate: Double = 16000.0) throws -> [Float] {
    let file = try AVAudioFile(forReading: url)
    let format = AVAudioFormat(commonFormat: .pcmFormatFloat32,
                               sampleRate: expectedSampleRate,
                               channels: 1,
                               interleaved: false)!

    let engine = AVAudioEngine()
    let player = AVAudioPlayerNode()
    engine.attach(player)

    let converter = AVAudioConverter(from: file.processingFormat, to: format)!
    let frameCapacity = AVAudioFrameCount(file.length)
    let buffer = AVAudioPCMBuffer(pcmFormat: file.processingFormat, frameCapacity: frameCapacity)!
    try file.read(into: buffer)

    let outputBuffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCapacity)!

    let inputBlock: AVAudioConverterInputBlock = { _, outStatus in
        outStatus.pointee = .haveData
        return buffer
    }

    try converter.convert(to: outputBuffer, error: nil, withInputFrom: inputBlock)

    guard let floatChannelData = outputBuffer.floatChannelData else {
        throw NSError(domain: "Audio", code: -1, userInfo: [NSLocalizedDescriptionKey: "Missing float data"])
    }

    let channelData = floatChannelData[0]
    let samples = Array(UnsafeBufferPointer(start: channelData, count: Int(outputBuffer.frameLength)))
    return samples
}

func chunkAndRunSegmentation(samples: [Float], chunkSize: Int = 160_000, model: MLModel, embeddingModel: MLModel) throws {
    let totalSamples = samples.count
    let numberOfChunks = Int(ceil(Double(totalSamples) / Double(chunkSize)))
    var annotations: [Segment: String] = [:]

    for i in 0 ..< numberOfChunks {
        let start = i * chunkSize
        let end = min((i + 1) * chunkSize, totalSamples)
        let chunk = Array(samples[start ..< end])

        // If chunk is shorter than 10s, pad with zeros
        var paddedChunk = chunk
        if chunk.count < chunkSize {
            paddedChunk += Array(repeating: 0.0, count: chunkSize - chunk.count)
        }

        let binarizedSegments = try getSegments(audioChunk: paddedChunk, model: model)
        let frames = SlidingWindow(start: Double(i) * 10.0, duration: 0.0619375, step: 0.016875)
        let slidingFeature = SlidingWindowFeature(data: binarizedSegments, slidingWindow: frames)
        if let embeddings = getEmbedding(audioChunk: paddedChunk,
                                         binarizedSegments: binarizedSegments,
                                         slidingWindowSegments: slidingFeature,
                                         embeddingModel: embeddingModel)
        {
            print("Embeddings shape: \(embeddings.shape.map { $0.intValue })")

            let shape = embeddings.shape.map { $0.intValue } // [3, 256]
            let numSpeakers = shape[0]
            let embeddingDim = shape[1]
            let strides = embeddings.strides.map { $0.intValue }

            var speakerSums = [Float](repeating: 0.0, count: numSpeakers)

            for s in 0 ..< numSpeakers {
                for d in 0 ..< embeddingDim {
                    let index = s * strides[0] + d * strides[1]
                    speakerSums[s] += embeddings[index].floatValue
                }
            }

            print("Sum along axis 1 (per speaker): \(speakerSums)") 

            // Step 3: Assign speaker label to each embedding
            var speakerLabels = [String]()
            for s in 0..<numSpeakers {
                var embeddingVec = [Float](repeating: 0.0, count: embeddingDim)
                for d in 0..<embeddingDim {
                    let index = s * strides[0] + d * strides[1]
                    embeddingVec[d] = embeddings[index].floatValue
                }
                let label = assignSpeaker(embedding: embeddingVec)
                speakerLabels.append(label)
            }

            print("Chunk \(i + 1): Assigned Speakers: \(speakerLabels)")

            // Step 4: Update annotations
            // Map speaker index 0,1,2 → assigned speakerLabels
            var labelMapping: [Int: Int] = [:]
            for (idx, label) in speakerLabels.enumerated() {
                if let spkNum = Int(label.components(separatedBy: " ").last ?? "") {
                    labelMapping[idx] = spkNum
                }
            }

            getAnnotation(annotation: &annotations,
                      speakerMapping: labelMapping,
                      binarizedSegments: binarizedSegments,
                      slidingWindow: frames)

            print("Chunk \(i + 1) → Segments shape: \(binarizedSegments[0].count) frames")
        }
    }

    // Final result
    print("\n=== Final Annotations ===")
    for (segment, speaker) in annotations.sorted(by: { $0.key.start < $1.key.start }) {
        print("\(speaker): \(segment.start) - \(segment.end)")
    }

}

func powersetConversion(_ segments: [[[Float]]]) -> [[[Float]]] {
    let powerset: [[Int]] = [
        [], // 0
        [0], // 1
        [1], // 2
        [2], // 3
        [0, 1], // 4
        [0, 2], // 5
        [1, 2], // 6
    ]

    let batchSize = segments.count
    let numFrames = segments[0].count
    let numCombos = segments[0][0].count // 7

    let numSpeakers = 3
    var binarized = Array(
        repeating: Array(
            repeating: Array(repeating: 0.0 as Float, count: numSpeakers),
            count: numFrames
        ),
        count: batchSize
    )

    for b in 0 ..< batchSize {
        for f in 0 ..< numFrames {
            let frame = segments[b][f]

            // Find index of max value in this frame
            guard let bestIdx = frame.indices.max(by: { frame[$0] < frame[$1] }) else {
                continue
            }

            // Mark the corresponding speakers as active
            for speaker in powerset[bestIdx] {
                binarized[b][f][speaker] = 1.0
            }
        }
    }

    return binarized
}

func getSegments(audioChunk: [Float], sampleRate _: Int = 16000, chunkSize: Int = 160_000, model: MLModel) throws -> [[[Float]]] {
    // Ensure correct shape: (1, 1, chunk_size)
    let audioArray = try MLMultiArray(shape: [1, 1, NSNumber(value: chunkSize)], dataType: .float32)
    for i in 0 ..< audioChunk.count {
        audioArray[i] = NSNumber(value: audioChunk[i])
    }

    // Prepare input
    let input = try MLDictionaryFeatureProvider(dictionary: ["audio": audioArray])

    // Run prediction
    let output = try model.prediction(from: input)

    // Extract segments output: shape assumed (1, frames, 7)
    guard let segmentOutput = output.featureValue(for: "segments")?.multiArrayValue else {
        throw NSError(domain: "ModelOutput", code: -1, userInfo: [NSLocalizedDescriptionKey: "Missing segments output"])
    }

    let frames = segmentOutput.shape[1].intValue
    let combinations = segmentOutput.shape[2].intValue

    // Convert MLMultiArray to [[[Float]]]
    var segments = Array(repeating: Array(repeating: Array(repeating: 0.0 as Float, count: combinations), count: frames), count: 1)

    for f in 0 ..< frames {
        for c in 0 ..< combinations {
            let index = f * combinations + c
            segments[0][f][c] = segmentOutput[index].floatValue
        }
    }

    // Apply powerset conversion
    let binarizedSegments = powersetConversion(segments)

    // Assume segments shape is (1, 589, 3)
    guard binarizedSegments.count == 1 else {
        fatalError("Expected batch size 1")
    }

    let b_frames = binarizedSegments[0]
    let numSpeakers = b_frames[0].count

    // Initialize sum array
    var speakerSums = Array(repeating: 0.0 as Float, count: numSpeakers)

    // Sum across axis 1 (frames)
    for frame in b_frames {
        for (i, value) in frame.enumerated() {
            speakerSums[i] += value
        }
    }

    print("Sum across axis 1 (frames): \(speakerSums)")

    return binarizedSegments
}

func loadModel(from path: String) throws -> MLModel {
    let url = URL(fileURLWithPath: path)
    let model = try MLModel(contentsOf: url)
    return model
}

do {
    let modelPath = "./pyannote_segmentation.mlmodelc"
    let embeddingPath = "./wespeaker.mlmodelc"
    let model = try loadModel(from: modelPath)
    let embeddingModel = try loadModel(from: embeddingPath)
    print("Model loaded successfully.")

    // let audioPath = "./first_10_seconds.wav"
    let audioPath = "./TS3003b_mix_headset.wav"

    let audioSamples = try loadAudioSamples(from: URL(fileURLWithPath: audioPath))
    try chunkAndRunSegmentation(samples: audioSamples, model: model, embeddingModel: embeddingModel)
} catch {
    print("Error: \(error)")
}