# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_functions.ipynb.

# %% auto 0
__all__ = ['Max', 'Min', 'Sum', 'Avg', 'Count', 'Abs', 'Round', 'First', 'Last', 'MonthsBetween', 'RowNumber', 'Rank',
           'DenseRank', 'PercentRank', 'CumeDist', 'Ntile', 'Lag', 'Lead', 'FirstValue', 'LastValue', 'NthValue',
           'CustomFunction', 'Cast', 'Coalesce', 'Concat', 'convert_date_format', 'Date', 'AddMonths', 'DateDiff',
           'DateTrunc']

# %% ../nbs/03_functions.ipynb 2
from fastcore.all import patch, delegates
from .utils import execute
from .terms import FieldBase, Field, custom_func, OverClause

# %% ../nbs/03_functions.ipynb 4
class CustomFunction(FieldBase):
    """A convenient class for creating custom functions."""

    def __init__(self, 
                 func_name: str,    # name of the function 
                 arg_names: list,   # list of arg names
                 window_func=False,  # whether this can be used as a window function
                 distinct_option=False # whether this function can be used with the distinct option
        ) -> None:
        import inspect
        super().__init__()
        self.func_name = func_name
        self.arg_names = arg_names
        self.window_func = window_func
        self.distinct_option = distinct_option
        self.distinct_ = False  # whether the distinct option is used
        self.get_sql = self.execute
        self.__qualname__ = func_name
        self.__signature__ = inspect.Signature(parameters=[inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD) for name in arg_names]) 
        self.__doc__ = None

    def __call__(self, *args):
        if len(args) != len(self.arg_names):
            raise ValueError(f"The number of args provided {len(args)} is not the same as the number of args expected by this function ({len(self.arg_names)})!")
        def func(*args):
            if self.distinct_:
                return f"{self.func_name}(DISTINCT {', '.join(args)})"
            else:
                return f"{self.func_name}({', '.join(args)})"

        self.func = func
        self.args = args
        return self

    def execute(self, **kwargs):
        args = [str(execute(arg, **kwargs)) for arg in self.args]
        return self.func(*args)

    def distinct(self):
        if self.distinct_option==True:
            self.distinct_ = True
            return self
        else:
            raise ValueError(f"This function does not support the distinct option!!")


@patch
def over(self:CustomFunction, partition_by):
    if self.window_func==True:
        return OverClause(self).over(partition_by)
    else:
        raise ValueError(f"This function is not a window function!!")

# %% ../nbs/03_functions.ipynb 11
Max = CustomFunction('MAX', ['field'], window_func=True, distinct_option=True)
Min = CustomFunction('MIN', ['field'], window_func=True, distinct_option=True)
Sum = CustomFunction('SUM', ['field'], window_func=True, distinct_option=True)
Avg = CustomFunction('AVG', ['field'], window_func=True, distinct_option=True)
Count = CustomFunction('COUNT', ['field'], window_func=True, distinct_option=True)
Abs = CustomFunction('ABS', ['field'])
Round = CustomFunction('ROUND', ['field', 'decimals'])
First = CustomFunction('FIRST', ['field'])
Last = CustomFunction('LAST', ['field'])

# %% ../nbs/03_functions.ipynb 21
@custom_func
def Cast(field, type):
    return f"CAST({field} AS {type})"

@custom_func
def Coalesce(*args):
    return f"COALESCE({', '.join(args)})"

@custom_func
def Concat(*args):
    return f"CONCAT({', '.join(args)})"


# date functions
def convert_date_format(format, dialect='sql'):
    if dialect == 'spark':
        return (format.replace('YYYY', 'yyyy')
                .replace('YY', 'yy')
                .replace('DD', 'dd')
                )
    elif dialect == 'snowflake':
        return (format.replace('HH', 'HH24')
                .replace('mm', 'MI')
        )
    else:
        return format

@custom_func
def Date(expression, format=None, dialect='sql'):
    if format:
        format = convert_date_format(format, dialect)
        return f"DATE('{expression}', '{format}')"
    else:
        return f"DATE('{expression}')"

@custom_func
def AddMonths(date, months, dialect='sql'):
    if dialect in ['snowflake', 'spark']:
        return f"ADD_MONTHS({date}, {months})"
    elif dialect == 'athena':
        return f"date_add('month', {months}, {date})"
    else:
        return f"DATE_ADD({date}, INTERVAL {months} MONTH)"

@custom_func
def DateDiff(interval, start_date, end_date, dialect='sql'):
    if dialect == 'athena':
        return f"date_diff('{interval}', {start_date}, {end_date})"
    else:
        return f"DATEDIFF({interval}, {start_date}, {end_date})"

@custom_func
def DateTrunc(interval, date, dialect='sql'):
    if dialect == 'spark':
        return f"TRUNC({date}, '{interval}')"
    else:
        return f"DATE_TRUNC('{interval}', {date})"


# %% ../nbs/03_functions.ipynb 22
MonthsBetween = CustomFunction('MONTHS_BETWEEN', ['start_date', 'end_date'])

# Window Functions
RowNumber = CustomFunction('ROW_NUMBER', [], window_func=True)
Rank = CustomFunction('RANK', [], window_func=True)
DenseRank = CustomFunction('DENSE_RANK', [], window_func=True)
PercentRank = CustomFunction('PERCENT_RANK', [], window_func=True)
CumeDist = CustomFunction('CUME_DIST', [], window_func=True)
Ntile = CustomFunction('NTILE', ['num_buckets'], window_func=True)
Lag = CustomFunction('LAG', ['field', 'offset'], window_func=True)
Lead = CustomFunction('LEAD', ['field', 'offset'], window_func=True)
FirstValue = CustomFunction('FIRST_VALUE', ['field'], window_func=True)
LastValue = CustomFunction('LAST_VALUE', ['field'], window_func=True)
NthValue = CustomFunction('NTH_VALUE', ['field', 'n'], window_func=True)
