"""
Main module of the neda inference package

This introduces the `Environment` class, which we use to facilitate repeating
tasks like running MCMC, calculating or estimating evidence. Finally, the
`main` function runs the whole scheme. Note that both of these are imported
into the ``tracklib.analysis.neda`` namespace, i.e. can be imported from there
(instead of ``tracklib.analysis.neda.neda``).
"""

import numpy as np
import scipy.optimize

from . import mcmc

class Environment:
    """
    Environment for inference runs

    This is essentially a semi-disguised way of using global variables for the
    `Trajectory` we are looking at, the inference `models.Model` we want to use, and
    the `MCMCscheme` together with the `MCMCconfig`. We use this approach
    mainly for code readability / useability.

    Attributes
    ----------
    traj : Trajectory
    model : models.Model
    MCMCconfig : dict
        see `tracklib.util.mcmc.Sampler`
    MCMCscheme : mcmc.MCMCScheme, optional
        defaults to an instance of `mcmc.TPWMCMC`
    """
    # Here we use ``env`` instead of ``self``.
    def __init__(self, traj, model, MCMCconfig, MCMCscheme=None):
        self.traj = traj
        self.model = model
        self.MCMCconfig = MCMCconfig
        if MCMCscheme is None:
            self.MCMCscheme = mcmc.TPWMCMC()
        else: # pragma: no cover
            self.MCMCscheme = MCMCscheme

    def runMCMC(env, prior):
        """
        Run the `MCMCscheme` with a given `prior`

        Parameters
        ----------
        prior : Prior

        Returns
        -------
        mcmc.MCMCRun
        """
        env.MCMCscheme.setup(env.traj, env.model, prior)
        env.MCMCscheme.configure(**env.MCMCconfig)
        return env.MCMCscheme.run()

    def posterior_density(env, mcmcrun, prior, trace_eval,
                          nSample_proposal=float('inf')):
        r"""
        Estimate the posterior density at a specific point from an `MCMCRun`

        Parameters
        ----------
        mcmcrun : MCMCRun
        prior : Prior
        trace_eval : Loopingtrace
            the point in parameter space for which to estimate the posterior
            density
        nSample_proposal : float, optional
            how many samples from the proposal distribution to use for
            evaluation of the exit rate. Defaults to exhaustive sampling, so
            may be reduced if that seems excessive.

        Returns
        -------
        float
            the estimated posterior density for `!trace_eval`. Might be
            ``np.inf`` if the MCMC ensemble collapsed to a single sample.

        Notes
        -----
        This estimation follows eq. (9) of [1]_. Essentially, we estimate the
        density at the sample point as the number of times it is entered from
        another point in the steady state ensemble, times the average number of
        steps it takes to leave this point again, divided by the total sample
        number :math:`M`. This yields

        .. math:: \hat{p}(\theta^*) = \frac{N_\text{enter}N_\text{stay}}{M} = \frac{P_\text{enter}}{k_\text{exit}} = \frac{M^{-1}\sum_{m=1}^M\, \alpha(\theta^{(m)}, \theta^*) q(\theta^{(m)}, \theta^*)}{J^{-1}\sum_{j=1}^J\, \alpha(\theta^*, \theta^{(j)})}\,,

        where the :math:`\theta^{(m)}` in the numerator are the samples from
        the given MCMC run (i.e. are assumed to be sampled from the posterior
        distribution), while :math:`\theta^{(j)}` in the denominator are
        independent samples from the proposal distribution
        :math:`q(\theta_\text{from} = \theta^*, \theta_\text{to})`
        around the evaluation point :math:`\theta^*`. Finally,
        :math:`\alpha(\theta_\text{from}, \theta_\text{to})` is the
        acceptance probability for a given step.

        References
        ----------
        .. [1] Chib, S. & Jeliazkov, I. Marginal Likelihood From the Metropolis-Hastings Output. Journal of the American Statistical Association 96, 270-281 (2001)
        """
        # Check whether the MCMC sample collapsed, in which case the posterior
        # estimation would diverge
        logLs = mcmcrun.logLs_trunc()
        if np.all(logLs == logLs[0]):
            print('MCMC ensemble collapsed to single sample')
            return np.inf

        L_eval = env.MCMCscheme.likelihood(env.traj, trace_eval, env.model, prior)

        # Chance to enter the target state from steady state
        def get_p_step_to_eval(trace): return env.MCMCscheme.stepping_probability(trace, trace_eval)
        p_step_to_eval = np.array(mcmcrun.evaluate(get_p_step_to_eval))

        # Note: p_accept = 1 *if* we evaluate at the maximum likelihood trace
        p_accept_move = env.MCMCscheme.acceptance_probability(mcmcrun.logLs_trunc(), L_eval)
        p_enter = np.mean(p_step_to_eval * p_accept_move)

        # Survival in target state
        neighbor_logLs = [env.MCMCscheme.likelihood(env.traj, trace, env.model, prior) \
                          for trace in env.MCMCscheme.gen_proposal_sample_from(trace_eval, nSample=nSample_proposal)]
        k_leave = np.mean(env.MCMCscheme.acceptance_probability(L_eval, neighbor_logLs))

        # Occupancy of the target state = p_enter * survival
        return p_enter / k_leave

    def evidence(env, prior, mcmcrun):
        """
        Calculate the (log-)evidence from an `MCMCRun`

        Parameters
        ----------
        prior : Prior
            the prior over `Loopingtraces <Loopingtrace>` used for the MCMC run
        mcmcrun : MCMCRun
            the MCMC sample

        Returns
        -------
        float
            the estimated log-evidence

        Notes
        -----
        The log-evidence is given by ``log(likelihood) + log(prior) -
        log(posterior)``.

        See also
        --------
        posterior_density
        """
        trace_eval, L_eval = mcmcrun.best_sample_L()
        log_post = np.log(env.posterior_density(mcmcrun, prior, trace_eval, nSample_proposal=len(mcmcrun.samples)))
        log_prior = prior.logpi(trace_eval)

        if np.isinf(log_post):
            return L_eval
        else:
            return L_eval + log_prior - log_post

    def evidence_differential(env, prior, ref_prior, ref_mcmcrun):
        r"""
        Estimate the relative evidence given a reference point

        Parameters
        ----------
        prior : Prior
            the prior whose evidence we want to estimate
        ref_prior : Prior
            the reference point in prior space
        ref_mcmcrun : MCMCRun
            an `MCMCRun` using the reference prior

        Returns
        -------
        float
            the estimated log-evidence, relative to the reference

        Notes
        -----
        The relative evidence is given by the expectation value of the prior
        ratio over the MCMC sample: ``E = < prior/ref_prior >`` where ``<.>``
        indicates an average over the MCMC sample.
        """
        return np.log(np.mean(np.exp(prior.logpi_vectorized(ref_mcmcrun.samples) - \
                                     ref_prior.logpi_vectorized(ref_mcmcrun.samples))))

def main(traj, model, priorfam,
         MCMCconfig, MCMCscheme=None,
         max_iterations=20, min_iterations=5,
         return_ = 'nothing', # 'traj', 'dict', or anything else
         show_progress=False, assume_notebook_for_progressbar=True,
        ):
    """
    Run the neda looping inference scheme

    The output of the inference run will be assembled into a dict whose fields
    are detailed below. Where exactly this dict will end up depends on the
    setting of `!return_`.

    Parameters
    ----------
    traj : Trajectory
        the `Trajectory` whose looping profile to infer
    model : models.Model
        the inference model to use
    priorfam : ParametricFamily
        a family of priors
    MCMCconfig : dict
        configuration for the MCMC runs. See `tracklib.util.mcmc.Sampler`
    MCMCscheme : mcmc.MCMCScheme, optional
        the sampling scheme to use, i.e. an object implementing the
        `mcmc.MCMCScheme` interface. Defaults to `mcmc.TPWMCMC'
    max_iterations : int, optional
        maximum number of iterations to run
    min_iterations : int, optional
        run at least this many iterations
    return_ : {'nothing', 'None', 'traj', 'dict'}
        what the return value of this function should be. Generally, the
        results of the inference run will be stored in a dict. If
        ``return_='dict'``, that dict is directly returned. Otherwise it is
        written to ``traj.meta['neda']``, and if ``return_='traj'`` the
        trajectory is returned (useful for parallelization). Otherwise this
        function returns nothing, i.e. the inference results can be accessed
        simply from ``traj.meta['neda']`` after calling this function.
    show_progress : bool, optional
        whether to show a progress bar. Note: for ``show_progress=True`` it
        might happen that the termination condition is fulfilled during the
        last "required" run, in which case the progress bar will stop before
        reaching its maximum.
    assume_notebook_for_progressbar : bool, optional
        set to ``False`` if running outside Jupyter notebook to show
        progressbar ASCII style

    Returns
    -------
    prior_prams : np.array
        the prior parameters for each iteration
    mcmcrun : list of MCMCRun
        the MCMC runs for each iteration
    evidence : np.array
        the evidence estimated from each `MCMCRun`
    evidence_diff : np.array
        the estimated evidence differential (relative evidence) at each
        iteration
    final : dict
        the corresponding values for the iteration that should be considered
        the final result. Has entries ``'prior_params'``, ``'mcmcrun'``,
        ``'evidence'``, and ``'iteration'`` where the last one is the index of
        the iteration the other three refer to.

    Example
    -------
    Assuming we have a `Trajectory` ``traj`` that we want to run the looping
    inference on:

    >>> # Set up inference scheme
    ... looppositions = [(0, 0), (0, -1)] # 2 states: unlooped (=0), fully looped (=1)
    ... model = neda.models.RouseModel(N=20, D=1, k=5, k_extra=1, looppositions=looppositions)
    ... priorfam = neda.ParametricFamily(start_params=(0), bounds=[(None, 0)])
    ... priorfam.get = lambda logq : neda.priors.GeometricPrior(logq, nStates=len(looppositions))
    ... MCMCconfig = {
    ...         'iterations' : 1000,
    ...         'burn_in'    :  100,
    ...         }
    ...
    ... # Run the inference
    ... neda.main(traj, model, priorfam, MCMCconfig, show_progress=True)
    ...
    ... # Visualize output
    ... from matplotlib import pyplot as plt
    ... neda.plot.butterfly(traj)
    ... plt.show()

    See also
    --------
    tracklib.analysis.kli.fit_RouseParams
    """
    assert min_iterations >= 2
    assert return_ in {'nothing', 'None', 'traj', 'dict'}

    # Set up environment
    env = Environment(traj, model, MCMCconfig, MCMCscheme)

    # Set up iterative scheme
    iterations = range(max_iterations)
    if show_progress: # pragma: no cover
        if assume_notebook_for_progressbar:
            from tqdm.notebook import tqdm
        else:
            from tqdm import tqdm
        iterations = tqdm(iterations, total=min(min_iterations, max_iterations))

    prior_params = [priorfam.start_params]
    mcmcruns = []
    evidences = []
    evidence_diffs = []

    # Run iterations
    for it in iterations:
        prior = priorfam.get(*prior_params[-1])
        mcmcruns.append(env.runMCMC(prior))
        evidences.append(env.evidence(prior, mcmcruns[-1]))

        if it+1 >= min_iterations:
            if np.argmax(evidences) < it:
                break

        # Maximize estimated evidence differential to find new prior parameters
        def minimization_target(*params):
            return -env.evidence_differential(priorfam.get(*params), prior, mcmcruns[-1])
        minimization_result = scipy.optimize.minimize(minimization_target,
                                                      x0=prior_params[-1],
                                                      bounds=priorfam.bounds)

        if not minimization_result.success: # pragma: no cover
            print(minimization_result)
            raise RuntimeError('Relative evidence maximization did not converge')

        evidence_diffs.append(-minimization_result.fun)
        prior_params.append(tuple(minimization_result.x))

    # Output
    best_it = np.argmax(evidences)
    output = {
        'prior_params'  : np.array(prior_params),
        'mcmcrun'       : mcmcruns,
        'evidence'      : np.array(evidences),
        'evidence_diff' : np.array(evidence_diffs),
        'final'         : {
            'prior_params' : prior_params[best_it],
            'mcmcrun'      : mcmcruns[best_it],
            'evidence'     : evidences[best_it],
            'iteration'    : best_it,
            'best_trace'   : mcmcruns[best_it].best_sample_L()[0],
            },
        }

    if return_ == 'dict':
        return output
    else:
        traj.meta['neda'] = output
        if return_ == 'traj':
            return traj
        else:
            return None


