package mhtree.benchmarking;

import cz.muni.fi.disa.similarityoperators.cover.AbstractRepresentation;
import messif.algorithms.AlgorithmMethodException;
import messif.buckets.BucketStorageException;
import messif.objects.LocalAbstractObject;
import messif.objects.util.AbstractObjectList;
import messif.objects.util.RankedAbstractObject;
import messif.objects.util.StreamGenericAbstractObjectIterator;
import messif.operations.Approximate;
import messif.operations.data.BulkInsertOperation;
import messif.operations.query.ApproxKNNQueryOperation;
import messif.operations.query.KNNQueryOperation;
import messif.statistics.Statistics;
import mhtree.InsertType;
import mhtree.MHTree;
import mhtree.MergingMethod;
import mhtree.ObjectToNodeDistance;
import mtree.MTree;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;
import messif.algorithms.Algorithm;
import messif.objects.keys.AbstractObjectKey;
import messif.operations.AnswerType;
import messif.statistics.OperationStatistics;

import static mhtree.ObjectToNodeDistance.AVERAGE;
import static mhtree.ObjectToNodeDistance.FURTHEST;
import static mhtree.ObjectToNodeDistance.NEAREST;

public class RunBenchmark {
    static final Logger log = Logger.getLogger("messif.algorithm");

    public static void main(String[] args) throws IOException, AlgorithmMethodException, InstantiationException, NoSuchMethodException, BucketStorageException, ClassNotFoundException {
        if (args.length != 8 && args.length != 11) {
            throw new IllegalArgumentException("Unexpected number of params");
        }

        boolean isMHtree = switch (args[0]) {
            case "MHULL-TREE" -> true;
            default -> false;
        };
        
        // e.g. messif.objects.impl.ObjectFloatVectorNeuralNetworkL2
        Class<? extends LocalAbstractObject> objClass = (Class<? extends LocalAbstractObject>) Class.forName(args[1]);   
        
//        Statistics.enableGlobally();
        Statistics.disableGlobally();
        AbstractRepresentation.PrecomputedDistances.COMPUTATION_THREADS = 16;

        List<LocalAbstractObject> objects = loadDataset(args[2], objClass);
        int leafCapacity = Integer.parseInt(args[3]);
        int nodeDegree = Integer.parseInt(args[4]);
        List<LocalAbstractObject> queries = loadDataset(args[5], objClass);
        
        InsertType insertType = args[6].equals("INCREMENTAL") ? InsertType.INCREMENTAL : InsertType.GREEDY;

        ObjectToNodeDistance objectToNodeDistance = switch (args[7]) {
            case "FURTHEST" -> FURTHEST;
            case "AVERAGE" -> AVERAGE;
            default -> NEAREST;
        };
        
        final MHTreeConfig cfg = new MHTreeConfig(
                leafCapacity,
                nodeDegree,
                insertType,
                objectToNodeDistance
        );
        final int[] ks = new int[]{1, 3, 5, 10, 20, 50, 100};
//        final int[] ks = new int[]{3};
        if (isMHtree) {
            percentageToRecallMHTree(cfg,
                    objects,
                    Arrays.asList(NEAREST),
                    //Arrays.asList(NEAREST, FURTHEST, AVERAGE),
                    queries, ks);
        } else if (args.length == 8) {      // Original M-tree
            percentageToRecallMTree(cfg,
                    objects,
                    queries, ks,
                    0, 0, null);
        } else {    // Pivoting M-tree
            List<LocalAbstractObject> pivots = loadDataset(args[10], objClass);
            percentageToRecallMTree(cfg,
                    objects,
                    queries, ks,
                    Integer.parseInt(args[8]), Integer.parseInt(args[9]), pivots);
        }
    }

