/*
 * Copyright 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 net.aavalla.utils.jobpool.parameters.FieldAnnotations;
import compid.data.FullDataset;
import compid.data.tuples.DatasetInfo;
import compid.data.tuples.Peptide;
import compid.data.tuples.PeptideHit;
import compid.data.tuples.ProtHitPeptHitJoin;
import compid.data.tuples.ProtSeq;
import compid.data.tuples.Protein;
import compid.data.tuples.ProteinHit;
import compid.data.tuples.ProteinHit.GroupMemberStatus;
import compid.data.tuples.Query;
import compid.importers.ParagonXmlImporter.ParagonXmlImporterParams;
import java.io.IOException;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import java.util.Deque;
import java.util.HashMap;
import java.util.List;
import java.util.NoSuchElementException;
import nu.xom.Attribute;
import nu.xom.Attribute.Type;
import nu.xom.Builder;
import nu.xom.Document;
import nu.xom.Element;
import nu.xom.NodeFactory;
import nu.xom.Nodes;
import nu.xom.ParsingException;

/**
 * Imports Paragon XML data including protein groups. Note that for each
 * unassigned peptide hit, only the best one for each query will be included.
 * Also queries without any peptide hit are included. The behaviour is alike
 * in all the importers.
 * 
 * @author Lari Natri
 */
public class ParagonXmlImporter extends DatasetFileImporter<ParagonXmlImporterParams> {

    public class ParagonXmlImporterParams extends DatasetFileImporterParams {

        @FieldAnnotations(
            name="Include Unassigned Peptides",
            description="Import also peptides which aren't assigned to any proteins. \n" +
                        "This should be chosen if one wishes to create peptide reports \n" +
                        "for all peptides. All queries are imported in any case.")
        private boolean includeUnassigned = true;

        public ParagonXmlImporterParams() {
            this.setDescription("Imports a dataset from Paragon in .xml format");
        }
    }

    public ParagonXmlImporter() {
        super();
        this.setName("Dataset importer for Paragon (xml)");
    }

    private static final String FILE_EXT = ".xml"; //NOI18N

    @Override
    protected FullDataset getAsDataSet(ParagonXmlImporterParams p) throws ImportException {

        compid.data.tuples.DatasetInfo cpDSInfo
                = new DatasetInfo(p.getName(),p.getDescr(),p.getDB(),p.getFile().getName(),
                                  DatasetInfo.AlgType.PARAGON, null, new Date());

        if (p.includeUnassigned) {
            cpDSInfo.setParams("Unassigned Peptides = "+p.includeUnassigned);
        }

        XOMStreamingParagonDatasetFactory factory = new XOMStreamingParagonDatasetFactory(cpDSInfo, p);
        Builder builder = new Builder(factory);

        try {
            builder.build(p.getFile());
        }
        catch (ParsingException ex) {
            ex.printStackTrace();
            throw new ImportException(p.getFile()+" is not a well-formed ProteinPilot 3 XML document");
        }
        catch (IOException ex) {
            throw new ImportException(ex);
        }
        if (factory.dataset.getProtHitPeptHitJoins().size()<1) {
            throw new ImportException("No protein hit data found in "+p.getFile());
        }
        return factory.dataset;


    }

    @Override
    public String getFileNameExtension() {
        return FILE_EXT;
    }

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

    private enum ParagonElementType {
        Spectrum, Match, Protein, Protein2Match,
        ProteinGroup, GroupMember, ParMinUnused, Other
    }

    private class ParagonElement extends Element {
        public final ParagonElementType elementType;
        public ParagonElement(String name, String namespace) {
            super(name.trim().toLowerCase(),namespace);
            name = this.getLocalName();
            if (name.equals("spectrum")) { //NOI18N
                this.elementType = ParagonElementType.Spectrum;
            } else if (name.equals("match")) { //NOI18N
                this.elementType = ParagonElementType.Match;
            } else if (name.equals("protein")) { //NOI18N
                this.elementType = ParagonElementType.Protein;
            } else if (name.equals("protein2match")) { //NOI18N
                this.elementType = ParagonElementType.Protein2Match;
            } else if (name.equals("proteingroup")) { //NOI18N
                this.elementType = ParagonElementType.ProteinGroup;
            } else if (name.equals("groupmember")) { //NOI18N
                this.elementType = ParagonElementType.GroupMember;
            } else if (name.equals("ui_min_unused_protscore")) { //NOI18N
                this.elementType = ParagonElementType.ParMinUnused;
            } else {
                this.elementType = ParagonElementType.Other;
            }
        }
    }


    private class XOMStreamingParagonDatasetFactory extends NodeFactory {