from __future__ import print_function
from __future__ import absolute_import
from __future__ import division
from past.utils import old_div
import anuga
import numpy as num
import math
from . import inlet_enquiry

from anuga.utilities.system_tools import log_to_file
from anuga.utilities.numerical_tools import ensure_numeric

class Structure_operator(anuga.Operator):
    """Structure Operator - transfer water from one rectangular box to another.
    Sets up the geometry of problem

    This is the base class for structures (culverts, pipes, bridges etc). Inherit from this class (and overwrite
    discharge_routine method for specific subclasses)

    Input: Two points, pipe_size (either diameter or width, depth),
    mannings_rougness,
    """

    counter = 0

    def __init__(self,
                 domain,
                 end_points=None,
                 exchange_lines=None,
                 enquiry_points=None,
                 invert_elevations=None,
                 width=None,
                 height=None,
                 diameter=None,
                 z1=None,
                 z2=None,
                 blockage=None,
                 barrels=None,
                 #culvert_slope=None,
                 apron=None,
                 manning=None,
                 enquiry_gap=None,
                 use_momentum_jet=False,
                 zero_outflow_momentum=True,
                 use_old_momentum_method=True,
                 always_use_Q_wetdry_adjustment=True,
                 force_constant_inlet_elevations=False,
                 description=None,
                 label=None,
                 structure_type=None,
                 logging=None,
                 verbose=None):

        """
        exchange_lines define the input lines for each inlet.

        If end_points = None, then the culvert_vector is calculated in the
        directions from the centre of echange_line[0] to centre of exchange_line[1}

        If end_points != None, then culvert_vector is unit vector in direction
        end_point[1] - end_point[0]
        """

        anuga.Operator.__init__(self,domain)

        self.master_proc = 0
        self.end_points = ensure_numeric(end_points)
        self.exchange_lines = ensure_numeric(exchange_lines)
        self.enquiry_points = ensure_numeric(enquiry_points)
        self.invert_elevations = ensure_numeric(invert_elevations)

        assert self.end_points is None or self.exchange_lines is None

        if height is None:
            height = width

        if width is None:
            width = diameter

        if apron is None:
            apron = width

        assert width is not None

        self.width  = width
        self.height = height
        self.diameter = diameter
        self.z1 = z1
        self.z2 = z2
        self.blockage = blockage
        self.barrels = barrels
        self.apron  = apron
        self.manning = manning
        self.enquiry_gap = enquiry_gap
        self.use_momentum_jet = use_momentum_jet
        self.zero_outflow_momentum = zero_outflow_momentum
        if use_momentum_jet and zero_outflow_momentum:
            msg = "Can't have use_momentum_jet and zero_outflow_momentum both True"
            raise Exception(msg)
        self.use_old_momentum_method = use_old_momentum_method
        self.always_use_Q_wetdry_adjustment = always_use_Q_wetdry_adjustment

        if description is None:
            self.description = ' '
        else:
            self.description = description

        if label is None:
            self.label = "structure_%g" % Structure_operator.counter
        else:
            self.label = label + '_%g' % Structure_operator.counter

        if structure_type is None:
            self.structure_type = 'generic structure'
        else:
            self.structure_type = structure_type

        self.verbose = verbose

        # Keep count of structures
        Structure_operator.counter += 1

        # Slots for recording current statistics
        self.accumulated_flow = 0.0
        self.discharge = 0.0
        self.discharge_abs_timemean = 0.0
        self.velocity = 0.0
        self.outlet_depth = 0.0
        self.delta_total_energy = 0.0
        self.driving_energy = 0.0

        if exchange_lines is not None:
            self.__process_skew_culvert()
        elif end_points is not None:
            self.__process_non_skew_culvert()
        else:
            raise Exception('Define either exchange_lines or end_points')

        self.inlets = []
        line0 = self.exchange_lines[0] #self.inlet_lines[0]
        if self.apron is None:
            poly0 = line0
        else:
            offset = -self.apron*self.outward_vector_0
            #print line0
            #print offset
            poly0 = num.array([ line0[0], line0[1], line0[1]+offset, line0[0]+offset])
            #print poly0
        if self.invert_elevations is None:
            invert_elevation0 = None
        else:
            invert_elevation0 = self.invert_elevations[0]

        enquiry_point0 = self.enquiry_points[0]

        #outward_vector0 = - self.culvert_vector
        self.inlets.append(inlet_enquiry.Inlet_enquiry(
                           self.domain,
                           poly0,
                           enquiry_point0,
                           invert_elevation = invert_elevation0,
                           outward_culvert_vector = self.outward_vector_0,
                           verbose = self.verbose))

        if force_constant_inlet_elevations:
            # Try to enforce a constant inlet elevation
            inlet_global_elevation = self.inlets[-1].get_average_elevation()
            self.inlets[-1].set_elevations(inlet_global_elevation)

        tris_0 = self.inlets[0].triangle_indices
        #print tris_0
        #print self.domain.centroid_coordinates[tris_0]

        line1 = self.exchange_lines[1]
        if self.apron is None:
            poly1 = line1
        else:
            offset = -self.apron*self.outward_vector_1
            #print line1
            #print offset
            poly1 = num.array([ line1[0], line1[1], line1[1]+offset, line1[0]+offset])
            #print poly1

        if self.invert_elevations is None:
            invert_elevation1 = None
        else:
            invert_elevation1 = self.invert_elevations[1]
        enquiry_point1 = self.enquiry_points[1]

        self.inlets.append(inlet_enquiry.Inlet_enquiry(
                           self.domain,
                           poly1,
                           enquiry_point1,
                           invert_elevation = invert_elevation1,
                           outward_culvert_vector = self.outward_vector_1,
                           verbose = self.verbose))

        if force_constant_inlet_elevations:
            # Try to enforce a constant inlet elevation
            inlet_global_elevation = self.inlets[-1].get_average_elevation()
            self.inlets[-1].set_elevations(inlet_global_elevation)

        tris_1 = self.inlets[1].triangle_indices

        self.set_logging(logging)

    def __call__(self):

        timestep = self.domain.get_timestep()

        Q, barrel_speed, outlet_depth = self.discharge_routine()

        old_inflow_depth = self.inflow.get_average_depth()
        old_inflow_stage = self.inflow.get_average_stage()
        old_inflow_xmom = self.inflow.get_average_xmom()
        old_inflow_ymom = self.inflow.get_average_ymom()

        # Implement the update of flow over a timestep by
        # using a semi-implict update. This ensures that
        # the update does not create a negative depth
        if old_inflow_depth > 0.0 :
            dt_Q_on_d = old_div(timestep*Q,old_inflow_depth)
        else:
            dt_Q_on_d = 0.0

        always_use_Q_wetdry_adjustment = self.always_use_Q_wetdry_adjustment
        use_Q_wetdry_adjustment = ((always_use_Q_wetdry_adjustment) |\
            (old_inflow_depth*self.inflow.get_area() <= Q*timestep))
        # If we use the 'Q_adjustment', then the discharge is rescaled so that
        # the depth update is:
        #    new_inflow_depth*inflow_area =
        #    old_inflow_depth*inflow_area -
        #    timestep*Q*(new_inflow_depth/old_inflow_depth)
        # The last term in () is a wet-dry improvement trick (which rescales Q)
        #
        # Before Feb 2015 this rescaling was always done (even if the flow was
        # not near wet-dry). Now if use_old_Q_adjustment=False, the rescaling
        # is only done if required to avoid drying
        #
        #
        factor = 1.0/(1.0 + old_div(dt_Q_on_d,self.inflow.get_area()))

        if use_Q_wetdry_adjustment:
            # FIXME:
            new_inflow_depth = old_inflow_depth*factor

            if(old_inflow_depth>0.):
                timestep_star = old_div(timestep*new_inflow_depth,old_inflow_depth)
            else:
                timestep_star = 0.

        else:
            new_inflow_depth = old_inflow_depth - old_div(timestep*Q,self.inflow.get_area())
            timestep_star = timestep

        if(self.use_old_momentum_method):
            # This method is here for consistency with the old version of the
            # routine
            new_inflow_xmom = old_inflow_xmom*factor
            new_inflow_ymom = old_inflow_ymom*factor

        else:
            # For the momentum balance, note that Q also advects the momentum,
            # The volumetric momentum flux should be Q*momentum/depth, where
            # momentum has an average value of new_inflow_mom (or
            # old_inflow_mom).  We use old_inflow_depth for depth
            #
            #     new_inflow_xmom*inflow_area =
            #     old_inflow_xmom*inflow_area -
            #     [timestep*Q]*new_inflow_xmom/old_inflow_depth
            # and:
            #     new_inflow_ymom*inflow_area =
            #     old_inflow_ymom*inflow_area -
            #     [timestep*Q]*new_inflow_ymom/old_inflow_depth
            #
            # The units balance: m^2/s*m^2 = m^2/s*m^2 - s*m^3/s*m^2/s *m^(-1)
            #
            if old_inflow_depth > 0.:
                if use_Q_wetdry_adjustment:
                    # Replace dt*Q with dt*Q*new_inflow_depth/old_inflow_depth = dt_Q_on_d*new_inflow_depth
                    factor2 = 1.0/(1.0 + old_div(dt_Q_on_d*new_inflow_depth,(old_inflow_depth*self.inflow.get_area())))
                else:
                    factor2 = 1.0/(1.0 + old_div(timestep*Q,(old_inflow_depth*self.inflow.get_area())))
            else:
                factor2 = 0.

            new_inflow_xmom = old_inflow_xmom*factor2
            new_inflow_ymom = old_inflow_ymom*factor2

        self.inflow.set_depths(new_inflow_depth)

        #inflow.set_xmoms(Q/inflow.get_area())
        #inflow.set_ymoms(0.0)

        self.inflow.set_xmoms(new_inflow_xmom)
        self.inflow.set_ymoms(new_inflow_ymom)

        loss = (old_inflow_depth - new_inflow_depth)*self.inflow.get_area()
        xmom_loss = (old_inflow_xmom - new_inflow_xmom)*self.inflow.get_area()
        ymom_loss = (old_inflow_ymom - new_inflow_ymom)*self.inflow.get_area()

        # set outflow
        outflow_extra_depth = old_div(Q*timestep_star,self.outflow.get_area())
        outflow_direction = - self.outflow.outward_culvert_vector
        #outflow_extra_momentum = outflow_extra_depth*barrel_speed*outflow_direction

        gain = outflow_extra_depth*self.outflow.get_area()

        #print gain, loss
        assert num.allclose(gain-loss, 0.0)

        # Stats
        self.accumulated_flow += gain
        self.discharge  = old_div(Q*timestep_star,timestep)
        self.discharge_abs_timemean += old_div(gain,self.domain.yieldstep)
        self.velocity =   barrel_speed
        self.outlet_depth = outlet_depth

        new_outflow_depth = self.outflow.get_average_depth() + outflow_extra_depth

        self.outflow.set_depths(new_outflow_depth)

        if self.use_momentum_jet:
            # FIXME (SR) Review momentum to account for possible hydraulic jumps at outlet
            # FIXME (GD) Depending on barrel speed I think this will be either
            # a source or sink of momentum (considering the momentum losses
            # above). Might not always be reasonable.
            #new_outflow_xmom = self.outflow.get_average_xmom() + outflow_extra_momentum[0]
            #new_outflow_ymom = self.outflow.get_average_ymom() + outflow_extra_momentum[1]
            new_outflow_xmom = barrel_speed*new_outflow_depth*outflow_direction[0]
            new_outflow_ymom = barrel_speed*new_outflow_depth*outflow_direction[1]

        elif self.zero_outflow_momentum:
            new_outflow_xmom = 0.0
            new_outflow_ymom = 0.0
            #new_outflow_xmom = outflow.get_average_xmom()
            #new_outflow_ymom = outflow.get_average_ymom()

        else:
            # Add the momentum lost from the inflow to the outflow. For
            # structures where barrel_speed is unknown + direction doesn't
            # change from inflow to outflow
            new_outflow_xmom = self.outflow.get_average_xmom() + old_div(xmom_loss,self.outflow.get_area())
            new_outflow_ymom = self.outflow.get_average_ymom() + old_div(ymom_loss,self.outflow.get_area())

        self.outflow.set_xmoms(new_outflow_xmom)
        self.outflow.set_ymoms(new_outflow_ymom)

    def set_culvert_height(self, height):

        self.culvert_height = height

    def set_culvert_width(self, width):

        self.culvert_width = width

    def set_culvert_z1(self, z1):

        self.culvert_z1 = z1

    def set_culvert_z2(self, z2):

        self.culvert_z2 = z2

    def set_culvert_blockage(self, blockage):

        self.culvert_blockage = blockage

    def set_culvert_barrels(self, barrels):

        self.culvert_barrels = barrels

    def __process_non_skew_culvert(self):

        """Create lines at the end of a culvert inlet and outlet.
        At either end two lines will be created; one for the actual flow to pass through and one a little further away
        for enquiring the total energy at both ends of the culvert and transferring flow.
        """

        self.culvert_vector = self.end_points[1] - self.end_points[0]
        self.culvert_length = math.sqrt(num.sum(self.culvert_vector**2))
        assert self.culvert_length > 0.0, 'The length of culvert is less than 0'

        self.culvert_vector /= self.culvert_length
        self.outward_vector_0 =   self.culvert_vector
        self.outward_vector_1 = - self.culvert_vector

        culvert_normal = num.array([-self.culvert_vector[1], self.culvert_vector[0]])  # Normal vector
        w = 0.5*self.width*culvert_normal # Perpendicular vector of 1/2 width

        self.exchange_lines = []

        # Build exchange polyline and enquiry point
        if self.enquiry_points is None:

            gap = (self.apron + self.enquiry_gap)*self.culvert_vector
            self.enquiry_points = []

            for i in [0, 1]:
                p0 = self.end_points[i] + w
                p1 = self.end_points[i] - w
                self.exchange_lines.append(num.array([p0, p1]))
                ep = self.end_points[i] + (2*i - 1)*gap #(2*i - 1) determines the sign of the points
                self.enquiry_points.append(ep)

        else:
            for i in [0, 1]:
                p0 = self.end_points[i] + w
                p1 = self.end_points[i] - w
                self.exchange_lines.append(num.array([p0, p1]))

    def __process_skew_culvert(self):

        """Compute skew culvert.
        If exchange lines are given, the enquiry points are determined. This is for enquiring
        the total energy at both ends of the culvert and transferring flow.
        """

        centre_point0 = 0.5*(self.exchange_lines[0][0] + self.exchange_lines[0][1])
        centre_point1 = 0.5*(self.exchange_lines[1][0] + self.exchange_lines[1][1])

        n_exchange_0 = len(self.exchange_lines[0])
        n_exchange_1 = len(self.exchange_lines[1])

        assert n_exchange_0 == n_exchange_1, 'There should be the same number of points in both exchange_lines'

        if n_exchange_0 == 2:

            if self.end_points is None:
                self.culvert_vector = centre_point1 - centre_point0
            else:
                self.culvert_vector = self.end_points[1] - self.end_points[0]

            self.outward_vector_0 =   self.culvert_vector
            self.outward_vector_1 = - self.culvert_vector

        elif n_exchange_0 == 4:

            self.outward_vector_0 = self.exchange_lines[0][3] - self.exchange_lines[0][2]
            self.outward_vector_1 = self.exchange_lines[1][3] - self.exchange_lines[1][2]

            self.culvert_vector = centre_point1 - centre_point0

        else:
            raise Exception('n_exchange_0 != 2 or 4')

        self.culvert_length = math.sqrt(num.sum(self.culvert_vector**2))
        assert self.culvert_length > 0.0, 'The length of culvert is less than 0'
        self.culvert_vector /= self.culvert_length

        outward_vector_0_length = math.sqrt(num.sum(self.outward_vector_0**2))
        assert outward_vector_0_length > 0.0, 'The length of outlet_vector_0 is less than 0'
        self.outward_vector_0 /= outward_vector_0_length

        outward_vector_1_length = math.sqrt(num.sum(self.outward_vector_1**2))
        assert outward_vector_1_length > 0.0, 'The length of outlet_vector_1 is less than 0'
        self.outward_vector_1 /= outward_vector_1_length

        if self.enquiry_points is None:

            gap = (self.apron + self.enquiry_gap)*self.culvert_vector

            self.enquiry_points = []

            self.enquiry_points.append(centre_point0 - gap)
            self.enquiry_points.append(centre_point1 + gap)

    def discharge_routine(self):

        msg = 'Need to implement '
        raise

    def statistics(self):

        message  = '=====================================\n'
        message += 'Structure Operator: %s\n' % self.label
        message += '=====================================\n'

        message += 'Structure Type: %s\n' % self.structure_type

        message += 'Description\n'
        message += '%s' % self.description

        #add the culvert dimensions, blockage factor here
        if self.structure_type == 'boyd_pipe':
            message += 'Culvert Diameter: %s\n'% self.diameter
            message += 'Culvert Blockage: %s\n'% self.blockage
            message += 'No.  of  barrels: %s\n'% self.barrels
        elif self.structure_type == 'boyd_box':
            message += 'Culvert  Height: %s\n'% self.height
            message += 'Culvert    Width: %s\n'% self.width
            message += 'Culvert Blockage: %s\n'% self.blockage
            message += 'No.  of  barrels: %s\n'% self.barrels
        else:
            message += 'Culvert Height: %s\n'% self.height
            message += 'Culvert  Width: %s\n'% self.width
            message += 'Batter Slope 1: %s\n'% self.z1
            message += 'Batter Slope 2: %s\n'% self.z2
            message += 'Culvert Blockage: %s\n'% self.blockage
            message += 'No.  of  barrels: %s\n'% self.barrels

        message += '\n'

        for i, inlet in enumerate(self.inlets):
            message += '-------------------------------------\n'
            message +=  'Inlet %i\n' % i
            message += '-------------------------------------\n'

            message += 'inlet triangle indices and centres and elevations\n'
            message += '%s' % inlet.triangle_indices
            message += '\n'

            message += '%s' % self.domain.get_centroid_coordinates()[inlet.triangle_indices]
            message += '\n'

            elev = self.domain.quantities['elevation'].centroid_values[inlet.triangle_indices]
            message += '%s' % elev
            message += '\n'

            elevation_range = elev.max() - elev.min()
            if not num.allclose(elevation_range, 0.):
                message += 'Warning: non-constant inlet elevation can cause well-balancing problems'

            message += 'region\n'
            message += '%s' % inlet.region
            message += '\n'

        message += '=====================================\n'

        return message

    def print_statistics(self):

        print(self.statistics())

    def print_timestepping_statistics(self):

        message = '---------------------------\n'
        message += 'Structure report for %s:\n' % self.label
        message += '--------------------------\n'
        message += 'Type: %s\n' % self.structure_type

        message += 'inlets[0]_enquiry_depth [m]:  %.2f\n' %self.inlets[0].get_enquiry_depth()
        message += 'inlets[0]_enquiry_speed [m/s]:  %.2f\n' %self.inlets[0].get_enquiry_speed()
        message += 'inlets[0]_enquiry_stage [m]:  %.2f\n' %self.inlets[0].get_enquiry_stage()
        message += 'inlets[0]_enquiry_elevation [m]:  %.2f\n' %self.inlets[0].get_enquiry_elevation()
        message += 'inlets[0]_average_depth [m]:  %.2f\n' %self.inlets[0].get_average_depth()
        message += 'inlets[0]_average_speed [m/s]:  %.2f\n' %self.inlets[0].get_average_speed()
        message += 'inlets[0]_average_stage [m]:  %.2f\n' %self.inlets[0].get_average_stage()
        message += 'inlets[0]_average_elevation [m]:  %.2f\n' %self.inlets[0].get_average_elevation()

        message += '\n'

        message += 'inlets[1]_enquiry_depth [m]:  %.2f\n' %self.inlets[1].get_enquiry_depth()
        message += 'inlets[1]_enquiry_speed [m/s]:  %.2f\n' %self.inlets[1].get_enquiry_speed()
        message += 'inlets[1]_enquiry_stage [m]:  %.2f\n' %self.inlets[1].get_enquiry_stage()
        message += 'inlets[1]_enquiry_elevation [m]:  %.2f\n' %self.inlets[1].get_enquiry_elevation()

        message += 'inlets[1]_average_depth [m]:  %.2f\n' %self.inlets[1].get_average_depth()
        message += 'inlets[1]_average_speed [m/s]:  %.2f\n' %self.inlets[1].get_average_speed()
        message += 'inlets[1]_average_stage [m]:  %.2f\n' %self.inlets[1].get_average_stage()
        message += 'inlets[1]_average_elevation [m]:  %.2f\n' %self.inlets[1].get_average_elevation()

        message += 'Discharge [m^3/s]: %.2f\n' % self.discharge
        message += 'Discharge_function_value [m^3/s]: %.2f\n' % self.discharge_abs_timemean
        message += 'Velocity  [m/s]: %.2f\n' % self.velocity
        message += 'Outlet Depth  [m]: %.2f\n' % self.outlet_depth
        message += 'Accumulated Flow [m^3]: %.2f\n' % self.accumulated_flow
        message += 'Inlet Driving Energy %.2f\n' % self.driving_energy
        message += 'Delta Total Energy %.2f\n' % self.delta_total_energy
        message += 'Control at this instant: %s\n' % self.case

        print(message)

    def set_logging(self, flag=True):

        self.logging = flag

        # If flag is true open file with mode = "w" to form a clean file for logging
        if self.logging:
            self.log_filename = self.domain.get_datadir() + '/' + self.label + '.log'
            log_to_file(self.log_filename, self.statistics(), mode='w')
            log_to_file(self.log_filename, 'time, discharge_instantaneous, discharge_abs_timemean, velocity, accumulated_flow, driving_energy_instantaneous, delta_total_energy_instantaneous')

            #log_to_file(self.log_filename, self.culvert_type)

    def timestepping_statistics(self):

        message  = '%.5f, ' % self.domain.get_time()
        message += '%.5f, ' % self.discharge
        message += '%.5f, ' % self.discharge_abs_timemean
        message += '%.5f, ' % self.velocity
        message += '%.5f, ' % self.accumulated_flow
        message += '%.5f, ' % self.driving_energy
        message += '%.5f' % self.delta_total_energy

        # Reset discharge_abs_timemean since last time this function was called
        # (FIXME: This assumes that the function is called only just after a
        # yield step)
        self.discharge_abs_timemean = 0.

        return message

    def get_inlets(self):

        return self.inlets

    def get_culvert_length(self):

        return self.culvert_length

    def get_culvert_slope(self):

        inlet0 = self.inlets[0]
        inlet1 = self.inlets[1]

        elev0 = inlet0.get_enquiry_invert_elevation()
        elev1 = inlet1.get_enquiry_invert_elevation()

        return old_div((elev1-elev0),self.get_culvert_length())

    def get_culvert_width(self):

        return self.width

    def get_culvert_diameter(self):

            return self.diameter

    def get_culvert_height(self):

        return self.height

    def get_culvert_z1(self):

        return self.z1

    def get_culvert_z2(self):

        return self.z2

    def get_culvert_blockage(self):

        return self.blockage

    def get_culvert_barrels(self):

        return self.barrels

    def get_culvert_apron(self):

        return self.apron

    def get_master_proc(self):

        return 0

    #--------------------------------------------------------
    # Set of enquiry functions so that in the sequential and paralle case
    # we can get equiry info fron the master Proc
    #---------------------------------------------------------

    def get_enquiry_stages(self):

        enq0 = self.inlets[0].get_enquiry_stage()
        enq1 = self.inlets[1].get_enquiry_stage()

        return [enq0, enq1]

    def get_enquiry_depths(self):

        enq0 = self.inlets[0].get_enquiry_depth()
        enq1 = self.inlets[1].get_enquiry_depth()

        return [enq0, enq1]

    def get_enquiry_positions(self):

        enq0 = self.inlets[0].get_enquiry_position()
        enq1 = self.inlets[1].get_enquiry_position()

        return [enq0, enq1]

    def get_enquiry_xmoms(self):

        enq0 = self.inlets[0].get_enquiry_xmom()
        enq1 = self.inlets[1].get_enquiry_xmom()

        return [enq0, enq1]

    def get_enquiry_ymoms(self):

        enq0 = self.inlets[0].get_enquiry_ymom()
        enq1 = self.inlets[1].get_enquiry_ymom()

        return [enq0, enq1]

    def get_enquiry_elevations(self):

        enq0 = self.inlets[0].get_enquiry_elevation()
        enq1 = self.inlets[1].get_enquiry_elevation()

        return [enq0, enq1]

    def get_enquiry_water_depths(self):

        enq0 = self.inlets[0].get_enquiry_water_depth()
        enq1 = self.inlets[1].get_enquiry_water_depth()

        return [enq0, enq1]

    def get_enquiry_invert_elevations(self):

        enq0 = self.inlets[0].get_enquiry_invert_elevation()
        enq1 = self.inlets[1].get_enquiry_invert_elevation()

        return [enq0, enq1]

    def get_enquiry_velocitys(self):

        enq0 = self.inlets[0].get_enquiry_velocity()
        enq1 = self.inlets[1].get_enquiry_velocity()

        return [enq0, enq1]

    def get_enquiry_xvelocitys(self):

        enq0 = self.inlets[0].get_enquiry_xvelocity()
        enq1 = self.inlets[1].get_enquiry_xvelocity()

        return [enq0, enq1]

    def get_enquiry_yvelocitys(self):

        enq0 = self.inlets[0].get_enquiry_yvelocity()
        enq1 = self.inlets[1].get_enquiry_yvelocity()

        return [enq0, enq1]

    def get_enquiry_speeds(self):

        enq0 = self.inlets[0].get_enquiry_speed()
        enq1 = self.inlets[1].get_enquiry_speed()

        return [enq0, enq1]

    def get_enquiry_velocity_heads(self):

        enq0 = self.inlets[0].get_enquiry_velocity_head()
        enq1 = self.inlets[1].get_enquiry_velocity_head()

        return [enq0, enq1]

    def get_enquiry_total_energys(self):

        enq0 = self.inlets[0].get_enquiry_total_energy()
        enq1 = self.inlets[1].get_enquiry_total_energy()

        return [enq0, enq1]

    def get_enquiry_specific_energys(self):

        enq0 = self.inlets[0].get_enquiry_specific_energy()
        enq1 = self.inlets[1].get_enquiry_specific_energy()

        return [enq0, enq1]


