import SwiftUI
import Vision
import CoreImage.CIFilterBuiltins

struct ContentView: View {
   @State private var picture = UIImage(named: "confused")!

   var body: some View {
      VStack {
         Button("Change Background") {
            Task {
               await detectPerson()
            }
         }.buttonStyle(.borderedProminent)

         Image(uiImage: picture)
            .resizable()
            .scaledToFit()
         Spacer()
      }
      .padding()
   }
   func detectPerson() async {
      let imageURL = Bundle.main.url(forResource: "confused", withExtension: "png")
      if let imageURL = imageURL {
         do {
            // Perform request
            let request = GeneratePersonSegmentationRequest()
            request.qualityLevel = .accurate
            let observation = try await request.perform(on: imageURL)
            let result = try observation.cgImage

            // Resize mask
            var mask = CIImage(cgImage: result)
            if let inputImage = CIImage(image: picture) {
               let maskScaleX = inputImage.extent.width / mask.extent.width
               let maskScaleY = inputImage.extent.height / mask.extent.height
               mask = mask.transformed(by: CGAffineTransformMake(maskScaleX, 0, 0, maskScaleY, 0, 0))
            }
            // Apply mask
            let context = CIContext()
            let maskFilter = CIFilter.blendWithMask()
            maskFilter.inputImage = CIImage(image: picture)
            maskFilter.backgroundImage = CIImage(image: UIImage(named: "office")!)
            maskFilter.maskImage = mask

            if let output = maskFilter.outputImage {
               if let image = context.createCGImage(output, from: output.extent) {
                  picture = UIImage(cgImage: image)
               }
            }
         } catch {
            print("Error performing the request: \(error)")
         }
      }
   }
}