from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from datetime import datetime, timedelta
from airflow.contrib.operators.spark_submit_operator import SparkSubmitOperator
from datetime import timedelta,datetime
import pendulum
import pytz

tz = pytz.timezone('Asia/Shanghai')

default_args = {
    'owner': 'Terry',
    'depends_on_past': False,
    'start_date': pendulum.datetime(year=2023, month=12, day=8).astimezone(tz),
    'email': ['tian23.li@tcl.com'],
    'email_on_failure': True,
    'email_on_retry': True,
    'retries': 2,
    'retry_delay': timedelta(minutes=8),
}

_config = {
    'conf': {
        'spark.executor.extraJavaOptions': '-XX:+UseConcMarkSweepGC',
        "spark.dynamicAllocation.minExecutors":8,
        "spark.dynamicAllocation.maxExecutors":8,
        "spark.yarn.executor.memoryOverhead":2648,
        "spark.driver.maxResultSize": "8g",
        "spark.yarn.maxAppAttempts": 1,
        "spark.network.timeout": 1000000,
        "spark.executor.heartbeatInterval" : 1000000
    },
    'executor_cores': 4,
    'executor_memory': '10g',
    'name': '{{ task_instance.task_id }}',
    
    'driver_memory': '10g',
}

dag = DAG(
    dag_id='multi_interest_2024_dag',
    description='flow for multi-interest modeling',
    default_args=default_args,
    schedule_interval='10 03 * * *',#30 05
)

kill_iPython_job = BashOperator(
    task_id='kill_iPython_job',
    bash_command="ps -ef | grep ipykernel | grep -v grep | awk '{print $2}' | xargs kill -9",
    #bash_command="kill $(nvidia-smi -g 0)",
    dag=dag,
)
# kill_large_processes = BashOperator(
#     task_id='kill_large_processes',
#     ## 
#     bash_command="cd /data/gangyanyuan/scheduled_jobs /data/gangyanyuan/scheduled_jobs/kill-large-processes.sh ",
#     dag=dag,
# )
update_jx_lst_media = BashOperator(
    task_id='update_jx_lst_media',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/recsys2/bin/python 0_update_jx_lst_media.py",
    depends_on_past =False,
    dag=dag,
)

generate_user_click_data = BashOperator(
    task_id='generate_user_click_data',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/recsys2/bin/python 1_get_user_click_data.py",
    depends_on_past =False,
    dag=dag,
)

get_ecpm_aids = BashOperator(
    task_id='get_ecpm_aids',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep && /data/dev/miniconda/envs/recsys2/bin/python get_ecpm_aids.py",
    depends_on_past =False,
    dag=dag,
)

# ## jx data flow. 
update_user_click_all_channelname_60days_012468 = BashOperator(
    task_id='update_user_click_all_channelname_60days_012468',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_july && /data/dev/miniconda/envs/deepmatch/bin/python 2_3_update_user_click_all_channelname_60days_012468.py",
    depends_on_past =False,
    dag=dag,
)
update_user_click_all_channelname_60days_3579 = BashOperator(
    task_id='update_user_click_all_channelname_60days_3579',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep && /data/dev/miniconda/envs/deepmatch/bin/python 2_4_update_user_click_all_channelname_60days_3579.py",
    depends_on_past =False,
    dag=dag,
)

generate_diffurec_dataset_all_channelname_012468 = BashOperator(
    task_id='generate_diffurec_dataset_all_channelname_012468',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_july && /data/dev/miniconda/envs/deepmatch/bin/python 3_generate_diffurec_dataset_all_channelname_012468.py",
    depends_on_past =False,
    dag=dag,
)

generate_diffurec_dataset_all_channelname_3579 = BashOperator(
    task_id='generate_diffurec_dataset_all_channelname_3579',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep && /data/dev/miniconda/envs/deepmatch/bin/python 3_generate_diffurec_dataset_all_channelname_3579.py",
    depends_on_past =False,
    dag=dag,
)

# get_pretrained_text_embedding_jx_lst = BashOperator(
#     task_id='get_pretrained_text_embedding_jx_lst',
#     bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/deepmatch/bin/python get_pretrained_text_embedding.py",
#     depends_on_past =False,
#     dag=dag,
# )
get_pretrained_content_embedding_jx_lst = BashOperator(
    task_id='get_pretrained_content_embedding_jx_lst',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/deepmatch/bin/python get_pretrained_content_embedding.py",
    depends_on_past =False,
    dag=dag,
)
get_pretrained_image_embedding_minicpm = BashOperator(
    task_id='get_pretrained_image_embedding_minicpm',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/deepmatch/bin/python get_pretrained_image_embedding_minicpm.py",
    depends_on_past =False,
    dag=dag,
)

pca_text_image = BashOperator(
    task_id='pca_text_image',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/deepmatch/bin/python pca.py --text_feat --image_feat",
    depends_on_past =False,
    dag=dag,
)

train_diffurec = BashOperator(
    task_id='train_diffurec',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/DiffuRec/src_v2 && /data/dev/miniconda/envs/pt1.0/bin/python main_daily.py --text_feat --image_feat --epochs 15",
    depends_on_past =False,
    dag=dag,
)