"""
Contains RDKit-related functions.
"""

# for handling ligand data and calculating ligand-related properties
from rdkit import Chem
from rdkit.Chem import Draw, AllChem, Descriptors
import numpy as np  # for some more functionalities when using Pandas (e.g. for handling NaN values)

def create_molecule_object(input_type, input_value):
    """
    This class is used to create an RDKit molecule object from various sources.
    It can be used to calculate or retrieve properties and descriptors (MW, h-bonds, etc.),
    generate images or SDF files or to perform similarity searches.

    Parameters
    ----------
    input_type : str
        Type of the input.
        Allowed input-types are: 'smiles', 'inchi', 'smarts', 'pdb_file'
    input_value : str
        Value of the corresponding input type.

    Returns
    -------
    rdkit.Chem.rdchem.Mol
        Structure as RDKit molecule object.
    """

    functions = {
        "smiles": Chem.MolFromSmiles,
        "inchi": Chem.MolFromInchi,
        "smarts": Chem.MolFromSmarts,
        "pdb_file": Chem.MolFromPDBFile,
    }
    Molobj = functions[input_type](input_value)
    Molobj.smiles = Chem.MolToSmiles(Molobj)
    return Molobj

def draw_molecules(
    list_mol_objs,
    list_legends=None,
    mols_per_row=3,
    sub_img_size=(350, 350),
    filepath=None,
):
    """
    Take a list of RDKit molecule objects and draws them as a grid image.

    Parameters
    ----------
    list_mol_objs: list
        List of RDKit molecule objects to be drawn.
    list_legends: list
        Optional; default: None
        List of legends for the molecules.
        If not provided, the list indices (+1) will be used as legends.
    mols_per_row : int
        Optional; default: 3
        Number of structures to show per row.
    sub_img_size : tuple (int, int)
        Optional; default: (350, 350)
        Size of each structure.
    filepath : str or pathlib.Path
        Full filepath to save the image in.

    Returns
    -------
    rdkit.Chem.Draw.MolsToGridImage
        Molecules shown as grid.
    """
    if list_legends is None:
        list_legends = list(map(str, range(1, len(list_mol_objs) + 1)))
    figure = Draw.MolsToGridImage(
        list_mol_objs,
        molsPerRow=mols_per_row,
        subImgSize=sub_img_size,
        legends=list(map(str, list_legends)),
    )
    if filepath is not None:
        with open(f"{filepath}.png", "wb") as f:
            f.write(figure.data)
    return figure

