/*
 *                    BioJava development code
 *
 * This code may be freely distributed and modified under the
 * terms of the GNU Lesser General Public Licence.  This should
 * be distributed with the code.  If you do not have a copy,
 * see:
 *
 *      http://www.gnu.org/copyleft/lesser.html
 *
 * Copyright for this code is held jointly by the individual
 * authors.  These should be listed in @author doc comments.
 *
 * For more information on the BioJava project and its aims,
 * or to join the biojava-l mailing list, visit the home page
 * at:
 *
 *      http://www.biojava.org/
 *
 * Created on 2012-12-01
 *
 */

package org.biojava.nbio.structure;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.TreeMap;

import org.biojava.nbio.structure.io.mmcif.chem.PolymerType;
import org.biojava.nbio.structure.io.mmcif.chem.ResidueType;
import org.biojava.nbio.structure.io.mmcif.model.ChemComp;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A map from {@link ResidueNumber ResidueNumbers} to ATOM record positions in a PDB file.
 *
 * <p>To use:
 *
 * <pre>
 * Atom[] atoms = new AtomCache().getAtoms("1w0p");
 * AtomPositionMap map = new AtomPositionMap(atoms);
 * ResidueNumber start = new ResidueNumber("A", 100, null);
 * ResidueNumber end = map.getEnd("A");
 * int pos = map.getPosition(start);
 * int length = map.calcLength(start, end);
 * </pre>
 *
 * <p>Note: The getLength() methods were introduced in BioJava 4.0.0 to replace
 * the calcLength methods. The new method returns the number of residues between
 * two residues, inclusive, whereas the previous method returned 1 less than that.
 * @author dmyerstu
 */
public class AtomPositionMap {

        private static final Logger logger = LoggerFactory.getLogger(AtomPositionMap.class);

        private HashMap<ResidueNumber, Integer> hashMap;
        private TreeMap<ResidueNumber, Integer> treeMap;


        /**
         * Used as a Predicate to indicate whether a particular Atom should be mapped
         */
        public static interface GroupMatcher {
                boolean matches(Group group);
        }

        /**
         * Matches CA atoms of protein groups
         */
        public static final GroupMatcher AMINO_ACID_MATCHER = new GroupMatcher() {
                @Override
                public boolean matches(Group group) {
                        if( group == null )
                                return false;
                        ChemComp chem = group.getChemComp();
                        if(chem == null)
                                return false;
                        // Get polymer type
                        PolymerType polyType = chem.getPolymerType();
                        if( polyType == null) {
                                ResidueType type = chem.getResidueType();
                                if(type != null ) {
                                        polyType = type.getPolymerType();
                                }
                        }
                        if( polyType == null ) {
                                return false;
                        }

                        return PolymerType.PROTEIN_ONLY.contains(polyType)
                                        && group.hasAtom(StructureTools.CA_ATOM_NAME);
                }
        };

        /**
         * Matches all atoms
         */
        public static final GroupMatcher ANYTHING_MATCHER = new GroupMatcher() {
                @Override
                public boolean matches(Group group) {
                        return true;
                }
        };

        /**
         * A map that is sorted by its values. Used for the treemap
         *
         * @author dmyerstu
         *
         * @param <T>
         *            The key type
         * @param <V>
         *            The value type
         */
        private static class ValueComparator<T, V extends Comparable<V>> implements Comparator<T> {

                private Map<T, V> map;

                public ValueComparator(Map<T, V> map) {
                        this.map = map;
                }

                @Override
                public int compare(T o1, T o2) {
                        return map.get(o1).compareTo(map.get(o2));
                }

        }

        /**
         * Creates a new AtomPositionMap containing peptide alpha-carbon atoms
         *
         * @param atoms
         */
        public AtomPositionMap(Atom[] atoms) {
                this(atoms, AMINO_ACID_MATCHER);
        }

        /**
         * Creates a new AtomPositionMap containing only atoms matched by {@code matcher}.
         *
         * If multiple atoms are present from a group, the first atom encountered will
         * be used.
         * @param atoms
         */
        public AtomPositionMap(Atom[] atoms, GroupMatcher matcher) {
                hashMap = new HashMap<ResidueNumber, Integer>();
                for (int i = 0; i < atoms.length; i++) {
                        Group group = atoms[i].getGroup();
                        ResidueNumber rn = group.getResidueNumber();
                        if (matcher.matches(group)) {
                                if (!hashMap.containsKey(rn)) {
                                        hashMap.put(rn, i);
                                }
                        }
                }
                Comparator<ResidueNumber> vc = new ValueComparator<ResidueNumber, Integer>(hashMap);
                treeMap = new TreeMap<ResidueNumber, Integer>(vc);
                treeMap.putAll(hashMap);
        }

