# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-2-Clause

from numba_cuda_mlir.numba_cuda.extending import (
    core_models,
    register_model,
    type_callable,
    unbox,
    NativeValue,
    make_attribute_wrapper,
    box,
)
from numba_cuda_mlir.numba_cuda import types
from numba_cuda_mlir.numba_cuda import cgutils
import warnings
from numba_cuda_mlir.numba_cuda.core.errors import (
    NumbaExperimentalFeatureWarning,
    NumbaValueError,
)
from numpy.polynomial.polynomial import Polynomial
from contextlib import ExitStack
from numba_cuda_mlir.numba_cuda.core.imputils import Registry

# The Polynomial(...) constructor @lower codegen built llvmlite IR and is
# filtered out on the MLIR path; the live datamodel, typing (@type_callable),
# attribute wrappers and box/unbox below do not. The registry is retained
# because the target context installs it.
registry = Registry("np.polynomial_core")


@register_model(types.PolynomialType)
class PolynomialModel(core_models.StructModel):
    def __init__(self, dmm, fe_type):
        members = [
            ("coef", fe_type.coef),
            ("domain", fe_type.domain),
            ("window", fe_type.window),
            # Introduced in NumPy 1.24, maybe leave it out for now
            # ('symbol', types.string)
        ]
        super().__init__(dmm, fe_type, members)


@type_callable(Polynomial)
def type_polynomial(context):
    def typer(coef, domain=None, window=None):
        default_domain = types.Array(types.int64, 1, "C")
        double_domain = types.Array(types.double, 1, "C")
        default_window = types.Array(types.int64, 1, "C")
        double_window = types.Array(types.double, 1, "C")
        double_coef = types.Array(types.double, 1, "C")

        warnings.warn(
            "Polynomial class is experimental",
            category=NumbaExperimentalFeatureWarning,
        )

        if isinstance(coef, types.Array) and all([a is None for a in (domain, window)]):
            if coef.ndim == 1:
                # If Polynomial(coef) is called, coef is cast to double dtype,
                # and domain and window are set to equal [-1, 1], i.e. have
                # integer dtype
                return types.PolynomialType(double_coef, default_domain, default_window, 1)
            else:
                msg = "Coefficient array is not 1-d"
                raise NumbaValueError(msg)
        elif all([isinstance(a, types.Array) for a in (coef, domain, window)]):
            if coef.ndim == 1:
                if all([a.ndim == 1 for a in (domain, window)]):
                    # If Polynomial(coef, domain, window) is called, then coef,
                    # domain and window are cast to double dtype
                    return types.PolynomialType(double_coef, double_domain, double_window, 3)
            else:
                msg = "Coefficient array is not 1-d"
                raise NumbaValueError(msg)

    return typer


make_attribute_wrapper(types.PolynomialType, "coef", "coef")
make_attribute_wrapper(types.PolynomialType, "domain", "domain")
make_attribute_wrapper(types.PolynomialType, "window", "window")
# Introduced in NumPy 1.24, maybe leave it out for now
# make_attribute_wrapper(types.PolynomialType, 'symbol', 'symbol')


@unbox(types.PolynomialType)
def unbox_polynomial(typ, obj, c):
    """
    Convert a Polynomial object to a native polynomial structure.
    """
    is_error_ptr = cgutils.alloca_once_value(c.builder, cgutils.false_bit)
    polynomial = cgutils.create_struct_proxy(typ)(c.context, c.builder)
    with ExitStack() as stack:
        natives = []
        for name in ("coef", "domain", "window"):
            attr = c.pyapi.object_getattr_string(obj, name)
            with cgutils.early_exit_if_null(c.builder, stack, attr):
                c.builder.store(cgutils.true_bit, is_error_ptr)
            t = getattr(typ, name)
            native = c.unbox(t, attr)
            c.pyapi.decref(attr)
            with cgutils.early_exit_if(c.builder, stack, native.is_error):
                c.builder.store(cgutils.true_bit, is_error_ptr)
            natives.append(native)

        polynomial.coef = natives[0]
        polynomial.domain = natives[1]
        polynomial.window = natives[2]

    return NativeValue(polynomial._getvalue(), is_error=c.builder.load(is_error_ptr))


@box(types.PolynomialType)
def box_polynomial(typ, val, c):
    """
    Convert a native polynomial structure to a Polynomial object.
    """
    ret_ptr = cgutils.alloca_once(c.builder, c.pyapi.pyobj)
    fail_obj = c.pyapi.get_null_object()

    with ExitStack() as stack:
        polynomial = cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val)
        coef_obj = c.box(typ.coef, polynomial.coef)
        with cgutils.early_exit_if_null(c.builder, stack, coef_obj):
            c.builder.store(fail_obj, ret_ptr)

        domain_obj = c.box(typ.domain, polynomial.domain)
        with cgutils.early_exit_if_null(c.builder, stack, domain_obj):
            c.builder.store(fail_obj, ret_ptr)

        window_obj = c.box(typ.window, polynomial.window)
        with cgutils.early_exit_if_null(c.builder, stack, window_obj):
            c.builder.store(fail_obj, ret_ptr)

        class_obj = c.pyapi.unserialize(c.pyapi.serialize_object(Polynomial))
        with cgutils.early_exit_if_null(c.builder, stack, class_obj):
            c.pyapi.decref(coef_obj)
            c.pyapi.decref(domain_obj)
            c.pyapi.decref(window_obj)
            c.builder.store(fail_obj, ret_ptr)

        if typ.n_args == 1:
            res1 = c.pyapi.call_function_objargs(class_obj, (coef_obj,))
            c.builder.store(res1, ret_ptr)
        else:
            res3 = c.pyapi.call_function_objargs(class_obj, (coef_obj, domain_obj, window_obj))
            c.builder.store(res3, ret_ptr)

        c.pyapi.decref(coef_obj)
        c.pyapi.decref(domain_obj)
        c.pyapi.decref(window_obj)
        c.pyapi.decref(class_obj)

    return c.builder.load(ret_ptr)