def save_molecule_image_to_file(mol_obj, filepath):
    """
    Save the image of a single molecule as a PNG file.

    Parameters
    ----------
    mol_obj : rdkit.Chem.rdchem.Mol
        The molecule to be saved as image.
    filepath : str or pathlib.Path
        Full filpath to save the image in.
    """
    Draw.MolToFile(mol_obj, f"{filepath}.png")

def save_3D_molecule_to_SDfile(mol_obj, filepath):
    """
    Generate a 3D conformer and save as SDF file.

    Parameters
    ----------
    mol_obj : rdkit.Chem.rdchem.Mol
        The molecule to be saved as SDF file.
    filepath : str or pathlib.Path
        Full filpath to save the image in.
    """
    mol = Chem.AddHs(mol_obj)
    embedding = AllChem.EmbedMolecule(mol, maxAttempts=1000, clearConfs=True)
    uffoptim = AllChem.UFFOptimizeMolecule(mol, maxIters=1000)
    # check if calculations converged (both should return 0 when converged)
    if embedding + uffoptim != 0:
        raise ValueError("Embedding/Optimization failed to converge.")
    session = Chem.SDWriter(f"{filepath}.sdf")
    session.write(mol)
    session.close()

def calculate_similarity_dice(mol_obj1, mol_obj2, morgan_radius=2, morgan_nbits=4096):
    """
    Calculate the Dice similarity between two molecules,
    based on 4096-bit Morgan fingerprints with a radius of 2.

    Parameters
    ----------
    mol_obj1 : rdkit.Chem.rdchem.Mol
        The first molecule.
    mol_obj2 : rdkit.Chem.rdchem.Mol
        The second molecule.
    morgan_radius : int
        Optional; default: 2
        The radius used to generate Morgan fingerprints.
    morgan_nbits : int
        Optional; default: 4096
        The number of bits in the Morgan fingerprints.
    Returns
    -------
    float
        Dice similarity between the two molecules rounded to two decimal places.
    """
    morgan_fp_mol1 = AllChem.GetMorganFingerprintAsBitVect(
        mol_obj1, radius=morgan_radius, nBits=morgan_nbits
    )
    morgan_fp_mol2 = AllChem.GetMorganFingerprintAsBitVect(
        mol_obj2, radius=morgan_radius, nBits=morgan_nbits
    )
    dice_similarity = round(
        AllChem.DataStructs.DiceSimilarity(morgan_fp_mol1, morgan_fp_mol2), 2
    )
    return dice_similarity

