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

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.regex.Pattern;

import org.apache.commons.lang3.StringUtils;
import org.apache.uima.UimaContext;
import org.apache.uima.resource.ResourceInitializationException;

import com.google.common.collect.ImmutableSet;

import uk.gov.dstl.baleen.core.pipelines.orderers.AnalysisEngineAction;
import uk.gov.dstl.baleen.types.common.Quantity;
import uk.gov.dstl.baleen.types.semantic.Entity;

/**
 * Find adjacent quantities of the same type and merge them
 *
 * <p>This annotator will find adjacent quantities of the same type and merge them into a single
 * quantity. For example, 7lb 4oz should be annotated as a single entity, not two.
 */
public class MergeAdjacentQuantities extends MergeAdjacent {
  @Override
  public void doInitialize(UimaContext aContext) throws ResourceInitializationException {
    separatorPattern = Pattern.compile(separator);

    classTypes = new ArrayList<>();
    classTypes.add(Quantity.class);
  }

  @Override
  public boolean shouldMerge(Entity e1, Entity e2) {
    if (!(e1 instanceof Quantity && e2 instanceof Quantity)) {
      return false;
    }

    Quantity q1 = (Quantity) e1;
    Quantity q2 = (Quantity) e2;

    return q1.getNormalizedQuantity() >= q2.getNormalizedQuantity()
        && StringUtils.equals(q1.getSubType(), q2.getSubType())
        && StringUtils.equals(q1.getNormalizedUnit(), q2.getNormalizedUnit());
  }

  @Override
  public boolean mergeAdditionalProperties(
      Entity merged, Class<? extends Entity> type, List<Entity> originalEntities) {
    if (type != Quantity.class || merged.getClass() != type) return false;

    Quantity qMerged = (Quantity) merged;
    Double normalizedQuantity = 0.0;

    Set<String> units = new HashSet<>();

    for (Entity e : originalEntities) {
      if (e.getClass() != type) continue;

      Quantity q = (Quantity) e;

      if (StringUtils.isNotBlank(q.getUnit())) {
        boolean newUnit = units.add(q.getUnit());
        if (!newUnit) return false;
      }

      normalizedQuantity += q.getNormalizedQuantity();

      setNormalizedUnit(qMerged, q.getNormalizedUnit());
      setSubType(qMerged, q.getSubType());
    }

    qMerged.setNormalizedQuantity(normalizedQuantity);

    return true;
  }

  private void setNormalizedUnit(Quantity qMerged, String unit) {
    if (StringUtils.isBlank(qMerged.getNormalizedUnit())) {
      qMerged.setNormalizedUnit(unit);
    }
  }

  private void setSubType(Quantity qMerged, String type) {
    if (StringUtils.isBlank(qMerged.getSubType())) {
      qMerged.setSubType(type);
    }
  }

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