    private static void percentageToRecallMHTree(MHTreeConfig config, List<LocalAbstractObject> objects, 
            final List<ObjectToNodeDistance> distFuncs,
            List<LocalAbstractObject> queries,
            int[] ks) throws BucketStorageException, RuntimeException {

        long buildingStartTimeStamp = System.currentTimeMillis();

        MHTree mhTree = new MHTree.MHTreeBuilder(objects, config.leafCapacity, config.nodeDegree)
                .objectToNodeDistance(config.objectToNodeDistance)
                .mergingMethod(MergingMethod.HULL_BASED_MERGE)
                .build();

        long buildinTime = System.currentTimeMillis() - buildingStartTimeStamp;
        //OperationStatistics.getLocalThreadStatistics().printStatistics();
        mhTree.printStatistics();
        System.out.println("Fat factor: " + mhTree.getFatFactor());
        System.out.println("Building time: " + buildinTime + " msec");

        System.out.println("kNN queries will be executed for k=" + Arrays.toString(ks));
        
        for (ObjectToNodeDistance dist : distFuncs) {
            System.gc();

            mhTree.getNodes().forEach(node -> node.objectToNodeDistance = dist);

            System.out.println("leafCapacity,nodeDegree,objectToNodeDistance,k,percentage,recall (min),recall (avg),recall (med),recall (max), time (msec)");

            Map<String, List<RankedAbstractObject>> kNNResults = prepareGroundTruth(ks, queries, mhTree);

            for (int k : ks) {
                List<ApproxKNNQueryOperation> approxOperations = queries
                        .parallelStream()
                        .map(object -> new ApproxKNNQueryOperation(object, k, 0, Approximate.LocalSearchType.PERCENTAGE, LocalAbstractObject.UNKNOWN_DISTANCE))
                        .collect(Collectors.toList());
                approxOperations
                        .parallelStream()
                        .forEach(op -> op.suppData = new SearchState(mhTree, op));

                double minimalRecall = 0;
                for (int percentage = 0; percentage <= 100; percentage += 5) {
                    final int approxLimit = Math.round((float) objects.size() * (float) percentage / 100f);
                    approxOperations
                            .parallelStream()
                            .filter(op -> ((SearchState) op.suppData).recall < 1d)
                            .forEach(op -> {
                                SearchState searchState = (SearchState) op.suppData;
                                searchState.approximateState.limit = approxLimit;
                                try {
                                    //mhTree.approxKNN(op);
                                    mhTree.executeOperation(op);
                                } catch (AlgorithmMethodException | NoSuchMethodException ex) {
                                }
                                searchState.time += op.getParameter("OperationTime", Long.class);
                                searchState.recall = PerformanceMeasures.measureRecall(op, kNNResults);
                                log.log(Level.INFO, "{0} processed: {1}; Recall: {2}", new Object[]{mhTree.getName(), op.toString(), searchState.recall});
                                // Vlasta: NO op.resetAnswer() here!!! Because we test it incrementally!!!! See the search state of operation! op.resetAnswer();
                            });

                    Stats recallStats = new Stats(
                            approxOperations
                                    .stream()
                                    .map(op -> ((SearchState) op.suppData).recall)
                                    .collect(Collectors.toList())
                    );
                    Stats timeStats = new Stats(
                            approxOperations
                                    .stream()
                                    .map(op -> (double)((SearchState) op.suppData).time)
                                    .collect(Collectors.toList())
                    );

                    System.out.println(String.join(",",
                            String.valueOf(config.leafCapacity),
                            String.valueOf(config.nodeDegree),
                            String.valueOf(dist),
                            String.valueOf(k),
                            String.valueOf(percentage),
                            String.format("%.2f,%.2f,%.2f,%.2f",
                                    recallStats.getMin(),
                                    recallStats.getAverage(),
                                    recallStats.getMedian(),
                                    recallStats.getMax()),
                            String.format("%.2f", timeStats.getAverage())));

                    minimalRecall = recallStats.getMin();
                    if (minimalRecall == 1.0)
                        break;
                }
            }
        }
    }

