package tudelft.healthpsychology.traumaontologies.owltree;

import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

import org.semanticweb.owlapi.model.AxiomType;
import org.semanticweb.owlapi.model.IRI;
import org.semanticweb.owlapi.model.OWLAnnotationAssertionAxiom;
import org.semanticweb.owlapi.model.OWLClass;
import org.semanticweb.owlapi.model.OWLDataProperty;
import org.semanticweb.owlapi.model.OWLDataPropertyDomainAxiom;
import org.semanticweb.owlapi.model.OWLDataPropertyRangeAxiom;
import org.semanticweb.owlapi.model.OWLOntology;
import org.semanticweb.owlapi.reasoner.InferenceType;
import org.semanticweb.owlapi.reasoner.OWLReasoner;
import org.semanticweb.owlapi.reasoner.OWLReasonerConfiguration;
import org.semanticweb.owlapi.reasoner.OWLReasonerFactory;
import org.semanticweb.owlapi.reasoner.SimpleConfiguration;
import org.semanticweb.owlapi.reasoner.structural.StructuralReasonerFactory;

import uk.ac.manchester.cs.owl.owlapi.OWLDatatypeImpl;
import uk.ac.manchester.cs.owl.owlapi.OWLLiteralImplString;

/**
 * OWLReasoner that assumes that the ontology contains a tree (not a graph)
 * <ul>
 * <li>each node has at most 1 parent
 * <li>no cycles
 * </ul>
 * *
 */
public class OwlTreeReasoner {
	private final OWLOntology ontology;
	private final OWLReasoner reasoner;
	public transient final OWLClass thing; // owl:Thing

	public OwlTreeReasoner(OWLOntology ontology) {
		this.ontology = ontology;
		/*
		 * use reasoner, to do class hierarchy the inference stuff. There might
		 * be some tricky restrictions or other OWL tricks that we don't know of
		 * so better leave this to the official reasoners.
		 */
		OWLReasonerFactory reasonerFactory = new StructuralReasonerFactory();
		OWLReasonerConfiguration config = new SimpleConfiguration();
		reasoner = reasonerFactory.createReasoner(ontology, config);
		reasoner.precomputeInferences(InferenceType.CLASS_HIERARCHY);
		// ignore if it's inconsistent. As long as we can get the superclasses,
		// don't care.
		thing = ontology.getOWLOntologyManager().getOWLDataFactory()
				.getOWLThing();
	}

	/**
	 * @param obj the {@link OWLClass} to get the children of
	 * @return direct children of the given obj
	 * 
	 */
	public List<OWLClass> getChildren(OWLClass obj) {
		return reasoner.getSubClasses(obj, true).entities()
				.filter(o -> !o.isOWLNothing()).collect(Collectors.toList());
	}

	/**
	 * @param maxdepth Must be positive or 0. 0 means node itself, 1 means
	 *                 direct children, 2 means children of the children, etc.
	 * @return list with only this node if this is a leaf node or depth<=0. If
	 *         this is not a leaf node, then returns a list joining all
	 *         getChildren(depth-1) of each child. The difference with getLeaves
	 *         is that this includes non-leave nodes at maxdepth.
	 */
	public List<OWLClass> getChildren(OWLClass obj, int maxdepth) {
		if (maxdepth <= 0) {
			return Arrays.asList(obj);
		}
		List<OWLClass> children = reasoner.getSubClasses(obj, true).entities()
				.filter(o -> !o.isOWLNothing()).collect(Collectors.toList());
		if (children.isEmpty()) {
			return Arrays.asList(obj);
		}
		return children.stream().map(child -> getChildren(child, maxdepth - 1))
				.flatMap(List::stream).collect(Collectors.toList());
	}

	/**
	 * @param obj      the {@link OWLClass} to get the leaf nodes from
	 * @param maxdepth the maximum depth of the leafs to be collected. 0 means
	 *                 obj itself.
	 * @return all leave nodes up do maxdepth deep. If given node is a leaf node
	 *         (no children), returns list with only that node. returns empty
	 *         list if maxdeptn<=0 and this node has children, or when the node
	 *         has no laves at given depth.
	 * 
	 */
	public List<OWLClass> getLeaves(OWLClass obj, int maxdepth) {
		// We have to remove "owl.NOTHING" objects, these are weird empty
		// solution "objects".
		List<OWLClass> children = reasoner.getSubClasses(obj, true).entities()
				.filter(o -> !o.isOWLNothing()).collect(Collectors.toList());
		if (children.isEmpty()) {
			return Arrays.asList(obj);
		}
		// this is not a leaf node.
		if (maxdepth <= 0) {
			return Collections.emptyList();
		}
		return children.stream().map(child -> getLeaves(child, maxdepth - 1))
				.flatMap(List::stream).collect(Collectors.toList());
	}

	/**
	 * @param obj the {@link OWLClass} to get the ancestors of
	 * @return all ancestors of the given obj, up to the root. Does not include
	 *         the class itself.
	 * 
	 */
	public List<OWLClass> getAncestors(OWLClass obj) {
		return new LinkedList<OWLClass>(
				reasoner.getSuperClasses(obj, false).entities()
						.filter(o -> !isThing(o)).collect(Collectors.toList()));
	}