        private class PeptideHitInfo {
            private final Peptide pep;
            private final double conf;
            private final String matchId;

            private PeptideHitInfo(Peptide peptide, double confidence, String matchId) {
                this.pep = peptide;
                this.conf = confidence;
                this.matchId = matchId;
            }
            private final Peptide getPeptide() { return this.pep; }
            private final double getConf() { return this.conf; }
            private final String getMatchId() { return this.matchId; }
        }

        private class ProteinHitInfo {

            private final Protein prot;
            private final double scoreUnused;
            private final double scoreTotal;
            private final List<PeptideHit> pepHits;

            private ProteinHitInfo(Protein prot, double scoreUnused,
                                   double scoreTotal, Collection<String> pephitIds) {
                this.prot = prot;
                this.scoreUnused = scoreUnused;
                this.scoreTotal = scoreTotal;
                this.pepHits = new ArrayList<PeptideHit>(pephitIds.size());
                for (String id : pephitIds) {
                    PeptideHit ph = XOMStreamingParagonDatasetFactory.this.idsToPepMatches.get(id);
                    if (ph!=null) {
                        this.pepHits.add(ph);
                    } else {
                        System.out.println("DEBUG: Peptide match id ("+id+") not found from id " +
                                "table (size: "+idsToPepMatches.size()+") in Paragon XML import");
                    }
                }
            }

            public final Protein getProt() { return this.prot; }
            public final double getScoreUnused() { return this.scoreUnused; }
            public final double getScoreTotal() { return this.scoreTotal; }
            public final List<PeptideHit> getPepHits() { return this.pepHits; }
        }


        private final Nodes empty = new Nodes();
        private final FullDataset dataset = new FullDataset();
        private final DatasetInfo datasetInfo;
        private final boolean includeUnassigned; // inc unassigned peptides
        private final Deque<ParagonElement> stack = new ArrayDeque<ParagonElement>();
        private final HashMap<String,PeptideHit> idsToPepMatches = new HashMap<String,PeptideHit>(512);
        private final HashMap<String,ProteinHitInfo> accToProteinHits = new HashMap<String,ProteinHitInfo>(128);
        private final List<PeptideHitInfo> currPepMatches = new ArrayList<PeptideHitInfo>();
        private final List<String> currProtPepMatchIds = new ArrayList<String>();
        private final HashMap<Query,PeptideHit> queryToBestHit = new HashMap<Query,PeptideHit>(512);
        private PeptideHitInfo bestPeptideForThisQuery = null;
        private double bestPeptideForThisQueryScore = Double.NEGATIVE_INFINITY;
        private int protGroupCounter = 1; // one-based
        private StringBuilder paramStrB = new StringBuilder();

        public XOMStreamingParagonDatasetFactory(final DatasetInfo dsi, 
                final ParagonXmlImporterParams p)
        {
            this.datasetInfo = dsi;
            this.includeUnassigned = p.includeUnassigned;
        }


        @Override
        public ParagonElement startMakingElement(String name, String namespace) {
            ParagonElement element = new ParagonElement(name, namespace);
            stack.push(element);
            return element;
        }

