//
//  ModelCreationUtils.mm
//  CoreML_tests
//
//  Created by Anil Katti on 4/8/19.
//  Copyright © 2019 Apple Inc. All rights reserved.
//

#include "ModelCreationUtils.hpp"
#include "ParameterTests.hpp"

using namespace CoreML;

Specification::NeuralNetwork* buildBasicUpdatableNeuralNetworkModel(Specification::Model& m) {
    TensorAttributes inTensorAttr = { "A", 3 };
    TensorAttributes outTensorAttr = { "B", 1 };
    
    return buildBasicNeuralNetworkModel(m, true, &inTensorAttr, &outTensorAttr);
}

Specification::NeuralNetwork* buildBasicNeuralNetworkModel(Specification::Model& m, bool isUpdatable, const TensorAttributes *inTensorAttr, const TensorAttributes *outTensorAttr, int numberOfLayers, bool areWeightsQuantized, bool isBiasQuantized) {
    auto inTensor = m.mutable_description()->add_input();
    inTensor->set_name(inTensorAttr->name);
    auto inTensorShape = inTensor->mutable_type()->mutable_multiarraytype();
    inTensorShape->set_datatype(Specification::ArrayFeatureType_ArrayDataType_FLOAT32);
    for (int i = 0; i < inTensorAttr->dimension; i++) {
        inTensorShape->add_shape(1);
    }
    
    auto outTensor = m.mutable_description()->add_output();
    outTensor->set_name(outTensorAttr->name);
    auto outTensorShape = outTensor->mutable_type()->mutable_multiarraytype();
    outTensorShape->set_datatype(Specification::ArrayFeatureType_ArrayDataType_FLOAT32);
    for (int i = 0; i < outTensorAttr->dimension; i++) {
        outTensorShape->add_shape(1);
    }
    
    m.set_specificationversion(MLMODEL_SPECIFICATION_VERSION_IOS13);
    
    auto neuralNet = m.mutable_neuralnetwork();
    
    for (int i = 0; i < numberOfLayers; i++) {
        auto layer = neuralNet->add_layers();
        
        std::string layerName = "inner_layer" + std::to_string(i);
        std::string input = "output" + std::to_string(i-1);
        std::string output = "output" + std::to_string(i);
        
        layer->set_name((numberOfLayers == 1) ? "inner_layer" : layerName.c_str());
        layer->add_input((i == 0) ? inTensorAttr->name : input.c_str());
        layer->add_output((i == numberOfLayers-1) ? outTensorAttr->name : output.c_str());
        
        Specification::InnerProductLayerParams *innerProductParams = layer->mutable_innerproduct();
        innerProductParams->set_inputchannels(1);
        innerProductParams->set_outputchannels(1);

        // set weight
        auto* weights = innerProductParams->mutable_weights();
        if (areWeightsQuantized) {
            auto *quant_params = weights->mutable_quantization();
            quant_params->set_numberofbits(1);
            auto *linear_quant_params = quant_params->mutable_linearquantization();
            linear_quant_params->add_scale(1.0f);
            linear_quant_params->add_bias(0.0f);
            weights->set_int8rawvalue("x01"); // this is [1]
        } else {
            weights->add_floatvalue(1.0);
        }

        // set bias
        innerProductParams->set_hasbias(true);
        auto* bias = innerProductParams->mutable_bias();
        if (isBiasQuantized) {
            auto *quant_params = bias->mutable_quantization();
            quant_params->set_numberofbits(1);
            auto *linear_quant_params = quant_params->mutable_linearquantization();
            linear_quant_params->add_scale(1.0f);
            linear_quant_params->add_bias(0.0f);
            bias->set_int8rawvalue("x01"); // this is [1]
        } else {
            bias->add_floatvalue(1.0);
        }
        
        if (isUpdatable) {
            layer->set_isupdatable(true);
            innerProductParams->mutable_weights()->set_isupdatable(true);
            innerProductParams->mutable_bias()->set_isupdatable(true);
        }
    }
    
    if (isUpdatable) {
        m.set_isupdatable(true);
        
        for (auto feature : m.description().input()) {
            auto trainingInput = m.mutable_description()->mutable_traininginput()->Add();
            trainingInput->CopyFrom(feature);
        }
    }
    
    return neuralNet;
}

