package hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage

import hu.bme.mit.inf.dslreasoner.logic.model.builder.TracedOutput
import hu.bme.mit.inf.dslreasoner.logic.model.builder.TypeScopes
import hu.bme.mit.inf.dslreasoner.logic.model.logiclanguage.LogiclanguageFactory
import hu.bme.mit.inf.dslreasoner.logic.model.logiclanguage.RelationDeclaration
import hu.bme.mit.inf.dslreasoner.logic.model.logiclanguage.Type
import hu.bme.mit.inf.dslreasoner.logic.model.logiclanguage.TypeDeclaration
import hu.bme.mit.inf.dslreasoner.logic.model.logiclanguage.TypeDefinition
import hu.bme.mit.inf.dslreasoner.logic.model.logicproblem.LogicProblem
import hu.bme.mit.inf.dslreasoner.logic.model.patterns.SupertypeStar
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partial2logicannotations.PartialModelRelation2Assertion
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.BooleanElement
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.IntegerElement
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialBooleanInterpretation
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialComplexTypeInterpretation
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialIntegerInterpretation
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialInterpretation
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialRealInterpretation
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialRelationInterpretation
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialStringInterpretation
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialinterpretationFactory
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.RealElement
import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.StringElement
import java.math.BigDecimal
import java.util.HashMap
import java.util.Map
import java.util.SortedSet
import org.eclipse.emf.ecore.util.EcoreUtil
import org.eclipse.viatra.query.runtime.api.ViatraQueryEngine
import org.eclipse.viatra.query.runtime.emf.EMFScope
import org.eclipse.xtend.lib.annotations.Data

import static extension hu.bme.mit.inf.dslreasoner.util.CollectionsUtil.*

@Data class Problem2PartialInterpretationTrace {
	Map<TypeDeclaration, PartialComplexTypeInterpretation> type2Interpretation
	PrimitiveValueTrace primitiveValues
	Map<RelationDeclaration, PartialRelationInterpretation> relation2Interpretation
}
@Data class PrimitiveValueTrace {
	PartialBooleanInterpretation booleanInterpretation
	Map<Boolean,BooleanElement> booleanMap
	PartialIntegerInterpretation integerInterpretation
	Map<Integer,IntegerElement> integerMap
	PartialRealInterpretation realInterpretation
	Map<BigDecimal,RealElement> realMap
	PartialStringInterpretation stringInterpretation
	Map<String,StringElement> stringMap
}

/**
 * Initializer class for an empty partial interpretation.
 */
class PartialInterpretationInitialiser {
	val extension PartialinterpretationFactory factory = PartialinterpretationFactory.eINSTANCE
	val extension LogiclanguageFactory factory2 = LogiclanguageFactory.eINSTANCE
	
	/**
	 * Initialises an empty partial interpretation from a logic problem and a scope.
	 */
	def TracedOutput<PartialInterpretation,Problem2PartialInterpretationTrace> initialisePartialInterpretation(
		LogicProblem problem,
		TypeScopes typeScopes) 
	{
		val res = createPartialInterpretation
		res.problem = problem
			
		// Initialise primitive elements
		val booleanType = initBooleans(res)
		val integerType = initIntegers(res, typeScopes.knownIntegers, typeScopes.minNewIntegers, typeScopes.maxNewIntegers)
		val realType = initReals(res, typeScopes.knownReals, typeScopes.minNewReals, typeScopes.maxNewReals)
		val stringType= initStrings(res,typeScopes.knownStrings, typeScopes.minNewStrings, typeScopes.maxNewStrings)
		val primitiveTrace = createPrimitiveTrace(booleanType, integerType, realType, stringType)
		// Initialise complex elements
		
		val type2Interpretation = initElements(res,
			typeScopes.minNewElementsByType,
			typeScopes.maxNewElementsByType,
			typeScopes.minNewElements,
			typeScopes.maxNewElements)
			
		// Initialise relations
		val relation2Interpretation = initRelations(res)
		
		val trace = new Problem2PartialInterpretationTrace(type2Interpretation,primitiveTrace,relation2Interpretation)
		return new TracedOutput(res,trace)
	}
	
	def createPrimitiveTrace(
		PartialBooleanInterpretation booleanType,
		PartialIntegerInterpretation integerType,
		PartialRealInterpretation realType,
		PartialStringInterpretation stringType)
	{
		return new PrimitiveValueTrace(
			booleanType,
			booleanType.elements.filter(BooleanElement).toMap[value],
			integerType,
			integerType.elements.filter(IntegerElement).toMap[value],
			realType,
			realType.elements.filter(RealElement).toMap[value],
			stringType,
			stringType.elements.filter(StringElement).toMap[value]
		)
	}
	
	def protected initBooleans(PartialInterpretation partialInterpretation) {
		val booleanInterpretation = createPartialBooleanInterpretation
		partialInterpretation.partialtypeinterpratation += booleanInterpretation
		
		val trueElement = createBooleanElement => [it.name = "true" it.value = true it.valueSet = true]
		booleanInterpretation.elements += trueElement
		partialInterpretation.newElements+=trueElement
		
		val falseElement = createBooleanElement => [it.name = "false" it.value = false it.valueSet = true]
		booleanInterpretation.elements += falseElement
		partialInterpretation.newElements+=falseElement
		
		return booleanInterpretation
	}
	
