#!/usr/bin/python3
# Halide tutorial lesson 14: The Halide type system

# This lesson more precisely describes Halide's type system.

# On linux, you can compile and run it like so:
# g++ lesson_14*.cpp -g -I ../include -L ../bin -lHalide -lpthread -ldl -o lesson_14 -std=c++11
# LD_LIBRARY_PATH=../bin ./lesson_14

# On os x:
# g++ lesson_14*.cpp -g -I ../include -L ../bin -lHalide -o lesson_14 -std=c++11
# DYLD_LIBRARY_PATH=../bin ./lesson_14

# If you have the entire Halide source tree, you can also build it by
# running:
#    make tutorial_lesson_14_types
# in a shell with the current directory at the top of the halide
# source tree.

#include "Halide.h"
#include <stdio.h>
#using namespace Halide
from halide import *

# This function is used to demonstrate generic code at the end of
# this lesson.
#Expr average(Expr a, Expr b)

def main():

    # All Exprs have a scalar type, and all Funcs evaluate to one or
    # more scalar types. The scalar types in Halide are unsigned
    # integers of various bit widths, signed integers of the same set
    # of bit widths, floating point numbers in single and double
    # precision, and opaque handles (equivalent to void *). The
    # following array contains all the legal types.

    valid_halide_types = [
        UInt(8), UInt(16), UInt(32), UInt(64),
        Int(8), Int(16), Int(32), Int(64),
        Float(32), Float(64), Handle() ]


    # Constructing and inspecting types.
    if True:
        # You can programmatically examine the properties of a Halide
        # type. This is useful when you write a C++ function that has
        # Expr arguments and you wish to check their types:
        assert UInt(8).bits() == 8
        assert Int(8).is_int()


        # You can also programmatically construct Types as a function of other Types.
        t = UInt(8)
        t = t.with_bits(t.bits() * 2)
        assert t == UInt(16)

        # Or construct a Type from a C++ scalar type
        #assert type_of<float>() == Float(32)

        # The Type struct is also capable of representing vector types,
        # but this is reserved for Halide's internal use. You should
        # vectorize code by using Func::vectorize, not by attempting to
        # construct vector expressions directly. You may encounter vector
        # types if you programmatically manipulate lowered Halide code,
        # but this is an advanced topic (see Func::add_custom_lowering_pass).

        # You can query any Halide Expr for its type. An Expr
        # representing a Var has type Int(32):
        x = Var("x")
        assert Expr(x).type() == Int(32)

        # Most transcendental functions in Halide cast their inputs to a
        # Float(32) and return a Float(32):
        assert sin(x).type() == Float(32)

        # You can cast an Expr from one Type to another using the cast operator:
        assert cast(UInt(8), x).type() == UInt(8)

        # This also comes in a template form that takes a C++ type.
        #assert cast<uint8_t>(x).type() == UInt(8)

        # You can also query any defined Func for the types it produces.
        f1 = Func("f1")
        f1[x] = cast(UInt(8), x)
        assert f1.output_types()[0] == UInt(8)

        f2 = Func("f2")
        f2[x] = Tuple(x, sin(x))
        assert f2.output_types()[0] == Int(32) and \
               f2.output_types()[1] == Float(32)



    # Type promotion rules.
    if True:
        # When you combine Exprs of different types (e.g. using '+',
        # '*', etc), Halide uses a system of type promotion
        # rules. These differ to C's rules. To demonstrate these
        # we'll make some Exprs of each type.
        x = Var("x")
        u8 = cast(UInt(8), x)
        u16 = cast(UInt(16), x)
        u32 = cast(UInt(32), x)
        u64 = cast(UInt(64), x)
        s8 = cast(Int(8), x)
        s16 = cast(Int(16), x)
        s32 = cast(Int(32), x)
        s64 = cast(Int(64), x)
        f32 = cast(Float(32), x)
        f64 = cast(Float(64), x)

        # The rules are as follows, and are applied in the order they are
        # written below.

        # 1) It is an error to cast or use arithmetic operators on Exprs of type Handle().

        # 2) If the types are the same, then no type conversions occur.
        for t in valid_halide_types:
            # Skip the handle type.
            if t.is_handle():
                continue
            e = cast(t, x)
            assert (e + e).type() == e.type()


        # 3) If one type is a float but the other is not, then the
        # non-float argument is promoted to a float (possibly causing a
        # loss of precision for large integers).
        assert (u8 + f32).type() == Float(32)
        assert (f32 + s64).type() == Float(32)
        assert (u16 + f64).type() == Float(64)
        assert (f64 + s32).type() == Float(64)

        # 4) If both types are float, then the narrower argument is
        # promoted to the wider bit-width.
        assert (f64 + f32).type() == Float(64)

        # The rules above handle all the floating-point cases. The
        # following three rules handle the integer cases.

        # 5) If one of the expressions is an integer constant, then it is
        # coerced to the type of the other expression.
        assert (u32 + 3).type() == UInt(32)
        assert (3 + s16).type() == Int(16)

        # If this rule would cause the integer to overflow, then Halide
        # will trigger an error, e.g. uncommenting the following line
        # will cause this program to terminate with an error.
        # Expr bad = u8 + 257

        # 6) If both types are unsigned integers, or both types are
        # signed integers, then the narrower argument is promoted to
        # wider type.
        assert (u32 + u8).type() == UInt(32)
        assert (s16 + s64).type() == Int(64)

        # 7) If one type is signed and the other is unsigned, both
        # arguments are promoted to a signed integer with the greater of
        # the two bit widths.
        assert (u8 + s32).type() == Int(32)
        assert (u32 + s8).type() == Int(32)

        # Note that this may silently overflow the unsigned type in the
        # case where the bit widths are the same.
        assert (u32 + s32).type() == Int(32)

        if False: # evaluate<X> not yet exposed to python
            # When an unsigned Expr is converted to a wider signed type in
            # this way, it is first widened to a wider unsigned type
            # (zero-extended), and then reinterpreted as a signed
            # integer. I.e. casting the UInt(8) value 255 to an Int(32)
            # produces 255, not -1.
            #int32_t result32 = evaluate<int>(cast<int32_t>(cast<uint8_t>(255)))
            assert result32 == 255

            # When a signed type is explicitly converted to a wider unsigned
            # type with the cast operator (the type promotion rules will
            # never do this automatically), it is first converted to the
            # wider signed type (sign-extended), and then reinterpreted as
            # an unsigned integer. I.e. casting the Int(8) value -1 to a
            # UInt(16) produces 65535, not 255.
            #uint16_t result16 = evaluate<uint16_t>(cast<uint16_t>(cast<int8_t>(-1)))
            assert result16 == 65535


    # The type Handle().
    if True:
        # Handle is used to represent opaque pointers. Applying
        # type_of to any pointer type will return Handle()

        #assert type_of<void *>() == Handle()
        #assert type_of<const char * const **>() == Handle()
        # (not clear what the proper python version would be)

        # Handles are always stored as 64-bit, regardless of the compilation
        # target.
        assert Handle().bits() == 64

        # The main use of an Expr of type Handle is to pass
        # it through Halide to other external code.


    # Generic code.
    if True:
        # The main explicit use of Type in Halide is to write Halide
        # code parameterized by a Type. In C++ you'd do this with
        # templates. In Halide there's no need - you can inspect and
        # modify the types dynamically at C++ runtime instead. The
        # function defined below averages two expressions of any
        # equal numeric type.
        x = Var("x")
        assert average(cast(Float(32), x), 3.0).type() == Float(32)
        assert average(x, 3).type() == Int(32)
        assert average(cast(UInt(8), x), cast(UInt(8), 3)).type() == UInt(8)


    print("Success!")

    return 0


def average(a, b):

    if type(a) is not Expr:
        a = Expr(a)

    if type(b) is not Expr:
        b = Expr(b)


    "Expr average(Expr a, Expr b)"
    # Types must match.
    assert a.type() == b.type()

    # For floating point types:
    if (a.type().is_float()):
        # The '2' will be promoted to the floating point type due to
        # rule 3 above.
        return (a + b)/2


    # For integer types, we must compute the intermediate value in a
    # wider type to avoid overflow.
    narrow = a.type()
    wider = narrow.with_bits(narrow.bits() * 2)
    a = cast(wider, a)
    b = cast(wider, b)
    return cast(narrow, (a + b)/2)



if __name__ == "__main__":
    main()
