package org.deeplearning4j.nearestneighbor.server;

import com.beust.jcommander.JCommander;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.ParameterException;
import org.apache.commons.io.IOUtils;
import org.deeplearning4j.clustering.sptree.DataPoint;
import org.deeplearning4j.clustering.vptree.VPTree;
import org.deeplearning4j.clustering.vptree.VPTreeFillSearch;
import org.deeplearning4j.nearestneighbor.model.*;
import org.jboss.netty.util.internal.ByteBufferUtil;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.serde.base64.Nd4jBase64;
import org.nd4j.serde.binary.BinarySerde;
import play.Mode;
import play.libs.Json;
import play.routing.RoutingDsl;
import play.server.Server;

import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import static play.mvc.Controller.request;
import static play.mvc.Results.badRequest;
import static play.mvc.Results.internalServerError;
import static play.mvc.Results.ok;

/**
 * A rest server for using an
 * {@link VPTree} based on loading an ndarray containing
 * the data points for the path
 * The input values are an {@link CSVRecord}
 * which (based on the input schema) will automatically
 * have their values transformed.
 *
 * @author Adam Gibson
 */
public class NearestNeighborsServer {
    @Parameter(names = {"--ndarrayPath"}, arity = 1, required = true)
    private String ndarrayPath = null;
    @Parameter(names = {"--nearestNeighborsPort"}, arity = 1)
    private int port = 9000;
    @Parameter(names = {"--similarityFunction"}, arity = 1)
    private String similarityFunction = "euclidean";
    @Parameter(names = {"--invert"}, arity = 1)
    private boolean invert = false;

    private Server server;

    public void runMain(String... args) throws Exception {
        JCommander jcmdr = new JCommander(this);

        try {
            jcmdr.parse(args);
        } catch (ParameterException e) {
            //User provides invalid input -> print the usage info
            jcmdr.usage();
            if (ndarrayPath == null)
                System.err.println("Json path parameter is missing.");
            try {
                Thread.sleep(500);
            } catch (Exception e2) {
            }
            System.exit(1);
        }

        final INDArray points = BinarySerde.readFromDisk(new File(ndarrayPath));
        VPTree tree = new VPTree(points, similarityFunction, invert);


        RoutingDsl routingDsl = new RoutingDsl();
        //return the host information for a given id
        routingDsl.POST("/knn").routeTo(FunctionUtil.function0((() -> {
            try {
                NearestNeighborRequest record = Json.fromJson(request().body().asJson(), NearestNeighborRequest.class);
                NearestNeighbor nearestNeighbor =
                                NearestNeighbor.builder().points(points).record(record).tree(tree).build();

                if (record == null)
                    return badRequest(Json.toJson(Collections.singletonMap("status", "invalid json passed.")));

                NearstNeighborsResults results =
                                NearstNeighborsResults.builder().results(nearestNeighbor.search()).build();


                return ok(Json.toJson(results));

            } catch (Exception e) {
                e.printStackTrace();
                return internalServerError();
            }
        })));

        routingDsl.POST("/knnnew").routeTo(FunctionUtil.function0((() -> {
            try {
                Base64NDArrayBody record = Json.fromJson(request().body().asJson(), Base64NDArrayBody.class);
                if (record == null)
                    return badRequest(Json.toJson(Collections.singletonMap("status", "invalid json passed.")));

                INDArray arr = Nd4jBase64.fromBase64(record.getNdarray());
                List<DataPoint> results;
                List<Double> distances;

                if (record.isForceFillK()) {
                    VPTreeFillSearch vpTreeFillSearch = new VPTreeFillSearch(tree, record.getK(), arr);
                    vpTreeFillSearch.search();
                    results = vpTreeFillSearch.getResults();
                    distances = vpTreeFillSearch.getDistances();
                } else {
                    results = new ArrayList<>();
                    distances = new ArrayList<>();
                    tree.search(arr, record.getK(), results, distances);


                }

                List<NearestNeighborsResult> nnResult = new ArrayList<>();
                for (DataPoint dataPoint : results) {
                    nnResult.add(new NearestNeighborsResult(dataPoint.getIndex()));
                }

                NearstNeighborsResults results2 = NearstNeighborsResults.builder().results(nnResult).build();
                return ok(Json.toJson(results2));

            } catch (Exception e) {
                e.printStackTrace();
                return internalServerError();
            }
        })));

        server = Server.forRouter(routingDsl.build(), Mode.DEV, port);


    }

    /**
     * Stop the server
     */
    public void stop() {
        if (server != null)
            server.stop();
    }

    public static void main(String[] args) throws Exception {
        new NearestNeighborsServer().runMain(args);
    }

}

@@ -2417,13 +2417,58 @@
     
-final
+String%5B%5D pathArr = ndarrayPath.split(%22,%22);%0A       
  IND
@@ -2476,64 +2476,233 @@
 rray
+%5B%5D
  points
- = BinarySerde.readFromDisk(new File(ndarrayPath)
+Arr = new INDArray%5BpathArr.length%5D;%0A        for (int i = 0; i %3C pathArr.length; i++)%0A            pointsArr%5Bi%5D = BinarySerde.readFromDisk(new File(pathArr%5Bi%5D));%0A        final INDArray points = Nd4j.concat(0, pointsArr
 );%0A 
@@ -2767,25 +2767,24 @@
 , invert);%0A%0A
-%0A
         Rout

Enable KNN server to load points from more than one NDArray. (#3663)