package org.deeplearning4j.optimize.listeners;

import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.eval.Evaluation;
import org.deeplearning4j.eval.IEvaluation;
import org.deeplearning4j.exception.DL4JInvalidInputException;
import org.deeplearning4j.nn.api.Model;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.api.IterationListener;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.MultiDataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;

import java.util.concurrent.atomic.AtomicLong;

/**
 *
 *
 * @author raver119@gmail.com
 */
@Slf4j
public class EvaluativeListener implements IterationListener {
    private ThreadLocal<AtomicLong> iterationCount = new ThreadLocal<>();
    private int frequency;

    private DataSetIterator dsIterator;
    private MultiDataSetIterator mdsIterator;
    private DataSet ds;
    private MultiDataSet mds;

    private IEvaluation[] evaluations;

    public EvaluativeListener(@NonNull DataSetIterator iterator, int frequency) {
        this(iterator, frequency, new Evaluation());
    }

    public EvaluativeListener(@NonNull MultiDataSetIterator iterator, int frequency) {
        this(iterator, frequency, new Evaluation());
    }

    public EvaluativeListener(@NonNull DataSetIterator iterator, int frequency, IEvaluation... evaluations) {
        this.dsIterator = iterator;
        this.frequency = frequency;
        this.evaluations = evaluations;
    }

    public EvaluativeListener(@NonNull MultiDataSetIterator iterator, int frequency, IEvaluation... evaluations) {
        this.mdsIterator = iterator;
        this.frequency = frequency;
        this.evaluations = evaluations;
    }

    public EvaluativeListener(@NonNull DataSet dataSet, int frequency) {
        this(dataSet, frequency, new Evaluation());
    }

    public EvaluativeListener(@NonNull MultiDataSet multiDataSet, int frequency) {
        this(multiDataSet, frequency, new Evaluation());
    }

    public EvaluativeListener(@NonNull DataSet dataSet, int frequency, IEvaluation... evaluations) {
        this.ds = dataSet;
        this.frequency = frequency;
        this.evaluations = evaluations;
    }

    public EvaluativeListener(@NonNull MultiDataSet multiDataSet, int frequency, IEvaluation... evaluations) {
        this.mds = multiDataSet;
        this.frequency = frequency;
        this.evaluations = evaluations;
    }

    /**
     * Get if listener invoked
     */
    @Override
    public boolean invoked() {
        return false;
    }

    /**
     * Change invoke to true
     */
    @Override
    public void invoke() {

    }

    /**
     * Event listener for each iteration
     *
     * @param model     the model iterating
     * @param iteration the iteration
     */
    @Override
    public void iterationDone(Model model, int iteration) {
        if (iterationCount.get() == null)
            iterationCount.set(new AtomicLong(0));

        if (iterationCount.get().getAndIncrement() % frequency != 0)
            return;

        log.info("Starting evaluation on iteration {}", iteration);
        if (model instanceof MultiLayerNetwork) {
            if (dsIterator != null) {
                ((MultiLayerNetwork) model).doEvaluation(dsIterator, evaluations);
            } else if (ds != null) {
                for (IEvaluation evaluation : evaluations)
                    evaluation.eval(ds.getLabels(), ((MultiLayerNetwork) model).output(ds.getFeatureMatrix()));
            }
        } else if (model instanceof ComputationGraph) {
            if (dsIterator != null) {
                ((ComputationGraph) model).doEvaluation(dsIterator, evaluations);
            } else if (mdsIterator != null) {
                ((ComputationGraph) model).doEvaluation(mdsIterator, evaluations);
            } else if (ds != null) {
                for (IEvaluation evaluation : evaluations)
                    evaluation.eval(ds.getLabels(), ((ComputationGraph) model).output(ds.getFeatureMatrix())[0]);
            } else if (mds != null) {
                for (IEvaluation evaluation : evaluations)
                    evaluation.eval(mds.getLabels()[0], ((ComputationGraph) model).output(mds.getFeatures())[0]);
            }
        } else throw new DL4JInvalidInputException("Model is unknown: " + model.getClass().getCanonicalName());

        // TODO: need to print out something now. Callback probably
    }
}

@@ -133,62 +133,9 @@
 val.
-Evaluation;%0Aimport org.deeplearning4j.eval.IEvaluation
+*
 ;%0Aim
@@ -4532,58 +4532,1080 @@
 DO: 
-need to print out something now. Callback probably
+maybe something better should be used here?%0A        log.info(%22Reporting evaluation results:%22);%0A        for (IEvaluation evaluation: evaluations)%0A            printEvalResults(evaluation);%0A    %7D%0A%0A%0A    protected void printEvalResults(IEvaluation evaluation) %7B%0A        if (evaluation instanceof Evaluation)%0A            log.info(%22%7B%7D:%5Cn%7B%7D%22, evaluation.getClass().getSimpleName(), ((Evaluation) evaluation).stats());%0A        else if (evaluation instanceof ROC)%0A            log.info(%22%7B%7D:%5Cn%7B%7D%22, evaluation.getClass().getSimpleName(), ((ROC) evaluation).calculateAUC());%0A        else if (evaluation instanceof ROCMultiClass)%0A            log.info(%22%7B%7D:%5Cn%7B%7D%22, evaluation.getClass().getSimpleName(), ((ROCMultiClass) evaluation).calculateAverageAUC());%0A        else if (evaluation instanceof ROCBinary)%0A            log.info(%22%7B%7D:%5Cn%7B%7D%22, evaluation.getClass().getSimpleName(), ((ROCBinary) evaluation).stats());%0A        else if (evaluation instanceof RegressionEvaluation)%0A            log.info(%22%7B%7D:%5Cn%7B%7D%22, evaluation.getClass().getSimpleName(), ((RegressionEvaluation) evaluation).stats());
 %0A   

print out evaluation results