/*
 *  * Copyright 2016 Skymind, Inc.
 *  *
 *  *    Licensed under the Apache License, Version 2.0 (the "License");
 *  *    you may not use this file except in compliance with the License.
 *  *    You may obtain a copy of the License at
 *  *
 *  *        http://www.apache.org/licenses/LICENSE-2.0
 *  *
 *  *    Unless required by applicable law or agreed to in writing, software
 *  *    distributed under the License is distributed on an "AS IS" BASIS,
 *  *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  *    See the License for the specific language governing permissions and
 *  *    limitations under the License.
 */

package org.datavec.nlp.reader;

import org.datavec.api.conf.Configuration;
import org.datavec.api.writable.IntWritable;
import org.datavec.api.records.reader.impl.FileRecordReader;
import org.datavec.api.split.InputSplit;
import org.datavec.api.vector.Vectorizer;
import org.datavec.api.writable.Writable;
import org.datavec.common.RecordConverter;
import org.datavec.nlp.vectorizer.TfidfVectorizer;
import org.nd4j.linalg.api.ndarray.INDArray;

import java.io.IOException;
import java.util.*;

/**
 * TFIDF record reader (wraps a tfidf vectorizer
 * for delivering labels and conforming to the record reader interface)
 *
 * @author Adam Gibson
 */
public class TfidfRecordReader extends FileRecordReader  {
    private TfidfVectorizer tfidfVectorizer;
    private List<List<Writable>> records = new ArrayList<>();
    private List<Integer> recordLabels = new ArrayList<>();
    private Iterator<Integer> labelIter;
    private Iterator<List<Writable>> recordIter;
    private int numFeatures;
    private boolean initialized = false;


    @Override
    public void initialize(InputSplit split) throws IOException, InterruptedException {
        initialize(new Configuration(),split);
    }

    @Override
    public void initialize(Configuration conf, InputSplit split) throws IOException, InterruptedException {
        super.initialize(conf,split);
        //train  a new one since it hasn't been specified
        if(tfidfVectorizer == null) {
            tfidfVectorizer = new TfidfVectorizer();
            tfidfVectorizer.initialize(conf);
            INDArray ret = tfidfVectorizer.fitTransform(this, new Vectorizer.RecordCallBack() {
                @Override
                public void onRecord(Collection<Writable> record) {
                    Iterator<Writable> writableIterator = record.iterator();
                    //skip the string
                    writableIterator.next();
                    recordLabels.add(writableIterator.next().toInt());
                }
            });

            //clear out old strings
            records.clear();
            for(int i = 0; i< ret.rows(); i++) {
                records.add(RecordConverter.toRecord(ret.getRow(i)));
            }

            //cache the number of features used for each document
            numFeatures = ret.columns();
            labelIter = recordLabels.iterator();
            recordIter = records.iterator();
        }
        else {
            records = new ArrayList<>();

            //the record reader has 2 phases, we are skipping the
            //document frequency phase and just using the super() to get the file contents
            //and pass it to the already existing vectorizer.
            while(hasNext()) {
                Collection<Writable> fileContents = next();
                if(appendLabel)
                    recordLabels.add(new IntWritable(getCurrentLabel()).toInt());
                records.add(RecordConverter.toRecord(tfidfVectorizer.transform(fileContents)));
            }

            labelIter = recordLabels.iterator();
            recordIter = records.iterator();
        }

        if(appendLabel){
            Iterator<List<Writable>> rIter = records.iterator();
            Iterator<Integer> lIter = recordLabels.iterator();
            while(rIter.hasNext()){
                List<Writable> record = rIter.next();
                Integer label = lIter.next();
                record.add(new IntWritable(label));
            }
        }
        this.initialized = true;
    }

    @Override
    public void reset() {
        if(inputSplit == null) throw new UnsupportedOperationException("Cannot reset without first initializing");
        labelIter = recordLabels.iterator();
        recordIter = records.iterator();
    }

    @Override
    public List<Writable> next() {
        if(recordIter == null)
            return super.next();
        List<Writable> record = recordIter.next();
        return record;
    }

    @Override
    public boolean hasNext() {
        //we aren't done vectorizing yet
        if(recordIter == null)
            return super.hasNext();
        return recordIter.hasNext();
    }

    @Override
    public void close() throws IOException {

    }

    @Override
    public void setConf(Configuration conf) {
        this.conf = conf;
    }

    @Override
    public Configuration getConf() {
        return conf;
    }

    public TfidfVectorizer getTfidfVectorizer() {
        return tfidfVectorizer;
    }

    public void setTfidfVectorizer(TfidfVectorizer tfidfVectorizer) {
        if(initialized){
            throw new IllegalArgumentException("Setting TfidfVectorizer after TfidfRecordReader initialization doesn't have an effect");
        }
        this.tfidfVectorizer = tfidfVectorizer;
    }

    public int getNumFeatures() {
        return numFeatures;
    }
}


@@ -5506,12 +5506,202 @@
 s;%0A    %7D
+%0A%0A    public void shuffle()%7B%0A        this.shuffle(new Random());%0A    %7D%0A%0A    public void shuffle(Random random)%7B%0A        Collections.shuffle(this.records, random);%0A        this.reset();%0A    %7D
 %0A%7D%0A%0A

Add shuffle method to TfidfRecordReader