        @Override
        public Nodes finishMakingElement(Element element) {
            ParagonElement e;
            try {
                e = stack.pop();
                if (stack.isEmpty()) return new Nodes(element);
            } catch (NoSuchElementException ex) {
                System.out.println("DEBUG: no such element in xml");
                return empty;
            }
            if (e.elementType.equals(ParagonElementType.Other)) {
                return empty;
            } else if (e.elementType.equals(ParagonElementType.Match)) {
                Peptide pep = new Peptide(e.getAttributeValue("seq"));
                PeptideHitInfo phi = new PeptideHitInfo(pep,Double.parseDouble(e.getAttributeValue("confidence")),e.getAttributeValue("id","http://www.w3.org/XML/1998/namespace"));
                currPepMatches.add(phi);

                double score = Double.parseDouble(e.getAttributeValue("score"));
                if (score>bestPeptideForThisQueryScore) {
                    this.bestPeptideForThisQuery = phi;
                    this.bestPeptideForThisQueryScore = score;
                }
            } else if (e.elementType.equals(ParagonElementType.Spectrum)) {
                try {
                    Query q = new Query(this.datasetInfo,e.getAttributeValue("id","http://www.w3.org/XML/1998/namespace"),
                            Double.parseDouble(e.getAttributeValue("precursormass")),
                            Double.parseDouble(e.getAttributeValue("charge")));
                    dataset.addQuery(q);
                    for (PeptideHitInfo pephitinfo : currPepMatches) {
                        PeptideHit ph = new PeptideHit(pephitinfo.getPeptide(),q,pephitinfo.getConf());
                        //dataset.addPeptideHit(ph);
                        if (this.includeUnassigned && pephitinfo.equals(bestPeptideForThisQuery)) {
                            this.queryToBestHit.put(q, ph);
                        }
                        idsToPepMatches.put(pephitinfo.getMatchId(), ph);
                    }
                } catch (Exception ex) {
                    ex.printStackTrace();
                }
                currPepMatches.clear();
                this.bestPeptideForThisQuery = null;
                this.bestPeptideForThisQueryScore = Double.NEGATIVE_INFINITY;
            } else if (e.elementType.equals(ParagonElementType.Protein2Match)) {
                currProtPepMatchIds.add(e.getAttributeValue("match_id"));
            } else if (e.elementType.equals(ParagonElementType.Protein)) {
                ProtSeq protSeq = null;
                String seqStr = e.getAttributeValue("sequence");
                if (seqStr!=null && seqStr.length()>0) {
                    protSeq = new ProtSeq(seqStr);
                }
                Protein prot = new Protein(e.getAttributeValue("id","http://www.w3.org/XML/1998/namespace"),
                                           e.getAttributeValue("name"),
                                           e.getAttributeValue("species"),
                                           null,protSeq);
                ProteinHitInfo phi = new ProteinHitInfo(prot,
                        Double.parseDouble(e.getAttributeValue("unused_protscore")),
                        Double.parseDouble(e.getAttributeValue("protscore")),
                        currProtPepMatchIds);
                this.currProtPepMatchIds.clear();
                this.accToProteinHits.put(prot.getAccession(), phi);

            } else if (e.elementType.equals(ParagonElementType.GroupMember)) {
                ProteinHitInfo phi = this.accToProteinHits.get(e.getAttributeValue("protein"));
                GroupMemberStatus status = null;
                String memberType = e.getAttributeValue("type").trim().toLowerCase();
                if (memberType.equals("winner")) {
                    status = GroupMemberStatus.WINNER;
                } else if (memberType.equals("sequencesubset")) {
                    status = GroupMemberStatus.SUBSEQSUBSET;
                } else if (memberType.equals("spectrasubset")) {
                    status = GroupMemberStatus.SUBPEPSUBSET;
                } else if (memberType.equals("isoformwinner")) {
                    status = GroupMemberStatus.SUBISOFORMWINNER;
                } else {
                    status = GroupMemberStatus.SUBUNKNOWN;
                }
                ProteinHit prothit = new ProteinHit(phi.getProt(),this.datasetInfo,
                        this.protGroupCounter,status,phi.getScoreUnused(),phi.getScoreTotal());
                for (PeptideHit pephit : phi.getPepHits()) {
                    this.dataset.addProteinHitPeptideHitJoin(new ProtHitPeptHitJoin(prothit, pephit));
                    if (this.includeUnassigned) {
                        this.queryToBestHit.remove(pephit.getQuery());
                    }
                }
            } else if (e.elementType.equals(ParagonElementType.ProteinGroup)) {
                ++this.protGroupCounter;
            }

            return empty;
        }

        @Override
        public Nodes makeAttribute(String name, String URI, String value, Type type) {
            ParagonElement e = stack.peekFirst();
            if (e==null || e.elementType.equals(ParagonElementType.Other)) {
                return empty;
            }

            //System.out.println("name="+name+" URI="+URI+" value="+value+" type="+type);
            e.addAttribute(new Attribute(name.trim().toLowerCase(),URI,value,type));
            return empty;
        }

        @Override
        public Nodes makeComment(String data) {
            return empty;
        }

        @Override
        public Nodes makeDocType(String rootElementName, String publicID, String systemID) {
            return empty;
        }

        @Override
        public Nodes makeProcessingInstruction(String target, String data) {
            return empty;
        }

        @Override
        public Nodes makeText(String data) {
//            ParagonElement e = stack.peek();
//            if (e!=null && e.elementType.equals(ParagonElementType.ParMinUnused)) {
//                if (this.paramStrB.length()>0) {
//                    this.paramStrB.append("; ");
//                }
//                this.paramStrB.append("Min Unused Protscore = "+data);
//            }
            return empty;
        }

        @Override
        public void finishMakingDocument(Document document) {
//            this.datasetInfo.setParams(this.paramStrB.toString());
            if (this.includeUnassigned) {
                for (PeptideHit ph : this.queryToBestHit.values()) {
                    this.dataset.addPeptideHit(ph);
                }
            }
            super.finishMakingDocument(document);
        }



    }


}
