from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import PythonOperator, BranchPythonOperator
from datetime import datetime, timedelta
import pandas as pd
import random

# Default args definition
default_args = {
    'owner': 'Rafael',
    'depends_on_past': False,
    'start_date': datetime(2020, 11, 29, 18, 20),
    'email': ['example1@example.com', 'example2@example.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'Retry_delay': timedelta(minutes=1)
}

# Dag definition
dag = DAG(
    'treino-03',
    description="Extrai dados do Titanic e calcula idade media para homens ou mulheres",
    default_args = default_args,
    schedule_interval='*/20 * * * *'
)

get_data = BashOperator(
    task_id='get-data',
    bash_command='curl https://raw.githubusercontent.com/A3Data/hermione/master/hermione/file_text/train.csv -o /usr/local/airflow/data/train.csv',
    dag=dag
)

def sorteia_h_m():
    return random.choice(['male', 'female'])

escolhe_h_m = PythonOperator(
    task_id='escolhe-h-m',
    python_callable=sorteia_h_m,
    dag=dag
)

def MouF(**context):
    value=context['task_instance'].xcom_pull(task_ids='escolhe-h-m')
    if value == 'male':
        return 'branch_homem'
    else:
        return 'branch_mulher'

male_female = BranchPythonOperator(
    task_id='condicional',
    python_callable=MouF,
    provide_context=True,
    dag=dag
)

def mean_homem():
    df = pd.read_csv('/usr/local/airflow/data/train.csv')
    med = df.loc[df.Sex == 'male'].Age.mean()
    print(f'Media de idade dos homens no Titanic: {med}')

branch_homem = PythonOperator(
    task_id='branch_homem',
    python_callable=mean_homem,
    dag=dag
)

def mean_mulher():
    df = pd.read_csv('/usr/local/airflow/data/train.csv')
    med = df.loc[df.Sex == 'female'].Age.mean()
    print(f'Media de idade das mulheres no Titanic: {med}')

branch_mulher = PythonOperator(
    task_id='branch_mulher',
    python_callable=mean_mulher,
    dag=dag
)



get_data >> escolhe_h_m >> male_female >> [branch_homem, branch_mulher]

[((593, 764), 'airflow.DAG', 'DAG', (['"""treino-03"""'], {'description': '"""Extrai dados do Titanic e calcula idade media para homens ou mulheres"""', 'default_args': 'default_args', 'schedule_interval': '"""*/20 * * * *"""'}), "('treino-03', description=\n    'Extrai dados do Titanic e calcula idade media para homens ou mulheres',\n    default_args=default_args, schedule_interval='*/20 * * * *')\n", (596, 764), False, 'from airflow import DAG\n'), ((788, 983), 'airflow.operators.bash_operator.BashOperator', 'BashOperator', ([], {'task_id': '"""get-data"""', 'bash_command': '"""curl https://raw.githubusercontent.com/A3Data/hermione/master/hermione/file_text/train.csv -o /usr/local/airflow/data/train.csv"""', 'dag': 'dag'}), "(task_id='get-data', bash_command=\n    'curl https://raw.githubusercontent.com/A3Data/hermione/master/hermione/file_text/train.csv -o /usr/local/airflow/data/train.csv'\n    , dag=dag)\n", (800, 983), False, 'from airflow.operators.bash_operator import BashOperator\n'), ((1068, 1143), 'airflow.operators.python_operator.PythonOperator', 'PythonOperator', ([], {'task_id': '"""escolhe-h-m"""', 'python_callable': 'sorteia_h_m', 'dag': 'dag'}), "(task_id='escolhe-h-m', python_callable=sorteia_h_m, dag=dag)\n", (1082, 1143), False, 'from airflow.operators.python_operator import PythonOperator, BranchPythonOperator\n'), ((1359, 1459), 'airflow.operators.python_operator.BranchPythonOperator', 'BranchPythonOperator', ([], {'task_id': '"""condicional"""', 'python_callable': 'MouF', 'provide_context': '(True)', 'dag': 'dag'}), "(task_id='condicional', python_callable=MouF,\n    provide_context=True, dag=dag)\n", (1379, 1459), False, 'from airflow.operators.python_operator import PythonOperator, BranchPythonOperator\n'), ((1671, 1746), 'airflow.operators.python_operator.PythonOperator', 'PythonOperator', ([], {'task_id': '"""branch_homem"""', 'python_callable': 'mean_homem', 'dag': 'dag'}), "(task_id='branch_homem', python_callable=mean_homem, dag=dag)\n", (1685, 1746), False, 'from airflow.operators.python_operator import PythonOperator, BranchPythonOperator\n'), ((1964, 2041), 'airflow.operators.python_operator.PythonOperator', 'PythonOperator', ([], {'task_id': '"""branch_mulher"""', 'python_callable': 'mean_mulher', 'dag': 'dag'}), "(task_id='branch_mulher', python_callable=mean_mulher, dag=dag)\n", (1978, 2041), False, 'from airflow.operators.python_operator import PythonOperator, BranchPythonOperator\n'), ((354, 384), 'datetime.datetime', 'datetime', (['(2020)', '(11)', '(29)', '(18)', '(20)'], {}), '(2020, 11, 29, 18, 20)\n', (362, 384), False, 'from datetime import datetime, timedelta\n'), ((546, 566), 'datetime.timedelta', 'timedelta', ([], {'minutes': '(1)'}), '(minutes=1)\n', (555, 566), False, 'from datetime import datetime, timedelta\n'), ((1019, 1052), 'random.choice', 'random.choice', (["['male', 'female']"], {}), "(['male', 'female'])\n", (1032, 1052), False, 'import random\n'), ((1502, 1550), 'pandas.read_csv', 'pd.read_csv', (['"""/usr/local/airflow/data/train.csv"""'], {}), "('/usr/local/airflow/data/train.csv')\n", (1513, 1550), True, 'import pandas as pd\n'), ((1790, 1838), 'pandas.read_csv', 'pd.read_csv', (['"""/usr/local/airflow/data/train.csv"""'], {}), "('/usr/local/airflow/data/train.csv')\n", (1801, 1838), True, 'import pandas as pd\n')]