conditional

# Atomwise
from atomwise.utils.flyte import aw_dynamic, aw_task, aw_workflow

from flytekit import Resources, conditional
from flytekit.extend import Promise


SETTINGS = {
    'small': '1000Mi',
    'medium': '2000Mi',
    'large': '3000Mi'
}


@aw_task
def example_task(input: str, idx: int) -> tuple[str, str]:
    # NOTE: this would be based on dataframe size
    if idx == 0:
        size = 'small'
    elif idx == 1:
        size = 'medium'
    elif idx == 2:
        size = 'large'
    
    return input, size


@aw_dynamic
def example_dynamic(input: str) -> str:
    size = 'small'
    for i in range(3):
        if type(size) == Promise:
            result, size = (conditional("test")
                .if_(size == 'small')
                .then(example_task(input=input, idx=i).with_overrides(requests=Resources(mem=SETTINGS['small'])))
                .elif_(size == 'medium')
                .then(example_task(input=input, idx=i).with_overrides(requests=Resources(mem=SETTINGS['medium'])))
                .elif_(size == 'large')
                .then(example_task(input=input, idx=i).with_overrides(requests=Resources(mem=SETTINGS['large'])))
                .else_()
                .fail("Blah")
            )
        else:
            result, size = example_task(input=input, idx=i).with_overrides(requests=Resources(mem=SETTINGS['small']))

    return result


@aw_workflow
def example_workflow(input: str) -> str:
    return example_dynamic(input=input)