/**
 * Copyright (c) 2011, The University of Southampton and the individual contributors.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 *   *  Redistributions of source code must retain the above copyright notice,
 *      this list of conditions and the following disclaimer.
 *
 *   *  Redistributions in binary form must reproduce the above copyright notice,
 *      this list of conditions and the following disclaimer in the documentation
 *      and/or other materials provided with the distribution.
 *
 *   *  Neither the name of the University of Southampton nor the names of its
 *      contributors may be used to endorse or promote products derived from this
 *      software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
/**
 *
 */
package org.openimaj.demos.sandbox.audio;

import java.io.DataOutputStream;
import java.io.File;
import java.io.FileFilter;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

import org.apache.commons.io.filefilter.WildcardFileFilter;
import org.kohsuke.args4j.Argument;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;
import org.kohsuke.args4j.Option;
import org.openimaj.audio.AudioAnnotator.AudioAnnotatorType;
import org.openimaj.audio.SampleChunk;
import org.openimaj.audio.conversion.MultichannelToMonoProcessor;
import org.openimaj.audio.features.MFCC;
import org.openimaj.feature.DoubleFV;
import org.openimaj.feature.FeatureExtractor;
import org.openimaj.io.IOUtils;
import org.openimaj.ml.annotation.Annotated;
import org.openimaj.ml.annotation.AnnotatedObject;
import org.openimaj.ml.training.IncrementalTrainer;
import org.openimaj.video.xuggle.XuggleAudio;

/**
 *      A trainer for the speech detector. Uses MFCCs to train a KNN classifier.
 *
 *      @author David Dupplaw (dpd@ecs.soton.ac.uk)
 *  @created 6 Dec 2012
 *      @version $Author$, $Revision$, $Date$
 */
public class SpeechDetectorTrainer
{
        /**
         *      Feature extractor for extracting MFCC features from a sample chunk.
         *      It assumes the sample chunk is mono, and will only return the MFCC
         *      features from the first channel.
         *
         *      @author David Dupplaw (dpd@ecs.soton.ac.uk)
         *  @created 6 Mar 2013
         */
        protected static class MFCCFeatureExtractor
                implements FeatureExtractor<DoubleFV,SampleChunk>
        {
                private final MFCC mfcc = new MFCC();

                @Override
                public DoubleFV extractFeature( final SampleChunk object )
                {
                        final double[] d = this.mfcc.calculateMFCC( object.getSampleBuffer() )[0];
                        return new DoubleFV(d);
                }
        }

        /**
         *      Options for the tool
         *
         *      @author David Dupplaw (dpd@ecs.soton.ac.uk)
         *  @created 6 Dec 2012
         *      @version $Author$, $Revision$, $Date$
         */
        protected static class Options
        {
                /** Directory containing input files */
                @Argument(required=true,usage="Directory containing positive examples",metaVar="INPUT-DIR")
                public File speechDir;

                /** Output filename for the classifier */
                @Argument(required=true,usage="Classifier output file",index=1,metaVar="CLASSIFIER-OUTPUT-FILE")
                public File outputFile;

                /** Directory containing negative input files */
                @Option(name="--negative",aliases="-n",usage="Directory containing negative examples")
                public File nonSpeechDir;

                /** The trainer to use */
                @Option(name="--annotator",aliases="-a",usage="Classifier type (default:KNN)")
                public AudioAnnotatorType trainer = AudioAnnotatorType.KNN;

                /** The filename filter to use */
                @Option(name="--filter",aliases="-f",usage="Filename filter (default:*.wav,*.mp3)")
                public List<String> filePattern = new ArrayList<String>(
                                Arrays.asList( new String[]{"*.wav","*.mp3"}) );

                /** Whether to recurse subdirectories */
                @Option(name="--recurse",aliases="-R",usage="Recurse subdirectories" )
                public boolean recurseSubdirectories = false;

                @Option(name="--limit",aliases="-l",usage="Limit number of example files")
                public int limitNumber = -1;
        }

        /** The options being used in this instance of the trainer */
        private final Options options;

        /** The number of files processed so far by this class */
        // Note this means this class should only be used to process
        // one directory at a time (don't call process from multiple Threads)
        private int numProcessed;

        /** MFCC Calculator */
        private final MFCC mfcc = new MFCC();