def calculate_druglikeness(mol_obj):
    """
    Calculate several molecular properties and drug-likeness scores,
    from an RDKit molecule object.

    Parameters
    ----------
    mol_obj: rdkit.Chem.rdchem.Mol
        Molecule object of interest.

    Returns
    -------
    dict
        The calculated values are returned in a dictionary with following keys:
        MolWt, NumHAcceptors, NumHDonors, MolLogP, TPSA, NumRotBonds, Saturation,
        lipinski_score, custom_drug_score, qed_score, total_drug_score
    """
    properties = {
        "mol_weight": round(Descriptors.MolWt(mol_obj), 3),
        "num_H_acceptors": Descriptors.NumHAcceptors(mol_obj),
        "num_H_donors": Descriptors.NumHDonors(mol_obj),
        "logp": round(Descriptors.MolLogP(mol_obj), 2),
        "tpsa": round(Descriptors.TPSA(mol_obj), 2),
        "num_rot_bonds": Descriptors.NumRotatableBonds(mol_obj),
        "saturation": round(Descriptors.FractionCSP3(mol_obj), 2),
        "drug_score_qed": round(Descriptors.qed(mol_obj), 2),
    }

    # Calculating Lipinski score
    l1 = int(properties["mol_weight"] < 500)
    l2 = int(properties["num_H_acceptors"] <= 10)
    l3 = int(properties["num_H_donors"] <= 5)
    l4 = int(properties["logp"] < 5)
    properties["drug_score_lipinski"] = round((l1 + l2 + l3 + l4) / 4, 2)

    # Calculating druglikeness score with custom scoring functions
    # derived from Hopkins paper: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3524573/
    def molWt_score(molWt):
        if molWt <= 440:
            return np.exp(-((molWt - 300) ** 2) / 15000)
        else:
            return np.exp(-(molWt - 180) / 190) + 0.01

    def molLogP_score(molLogP):
        return np.exp(-((molLogP - 2.5) ** 2) / 9)

    def numHDonors_score(numHDonors):
        if numHDonors == 0:
            return 0.6
        elif numHDonors < 5:
            return np.exp(-((numHDonors - 1) ** 2) / 5)
        else:
            return np.exp(-((numHDonors - 1) ** 2) / 5) + (0.4 / numHDonors)

    def numHAcceptors_score(numHAcceptors):
        if numHAcceptors < 4:
            return np.exp(-((numHAcceptors - 3) ** 2) / 3)
        else:
            return np.exp(-0.3 * numHAcceptors / 0.8 + 1.4)

    def TPSA_score(TPSA):
        if TPSA < 50:
            return 0.015 * TPSA + 0.25
        else:
            return np.exp(-((TPSA - 50) ** 2) / 8000)

    def numRotBonds_score(numRotBonds):
        if numRotBonds < 10:
            return np.exp(-((numRotBonds - 4) ** 2) / 19)
        else:
            return np.exp(-((numRotBonds - 4) ** 2) / 19) + (1.5 / numRotBonds ** 1.5)

    def saturation_score(saturation):
        return np.exp(-((saturation - 0.625) ** 2) / 0.05)

    d1 = molWt_score(properties["mol_weight"])
    d2 = numHAcceptors_score(properties["num_H_acceptors"])
    d3 = numHDonors_score(properties["num_H_donors"])
    d4 = molLogP_score(properties["logp"])
    d5 = TPSA_score(properties["tpsa"])
    d6 = numRotBonds_score(properties["num_rot_bonds"])
    d7 = saturation_score(properties["saturation"])
    properties["drug_score_custom"] = round((d1 + d2 + d3 + d4 + d5 + d6 + d7) / 7, 2)

    properties["drug_score_total"] = round(
        (
            3 * properties["drug_score_qed"]
            + 2 * properties["drug_score_custom"]
            + properties["drug_score_lipinski"]
        )
        / 6,
        2,
    )
    return properties


import random

import numpy as np

from .meta_gf import GFMeta
from .linalg import dot, inner, outer, matrix_rank, solve, inv, det, row_reduce, lu_decompose, lup_decompose
from .overrides import set_module
from .poly_conversion import integer_to_poly, poly_to_str, str_to_integer

__all__ = ["GFArray"]

UNSUPPORTED_ONE_ARG_FUNCTIONS = [
    np.packbits, np.unpackbits,
    np.unwrap,
    np.around, np.round_, np.fix,
    np.gradient, np.trapz,
    np.i0, np.sinc,
    np.angle, np.real, np.imag, np.conj, np.conjugate,
]

UNSUPPORTED_TWO_ARG_FUNCTIONS = [
    np.lib.scimath.logn,
    np.cross,
]

UNSUPPORTED_FUNCTIONS = UNSUPPORTED_ONE_ARG_FUNCTIONS + UNSUPPORTED_TWO_ARG_FUNCTIONS

OVERRIDDEN_FUNCTIONS = {
    np.dot: dot,
    np.inner: inner,
    np.outer: outer,
    # np.tensordot: "tensordot",
    np.linalg.matrix_rank: matrix_rank,
    np.linalg.inv: inv,
    np.linalg.det: det,
    np.linalg.solve: solve
}

