numpy.jl

const CANONICALIZE_TIMEDELTA64 = Ref(false)

struct pyconvert_rule_numpysimplevalue{R,S} <: Function end

function (::pyconvert_rule_numpysimplevalue{R,SAFE})(::Type{T}, x::Py) where {R,SAFE,T}
    ans = Base.GC.@preserve x C.PySimpleObject_GetValue(R, getptr(x))
    if SAFE
        pyconvert_return(convert(T, ans))
    else
        pyconvert_tryconvert(T, ans)
    end
end

const NUMPY_SIMPLE_TYPES = [
    ("bool_", Bool),
    ("int8", Int8),
    ("int16", Int16),
    ("int32", Int32),
    ("int64", Int64),
    ("uint8", UInt8),
    ("uint16", UInt16),
    ("uint32", UInt32),
    ("uint64", UInt64),
    ("float16", Float16),
    ("float32", Float32),
    ("float64", Float64),
    ("complex32", ComplexF16),
    ("complex64", ComplexF32),
    ("complex128", ComplexF64),
]

"""
    pydatetime64([year, month, day, hour, minute, second, millisecond, microsecond, nanosecond, weeks])
    pydatetime64(; [year, month, day, hour, minute, second, millisecond, microsecond, nanosecond, weeks])

Create a `numpy.pydatetime64` object from the given time arguments.
Arguments can be supplied either as keyword arguments or positional arguments in the above order.

Examples:
```julia
pydatetime64(2025, 1, 2, 3, 4, 5, 6, 7, 8)
# Python: np.datetime64('2025-01-02T03:04:05.006007008')

pydatetime64(year = 2025, month = 5, day = 20, nanosecond = 1)
# Python: np.datetime64('2025-05-20T00:00:00.000000001')
```
"""
function pydatetime64(
    _year::Integer=0, _month::Integer=1, _day::Integer=1, _hour::Integer=0, _minute::Integer=0,_second::Integer=0, _millisecond::Integer=0, _microsecond::Integer=0, _nanosecond::Integer=0;
    year::Integer=_year, month::Integer=_month, day::Integer=_day, hour::Integer=_hour, minute::Integer=_minute, second::Integer=_second,
    millisecond::Integer=_millisecond, microsecond::Integer=_microsecond, nanosecond::Integer=_nanosecond
)
    pyimport("numpy").datetime64("$(DateTime(year, month, day, hour, minute, second))") + 
        pytimedelta64(; milliseconds = millisecond, microseconds = microsecond, nanoseconds = nanosecond)
end
function pydatetime64(@nospecialize(x::T)) where T <: Period
    T <: Union{Week, Day, Hour, Minute, Second, Millisecond, Microsecond} || 
        error("Unsupported Period type: `$x::$T`. Consider using pytimedelta64 instead.")
    args = map(Base.Fix1(isa, x), (Day, Second, Millisecond, Microsecond,  Minute, Hour, Week))
    pydatetime64(map(Base.Fix1(*, x.value), args)...)
end
function pydatetime64(x::Union{Date, DateTime})
    pyimport("numpy").datetime64("$x")
end
export pydatetime64