get_pretrained_content_embedding_sep = BashOperator(
    task_id='get_pretrained_content_embedding_sep',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep && /data/dev/miniconda/envs/deepmatch/bin/python get_pretrained_content_embedding_sep.py",
    depends_on_past =False,
    dag=dag,
)
get_pretrained_image_embedding_minicpm_sep = BashOperator(
    task_id='get_pretrained_image_embedding_minicpm_sep',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep && /data/dev/miniconda/envs/deepmatch/bin/python get_pretrained_image_embedding_minicpm_sep.py",
    depends_on_past =False,
    dag=dag,
)

pca_text_image_v2 = BashOperator(
    task_id='pca_text_image_v2',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/deepmatch/bin/python pca_v2.py --text_feat --image_feat",
    depends_on_past =False,
    dag=dag,
)

train_llm_diffurec_v2 = BashOperator(
    task_id='train_llm_diffurec_v2',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/DiffuRec/src_v2 && /data/dev/miniconda/envs/pt1.0/bin/python main_daily_llm_v2.py --text_feat --image_feat --epochs 15",
    depends_on_past =False,
    dag=dag,
)
generate_inference_dataset_012468 = BashOperator(
    task_id='generate_inference_dataset_012468',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_july && /data/dev/miniconda/envs/deepmatch/bin/python 4_1_generate_inference_dataset_012468.py",
    depends_on_past =False,
    dag=dag,
)
generate_inference_dataset_3579 = BashOperator(
    task_id='generate_inference_dataset_3579',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep && /data/dev/miniconda/envs/deepmatch/bin/python 4_2_generate_inference_dataset_3579.py",
    depends_on_past =False,
    dag=dag,
)


generate_diffurec_recall_result_2468 = BashOperator(
    task_id='generate_diffurec_recall_result_2468',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_july && /data/dev/miniconda/envs/pt1.0/bin/python 6_1_generate_diffurec_recall_result_2468.py --text_feat --image_feat",
    depends_on_past =False,
    dag=dag,
)

generate_diffurec_recall_result_3579 = BashOperator(
    task_id='generate_diffurec_recall_result_3579',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep && /data/dev/miniconda/envs/pt1.0/bin/python 6_2_generate_diffurec_recall_result_3579.py --text_feat --image_feat",
    depends_on_past =False,
    dag=dag,
)
# generate_diffurec_recall_result_7 = BashOperator(
#     task_id='generate_diffurec_recall_result_7',
#     bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep && /data/dev/miniconda/envs/pt1.0/bin/python 6_3_generate_diffurec_recall_result_7.py --text_feat --image_feat",
#     depends_on_past =False,
#     dag=dag,
# )


# vip model
get_pretrained_content_embedding_vip = BashOperator(
    task_id='get_pretrained_content_embedding_vip',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep/vip && /data/dev/miniconda/envs/deepmatch/bin/python get_pretrained_content_embedding_vip.py",
    depends_on_past =False,
    dag=dag,
)
get_pretrained_image_embedding_minicpm_vip = BashOperator(
    task_id='get_pretrained_image_embedding_minicpm_vip',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep/vip && /data/dev/miniconda/envs/deepmatch/bin/python get_pretrained_image_embedding_minicpm_vip.py",
    depends_on_past =False,
    dag=dag,
)
generate_diffurec_dataset_all_channelname_vip = BashOperator(
    task_id='generate_diffurec_dataset_all_channelname_vip',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep/vip && /data/dev/miniconda/envs/deepmatch/bin/python 3_generate_diffurec_dataset_all_channelname_vip.py",
    depends_on_past =False,
    dag=dag,
)

pca_text_image_vip = BashOperator(
    task_id='pca_text_image_vip',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/deepmatch/bin/python pca_vip.py --text_feat --image_feat",
    depends_on_past =False,
    dag=dag,
)

train_diffurec_vip = BashOperator(
    task_id='train_diffurec_vip',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/DiffuRec/src_v2 && /data/dev/miniconda/envs/pt1.0/bin/python main_daily_vip.py --max_len 10 --text_feat --image_feat --epochs 15",
    depends_on_past =False,
    dag=dag,
)

generate_inference_dataset_vip = BashOperator(
    task_id='generate_inference_dataset_vip',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep/vip && /data/dev/miniconda/envs/deepmatch/bin/python 4_generate_inference_dataset_vip.py",
    depends_on_past =False,
    dag=dag,
)
generate_diffurec_recall_result_vip = BashOperator(
    task_id='generate_diffurec_recall_result_vip',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_sep/vip && /data/dev/miniconda/envs/pt1.0/bin/python 5_generate_diffurec_recall_result_vip.py --text_feat --image_feat",
    depends_on_past =False,
    dag=dag,
)

