from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta
import pandas as pd
import google.auth
from google.cloud import bigquery, storage
from io import BytesIO
import os, sys, pickle, re
import numpy as np
from sklearn.dummy import DummyRegressor
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error, median_absolute_error, root_mean_squared_error
from sklearn.model_selection import KFold
from xgboost import XGBRegressor

# Import model_functions (ensure this file is in the same directory or adjust sys.path)
try:
    from model_functions import train_model, assess, importance, load_data, split, reverse_log_y
    print("Successfully imported model_functions")
except ImportError as e:
    print(f"Error importing model_functions: {e}")
    sys.path.append('/path/to/model_functions_dir')
    from model_functions import train_model, assess, importance, load_data, split, reverse_log_y

# Email libraries
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.base import MIMEBase
from email import encoders

# Define default arguments
default_args = {
    'owner': 'arun_samria_cvshealth_com',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

# Define the DAG
with DAG(
    'tat_model_pipeline',
    default_args=default_args,
    description='A pipeline for TAT model data preparation, training, evaluation, and prediction',
    schedule_interval=None,
    start_date=datetime(2025, 3, 9),
    catchup=False,
) as dag:

    # Initialize GCS and BigQuery clients
    credentials, project = google.auth.default()
    bigquery_client = bigquery.Client(credentials=credentials)
    storage_client = storage.Client()

    # Task 1: Data Preparation
    def run_data_preparation():
        user = 'arun_samria_cvshealth_com'
        name = "TAT_11262024"

        query_ = f"""DROP TABLE IF EXISTS `anbc-pss-dev.anbc_sga_pss_dev.task_list_resp_{user}`"""
        query_job = bigquery_client.query(query_)
        query_job.result()

        query_ = f"""
        CREATE TABLE `anbc-pss-dev.anbc_sga_pss_dev.task_list_resp_{user}`
        OPTIONS (labels=[("owner", "{user}")]) AS
        WITH task_list AS (
            SELECT DISTINCT SBP_WRK_QUE_ID
            FROM `edp-prod-pbmstorage.edp_pbm_spclt_cnsv.V_SPCLT_BP_WRK_QUE`
            WHERE CURR_IND = 'Y'
        ), dataset AS (
            SELECT DISTINCT
                prc.SPCLT_PTNT_GID, prc.PTNT_ID,
                prc.SBP_ENTY_ID, prc.SBP_ENTY_TYP_CD,
                tsk.SBP_WRK_QUE_TASK_BPM_TASK_ID, tsk_map.SBP_WRK_QUE_TASK_NM, wrk_que.SBP_QUE_NM AS QUE_NM, tsk.SBP_QUE_TASK_STUS_CD,
                tsk.SRC_CRTE_TS, tsk.CMPLTN_TS,
                DATETIME_DIFF(tsk.CMPLTN_TS, tsk.SRC_CRTE_TS, SECOND)/86400 AS TAT_DAYS,
                CASE WHEN EXTRACT(DATE FROM tsk.SRC_CRTE_TS) >= '2024-06-01' THEN 'TEST' ELSE 'TRAIN' END AS COHORT
            FROM `edp-prod-pbmstorage.edp_pbm_spclt_cnsv.V_SPCLT_BP_ENTY_BPM_PRCS_MAP` prc
            INNER JOIN `edp-prod-pbmstorage.edp_pbm_spclt_cnsv.V_SPCLT_BP_WRK_QUE_TASK_BPM_TASK` tsk
            INNER JOIN `edp-prod-pbmstorage.edp_pbm_spclt_cnsv.V_SPCLT_BP_WRK_QUE_TASK_MAP` tsk_map
            ON tsk.SBP_WRK_QUE_TASK_BPM_TASK_ID = tsk_map.SBP_WRK_QUE_TASK_ID AND tsk_map.CURR_IND = 'Y'
            INNER JOIN `edp-prod-pbmstorage.edp_pbm_spclt_cnsv.V_SPCLT_BP_WRK_QUE` wrk_que
            ON tsk_map.SBP_WRK_QUE_ID = wrk_que.SBP_WRK_QUE_ID AND wrk_que.CURR_IND = 'Y'
            INNER JOIN task_list WRK_QUE_MAP
            ON wrk_que.SBP_PRNT_WRK_QUE_ID = WRK_QUE_MAP.SBP_WRK_QUE_ID
            WHERE prc.CURR_IND = 'Y'
            AND EXTRACT(DATE FROM tsk.SRC_CRTE_TS) >= DATE('2024-05-01')
            AND EXTRACT(DATE FROM tsk.SRC_CRTE_TS) <= DATE('2024-06-30')
        )
        SELECT * FROM dataset WHERE COHORT = 'TRAIN' AND RAND() < 1.0
        UNION ALL
        SELECT * FROM dataset WHERE COHORT = 'TEST' AND RAND() < 1.0
        """
        query_job = bigquery_client.query(query_)
        query_job.result()

        q1 = f"""
        CREATE OR REPLACE TABLE `anbc-pss-dev.anbc_sga_pss_dev.{name}_Train_Test_Combo`
        OPTIONS (labels=[("owner", "arun_samria_cvshealth_com")]) AS
        WITH A AS (
            SELECT DISTINCT
                mapp.SBP_WRK_QUE_TASK_BPM_TASK_ID, mapp.COHORT, mapp.SPCLT_PTNT_GID, mapp.SBP_ENTY_ID,
                mapp.SRC_CRTE_TS, mapp.SBP_ENTY_TYP_CD, atr.CLNT_ID, atr.DID_ID, atr.DRUG_LIST, atr.EHR_IND,
                atr.MD_NPI, atr.NEEDS_DT, atr.NEWFILL_IND, atr.ORG_ID, atr.PAT_ST_CD, atr.PT_BU_CD, atr.PYR_ID_LIST,
                atr.RX_BU_CD, atr.SAC_LIST, atr.SRC_DEL_FLG, atr.THP_CD, atr.TRGT_DT, atr.VIP_IND,
                RANK() OVER (PARTITION BY SBP_WRK_QUE_TASK_BPM_TASK_ID ORDER BY DATETIME_DIFF(CAST(atr.SRC_ADD_TS AS DATETIME), CAST(mapp.SRC_CRTE_TS AS DATETIME), MINUTE) DESC) AS RANKING
            FROM `anbc-pss-dev.anbc_sga_pss_dev.task_list_resp_{user}` mapp
            INNER JOIN `anbc-pss-dev.anbc_sga_pss_dev.SBP_ACTV_ENTY_ATRB_HIST` atr
            ON SAFE_CAST(mapp.SBP_ENTY_ID AS FLOAT64) = SAFE_CAST(atr.SBP_ENTY_ID AS FLOAT64)
            AND mapp.SBP_ENTY_TYP_CD = atr.SBP_ENTY_TYP_CD AND atr.CURR_IND = 'Y'
        ), B AS (
            SELECT * FROM A WHERE RANKING = 1
        )
        SELECT
            T.SBP_WRK_QUE_TASK_BPM_TASK_ID, T.COHORT, T.SPCLT_PTNT_GID, T.SBP_WRK_QUE_TASK_NM, T.QUE_NM, T.TAT_DAYS,
            B.CLNT_ID, B.DID_ID, B.DRUG_LIST, B.EHR_IND, B.MD_NPI, B.NEEDS_DT, B.NEWFILL_IND, B.ORG_ID, B.PAT_ST_CD,
            B.PT_BU_CD, B.PYR_ID_LIST, B.RX_BU_CD, B.SAC_LIST, B.SRC_DEL_FLG, B.THP_CD, B.TRGT_DT, B.VIP_IND,
            CASE WHEN B.SBP_WRK_QUE_TASK_BPM_TASK_ID IS NULL THEN 'NO_RECORD' ELSE 'RECORD' END AS MATCH_DEBUG
        FROM `anbc-pss-dev.anbc_sga_pss_dev.task_list_resp_{user}` T
        LEFT JOIN B
        ON T.SBP_WRK_QUE_TASK_BPM_TASK_ID = B.SBP_WRK_QUE_TASK_BPM_TASK_ID
        AND T.COHORT = B.COHORT AND T.SPCLT_PTNT_GID = B.SPCLT_PTNT_GID
        WHERE T.TAT_DAYS >= 0 AND CASE WHEN B.SBP_WRK_QUE_TASK_BPM_TASK_ID IS NULL THEN 'NO_RECORD' ELSE 'RECORD' END = 'RECORD'
        """
        query_job_2 = bigquery_client.query(q1)
        query_job_2.result()

        query_job_3 = bigquery_client.query(f"SELECT * FROM `anbc-pss-dev.anbc_sga_pss_dev.{name}_Train_Test_Combo`")
        sbp_data = query_job_3.to_dataframe()

        sbp_data = sbp_data.drop(columns=['SBP_WRK_QUE_TASK_BPM_TASK_ID', 'SPCLT_PTNT_GID', 'MATCH_DEBUG', 'TRGT_DT', 'NEEDS_DT', 'SRC_DEL_FLG']).copy()
        sbp_data['MD_NPI'] = sbp_data['MD_NPI'].astype(str)
        for i in [i for i in sbp_data.columns if i not in ['TAT_DAYS', 'COHORT', 'SBP_WRK_QUE_TASK_NM', 'QUE_NM']]:
            sbp_data[i] = sbp_data[i].map(lambda x: re.sub('[^0-9a-zA-Z|]+', '_', x.strip())).replace("", "UNK")

        bucket = storage_client.bucket('anbc-sga-dev')
        prefix = 'data/pickles/'
        for que in sbp_data['QUE_NM'].unique():
            train_ = sbp_data[(sbp_data['QUE_NM'] == que) & (sbp_data['COHORT'] == 'TRAIN')]
            train_blob_name = f"{prefix}TRAIN_{que.replace(' ', '_')}.pkl"
            train_blob = bucket.blob(train_blob_name)
            with train_blob.open(mode='wb') as f:
                pickle.dump(train_.drop(columns=['QUE_NM', 'COHORT']), f)
            print(f"Uploaded {train_blob_name}")

            test_ = sbp_data[(sbp_data['QUE_NM'] == que) & (sbp_data['COHORT'] == 'TEST')]
            test_blob_name = f"{prefix}OOS_{que.replace(' ', '_')}.pkl"
            test_blob = bucket.blob(test_blob_name)
            with test_blob.open(mode='wb') as f:
                pickle.dump(test_.drop(columns=['QUE_NM', 'COHORT']), f)
            print(f"Uploaded {test_blob_name}")

        print("Data preparation completed")

    # Task 2: Model Training
    def run_model_training():
        storage_client = storage.Client()
        bucket = storage_client.bucket('anbc-sga-dev')
        prefix = 'data/pickles/'
        blobs = storage_client.list_blobs(bucket.name, prefix=prefix)
        tasks = set()
        for blob in blobs:
            blob_name = blob.name
            if blob_name.startswith(prefix) and blob_name.endswith('.pkl'):
                file_name = blob_name[len(prefix):]
                task_part = file_name.replace('TRAIN_', '').replace('OOS_', '').replace('.pkl', '')
                tasks.add(task_part)

        PREFIX = 'projects/specialty_ops/tat_model/EVAL_11262024/'
        res = []
        for Q in tasks:
            avg_desc = f"Average Model {str(date.today()).replace('-', '_')}"
            average_model = DummyRegressor(strategy='mean')
            avg_mod = train_model(
                desc=avg_desc, model=average_model, params={},
                queue=Q, feature_pipe=split, transformy=False,
                splits=10, feat_cols=[], threshold=1, dropna=False
            )
            avg_table, avg_pred, _ = assess(avg_mod, 5)
            avg_ = {
                'feat': avg_mod['features'], 'model': avg_mod['mlm'], 'desc': avg_mod['MODEL'],
                'eval': avg_table, 'pred_truth': avg_pred, 'feat_importance': [],
                'datetime': str(datetime.now()), 'task': Q
            }
            bucket = storage_client.bucket('anbc-sga-data-pss-dev')
            blob = bucket.blob(f"{PREFIX}RESULT_{avg_desc.replace(' ', '_')}_{Q.replace(' ', '_')}_{str(date.today()).replace('-', '_')}.pkl")
            with blob.open(mode='wb') as f:
                pickle.dump(avg_, f)

            xgr = XGBRegressor(random_state=0, n_jobs=30)
            xgr_params = {
                'n_estimators': [500, 100, 10, 5], 'learning_rate': [0.1, 0.001, 0.01, 0.001],
                'subsample': [0.4, 0.8, 0.9, 1], 'colsample_bytree': [0.3, 0.5, 0.8], 'gamma': [0, 1, 5, 10]
            }
            pipe_desc = 'split values using pipes and one hot encoding'
            xgr_y_desc = f"XG Boost Model (Y Transform) - {pipe_desc} Model {str(date.today()).replace('-', '_')}"
            xgb_yt_fit = train_model(
                desc=xgr_y_desc, model=xgr, params=xgr_params,
                queue=Q, feature_pipe=split, transform_y=True,
                splits=10, feat_cols=[], threshold=1, dropna=False
            )
            xgb_yt_table, xgb_yt_pred, _ = assess(xgb_yt_fit, 5)
            xgb_yt_feat = importance(xgb_yt_fit)
            xgb_yt_ = {
                'feat': xgb_yt_fit['features'], 'model': xgb_yt_fit['mlm'], 'desc': xgb_yt_fit['MODEL'],
                'eval': xgb_yt_table, 'pred_truth': xgb_yt_pred, 'feat_importance': xgb_yt_feat,
                'datetime': str(datetime.now()), 'task': Q
            }
            blob = bucket.blob(f"{PREFIX}RESULT_{xgr_y_desc.replace(' ', '_')}_{Q.replace(' ', '_')}_{str(date.today()).replace('-', '_')}.pkl")
            with blob.open(mode='wb') as f:
                pickle.dump(xgb_yt_, f)

        print("Model training completed")

    # Task 3: Model Evaluation
    def run_model_evaluation():
        task_in = 'PRESCRIBER_ENGAGEMENT'
        storage_client = storage.Client()
        bucket = storage_client.bucket('anbc-sga-data-pss-dev')
        PREFIX = 'projects/specialty_ops/tat_model/EVAL_11262024/'
        loaded = {}
        for fblob in storage_client.list_blobs('anbc-sga-data-pss-dev', prefix=PREFIX):
            path = fblob.name
            if task_in in path:
                blob = bucket.blob(path)
                with blob.open(mode='rb') as f:
                    load = pickle.load(f)
                loaded[path] = load

        evaluation = []
        feat_import = []
        for i in loaded.keys():
            obj = loaded[i]
            frame = obj['eval']
            evaluation.append(frame.assign(task=obj['task'], item=i))
            if len(obj['feat_importance']) > 0:
                feat_import.append(obj['feat_importance'].assign(task=obj['task'], item=i))

        evaluation = pd.concat(evaluation, axis=0)
        mae_oos = evaluation.loc['test_mean_absolute_error'].pivot(index='task', columns='MODEL', values='OOS')
        mae_oos.index = ['test_mean_absolute_error']
        r2_oos = evaluation.loc['test_r2'].pivot(index='task', columns='MODEL', values='OOS')
        r2_oos.index = ['test_r2']
        rmse_oos = evaluation.loc['test_root_mean_squared_error'].pivot(index='task', columns='MODEL', values='OOS')
        rmse_oos.index = ['test_root_mean_squared_error']
        eval_df = pd.concat([r2_oos, mae_oos, rmse_oos], axis=0)

        if feat_import:
            feat_df = pd.concat(feat_import, axis=0).reset_index()
            pivot_df = feat_df.pivot(index='index', columns='MODEL', values='IMPORTANCE').fillna(0)
            xgb_model_name = f"XG Boost Model (Y Transform) - split values using pipes and one hot encoding Model {str(date.today()).replace('-', '_')}"
            if xgb_model_name in pivot_df.columns:
                sorted_df = pivot_df.sort_values(by=xgb_model_name, ascending=False)
            else:
                print(f"Model {xgb_model_name} not found. Available models: {pivot_df.columns.tolist()}")
                sorted_df = pivot_df
            feat_import_df = sorted_df.head(49)

        print("Model evaluation completed")
        return eval_df

    # Task 4: Model Prediction (Updated with Excel and Email)
    def run_model_prediction():
        task_in = 'PRESCRIBER_ENGAGEMENT'
        storage_client = storage.Client()
        bucket = storage_client.bucket('anbc-sga-data-pss-dev')
        PREFIX = 'projects/specialty_ops/tat_model/EVAL_11262024/'
        models = {}
        for fblob in storage_client.list_blobs('anbc-sga-data-pss-dev', prefix=PREFIX):
            path = fblob.name
            if task_in in path:
                print(path)
                blob = bucket.blob(path)
                with blob.open(mode='rb') as f:
                    load = pickle.load(f)
                models[path] = load

        # Update model path with current date
        model_path = f"projects/specialty_ops/tat_model/EVAL_11262024/RESULT_XG_Boost_Model_(Y_Transform)_-_split_values_using_pipes_and_one_hot_encoding_Model_{str(date.today()).replace('-', '_')}_PRESCRIBER_ENGAGEMENT_{str(date.today()).replace('-', '_')}.pkl"
        if model_path not in models:
            raise ValueError(f"Model path {model_path} not found in loaded models. Available paths: {list(models.keys())}")
        model_obj = models[model_path]['model']

        # Load and transform data
        bucket = storage_client.bucket('anbc-sga-dev')
        train_blob = bucket.blob(f'data/pickles/TRAIN_{task_in.replace(" ", "_")}.pkl')
        with train_blob.open(mode='rb') as f:
            X_TRN, Y_TRN = load_data(BytesIO(f.read()))
        keep = X_TRN.head(0).columns
        feature_pipeline = split(task_in, keep)
        X_T = feature_pipeline.fit_transform(X_TRN)

        oos_blob = bucket.blob(f'data/pickles/OOS_{task_in.replace(" ", "_")}.pkl')
        with oos_blob.open(mode='rb') as f:
            X_OOS, Y_OOS = load_data(BytesIO(f.read()))
        X_O = feature_pipeline.transform(X_OOS)

        # Predict and evaluate
        TRN_PRED = reverse_log_y(model_obj.predict(X_T))
        OOS_PRED = reverse_log_y(model_obj.predict(X_O))
        OOS_PRED = np.where(OOS_PRED < 0, 0, OOS_PRED)

        test_mlm_cvs_r = r2_score(Y_OOS, OOS_PRED)
        test_mlm_cvs_rmse = root_mean_squared_error(Y_OOS, OOS_PRED)
        test_mlm_cvs_mse = mean_squared_error(Y_OOS, OOS_PRED)
        test_mlm_cvs_mae = mean_absolute_error(Y_OOS, OOS_PRED)
        test_mlm_cvs_medae = median_absolute_error(Y_OOS, OOS_PRED)

        # Step 1: Create a dictionary of the metrics
        metrics = {
            'r2': test_mlm_cvs_r,
            'rmse': test_mlm_cvs_rmse,
            'mse': test_mlm_cvs_mse,
            'mae': test_mlm_cvs_mae,
            'medae': test_mlm_cvs_medae
        }

        # Step 2: Convert metrics to a DataFrame
        metrics_df = pd.DataFrame([metrics])

        # Step 3: Save the DataFrame to an Excel file
        excel_file = f"prediction_metrics_{str(date.today()).replace('-', '_')}.xlsx"
        metrics_df.to_excel(excel_file, index=False)
        print(f"Saved prediction metrics to {excel_file}")

        # Step 4: Send the Excel file via email
        # Email configuration
        sender_email = "your_email@gmail.com"  # Replace with your email
        sender_password = "your_password"  # Replace with your email password or app-specific password
        recipient_email = "arunSamria619@gmail.com"

        # Create the email message
        msg = MIMEMultipart()
        msg['From'] = sender_email
        msg['To'] = recipient_email
        msg['Subject'] = f"Prediction Metrics for {task_in} - {str(date.today())}"

        # Email body
        body = f"Dear Arun,\n\nAttached are the prediction metrics for the task {task_in} generated on {str(date.today())}.\n\nBest regards,\nYour Team"
        msg.attach(MIMEText(body, 'plain'))

        # Attach the Excel file
        with open(excel_file, 'rb') as f:
            part = MIMEBase('application', 'octet-stream')
            part.set_payload(f.read())
            encoders.encode_base64(part)
            part.add_header('Content-Disposition', f"attachment; filename= {excel_file}")
            msg.attach(part)

        # Send the email
        try:
            server = smtplib.SMTP('smtp.gmail.com', 587)
            server.starttls()
            server.login(sender_email, sender_password)
            server.sendmail(sender_email, recipient_email, msg.as_string())
            server.quit()
            print(f"Email with prediction metrics sent to {recipient_email}")
        except Exception as e:
            print(f"Failed to send email: {e}")
            raise

        # Clean up: Optionally delete the local Excel file after sending
        os.remove(excel_file)
        print(f"Deleted local file {excel_file}")

        print("Model prediction completed")
        return metrics

    # Define tasks
    task_data_preparation = PythonOperator(
        task_id='data_preparation',
        python_callable=run_data_preparation,
    )

    task_model_training = PythonOperator(
        task_id='model_training',
        python_callable=run_model_training,
    )

    task_model_evaluation = PythonOperator(
        task_id='model_evaluation',
        python_callable=run_model_evaluation,
    )

    task_model_prediction = PythonOperator(
        task_id='model_prediction',
        python_callable=run_model_prediction,
    )

    # Set task dependencies
    task_data_preparation >> task_model_training >> task_model_evaluation >> task_model_prediction




 X_TRN, Y_TRN = load_data(BytesIO(f.read()))

says error : expected str, bytes or os.pathlike object, not bytesIO