
from __future__ import division
import csv
import time
from pathlib import Path
from queue import PriorityQueue
import distance
import numpy as np
import pandas as pd
import keras
from jellyfish import damerau_levenshtein_distance
from NeSy4PPM.Data_preprocessing.log_utils import LogData
from NeSy4PPM.Data_preprocessing.utils import Encodings, BK_type
from NeSy4PPM.Prediction.prepare_data import get_beam_size, encode, get_pn_fitness, compliance_checking
from NeSy4PPM.Training.train_common import CustomTransformer
import absl.logging
absl.logging.set_verbosity(absl.logging.ERROR)
from tqdm import tqdm


def run_experiments(log_data: LogData, evaluation_traces: pd.DataFrame, maxlen, encoder: Encodings, char_indices,
                    target_indices_char, char_indices_group, target_indices_char_group, model_file: Path, output_file: Path, bk_model,
                    method_fitness: str=None, resource: bool=False, weight: float=0.0, bk_end:bool=False, beam_size:int=1):
    # Load model, set this to the model generated by train.py
    model = keras.models.load_model(model_file, custom_objects={'CustomTransformer': CustomTransformer})
    class NodePrediction:
        def __init__(self, crop_trace: pd.DataFrame, probability_of=0):
            self.cropped_trace = crop_trace
            self.cropped_line = ''.join(crop_trace[log_data.act_name_key].tolist())
            if resource:
                self.cropped_line_group = ''.join(crop_trace[log_data.res_name_key].tolist())
            if (not resource and self.cropped_line[-1] != '!') or (resource and self.cropped_line_group [-1] != '!' and self.cropped_line[-1] != '!') :
                self.model_input = encode(crop_trace, log_data,encoder, maxlen, char_indices, char_indices_group, resource)
            self.probability_of = probability_of

        def __str__(self):
            return f"Prefix: {self.cropped_line}, prob. {self.probability_of}"

        def __lt__(self, other):
            return -self.probability_of < -other.probability_of
        def get_cropped_trace(self):
            return self.cropped_trace

    class CacheFitness:
        def __init__(self):
            self.trace = {}
            
        def add(self, crop_trace:str, fitness: float):
            self.trace[crop_trace] = fitness
            
        def get(self, crop_trace:str):
            if crop_trace not in self.trace.keys():
                return None
            else:
                return self.trace[crop_trace]

    class CacheTrace:
        def __init__(self):
            self.trace = {}
            
        def add(self, crop_trace:str, output: list):
            self.trace[crop_trace] = output
            
        def get(self, crop_trace:str):
            if crop_trace not in self.trace.keys():
                return None
            else:
                return self.trace[crop_trace]

    def ensure_node_prediction(prefix, log_data, resource):
        if isinstance(prefix, NodePrediction):
            return prefix
        else:
            if isinstance(prefix, dict):
                crop_trace = prefix.get('crop_trace', pd.DataFrame())
                probability_of = prefix.get('probability_of', 0)
                return NodePrediction(crop_trace, probability_of)
            elif isinstance(prefix, str):
                crop_trace = pd.DataFrame({log_data.act_name_key: list(prefix)})
                if resource:
                    crop_trace[log_data.res_name_key] = [''] * len(crop_trace)  # Add empty 'Resource' column
                return NodePrediction(crop_trace, 0)
            elif isinstance(prefix, pd.DataFrame):
                return NodePrediction(prefix, 0)
            else:
                raise ValueError(f"Cannot convert {type(prefix)} to NodePrediction")
            
    def apply_trace(trace, prefix_size, log_data, predict_size, bk_file, target_indices_char,
                    target_indices_char_group,bk_model, method_fitness, resource, weight,bk_end,beam_size):

        if len(trace) > prefix_size:
            start_time = time.time()
            trace_name = trace[log_data.case_name_key].iloc[0]
            trace_prefix = trace.head(prefix_size)

            # Concatenate activities and resources in the trace prefix
            trace_prefix_act = ''.join(trace_prefix[log_data.act_name_key].tolist())
            trace_prefix_res = ''.join(trace_prefix[log_data.res_name_key].tolist()) if resource else None

            act_prefix = ''.join(trace_prefix[log_data.act_name_key].tolist()) + "_" + str(weight)
            res_prefix = ''.join(trace_prefix[log_data.res_name_key].tolist()) if resource else None
            check_prefix = cache_trace.get(act_prefix+""+res_prefix) if resource else cache_trace.get(act_prefix)
            if check_prefix == None:

                trace_ground_truth = trace.tail(trace.shape[0] - prefix_size)
                act_ground_truth = ''.join(trace_ground_truth[log_data.act_name_key].tolist())


                if resource:
                    res_ground_truth = ''.join(trace_ground_truth[log_data.res_name_key].tolist())
                # Initialize queue for beam search, put root of the tree inside
                visited_nodes: PriorityQueue[NodePrediction] = PriorityQueue()
                visited_nodes.put(NodePrediction(trace_prefix))
                frontier_nodes: PriorityQueue[NodePrediction] = PriorityQueue()

                child_node = None
                best_child = None
                is_violated = False
                for i in range(predict_size - prefix_size): #copia al posto di ground truth
                    if visited_nodes.empty():
                        break
                    violated_nodes = {}
                    for k in range(min(beam_size, len(visited_nodes.queue))):
                        child_node = visited_nodes.get()
                        temp_cropped_trace = child_node.cropped_trace
                        if child_node.cropped_line[-1] == "!" or (resource and child_node.cropped_line_group[-1] == "!"):
                            if bk_end:
                                prefix_trace = ensure_node_prediction(temp_cropped_trace, log_data, resource)
                                prefix_trace = prefix_trace.cropped_trace if isinstance(prefix_trace,
                                                                                    NodePrediction) else prefix_trace
                                prefix_trace = prefix_trace[:-1]
                                if resource:
                                    BK_res = compliance_checking(log_data, child_node.cropped_line[-1],child_node.cropped_line_group[-1], \
                                                                 bk_model["model"], prefix_trace,resource) if bk_model["type"] == BK_type.Declare_End else 0
                                else:
                                    if bk_model["type"] == BK_type.Declare_End:
                                        BK_res = compliance_checking(log_data, child_node.cropped_line[-1],
                                                         None,bk_model["model"], prefix_trace,resource)
                                    if bk_model["type"] == BK_type.Procedural_End:
                                        fitness = get_pn_fitness(bk_file, method_fitness, prefix_trace, log_data)[
                                                trace_name]
                                        BK_res = fitness if fitness == 1.0 else np.NINF
                            if k == 0:
                                if bk_end and BK_res == np.NINF:  # violated: continue the search
                                    if not is_violated:
                                        best_child = child_node
                                        is_violated = True
                                    violated_nodes[k] = child_node
                                    continue
                                else: # satisfied or not using BK
                                    visited_nodes = PriorityQueue()
                                    break
                            else:
                                if bk_end:
                                    if BK_res == np.NINF:  # violated: continue the search
                                        if all(violated_nodes.get(i) for i in range(k)) and k == min(beam_size, len(visited_nodes.queue)) -1:
                                            visited_nodes = PriorityQueue()
                                            child_node = best_child #violated_nodes.get(0)
                                            break
                                        else:
                                            violated_nodes[k] = child_node
                                            continue
                                    else: # satisfied
                                        if all(violated_nodes.get(i) for i in range(k)):
                                            is_violated = False
                                            visited_nodes = PriorityQueue()
                                            break
                                        else:
                                            continue
                                else:
                                    continue
                        enc = child_node.model_input
                        if encoder == Encodings.Multi_encoders:
                            y = model.predict([enc["x_act"], enc["x_group"]], verbose=0)
                        else:
                            y = model.predict(enc, verbose=0) # make predictions

                        if resource :
                            y_char = y[0][0]
                            y_group = y[1][0]
                        else:
                            y_char = y[0]
                            y_group = None

                        # combine with fitness
                        if method_fitness is None or weight == 0:
                            fitness = []
                            score = y_char
                            
                        else:
                            fitness = []
                            
                            for f in range(1,len(target_indices_char)+1):
                                
                                if f <= len(target_indices_char):
                                    temp_prediction = target_indices_char[f]               
                                else:
                                    temp_prediction = log_data.new_chars[f-len(target_indices_char)]
                                    
                                predicted_row = temp_cropped_trace.tail(1).copy()
                                predicted_row.loc[:, log_data.act_name_key] = temp_prediction
                                temp_cropped_trace_next= pd.concat([temp_cropped_trace, predicted_row])                        
                                temp_cropped_line_next = ''.join(temp_cropped_trace_next[log_data.act_name_key].tolist()) 
                                                                
                                check_cache = cache_fitness.get(temp_cropped_line_next )
                                if check_cache == None:
                                    fitness_current = get_pn_fitness(bk_file, method_fitness, temp_cropped_trace_next,
                                                                     log_data)[trace_name]
                                    cache_fitness.add(temp_cropped_line_next, fitness_current)
                                else:
                                    fitness_current = check_cache
                                    
                                fitness = fitness + [np.exp(fitness_current) ]
                                                            
                            if sum(fitness) > 0:
                                fitness = [f/sum(fitness) for f in fitness] 
                            else:
                                fitness = np.repeat(1/len(fitness),len(fitness)).tolist()

                            score = [pow(a,1-weight)*pow(b,weight) for a,b in zip(y_char, fitness)]
                        temp_cropped_trace = ensure_node_prediction(temp_cropped_trace, log_data, resource)
                        # put top 3 based on score
                        frontier_nodes = get_beam_size(frontier_nodes, NodePrediction, child_node, bk_model, weight,
                                                       temp_cropped_trace, score, y_group, target_indices_char,
                                                       target_indices_char_group, log_data, resource,
                                                       beam_size=beam_size)
                        
                    visited_nodes = frontier_nodes
                    frontier_nodes = PriorityQueue()
                if is_violated: child_node = best_child
                predicted = child_node.cropped_line[prefix_size:-1]
                if resource:
                    predicted_group = child_node.cropped_line_group[prefix_size:-1]

                output = []
                if len(act_ground_truth) > 0:
                    output.append(trace_name)
                    output.append(prefix_size)
                    output.append('>>'.join([log_data.act_enc_mapping[i] if i != "!" else "" for i in trace_prefix_act]))
                    output.append('>>'.join([log_data.act_enc_mapping[i] if i != "!" else "" for i in act_ground_truth]))
                    output.append('>>'.join([log_data.act_enc_mapping[i] if i != "!" else "" for i in predicted]))
                    dls = 1 - \
                        (damerau_levenshtein_distance(predicted, act_ground_truth) / max(len(predicted), len(act_ground_truth)))
                    if dls < 0:
                        dls = 0
                    output.append(dls)
                    output.append(1 - distance.jaccard(predicted, act_ground_truth))
                    
                    if resource:
                        output.append('>>'.join([log_data.res_enc_mapping[i] if i != "!" else "" for i in trace_prefix_res]))
                        output.append('>>'.join([log_data.res_enc_mapping[i] if i != "!" else "" for i in res_ground_truth]))
                        output.append('>>'.join([log_data.res_enc_mapping[i] if i != "!" else "" for i in predicted_group]))
                        dls_res = 1 - \
                            (damerau_levenshtein_distance(predicted_group, res_ground_truth)
                                / max(len(predicted_group), len(res_ground_truth)))
                        if dls_res < 0:
                            dls_res = 0
                        output.append(dls_res)
                        output.append(1 - distance.jaccard(predicted_group, res_ground_truth))
                    output.append(weight)
                    output.append(time.time() - start_time) # Prediction time
                    output_cache=output.copy()
                    if resource: output_cache.append(predicted_group)
                    output_cache.append(predicted)
                    cache_trace.add(act_prefix+""+res_prefix, output_cache) if resource else cache_trace.add(act_prefix, output_cache)
                    print(output)
            else:
                trace_ground_truth = trace.tail(trace.shape[0] - prefix_size)
                act_ground_truth = ''.join(trace_ground_truth[log_data.act_name_key].tolist())
                output = []
                
                output.append(trace_name)
                output.append(prefix_size)
                output.append(check_prefix[2]) # Prefix
                output.append('>>'.join([log_data.act_enc_mapping[i] if i != "!" else "" for i in act_ground_truth])) #Ground_truth
                predicted = check_prefix[-1] # Predicted acts symbols
                output.append(check_prefix[4]) # Predicted acts labels
                dls = 1 - \
                    (damerau_levenshtein_distance(predicted, act_ground_truth) / max(len(predicted), len(act_ground_truth)))
                if dls < 0:
                    dls = 0
                output.append(dls)
                output.append(1 - distance.jaccard(predicted, act_ground_truth))
                if resource:
                    trace_prefix_res = ''.join(trace_prefix[log_data.res_name_key].tolist())
                    res_ground_truth = ''.join(trace_ground_truth[log_data.res_name_key].tolist())
                    predicted_group = check_prefix[-2]
                    output.append('>>'.join([log_data.res_enc_mapping[i] if i != "!" else "" for i in trace_prefix_res]))
                    output.append('>>'.join([log_data.res_enc_mapping[i] if i != "!" else "" for i in res_ground_truth]))
                    output.append(check_prefix[9]) # Predicted res labels

                    dls_res = 1 - (damerau_levenshtein_distance(predicted_group, res_ground_truth) / max(len(predicted_group), len(res_ground_truth)))
                    dls_res = max(dls_res, 0)  # Ensure non-negative
                    output.append(dls_res)
                    output.append(1 - distance.jaccard(predicted_group, res_ground_truth))
                output.append(check_prefix[12]) if resource else check_prefix[7] # weight
                output.append(check_prefix[13]) if resource else check_prefix[8]  # Prediction time
                print(output)

            if output:
                with open(output_file, 'a', encoding='utf-8', newline='') as csvfile:
                    spamwriter = csv.writer(csvfile, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL)
                    spamwriter.writerow(output)


