import numpy as np

from domiknows.graph import Concept

from .ilpOntSolver import ilpOntSolver
from .session.solver_session import SolverSession
from .session.gurobi_session import GurobiSession
from .constructor.constructor import ScoreConstructor, ProbConstructor


class MiniSolverDebug(ilpOntSolver):
    ilpSolver = 'mini_debug'

    def __init__(self, graph, ontologiesTuple, _ilpConfig, constructor=None, SessionType=None, **kwargs):
        super().__init__(graph, ontologiesTuple, _ilpConfig)
        def func(node):
            if isinstance(node, Concept):
                return node
        self.names = {}
        for cur_graph in graph:
            self.names.update({concept.name: concept for concept in cur_graph.traversal_apply(func)})
        self.constructor = constructor or ScoreConstructor(**kwargs)
        self.SessionType = SessionType or GurobiSession

    def set_predication(self, predicate, idx, value):
        # predicates_result[concept][idx] = session.get_value(variables[concept, x])
        predicate[idx] = value

    def calculateILPSelection(self, data, *predicates_list):
        concepts_list = []
        for predicates in predicates_list:
            concept_dict = {}
            for predicate, v in predicates.items():
                concept = self.names[predicate]
                concept_dict[concept] = v
            concepts_list.append(concept_dict)

        # call to solve
        concepts_results = self.solve_legacy(data, *concepts_list)

        # prepare result
        results = [{k.name:v for k, v in concepts_result.items()}
                for concepts_result in concepts_results]
        return results

    def solve_legacy(self, data, *predicates_list):
        # data is a list of objects of the base type
        # predicates_list is a list of predicates
        # predicates_list[i] is the dict of predicates of the objects of the base type to the power of i
        # predicates_list[i][concept] is the prediction result of the predicate for concept
        self.myLogger.debug('Start for data %s', data)
        candidates = self.constructor.candidates(data, *predicates_list)

        session = self.SessionType()
        variables, predictions, variables_not, predictions_not = self.constructor.variables(session, candidates, * predicates_list)
        self.constructor.constraints(session, candidates, variables, variables_not, *predicates_list)

        # Set objective
        objective = self.constructor.objective(candidates, variables, predictions, variables_not, predictions_not, *predicates_list)
        session.obj(SolverSession.OTYPE.MAX, objective)
        self.myLogger.debug(' - model %s', session)

        # solve
        self.myLogger.debug('start optimize')
        session.optimize()
        self.myLogger.debug('end optimize')
        self.myLogger.debug(' - model %s', session)

        # collect result
        length = len(data)
        retval = []
        for arity, predicates in enumerate(predicates_list, 1):
            predicates_result = {}
            retval.append(predicates_result)
            for concept, _ in predicates.items():
                self.myLogger.debug('for %s', concept.name)
                predicates_result[concept] = np.zeros((length,) * arity)
                for x in candidates[concept]:
                    #import pdb; pdb.set_trace()
                    if (concept, x) not in variables: continue
                    # NB: candidates generated by 'C' order
                    idx, _ = zip(*x)
                    value = session.get_value(variables[concept, x])
                    self.set_predication(predicates_result[concept], idx, value)
                #import pdb;pdb.set_trace()
                self.myLogger.debug(predicates_result[concept])

        if len(retval) == 1:
            return retval[0]

        return retval

class MiniProbSolverDebug(MiniSolverDebug):
    ilpSolver = 'mini_prob_debug'

    def __init__(self, graph, ontologiesTuple, _ilpConfig, SessionType=None, **kwargs):
        constructor = ProbConstructor(**kwargs)
        super().__init__(graph, ontologiesTuple, _ilpConfig, constructor=constructor, SessionType=SessionType)
