import ecole
from ecole.typing import ObservationFunction


class NodeDepth(ObservationFunction):
    def __init__(self):
        super().__init__()

    def before_reset(self, model: ecole.environment.Environment):
        super().before_reset(model)

    def extract(self, model: ecole.environment.Environment, done: bool) -> int:
        if done:
            return -1
        else:
            m = model.as_pyscipopt()
            curr_node = m.getCurrentNode()
            depth = curr_node.getDepth()
            return depth