##############################################################

    with open(output_file, 'w', encoding='utf-8', newline='') as csvfile:
        spamwriter = csv.writer(csvfile, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL)
        # Headers for the new file
        if resource :
            headers = ["Case ID", "Prefix length",
                 "Trace Prefix Act", "Ground truth", "Predicted Acts", "Damerau-Levenshtein Acts", "Jaccard Acts",
                 "Trace Prefix Res", "Ground Truth Resources", "Predicted Resources", "Damerau-Levenshtein Resources",
                 "Jaccard Resources","Weight","Time"]
        else:
            headers = ["Case ID", "Prefix length","Trace Prefix Act", "Ground truth", "Predicted",
                 "Damerau-Levenshtein Acts", "Jaccard Acts", "Weight","Time"]
        spamwriter.writerow(headers)
    cache_fitness = CacheFitness()
    cache_trace = CacheTrace()
    print(headers)
    for prefix_size in range(log_data.evaluation_prefix_start, log_data.evaluation_prefix_end+1):
        evaluation_traces = evaluation_traces.reset_index(drop=True)
        tqdm.pandas()
        evaluation_traces.groupby(log_data.case_name_key, group_keys=False).progress_apply(lambda x: apply_trace(x,
                                                                                                                     prefix_size,
                                                                                                                     log_data,
                                                                                                                     maxlen,
                                                                                                                     encoder,
                                                                                                                     target_indices_char,
                                                                                                                     target_indices_char_group,
                                                                                                                     bk_model,
                                                                                                                     method_fitness,
                                                                                                                     resource,
                                                                                                                     weight,
                                                                                                                     bk_end,beam_size))

