package uk.ac.ox.cs.pagoda.approx;

import org.apache.commons.io.FilenameUtils;
import org.semanticweb.HermiT.Configuration;
import org.semanticweb.HermiT.model.DLClause;
import org.semanticweb.HermiT.model.DLOntology;
import org.semanticweb.HermiT.structural.OWLClausification;
import org.semanticweb.owlapi.model.*;
import org.semanticweb.owlapi.profiles.OWL2RLProfile;
import org.semanticweb.owlapi.profiles.OWLProfileReport;
import org.semanticweb.owlapi.profiles.OWLProfileViolation;
import uk.ac.ox.cs.pagoda.constraints.NullaryBottom;
import uk.ac.ox.cs.pagoda.constraints.UnaryBottom;
import uk.ac.ox.cs.pagoda.owl.OWLHelper;
import uk.ac.ox.cs.pagoda.util.Namespace;
import uk.ac.ox.cs.pagoda.util.Utility;

import java.io.*;
import java.nio.file.Paths;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class RLPlusOntology implements KnowledgeBase {

    private static final String DEFAULT_ONTOLOGY_FILE_EXTENSION = "owl";
    private static final Pattern ONTOLOGY_ID_REGEX = Pattern.compile("ontologyid\\((?<id>[a-z0-9\\-]+)\\)");
    OWLOntologyManager manager;
    OWLDataFactory factory;
    String ontologyIRI;
    String corrFileName = null;
    String outputPath, aBoxPath;
    OWLOntology inputOntology = null;
    OWLOntology tBox = null;
    OWLOntology aBox = null;
    OWLOntology restOntology = null;
    OWLOntology outputOntology = null; //RL ontology
    DLOntology dlOntology = null;
    int rlCounter = 0;
    LinkedList<Clause> clauses;
    Map<OWLAxiom, OWLAxiom> correspondence;
    BottomStrategy botStrategy;
    Random random = new Random(19900114);
    private Map<OWLClassExpression, Integer> subCounter = null;
    private Map<OWLClass, OWLClass> atomic2negation = new HashMap<OWLClass, OWLClass>();

    // TODO don't know if it is correct
    @Override
    public void load(OWLOntology ontology, uk.ac.ox.cs.pagoda.constraints.BottomStrategy bottomStrategy) {
        if(bottomStrategy instanceof UnaryBottom)
            botStrategy = BottomStrategy.UNARY;
        else if(bottomStrategy instanceof NullaryBottom)
            botStrategy = BottomStrategy.NULLARY;
        else
            botStrategy = BottomStrategy.TOREMOVE;

        if(corrFileName == null)
            corrFileName = "rlplus.crr";
        manager = ontology.getOWLOntologyManager();
//		manager = OWLManager.createOWLOntologyManager();
        factory = manager.getOWLDataFactory();
        inputOntology = ontology;

        try {
            IRI ontologyIri;
            if(ontology.isAnonymous()) {
                Matcher matcher = ONTOLOGY_ID_REGEX.matcher(ontology.getOntologyID().toString().toLowerCase());
                if(!matcher.matches()) throw new Error("Anonymous ontology without internal id");
                ontologyIri =
                        IRI.create("http://www.example.org/", matcher.group("id") + "." + DEFAULT_ONTOLOGY_FILE_EXTENSION);
            }
            else
                ontologyIri = inputOntology.getOntologyID().getOntologyIRI();

            String ontologyIriPrefix = ontologyIri.getNamespace();
            ontologyIRI = ontologyIri.toString();
            String ontologyIriFragment = ontologyIri.getFragment();
            String originalFileName = FilenameUtils.removeExtension(ontologyIriFragment);
            String originalExtension = FilenameUtils.getExtension(ontologyIriFragment);
            if(originalExtension.isEmpty()) originalExtension = DEFAULT_ONTOLOGY_FILE_EXTENSION;


            IRI rlOntologyIRI = IRI.create(ontologyIriPrefix, originalFileName + "-RL." + originalExtension);
            outputPath = Paths.get(Utility.getGlobalTempDirAbsolutePath(),
                                   originalFileName + "-RL." + originalExtension).toString();
            IRI rlDocumentIRI = IRI.create("file://" + outputPath);
            outputOntology = manager.createOntology(rlOntologyIRI);
            manager.setOntologyDocumentIRI(outputOntology, rlDocumentIRI);

            String tBoxOntologyFragment = originalFileName + "-TBox." + originalExtension;
            IRI tBoxOntologyIRI = IRI.create(ontologyIriPrefix, tBoxOntologyFragment);
            IRI tBoxDocumentIRI =
                    IRI.create("file://" + Paths.get(Utility.getGlobalTempDirAbsolutePath(), tBoxOntologyFragment));

            String aBoxOntologyFragment = originalFileName + "-ABox." + originalExtension;
            IRI aBoxOntologyIRI = IRI.create(ontologyIriPrefix, aBoxOntologyFragment);
            aBoxPath = Paths.get(Utility.getGlobalTempDirAbsolutePath(), aBoxOntologyFragment).toString();
            IRI aBoxDocumentIRI =
                    IRI.create("file://" + Paths.get(Utility.getGlobalTempDirAbsolutePath(), aBoxOntologyFragment));

            tBox = manager.createOntology(tBoxOntologyIRI);
            aBox = manager.createOntology(aBoxOntologyIRI);
            manager.setOntologyDocumentIRI(tBox, tBoxDocumentIRI);
            manager.setOntologyDocumentIRI(aBox, aBoxDocumentIRI);

            FileOutputStream aBoxOut = new FileOutputStream(aBoxPath);
            manager.saveOntology(aBox, aBoxOut);
            aBoxOut.close();

            restOntology = manager.createOntology();
        } catch(OWLOntologyCreationException | OWLOntologyStorageException | IOException e) {
            e.printStackTrace();
            System.exit(1);
        }
    }

    public OWLOntology getTBox() {
        return tBox;
    }

    public String getABoxPath() {
        return aBoxPath;
    }

    public void simplify() {
        if(simplifyABox()) {
            save(aBox);
//			save(tBox);
        }
        else
            tBox = inputOntology;
    }

    @Override
    public void transform() {
        simplify();
        filter();
        clausify();

        subCounter = new HashMap<OWLClassExpression, Integer>();
        clauses = new LinkedList<Clause>();
        Clausifier clausifier = Clausifier.getInstance(restOntology);

        for(DLClause c : dlOntology.getDLClauses()) {
            Clause clause = new Clause(clausifier, c);
            clauses.add(clause);

			/*
             * count the expressions in the left
			 */
            for(OWLClassExpression exp : clause.getSubClasses()) {
                if(exp instanceof OWLClass)
                    add2SubCounter(exp);
                else if(exp instanceof OWLObjectSomeValuesFrom) {
                    OWLObjectSomeValuesFrom someValue = (OWLObjectSomeValuesFrom) exp;
                    add2SubCounter(factory.getOWLObjectSomeValuesFrom(someValue.getProperty(), factory.getOWLThing()));
                    add2SubCounter(someValue.getFiller());
                }
                else if(exp instanceof OWLObjectMinCardinality) {
                    OWLObjectMinCardinality minCard = (OWLObjectMinCardinality) exp;
                    add2SubCounter(factory.getOWLObjectSomeValuesFrom(minCard.getProperty(), factory.getOWLThing()));
                    add2SubCounter(minCard.getFiller());
                }
                else
                    Utility.logError("strange class expression: " + exp);

            }
        }

        correspondence = new HashMap<OWLAxiom, OWLAxiom>();
        Set<OWLAxiom> addedAxioms = new HashSet<OWLAxiom>();
        OWLClassExpression subExp;
        for(Clause clause : clauses) {
            subExp = uk.ac.ox.cs.pagoda.owl.OWLHelper.getSimplifiedConjunction(factory, clause.getSubClasses());
            addedAxioms.clear();
            for(OWLClassExpression exp : getDisjunctionApprox0(clause.getSuperClasses())) {
                addedAxioms.add(factory.getOWLSubClassOfAxiom(subExp, transform(exp, addedAxioms)));
                for(OWLAxiom a : addedAxioms)
                    addAxiom2output(a, factory.getOWLSubClassOfAxiom(subExp,
                                                                     OWLHelper.getSimplifiedDisjunction(factory, clause.getSuperClasses())));
            }
        }

        subCounter.clear();
    }

    @Override
    public void save() {
        if(corrFileName != null)
            save(correspondence, corrFileName);
        save(outputOntology);
    }

    public OWLOntologyManager getOWLOntologyManager() {
        return inputOntology.getOWLOntologyManager();
    }

    public String getOntologyIRI() {
        return ontologyIRI;
    }

    public OWLOntology getOutputOntology() {
        return outputOntology;
    }

    @Override
    public String getOutputPath() {
        return outputPath;
    }

    @Override
    public String getDirectory() {
        return outputPath.substring(0, outputPath.lastIndexOf(Utility.FILE_SEPARATOR));
    }

    public void setCorrespondenceFileLoc(String path) {
        corrFileName = path;
    }

    private void add2SubCounter(OWLClassExpression exp) {
        Integer count = subCounter.get(exp);
        if(count == null) count = 0;
        ++count;
        subCounter.put(exp, count);
    }

    private void save(Map<OWLAxiom, OWLAxiom> map, String corrFileName) {
        if(corrFileName == null) return;
        ObjectOutput output;
        try {
            output = new ObjectOutputStream(new BufferedOutputStream(new FileOutputStream(corrFileName)));
            output.writeObject(map);
            output.close();
        } catch(IOException e) {
            e.printStackTrace();
        }
    }

    protected void save(OWLOntology onto) {
        try {
            onto.getOWLOntologyManager().saveOntology(onto);
        } catch(OWLOntologyStorageException e) {
            e.printStackTrace();
        }
    }

    /*
     * treat disjunction as conjunction
     */
    private Set<OWLClassExpression> getDisjunctionApprox0(Set<OWLClassExpression> superClasses) {
        return superClasses;
    }

    /*
     * choose one simple class disjunct
     */
    @SuppressWarnings("unused")
    private Set<OWLClassExpression> getDisjunctionApprox1(Set<OWLClassExpression> superClasses) {
        if(superClasses.isEmpty() || superClasses.size() == 1)
            return superClasses;

        OWLClassExpression rep = null;
        int min = Integer.MAX_VALUE, o;
        for(OWLClassExpression exp : superClasses)
            if(exp instanceof OWLClass && (o = getOccurrence(exp)) < min) {
                min = o;
                rep = exp;
            }

        if(rep == null) rep = superClasses.iterator().next();

        return Collections.singleton(rep);
    }

    /*
     * randomly choose a class expression to represent this disjunction
     */
    @SuppressWarnings("unused")
    private Set<OWLClassExpression> getDisjunctionApprox2(Set<OWLClassExpression> superClasses) {
        if(superClasses.isEmpty() || superClasses.size() == 1)
            return superClasses;

        int index = random.nextInt() % superClasses.size();
        if(index < 0) index += superClasses.size();

        int i = 0;
        for(OWLClassExpression exp : superClasses)
            if(i++ == index)
                return Collections.singleton(exp);
        return null;
    }

    /*
     * choose the one that appears least in the l.h.s.
     */
    @SuppressWarnings("unused")
    private Set<OWLClassExpression> getDisjunctionApprox3(Set<OWLClassExpression> superClasses) {
        if(superClasses.isEmpty() || superClasses.size() == 1)
            return superClasses;

        OWLClassExpression rep = null, exp1;
        int occurrence = Integer.MAX_VALUE, o;
        for(OWLClassExpression exp : superClasses) {
            o = 0;
            exp1 = exp;
            if(exp instanceof OWLObjectMinCardinality) {
                OWLObjectMinCardinality minCard = (OWLObjectMinCardinality) exp;
                if(minCard.getCardinality() == 1)
                    exp1 = factory.getOWLObjectSomeValuesFrom(minCard.getProperty(), minCard.getFiller());
            }

            if(!subCounter.containsKey(exp1) || (o = subCounter.get(exp1)) < occurrence) {
                rep = exp;
                occurrence = o;
            }
        }

        return Collections.singleton(rep);
    }

    private int getOccurrence(OWLClassExpression exp) {
        if(!subCounter.containsKey(exp))
            return 0;
        return subCounter.get(exp);
    }

    @SuppressWarnings("unused")
    private Set<OWLClassExpression> getDisjunctionApprox4(Set<OWLClassExpression> superClasses) {
        if(superClasses.isEmpty() || superClasses.size() == 1)
            return superClasses;

        OWLClassExpression rep = null;
        int occurrence = Integer.MAX_VALUE, o;
        for(OWLClassExpression exp : superClasses) {
            o = 0;
            if(exp instanceof OWLObjectMinCardinality) {
                OWLObjectMinCardinality minCard = (OWLObjectMinCardinality) exp;
                if(minCard.getCardinality() == 1) {
                    o =
                            getOccurrence((factory.getOWLObjectSomeValuesFrom(minCard.getProperty(), factory.getOWLThing())));
                    o += getOccurrence(minCard.getFiller());
//					if (o < o1) o = o1;
                }
            }
            else
                o = getOccurrence(exp);

            if(o < occurrence || o == occurrence && !(rep instanceof OWLClass)) {
                rep = exp;
                occurrence = o;
            }
        }

        return Collections.singleton(rep);
    }

    private boolean simplifyABox() {
        boolean flag = false;
        Map<OWLClassExpression, OWLClass> complex2atomic = new HashMap<OWLClassExpression, OWLClass>();

        OWLDatatype anyURI = factory.getOWLDatatype(IRI.create(Namespace.XSD_NS + "anyURI"));
        OWLObjectProperty sameAs = factory.getOWLObjectProperty(IRI.create(Namespace.EQUALITY));
        OWLObjectProperty differentFrom = factory.getOWLObjectProperty(IRI.create(Namespace.INEQUALITY));

        for(OWLOntology imported : inputOntology.getImportsClosure())
            for(OWLAxiom axiom : imported.getAxioms()) {
                if(axiom instanceof OWLClassAssertionAxiom) {
                    flag = true;
                    OWLClassAssertionAxiom assertion = (OWLClassAssertionAxiom) axiom;
                    OWLClassExpression clsExp = assertion.getClassExpression();
                    OWLClass cls;
                    if(clsExp instanceof OWLClass) {
                        if(((OWLClass) clsExp).toStringID().startsWith("owl:"))
                            manager.addAxiom(tBox, axiom);
                        else manager.addAxiom(aBox, axiom);
                    }
                    else {
                        if((cls = complex2atomic.get(clsExp)) == null) {
                            complex2atomic.put(clsExp, cls = getNewConcept(tBox, rlCounter++));
                            manager.addAxiom(tBox, factory.getOWLSubClassOfAxiom(cls, clsExp));
                        }
                        manager.addAxiom(aBox, factory.getOWLClassAssertionAxiom(cls, assertion.getIndividual()));
                    }
                }
                else if(axiom instanceof OWLObjectPropertyAssertionAxiom || axiom instanceof OWLDataPropertyAssertionAxiom || axiom instanceof OWLAnnotationAssertionAxiom) {
                    if(axiom.getDataPropertiesInSignature().contains(anyURI)) continue;
                    flag = true;
                    manager.addAxiom(aBox, axiom);
                }
                else if(axiom instanceof OWLSameIndividualAxiom) {
                    OWLIndividual firstIndividual = null, previousIndividual = null, lastIndividual = null;
                    for(OWLIndividual next : ((OWLSameIndividualAxiom) axiom).getIndividuals()) {
                        if(firstIndividual == null) firstIndividual = previousIndividual = next;
                        else
                            manager.addAxiom(aBox, factory.getOWLObjectPropertyAssertionAxiom(sameAs, previousIndividual, next));
                        previousIndividual = lastIndividual = next;
                    }
                    manager.addAxiom(aBox, factory.getOWLObjectPropertyAssertionAxiom(sameAs, lastIndividual, firstIndividual));
                }
                else if(axiom instanceof OWLDifferentIndividualsAxiom) {
                    int index1 = 0, index2;
                    for(OWLIndividual individual1 : ((OWLDifferentIndividualsAxiom) axiom).getIndividuals()) {
                        ++index1;
                        index2 = 0;
                        for(OWLIndividual individual2 : ((OWLDifferentIndividualsAxiom) axiom).getIndividuals()) {
                            if(index2++ < index1) {
                                manager.addAxiom(aBox, factory.getOWLObjectPropertyAssertionAxiom(differentFrom, individual1, individual2));
                            }
                            else break;
                        }
                    }
                }
                else
                    manager.addAxiom(tBox, axiom);
            }

        return flag;
    }

    private void filter() {
        OWL2RLProfile profile = new OWL2RLProfile();
        OWLProfileReport report = profile.checkOntology(tBox);
        Set<OWLAxiom> rlAxioms = tBox.getAxioms();
        OWLAxiom axiom;

        for(OWLProfileViolation violation : report.getViolations()) {
            manager.addAxiom(restOntology, axiom = violation.getAxiom());
            rlAxioms.remove(axiom);
        }

        for(Iterator<OWLAxiom> iter = rlAxioms.iterator(); iter.hasNext(); )
            addAxiom2output(iter.next(), null);
    }

    private void clausify() {
        Configuration conf = new Configuration();
        OWLClausification clausifier = new OWLClausification(conf);
        dlOntology = (DLOntology) clausifier.preprocessAndClausify(restOntology, null)[1];
        clausifier = null;
    }

    protected void addAxiom2output(OWLAxiom axiom, OWLAxiom correspondingAxiom) {
        manager.addAxiom(outputOntology, axiom);
        if(correspondingAxiom != null)
            correspondence.put(axiom, correspondingAxiom);
    }

    private OWLClassExpression transform(OWLClassExpression exp, Set<OWLAxiom> addedAxioms) {
        if(exp instanceof OWLClass)
            return exp;

        if(exp instanceof OWLObjectHasValue)
            return exp;

        if(exp instanceof OWLObjectSomeValuesFrom) {
            OWLObjectSomeValuesFrom someValueExp = (OWLObjectSomeValuesFrom) exp;

            OWLClassExpression tExp = someValueExp.getFiller();
            if(tExp.equals(factory.getOWLThing()))
                exp = factory.getOWLObjectMinCardinality(1, someValueExp.getProperty());
            else
                exp = factory.getOWLObjectMinCardinality(1, someValueExp.getProperty(), someValueExp.getFiller());
        }

        if(exp instanceof OWLObjectMinCardinality) {
            OWLObjectMinCardinality minExp = (OWLObjectMinCardinality) exp;
            OWLObjectPropertyExpression r;

            if(minExp.getFiller().equals(factory.getOWLThing())) {
                r = minExp.getProperty();
            }
            //TODO to be restored ...
            //else if ((r = exists2role.get(someValueExp)) == null) {
            // deal with r' \subseteq r & range(r') \subseteq C
            else {
                r = getNewRole(outputOntology, rlCounter);
                addedAxioms.add(factory.getOWLSubObjectPropertyOfAxiom(r, minExp.getProperty()));
                OWLClassExpression tExp = minExp.getFiller();
                if(!(tExp instanceof OWLObjectComplementOf)) {
                    if(tExp.equals(factory.getOWLThing())) ;
                    else
                        addedAxioms.add(factory.getOWLObjectPropertyRangeAxiom(r, tExp));
                }
                else if(botStrategy != BottomStrategy.TOREMOVE) {
                    OWLClass cls = (OWLClass) tExp.getComplementNNF();
                    OWLClass neg;
                    if((neg = atomic2negation.get(cls)) == null) {
                        neg = getNewConcept(outputOntology, rlCounter);
                        addedAxioms.add(factory.getOWLDisjointClassesAxiom(neg, cls));
                        atomic2negation.put(cls, neg);
                    }
                    addedAxioms.add(factory.getOWLObjectPropertyRangeAxiom(r, neg));
                }
//				exists2role.put(someValueExp, (OWLObjectProperty) r);
            }

            // deal with r'(x,c)
            Set<OWLClassExpression> ret = new HashSet<OWLClassExpression>();
            int num = minExp.getCardinality();

            Set<OWLNamedIndividual> cs = new HashSet<OWLNamedIndividual>();
            OWLNamedIndividual c;
            for(int i = 0; i < num; ++i) {
                c = getNewIndividual(outputOntology, rlCounter++);
                ret.add(factory.getOWLObjectHasValue(r, c));
                cs.add(c);
            }

            if(botStrategy != BottomStrategy.TOREMOVE && cs.size() > 1) {
                addedAxioms.add(factory.getOWLDifferentIndividualsAxiom(cs));
            }

            return OWLHelper.getSimplifiedConjunction(factory, ret);
        }

        if(exp instanceof OWLObjectMaxCardinality) {
            OWLObjectMaxCardinality maxExp = (OWLObjectMaxCardinality) exp;
            OWLClassExpression tExp = maxExp.getFiller();
            int card = maxExp.getCardinality() >= 1 ? 1 : 0;
            if(!(tExp instanceof OWLObjectComplementOf))
                return factory.getOWLObjectMaxCardinality(card, maxExp.getProperty(), tExp);
            else {
                Utility.logDebug("oh, to be tested ... ");
                OWLClassExpression tExp1 =
                        factory.getOWLObjectAllValuesFrom(maxExp.getProperty(), tExp.getComplementNNF());
                if(card == 0)
                    return tExp1;
                else {
                    OWLClassExpression tExp2 = factory.getOWLObjectMaxCardinality(1, maxExp.getProperty());
                    return factory.getOWLObjectIntersectionOf(tExp1, tExp2);
                }
            }
        }

        if(exp instanceof OWLObjectAllValuesFrom)
            return exp;

        if(exp instanceof OWLObjectOneOf)
            if(((OWLObjectOneOf) exp).getIndividuals().size() == 1)
                return exp;
            else
                return null;

        if(exp instanceof OWLDataHasValue)
            return exp;

        //TODO overapproximation - dealing with OWLDataMinCardinality

        if(exp instanceof OWLDataSomeValuesFrom) {
            return exp;
        }

        if(exp instanceof OWLDataMinCardinality) {
            return exp;
        }

        if(exp instanceof OWLDataMaxCardinality) {
            return exp;
        }


        Set<OWLClassExpression> exps = exp.asConjunctSet();
        if(exps.size() == 1 && exps.iterator().next() == exp) {
            Utility.logError(exp, "error in transform of Ontology~~~~");
        }
        Set<OWLClassExpression> nexps = new HashSet<OWLClassExpression>();
        OWLClassExpression ne;
        boolean changes = false;
        for(OWLClassExpression e : exps) {
            ne = transform(e, addedAxioms);
            if(ne != e) changes = true;
            nexps.add(ne);
        }
        if(changes)
            return OWLHelper.getSimplifiedConjunction(factory, nexps);
        else
            return exp;
    }

    protected OWLNamedIndividual getNewIndividual(OWLOntology onto, int number) {
        OWLOntologyManager manager = onto.getOWLOntologyManager();
        OWLDataFactory factory = manager.getOWLDataFactory();
        OWLNamedIndividual newIndividual =
                factory.getOWLNamedIndividual(IRI.create(Namespace.PAGODA_ANONY + "NI" + number));
        manager.addAxiom(onto, factory.getOWLDeclarationAxiom(newIndividual));
        return newIndividual;
    }

    protected OWLObjectProperty getNewRole(OWLOntology onto, int number) {
        OWLOntologyManager manager = onto.getOWLOntologyManager();
        OWLDataFactory factory = manager.getOWLDataFactory();
        OWLObjectProperty newProperty = factory.getOWLObjectProperty(IRI.create(Namespace.PAGODA_AUX + "NR" + number));
        manager.addAxiom(onto, factory.getOWLDeclarationAxiom(newProperty));
        return newProperty;
    }

    private OWLClass getNewConcept(OWLOntology onto, int number) {
        OWLOntologyManager manager = onto.getOWLOntologyManager();
        OWLDataFactory factory = manager.getOWLDataFactory();
        OWLClass newClass = factory.getOWLClass(IRI.create(Namespace.PAGODA_AUX + "NC" + number));
        manager.addAxiom(onto, factory.getOWLDeclarationAxiom(newClass));
        return newClass;
    }

    private enum BottomStrategy {TOREMOVE, NULLARY, UNARY}
}