import logging

logger = logging.getLogger('nlu')
from nlu.pipe.pipe_logic import PipeUtils
import pandas as pd
from nlu.pipe.utils.data_conversion_utils import DataConversionUtils


def __predict_standard_spark(pipe, data, output_level, positions, keep_stranger_features, metadata,
                             drop_irrelevant_cols, return_spark_df,get_embeddings):
    # 1. Convert data to Spark DF
    data, stranger_features, output_datatype = DataConversionUtils.to_spark_df(data, pipe.spark, pipe.raw_text_column)

    # 3. Apply Spark Pipeline
    data = pipe.spark_transformer_pipe.transform(data)

    # 4. Convert resulting spark DF into nicer format and by default into pandas.
    if return_spark_df: return data  # Returns RAW  Spark Dataframe result of pipe prediction
    return pipe.pythonify_spark_dataframe(data,
                                          keep_stranger_features=keep_stranger_features,
                                          stranger_features=stranger_features,
                                          output_metadata=metadata,
                                          drop_irrelevant_cols=drop_irrelevant_cols,
                                          positions=positions,
                                          output_level=output_level,
                                          )


def __predict_multi_threaded_light_pipe(pipe, data, output_level, positions, keep_stranger_features, metadata,
                                        drop_irrelevant_cols, ):
    # 1. Try light pipe predcit
    # 2. if vanilla Fails use Vanilla
    # 3.  if vanilla fails raise error
    data, stranger_features, output_datatype = DataConversionUtils.to_pandas_df(data, pipe.raw_text_column)

    # Predict -> Cast to PDF -> Join with original inputs. It does NOT yield EMBEDDINGS.
    data = data.join(pd.DataFrame(pipe.spark_transformer_pipe.fullAnnotate(data.text.values)))

    return pipe.pythonify_spark_dataframe(data,
                                          keep_stranger_features=keep_stranger_features,
                                          stranger_features=stranger_features,
                                          output_metadata=metadata,
                                          drop_irrelevant_cols=drop_irrelevant_cols,
                                          positions=positions,
                                          output_level=output_level,
                                          get_embeddings=False
                                          )


def predict_help(pipe,data,output_level,positions,keep_stranger_features,metadata,multithread,drop_irrelevant_cols,return_spark_df,get_embeddings):
    '''
    Annotates a Pandas Dataframe/Pandas Series/Numpy Array/Spark DataFrame/Python List strings /Python String
    :param data: Data to predict on
    :param output_level: output level, either document/sentence/chunk/token
    :param positions: wether to output indexes that map predictions back to position in origin string
    :param keep_stranger_features: wether to keep columns in the dataframe that are not generated by pandas. I.e. when you s a dataframe with 10 columns and only one of them is named text, the returned dataframe will only contain the text column when set to false
    :param metadata: wether to keep additonal metadata in final df or not like confidiences of every possible class for preidctions.
    :param multithread: Whether to use multithreading based lightpipeline. In some cases, this may cause errors.
    :param drop_irellevant_cols: Wether to drop cols of different output levels, i.e. when predicting token level and dro_irrelevant_cols = True then chunk, sentence and Doc will be dropped
    :param return_spark_df: Prediction results will be returned right after transforming with the Spark NLP pipeline
    :return:
    '''
    if output_level != '': pipe.output_level = output_level
    if output_level == 'sentence' or output_level == 'document': pipe.components = PipeUtils.configure_component_output_levels(
        pipe)
    if output_level in ['token', 'chunk', 'relation']: pipe.components = PipeUtils.configure_component_output_levels(
        pipe, 'document')

    if not pipe.is_fitted:
        if pipe.has_trainable_components:
            pipe.fit(data)
        else:
            pipe.fit()

    # 2. configure Lightpipline usage
    pipe.configure_light_pipe_usage(DataConversionUtils.size_of(data), multithread)
    if return_spark_df:
        try:
            return __predict_standard_spark(pipe, data, output_level, positions, keep_stranger_features, metadata,
                                            drop_irrelevant_cols, return_spark_df)
        except Exception as err:
            logger.warning(f"Predictions Failed={err}")
            pipe.print_exception_err(err)
            raise Exception("Failure to process data with NLU")
    elif not get_embeddings and multithread:
        # Try Multithreaded with Fallback vanilla as option. No Embeddings in this mode
        try:
            return __predict_multi_threaded_light_pipe(pipe, data, output_level, positions, keep_stranger_features,
                                                       metadata, drop_irrelevant_cols)
        except Exception as err:
            logger.warning(
                f"Multithreaded mode with Lightpipeline failed. trying to predict again with non multithreaded mode, err={err}")
            try:
                return __predict_standard_spark(pipe, data, output_level, positions, keep_stranger_features, metadata,
                                                drop_irrelevant_cols, return_spark_df, get_embeddings)
            except Exception as err:
                logger.warning(f"Predictions Failed={err}")
                pipe.print_exception_err(err)
                raise Exception("Failure to process data with NLU")
    else :
        # Standard predict with no fallback
        try:
            return __predict_standard_spark(pipe, data, output_level, positions, keep_stranger_features, metadata,
                                            drop_irrelevant_cols, return_spark_df,get_embeddings)
        except Exception as err:
            logger.warning(f"Predictions Failed={err}")
            pipe.print_exception_err(err)
            raise Exception("Failure to process data with NLU")

    # except Exception as err:
    #
    #     import sys
    #     if multithread == True:
    #         # return pipe.predict(data, output_level=output_level, positions=positions,keep_stranger_features=keep_stranger_features, metadata=metadata, multithread=False,return_spark_df=return_spark_df)
    #         return pipe.predict(data=data, output_level=output_level, positions=positions,
    #                             keep_stranger_features=keep_stranger_features, metadata=metadata, multithread=False,
    #                             drop_irrelevant_cols=drop_irrelevant_cols, return_spark_df=return_spark_df)
    #     else:
    #         pipe.print_exception_err(err)
