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

from numba.cuda.extending import (
    core_models,
    register_model,
    type_callable,
    unbox,
    NativeValue,
    make_attribute_wrapper,
    box,
)
from numba.cuda import types
from numba.cuda import cgutils
import warnings
from numba.cuda.core.errors import (
    NumbaExperimentalFeatureWarning,
    NumbaValueError,
)
from numpy.polynomial.polynomial import Polynomial
from contextlib import ExitStack
import numpy as np
from llvmlite import ir
from numba.cuda.core.imputils import Registry

registry = Registry("np.polynomial_core")
lower = registry.lower


@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')


@lower(Polynomial, types.Array)
def impl_polynomial1(context, builder, sig, args):
    def to_double(arr):
        return np.asarray(arr, dtype=np.double)

    def const_impl():
        return np.asarray([-1, 1])

    typ = sig.return_type
    polynomial = cgutils.create_struct_proxy(typ)(context, builder)
    sig_coef = sig.args[0].copy(dtype=types.double)(sig.args[0])
    coef_cast = context.compile_internal(builder, to_double, sig_coef, args)
    sig_domain = sig.args[0].copy(dtype=types.intp)()
    sig_window = sig.args[0].copy(dtype=types.intp)()
    domain_cast = context.compile_internal(builder, const_impl, sig_domain, ())
    window_cast = context.compile_internal(builder, const_impl, sig_window, ())
    polynomial.coef = coef_cast
    polynomial.domain = domain_cast
    polynomial.window = window_cast

    return polynomial._getvalue()


@lower(Polynomial, types.Array, types.Array, types.Array)
def impl_polynomial3(context, builder, sig, args):
    def to_double(coef):
        return np.asarray(coef, dtype=np.double)

    typ = sig.return_type
    polynomial = cgutils.create_struct_proxy(typ)(context, builder)

    coef_sig = sig.args[0].copy(dtype=types.double)(sig.args[0])
    domain_sig = sig.args[1].copy(dtype=types.double)(sig.args[1])
    window_sig = sig.args[2].copy(dtype=types.double)(sig.args[2])
    coef_cast = context.compile_internal(
        builder, to_double, coef_sig, (args[0],)
    )
    domain_cast = context.compile_internal(
        builder, to_double, domain_sig, (args[1],)
    )
    window_cast = context.compile_internal(
        builder, to_double, window_sig, (args[2],)
    )

    domain_helper = context.make_helper(
        builder, domain_sig.return_type, value=domain_cast
    )
    window_helper = context.make_helper(
        builder, window_sig.return_type, value=window_cast
    )

    i64 = ir.IntType(64)
    two = i64(2)

    s1 = builder.extract_value(domain_helper.shape, 0)
    s2 = builder.extract_value(window_helper.shape, 0)
    pred1 = builder.icmp_signed("!=", s1, two)
    pred2 = builder.icmp_signed("!=", s2, two)

    with cgutils.if_unlikely(builder, pred1):
        context.fndesc.call_conv.return_user_exc(
            builder,
            ValueError,
            ("Domain has wrong number of elements.",),
        )

    with cgutils.if_unlikely(builder, pred2):
        context.fndesc.call_conv.return_user_exc(
            builder,
            ValueError,
            ("Window has wrong number of elements.",),
        )

    polynomial.coef = coef_cast
    polynomial.domain = domain_helper._getvalue()
    polynomial.window = window_helper._getvalue()

    return polynomial._getvalue()


@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)
