/*
 * Copyright 2009, 2010 Lari Natri <lari.natri@gmail.com>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 * 
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 * 
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */


package compid.importers;

import compid.data.DAOException;
import compid.data.DAOFactory;
import compid.data.UpdateHandler;
import compid.data.tuples.ProtSeq;
import compid.data.tuples.Protein;
import compid.importers.FASTAImporter.FASTAImporterParams;
import compid.importers.ImportException.ExceptionType;
import net.aavalla.utils.jobpool.JobPriority;
import net.aavalla.utils.jobpool.ParameterizedJob;
import net.aavalla.utils.jobpool.parameters.FieldAnnotations;
import net.aavalla.utils.jobpool.parameters.InputFile;
import net.aavalla.utils.jobpool.parameters.validators.InputFileValidator;
import net.aavalla.utils.jobpool.parameters.ParameterCollection;
import net.aavalla.utils.general.FastaUtils;
import net.aavalla.utils.general.FastaUtils.FastaLineType;
import net.aavalla.utils.general.RetrievalHashSet;
import net.aavalla.utils.jobpool.Job;
import net.aavalla.utils.jobpool.JobCategory;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;



/**
 *
 * @author Lari Natri
 */
public class FASTAImporter extends ParameterizedJob<Object,FASTAImporterParams> {

    public class FASTAImporterParams extends ParameterCollection {

        private static final boolean INCLUDE_REFS = true;

        @FieldAnnotations(
            name="Input file (FASTA)",
            description="A collection of protein sequences in FASTA format to be imported",
            validator=InputFileValidator.class)
        private InputFile fastaFile;

        public FASTAImporterParams() {
            this.setDescription("Scans a given protein sequence database in FASTA format and<br/>" +
                                "imports sequences for those proteins in Compid's internal<br/>" +
                                "database which are missing them.");
            this.setHelpID("importseqs");
        }

        public File getFastaFile() { return this.fastaFile; }


    }


    private static final int BATCHSIZE = 30000; // TODO: loop through protdb in batches to conserve memory
    public static final int FILEBUFFERSIZE = 8388608; // 8 MiB
    private static final double MINAARATIOINSEQUENCE = 0.95d;
    private static final int MAXERRONEUSSEQUENCES = 500;


    public FASTAImporter() {
        super();
        this.setName("Protein Sequence importer (FASTA)");
        this.setPriority(JobPriority.LOW);
        this.addCategory(JobCategory.MEMORY);
        this.setSupportsCancellation(true);
    }


    @Override
    protected FASTAImporterParams getNonfinalEditableParameters() {
        return new FASTAImporterParams();
    }


    @Override
    protected Object work(final FASTAImporterParams p) throws Throwable {

        SequencelessProteinFinder job1 = new SequencelessProteinFinder();

        FastaParserParams fpp = new FastaParserParams();
        fpp.protsToUpdate = this.waitSubJob(job1);
        if (this.cancelPending()) return CANCELLEDRESULT;
        fpp.importerParams = p;
        int sequencelessCount = fpp.protsToUpdate.size();
        if (sequencelessCount==0) {
            return "Database contains no proteins without a sequence.";
        }
        FastaParser parser = new FastaParser(fpp);
        RetrievalHashSet<Protein> newSeqs = this.waitSubJob(parser);

        if (this.cancelPending()) return CANCELLEDRESULT;
        if (newSeqs.size()==0) {
            return "No sequences added. Database still contains \n"+
                   sequencelessCount+" proteins without a sequence.";
        }
        SequenceInjecter injecter = new SequenceInjecter(newSeqs);
        this.setSupportsCancellation(false);
        this.waitSubJob(injecter);
        return newSeqs.size()+" protein sequences added. Database now " +
               "contains "+(sequencelessCount-newSeqs.size())+" proteins without sequence.";
    }


    private class SequencelessProteinFinder extends Job<RetrievalHashSet<Protein>,Object> {
        public SequencelessProteinFinder() {
            super();
            this.setName("Sequenceless protein finder");
            this.addCategory(JobCategory.DB_READ);
            this.addConflictingCategory(JobCategory.DB_WRITE);
        }
        @Override
        protected RetrievalHashSet<Protein> work(Object p) {
            return DAOFactory.getInstance().getQueryHandler().getAllProteinsWithoutSequence();
        }
    }

    private class FastaParserParams {
        private RetrievalHashSet<Protein> protsToUpdate;
        private FASTAImporterParams importerParams;
    }

    private class FastaParser extends Job<RetrievalHashSet<Protein>,FastaParserParams> {

        private FastaParser(FastaParserParams p) {
            super(p);
            this.setName("FASTA parser");
            this.setSupportsCancellation(true);
            this.addCategory(JobCategory.MEMORY);
            this.addCategory(JobCategory.FILESYSTEM);
        }