"""
    pytimedelta64([years, months, days, hours, minutes, seconds, milliseconds, microseconds, nanoseconds, weeks]; canonicalize=false)
    pytimedelta64(; [years, months, days, hours, minutes, seconds, milliseconds, microseconds, nanoseconds, weeks], canonicalize=false)

Create a `numpy.timedelta64` object from the given time arguments.
Arguments can be supplied either as keyword arguments or positional arguments in the above order.
`years` and `months` cannot be mixed with other units.

Parameters:
- `canonicalize=false`: If `false`, the unit of the result will be the unit of the least unit that was specified, 
    if `true`, the result will be converted to least possible unit, e.g. 60s will be converted to 1m.

Examples:
```julia
pytimedelta64(1, 2)
# Python: np.timedelta64(14,'M')

pytimedelta64(years = 1, months = 12; canonicalize = true)
# Python: np.timedelta64(2,'Y')

pytimedelta64(hours = 3, minutes = 60)
# Python: np.timedelta64(240,'m')
```
"""
function pytimedelta64(
    _years::Union{Nothing,Integer}=nothing, _months::Union{Nothing,Integer}=nothing,
    _days::Union{Nothing,Integer}=nothing, _hours::Union{Nothing,Integer}=nothing,
    _minutes::Union{Nothing,Integer}=nothing, _seconds::Union{Nothing,Integer}=nothing,
    _milliseconds::Union{Nothing,Integer}=nothing, _microseconds::Union{Nothing,Integer}=nothing,
    _nanoseconds::Union{Nothing,Integer}=nothing, _weeks::Union{Nothing,Integer}=nothing;
    years::Union{Nothing,Integer}=_years, months::Union{Nothing,Integer}=_months,
    days::Union{Nothing,Integer}=_days, hours::Union{Nothing,Integer}=_hours,
    minutes::Union{Nothing,Integer}=_minutes, seconds::Union{Nothing,Integer}=_seconds,
    milliseconds::Union{Nothing,Integer}=_milliseconds, microseconds::Union{Nothing,Integer}=_microseconds,
    nanoseconds::Union{Nothing,Integer}=_nanoseconds, weeks::Union{Nothing,Integer}=_weeks,
    canonicalize::Bool = false)

    y::Integer = something(years, 0)
    m::Integer = something(months, 0)
    d::Integer = something(days, 0)
    h::Integer = something(hours, 0)
    min::Integer = something(minutes, 0)
    s::Integer = something(seconds, 0)
    ms::Integer = something(milliseconds, 0)
    µs::Integer = something(microseconds, 0)
    ns::Integer = something(nanoseconds, 0)
    w::Integer = something(weeks, 0)
    cp = sum((
        Year(y), Month(m), Week(w), Day(d), Hour(h), Minute(min), Second(s), Millisecond(ms), Microsecond(µs), Nanosecond(ns))
    )
    # make sure the correct unit is used when value is 0
    if isempty(cp.periods)
        Units = (Second, Year, Month, Week, Day, Hour, Minute, Second, Millisecond, Microsecond, Nanosecond)
        index::Integer = findlast(!isnothing, (0, years, months, weeks, days, hours, minutes, seconds, milliseconds, microseconds, nanoseconds));
        pytimedelta64(Units[index](0))
    else
        pytimedelta64(cp; canonicalize)
    end
end
"""
    pytimedelta64(x::Union{Period, CompoundPeriod}; canonicalize=false)

Convert a Julia `Period` or `CompoundPeriod` to a `numpy.timedelta64` object.
"""
function pytimedelta64(@nospecialize(x::T); canonicalize::Bool = false) where T <: Period
    canonicalize && return pytimedelta64(@__MODULE__().canonicalize(x))

    index = findfirst(T .== (Year, Month, Week, Day, Hour, Minute, Second, Millisecond, Microsecond, Nanosecond, T))::Int
    unit = ("Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns", "")[index]
    pyimport("numpy").timedelta64(x.value, unit)
end
function pytimedelta64(x::CompoundPeriod; canonicalize::Bool = false)
    canonicalize && (x = @__MODULE__().canonicalize(x))
    isempty(x.periods) ? pytimedelta64(Second(0)) : sum(pytimedelta64.(x.periods))
end
"""
    pytimedelta64(x::Integer)

Create a dimensionless `numpy.timedelta64` object. If added to another `numpy.timedelta64` object, it will be interpreted as being of the same unit.

Examples:
```julia
pytimedelta64(2)
# Python: np.timedelta64(2)

pytimedelta64(Hour(1)) + pytimedelta64(2)
# Python: np.timedelta64(3,'h')

pytimedelta64(2) + pytimedelta64(Hour(1)) + pytimedelta64(Minute(1))
# Python: np.timedelta64(181,'m')
```
"""
function pytimedelta64(x::Integer)
    pyimport("numpy").timedelta64(x)
