//Dstl (c) Crown Copyright 2017
package uk.gov.dstl.baleen.annotators.coreference.impl.enhancers;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

import uk.gov.dstl.baleen.annotators.coreference.impl.data.Mention;
import uk.gov.dstl.baleen.resources.SharedGenderMultiplicityResource;
import uk.gov.dstl.baleen.resources.data.Multiplicity;
import uk.gov.dstl.baleen.types.common.Organisation;
import uk.gov.dstl.baleen.types.language.WordToken;

/**
 * Add multiplicity information to mentions.
 */
public class MultiplicityEnhancer implements MentionEnhancer {

	private static final Map<String, Multiplicity> PRONOUN_MAP = new HashMap<>();
	private final SharedGenderMultiplicityResource multiplicityResource;

	static {
		// See Person

		Arrays.asList("i", "he", "her", "herself", "hers", "her", "him", "himself", "his", "it", "its", "itself", "me",
				"myself", "mine", "my", "one", "oneself", "one's", "she", "yourself").stream()
				.forEach(s -> PRONOUN_MAP.put(s, Multiplicity.SINGULAR));
		Arrays.asList("ours", "our", "ourself", "ourselves", "their", "theirs", "them", "themself", "themselves",
				"they", "us", "we", "yourself", "yourselves").stream()
				.forEach(s -> PRONOUN_MAP.put(s, Multiplicity.PLURAL));
	}

	/**
	 * Constructor for MultiplicityEnhancer
	 */
	public MultiplicityEnhancer(SharedGenderMultiplicityResource multiplicityResource) {
		this.multiplicityResource = multiplicityResource;
	}

	@Override
	public void enhance(Mention mention) {
		switch (mention.getType()) {
		case PRONOUN:
			mention.setMultiplicity(PRONOUN_MAP.getOrDefault(mention.getText().toLowerCase(), Multiplicity.UNKNOWN));
			return;
		case ENTITY:
			mention.setMultiplicity(getEntityMultiplicity(mention));
			break;
		case NP:
			mention.setMultiplicity(getNounPhraseMultiplicity(mention));
			break;
		default:
			return;
		}

		// TODO: Should we always check our resource and then override the multiplicity?

		if (mention.getMultiplicity() == Multiplicity.UNKNOWN) {
			final Multiplicity assignedMultiplicity = multiplicityResource.lookupMultiplicity(mention.getText());
			mention.setMultiplicity(assignedMultiplicity);
		}
	}
	
	private Multiplicity getEntityMultiplicity(Mention mention){
		// Assumed singular, unless organisation
		if (mention.getAnnotation() instanceof Organisation) {
			return Multiplicity.UNKNOWN;
		} else {
			return Multiplicity.SINGULAR;
		}
	}
	
	private Multiplicity getNounPhraseMultiplicity(Mention mention){
		Multiplicity m = Multiplicity.UNKNOWN;
		final WordToken head = mention.getHeadWordToken();
		if (head != null) {
			if ("NNS".equalsIgnoreCase(head.getPartOfSpeech())
					|| "NPS".equalsIgnoreCase(mention.getHeadWordToken().getPartOfSpeech())) {
				m = Multiplicity.PLURAL;
			} else {
				m = Multiplicity.SINGULAR;
			}
		}
		
		return m;
	}

}