import base64
import json
import logging
import os
import pickle

import pandas as pd
import sqlalchemy
from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance
from airflow.utils.session import provide_session
from sklearn.linear_model import Ridge
from sqlalchemy import func
from sqlalchemy.sql.expression import cast

from fansifter.hooks.enrichment import EnrichmentDependenciesHook
from fansifter.hooks.runtime_estimation import SaveRuntimeEstimationModelsHook
from fansifter.operators.base import FanSifterBaseOperator

DAGS_TO_ESTIMATE = ["enrichments_dag", "algorithms_dag"]


class RuntimeEstimationModelTrainingOperator(FanSifterBaseOperator):
    @provide_session
    def get_training_data(self, session):
        task_data = (
            session.query(TaskInstance)
            .add_columns(
                DagRun.run_id,
                func.encode(DagRun.conf, "base64"),  # cannot read bytea directly, encoding into base64
                cast(
                    func.split_part(DagRun.run_id, ":", 3),  # run_id has uuid:schema:fan_count format
                    sqlalchemy.Integer,
                ),
                TaskInstance.task_id,
                TaskInstance.duration,
            )
            .join(DagRun, TaskInstance.execution_date == DagRun.execution_date)
            .filter(
                TaskInstance.dag_id.in_(DAGS_TO_ESTIMATE),
                TaskInstance.state == "success",
                TaskInstance.task_id.like("Enrich%"),
                # change line below after breaking changes in enrichment speed
                TaskInstance.execution_date >= func.date('2021-10-17'),
                # use data from last monthly (we might have improved the enrichment speed during this period)
                TaskInstance.execution_date >= sqlalchemy.text("current_date - interval '2' month"),
            )
            .all()
        )
        result = []
        for i, row in enumerate(task_data):
            (ti, run_id, conf_base64, fan_count, task_id, duration) = row
            # conf is a pickled json stored in bytes, decoding from base64 and unpickling:
            conf_bytes = base64.b64decode(conf_base64.encode("ascii"))
            conf = pickle.loads(conf_bytes)

            result.append([run_id, conf, conf.get("source_attribute_ids"), fan_count, task_id, duration])

        return result

    def execute(self, context):
        training_data = self.get_training_data()
        # pickle.dump(training_data, open("training_data.pkl", "wb"))

        df = pd.DataFrame(
            training_data, columns=["run_id", "conf", "source_attribute_ids", "fan_count", "task_id", "duration"]
        )
        tasks = sorted([task for task in df["task_id"].unique()])

        dependencies = EnrichmentDependenciesHook().get_dependencies()
        source_attribute_ids_by_enrichment = dependencies["source_attribute_ids"]
        training_result = []

        for task in tasks:
            df_task = df[df["task_id"] == task]
            all_task_source_attribute_ids = sorted(source_attribute_ids_by_enrichment[task.split(".")[0]])

            X = df_task[["fan_count"]].copy()
            y = df_task["duration"]
            for attribute_id in all_task_source_attribute_ids:
                X[attribute_id] = 0
            for i, row in df_task.iterrows():
                present_attributes = list(
                    set(row["source_attribute_ids"]).intersection(set(all_task_source_attribute_ids))
                )
                X.loc[i, present_attributes] = 1

            # todo CV and model selection
            Model = Ridge
            model_params = dict(alpha=1.0)

            model = Model(**model_params)
            model.fit(X.values, y)

            training_result.append(
                [
                    task,
                    all_task_source_attribute_ids,
                    Model.__name__,
                    json.dumps(model_params),
                    list(model.coef_),
                    model.intercept_,
                ]
            )

        training_result_df = pd.DataFrame(
            training_result, columns=["task_id", "source_attribute_ids", "model", "model_params", "coef_", "intercept_"]
        )

        if os.environ.get("IS_LOCAL", "false") == "false":
            SaveRuntimeEstimationModelsHook().save_models(training_result_df)
        else:
            logging.info("Running Airflow locally, no need to save model weights to DEV db...")
            logging.info(training_result_df.to_dict())