import pandas as pd

import numpy as np

from tqdm import tqdm

import transformers

from datasets import Dataset, ClassLabel, Sequence, load_dataset, load_metric

from spacy import displacy

from transformers import (AutoModelForTokenClassification, 

                          AutoTokenizer, 

                          DataCollatorForTokenClassification,

                          pipeline,

                          TrainingArguments, 

                          Trainer,

                          AutoConfig,

                        AutoModelForSequenceClassification,

                        AutoTokenizer,

                        DataCollatorWithPadding,

                        EvalPrediction,

                        Trainer,

                        TrainingArguments,

                        default_data_collator,

                        set_seed,)

assert transformers.__version__ >= "4.11.0"

# from src.utils.parse_data import parse_ast, parse_concept, parse_relation

# ---------------------------------------------------------------------------- #

#                              CONCEPTS DETECTIONS                             #

# ---------------------------------------------------------------------------- #

label_names = ["O", "B-PROBLEM", "I-PROBLEM", "B-TEST", "I-TEST", "B-TREATMENT", "I-TREATMENT"]

id2label = {i: label for i, label in enumerate(label_names)}

label2id = {v: k for k, v in id2label.items()}

model_folder_name = "debru 3la path dyal modèle w7ettuh hna"

model_checkpoint = f"models/{model_folder_name}"

model = AutoModelForTokenClassification.from_pretrained(model_checkpoint, label2id=label2id, id2label=id2label)

tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

effect_ner_model = pipeline(task="ner", model=model, tokenizer=tokenizer)

def detect_concept(raw_text):

    outputs = effect_ner_model(raw_text, aggregation_strategy ="simple")

    entities = []

    params = [{"text": sentence, "ents": entities, "title": None}]

    html = displacy.render(

        params,

        style="ent",

        manual=True,

        # jupyter=True,

        options={

            "colors": {

                "PROBLEM": "#f08080",

                "TEST": "#9bddff",

                "TREATMENT": "#ffdab9",

            },

        },

    )

    return outputs

# ---------------------------------------------------------------------------- #

#                           ASSERTIONS CLASSIFICATION                          #

# ---------------------------------------------------------------------------- #

label_list = ['present',

 'possible',

 'absent',

 'conditional',

 'hypothetical',

 'associated_with_someone_else']

id2label = {i: label for i, label in enumerate(label_list)}

label2id = {v: k for k, v in id2label.items()}

model_name_or_path = "..."

tokenizer = AutoTokenizer.from_pretrained(

    model_name_or_path,

    # cache_dir=cache_dir,

)

model = AutoModelForSequenceClassification.from_pretrained(

    model_name_or_path,

    from_tf=bool(".ckpt" in model_name_or_path),

    # cache_dir=cache_dir,

    label2id=label2id,

    id2label=id2label

)

def detect_assertions(raw_text):

    lines = raw_text.split('\n')

    df = pd.DataFrame({"text": lines, "line_number": range(len(lines))})

    concept_df = pd.DataFrame(detect_concept(raw_text))

    if concept_type == "problem":

        text = df[(df["filename"] == fname) & (df["line_number"] == start_line-1)].text.values[0]

        concept_df.append({"concept_text": concept_text, "text": text, "line_number":start_line})

    concept_df = pd.DataFrame(concept_df)

    df = concept_df[["line_number", "text", "concept_text"]]

    df.rename(columns={"text":"sentence1", "concept_text":"sentence2"}, inplace=True)

    predict_dataset = Dataset.from_pandas(df, preserve_index=False)

    predict_dataset = predict_dataset.map(

                    preprocess_function,

                    batched=True,

                    desc="Running tokenizer on prediction dataset",

                )

    predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict")

    predictions = np.argmax(predictions, axis=1)

    df["prediction"] = [label2ast[label] for label in predictions]

    return df

# ---------------------------------------------------------------------------- #

#                             RELATIONS EXTRACTION                             #

# ---------------------------------------------------------------------------- #

model_folder_name = "......."

model_checkpoint = f"models/{model_folder_name}"