end
export pytimedelta64

function pyconvert_rule_datetime64(::Type{DateTime}, x::Py)
    unit, count = pyconvert(Tuple, pyimport("numpy").datetime_data(x))
    value = reinterpret(Int64, pyconvert(Vector, x))[1]
    units = ("Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns")
    types = (Year, Month, Week, Day, Hour, Minute, Second, Millisecond, Microsecond, Nanosecond)
    T = types[findfirst(==(unit), units)::Int]
    pyconvert_return(DateTime(_base_datetime) + T(value * count))
end

function pyconvert_rule_timedelta64(::Type{CompoundPeriod}, x::Py)
    unit, count = pyconvert(Tuple, pyimport("numpy").datetime_data(x))
    value = reinterpret(Int64, pyconvert(Vector, x))[1]
    units = ("Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns")
    types = (Year, Month, Week, Day, Hour, Minute, Second, Millisecond, Microsecond, Nanosecond)
    T = types[findfirst(==(unit), units)::Int]
    cp = CompoundPeriod(T(value * count))
    CANONICALIZE_TIMEDELTA64[] && (cp = @__MODULE__().canonicalize(cp))
    pyconvert_return(cp)
end

function pyconvert_rule_timedelta64(::Type{T}, x::Py) where T<:Period
    pyconvert_return(convert(T, pyconvert_rule_timedelta64(CompoundPeriod, x)))
end

function init_numpy()
    for (t, T) in NUMPY_SIMPLE_TYPES
        isbool = occursin("bool", t)
        isint = occursin("int", t) || isbool
        isuint = occursin("uint", t) || isbool
        isfloat = occursin("float", t)
        iscomplex = occursin("complex", t)
        isreal = isint || isfloat
        isnumber = isreal || iscomplex

        name = "numpy:$t"
        rule = pyconvert_rule_numpysimplevalue{T,false}()
        saferule = pyconvert_rule_numpysimplevalue{T,true}()

        pyconvert_add_rule(name, T, saferule, PYCONVERT_PRIORITY_ARRAY)
        isuint && pyconvert_add_rule(name, UInt, sizeof(T) ≤ sizeof(UInt) ? saferule : rule)
        isuint && pyconvert_add_rule(name, Int, sizeof(T) < sizeof(Int) ? saferule : rule)
        isint &&
            !isuint &&
            pyconvert_add_rule(name, Int, sizeof(T) ≤ sizeof(Int) ? saferule : rule)
        isint && pyconvert_add_rule(name, Integer, rule)
        isfloat && pyconvert_add_rule(name, Float64, saferule)
        isreal && pyconvert_add_rule(name, Real, rule)
        iscomplex && pyconvert_add_rule(name, ComplexF64, saferule)
        iscomplex && pyconvert_add_rule(name, Complex, rule)
        isnumber && pyconvert_add_rule(name, Number, rule)
    end

    priority = PYCONVERT_PRIORITY_ARRAY
    pyconvert_add_rule("numpy:datetime64", DateTime, pyconvert_rule_datetime64, priority)
    let TT = (CompoundPeriod, Year, Month, Day, Hour, Minute, Second, Millisecond, Microsecond, Nanosecond, Week)
        Base.Cartesian.@nexprs 11 i -> pyconvert_add_rule("numpy:timedelta64", TT[i], pyconvert_rule_timedelta64, priority)
    end

    priority = PYCONVERT_PRIORITY_CANONICAL
    pyconvert_add_rule("numpy:datetime64", DateTime, pyconvert_rule_datetime64, priority)
    pyconvert_add_rule("numpy:timedelta64", Nanosecond, pyconvert_rule_timedelta, priority)    
end