Specification::NeuralNetwork* addInnerProductLayer(Specification::Model& m, bool isUpdatable, const char *name, const TensorAttributes *inTensorAttr, const TensorAttributes *outTensorAttr, bool areWeightsQuantized, bool isBiasQuantized) {
    
    auto neuralNet = m.mutable_neuralnetwork();
    auto layer = neuralNet->add_layers();
    
    layer->set_name(name);
    layer->add_input(inTensorAttr->name);
    layer->add_output(outTensorAttr->name);
    Specification::InnerProductLayerParams *innerProductParams = layer->mutable_innerproduct();
    innerProductParams->set_inputchannels(1);
    innerProductParams->set_outputchannels(1);

    // set weight
    auto* weights = innerProductParams->mutable_weights();
    if (areWeightsQuantized) {
        auto *quant_params = weights->mutable_quantization();
        quant_params->set_numberofbits(1);
        auto *linear_quant_params = quant_params->mutable_linearquantization();
        linear_quant_params->add_scale(1.0f);
        linear_quant_params->add_bias(0.0f);
        weights->set_int8rawvalue("x01"); // this is [1]
    } else {
        weights->add_floatvalue(1.0);
    }

    // set bias
    innerProductParams->set_hasbias(true);
    auto* bias = innerProductParams->mutable_bias();
    if (isBiasQuantized) {
        auto *quant_params = bias->mutable_quantization();
        quant_params->set_numberofbits(1);
        auto *linear_quant_params = quant_params->mutable_linearquantization();
        linear_quant_params->add_scale(1.0f);
        linear_quant_params->add_bias(0.0f);
        bias->set_int8rawvalue("x01"); // this is [1]
    } else {
        bias->add_floatvalue(1.0);
    }
    
    if (isUpdatable) {
        layer->set_isupdatable(true);
        innerProductParams->mutable_weights()->set_isupdatable(true);
        innerProductParams->mutable_bias()->set_isupdatable(true);
    }

    return neuralNet;
}

Specification::NeuralNetwork* addSoftmaxLayer(Specification::Model& m, const char *name,  const char *input, const char *output) {
    
    auto neuralNet = m.mutable_neuralnetwork();
    auto softmaxLayer = neuralNet->add_layers();
    
    softmaxLayer->set_name(name);
    softmaxLayer->add_input(input);
    softmaxLayer->add_output(output);
    softmaxLayer->mutable_softmax();
    
    return neuralNet;
}