	def protected initIntegers(PartialInterpretation partialInterpretation, SortedSet<Integer> knownIntegers, int minNewIntegers, int maxNewIntegers) {
		val integerInterpretation = createPartialIntegerInterpretation
		partialInterpretation.partialtypeinterpratation += integerInterpretation
		
		for(knownInteger : knownIntegers) {
			val integerElement = createIntegerElement => [it.name = knownInteger.toString it.value = knownInteger it.valueSet = true]
			integerInterpretation.elements += integerElement
			partialInterpretation.newElements += integerElement
		}
		if(maxNewIntegers>0) {
			throw new UnsupportedOperationException('''Unspecified Integers are currently not supported!''')
		}
		
		return integerInterpretation
	}
	
	def protected initReals(PartialInterpretation partialInterpretation, SortedSet<BigDecimal> knownReals, int minNewReals, int maxNewReals) {
		val realInterpretation = createPartialRealInterpretation
		partialInterpretation.partialtypeinterpratation += realInterpretation
		
		for(knownReal : knownReals) {
			val realElement = createRealElement => [it.name = knownReal.toString it.value = knownReal it.valueSet = true]
			realInterpretation.elements += realElement
			partialInterpretation.newElements += realElement
		}
		if(maxNewReals>0) {
			throw new UnsupportedOperationException('''Unspecified Real values are currently not supported!''')
		}
		return realInterpretation
	}
	
	def protected initStrings(PartialInterpretation partialInterpretation, SortedSet<String> knownStrings, int minNewStrings, int maxNewStrings) {
		val stringInterpretation = createPartialStringInterpretation
		partialInterpretation.partialtypeinterpratation += stringInterpretation
		
		for(knownString : knownStrings) {
			val stringElement = createStringElement => [it.name = '''"«knownString»"''' it.value = knownString it.valueSet = true]
			stringInterpretation.elements += stringElement
			partialInterpretation.newElements += stringElement
		}
		if(maxNewStrings>0) {
			throw new UnsupportedOperationException('''Unspecified String values are currently not supported!''')
		}
		
		return stringInterpretation
	}
	
	def protected initElements(PartialInterpretation interpretation,
		Map<Type, Integer> minNewElementsByType, Map<Type, Integer> maxNewElementsByType,
		int minNewElements, int maxNewElements)
	{
		val Map<TypeDeclaration, PartialComplexTypeInterpretation> type2Interpretation = new HashMap
		val engine = ViatraQueryEngine.on(new EMFScope(interpretation.problem))
		// Elements
		interpretation.minNewElements = minNewElements
		interpretation.maxNewElements = maxNewElements
		// elements from problem are included
		if(maxNewElements>0) {
			val newElements = createDefinedElement => [it.name = "New Objects"]
			interpretation.openWorldElements += newElements 
		}
		
		for(typeDeclaration : interpretation.problem.types.filter(TypeDeclaration)) {
			val typeInterpretation = typeDeclaration.initialisePartialTypeInterpretation(engine)
			interpretation.partialtypeinterpratation += typeInterpretation
			type2Interpretation.put(typeDeclaration,typeInterpretation)
		}
		interpretation.problem.connectSuperypes(type2Interpretation)
		return type2Interpretation
	}
	
	def private connectSuperypes(LogicProblem problem, Map<TypeDeclaration, PartialComplexTypeInterpretation> trace) {
		for(typeDeclaration : problem.types.filter(TypeDeclaration)) {
			val supertypes = typeDeclaration.<Type>transitiveClosurePlus[it.supertypes]
			for(supertype : supertypes.filter(TypeDeclaration)) {
				typeDeclaration.lookup(trace).supertypeInterpretation += supertype.lookup(trace)
				//println('''«typeDeclaration.name» --> «supertype.name»''')
			}
		}
	}
	
	def private initRelations(PartialInterpretation interpretation) {
		val Map<RelationDeclaration, PartialRelationInterpretation> relation2Interpretation = new HashMap
		for(relation : interpretation.problem.relations.filter(RelationDeclaration)) {
			val partialInterpretation = relation.initialisePartialRelationInterpretation
			interpretation.partialrelationinterpretation += partialInterpretation
			relation2Interpretation.put(relation,partialInterpretation)
		}
		for(pMR2A : interpretation.problem.annotations.filter(PartialModelRelation2Assertion)) {
			val relation = pMR2A.targetRelation
			val r = relation.lookup(relation2Interpretation)
			r.relationlinks+=pMR2A.links.map[EcoreUtil.copy(it)]
		}
//		interpretation.partialfunctioninterpretation += interpretation.problem.functions.filter(FunctionDeclaration)
//			.map[initialisePartialFunctionInterpretation(trace)]
//		interpretation.partialconstantinterpretation += interpretation.problem.constants.filter(ConstantDeclaration)
//			.map[initialisePartialConstantDeclaration(trace)]

		return relation2Interpretation
	}
	
	def private initialisePartialTypeInterpretation(TypeDeclaration t, ViatraQueryEngine engine) {
		val supertypeMatcher = SupertypeStar.Matcher.on(engine)
		val res = createPartialComplexTypeInterpretation => [
			it.interpretationOf = t
			it.elements += supertypeMatcher.getAllValuesOfsubtype(t)
				.filter(TypeDefinition)
				.map[elements].flatten
		]
		return res
	}
	
	def private initialisePartialRelationInterpretation(RelationDeclaration r) {
		val res = createPartialRelationInterpretation => [
			it.interpretationOf = r
			if(r.parameters.size==2) {
				it.param1 = r.parameters.get(0)
				it.param2 = r.parameters.get(1)
			} else throw new UnsupportedOperationException
		]
		return res
	}
}