FUNCTIONS_REQUIRING_VIEW = [
    np.copy, np.concatenate,
    np.broadcast_to,
    np.trace,
]

UNSUPPORTED_ONE_ARG_UFUNCS = [
    np.invert, np.sqrt,
    np.log2, np.log10,
    np.exp, np.expm1, np.exp2,
    np.sin, np.cos, np.tan,
    np.sinh, np.cosh, np.tanh,
    np.arcsin, np.arccos, np.arctan,
    np.arcsinh, np.arccosh, np.arctanh,
    np.degrees, np.radians,
    np.deg2rad, np.rad2deg,
    np.floor, np.ceil, np.trunc, np.rint,
]

UNSUPPORTED_TWO_ARG_UFUNCS = [
    np.hypot, np.arctan2,
    np.logaddexp, np.logaddexp2,
    np.remainder,
    np.fmod, np.modf,
    np.fmin, np.fmax,
]

UNSUPPORTED_UFUNCS = UNSUPPORTED_ONE_ARG_UFUNCS + UNSUPPORTED_TWO_ARG_UFUNCS

OVERRIDDEN_UFUNCS = {
    np.add: "_ufunc_add",
    np.subtract: "_ufunc_subtract",
    np.multiply: "_ufunc_multiply",
    np.floor_divide: "_ufunc_divide",
    np.true_divide: "_ufunc_divide",
    np.negative: "_ufunc_negative",
    np.reciprocal: "_ufunc_reciprocal",
    np.power: "_ufunc_power",
    np.square: "_ufunc_square",
    np.log: "_ufunc_log",
    np.matmul: "_ufunc_matmul",
}

UFUNCS_REQUIRING_VIEW = [
    np.bitwise_and, np.bitwise_or, np.bitwise_xor,
    np.left_shift, np.right_shift,
]

