# Author: Stefan Wunsch CERN  09/2019

################################################################################
# Copyright (C) 1995-2019, Rene Brun and Fons Rademakers.                      #
# All rights reserved.                                                         #
#                                                                              #
# For the licensing terms see $ROOTSYS/LICENSE.                                #
# For the list of contributors see $ROOTSYS/README/CREDITS.                    #
################################################################################

from .. import pythonization


def Compute(self, x):
    import numpy as np

    import ROOT

    # numpy.array is a factory and the actual type of a numpy array is numpy.ndarray
    if isinstance(x, np.ndarray):
        if len(x.shape) == 1:
            x_ = ROOT.VecOps.AsRVec(x)
            y = self._OriginalCompute(x_)
            return np.asarray(y)
        elif len(x.shape) == 2:
            x_ = ROOT.TMVA.Experimental.AsRTensor(x)
            y = self._OriginalCompute(x_)
            return np.asarray(y)
        else:
            raise Exception("Call to Compute can process only numpy arrays of rank 1 or 2.")

    # As fall-through we go to the original compute function and use the error-handling from cppyy
    return self._OriginalCompute(x)


@pythonization("RBDT", ns="TMVA::Experimental", is_prefix=True)
def pythonize_rbdt(klass):
    # Parameters:
    # klass: class to be pythonized

    klass._OriginalCompute = klass.Compute
    klass.Compute = Compute