Specification::NeuralNetworkClassifier* buildBasicNeuralNetworkClassifierModel(Specification::Model& m, bool isUpdatable, const TensorAttributes *inTensorAttr, std::vector<std::string> stringClassLabels, std::vector<int64_t> intClassLabels, bool includeBias) {
    
    bool usesStringClassLabels = stringClassLabels.size() > 0;
    
    auto input = m.mutable_description()->add_input();
    auto inputType = new Specification::FeatureType;
    auto multiArray = inputType->mutable_multiarraytype();
    multiArray->mutable_shape()->Add(inTensorAttr->dimension);
    multiArray->set_datatype(::CoreML::Specification::ArrayFeatureType_ArrayDataType_DOUBLE);
    input->set_name(inTensorAttr->name);
    input->set_allocated_type(inputType);
    
    auto output = m.mutable_description()->add_output();
    auto outputType = new Specification::FeatureType;
    if (usesStringClassLabels) {
        outputType->mutable_stringtype();
    } else {
        outputType->mutable_int64type();
    }
    
    output->set_name("predictedClass");
    output->set_allocated_type(outputType);
    
    output = m.mutable_description()->add_output();
    outputType = new Specification::FeatureType;
    auto dictionary = outputType->mutable_dictionarytype();
    if (usesStringClassLabels) {
        dictionary->mutable_stringkeytype();
    } else {
        dictionary->mutable_int64keytype();
    }
    
    output->set_name("classProbabilities");
    output->set_allocated_type(outputType);

    m.mutable_description()->set_predictedfeaturename("predictedClass");
    m.mutable_description()->set_predictedprobabilitiesname("classProbabilities");
    
    auto classifier = m.mutable_neuralnetworkclassifier();
    
    if (usesStringClassLabels) {
        for (auto className : stringClassLabels) {
            classifier->mutable_stringclasslabels()->add_vector(className);
        }
    } else {
        for (auto className : intClassLabels) {
            classifier->mutable_int64classlabels()->add_vector(className);
        }
    }
    
    // Add inner product layer
    auto innerProductLayer = classifier->add_layers();
    innerProductLayer->set_name("inner_product");
    innerProductLayer->add_input(inTensorAttr->name);
    innerProductLayer->add_output("intermediateOutput");
    
    uint64_t C_in = (uint64_t)inTensorAttr->dimension;
    uint64_t C_out = 0;
    
    if (usesStringClassLabels) {
        C_out = (uint64_t)stringClassLabels.size();
    } else {
        C_out = (uint64_t)intClassLabels.size();
    }
    
    auto innerProductParams = innerProductLayer->mutable_innerproduct();
    innerProductParams->set_inputchannels(C_in);
    innerProductParams->set_outputchannels(C_out);
    
    ::google::protobuf::RepeatedField<float>* weightsWrite = innerProductParams->mutable_weights()->mutable_floatvalue();
    weightsWrite->Resize((int)C_in * (int)C_out, 0.0);
    float *destination_weights = weightsWrite->mutable_data();
    for (uint64_t i = 0; i < C_in; i++) {
        for (uint64_t j = 0; j < C_out; j++) {
            float random = float(rand())/float(RAND_MAX);
            destination_weights[i * C_out + j] = random;
        }
    }
    
    if (includeBias) {
        innerProductParams->set_hasbias(true);
        
        ::google::protobuf::RepeatedField<float>* biasWrite = innerProductParams->mutable_bias()->mutable_floatvalue();
        biasWrite->Resize((int)1 * (int)C_out, 0.0);
        float *destination_bias = biasWrite->mutable_data();
        for (uint64_t i = 0; i < 1; i++) {
            for (uint64_t j = 0; j < C_out; j++) {
                float random = float(rand())/float(RAND_MAX);
                destination_bias[i * C_out + j] = random;
            }
        }
    }
    
    // Add inner product layer
    auto softmaxLayer = classifier->add_layers();
    softmaxLayer->set_name("softmax");
    softmaxLayer->add_input("intermediateOutput");
    softmaxLayer->add_output("scoreVector");
    softmaxLayer->mutable_softmax();
    
    if (isUpdatable) {
        m.set_isupdatable(true);
        innerProductLayer->set_isupdatable(true);
        innerProductParams->mutable_weights()->set_isupdatable(true);
        
        if (includeBias) {
            innerProductParams->mutable_bias()->set_isupdatable(true);
        }
        
        addCategoricalCrossEntropyLoss(m, classifier, "cross_entropy_loss", "scoreVector", "target");
        
        addLearningRate(classifier, Specification::Optimizer::kSgdOptimizer, 0.7f, 0.0f, 1.0f);
        addMiniBatchSize(classifier, Specification::Optimizer::kSgdOptimizer, 32, 1, 100, {16, 32, 64, 128});
        addEpochs(classifier, 100, 1, 100, std::set<int64_t>());
        addShuffleAndSeed(classifier, 2019, 0, 2019, std::set<int64_t>());

        m.mutable_description()->clear_traininginput();
        
        for (auto feature : m.description().input()) {
            auto trainingInput = m.mutable_description()->mutable_traininginput()->Add();
            trainingInput->CopyFrom(feature);
        }
        for (auto feature : m.description().output()) {
            if (feature.name() == m.description().predictedfeaturename()) {
                auto trainingInput = m.mutable_description()->mutable_traininginput()->Add();
                trainingInput->CopyFrom(feature);
            }
        }
    }
    
    return classifier;
}

