// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#import "OnnxruntimeModule.h"
#import "FakeRCTBlobManager.h"
#import "TensorHelper.h"

#import <XCTest/XCTest.h>
#import <onnxruntime/onnxruntime_cxx_api.h>

@interface OnnxruntimeModuleTest : XCTestCase

@end

@implementation OnnxruntimeModuleTest

FakeRCTBlobManager *fakeBlobManager = nil;

+ (void)initialize {
  if (self == [OnnxruntimeModuleTest class]) {
    fakeBlobManager = [FakeRCTBlobManager new];
  }
}

- (void)testOnnxruntimeModule {
  NSBundle *bundle = [NSBundle bundleForClass:[OnnxruntimeModuleTest class]];
  NSString *dataPath = [bundle pathForResource:@"test_types_float" ofType:@"ort"];
  NSString *sessionKey = @"";
  NSString *sessionKey2 = @"";

  OnnxruntimeModule *onnxruntimeModule = [OnnxruntimeModule new];
  [onnxruntimeModule setBlobManager:fakeBlobManager];

  {
    // test loadModelFromBuffer()
    NSMutableDictionary *options = [NSMutableDictionary dictionary];
    NSData *fileData = [NSData dataWithContentsOfFile:dataPath];

    NSDictionary *resultMap = [onnxruntimeModule loadModelFromBuffer:fileData options:options];
    sessionKey = resultMap[@"key"];
    NSArray *inputNames = resultMap[@"inputNames"];
    XCTAssertEqual([inputNames count], 1);
    XCTAssertEqualObjects(inputNames[0], @"input");
    NSArray *outputNames = resultMap[@"outputNames"];
    XCTAssertEqual([outputNames count], 1);
    XCTAssertEqualObjects(outputNames[0], @"output");

    // test loadModel()
    NSDictionary *resultMap2 = [onnxruntimeModule loadModel:dataPath options:options];
    sessionKey2 = resultMap2[@"key"];
    NSArray *inputNames2 = resultMap2[@"inputNames"];
    XCTAssertEqual([inputNames2 count], 1);
    XCTAssertEqualObjects(inputNames2[0], @"input");
    NSArray *outputNames2 = resultMap2[@"outputNames"];
    XCTAssertEqual([outputNames2 count], 1);
    XCTAssertEqualObjects(outputNames2[0], @"output");
  }

  // test run()
  {
    NSMutableDictionary *inputTensorMap = [NSMutableDictionary dictionary];

    // dims
    NSArray *dims = @[ [NSNumber numberWithLong:1], [NSNumber numberWithLong:5] ];
    inputTensorMap[@"dims"] = dims;

    // type
    inputTensorMap[@"type"] = JsTensorTypeFloat;

    // data
    std::array<float, 5> outValues{std::numeric_limits<float>::min(), 1.0f, -2.0f, 3.0f,
                                   std::numeric_limits<float>::max()};

    const NSInteger byteBufferSize = outValues.size() * sizeof(float);
    unsigned char *byteBuffer = static_cast<unsigned char *>(malloc(byteBufferSize));
    NSData *byteBufferRef = [NSData dataWithBytesNoCopy:byteBuffer length:byteBufferSize];
    float *floatPtr = (float *)[byteBufferRef bytes];
    for (NSUInteger i = 0; i < outValues.size(); ++i) {
      *floatPtr++ = outValues[i];
    }
    floatPtr = (float *)[byteBufferRef bytes];

    XCTAssertNotNil(fakeBlobManager);
    inputTensorMap[@"data"] = [fakeBlobManager testCreateData:byteBufferRef];

    NSMutableDictionary *inputDataMap = [NSMutableDictionary dictionary];
    inputDataMap[@"input"] = inputTensorMap;

    NSMutableDictionary *options = [NSMutableDictionary dictionary];

    NSMutableArray *output = [NSMutableArray array];
    [output addObject:@"output"];

    NSDictionary *resultMap = [onnxruntimeModule run:sessionKey input:inputDataMap output:output options:options];
    NSDictionary *resultMap2 = [onnxruntimeModule run:sessionKey2 input:inputDataMap output:output options:options];

    // Compare output & input, but data.blobId is different
    // dims
    XCTAssertTrue([[resultMap objectForKey:@"output"][@"dims"] isEqualToArray:inputTensorMap[@"dims"]]);
    XCTAssertTrue([[resultMap2 objectForKey:@"output"][@"dims"] isEqualToArray:inputTensorMap[@"dims"]]);

    // type
    XCTAssertEqual([resultMap objectForKey:@"output"][@"type"], JsTensorTypeFloat);
    XCTAssertEqual([resultMap2 objectForKey:@"output"][@"type"], JsTensorTypeFloat);

    // data ({ blobId, offset, size })
    XCTAssertEqual([[resultMap objectForKey:@"output"][@"data"][@"offset"] longValue], 0);
    XCTAssertEqual([[resultMap2 objectForKey:@"output"][@"data"][@"size"] longValue], byteBufferSize);
  }

  // test dispose
  {
    [onnxruntimeModule dispose:sessionKey];
    [onnxruntimeModule dispose:sessionKey2];
  }
}

@end