/*
 *                    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/
 *
 */
package org.biojava.nbio.structure.symmetry.core;

import org.biojava.nbio.structure.*;
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;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * Extracts information about all the chains in a structure, including
 * chain Ids, sequences, and atoms. Includes both protein and nucleic acid chains.
 */
public class ProteinChainExtractor  {

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

	private Structure structure = null;
	private QuatSymmetryParameters parameters = null;
	private boolean modified = true;
	private int adjustedMinimumSequenceLength = 0;

	private List<Atom[]> cAlphaTrace = new ArrayList<Atom[]>();
	private List<String> chainIds = new ArrayList<String>();
	private List<Integer> modelNumbers = new ArrayList<Integer>();
	private List<String> sequences = new ArrayList<String>();
	private int nucleicAcidChainCount = 0;

	public ProteinChainExtractor(Structure structure, QuatSymmetryParameters parameters) {
		this.structure = structure;
		this.parameters = parameters;
		modified = true;
	}

	public List<Atom[]> getCalphaTraces() {
		run();
		return cAlphaTrace;
	}

	public List<String> getChainIds() {
		run();
		return chainIds;
	}

	public List<Integer> getModelNumbers() {
		run();
		return modelNumbers;
	}

	public List<String> getSequences() {
		run();
		return sequences;
	}

	/**
	 * @return the nucleicAcidChainCount
	 */
	public int getNucleicAcidChainCount() {
		run();
		return nucleicAcidChainCount;
	}

	public int getAdjustedMinimumSequenceLength() {
		run();
		return adjustedMinimumSequenceLength;
	}

	private void run() {
		if (modified) {
			calcAdjustedMinimumSequenceLength();
			extractProteinChains();
			modified = false;
		}
	}

	private void extractProteinChains() {
		int models = 1;
		if (structure.isBiologicalAssembly()) {
			models = structure.nrModels();
		}

		if (parameters.isVerbose()) {
			System.out.println("Protein chains used in calculation:");
			System.out.println("Adjusted minimum sequence length: " + adjustedMinimumSequenceLength);
		}

		for (int i = 0; i < models; i++) {
			for (Chain c : structure.getChains(i)) {
				if (isNucleicAcidChain(c)) {
					nucleicAcidChainCount++;
				} //TODO Should we break here for DNA? If "CA" atoms are present, could cause bugs. -Spencer 9-2015
				Atom[] ca = StructureTools.getAtomCAArray(c);
				ca = retainStandardAminoAcidResidues(ca);

				if (ca.length >= adjustedMinimumSequenceLength) {
					if (parameters.isVerbose()) {
						System.out.println("Chain " + c.getChainID() + " Calpha atoms: " + ca.length + " seqres: " + c.getSeqResSequence());
					}

					cAlphaTrace.add(ca);
					chainIds.add(c.getChainID());
					modelNumbers.add(i);
					sequences.add(replaceQuestionMarks(c.getSeqResSequence()));
				}
			}
		}
	}

	/**
	 * Returns an adapted minimum sequence length. This method ensure that
	 * structure that only have short chains are not excluded by the
	 * minimumSequenceLength cutoff value;
	 * @return
	 */
	private void calcAdjustedMinimumSequenceLength() {
		int models = 1;
		if (structure.isBiologicalAssembly()) {
			models = structure.nrModels();
		}

		int maxLength = Integer.MIN_VALUE;
		int minLength = Integer.MAX_VALUE;

		List<Integer> lengths = new ArrayList<Integer>();
		for (int i = 0; i < models; i++) {
			for (Chain c : structure.getChains(i)) {
				Atom[] ca = StructureTools.getAtomCAArray(c);
				ca = retainStandardAminoAcidResidues(ca);
				if (ca.length >= parameters.getAbsoluteMinimumSequenceLength()) {
					maxLength = Math.max(ca.length, maxLength);
					minLength = Math.min(ca.length, minLength);
					lengths.add(ca.length);
				}
			}
		}

		adjustedMinimumSequenceLength = parameters.getMinimumSequenceLength();

		if (lengths.size() < 2) {
			return;
		}

		double median = 0;
		Collections.sort(lengths);
		if (lengths.size() %2 == 1) {
			int middle = (lengths.size()-1) / 2;
			median = lengths.get(middle);
		} else {
			int middle2 = lengths.size()/2;
			int middle1 = middle2-1;
			median = 0.5 * (lengths.get(middle1) + lengths.get(middle2));
		}

		if (minLength >= median * parameters.getMinimumSequenceLengthFraction()) {
			adjustedMinimumSequenceLength = Math.min(minLength, parameters.getMinimumSequenceLength());
		}
	}

	private boolean isNucleicAcidChain(Chain chain) {
		int count = 0;
		for (Group group: chain.getAtomGroups()) {
			PolymerType type = group.getChemComp().getPolymerType();
			if (type != null && (type.equals(PolymerType.dna) || type.equals(PolymerType.rna) || type.equals(PolymerType.dnarna))) {
				count++;
			}
		}
		return count > 3;
	}

	// In some cases "?" are in the sequence string. Example 2WS1. This is caused
	// because the chemical component file YNM doesn't contain a one-letter code.
	private String replaceQuestionMarks(String sequence) {
		return sequence.replaceAll("\\?", "X");
	}

	private Atom[] retainStandardAminoAcidResidues(Atom[] atoms) {
		List<Atom> atomList = new ArrayList<Atom>(atoms.length);
		for (Atom atom: atoms) {
			Group group = atom.getGroup();
			if (group.getPDBName().equalsIgnoreCase("UNK")) {
				continue;
			}
			if (! isAminoAcid(group)) {
				continue;
			}
			atomList.add(atom);
		}
		return atomList.toArray(new Atom[atomList.size()]);
	}

	private boolean isAminoAcid(Group group) {
		ChemComp cc= group.getChemComp();
		if (cc.getResidueType() == null) {
			logger.warn("null residue type for: " + group.getPDBName());
			return false;
		}
		return (cc.getResidueType().equals(ResidueType.lPeptideLinking) || cc.getResidueType().equals(ResidueType.glycine));
	}
}
