package uk.ac.man.cs.lethe.internal.application.benchmarks

import uk.ac.man.cs.lethe.internal.dl.owlapi.ModuleExtractor
import uk.ac.man.cs.lethe.internal.dl.owlapi.{OWLApiConverter, OWLApiInterface, OWLReasonerFacade}
import uk.ac.man.cs.lethe.internal.dl.filters.OWLOntologyFilters

import java.io.File
import java.util.Date
import java.util.concurrent.TimeoutException
import org.semanticweb.owlapi.model.OWLAxiom
import org.semanticweb.owlapi.model.OWLOntology

import scala.concurrent._
import scala.concurrent.duration._
import scala.util.Random
import uk.ac.man.cs.lethe.internal.dl.datatypes.DLHelpers
import uk.ac.man.cs.lethe.internal.dl.datatypes.DLStatement
import uk.ac.man.cs.lethe.internal.dl.datatypes.Ontology
import uk.ac.man.cs.lethe.internal.dl.forgetting.DirectALCForgetter
import uk.ac.man.cs.lethe.internal.forgetting.ForgetterWithBackgroundKB



object BackgroundOntologyForgettingExperimenter { 

  implicit val exec = new CustomExecutionContext()
  
  def experimentSequence(owlOntology: OWLOntology, subsetSize: Int, runs: Int) = { 
    val forgetter: ForgetterWithBackgroundKB[Ontology, String] = DirectALCForgetter

    for(i <- 0 to runs){ 
      val ontologyPart = pickOntologySubset(owlOntology, subsetSize)
      val background = OWLApiConverter.convert(owlOntology)
      val forget = pickSignature(ontologyPart)

      val start = new Date().getTime
	  
      val result = forgetter.forget(ontologyPart, forget, background)
      val duration = new Date().getTime - start

      println(i)
      println("\tStatements: "+result.statements.size+ "("+ontologyPart.statements.size+")")
      println("\tSize: "+result.size+" ("+ontologyPart.size+")")
      println("\tSigSize: "+result.signature.size+ "("+ontologyPart.signature.size+")")
      println("\tDuration: "+(duration.toDouble/1000)+" sec. ")
      println("\tForgot "+forget.size+" symbols.")
      println()

      if(result.signature.exists(forget)){ 
	println("not in signature: "+result.signature.filter(forget).mkString(", "))
      }
    }
  }

  def experiment(owlOntology: OWLOntology, 
		 partSize: Int, 
		 sigSize: Int,
		 timeout: Duration) = { 

    val forgetter: ForgetterWithBackgroundKB[Ontology, String] = DirectALCForgetter


    val sample = pickRandomCoherentSample(owlOntology, partSize, sigSize)
    val ontology = sample.ontologyPart
    val signature = sample.signature
    val background = sample.background
    val toForget = sample.forget

    println("Input statements: "+ontology.statements.size)
    println("Input avg axiom size: "+(ontology.size.toDouble/ontology.statements.size))
    println("Input max statement size: "+ontology.statements.map(_.size).max)
    println("Input signature: "+ontology.atomicConcepts.size)
    println("Background ontology statements: "+background.statements.size)
    println("Desired signature size: "+signature.size)
    println()

    // var wholeOntology = OWLApiConverter.convert(owlOntology)
    // val toForget = wholeOntology.atomicConcepts--signature
    // val background = Ontology.buildFrom(wholeOntology.statements.toSet--ontology.statements)

    val startTime = new Date().getTime

    val futureResult = future[Ontology]{ 
      forgetter.forget(ontology, toForget, background)
    }
    try { 
      val result = Await.result(futureResult, timeout)
      val duration = new Date().getTime-startTime

      println("Output statements: "+result.statements.size)
      println("Output avg axiom size: "+(result.size.toDouble/result.statements.size))
      if(result.statements.size>0)
	println("Output max statement size: "+result.statements.map(_.size).max)
      println("Output signature: "+result.atomicConcepts.size)
      println("Duration: "+(duration.toDouble/1000)+"sec.")
      if(result.atomicConcepts.exists(!signature(_)))
	println("not in signature: "+result.atomicConcepts.filter(toForget).mkString(", "))
    } catch { 
      case _: TimeoutException => 
	println("Timeout.")
      exec.lastThread.getOrElse(throw new RuntimeException("Not started")).stop()
    }
      println()
      println()
  }