@set_module("galois")
class GFArray(np.ndarray, metaclass=GFMeta):
    """
    Create an array over :math:`\\mathrm{GF}(p^m)`.

    The :obj:`galois.GFArray` class is a parent class for all Galois field array classes. Any Galois field :math:`\\mathrm{GF}(p^m)`
    with prime characteristic :math:`p` and positive integer :math:`m`, can be constructed by calling the class factory
    `galois.GF(p**m)`.

    Warning
    -------
        This is an abstract base class for all Galois field array classes. :obj:`galois.GFArray` cannot be instantiated
        directly. Instead, Galois field array classes are created using :obj:`galois.GF`.

        For example, one can create the :math:`\\mathrm{GF}(7)` field array class as follows:

        .. ipython:: python

            GF7 = galois.GF(7)
            print(GF7)

        This subclass can then be used to instantiate arrays over :math:`\\mathrm{GF}(7)`.

        .. ipython:: python

            GF7([3,5,0,2,1])
            GF7.Random((2,5))

    :obj:`galois.GFArray` is a subclass of :obj:`numpy.ndarray`. The :obj:`galois.GFArray` constructor has the same syntax as
    :obj:`numpy.array`. The returned :obj:`galois.GFArray` object is an array that can be acted upon like any other
    numpy array.

    Parameters
    ----------
    array : array_like
        The input array to be converted to a Galois field array. The input array is copied, so the original array
        is unmodified by changes to the Galois field array. Valid input array types are :obj:`numpy.ndarray`,
        :obj:`list` or :obj:`tuple` of ints or strs, :obj:`int`, or :obj:`str`.
    dtype : numpy.dtype, optional
        The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
        dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.
    copy : bool, optional
        The `copy` keyword argument from :obj:`numpy.array`. The default is `True` which makes a copy of the input
        object is it's an array.
    order : str, optional
        The `order` keyword argument from :obj:`numpy.array`. Valid values are `"K"` (default), `"A"`, `"C"`, or `"F"`.
    ndmin : int, optional
        The `ndmin` keyword argument from :obj:`numpy.array`. The minimum number of dimensions of the output.
        The default is 0.

    Returns
    -------
    galois.GFArray
        The copied input array as a :math:`\\mathrm{GF}(p^m)` field array.

    Examples
    --------
    Construct various kinds of Galois fields using :obj:`galois.GF`.

    .. ipython:: python

        # Construct a GF(2^m) class
        GF256 = galois.GF(2**8); print(GF256)

        # Construct a GF(p) class
        GF571 = galois.GF(571); print(GF571)

        # Construct a very large GF(2^m) class
        GF2m = galois.GF(2**100); print(GF2m)

        # Construct a very large GF(p) class
        GFp = galois.GF(36893488147419103183); print(GFp)

    Depending on the field's order (size), only certain `dtype` values will be supported.

    .. ipython:: python

        GF256.dtypes
        GF571.dtypes

    Very large fields, which can't be represented using `np.int64`, can only be represented as `dtype=np.object_`.

    .. ipython:: python

        GF2m.dtypes
        GFp.dtypes

    Newly-created arrays will use the smallest, valid dtype.

    .. ipython:: python

        a = GF256.Random(10); a
        a.dtype

    This can be explicitly set by specifying the `dtype` keyword argument.

    .. ipython:: python

        a = GF256.Random(10, dtype=np.uint32); a
        a.dtype

    Arrays can also be created explicitly by converting an "array-like" object.

    .. ipython:: python

        # Construct a Galois field array from a list
        l = [142, 27, 92, 253, 103]; l
        GF256(l)

        # Construct a Galois field array from an existing numpy array
        x_np = np.array(l, dtype=np.int64); x_np
        GF256(l)

    Arrays can also be created by "view casting" from an existing numpy array. This avoids
    a copy operation, which is especially useful for large data already brought into memory.

    .. ipython:: python

        a = x_np.view(GF256); a

        # Changing `x_np` will change `a`
        x_np[0] = 0; x_np
        a
    """

    def __new__(cls, array, dtype=None, copy=True, order="K", ndmin=0):
        if cls is GFArray:
            raise NotImplementedError("GFArray is an abstract base class that cannot be directly instantiated. Instead, create a GFArray subclass using `galois.GF`.")
        return cls._array(array, dtype=dtype, copy=copy, order=order, ndmin=ndmin)

    @classmethod
    def _get_dtype(cls, dtype):
        if dtype is None:
            return cls.dtypes[0]

        # Convert "dtype" to a numpy dtype. This does platform specific conversion, if necessary.
        # For example, np.dtype(int) == np.int64 (on some systems).
        dtype = np.dtype(dtype)
        if dtype not in cls.dtypes:
            raise TypeError(f"{cls.name} arrays only support dtypes {[np.dtype(d).name for d in cls.dtypes]}, not '{dtype.name}'.")

        return dtype

    @classmethod
    def _array(cls, array_like, dtype=None, copy=True, order="K", ndmin=0):
        dtype = cls._get_dtype(dtype)
        array_like = cls._check_array_like_object(array_like)
        array = np.array(array_like, dtype=dtype, copy=copy, order=order, ndmin=ndmin)
        return array.view(cls)

    @classmethod
    def _check_array_like_object(cls, array_like):
        if isinstance(array_like, str):
            # Convert the string to an integer
            array_like = str_to_integer(array_like, cls.prime_subfield)

        if isinstance(array_like, (int, np.integer)):
            # Just check that the single int is in range
            cls._check_array_values(array_like)

        elif isinstance(array_like, (list, tuple)):
            # Recursively check the items in the iterable to ensure they're of the correct type
            # and that their values are in range
            array_like = cls._check_iterable_types_and_values(array_like)

        elif isinstance(array_like, np.ndarray):
            if array_like.dtype == np.object_:
                array_like = cls._check_array_types_dtype_object(array_like)
            elif not np.issubdtype(array_like.dtype, np.integer):
                raise TypeError(f"{cls.name} arrays must have integer dtypes, not {array_like.dtype}.")
            cls._check_array_values(array_like)

        else:
            raise TypeError(f"{cls.name} arrays can be created with scalars of type int, not {type(array_like)}.")

        return array_like

    @classmethod
    def _check_iterable_types_and_values(cls, iterable):
        new_iterable = []
        for item in iterable:
            if isinstance(item, (list, tuple)):
                item = cls._check_iterable_types_and_values(item)
                new_iterable.append(item)
                continue

            if isinstance(item, str):
                item = str_to_integer(item, cls.prime_subfield)
            elif not isinstance(item, (int, np.integer, cls)):
                raise TypeError(f"When {cls.name} arrays are created/assigned with an iterable, each element must be an integer. Found type {type(item)}.")

            if not 0 <= item < cls.order:
                raise ValueError(f"{cls.name} arrays must have elements in 0 <= x < {cls.order}, not {item}.")

            # Ensure the type is int so dtype=object classes don't get all mixed up
            new_iterable.append(int(item))

        return new_iterable

    @classmethod
    def _check_array_types_dtype_object(cls, array):
        if array.size == 0:
            return array
        if array.ndim == 0:
            if not isinstance(array[()], (int, np.integer, cls)):
                raise TypeError(f"When {cls.name} arrays are created/assigned with a numpy array with dtype=object, each element must be an integer. Found type {type(array[()])}.")
            return int(array)

        iterator = np.nditer(array, flags=["multi_index", "refs_ok"])
        for _ in iterator:
            a = array[iterator.multi_index]
            if not isinstance(a, (int, np.integer, cls)):
                raise TypeError(f"When {cls.name} arrays are created/assigned with a numpy array with dtype=object, each element must be an integer. Found type {type(a)}.")

            # Ensure the type is int so dtype=object classes don't get all mixed up
            array[iterator.multi_index] = int(a)

        return array

    @classmethod
    def _check_array_values(cls, array):
        if not isinstance(array, np.ndarray):
            # Convert single integer to array so next step doesn't fail
            array = np.array(array)

        # Check the value of the "field elements" and make sure they are valid
        if np.any(array < 0) or np.any(array >= cls.order):
            idxs = np.logical_or(array < 0, array >= cls.order)
            raise ValueError(f"{cls.name} arrays must have elements in 0 <= x < {cls.order}, not {array[idxs]}.")

    ###############################################################################
    # Alternate constructors
    ###############################################################################

    @classmethod
    def Zeros(cls, shape, dtype=None):
        """
        Creates a Galois field array with all zeros.

        Parameters
        ----------
        shape : tuple
            A numpy-compliant `shape` tuple, see :obj:`numpy.ndarray.shape`. An empty tuple `()` represents a scalar.
            A single integer or 1-tuple, e.g. `N` or `(N,)`, represents the size of a 1-dim array. An n-tuple, e.g.
            `(M,N)`, represents an n-dim array with each element indicating the size in each dimension.
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            A Galois field array of zeros.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(31)
            GF.Zeros((2,5))
        """
        dtype = cls._get_dtype(dtype)
        array = np.zeros(shape, dtype=dtype)
        return array.view(cls)

    @classmethod
    def Ones(cls, shape, dtype=None):
        """
        Creates a Galois field array with all ones.

        Parameters
        ----------
        shape : tuple
            A numpy-compliant `shape` tuple, see :obj:`numpy.ndarray.shape`. An empty tuple `()` represents a scalar.
            A single integer or 1-tuple, e.g. `N` or `(N,)`, represents the size of a 1-dim array. An n-tuple, e.g.
            `(M,N)`, represents an n-dim array with each element indicating the size in each dimension.
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            A Galois field array of ones.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(31)
            GF.Ones((2,5))
        """
        dtype = cls._get_dtype(dtype)
        array = np.ones(shape, dtype=dtype)
        return array.view(cls)

    @classmethod
    def Identity(cls, size, dtype=None):
        """
        Creates an :math:`n \\times n` identity matrix over :math:`\\mathrm{GF}(q)`.

        Parameters
        ----------
        size : int
            The size :math:`n` along one axis of the matrix. The resulting array has shape `(size,size)`.
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            A Galois field identity matrix of shape `(size, size)`.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(31)
            GF.Identity(4)
        """
        dtype = cls._get_dtype(dtype)
        array = np.identity(size, dtype=dtype)
        return array.view(cls)

    @classmethod
    def Vandermonde(cls, a, m, n, dtype=None):
        """
        Creates a :math:`m \\times n` Vandermonde matrix of :math:`a \\in \\mathrm{GF}(q)`.

        Parameters
        ----------
        a : int, galois.GFArray
            An element of :math:`\\mathrm{GF}(q)`.
        m : int
            The number of rows in the Vandermonde matrix.
        n : int
            The number of columns in the Vandermonde matrix.
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            The :math:`m \\times n` Vandermonde matrix.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(2**3)
            a = GF.primitive_element
            V = GF.Vandermonde(a, 7, 7)
            with GF.display("power"):
                print(V)
        """
        if not isinstance(a, (int, np.integer,cls)):
            raise TypeError(f"Argument `a` must be an integer or element of {cls.name}, not {type(a)}.")
        if not isinstance(m, (int, np.integer)):
            raise TypeError(f"Argument `m` must be an integer, not {type(m)}.")
        if not isinstance(n, (int, np.integer)):
            raise TypeError(f"Argument `n` must be an integer, not {type(n)}.")
        if not m > 0:
            raise ValueError(f"Argument `m` must be non-negative, not {m}.")
        if not n > 0:
            raise ValueError(f"Argument `n` must be non-negative, not {n}.")

        dtype = cls._get_dtype(dtype)
        a = cls(a, dtype=dtype)
        if not a.ndim == 0:
            raise ValueError(f"Argument `a` must be a scalar, not {a.ndim}-D.")

        v = a ** np.arange(0, m)
        V = np.power.outer(v, np.arange(0, n))

        return V

    @classmethod
    def Range(cls, start, stop, step=1, dtype=None):
        """
        Creates a Galois field array with a range of field elements.

        Parameters
        ----------
        start : int
            The starting value (inclusive).
        stop : int
            The stopping value (exclusive).
        step : int, optional
            The space between values. The default is 1.
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            A Galois field array of a range of field elements.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(31)
            GF.Range(10,20)
        """
        dtype = cls._get_dtype(dtype)
        if not stop <= cls.order:
            raise ValueError(f"The stopping value must be less than the field order of {cls.order}, not {stop}.")

        if dtype != np.object_:
            array = np.arange(start, stop, step=step, dtype=dtype)
        else:
            array = np.array(range(start, stop, step), dtype=dtype)

        return array.view(cls)

    @classmethod
    def Random(cls, shape=(), low=0, high=None, dtype=None):
        """
        Creates a Galois field array with random field elements.

        Parameters
        ----------
        shape : tuple
            A numpy-compliant `shape` tuple, see :obj:`numpy.ndarray.shape`. An empty tuple `()` represents a scalar.
            A single integer or 1-tuple, e.g. `N` or `(N,)`, represents the size of a 1-dim array. An n-tuple, e.g.
            `(M,N)`, represents an n-dim array with each element indicating the size in each dimension.
        low : int, optional
            The lowest value (inclusive) of a random field element. The default is 0.
        high : int, optional
            The highest value (exclusive) of a random field element. The default is `None` which represents the
            field's order :math:`p^m`.
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            A Galois field array of random field elements.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(31)
            GF.Random((2,5))
        """
        dtype = cls._get_dtype(dtype)
        high = cls.order if high is None else high
        if not 0 <= low < high <= cls.order:
            raise ValueError(f"Arguments must satisfy `0 <= low < high <= order`, not `0 <= {low} < {high} <= {cls.order}`.")

        if dtype != np.object_:
            array = np.random.randint(low, high, shape, dtype=dtype)
        else:
            array = np.empty(shape, dtype=dtype)
            iterator = np.nditer(array, flags=["multi_index", "refs_ok"])
            for _ in iterator:
                array[iterator.multi_index] = random.randint(low, high - 1)

        return array.view(cls)

    @classmethod
    def Elements(cls, dtype=None):
        """
        Creates a Galois field array of the field's elements :math:`\\{0, \\dots, p^m-1\\}`.

        Parameters
        ----------
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            A Galois field array of all the field's elements.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(31)
            GF.Elements()
        """
        return cls.Range(0, cls.order, step=1, dtype=dtype)

    @classmethod
    def Vector(cls, array, dtype=None):
        """
        Creates a Galois field array over :math:`\\mathrm{GF}(p^m)` from length-:math:`m` vectors over the prime subfield :math:`\\mathrm{GF}(p)`.

        Parameters
        ----------
        array : array_like
            The input array with field elements in :math:`\\mathrm{GF}(p)` to be converted to a Galois field array in :math:`\\mathrm{GF}(p^m)`.
            The last dimension of the input array must be :math:`m`. An input array with shape `(n1, n2, m)` has output shape `(n1, n2)`.
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            A Galois field array over :math:`\\mathrm{GF}(p^m)`.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(2**6)
            vec = galois.GF2.Random((3,6)); vec
            a = GF.Vector(vec); a
            with GF.display("poly"):
                print(a)
            a.vector()
        """
        order = cls.prime_subfield.order
        degree = cls.degree
        array = cls.prime_subfield(array).view(np.ndarray).astype(cls.dtypes[-1])  # Use the largest dtype so computation doesn't overflow
        if not array.shape[-1] == degree:
            raise ValueError(f"The last dimension of `array` must be the field extension dimension {cls.degree}, not {array.shape[-1]}.")
        degrees = np.arange(degree - 1, -1, -1, dtype=cls.dtypes[-1])
        array = np.sum(array * order**degrees, axis=-1)
        return cls(array, dtype=dtype)

    ###############################################################################
    # Array methods
    ###############################################################################

    def vector(self, dtype=None):
        """
        Converts the Galois field array over :math:`\\mathrm{GF}(p^m)` to length-:math:`m` vectors over the prime subfield :math:`\\mathrm{GF}(p)`.

        For an input array with shape `(n1, n2)`, the output shape is `(n1, n2, m)`.

        Parameters
        ----------
        dtype : numpy.dtype, optional
            The :obj:`numpy.dtype` of the array elements. The default is `None` which represents the smallest valid
            dtype for this class, i.e. the first element in :obj:`galois.GFMeta.dtypes`.

        Returns
        -------
        galois.GFArray
            A Galois field array of length-:math:`m` vectors over :math:`\\mathrm{GF}(p)`.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(2**6)
            a = GF.Random(3); a
            vec = a.vector(); vec
            GF.Vector(vec)
        """
        order = type(self).prime_subfield.order
        degree = type(self).degree
        array = self.view(np.ndarray)
        array = np.repeat(array, degree).reshape(*array.shape, degree)
        x = 0
        for i in range(degree):
            q = (array[...,i] - x) // order**(degree - 1 - i)
            array[...,i] = q
            x += q*order**(degree - 1 - i)
        return type(self).prime_subfield(array, dtype=dtype)

    def row_reduce(self, ncols=None):
        """
        Performs Gaussian elimination on the matrix to achieve reduced row echelon form.

        **Row reduction operations**

        1. Swap the position of any two rows.
        2. Multiply a row by a non-zero scalar.
        3. Add one row to a scalar multiple of another row.

        Parameters
        ----------
        ncols : int, optional
            The number of columns to perform Gaussian elimination over. The default is `None` which represents
            the number of columns of the input array.

        Returns
        -------
        galois.GFArray
            The reduced row echelon form of the input array.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(31)
            A = GF.Random((4,4)); A
            A.row_reduce()
            np.linalg.matrix_rank(A)

        One column is a linear combination of another.

        .. ipython:: python

            GF = galois.GF(31)
            A = GF.Random((4,4)); A
            A[:,2] = A[:,1] * GF(17); A
            A.row_reduce()
            np.linalg.matrix_rank(A)

        One row is a linear combination of another.

        .. ipython:: python

            GF = galois.GF(31)
            A = GF.Random((4,4)); A
            A[3,:] = A[2,:] * GF(8); A
            A.row_reduce()
            np.linalg.matrix_rank(A)
        """
        return row_reduce(self, ncols=ncols)

    def lu_decompose(self):
        """
        Decomposes the input array into the product of lower and upper triangular matrices.

        Returns
        -------
        galois.GFArray
            The lower triangular matrix.
        galois.GFArray
            The upper triangular matrix.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(5)

            # Not every square matrix has an LU decomposition
            A = GF([[2, 4, 4, 1], [3, 3, 1, 4], [4, 3, 4, 2], [4, 4, 3, 1]])
            L, U = A.lu_decompose()
            L
            U

            # A = L U
            np.array_equal(A, L @ U)
        """
        return lu_decompose(self)

    def lup_decompose(self):
        """
        Decomposes the input array into the product of lower and upper triangular matrices using partial pivoting.

        Returns
        -------
        galois.GFArray
            The lower triangular matrix.
        galois.GFArray
            The upper triangular matrix.
        galois.GFArray
            The permutation matrix.

        Examples
        --------
        .. ipython:: python

            GF = galois.GF(5)
            A = GF([[1, 3, 2, 0], [3, 4, 2, 3], [0, 2, 1, 4], [4, 3, 3, 1]])
            L, U, P = A.lup_decompose()
            L
            U
            P

            # P A = L U
            np.array_equal(P @ A, L @ U)
        """
        return lup_decompose(self)

    ###############################################################################
    # Overridden numpy methods
    ###############################################################################

    def astype(self, dtype, **kwargs):  # pylint: disable=arguments-differ
        if dtype not in type(self).dtypes:
            raise TypeError(f"{type(self).name} arrays can only be cast as integer dtypes in {type(self).dtypes}, not {dtype}.")
        return super().astype(dtype, **kwargs)

    def __array_finalize__(self, obj):
        """
        A numpy dunder method that is called after "new", "view", or "new from template". It is used here to ensure
        that view casting to a Galois field array has the appropriate dtype and that the values are in the field.
        """
        if obj is not None and not isinstance(obj, GFArray):
            # Only invoked on view casting
            if obj.dtype not in type(self).dtypes:
                raise TypeError(f"{type(self).name} can only have integer dtypes {type(self).dtypes}, not {obj.dtype}.")
            if np.any(obj < 0) or np.any(obj >= type(self).order):
                idxs = np.logical_or(obj < 0, obj >= type(self).order)
                raise ValueError(f"{type(self).name} arrays must have values in 0 <= x < {type(self).order}, not {obj[idxs]}.")

    def __getitem__(self, key):
        item = super().__getitem__(key)
        if np.isscalar(item):
            # Return scalar array elements as 0-dimension Galois field arrays. This enables Galois field arithmetic
            # on scalars, which would otherwise be implemented using standard integer arithmetic.
            item = self.__class__(item, dtype=self.dtype)
        return item

    def __setitem__(self, key, value):
        # Verify the values to be written to the Galois field array are in the field
        value = self._check_array_like_object(value)
        super().__setitem__(key, value)

    def __array_function__(self, func, types, args, kwargs):
        if func in OVERRIDDEN_FUNCTIONS:
            output = OVERRIDDEN_FUNCTIONS[func](*args, **kwargs)

        elif func in UNSUPPORTED_FUNCTIONS:
            raise NotImplementedError(f"The numpy function '{func.__name__}' is not supported on Galois field arrays. If you believe this function should be supported, please submit a GitHub issue at https://github.com/mhostetter/galois/issues.\n\nIf you'd like to perform this operation on the data (but not necessarily a Galois field array), you should first call `array = array.view(np.ndarray)` and then call the function.")

        else:
            if func is np.insert:
                args = list(args)
                args[2] = self._check_array_like_object(args[2])
                args = tuple(args)

            output = super().__array_function__(func, types, args, kwargs)  # pylint: disable=no-member

            if func in FUNCTIONS_REQUIRING_VIEW:
                if np.isscalar(output):
                    output = type(self)(output, dtype=self.dtype)
                else:
                    output = output.view(type(self))

        return output

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):  # pylint: disable=too-many-branches
        # print(ufunc, method, inputs, kwargs)
        meta = {}
        meta["types"] = [type(inputs[i]) for i in range(len(inputs))]
        meta["operands"] = list(range(len(inputs)))
        if method in ["at", "reduceat"]:
            # Remove the second argument for "at" ufuncs which is the indices list
            meta["operands"].pop(1)
        meta["field_operands"] = [i for i in meta["operands"] if isinstance(inputs[i], self.__class__)]
        meta["non_field_operands"] = [i for i in meta["operands"] if not isinstance(inputs[i], self.__class__)]
        meta["field"] = self.__class__
        meta["dtype"] = self.dtype
        # meta["ufuncs"] = self._ufuncs

        if ufunc in OVERRIDDEN_UFUNCS:
            # Set all ufuncs with "casting" keyword argument to "unsafe" so we can cast unsigned integers
            # to integers. We know this is safe because we already verified the inputs.
            if method not in ["reduce", "accumulate", "at", "reduceat"]:
                kwargs["casting"] = "unsafe"

            # Need to set the intermediate dtype for reduction operations or an error will be thrown. We
            # use the largest valid dtype for this field.
            if method in ["reduce"]:
                kwargs["dtype"] = type(self).dtypes[-1]

            return getattr(type(self), OVERRIDDEN_UFUNCS[ufunc])(ufunc, method, inputs, kwargs, meta)

        elif ufunc in UNSUPPORTED_UFUNCS:
            raise NotImplementedError(f"The numpy ufunc '{ufunc.__name__}' is not supported on Galois field arrays. If you believe this ufunc should be supported, please submit a GitHub issue at https://github.com/mhostetter/galois/issues.")

        else:
            inputs, kwargs = type(self)._view_inputs_as_ndarray(inputs, kwargs)
            output = super().__array_ufunc__(ufunc, method, *inputs, **kwargs)  # pylint: disable=no-member

            if ufunc in UFUNCS_REQUIRING_VIEW:
                output = output.view(type(self))

            return output

    ###############################################################################
    # Display methods
    ###############################################################################

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        # pylint: disable=attribute-defined-outside-init
        formatter = {}
        if type(self).display_mode == "poly":
            formatter["int"] = self._print_poly
            formatter["object"] = self._print_poly
        elif type(self).display_mode == "power":
            nonzero_idxs = np.nonzero(self)
            if self.ndim > 1:
                self._display_power_pre_width = 0 if nonzero_idxs[0].size == self.size else 1
                max_power = np.max(np.log(self[nonzero_idxs]))
                if max_power > 1:
                    self._display_power_width = self._display_power_pre_width + 2 + len(str(max_power))
                else:
                    self._display_power_width = self._display_power_pre_width + 1
            else:
                self._display_power_pre_width = None
                self._display_power_width = None
            formatter["int"] = self._print_power
            formatter["object"] = self._print_power
        elif self.dtype == np.object_:
            formatter["object"] = self._print_int

        cls = type(self)
        class_name = cls.__name__
        with np.printoptions(formatter=formatter):
            cls.__name__ = "GF"  # Rename the class so very large fields don't create large indenting
            string = super().__repr__()
        cls.__name__ = class_name

        if cls.degree == 1:
            order = "{}".format(cls.order)
        else:
            order = "{}^{}".format(cls.characteristic, cls.degree)

        # Remove the dtype from the repr and add the Galois field order
        dtype_idx = string.find("dtype")
        if dtype_idx == -1:
            string = string[:-1] + f", order={order})"
        else:
            string = string[:dtype_idx] + f"order={order})"

        return string

    @staticmethod
    def _print_int(element):
        return "{:d}".format(int(element))

    def _print_poly(self, element):
        poly = integer_to_poly(element, type(self).characteristic)
        poly_var = "α" if type(self).primitive_element == type(self).characteristic else "x"
        return poly_to_str(poly, poly_var=poly_var)

    def _print_power(self, element):
        if element == 0:
            s = "-∞"
        else:
            power = type(self)._ufuncs["log"](element)
            if power > 1:
                s = f"α^{power}"
            elif power == 1:
                s = "α"
            else:
                s = "1"

            if self._display_power_pre_width:
                s = " " + s

        if self._display_power_width:
            return s + " "*(self._display_power_width - len(s))
        else:
            return s

    @classmethod
    def _poly_eval(cls, coeffs, x):
        coeffs = cls(coeffs)  # Convert coefficient into the field
        coeffs = coeffs.view(np.ndarray)  # View cast to normal integers so ufunc_poly_eval call uses normal arithmetic
        coeffs = np.atleast_1d(coeffs)
        if coeffs.size == 1:
            # TODO: Why must coeffs have atleast 2 elements otherwise it will be converted to a scalar, not 1d array?
            coeffs = np.insert(coeffs, 0, 0)

        x = cls(x)  # Convert evaluation values into the field (checks that values are in the field)
        x = x.view(np.ndarray)  # View cast to normal integers so ufunc_poly_eval call uses normal arithmetic
        x = np.atleast_1d(x)

        if cls.dtypes[-1] == np.object_:
            # For object dtypes, call the vectorized classmethod
            y = cls._ufuncs["poly_eval"](coeffs=coeffs, values=x)  # pylint: disable=not-callable
        else:
            # For integer dtypes, call the JIT-compiled gufunc
            y = np.copy(x)
            cls._ufuncs["poly_eval"](coeffs, x, y, casting="unsafe")  # pylint: disable=not-callable

        y = cls(y)
        if y.size == 1:
            y = y[0]

        return y