//
//  KeypointsUtils.swift
//  PickleBallUI
//
//  Created by Michael Zhang on 1/20/25.
//
import Foundation
import SwiftUI

typealias XYXY = (x1: Float, y1: Float, x2: Float, y2: Float)

struct Prediction {
    let id = UUID()
    let classIndex: Int
    let score: Float
    let xyxy: (x1: Float, y1: Float, x2: Float, y2: Float)
    let maskCoefficients: [Float]
    let inputImgSize: CGSize
}

struct MaskPrediction: Identifiable {
    let id = UUID()
    
    let classIndex: Int
    
    let mask: [UInt8]
    let maskSize: (width: Int, height: Int)
    
    func getMaskImage() -> UIImage? {
        guard !mask.isEmpty else { return nil }
        
        let coloredMask = colorizeMask(mask, color: UIColor.white)
        
        let numComponents = 4
        let numBytes = maskSize.width * maskSize.height * numComponents
        
        let colorspace = CGColorSpaceCreateDeviceRGB()
        let rgbData = CFDataCreate(nil, coloredMask, numBytes)!
        let provider = CGDataProvider(data: rgbData)!
        guard let rgbImageRef = CGImage(
            width: maskSize.width,
            height: maskSize.height,
            bitsPerComponent: 8,
            bitsPerPixel: 8 * numComponents,
            bytesPerRow: maskSize.width * numComponents,
            space: colorspace,
            bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue),
            provider: provider,
            decode: nil,
            shouldInterpolate: false,
            intent: CGColorRenderingIntent(rawValue: 0)!
        ) else { return nil }
        
        return UIImage(cgImage: rgbImageRef)
    }
}

func colorizeMask(_ mask: [UInt8], color: UIColor) -> [UInt8] {
    var red: CGFloat = 0
    var green: CGFloat = 0
    var blue: CGFloat = 0
    var alpha: CGFloat = 0
    
    color.getRed(&red, green: &green, blue: &blue, alpha: &alpha)
    
    var coloredMask: [UInt8] = []
    for value in mask {
        if value == 0 {
            coloredMask.append(contentsOf: [0, 0, 0, 0])
        } else {
            coloredMask.append(
                contentsOf: [
                    UInt8(truncating: (red * 255) as NSNumber),
                    UInt8(truncating: (green * 255) as NSNumber),
                    UInt8(truncating: (blue * 255) as NSNumber),
                    255,
                ]
            )
        }
    }
    
    return coloredMask
}

extension Collection where Element == MaskPrediction {
    func combineToSingleImage() -> UIImage? {
        guard let firstMask = self.first else { return nil }
        
        let size = CGSize(width: firstMask.maskSize.width, height: firstMask.maskSize.height)
        UIGraphicsBeginImageContext(size)
        
        let areaSize = CGRect(x: 0, y: 0, width: size.width, height: size.height)
        firstMask.getMaskImage()?.draw(in: areaSize)
        
        for mask in self.dropFirst() {
            mask.getMaskImage()?.draw(in: areaSize, blendMode: .normal, alpha: 1.0)
        }
        
        let newImage = UIGraphicsGetImageFromCurrentImageContext()
        UIGraphicsEndImageContext()
        
        return newImage
    }
}

extension MaskPrediction {
    func toCGImage() -> CGImage? {
        let bitsPerComponent = 8
        let bitsPerPixel = 8
        let bytesPerRow = maskSize.width
        
        guard let provider = CGDataProvider(data: Data(mask) as CFData),
              let cgImage = CGImage(
                width: maskSize.width,
                height: maskSize.height,
                bitsPerComponent: bitsPerComponent,
                bitsPerPixel: bitsPerPixel,
                bytesPerRow: bytesPerRow,
                space: CGColorSpaceCreateDeviceGray(),
                bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.none.rawValue),
                provider: provider,
                decode: nil,
                shouldInterpolate: false,
                intent: .defaultIntent
              ) else {
            return nil
        }
        return cgImage
    }
}