        /**
         * Creates a new AtomPositionMap containing representative atoms
         * from a structure.
         * @param s
         */
        public AtomPositionMap(Structure s) {
                this(StructureTools.getRepresentativeAtomArray(s));
        }

        /**
         * Calculates the number of residues of the specified chain in a given range, inclusive.
         * @param positionA index of the first atom to count
         * @param positionB index of the last atom to count
         * @param startingChain Case-sensitive chain
         * @return The number of atoms between A and B inclusive belonging to the given chain
         */
        public int getLength(int positionA, int positionB, String startingChain) {

                int positionStart, positionEnd;
                if (positionA <= positionB) {
                        positionStart = positionA;
                        positionEnd = positionB;
                } else {
                        positionStart = positionB;
                        positionEnd = positionA;
                }

                int count = 0;
                // Inefficient search
                for (Map.Entry<ResidueNumber, Integer> entry : treeMap.entrySet()) {
                        if (entry.getKey().getChainId().equals(startingChain)
                                        && positionStart <= entry.getValue()
                                        && entry.getValue() <= positionEnd)
                        {
                                count++;
                        }
                }
                return count;
        }


        /**
         * Calculates the number of residues of the specified chain in a given range.
         * Will return a negative value if the start is past the end.
         * @param positionStart index of the first atom to count
         * @param positionEnd index of the last atom to count
         * @param startingChain Case-sensitive chain
         * @return The number of atoms from A to B inclusive belonging to the given chain
         */
        public int getLengthDirectional(int positionStart, int positionEnd, String startingChain) {
                int count = getLength(positionStart,positionEnd,startingChain);
                if(positionStart <= positionEnd) {
                        return count;
                } else {
                        return -count;
                }
        }

        /**
         * Calculates the number of atoms between two ResidueNumbers, inclusive. Both residues
         * must belong to the same chain.
         * @param start First residue
         * @param end Last residue
         * @return The number of atoms from A to B inclusive
         * @throws IllegalArgumentException if start and end are on different chains,
         *  or if either of the residues doesn't exist
         */
        public int getLength(ResidueNumber start, ResidueNumber end) {
                if( ! start.getChainId().equals(end.getChainId())) {
                        throw new IllegalArgumentException(String.format(
                                        "Chains differ between %s and %s. Unable to calculate length.",
                                        start,end));
                }
                Integer startPos = getPosition(start);
                Integer endPos = getPosition(end);
                if(startPos == null) {
                        throw new IllegalArgumentException("Residue "+start+" was not found.");
                }
                if(endPos == null) {
                        throw new IllegalArgumentException("Residue "+start+" was not found.");
                }
                return getLength(startPos, endPos, start.getChainId());
        }

        /**
         * Calculates the number of atoms between two ResidueNumbers, inclusive. Both residues
         * must belong to the same chain.
         * Will return a negative value if the start is past the end.
         * @param start First residue
         * @param end Last residue
         * @return The number of atoms from A to B inclusive
         * @throws IllegalArgumentException if start and end are on different chains,
         *  or if either of the residues doesn't exist
         */
        public int getLengthDirectional(ResidueNumber start, ResidueNumber end) {
                if( ! start.getChainId().equals(end.getChainId())) {
                        throw new IllegalArgumentException(String.format(
                                        "Chains differ between %s and %s. Unable to calculate length.",
                                        start,end));
                }
                Integer startPos = getPosition(start);
                Integer endPos = getPosition(end);
                if(startPos == null) {
                        throw new IllegalArgumentException("Residue "+start+" was not found.");
                }
                if(endPos == null) {
                        throw new IllegalArgumentException("Residue "+start+" was not found.");
                }
                return getLengthDirectional(startPos, endPos, start.getChainId());
        }


        public NavigableMap<ResidueNumber, Integer> getNavMap() {
                return treeMap;
        }

        /**
         * Gets the 0-based index of residueNumber to the matched atoms
         * @param residueNumber
         * @return The position of the ATOM record in the PDB file corresponding to
         * the {@code residueNumber}, or null if the residueNumber was not found
         */
        public Integer getPosition(ResidueNumber residueNumber) {
                return hashMap.get(residueNumber);
        }

        /**
         * @param chainId
         * @return The first {@link ResidueNumber} of the specified chain (the one highest down in the PDB file)
         */
        public ResidueNumber getFirst(String chainId) {
                Map.Entry<ResidueNumber,Integer> entry = treeMap.firstEntry();
                while (true) {
                        if (entry.getKey().getChainId().equals(chainId)) return entry.getKey();
                        entry = treeMap.higherEntry(entry.getKey());
                        if (entry == null) return null;
                }
        }

