# -*- coding: utf-8 -*-

# (c) Copyright by Pierre-Henri Wuillemin, UPMC, 2017
# (pierre-henri.wuillemin@lip6.fr)

# Permission to use, copy, modify, and distribute this
# software and its documentation for any purpose and
# without fee or royalty is hereby granted, provided
# that the above copyright notice appear in all copies
# and that both that copyright notice and this permission
# notice appear in supporting documentation or portions
# thereof, including modifications, that you make.

# THE AUTHOR P.H. WUILLEMIN  DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE, INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS, IN NO EVENT
# SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT
# OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER
# RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER
# IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
# ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE
# OR PERFORMANCE OF THIS SOFTWARE!

from __future__ import print_function

import sys
import os
import math

import pyAgrum as gum
import pydotplus as dot
from itertools import product

from ._utils.pyAgrum_header import pyAgrum_header


class GraphicalBNComparator:
  """
  BNGraphicalComparator allows to compare in multiple way 2 BNs...The smallest assumption is that the names of the variables are the same in the 2 BNs.
  But some comparisons will have also to check the type and domainSize of the variables. The bns have not exactly the
  same role :
  _bn1 is rather the referent model for the comparison whereas _bn2 is the compared one to the referent model

  Parameters
  ----------
  name1 : str or pyAgrum.BayesNet
  	a BN or a filename for reference
  name2 : str or pyAgrum.BayesNet
  	another BN or antoher filename for comparison
  """

  def __init__(self, name1, name2, delta=1e-6):
    self.DELTA_ERROR = delta
    if isinstance(name1, str):
      self._bn1 = gum.loadBN(name1)
      self._bn1.setProperty('name', '"' + os.path.basename(self._bn1.property('name') + '"'))
    else:
      self._bn1 = name1

    if isinstance(name2, str):
      self._bn2 = gum.loadBN(name2)
      self._bn2.setProperty('name', '"' + os.path.basename(self._bn2.property('name') + '"'))
    else:
      self._bn2 = name2

    s1 = set(self._bn1.names())
    s2 = set(self._bn2.names())

    if s1 != s2:
      raise ValueError(
          "The 2 BNs are not comparable! There are names not present in the 2 BNs : " + str(
              s1.symmetric_difference(s2)))

  def _compareBNVariables(self):
    """
    Checks if the two BNs have the same set of variables

    Returns
    -------
    str
      'OK' if the BNs have composed of the same variables, indicates problematic variables otherwise
  
    """
    # it is assumed (checked by the constructor) that _bn1 and _bn2 share the same set of variable names
    for i in self._bn1.nodes():
      v1 = self._bn1.variable(i)
      v2 = self._bn2.variableFromName(v1.name())
      if v2.domainSize() != v1.domainSize():
        return v1.name() + " has not the same domain size in the two bns"

    return "OK"

  def _compareBNParents(self):
    """
    Returns
    -------
    str
      'OK' if _bn2 have (at least) the same variable as b1 and their parents are the same.

    """
    for id1 in self._bn1.nodes():
      id2 = self._bn2.idFromName(self._bn1.variable(id1).name())

      p1 = _parents_name(self._bn1, id1)
      p2 = _parents_name(self._bn2, id2)
      if p1 != p2:
        return self._bn1.variable(id1).name() + " has different parents in the two bns whose names are in " + str(
            p1.symmetric_difference(p2))

    return "OK"

  def _comparePotentials(self, pot1, pot2):

    """
    Compare 2 potentials one in each Bayesian network

    Parameters
    ----------
    pot1 : pyAgrum.Potential
      one of b1's cpts
    pot2 : pyAgrum.Potential
      one of _bn2's cpts

    Returns
    -------
    str
      'OK' if CPTs are the same

    Raises
    ------
    gum.KeyError
      If cpts are not from the same variable
    """
    I1 = gum.Instantiation(pot1)
    I2 = gum.Instantiation(pot2)
    I1.setFirst()
    while not I1.end():
      I2.fromdict(I1.todict())  # copy value on the base of names
      if abs(pot1.get(I1) - pot2.get(I2)) > self.DELTA_ERROR:
        return "Different CPTs for " + pot1.variable(0).name()
      I1 += 1
    return "OK"

  def _compareBNCPT(self):
    """
    Returns
    -------
    str
      'OK' if _bn2 have (at least) the same variable as b1 and their cpts are the same
    """
    for i in self._bn1.nodes():
      res = self._comparePotentials(self._bn1.cpt(i), self._bn2.cpt(self._bn1.variable(i).name()))
      if res != "OK":
        return res
    return "OK"

  def equivalentBNs(self):
    """
    Check if the 2 BNs are equivalent :

    * same variables
    * same graphical structure
    * same parmaeters 
    
    Returns
    -------
    str
      "OK" if bn are the same, a description of the error otherwise

    """

    ret = self._compareBNVariables()
    if ret != "OK":
      return ret

    ret = self._compareBNParents()
    if ret != "OK":
      return ret

    ret = self._compareBNCPT()
    return ret

  def dotDiff(self):
    """ Return a pydotplus graph that compares the arcs of _bn1 (reference) with those of self._bn2.
    full black line: the arc is common for both
    full red line: the arc is common but inverted in _bn2
    dotted black line: the arc is added in _bn2
    dotted red line: the arc is removed in _bn2

    Returns
    -------
    pydotplus.Dot
      the result dot graph
    """
    graph = dot.Dot(graph_type='digraph', bgcolor="transparent")

    # Nodes
    for n in self._bn1.names():
      node = dot.Node('"' + n + '"', style="filled",
                      bgcol="#444444",
                      fgcol="#FFFFFF",
                      tooltip='"({0}) {1}{2}"'.format(self._bn1.idFromName(n), n, ""))
      graph.add_node(node)

    # Arcs
    for n1 in self._bn2.names():
      for n2 in self._bn2.names():
        id1in2 = self._bn2.idFromName(n1)
        id2in2 = self._bn2.idFromName(n2)
        id1in1 = self._bn1.idFromName(n1)
        id2in1 = self._bn1.idFromName(n2)

        if self._bn1.dag().existsArc(id1in1, id2in1):
          if self._bn2.dag().existsArc(id1in2, id2in2):
            # If present in both
            edge = dot.Edge('"' + n1 + '"', '"' + n2 + '"')
            graph.add_edge(edge)
          elif self._bn2.dag().existsArc(id2in2, id1in2):
            # If inversed in _bn2
            edge = dot.Edge('"' + n2 + '"', '"' + n1 + '"')
            edge.set_color("red")
            graph.add_edge(edge)
          else:
            # If removed from _bn2
            edge = dot.Edge('"' + n1 + '"', '"' + n2 + '"', style='dashed')
            edge.set_color("red")
            graph.add_edge(edge)
        else:
          if not self._bn1.dag().existsArc(id2in1, id1in1) and self._bn2.dag().existsArc(id1in2, id2in2):
            # If added to _bn2
            edge = dot.Edge('"' + n1 + '"', '"' + n2 + '"', style='dashed')
            graph.add_edge(edge)

    return graph

  def scores(self):
    """
    Compute Precision, Recall, F-score for self._bn2 compared to self._bn1

    precision and recall are computed considering BN1 as the reference

    Fscor is 2*(recall* precision)/(recall+precision) and is the weighted average of Precision and Recall.

    dist2opt=square root of (1-precision)^2+(1-recall)^2 and represents the euclidian distance to the ideal point (precision=1, recall=1)

    Returns
    -------
    dict[str,double]
      A dictionnary containing 'precision', 'recall', 'fscore', 'dist2opt' and so on.
    """
    # t: True, f: False, p: Positive, n: Negative
    count = {"tp": 0, "tn": 0, "fp": 0, "fn": 0}

    # We look at all combination
    listVariables = self._bn1.names()
    for head, tail in product(listVariables, listVariables):
      if head != tail:
        idHead_1 = self._bn1.idFromName(head)
        idTail_1 = self._bn1.idFromName(tail)

        idHead_2 = self._bn2.idFromName(head)
        idTail_2 = self._bn2.idFromName(tail)

        if self._bn1.dag().existsArc(idHead_1, idTail_1):  # Check arcs head->tail
          if self._bn2.dag().existsArc(idHead_2, idTail_2):  # if arc: Good Direction
            count["tp"] += 1
          elif self._bn2.dag().existsArc(idTail_2, idHead_2):  # If arc: Bad Direction
            count["fn"] += 1
          else:  # If no arc:
            count["fn"] += 1
        elif self._bn1.dag().existsArc(idTail_1, idHead_1):  # Check arcs n2->n1
          if self._bn2.dag().existsArc(idHead_2, idTail_2):  # if arc: Bad Direction
            count["fn"] += 1
          elif self._bn2.dag().existsArc(idTail_2, idHead_2):  # If arc: Good Direction
            count["tp"] += 1
          else:  # If no arc:
            count["fn"] += 1
        else:  # Check if no arc
          if self._bn2.dag().existsArc(idHead_2, idTail_2) or self._bn2.dag().existsArc(idTail_2, idHead_2):  # If arc
            count["fp"] += 1
          else:  # If no arc
            count["tn"] += 1

    # Compute the scores
    if count["tp"] + count["fn"] != 0:
      recall = (1.0 * count["tp"]) / (count["tp"] + count["fn"])
    else:
      recall = 0.0

    if count["tp"] + count["fp"] != 0:
      precision = (1.0 * count["tp"]) / (count["tp"] + count["fp"])
    else:
      precision = 0.0

    if precision + recall != 0.0:
      Fscore = (2 * recall * precision) / (recall + precision)
    else:
      Fscore = 0.0

    return {
      'count'    : count,
      'recall'   : recall,
      'precision': precision,
      'fscore'   : Fscore,
      'dist2opt' : math.sqrt((1 - precision) ** 2 + (1 - recall) ** 2)
    }


def _nodeId(bn, n):
  if type(n) == str:
    return bn.idFromName(n)
  else:
    return n


def _parents_name(bn, n):
  return {bn.variable(n).name() for n in bn.parents(_nodeId(bn, n))}


def module_help(exit_value=1):
  """
  defines help viewed if args are not OK on command line, and exit with exit_value
  """
  print(os.path.basename(sys.argv[0]),
        "_bn1.{" + gum.availableBNExts() + "} self._bn2.{" + gum.availableBNExts() + "}")
  sys.exit(exit_value)


if __name__ == "__main__":
  pyAgrum_header("2011-19")

  if len(sys.argv) != 3:
    module_help()

  cmp = GraphicalBNComparator(sys.argv[1], sys.argv[2])
  print(cmp.equivalentBNs())
