import ibis.expr.datatypes as dt
import ibis.expr.rules as rlz
from ibis.expr.operations import ValueOp


class JulianDay(ValueOp):
    arg = rlz.string

    output_dtype = dt.float32
    output_shape = rlz.shape_like('arg')

"""
Explanation: Adding an Elementwise Operation
This notebook will show you how to add a new elementwise operation to an existing backend.
We are going to add julianday, a function supported by the SQLite database, to the SQLite Ibis backend.
The Julian day of a date, is the number of days since January 1st, 4713 BC. For more information check the Julian day wikipedia page.
Step 1: Define the Operation
Let's define the julianday operation as a function that takes one string input argument and returns a float.
python
def julianday(date: str) -&gt; float:
    """Julian date"""
End of explanation
"""


from ibis.expr.types import BinaryValue, StringValue


def julianday(string_value):
    return JulianDay(string_value).to_expr()


StringValue.julianday = julianday

"""
Explanation: We just defined a JulianDay class that takes one argument of type string or binary, and returns a float.
Step 2: Define the API
Because we know the output type of the operation, to make an expression out of JulianDay we simply need to construct it and call its ibis.expr.types.Node.to_expr method.
We still need to add a method to StringValue and BinaryValue (this needs to work on both scalars and columns).
When you add a method to any of the expression classes whose name matches *Value both the scalar and column child classes will pick it up, making it easy to define operations for both scalars and columns in one place.
We can do this by defining a function and assigning it to the appropriate class
of expressions.
End of explanation
"""


import ibis

t = ibis.table([('string_col', 'string')], name='t')

t.string_col.julianday()

"""
Explanation: Interlude: Create some expressions with sha1
End of explanation
"""


import sqlalchemy as sa


@ibis.sqlite.add_operation(JulianDay)
def _julianday(translator, expr):
    # pull out the arguments to the expression
    (arg,) = expr.op().args

    # compile the argument
    compiled_arg = translator.translate(arg)

    # return a SQLAlchemy expression that calls into the SQLite julianday function
    return sa.func.julianday(compiled_arg)

"""
Explanation: Step 3: Turn the Expression into SQL
End of explanation
"""


!curl -LsS -o $TEMPDIR/geography.db 'https://storage.googleapis.com/ibis-tutorial-data/geography.db'

import os
import tempfile

import ibis

db_fname = os.path.join(tempfile.gettempdir(), 'geography.db')

con = ibis.sqlite.connect(db_fname)

"""
Explanation: Step 4: Putting it all Together
End of explanation
"""


independence = con.table('independence')
independence

day = independence.independence_date.cast('string')
day

julianday_expr = day.julianday().name("jday")
julianday_expr

sql_expr = julianday_expr.compile()
print(sql_expr)

result = julianday_expr.execute()
result.head()

"""
Explanation: Create and execute a julianday expression
End of explanation
"""


scalar = ibis.literal('2010-03-14')
scalar

julianday_scalar = scalar.julianday()

con.execute(julianday_scalar)

"""
Explanation: Because we've defined our operation on StringValue, and not just on StringColumn we get operations on both string scalars and string columns for free
End of explanation
"""