Specification::KNearestNeighborsClassifier* buildBasicNearestNeighborClassifier(Specification::Model& m, bool isUpdatable, const TensorAttributes *inTensorAttr, const char *outTensorName) {
    auto inTensor = m.mutable_description()->add_input();
    inTensor->set_name(inTensorAttr->name);
    auto inTensorShape = inTensor->mutable_type()->mutable_multiarraytype();
    inTensorShape->set_datatype(Specification::ArrayFeatureType_ArrayDataType_FLOAT32);
    for (int i = 0; i < inTensorAttr->dimension; i++) {
        inTensorShape->add_shape(1);
    }
    
    auto outTensor = m.mutable_description()->add_output();
    outTensor->set_name(outTensorName);
    outTensor->mutable_type()->mutable_stringtype();
    
    m.mutable_description()->set_predictedfeaturename(outTensorName);
    
    m.set_specificationversion(MLMODEL_SPECIFICATION_VERSION_IOS13);
    
    auto nearestNeighborClassifier = m.mutable_knearestneighborsclassifier();
    int numberOfNeighbors = 3;
    nearestNeighborClassifier->mutable_numberofneighbors()->mutable_set()->add_values(numberOfNeighbors);
    nearestNeighborClassifier->mutable_numberofneighbors()->set_defaultvalue(numberOfNeighbors);

    auto nearestNeighborIndex = nearestNeighborClassifier->mutable_nearestneighborsindex();
    
    nearestNeighborIndex->mutable_singlekdtreeindex()->set_leafsize(30);
    nearestNeighborIndex->mutable_squaredeuclideandistance();
    
    nearestNeighborIndex->set_numberofdimensions(inTensorAttr->dimension);
    
    auto floatVector = nearestNeighborIndex->add_floatsamples();
    auto pointVector = floatVector->mutable_vector();
    for (int i = 0; i < inTensorAttr->dimension; i++) {
        pointVector->Add((float)i);
    }
    
    nearestNeighborClassifier->mutable_uniformweighting();
    nearestNeighborClassifier->mutable_stringclasslabels()->add_vector(std::string("zero"));
    
    if (isUpdatable) {
        m.set_isupdatable(true);
    }
    
    return nearestNeighborClassifier;
}

Specification::Pipeline* buildEmptyPipelineModel(Specification::Model& m, bool isUpdatable, const TensorAttributes *inTensorAttr, const TensorAttributes *outTensorAttr) {
    auto inTensor = m.mutable_description()->add_input();
    inTensor->set_name(inTensorAttr->name);
    auto inTensorShape = inTensor->mutable_type()->mutable_multiarraytype();
    inTensorShape->set_datatype(Specification::ArrayFeatureType_ArrayDataType_FLOAT32);
    for (int i = 0; i < inTensorAttr->dimension; i++) {
        inTensorShape->add_shape(1);
    }
    
    auto outTensor = m.mutable_description()->add_output();
    outTensor->set_name(outTensorAttr->name);
    auto outTensorShape = outTensor->mutable_type()->mutable_multiarraytype();
    outTensorShape->set_datatype(Specification::ArrayFeatureType_ArrayDataType_FLOAT32);
    for (int i = 0; i < outTensorAttr->dimension; i++) {
        outTensorShape->add_shape(1);
    }
    
    m.set_specificationversion(MLMODEL_SPECIFICATION_VERSION_IOS13);
    
    auto pipeline = m.mutable_pipeline();
    
    if (isUpdatable) {
        m.set_isupdatable(true);
    }
    
    return pipeline;
}

Specification::Pipeline* buildEmptyPipelineModelWithStringOutput(Specification::Model& m, bool isUpdatable, const TensorAttributes *inTensorAttr, const char *outTensorName) {
    auto inTensor = m.mutable_description()->add_input();
    inTensor->set_name(inTensorAttr->name);
    auto inTensorShape = inTensor->mutable_type()->mutable_multiarraytype();
    inTensorShape->set_datatype(Specification::ArrayFeatureType_ArrayDataType_FLOAT32);
    for (int i = 0; i < inTensorAttr->dimension; i++) {
        inTensorShape->add_shape(1);
    }
    
    auto outTensor = m.mutable_description()->add_output();
    outTensor->set_name(outTensorName);
    outTensor->mutable_type()->mutable_stringtype();
    
    m.set_specificationversion(MLMODEL_SPECIFICATION_VERSION_IOS13);
    
    auto pipeline = m.mutable_pipeline();
    
    if (isUpdatable) {
        m.set_isupdatable(true);
    }
    
    return pipeline;
}

void addCategoricalCrossEntropyLossWithSoftmaxAndSGDOptimizer(Specification::Model& m, const char *softmaxInputName) {
    
    auto neuralNets = m.mutable_neuralnetwork();

    // set a softmax layer
    auto softmaxLayer = neuralNets->add_layers();
    softmaxLayer->set_name("softmax");
    softmaxLayer->add_input(softmaxInputName);
    softmaxLayer->add_output("softmax_out");
    softmaxLayer->mutable_softmax();

    addCategoricalCrossEntropyLoss(m, neuralNets, "cross_entropy_loss_layer", "softmax_out", "target");
    
    addLearningRate(neuralNets, Specification::Optimizer::kSgdOptimizer, 0.7f, 0.0f, 1.0f);
    addMiniBatchSize(neuralNets, Specification::Optimizer::kSgdOptimizer, 10, 5, 100, std::set<int64_t>());
    addEpochs(neuralNets, 100, 1, 100, std::set<int64_t>());
    addShuffleAndSeed(neuralNets, 2019, 0, 2019, std::set<int64_t>());
}

