import * as tf from '@tensorflow/tfjs';
import {bundleResourceIO} from '@tensorflow/tfjs-react-native';
import {Image} from 'react-native';

const myModelJson = require('../assets/my_model.json');
const myModelWeights = require('../assets/my_model_weights.bin');

class LocalFileHandler {
  constructor(modelJson, modelWeightsId) {
    this.modelJson = modelJson;
    this.modelWeightsId = modelWeightsId;
  }

  async load() {
    const localFilePath = unescape(
      Image.resolveAssetSource(this.modelWeightsId).uri.substr(
        'file://'.length,
      ),
    );
    const base64Weights = await RNFS.readFile(localFilePath, 'base64');
    const weightData = tf.util.encodeString(base64Weights, 'base64').buffer;

    const modelArtifacts = Object.assign({}, this.modelJson);
    modelArtifacts.weightSpecs = this.modelJson.weightsManifest[0].weights;
    delete modelArtifacts.weightManifest;
    modelArtifacts.weightData = weightData;
    return modelArtifacts;
  }
}

let resourceIO = bundleResourceIO(
  myModelJson,
  myModelWeights,
);

if (
  Image.resolveAssetSource(myModelWeights).uri.match('^file://')
) {
  resourceIO = new LocalFileHandler(
    myModelJson,
    myModelWeights,
  );
}

const model = await tf.loadGraphModel(resourceIO);