import SwiftUI
import CoreImage
import CoreImage.CIFilterBuiltins
import UIKit
import CoreGraphics
import Accelerate
import Foundation
import CoreML
import CreateML
import CreateMLComponents


func createMLShapedArray(from array: [Double], shape: [Int]) -> MLShapedArray {
    return MLShapedArray(scalars: array, shape: shape)
}


func calculateTransformationMatrixWithNonlinearity(sourceRGB: [[Double]], referenceRGB: [[Double]], degree: Int = 3) async throws -> MultivariateLinearRegressor.Model {

    
    
    
   let annotatedFeatures2 = zip(sourceRGB, referenceRGB).map { (featureArray, targetArray) -> AnnotatedFeature, MLShapedArray> in
        let featureMLShapedArray = createMLShapedArray(from: featureArray, shape: [featureArray.count])
        let targetMLShapedArray = createMLShapedArray(from: targetArray, shape: [targetArray.count])
        return AnnotatedFeature(feature: featureMLShapedArray, annotation: targetMLShapedArray)
    }

    // Flatten the sourceRGBPoly into a single-dimensional array
    var flattenedArray = sourceRGB.flatMap { $0 }
    let featuresMLShapedArray = createMLShapedArray(from: flattenedArray, shape: [2, 4])
    flattenedArray = referenceRGB.flatMap { $0 }
    let targetMLShapedArray = createMLShapedArray(from: flattenedArray, shape: [2, 2])

    // Create AnnotatedFeature instances
   /* let annotatedFeatures2: [AnnotatedFeature, MLShapedArray>] = [
        AnnotatedFeature(feature: featuresMLShapedArray, annotation: targetMLShapedArray)
    ]*/
 
    let annotatedBatch = AnnotatedBatch(features: featuresMLShapedArray, annotations: targetMLShapedArray)

    
    var regressor = MultivariateLinearRegressor()
    regressor.configuration.learningRate = 0.1
    regressor.configuration.maximumIterationCount=5000
    regressor.configuration.batchSize=2
    

    let model = try await regressor.fitted(to: annotatedBatch,validateOn: nil)

    
    //var model = try await regressor.fitted(to: annotatedFeatures2)

    // Proceed to prediction once the model is fitted
    let predictions = try await model.prediction(from: annotatedFeatures2)
    // Process or use the predictions
    print(predictions)
    print("Predictions:", predictions)
    return model
}




struct ContentView: View {
  var body: some View {
    VStack {}
    .onAppear {
      Task {
        do {
            let sourceMatrix: [[Double]] = [
                [0,0.1,0.2,0.3],
                [0.5,0.2,0.6,0.2]
                
            ]
            
            let referenceMatrix: [[Double]] = [
                [0.2,0.7],
                [0.9,0.1]
            ]

          let model = try await calculateTransformationMatrixWithNonlinearity(sourceRGB: sourceMatrix, referenceRGB: referenceMatrix, degree: 2
          )
          print("Model fitted successfully:", model)
        } catch {
          print("Error:", error)
        }
      }
    }
  }
}