def extract_relations(raw_text):

    # split lines

    lines = raw_text.split('\n')

    df = pd.DataFrame({"text": lines, "line_number": range(len(lines))})

    # add concepts  

    concepts = detect_concept(raw_text)

    rel_df = pd.DataFrame()

    concept_df = pd.DataFrame(concepts)

    test_concept_df = concept_df[concept_df["concept_type"] == "test"]

    problem_concept_df = concept_df[concept_df["concept_type"] == "problem"]

    treatment_concept_df = concept_df[concept_df["concept_type"] == "treatment"]

    # class test --> problem

    test_problem_df = pd.merge(test_concept_df, problem_concept_df, how="inner", on="start_line")

    # class treatment --> problem

    treatment_problem_df = pd.merge(treatment_concept_df, problem_concept_df, how="inner", on="start_line")

    # class problem --> problem

    problem_problem_df = pd.merge(problem_concept_df, problem_concept_df, how="inner", on="start_line")

    problem_problem_df = problem_problem_df[problem_problem_df["concept_text_x"] != problem_problem_df["concept_text_y"]] # TODO: remove duplicates ?

    rel_df = pd.concat([test_problem_df, treatment_problem_df, problem_problem_df], axis=0)

    rel_df = rel_df.sort_values(by=["filename", "start_line"])

    rel_df = rel_df.reset_index(drop=True)

    def preprocess_text(row):

        line =  df[(df["filename"] == row["filename"]) & (df["line_number"] == row["start_line"]-1)]["text"].values[0]

        # line = line.lower()

        line = " ".join(line.split()) # remove multiple spaces

        concept_text_x = "<< "+ " ".join(line.split()[row["start_word_number_x"]:row["end_word_number_x"]+1]) + " >>"

        concept_text_y = "[[ " + " ".join(line.split()[row["start_word_number_y"]:row["end_word_number_y"]+1]) + " ]]"

        start_word_number_x = row["start_word_number_x"]

        end_word_number_x = row["end_word_number_x"]

        start_word_number_y = row["start_word_number_y"]

        end_word_number_y = row["end_word_number_y"]

        if row["start_word_number_x"] > row["start_word_number_y"]:

            concept_text_x, concept_text_y = concept_text_y, concept_text_x

            start_word_number_x, start_word_number_y = start_word_number_y, start_word_number_x

            end_word_number_x, end_word_number_y = end_word_number_y, end_word_number_x

        text = " ".join(line.split()[: start_word_number_x] + [concept_text_x] + line.split()[end_word_number_x+1: start_word_number_y] + [concept_text_y] + line.split()[end_word_number_y+1:])

        row["text"] = text

        return row

    predict_df = rel_df.apply(preprocess_text, axis=1)

    predict_dataset = Dataset.from_pandas(predict_df, preserve_index=False)

    # Preprocessing the dataset

    # Padding strategy

    def preprocess_function(examples):

        # Tokenize the texts

        return tokenizer(

            examples["text"],

            padding=False, # We will pad later, dynamically at batch creation, to the max sequence length in each batch

            truncation=True,

        )

    predict_dataset = predict_dataset.map(

            preprocess_function,

            batched=True,

            desc="Running tokenizer on prediction dataset",

        ) 

    trainer = Trainer(

        model=model,

        tokenizer=tokenizer,

        data_collator=DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8),

    )

    predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict")

    predictions = np.argmax(predictions, axis=1)

    rel_df["prediction"] = [id2label[label] for label in predictions]

    rel_df

    pred_relations = []

    for i, row in tqdm(rel_df.iterrows()):

        filename = row["filename"]

        concept_text_x = row["concept_text_x"]

        concept_text_y = row["concept_text_y"]

        concept_type_x = row["concept_type_x"]

        concept_type_y = row["concept_type_y"]

        start_word_number_x = row["start_word_number_x"]

        end_word_number_x = row["end_word_number_x"]

        start_word_number_y = row["start_word_number_y"]

        end_word_number_y = row["end_word_number_y"]

        line_number = row["start_line"]

        prediction = row["prediction"]

        if prediction != "Other":

            pred_relations.append({"concept_text_x":concept_text_x, "concept_text_y":concept_text_y, "concept_type_x":concept_type_x, "concept_type_y":concept_type_y, "start_word_number_x":start_word_number_x, "end_word_number_x":end_word_number_x, "start_word_number_y":start_word_number_y, "end_word_number_y":end_word_number_y, "line_number":line_number, "filename":filename, "prediction":prediction})