void createSimpleNeuralNetworkClassifierModel(Specification::Model *spec, const char *inputName, const char *outputName) {
    
    spec->set_specificationversion(MLMODEL_SPECIFICATION_VERSION);

    auto *topIn = spec->mutable_description()->add_input();
    topIn->set_name(inputName);
    auto *shape = topIn->mutable_type()->mutable_multiarraytype();
    shape->add_shape(3);
    topIn->mutable_type()->mutable_multiarraytype()->set_datatype(::CoreML::Specification::ArrayFeatureType_ArrayDataType_INT32);

    auto *out = spec->mutable_description()->add_output();
    out->set_name(outputName);
    auto *outshape = out->mutable_type()->mutable_multiarraytype();
    outshape->add_shape(3);
    out->mutable_type()->mutable_multiarraytype()->set_datatype(::CoreML::Specification::ArrayFeatureType_ArrayDataType_INT32);

    auto *out2 = spec->mutable_description()->add_output();
    out2->set_name("features");
    out2->mutable_type()->mutable_stringtype();

    auto *out3 = spec->mutable_description()->add_output();
    out3->set_name("probs");
    out3->mutable_type()->mutable_dictionarytype();
    out3->mutable_type()->mutable_dictionarytype()->mutable_stringkeytype();

    std::string featureName = "features";
    spec->mutable_description()->set_predictedfeaturename(featureName);
    std::string probsName = "probs";
    spec->mutable_description()->set_predictedprobabilitiesname(probsName);

    const auto nn = spec->mutable_neuralnetworkclassifier();
    auto labels = nn->mutable_stringclasslabels();
    labels->add_vector("label1");

    Specification::NeuralNetworkLayer *innerProductLayer = nn->add_layers();
    innerProductLayer->add_input(inputName);
    innerProductLayer->add_output("middle");
    Specification::InnerProductLayerParams *innerProductParams = innerProductLayer->mutable_innerproduct();
    innerProductParams->set_hasbias(false);
    innerProductParams->set_inputchannels(3);
    innerProductParams->set_outputchannels(3);
    for (int i = 0; i < 9; i++) {
        innerProductParams->mutable_weights()->add_floatvalue(1.0);
    }

    Specification::NeuralNetworkLayer *innerProductLayer2 = nn->add_layers();
    innerProductLayer2->add_input("middle");
    innerProductLayer2->add_output(outputName);
    Specification::InnerProductLayerParams *innerProductParams2 = innerProductLayer2->mutable_innerproduct();
    innerProductParams2->set_hasbias(false);
    innerProductParams2->set_inputchannels(3);
    innerProductParams2->set_outputchannels(3);
    for (int i = 0; i < 9; i++) {
        innerProductParams2->mutable_weights()->add_floatvalue(1.0);
    }
}

void createSimpleFeatureVectorizerModel(Specification::Model *spec, const char *outputName, Specification::ArrayFeatureType_ArrayDataType arrayType, int inputSize) {
    
    spec->set_specificationversion(MLMODEL_SPECIFICATION_VERSION);
    
    auto interface = spec->mutable_description();
    auto metadata = interface->mutable_metadata();
    metadata->set_shortdescription(std::string("Testing serialization"));
    
    for (int d = 0; d < inputSize; d++) {
        auto input = interface->add_input();
        auto inputType = new Specification::FeatureType;
        inputType->mutable_doubletype();
        input->set_name("input" + std::to_string(d));
        input->set_allocated_type(inputType);
    }
    
    auto output = interface->add_output();
    auto outputType = new Specification::FeatureType;
    outputType->mutable_multiarraytype()->mutable_shape()->Add(inputSize);
    outputType->mutable_multiarraytype()->set_datatype(arrayType);
    output->set_name(outputName);
    output->set_allocated_type(outputType);
    
    auto featurVectorizerModel = spec->mutable_featurevectorizer();
    
    for (int d = 0; d < inputSize; d++) {
        auto inputlist = featurVectorizerModel->add_inputlist();
        inputlist->set_inputcolumn("input" + std::to_string(d));
        inputlist->set_inputdimensions(1);
    }
    
}


