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

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.fit.descriptor.ConfigurationParameter;
import org.apache.uima.fit.util.JCasUtil;
import org.apache.uima.jcas.JCas;

import com.google.common.base.Strings;
import com.google.common.collect.ImmutableSet;

import uk.gov.dstl.baleen.core.pipelines.orderers.AnalysisEngineAction;
import uk.gov.dstl.baleen.types.semantic.Entity;
import uk.gov.dstl.baleen.types.semantic.ReferenceTarget;
import uk.gov.dstl.baleen.uima.BaleenAnnotator;

/**
 * Find groups of entities of the same type and same value, ignoring case and apostrophes, and point
 * them at the same ReferenceTarget.
 *
 * <p>If two or more existing ReferenceTargets exist within a group, then the group is either
 * ignored or it's ReferenceTargets merged depending on the mergeReferents parameter. If there is
 * one existing ReferenceTarget within a group then that is used as the ReferenceTarget for all
 * entities in that group.
 *
 * @baleen.javadoc
 */
public class CorefCapitalisationAndApostrophe extends BaleenAnnotator {
  /**
   * If two location entities are thought to be coreferences, but they have different existing
   * reference targets, should we merge them?
   *
   * @baleen.config false
   */
  public static final String PARAM_MERGE_REFERENTS = "mergeReferents";

  @ConfigurationParameter(name = PARAM_MERGE_REFERENTS, defaultValue = "false")
  boolean mergeReferents = false;

  @Override
  public void doProcess(JCas jCas) throws AnalysisEngineProcessException {
    Map<String, List<Entity>> groups = processEntities(jCas);
    processGroups(jCas, groups);
  }

  private void processGroups(JCas jCas, Map<String, List<Entity>> groups) {
    for (List<Entity> group : groups.values()) {
      if (group.size() <= 1) {
        continue;
      }

      Set<ReferenceTarget> rts = new HashSet<>();
      for (Entity e : group) {
        if (e.getReferent() != null) {
          rts.add(e.getReferent());
        }
      }

      ReferenceTarget rt = selectAppropriateReferenceTarget(jCas, rts);
      if (rt == null) {
        getMonitor()
            .debug(
                "Unable to coreference capitalised entities '{}' as they have different existing referents",
                getEntityValue(group.get(0)));
      } else {
        for (Entity e : group) {
          e.setReferent(rt);
        }
      }
    }
  }

  private Map<String, List<Entity>> processEntities(JCas jCas) {
    Map<String, List<Entity>> groups = new HashMap<>();
    Collection<Entity> entities = JCasUtil.select(jCas, Entity.class);

    for (Entity entity : entities) {
      String value = getEntityValue(entity);
      value = normalizeValue(value);

      String key = entity.getType().getName().toUpperCase() + "::" + value;
      List<Entity> groupEntities = groups.containsKey(key) ? groups.get(key) : new ArrayList<>();

      groupEntities.add(entity);
      groups.put(key, groupEntities);
    }
    return groups;
  }

  private String getEntityValue(Entity e) {
    String val = e.getValue();
    if (Strings.isNullOrEmpty(val)) {
      val = e.getCoveredText();
    }

    return val;
  }

  private String normalizeValue(String value) {
    String s = value.trim().toLowerCase();

    if (s.endsWith("'s")) {
      s = s.substring(0, s.length() - 2);
    } else if (s.endsWith("s'")) {
      s = s.substring(0, s.length() - 1);
    }

    return s;
  }

  private ReferenceTarget selectAppropriateReferenceTarget(
      JCas jCas, Collection<ReferenceTarget> referenceTargets) {
    ReferenceTarget rt = null;

    if (referenceTargets.isEmpty()) {
      rt = createReferenceTarget(jCas);
    } else if (referenceTargets.size() == 1) {
      rt = referenceTargets.toArray(new ReferenceTarget[0])[0];
    } else if (mergeReferents) {
      rt = createReferenceTarget(jCas);

      removeFromJCasIndex(referenceTargets);
    }

    return rt;
  }

  private ReferenceTarget createReferenceTarget(JCas jCas) {
    ReferenceTarget rt = new ReferenceTarget(jCas);
    rt.setBegin(0);
    rt.setEnd(jCas.getDocumentText().length());
    addToJCasIndex(rt);

    return rt;
  }

  @Override
  public AnalysisEngineAction getAction() {
    return new AnalysisEngineAction(ImmutableSet.of(Entity.class), Collections.emptySet());
  }
}