	/**
	 * 
	 * @param name the name for which an IRI is needed
	 * @return URI for the given object name in this ontology: the ontolyg
	 *         basename followed by "#" and name.
	 */
	public IRI getIri(String name) {
		return IRI.create(getBaseName() + "#" + name);
	}

	public String getBaseName() {
		return ontology.getOntologyID().getOntologyIRI().get().toString();
	}

	/**
	 * 
	 * @param name the object name. See also {@link #getIri(String)}
	 * @return {@link OWLClass} with the given name
	 */
	public OWLClass getObject(String name) {
		IRI iri = getIri(name);
		if (!ontology.containsClassInSignature(iri)) {
			throw new IllegalArgumentException(
					"Ontology does not contain object '#" + name + "'");
		}
		OWLClass obj = ontology.getOWLOntologyManager().getOWLDataFactory()
				.getOWLClass(iri);
		return obj;
	}

	/**
	 * @param obj the {@link OWLClass} to get the parent of
	 * @return direct parent of the given obj, or null if no parent
	 * @throws IllegalStateException if object has more than 1 parent.
	 * 
	 */
	public OWLClass getParent(OWLClass obj) {
		Set<OWLClass> superClass = reasoner.getSuperClasses(obj, true)
				.getFlattened();
		if (superClass.size() > 1) {
			throw new IllegalStateException(
					"There is a bug in the ontology: class " + obj
							+ " has multiple parents: " + superClass);
		}
		if (superClass.isEmpty())
			return null;
		OWLClass s = superClass.iterator().next();
		if (thing.equals(s)) // OWL adds "thing" at root. Remove it..
			return null;
		return s;
	}

	/**
	 * @param object
	 * @return all {@link OWLDataPropertyDomainAxiom}s of object and all
	 *         ancestor classes. These axioms are basically URIs about which
	 *         other info like domain, range and comments are stored.
	 */
	public List<OWLDataPropertyDomainAxiom> getProperties(OWLClass object) {
		List<OWLClass> thisAndAncestors = getAncestors(object);
		thisAndAncestors.add(object);
		List<OWLDataPropertyDomainAxiom> properties = new LinkedList<>();
		for (OWLClass clas : thisAndAncestors) {
			ontology.axioms(AxiomType.DATA_PROPERTY_DOMAIN)
					.filter(dp -> clas.equals(dp.getDomain()))
					.forEach(dp -> properties.add(dp));
		}
		return properties;
	}

	/**
	 * 
	 * @param axiom an {@link OWLDataProperty}, typically an
	 *              {@link OWLDataPropertyDomainAxiom}
	 * @return list of comment annotations for this axiom
	 */
	public String getComment(OWLDataPropertyDomainAxiom axiom) {
		List<String> list = axiom.dataPropertiesInSignature()
				.map(odp -> getCommentAnnotation(odp))
				.collect(Collectors.toList());
		if (list.size() != 1) {
			throw new IllegalStateException(
					"There must be exactly 1 comment for  " + axiom
							+ "  but found " + list);
		}
		return list.get(0);
	}

	/**
	 * @param odp the {@link OWLDataPropertyDomainAxiom}
	 * @return the data ranges set for the given odp
	 */
	public String getRange(OWLDataPropertyDomainAxiom odp) {
		IRI property = (IRI) odp.getProperty().components().findFirst().get();

		List<OWLDataPropertyRangeAxiom> axioms = ontology
				.axioms(AxiomType.DATA_PROPERTY_RANGE)
				.filter(p -> property
						.equals(p.getProperty().components().findFirst().get()))
				.collect(Collectors.toList());
		if (axioms.size() != 1) {
			throw new IllegalStateException(
					"There are must be exactly 1 range axiom for  " + odp
							+ " but found " + axioms);
		}

		// there is no doc on how to do this without nasty casts...
		return ((OWLDatatypeImpl) axioms.get(0).getRange()).getIRI().toString();
	}

	/**
	 * 
	 * @param c an {@link OWLClass} object
	 * @return true iff c equals to Owl:Thing.
	 */
	public boolean isThing(OWLClass c) {
		return thing.equals(c);
	}

	/**
	 * @param odp
	 * @return the annotation for this OWLDataProperty.
	 */
	private String getCommentAnnotation(OWLDataProperty odp) {
		List<OWLAnnotationAssertionAxiom> annotationaxioms = ontology
				.annotationAssertionAxioms(odp.getIRI())
				.collect(Collectors.toList());
		if (annotationaxioms.size() != 1) {
			// we don't allow because we can't figure out the range in that
			// case.
			throw new IllegalStateException(
					"There must be exactly 1 comment for  " + odp
							+ " but found " + annotationaxioms);
		}
		// there is no doc on how to do this without nasty casts...
		return ((OWLLiteralImplString) annotationaxioms.get(0).getAnnotation()
				.annotationValue()).getLiteral();

	}

}