    private static void percentageToRecallMTree(MHTreeConfig config, List<LocalAbstractObject> objects, 
            List<LocalAbstractObject> queries,
            int[] ks,
            int pmTreeNPD, int pmTreeNHR, List<LocalAbstractObject> pmTreePivots) throws NoSuchMethodException, AlgorithmMethodException, RuntimeException, InstantiationException {

        long buildingStartTimeStamp = System.currentTimeMillis();

        System.out.println("Building M-tree with NPD=" + pmTreeNPD + " and NHR=" + pmTreeNHR);
        MTree mTree;
        if (pmTreeNHR == 0 && pmTreeNHR == 0)
            mTree = new MTree(config.nodeDegree, config.leafCapacity);
        else        // Pivoting M-tree
            mTree = new MTree(config.nodeDegree, config.leafCapacity, 
                                Math.max(pmTreeNPD, pmTreeNHR), pmTreePivots.iterator(), pmTreeNPD, pmTreeNHR);

        Collections.shuffle(objects);
        System.out.println("Shuffling objects done. First is now " + objects.get(0).getLocatorURI());

        BulkInsertOperation opIns = new BulkInsertOperation(objects);

        //mTree.setMaxSpanningTree(1);
        mTree.insert(opIns);
        long buildingTime = System.currentTimeMillis() - buildingStartTimeStamp;

        mTree.markObjectsWithBucketIds();
        mTree.printStatistics();
        //mTree.checkConsistency();
        System.out.println("Fat factor: " + mTree.getFatFactor());
        System.out.println("Building time: " + buildingTime + " msec");

        System.out.println("kNN queries will be executed for k=" + Arrays.toString(ks));
        
        System.out.println("leafCapacity,nodeDegree,objectToNodeDistance,k,percentage,recall (min),recall (avg),recall (med),recall (max), time (msec)");

        Map<String, List<RankedAbstractObject>> kNNResults = prepareGroundTruth(ks, queries, mTree);
        

//        int numberOfQueries = queries.size();
        for (int k : ks) {
            double minimalRecall = 0;
//            for (int percentage = 55; percentage <= 55; percentage += 1) {
//            for (int percentage = 0; percentage <= 5; percentage += 1) {
            for (int percentage = 0; percentage <= 100; percentage += 5) {
                final int approxLimit = percentage;
                List<ApproxKNNQueryOperation> approxOperations = queries
                        .parallelStream()
                        .map(object -> new ApproxKNNQueryOperation(object, k, AnswerType.ORIGINAL_OBJECTS, approxLimit, Approximate.LocalSearchType.PERCENTAGE, LocalAbstractObject.UNKNOWN_DISTANCE))
                        .collect(Collectors.toList());
                approxOperations
                        .parallelStream()
                        .forEach(op -> op.suppData = new SearchInfoState());
                
                approxOperations
                        .parallelStream()
                        .forEach(op -> {
                            SearchInfoState searchState = (SearchInfoState) op.suppData;
                            try {
                                mTree.executeOperation(op);
                            } catch (AlgorithmMethodException | NoSuchMethodException ex) {
                            }
                            searchState.time = op.getParameter("OperationTime", Long.class);
                            searchState.recall = PerformanceMeasures.measureRecall(op, kNNResults);
                            log.log(Level.INFO, "{0} processed: {1}; Recall: {2}", new Object[]{mTree.getName(), op.toString(), searchState.recall});
                            log.log(Level.INFO, "{0} processed: {1}; Answer: {2}", new Object[]{mTree.getName(), op.toString(), iterToString(op.getAnswer())});
                            OperationStatistics.getLocalThreadStatistics().printStatistics();
//                            if (searchState.recall != 1.0) {
//                                mTree.checkConsistency();
//                                try {
//                                    mTree.storeToFile("mtree-bad.bin");
//                                } catch (IOException ex) {
//                                    Logger.getLogger(RunBenchmark.class.getName()).log(Level.SEVERE, null, ex);
//                                }
//                            } else {
//                                try {
//                                    mTree.storeToFile("mtree-ok.bin");
//                                } catch (IOException ex) {
//                                    Logger.getLogger(RunBenchmark.class.getName()).log(Level.SEVERE, null, ex);
//                                }
//                            }
                        });

                Stats recallStats = new Stats(
                        approxOperations
                                .stream()
                                .map(op -> ((SearchInfoState) op.suppData).recall)
                                .collect(Collectors.toList())
                );
                Stats timeStats = new Stats(
                        approxOperations
                                .stream()
                                .map(op -> (double)((SearchInfoState) op.suppData).time)
                                .collect(Collectors.toList())
                );

                System.out.println(String.join(",",
                        String.valueOf(config.leafCapacity),
                        String.valueOf(config.nodeDegree),
                        "",
                        String.valueOf(k),
                        String.valueOf(percentage),
                        String.format(Locale.ENGLISH, "%.2f,%.2f,%.2f,%.2f",
                                recallStats.getMin(),
                                recallStats.getAverage(),
                                recallStats.getMedian(),
                                recallStats.getMax()),
                        String.format(Locale.ENGLISH, "%.2f", timeStats.getAverage())));

                minimalRecall = recallStats.getMin();
                //if (minimalRecall == 1.0)
                //    break;
            }
        }
    }
    
    private static String iterToString(Iterator it) {
        StringBuilder sb = new StringBuilder();
        while (it.hasNext()) {
            if (sb.length() > 0)
                sb.append(", ");
            sb.append(it.next());
        }
        return sb.toString();
    }

    private static Map<String, List<RankedAbstractObject>> prepareGroundTruth(int[] ks, List<LocalAbstractObject> queries, Algorithm alg) {
        int maxK = Arrays.stream(ks).max().getAsInt();
        
        List<KNNQueryOperation> kNNOperations = queries
                .parallelStream()
                .map(object -> new KNNQueryOperation(object, maxK, AnswerType.ORIGINAL_OBJECTS))
                .collect(Collectors.toList());
        Map<String, List<RankedAbstractObject>> kNNResults = kNNOperations
                .parallelStream()
                .map((KNNQueryOperation op) -> {
                    try {
                        alg.executeOperation(op);
                    } catch (AlgorithmMethodException | NoSuchMethodException ex) { }
                    KnnResultPair pair = new KnnResultPair();
                    if (!AbstractObjectKey.class.equals(op.getQueryObject().getObjectKey().getClass()))
                        System.out.println("ERROR: Ground truth query has a non-expected abstract key! " + op.getQueryObject().getObjectKey());
                    pair.id = op.getQueryObject().getLocatorURI();
                    pair.kNNObjects = new ArrayList<>(maxK);
                    for (RankedAbstractObject o : op)
                        pair.kNNObjects.add(o);
                    if (pair.kNNObjects.size() != maxK)
                        System.out.println("ERROR: Ground truth for query " + pair.id + " constains " + pair.kNNObjects.size() + " objects, Expected " + maxK);
                    return pair;
                })
                .collect(Collectors.toMap(entry -> entry.id, entry -> entry.kNNObjects));
        return kNNResults;
    }

    private static List<LocalAbstractObject> loadDataset(String path, Class<? extends LocalAbstractObject> objCls) throws IOException {
        return new AbstractObjectList<>(new StreamGenericAbstractObjectIterator<>(objCls, path));
    }
}
