import shutil
import mlflow
import os
import json

from splink.charts import save_offline_chart
import splink

import mlflow
from mlflow.models.signature import infer_signature
from pyspark.sql import SparkSession
from splink.spark.spark_linker import SparkLinker

# ======================================================================
# ======================================================================
def _log_comparison_details(comparison:dict) -> None:
    """
    Traverse a single comparison item from the splink settings JSON and log the sql condition the comparison was made
    under and the m and u values

    :param comparison: dictionary element from a Splink settings object

    """
    output_column_name = comparison['output_column_name']

    comparison_levels = comparison['comparison_levels']
    for _ in comparison_levels:
        sql_condition = _.get('sql_condition')
        m_probability = _.get('m_probability')
        u_probability = _.get('u_probability')
        log_dict = {
            f"output column {output_column_name} compared through condition {sql_condition} m_probability": m_probability,
            f"output column {output_column_name} compared through condition {sql_condition} u_probability": u_probability

        }
        if m_probability: # this means we don't log the null comparison as it doesnt have a m and u value
            mlflow.log_params(log_dict)


# ======================================================================
# ======================================================================
def _log_comparisons(splink_model_json:dict) -> None:
    """
    Traverse the comparisons part of the splink settings and extract the learned values.
    This allows you to easily compare values across different training conditions.

    :param splink_model_json: the settings json from a splink model.

    """
    comparisons = splink_model_json['comparisons']
    for _ in comparisons:
        _log_comparison_details(_)


# ======================================================================
# ======================================================================
def _log_hyperparameters(splink_model_json:dict) -> None:
    """
    Simple method for extracting parameters from the splink settings and logging as parameters in MLFlow
    
    :param splink_model_json: the settings json from a splink model.

    """
    hyper_param_keys = ['link_type', 'probability_two_random_records_match', 'sql_dialect', 'unique_id_column_name',
                        'em_convergence', 'retain_matching_columns', 'blocking_rules_to_generate_predictions']
    for key in hyper_param_keys:
        mlflow.log_param(key, splink_model_json[key])


# ======================================================================
# ======================================================================
def _log_splink_model_json(splink_model_json:dict) -> None:
    """
    Log the splink settings json in its entirery under the name "linker.json"

    :param splink_model_json: the settings json from a splink model.

    """
    mlflow.log_dict(splink_model_json, "linker.json")


# ======================================================================
# ======================================================================
class splinkSparkMLFlowWrapper(mlflow.pyfunc.PythonModel):
    """
    Wrapper class necessary for storing and retrieving the splink model JSON from the MLFlow model repository.
    """

    def load_context(self, context:mlflow.pyfunc.PythonModelContext) -> None:
        """
        Loads linker settings json into context

        :param context: Python Model context instance

        """
        # this simply stores the json with the MLFLow model.
        self.linker_json = context.artifacts['linker_json_path']



# ======================================================================
# ======================================================================
def _save_splink_model_to_mlflow(
    linker:splink.spark.spark_linker.SparkLinker,
    model_name:str) -> None:
    """
    Simple method for logging a splink model
    
    :param linker: Splink model object
    :param model_name: name to save model under

    """
    path = "linker.json"
    if os.path.isfile(path):
        os.remove(path)
    linker.save_settings_to_json(path)
    artifacts = {"linker_json_path": path}
    mlflow.pyfunc.log_model(model_name, python_model=splinkSparkMLFlowWrapper(), artifacts=artifacts)
    os.remove(path)


# ======================================================================
# ======================================================================
def log_splink_model_to_mlflow(
    linker:splink.spark.spark_linker.SparkLinker,
    model_name:str,
    log_parameters_charts:bool=True,
    log_profiling_charts:bool=False,
    params:dict=None,
    metrics:dict=None,
    artifacts:dict=None
    ) -> None:
    """
    Comprehensive logging of Splink attributes, parameters, charts and model JSON to MLFlow to provide easy
    reproducability and tracking during model development and deployment. This will create a new run under which
    to log everything. Additional parameters and metrics can be logged through the params and metrics arguments.

    For details on MLFlow please visit https://mlflow.org/.

    This will log
    - All the m and u values trained along with their training condition
    - All the hyperparameters set during training (blocking rules, link type, sql_dialect etc)

    :param run: an mlflow.start_run() instance
    :param linker: a trained Splink linkage model
    :param model_name: str, the name to log the model under
    :param log_parameters_charts: boolean, whether to log parameter charts or not. Default to True as this is a quick operation.
    :param log_profiling_charts: boolean, whether to log data profiling charts or not. Default to False as this requires running Spark jobs and can take time.
    :param params: dict[str: str]. Dictionary is of param_name :param_value. Optional argument for logging arbitrary parameters
    :param metrics: dict[str: double or int]. Dictionary is of metric_name :metric_value. Optional argument for logging arbitrary metrics
    :param artifacts: dict[str: str]. Dictionary is of artifact_name :artifact_filepath. Optional argument for logging arbitrary artifacts (NB - artifacts must be written to disk before logging).


    """

    splink_model_json = linker._settings_obj.as_dict()

   
    _log_splink_model_json(splink_model_json)
    _log_hyperparameters(splink_model_json)
    #_log_comparisons(splink_model_json)
    #_save_splink_model_to_mlflow(linker, model_name)
    model = SplinkMLFlowWrapper()
    model.set_spark_linker(linker)
    mlflow.pyfunc.log_model(model_name, python_model=model)
    # if log_profiling_charts or log_parameters_charts:
        # _log_linker_charts(linker, log_parameters_charts, log_profiling_charts)
    if params:
        mlflow.log_params(params)
    if metrics:
        mlflow.log_metrics(metrics)
    if artifacts:
        mlflow.log_artifacts(artifacts)