  def experiment(owlOntology: OWLOntology, probability: Double, timeout: Duration): Unit = { 

    val forgetter: ForgetterWithBackgroundKB[Ontology, String] = DirectALCForgetter


    val preparedOnt = ModuleExtractor.prepare(owlOntology)

    val sample = pickRandomCoherentSample(preparedOnt, probability)
    val ontology = sample.ontologyPart
    if(ontology.size<1){ 
      println("Empty sample.")
      return
    }
    val signature = sample.signature
    val background = sample.background
    val toForget = sample.forget

    println("Input statements: "+ontology.statements.size)
    println("Input avg axiom size: "+(ontology.size.toDouble/ontology.statements.size))
    println("Input max statement size: "+ontology.statements.map(_.size).max)
    println("Input signature: "+ontology.atomicConcepts.size)
    println("Desired signature size: "+signature.size)
    println()

    // var wholeOntology = OWLApiConverter.convert(owlOntology)
    // val toForget = wholeOntology.atomicConcepts--signature
    // val background = Ontology.buildFrom(wholeOntology.statements.toSet--ontology.statements)

    val startTime = new Date().getTime
    
    val futureResult = future[Ontology]{ 
      forgetter.forget(ontology, toForget, background)
    }
    try{ 
      var result = Await.result(futureResult, timeout)
      val duration = new Date().getTime-startTime

      result = DLHelpers.splitConjunctions(result)

      println("Output statements: "+result.statements.size)
      println("Output avg axiom size: "+(result.size.toDouble/result.statements.size))
      println("Output max statement size: "+result.statements.map(_.size).max)
      println("Output signature: "+result.atomicConcepts.size)
      println("Duration: "+(duration.toDouble/1000)+"sec.")
      if(result.signature.exists(toForget))
	println("not in signature: "+result.signature.filter(toForget).mkString(", "))
      println()
      println()
    } catch { 
      case _: TimeoutException => println("Timeout.")
      exec.lastThread.getOrElse(throw new RuntimeException("Not started")).stop()
      case e: Throwable => e.printStackTrace()
    }
  }

  def experiment(owlOntology: OWLOntology, subsetSize: Int) = { 

    val forgetter: ForgetterWithBackgroundKB[Ontology, String] = DirectALCForgetter

    val ontologyPart = pickOntologySubset(owlOntology, subsetSize)
    val background = OWLApiConverter.convert(owlOntology)
    val forget = pickSignature(ontologyPart)


    println("Input: \n"+ontologyPart)
    println("Forget "+forget.mkString(", "))
    println(new Date())
	  
    val result = forgetter.forget(ontologyPart, forget, background)

    println("Done: "+new Date())
    println("Input: \n"+ontologyPart)
    println("Forget "+forget.mkString(", "))
    println("Result: \n"+result)
    if(result.signature.exists(forget))
      println("not in signature: "+result.signature.filter(forget).mkString(", "))
  }


  case class Sample(ontologyPart: Ontology, signature: Set[String], background: Ontology, forget: Set[String])