        /**
         * @param chainId
         * @return The last {@link ResidueNumber} of the specified chain (the one farthest down in the PDB file)
         */
        public ResidueNumber getLast(String chainId) {
                Map.Entry<ResidueNumber,Integer> entry = treeMap.lastEntry();
                while (true) {
                        if (entry.getKey().getChainId().equals(chainId)) return entry.getKey();
                        entry = treeMap.lowerEntry(entry.getKey());
                        if (entry == null) return null;
                }
        }

        /**
         * @return The first {@link ResidueNumber} of any chain (the one farthest down in the PDB file)
         */
        public ResidueNumber getFirst() {
                return treeMap.firstKey();
        }

        /**
         * @return The last {@link ResidueNumber} of any chain (the one farthest down in the PDB file)
         */
        public ResidueNumber getLast() {
                return treeMap.lastKey();
        }

        /**
         * Returns a list of {@link ResidueRange ResidueRanges} corresponding to this entire AtomPositionMap.
         */
        public List<ResidueRangeAndLength> getRanges() {
                String currentChain = "";
                ResidueNumber first = null;
                ResidueNumber prev = null;
                List<ResidueRangeAndLength> ranges = new ArrayList<ResidueRangeAndLength>();
                for (ResidueNumber rn : treeMap.keySet()) {
                        if (!rn.getChainId().equals(currentChain)) {
                                if (first != null) {
                                        ResidueRangeAndLength newRange = new ResidueRangeAndLength(currentChain, first, prev, this.getLength(first, prev));
                                        ranges.add(newRange);
                                }
                                first = rn;
                        }
                        prev = rn;
                        currentChain = rn.getChainId();
                }
                ResidueRangeAndLength newRange = new ResidueRangeAndLength(currentChain, first, prev, this.getLength(first, prev));
                ranges.add(newRange);
                return ranges;
        }

        /**
         * Trims a residue range so that both endpoints are contained in this map.
         * @param rr
         * @param map
         * @return
         */
        public ResidueRangeAndLength trimToValidResidues(ResidueRange rr) {
                ResidueNumber start = rr.getStart();
                ResidueNumber end = rr.getEnd();
                String chain = rr.getChainId();
                // Add chainId
                if(start.getChainId() == null) {
                        start = new ResidueNumber(chain,start.getSeqNum(),start.getInsCode());
                }
                if(end.getChainId() == null) {
                        end = new ResidueNumber(chain,end.getSeqNum(),end.getInsCode());
                }
                // Check that start and end are present in the map.
                // If not, try to find the next valid residue
                // (terminal residues sometimes lack CA atoms, so they don't appear)
                Integer startIndex = getPosition(start);
                if( startIndex == null) {
                        // Assume that the residue numbers are sequential
                        // Find startIndex such that the SeqNum is bigger than start's seqNum
                        for(ResidueNumber key :  treeMap.keySet()) {
                                if( !key.getChainId().equals(chain) )
                                        continue;
                                if( start.getSeqNum() <= key.getSeqNum() ) {
                                        start = key;
                                        startIndex = getPosition(key);
                                        break;
                                }
                        }
                        if( startIndex == null ) {
                                logger.error("Unable to find Residue {} in AtomPositionMap, and no plausible substitute.",start);
                                return null;
                        } else {
                                logger.warn("Unable to find Residue {}, so substituting {}.",rr.getStart(),start);
                        }
                }
                Integer endIndex = getPosition(end);
                if( endIndex == null) {
                        // Assume that the residue numbers are sequential
                        // Find startIndex such that the SeqNum is bigger than start's seqNum
                        for(ResidueNumber key :  treeMap.descendingKeySet()) {
                                if( !key.getChainId().equals(chain) )
                                        continue;
                                Integer value = getPosition(key);
                                if( value < startIndex ) {
                                        // start is before the end!
                                        break;
                                }
                                if( end.getSeqNum() >= key.getSeqNum() ) {
                                        end = key;
                                        endIndex = value;
                                        break;
                                }
                        }
                        if( endIndex == null ) {
                                logger.error("Unable to find Residue {} in AtomPositionMap, and no plausible substitute.",end);
                                return null;
                        } else {
                                logger.warn("Unable to find Residue {}, so substituting {}.",rr.getEnd(),end);
                        }
                }

                // now use those to calculate the length
                // if start or end is null, will throw NPE
                int length = getLength(startIndex, endIndex,chain);

                return new ResidueRangeAndLength(chain, start, end, length);
        }
}