        @Override
        protected RetrievalHashSet<Protein> work(FastaParserParams p) throws Throwable {
            RetrievalHashSet<Protein> protsToUpdate = p.protsToUpdate;
            FASTAImporterParams params = p.importerParams;
            System.out.println("DEBUG: Starting update now. There are "+protsToUpdate.size()+" proteins needing sequences");

            FileReader fileReader = null;
            BufferedReader bufReader = null;
            int erroneusSequencesCount = 0;
            int fastaHeaderCount = 0;

            if (!params.getFastaFile().exists()) {
                throw new ImportException(ExceptionType.FILENOTFOUND,
                                          params.getFastaFile().toString());
            }
            if (!params.getFastaFile().canRead()) {
                throw new ImportException(ExceptionType.FILENOTREADABLE,
                                          params.getFastaFile().toString());
            }

            try {
                RetrievalHashSet<Protein> updatedProts = new RetrievalHashSet<Protein>(512);
                String nextFullHeaderLine;

                long fileSize = params.getFastaFile().length();
                long curBytes = 0;

                fileReader = new FileReader(params.getFastaFile());
                bufReader = new BufferedReader(fileReader, FILEBUFFERSIZE);

                do { // read file until first header is found
                    nextFullHeaderLine = bufReader.readLine();
                    if (nextFullHeaderLine==null) { // EOF reached
                        throw new ImportException(ExceptionType.FASTA_INCORRECTFORMAT,
                                "File "+params.getFastaFile().getName()+" is not in " +
                                "FASTA format. \n(Didn't find any headers.)");
                    }
                    curBytes+=nextFullHeaderLine.length();
                    nextFullHeaderLine = nextFullHeaderLine.trim();
                } while (FastaUtils.getFastaLineType(nextFullHeaderLine)!=FastaLineType.HEADER);


                while (nextFullHeaderLine!=null && protsToUpdate.isEmpty()==false) {
                    this.reportProgress((int)Math.round(100*curBytes/(double)fileSize));

                    boolean entryComplete = false;
                    StringBuilder sequenceBuilder = new StringBuilder();
                    String fullHeaderLine = nextFullHeaderLine;

                    while (!entryComplete) { // read lines until entry is complete
                        if (this.cancelPending()) throw CANCELLEDRESULT;
                        String line = bufReader.readLine();
                        
                        if (line==null) { // EOF reached
                            entryComplete = true;
                            nextFullHeaderLine = null; // to signify EOF
                            curBytes = fileSize;
                        } else { // a line was read
                            curBytes+=line.length();
                            line = line.trim();
                            switch(FastaUtils.getFastaLineType(line)) {
                                case HEADER:
                                    nextFullHeaderLine = line;
                                    entryComplete = true;
                                    break;
                                case SEQUENCE:
                                    sequenceBuilder.append(line);
                                    break;
                                case COMMENT:
                                case OTHER:
                                default:
                                    break;
                            }
                        } // else
                    }

                    String[] headers = fullHeaderLine.split("\u0001"); // split by Ctrl-A
                    headers[0] = headers[0].substring(1); // remove '>' from beginning

                    RetrievalHashSet<String> accs = new RetrievalHashSet<String>(headers.length*3);
                    ProtSeq protSeq = null;

                    for (int h = 0; h < headers.length; ++h) { // loop through headers for the same sequence
                        String descr = "";
                        String idFull = "";

                        String[] idAndDescrParts = headers[h].split("\\s",2); // split by first whitespace
                        if (idAndDescrParts.length >= 1) {
                            idFull = idAndDescrParts[0];
                            if (idAndDescrParts.length >= 2) {
                                descr = idAndDescrParts[1];
                            }
                        }

                        accs.add(idFull); // add the full unsplit id also

                        String[] idSplitted = idFull.split("[|]"); // split by | (at least NCBInr has ginumber + source id)

                        for (int s = 0; s < idSplitted.length-1; ++s) {
                            if (idSplitted[s].equals("gi")) {
                                if (s+1 > idSplitted.length-1) continue;
                                accs.add(idSplitted[s]+"|"+idSplitted[s+1]);
                                ++s;
                            } else if (idSplitted[s].equals("gb")
                                    || idSplitted[s].equals("emb")
                                    || idSplitted[s].equals("dbj")
                                    || idSplitted[s].equals("ref")
                                    || idSplitted[s].equals("bbs")
                                    || idSplitted[s].equals("lcl")
                                    || idSplitted[s].equals("tpg")
                                    || idSplitted[s].equals("tpd")
                                    || idSplitted[s].equals("tpe")) {
                                if (s+1 > idSplitted.length-1) continue;
                                accs.add(idSplitted[s]+"|"+idSplitted[s+1]);
                                ++s;
                            } else if (idSplitted[s].equals("pir")
                                    || idSplitted[s].equals("prf")) {
                                if (s+2 > idSplitted.length-1) continue;
                                accs.add(idSplitted[s]+"||"+idSplitted[s+2]);
                                s+=2;
                            } else if (idSplitted[s].equals("sp")
                                    || idSplitted[s].equals("pdb")) {
                                if (s+1 > idSplitted.length-1) continue;
                                accs.add(idSplitted[s]+"|"+idSplitted[s+1]); // without "name"/chain
                                ++s;
                                if (s+1 > idSplitted.length-1) continue;
                                accs.add(idSplitted[s-1]+"|"+idSplitted[s]+"|"+idSplitted[s+1]); // with "name"/chain
                                ++s;
                            } else if (idSplitted[s].equals("pat")
                                    || idSplitted[s].equals("gnl")) {
                                if (s+2 > idSplitted.length-1) continue;
                                accs.add(idSplitted[s]+"|"+idSplitted[s+1]+"|"+idSplitted[s+2]);
                                s+=2;
                            } else {
                                if (idSplitted[s].length()>2) {
                                    accs.add(idSplitted[s]);
                                }
                                if (s != 0 && s < idSplitted.length-1) { // s==0: added before for
                                    StringBuilder unknownIdBuilder = new StringBuilder();
                                    unknownIdBuilder.append(idSplitted[s]);
                                    for (int s2=s+1; s2<idSplitted.length-1; ++s2) {
                                        unknownIdBuilder.append("|"+idSplitted[s2]);
                                        ++s;
                                    }
                                    accs.add(unknownIdBuilder.toString());
                                }
                            } // else

                        } // for idSplitted

                        //System.out.println("-------------------------------------");
                        for (String acc : accs) {
                            //System.out.println("DEBUG: testing "+acc);
                            // proteins are equal if their accession codes are
                            Protein tmpProt = new Protein(acc,descr,"<tmp>");
                            if (protsToUpdate.contains(tmpProt)) {
                                if (protSeq==null) { // create only once for each distinct sequence
                                    protSeq = new ProtSeq(sequenceBuilder.toString());
                                    double ratioOfAA = ((double)protSeq.getLength()) / ((double)sequenceBuilder.length());
                                    if (ratioOfAA < MINAARATIOINSEQUENCE && protSeq.getLength()<ProtSeq.MAXSEQUENCELENGTH) { // too many non-AA symbols in sequence
                                        System.out.println("DEBUG: sequence has "+(sequenceBuilder.length()-protSeq.getLength())+" non-AA symbols (they were removed)");
                                        System.out.println("       "+sequenceBuilder.toString());
                                        ++erroneusSequencesCount;
                                        if (erroneusSequencesCount > MAXERRONEUSSEQUENCES) {
                                            throw new ImportException(ExceptionType.FASTA_INCORRECTFORMAT,
                                                    "File "+params.getFastaFile().getName()+" is not in " +
                                                    "FASTA format. (Too many erroneus sequences.)");
                                        }
                                    }
                                }

                                Protein updatableProt = protsToUpdate.getKey(tmpProt);
                                protsToUpdate.remove(updatableProt);
                                updatableProt.setProtSeq(protSeq);
                                updatedProts.add(updatableProt); // added only if doesn't already contain same accession

                            }
                        } // for accs

                    } // for headers with same seq

                    ++fastaHeaderCount;

                    if (fastaHeaderCount % 100000 == 0) {
                        System.out.println("DEBUG: "+updatedProts.size()+" updates found, "+fastaHeaderCount+" headers processed");
                    }

                } // end: while there is something to read in the FASTA file

                System.out.println("DEBUG: "+updatedProts.size()+" updates found, "+fastaHeaderCount+" headers processed");

                return updatedProts;


            } catch (FileNotFoundException ex) {
                throw new ImportException(ExceptionType.FILENOTFOUND,
                                          params.getFastaFile().getAbsolutePath());
            } catch (IOException ex) {
                throw new ImportException(ex);
            } catch (DAOException ex) {
                throw new ImportException(ex);
            } finally {
                try {
                    bufReader.close();
                    fileReader.close();
                } catch (Exception ex) {} // never mind this
            }

        }

    }

    private class SequenceInjecter extends Job<Integer,RetrievalHashSet<Protein>> {
        private SequenceInjecter(RetrievalHashSet<Protein> newSeqs) {
            super(newSeqs);
            this.setName("Sequence injecter");
            this.addCategory(JobCategory.DB_WRITE);
            this.addConflictingCategory(JobCategory.DB_WRITE);
            this.addConflictingCategory(JobCategory.DB_READ);
        }
        @Override protected Integer work(RetrievalHashSet<Protein> newSeqs) {
            UpdateHandler uh = DAOFactory.getInstance().getUpdateHandler();
            return uh.updateSequenceRefs((RetrievalHashSet<Protein>)newSeqs);
        }
    }



}