save_recall_to_hive = SparkSubmitOperator(
    task_id='save_recall_to_hive',
    dag=dag,
    **_config,
    application= '/data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/save_recall_to_hive.py',
    conn_id = "spark_local",
)
# save_recall_to_hive_v2 = SparkSubmitOperator(
#     task_id='save_recall_to_hive_v2',
#     dag=dag,
#     **_config,
#     application= '/data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/save_recall_to_hive_v2.py',
#     conn_id = "spark_local",
# )
# generate_diffurec_test_result = BashOperator(
#     task_id='generate_diffurec_test_result',
#     bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/pt1.0/bin/python 5_generate_diffurec_test_result.py --text_feat",
#     depends_on_past =False,
#     dag=dag,
# )
save_user_click_seq_to_hive = SparkSubmitOperator(
    task_id='save_user_click_seq_to_hive',
    dag=dag,
    **_config,
    application= '/data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/jx_lst_july/save_user_click_seq_to_hive.py',
    conn_id = "spark_local",
)

# modal embedding. 
update_text_embedding = BashOperator(
    task_id='update_text_embedding',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/script && /data/dev/miniconda/envs/deepmatch/bin/python update_text_embedding_v2.py",
    depends_on_past =False,
    dag=dag,
)
update_content_embedding = BashOperator(
    task_id='update_content_embedding',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/script && /data/dev/miniconda/envs/easyrec/bin/python update_content_embedding.py",
    depends_on_past =False,
    dag=dag,
)


# jx feed 
update_jx_feed_media = BashOperator(
    task_id='update_jx_feed_media',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/recsys2/bin/python 0_1_update_feed_media.py",
    depends_on_past =False,
    dag=dag,
)
update_jx_feed_media_v2 = BashOperator(
    task_id='update_jx_feed_media_v2',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/recsys2/bin/python 0_3_update_feed_media_v2.py",
    depends_on_past =False,
    dag=dag,
)
update_content_embedding_jx_feed = BashOperator(
    task_id='update_content_embedding_jx_feed',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/script && /data/dev/miniconda/envs/easyrec/bin/python update_content_embedding_jx_feed.py",
    depends_on_past =False,
    dag=dag,
)
# i2i recommendation
generate_inference_dataset_i2i = BashOperator(
    task_id='generate_inference_dataset_i2i',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/i2i && /data/dev/miniconda/envs/deepmatch/bin/python generate_inference_dataset_i2i.py",
    depends_on_past =False,
    dag=dag,
)
generate_diffurec_recall_result_i2i = BashOperator(
    task_id='generate_diffurec_recall_result_i2i',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/i2i && /data/dev/miniconda/envs/pt1.0/bin/python generate_diffurec_recall_result_i2i.py --text_feat --image_feat",
    depends_on_past =False,
    dag=dag,
)

# save_i2i_recall_to_hive = SparkSubmitOperator(
#     task_id='save_i2i_recall_to_hive',
#     dag=dag,
#     **_config,
#     application= '/data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script/i2i/save_i2i_recall_to_hive.py',
#     conn_id = "spark_local",
# )

# for summary & content embedding
update_jx_lst_media_v2 = BashOperator(
    task_id='update_jx_lst_media_v2',
    bash_command="cd /data/gangyanyuan/terry/terry-recsys-2024/multi-interest/script && /data/dev/miniconda/envs/recsys2/bin/python 0_2_update_jx_lst_media_v2.py",
    depends_on_past =False,
    dag=dag,
)

kill_iPython_job 

update_jx_lst_media >> update_jx_feed_media >> update_text_embedding >> update_jx_lst_media_v2 >> update_content_embedding >> update_jx_feed_media_v2 >> update_content_embedding_jx_feed

# jx lst
update_jx_lst_media >> get_ecpm_aids >> generate_user_click_data  >> update_user_click_all_channelname_60days_012468 >> generate_diffurec_dataset_all_channelname_012468  >> get_pretrained_content_embedding_jx_lst >> get_pretrained_image_embedding_minicpm >> pca_text_image >> train_diffurec >> generate_inference_dataset_012468 >> generate_diffurec_recall_result_2468

generate_user_click_data >> update_user_click_all_channelname_60days_3579 >> generate_diffurec_dataset_all_channelname_3579 >> get_pretrained_content_embedding_sep >> get_pretrained_image_embedding_minicpm_sep >> pca_text_image_v2 >> train_llm_diffurec_v2 >> generate_inference_dataset_3579 >> generate_diffurec_recall_result_3579  >> train_diffurec_vip >> generate_inference_dataset_vip >> generate_diffurec_recall_result_vip

# vip
generate_inference_dataset_3579 >> generate_diffurec_dataset_all_channelname_vip >> get_pretrained_content_embedding_vip >> get_pretrained_image_embedding_minicpm_vip >> pca_text_image_vip


[generate_diffurec_recall_result_2468,generate_diffurec_recall_result_vip] >> save_recall_to_hive >> save_user_click_seq_to_hive

save_recall_to_hive >> generate_inference_dataset_i2i >> generate_diffurec_recall_result_i2i #>> save_i2i_recall_to_hive
# movie
# generate_user_click_data >> update_training_data_movie_tv_60days >> generate_diffurec_dataset_movie >> get_pretrained_text_embedding >> get_pretrained_image_embedding >> train_diffurec_movie_text >> generate_inference_dataset_movie >> generate_diffurec_recall_result_movie_text >> merge_xsimgcl_recall_movie >> save_movie_recall_to_hive