#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jul 22 09:21:54 2025

@author: samd
"""
import bilbystats as bs
import pandas as pd
from collections import defaultdict
from seqeval.metrics import accuracy_score, precision_score, recall_score, f1_score
import numpy as np

def extract_entities_by_sentence(df):
    """
    Extract named entities from NER-tagged data and group by sentence.
    
    Args:
        df: DataFrame with columns ['sentence_id', 'word', 'pos', 'chunk', 'ner']
    
    Returns:
        list: List of dictionaries where each dict contains {entity: entity_type} for each sentence.
              Empty dictionaries are included for sentences with no entities.
    """
    sentence_entities = defaultdict(dict)
    
    # Group by sentence_id
    for sentence_id, group in df.groupby('sentence_id'):
        current_entity = []
        current_entity_type = None
        
        for _, row in group.iterrows():
            ner_tag = row['ner']
            word = row['word']
            
            if ner_tag == 'O':  # Outside any entity
                # If we were building an entity, save it
                if current_entity and current_entity_type:
                    entity_text = ' '.join(current_entity)
                    # Remove B- or I- prefix from entity type
                    clean_entity_type = current_entity_type.split('-')[-1]
                    sentence_entities[sentence_id][entity_text] = clean_entity_type
                
                # Reset for next entity
                current_entity = []
                current_entity_type = None
                
            elif ner_tag.startswith('B-'):  # Beginning of entity
                # If we were building a previous entity, save it first
                if current_entity and current_entity_type:
                    entity_text = ' '.join(current_entity)
                    clean_entity_type = current_entity_type.split('-')[-1]
                    sentence_entities[sentence_id][entity_text] = clean_entity_type
                
                # Start new entity
                current_entity = [word]
                current_entity_type = ner_tag
                
            elif ner_tag.startswith('I-'):  # Inside entity
                # Continue building current entity
                if current_entity and current_entity_type:
                    current_entity.append(word)
                else:
                    # This shouldn't happen in well-formed data, but handle it
                    current_entity = [word]
                    current_entity_type = ner_tag
        
        # Don't forget the last entity if the sentence ends with one
        if current_entity and current_entity_type:
            entity_text = ' '.join(current_entity)
            clean_entity_type = current_entity_type.split('-')[-1]
            sentence_entities[sentence_id][entity_text] = clean_entity_type
    
    # Get all unique sentence_ids from the dataframe to ensure we include sentences with no entities
    all_sentence_ids = df['sentence_id'].unique()
    
    # Create result list ensuring all sentences are represented (even those with no entities)
    result = []
    for sentence_id in sorted(all_sentence_ids):
        result.append(sentence_entities[sentence_id])  # defaultdict will return {} for missing keys
    
    return result

def BI2sentences(filename, sent_id = 'sentence_id', reverse = True):
    df = read_conll_to_dataframe(filename)
    sentences = df.groupby(sent_id)['word'].apply(lambda x: ' '.join(x)).to_dict()
    sentences = list(sentences.values())
    sentence_entities = extract_entities_by_sentence(df)
    
    if reverse:
        sentence_entities = bs.reverse_dicts(sentence_entities)
    
        all_entries = []
        for i in range(len(sentence_entities)):
            entry = sentence_entities[i]
            if 'PER' not in entry:
                entry['PER'] = []

            if 'ORG' not in entry:
                entry['ORG'] = []

            if 'LOC' not in entry:
                entry['LOC'] = []

            if 'MISC' not in entry:
                entry['MISC'] = []

            all_entries.append(entry)
            
        sentence_entities = all_entries
    
    return sentences, sentence_entities
    
def read_conll_to_dataframe(filepath):
    data = []
    sentence_id = 0
    
    with open(filepath, 'r', encoding='utf-8') as file:
        for line in file:
            line = line.strip()
            
            if line.startswith('-DOCSTART-'):
                continue
            
            if not line:
                sentence_id += 1
                continue
            
            parts = line.split()
            if len(parts) >= 4:
                data.append({
                    'sentence_id': sentence_id,
                    'word': parts[0],
                    'pos': parts[1],
                    'chunk': parts[2],
                    'ner': parts[3]
                })
    
    return pd.DataFrame(data)

def make_entity_metrics(id2label):
    """Create a compute_metrics function with id2label captured in closure"""
    
    def compute_entity_metrics(p):
        """Compute accuracy, precision, recall, and F1 for NER"""
        predictions, labels = p
        predictions = np.argmax(predictions, axis=2)
        
        # Remove ignored index (special tokens)
        true_predictions = [
            [id2label[p] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]
        true_labels = [
            [id2label[l] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]
        
        # Per-entity metrics
        entity_metrics = {
            "entity_accuracy": accuracy_score(true_labels, true_predictions),
            "entity_precision": precision_score(true_labels, true_predictions),
            "entity_recall": recall_score(true_labels, true_predictions),
            "entity_f1": f1_score(true_labels, true_predictions),
            }

        # Per-sentence accuracy
        text_accuracy = sum(
            [pred == gold for pred, gold in zip(true_predictions, true_labels)]
            ) / len(true_labels)

        entity_metrics["text_accuracy"] = text_accuracy
        
        return entity_metrics
    
    return compute_entity_metrics

def transform_entities(entities, text, startat = 0, id_counters = None, rm_pun = True, doc_id = None):
    # Track counts of each (word, entity_group) combination
    if not id_counters:
        id_counters = defaultdict(int)
    result = []
    
    for ent in entities:
        #word = ent['word']
        entity_name = text[ent['start']:ent['end']]

        # Skip empty entities (model extraction errors)
        if not entity_name:
            continue

        diff = 0
        if rm_pun:
            check_pun = 1

            while check_pun:
                if len(entity_name) >= 4 and entity_name[-4:] == 'Ltd.':
                    check_pun = 0
                elif entity_name and entity_name[-1] in {',', '.', '?', '!', ':', ';', '\n', '。', '！', '？', '；', '…'}:
                    entity_name = entity_name[:-1]
                    diff += -1
                elif len(entity_name) >= 2 and entity_name[-2:] in {"'s", "'s"}:
                    entity_name = entity_name[:-2]
                    diff += -2
                elif entity_name and ('(' not in entity_name) and entity_name[-1] == ')':
                    entity_name = entity_name[:-1]
                    diff += -1
                elif entity_name and (entity_name.count('"') == 1) and entity_name[-1] == '"':
                    entity_name = entity_name[:-1]
                    diff += -1
                elif entity_name and (entity_name.count("'") == 1) and entity_name[-1] == "'":
                    entity_name = entity_name[:-1]
                    diff += -1
                else:
                    check_pun = 0 
        
        ent_type = ent['entity_group']
        id_counters[(entity_name, ent_type)] += 1
        new_ent = {
            'extracted_entity_text': entity_name,
            'extracted_entity_type': ent_type,
            'score': ent['score'],
            'start': ent['start'] + startat,
            'end': ent['end'] + startat + diff,
            'occurrence_count': id_counters[(entity_name, ent_type)],
            'model': None,
            'timestamp': None,
            'source_document_uuid': doc_id
        }
        result.append(new_ent)
        
    return result, id_counters

def document_entities(set_of_entities, texts, rm_pun = True, doc_id = None):
    out = []
    startat = 0
    id_counters = defaultdict(int)
    
    #texts = bs.rep_texts(texts, '\n', ' ')
    
    for i in range(len(set_of_entities)):
        entities = set_of_entities[i]
        text = texts[i]
        t_entities, id_counters = transform_entities(entities, text, startat, id_counters, rm_pun, doc_id = doc_id)
        out += t_entities
        startat += len(text)

    return out