# ======================================================================
# ======================================================================
def _log_linker_charts(
    linker:splink.spark.spark_linker.SparkLinker,
    log_parameters_charts:bool,
    log_profiling_charts:bool
    ) -> None:
    '''
    Log all the non-data related charts to MLFlow

    :param linker: a Splink linker object
    :param log_parameters_charts: boolean, whether to log parameter charts or not
    :param log_profiling_charts: boolean, whether to log data profiling charts or not


    '''

    charts_dict = {}
    if log_parameters_charts:
        weights_chart = linker.match_weights_chart()
        mu_chart = linker.m_u_parameters_chart()
        compare_chart = linker.parameter_estimate_comparisons_chart()
        charts_dict["weights_chart"] = weights_chart
        charts_dict["mu_chart"] = mu_chart
        charts_dict["compare_chart"] = compare_chart

    if log_profiling_charts:
        missingness_chart = linker.missingness_chart()
        # completeness = linker.completeness_chart() # Note to Splink team - this method behaves differently to the others
        unlinkables = linker.unlinkables_chart()
        blocking_rules_chart = linker.cumulative_num_comparisons_from_blocking_rules_chart()
        charts_dict["missingness_chart"] = missingness_chart
       # charts_dict["completeness"] = completeness
        charts_dict["unlinkables"] = unlinkables
        charts_dict["blocking_rules_chart"] = blocking_rules_chart

    for name, chart in charts_dict.items():
        _log_chart(name, chart)


# ======================================================================
# ======================================================================
def _log_chart(chart_name:str, chart):
    '''
    Save a chart to MLFlow. This writes the chart out temporarily for MLFlow to pick up and save.

    :param chart_name: the name the chart will be given in MLFlow
    :param chart: chart object from Splink

    '''
    path = f"{chart_name}.html"
    if os.path.isfile(path):
        os.remove(path)
    save_offline_chart(chart.spec, path)
    mlflow.log_artifact(path)
    os.remove(path)

# ======================================================================
# ======================================================================
def get_model_json(artifact_uri:str) -> dict:
    """

    Returns a Splink model JSON from a model tracked with MLFlow. This can be retrieved either from a run or
    from the MLFlow Model Registry, if and only if the model was logged to MLFlow using the method
    splink.mlflow.log_splink_model_to_mlflow().

    Extracting the linker's JSON specification requires temporarily writing to /tmp on the local file system. All
    written files are deleted before the function returns.

    Returns a splink setting json as a python dict.

    :param artifact_uri: URI pointing to the artifacts, such as "runs:/500cf58bee2b40a4a82861cc31a617b1/my_model.pkl", "models:/my_model/Production", or "s3://my_bucket/my/file.txt".

    """
    temp_file_path = f"/tmp/{artifact_uri.split('/')[1]}/splink_mlflow_artifacts_download"
    os.makedirs(temp_file_path)
    mlflow.artifacts.download_artifacts(
        artifact_uri=artifact_uri
        , dst_path=temp_file_path
    )

    with open(f"{temp_file_path}/artifacts/linker.json", "r") as f:
        linker_json = json.load(f)
    shutil.rmtree(temp_file_path)

    return linker_json

# ======================================================================
# ======================================================================
class SplinkMLFlowWrapper(mlflow.pyfunc.PythonModel):
    """
    MLFlow compatibility wrapper for splink linker.

    """
 
    def __init__(self, **kwargs):
        self.settings = {}
        self.linker = None
        for key, value in kwargs.items():
            setattr(self, key, value)
            
    def __getstate__(self):
        json_result = self.get_settings_object()
        return json_result
            
    def __setstate__(self, json_dict):
        self.settings = json_dict
    
    def load_context(self, context):
        # this simply stores the json with the MLFLow model.
        self.context = context
        
    def clear_context(self):
        # model cannot be pickled, so disposing it
        self.linker.spark = None
        
    def set_should_evaluate(self, flag):
        self.should_evaluate = flag
        
    def spark_linker(self, data):
        spark = SparkSession.builder.getOrCreate()
        self.linker = SparkLinker(data, spark=spark)
        spark = None
        return self.linker
        
    def set_spark_linker(self, linker: SparkLinker):
        self.linker = linker
    
    def set_settings(self, settings):
        self.settings = settings
        if self.linker:
            self.linker.initialise_settings(self.settings)
        
    def get_linker(self):
        return self.linker
    
    def get_settings(self):
        return self.settings
   
    def get_settings_object(self):
        return self.get_linker()._settings_obj.as_dict()
    
 
    def predict(self, context, X):
        """
        Predict labels on provided data
        """
        self.linker = self.spark_linker(X)
        self.linker.initialise_settings(self.settings)
        result = self.linker.predict()
        return result