  def pickRandomCoherentSample(owlOntology: OWLOntology, probability: Double): Sample = { 
    var owlAxioms = ModuleExtractor.randomGenuineModule(owlOntology)

    var moduleStatements = owlAxioms.flatMap(OWLApiConverter.convert)

    var allSymbols = moduleStatements.flatMap(_.atomicConcepts)

    var statements = Set[DLStatement]()
    var symbols = Set[String]()

    var rest = Set[DLStatement]()
    var forget = Set[String]()
    
    moduleStatements.foreach(st => if(Random.nextDouble<probability) statements+=st else rest+=st)
    allSymbols.foreach(s => if(Random.nextDouble<probability) symbols+=s else forget+=s)

    val resultingOntology = Ontology.buildFrom(statements)
    val background = Ontology.buildFrom(rest)

    Sample(resultingOntology, symbols, background, forget)
  }

  def pickRandomCoherentSample(owlOntology: OWLOntology, 
			       partSize: Int, 
			       sigSize: Int): Sample = { 
    var moduleStatements = ModuleExtractor.randomGenuineModule(owlOntology).flatMap(OWLApiConverter.convert)
    while(moduleStatements.size<partSize)
      moduleStatements++=ModuleExtractor.randomGenuineModule(owlOntology).flatMap(OWLApiConverter.convert)


    val statements = Random.shuffle(moduleStatements.toSeq).slice(0,partSize)
    var allSymbols = statements.flatMap(_.atomicConcepts).toSet
    val symbols = Random.shuffle(allSymbols.toSeq).slice(0, sigSize).toSet

    var rest = moduleStatements -- statements
    var forget = moduleStatements.flatMap(_.signature).toSet --symbols

    val resultingOntology = Ontology.buildFrom(statements)
    val background = Ontology.buildFrom(rest)

    Sample(resultingOntology, symbols, background, forget)
  }


  // Methods based on supplied sizes

  def pickOntologySubset(owlOntology: OWLOntology, subsetSize: Int): Ontology = { 
    var owlAxioms = Set[OWLAxiom]()

    while(owlAxioms.size<subsetSize)
      owlAxioms++=ModuleExtractor.randomGenuineModule(owlOntology)

    val statements = owlAxioms.flatMap(OWLApiConverter.convert)

    val subset = Random.shuffle(statements.toSeq).slice(0,subsetSize)

    val result = new Ontology()
    subset.foreach(result.addStatement)

    result
  }

  def pickSignature(ontology: Ontology) = { 
    val symbols = ontology.atomicConcepts
    var result = Set[String]()

    symbols.foreach{ symbol => 
      if(Random.nextBoolean)
	result += symbol
    }

    result
  }


  def main(args: Array[String]): Unit = { 
    print("Loading ontology...")
    val ontology = OWLApiInterface.getOWLOntology(new File(args(0)))
    OWLOntologyFilters.restrictToALCH(ontology)
    OWLOntologyFilters.clearABox(ontology)

    if(!OWLReasonerFacade.isConsistent(ontology))
       return 


    if(args(1)=="prob"){ 
      val repeat = args(3).toInt
      val probability = args(2).toDouble
      val timeout = args(4).toInt
    
      println("Running sequence of "+repeat+" runs, probability "+probability)
      println("Ontology: "+ontology)

      for(i <- 1 to repeat){ 
      	println("Experiment "+i)
      	println("-----------------")

      	experiment(ontology, probability, timeout seconds)
      }
    } 
    else if(args(1)=="num"){ 
      val repeat = args(4).toInt
      val partSize = args(2).toInt
      val sigSize = args(3).toInt
      val timeout = args(5).toInt

      val restricted = ModuleExtractor.prepare(ontology)

      if(restricted.getLogicalAxioms().size<partSize || OWLApiConverter.convert(restricted).tbox.axioms.size<partSize){
	println("Too small")
	return
      }
    
      println("Running sequence of "+repeat+" runs, partSize: "+partSize+", sigSize: "+sigSize)
      println("Ontology: "+ontology)

      for(i <- 1 to repeat){ 
      	println("Experiment "+i)
      	println("-----------------")

      	experiment(ontology, partSize, sigSize, timeout seconds)
      }
    }
    else 
      experiment(ontology, args(1).toInt)
  }

  
}