        /**
         *      Constructor for the trainer tool
         *      @param options the options to use.
         * @throws FileNotFoundException
         */
        public SpeechDetectorTrainer( final Options options )
                        throws FileNotFoundException
        {
                this.options = options;

                // Create the trainer
                final IncrementalTrainer<Annotated<DoubleFV, String>> t =
                                options.trainer.getAnnotator();

                // Add the speech data to the space
                this.processDirectory( options.speechDir, true, t, options.limitNumber );

                // Add the non-speech stuff to the space
                if( options.nonSpeechDir != null )
                        this.processDirectory( options.nonSpeechDir, false, t, options.limitNumber );

                // Write the classifier to a file
                try
                {
                        IOUtils.write( t, new DataOutputStream( new FileOutputStream( options.outputFile ) ) );
                }
                catch( final IOException e )
                {
                        e.printStackTrace();
                }
        }

        /**
         *      Processes the input directory as a set of files whose speech content
         *      is given by the boolean parameter
         *      @param inputDir The input directory
         *      @param speechFiles Whether the files contain speech
         *      @param t The trainer to use
         *      @param limit The number of files to limit to
         *      @throws FileNotFoundException if the input directory cannot be found
         */
        private void processDirectory( final File inputDir, final boolean speechFiles,
                        final IncrementalTrainer<Annotated<DoubleFV, String>> t, final int limit )
                        throws FileNotFoundException
        {
                if( !inputDir.exists() )
                        throw new FileNotFoundException( inputDir+" does not exist." );

                System.out.println( "Entering directory "+inputDir );

                // Instantiate our filename filter
                final FileFilter fileFilter = new WildcardFileFilter(
                                this.options.filePattern );

                // Go through all the files in the directory
                int chunkCount = 0;
                final File[] files = inputDir.listFiles( fileFilter );
                this.numProcessed = 0;
                for( final File file: files )
                {
                        if( limit > 0 && this.numProcessed >= limit )
                                break;

                        if( file.isDirectory() )
                        {
                                if( this.options.recurseSubdirectories )
                                        this.processDirectory( file, speechFiles, t, limit );

                                continue;
                        }

                        System.out.println( "Processing "+file );

                        try
                        {
                                int fileChunkCount = 0;

                                // Create an audio object for the input file
                                final XuggleAudio xa = new XuggleAudio( file );
                                final MultichannelToMonoProcessor mtm =
                                                new MultichannelToMonoProcessor( xa );

                                // Loop through all the chunks in the audio file
                                SampleChunk sc = null;
                                while( (sc = mtm.nextSampleChunk()) != null )
                                {
                                        // Calculate the MFCC for this frame.
                                        final double[][] calculatedMFCC =
                                                        this.mfcc.calculateMFCC( sc.getSampleBuffer() );

                                        // Create an annotated object that says that this sample chunk
                                        // either does or doesn't represent speech
                                        final AnnotatedObject<DoubleFV,String> o =
                                                        new AnnotatedObject<DoubleFV, String>(
                                                                        new DoubleFV(calculatedMFCC[0]),
                                                                        speechFiles?"Speech":"Non-Speech" );

                                        // Now train on that data. Training will involve extracting
                                        // a feature and inserting it into some feature space.
                                        t.train( Collections.singleton( o ) );

                                        // Show a little counter thingy
                                        if( fileChunkCount % 1000 == 0 )
                                                System.out.print( fileChunkCount+"..." );

                                        // count how many items we've trained on
                                        fileChunkCount++;
                                }

                                chunkCount += fileChunkCount;
                                this.numProcessed++;
                        }
                        catch( final Exception e )
                        {
                                e.printStackTrace();
                        }
                        finally
                        {
                        }
                }

                System.out.println( "Trained on "+chunkCount+" sample frames.");
        }

        /**
         *      Returns the trainer used for this trainer
         *      @return The trainer
         */
        public AudioAnnotatorType getTrainer()
        {
                return this.options.trainer;
        }



        // ======================================================================
        /**
         *      Parses the command line arguments to create an options object.
         *      @param args
         */
        private static Options parseArgs( final String args[] )
        {
                final Options o = new Options();
                final CmdLineParser p = new CmdLineParser( o );
                try
                {
                        p.parseArgument( args );
                }
                catch( final CmdLineException e )
                {
                System.err.println( e.getMessage() );
                System.err.println( "java SpeechDetectorTrainer INPUT-DIR CLASSIFIER-OUTPUT-FILE");
                p.printUsage( System.err );
                System.exit(1);
                }

                return o;
        }

        /**
         *
         *      @param args
         */
        public static void main( final String[] args )
        {
                try
                {
                        final Options options = SpeechDetectorTrainer.parseArgs( args );
                        new SpeechDetectorTrainer( options );
                }
                catch( final FileNotFoundException e )
                {
                        e.printStackTrace();
                }
        }
}