import streamlit as st
import torch
from transformers import AutoConfig, AutoTokenizer, AutoModel
from huggingface_hub import login
import re
import copy
from modeling import ST2ModelV2
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

hf_token = st.secrets["HUGGINGFACE_TOKEN"]
login(token=hf_token)


@st.cache_resource
def load_model():
    
    config = AutoConfig.from_pretrained("roberta-large")

    tokenizer = AutoTokenizer.from_pretrained("roberta-large", use_fast=True, add_prefix_space=True)
    
    class Args:
        def __init__(self):
            
            self.dropout = 0.3
            self.signal_classification = True
            self.pretrained_signal_detector = False
                  

    args = Args()

    model = ST2ModelV2(args)

    #repo_id = "anamargarida/MultiSpanTraining_ROBERTA"
    repo_id = "anamargarida/MultiSpanExtractionRoberta2"  
    filename = "model.safetensors"  

    
    model_path = hf_hub_download(repo_id=repo_id, filename=filename)

    state_dict = load_file(model_path)

    model.load_state_dict(state_dict)
    
    return tokenizer, model


tokenizer, model = load_model()

model.eval()  
def extract_arguments(text, tokenizer, model):

    class Args:
        def __init__(self):
            self.signal_classification = True
            self.pretrained_signal_detector = False
        
    args = Args()
   
    inputs = tokenizer(text, return_offsets_mapping=True, return_tensors="pt")
     
    word_ids = inputs.word_ids()
    
    with torch.no_grad():
        outputs = model(**inputs)

    
    start_cause_logits = outputs["start_arg0_logits"].squeeze(0)
    end_cause_logits = outputs["end_arg0_logits"].squeeze(0)
    start_effect_logits = outputs["start_arg1_logits"].squeeze(0)
    end_effect_logits = outputs["end_arg1_logits"].squeeze(0)
    start_signal_logits = outputs["start_sig_logits"].squeeze(0)
    end_signal_logits = outputs["end_sig_logits"].squeeze(0)

    
    start_cause_all = []
    end_cause_all = []
    start_effect_all = []
    end_effect_all = []
    start_signal_all = []
    end_signal_all = []

    num_relations = 3
    #num_relations = 2

    has_signal = 1
    if args.signal_classification:
        if not args.pretrained_signal_detector:
            has_signal = outputs["signal_classification_logits"].argmax().item()
        else:
            has_signal = signal_detector.predict(text=batch["text"])

    
    
    for rel_idx in range(num_relations):
      
      

        start_cause, end_cause, start_effect, end_effect = model.position_selector(
                                start_cause_logits=start_cause_logits,
                                end_cause_logits=end_cause_logits,
                                start_effect_logits=start_effect_logits,
                                end_effect_logits=end_effect_logits,
                                word_ids=word_ids,
                                rel_idx=rel_idx
                            )
    
        start_cause_all.append(start_cause)
        end_cause_all.append(end_cause)
        start_effect_all.append(start_effect)
        end_effect_all.append(end_effect)

        
    
        if has_signal:
            start_signal, end_signal = model.get_predicted_span_signal(start_signal_logits, end_signal_logits, word_ids, rel_idx=rel_idx, max_window=5)

        if not has_signal:
            start_signal, end_signal = None, None
        
        start_signal_all.append(start_signal)
        end_signal_all.append(end_signal)

   

    def scores(rel_idx, start_cause, end_cause, start_effect, end_effect, start_signal, end_signal):
        
        start_cause_probs = torch.softmax(start_cause_logits[rel_idx], dim=-1)
        end_cause_probs = torch.softmax(end_cause_logits[rel_idx], dim=-1)
        start_effect_probs = torch.softmax(start_effect_logits[rel_idx], dim=-1)
        end_effect_probs = torch.softmax(end_effect_logits[rel_idx], dim=-1)
        
        start_signal_probs = torch.softmax(start_signal_logits[rel_idx], dim=-1)
        end_signal_probs = torch.softmax(end_signal_logits[rel_idx], dim=-1)
        
        
        sc = start_cause_probs[start_cause]
        ec = end_cause_probs[end_cause]
        se = start_effect_probs[start_effect]
        ee = end_effect_probs[end_effect]
    
        if has_signal:
            ss = start_signal_probs[start_signal]
            es = end_signal_probs[end_signal]
            total_score = round((sc * ec * se * ee * ss * es).item(), 4)
            ss_score = round(ss.item(), 4)
            es_score = round(es.item(), 4)
        else:
            ss_score = None
            es_score = None
            total_score = None
    
        total_score_nosignal = round((sc * ec * se * ee).item(), 4)
    
        # Log probs
        log_sc = torch.log(sc + 1e-12)
        log_ec = torch.log(ec + 1e-12)
        log_se = torch.log(se + 1e-12)
        log_ee = torch.log(ee + 1e-12)
        sum_log_probs = log_sc + log_ec + log_se + log_ee
    
        return {
            'start_cause_score': round(sc.item(), 4),
            'end_cause_score': round(ec.item(), 4),
            'start_effect_score': round(se.item(), 4),
            'end_effect_score': round(ee.item(), 4),
            'start_signal_score': ss_score,
            'end_signal_score': es_score,
            'relation probability (sc * ec * se * ee * ss * es)': total_score,
            'relation probability (sc * ec * se * ee)': total_score_nosignal,
            'sum of log-probability scores (log_sc + log_ec + log_se + log_ee)': round(sum_log_probs.item(), 4),
        }

    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    token_ids = inputs["input_ids"][0]
    offset_mapping = inputs["offset_mapping"][0].tolist()
   
    def mark_text_by_position(original_text, start_token, end_token, color):
        """Marks text in the original string based on character positions."""
        # Inserts tags into the original text based on token offsets.
        if start_token is not None and end_token is not None:

            #st.write(f"Start: {start_token}, End: {end_token}")
            

            if start_token > end_token:
                return None
                
            if start_token <= end_token:
                
        
                start_idx, end_idx = offset_mapping[start_token][0], offset_mapping[end_token][1]
                 
                if start_idx is not None and end_idx is not None and start_idx < end_idx:
                    #st.write(f"Start_idx: {start_idx}, End_idx: {end_idx}")
                    
                    return (
                        original_text[:start_idx]
                        + f"<mark style='background-color:{color}; padding:2px; border-radius:4px;'>"
                        + original_text[start_idx:end_idx]
                        + "</mark>"
                        + original_text[end_idx:]
                    )
    
        return original_text  
    
        
        
    cause_text1 = mark_text_by_position(input_text, start_cause_all[0], end_cause_all[0], "#FFD700")  # yellow for cause
    effect_text1 = mark_text_by_position(input_text, start_effect_all[0], end_effect_all[0], "#90EE90")  # green for effect

    if has_signal:
        signal_text1 = mark_text_by_position(input_text, start_signal_all[0], end_signal_all[0], "#FF6347")  # red for signal
    else: 
        signal_text1 = None
    
    cause_text2 = mark_text_by_position(input_text, start_cause_all[1], end_cause_all[1], "#FFD700")  # yellow for cause
    effect_text2 = mark_text_by_position(input_text, start_effect_all[1], end_effect_all[1], "#90EE90")  # green for effect

    if has_signal:
        signal_text2 = mark_text_by_position(input_text, start_signal_all[1], end_signal_all[1], "#FF6347")  # red for signal
    else: 
        signal_text2 = None
        
    cause_text3 = mark_text_by_position(input_text, start_cause_all[2], end_cause_all[2], "#FFD700")  # yellow for cause
    effect_text3 = mark_text_by_position(input_text, start_effect_all[2], end_effect_all[2], "#90EE90")  # green for effect

    if has_signal:
        signal_text3 = mark_text_by_position(input_text, start_signal_all[2], end_signal_all[2], "#FF6347")  # red for signal
    else: 
        signal_text3 = None
        
    scores_rel1 = scores(0, start_cause_all[0], end_cause_all[0], start_effect_all[0], end_effect_all[0], start_signal_all[0], end_signal_all[0])
    scores_rel2 = scores(1, start_cause_all[1], end_cause_all[1], start_effect_all[1], end_effect_all[1], start_signal_all[0], end_signal_all[1])
    scores_rel3 = scores(2, start_cause_all[2], end_cause_all[2], start_effect_all[2], end_effect_all[2], start_signal_all[2], end_signal_all[2])
        
    return cause_text1, effect_text1, signal_text1, cause_text2, effect_text2, signal_text2, cause_text3, effect_text3, signal_text3, scores_rel1, scores_rel2, scores_rel3
    #return cause_text1, effect_text1, signal_text1, cause_text2, effect_text2, signal_text2

st.title("Causal Relation Extraction")
st.subheader("Multi-Relation Model - With joint signal classification & Without overlap penalization")
input_text = st.text_area("Enter your text here:", height=100)


if st.button("Extract"):
    if input_text:
        cause_text1, effect_text1, signal_text1, cause_text2, effect_text2, signal_text2, cause_text3, effect_text3, signal_text3, scores_rel1, scores_rel2, scores_rel3 = extract_arguments(input_text, tokenizer, model)
        #cause_text1, effect_text1, signal_text1, cause_text2, effect_text2, signal_text2 = extract_arguments(input_text, tokenizer, model)

        
        log_scores = {
            "Relation_1": scores_rel1['sum of log-probability scores (log_sc + log_ec + log_se + log_ee)'],
            "Relation_2": scores_rel2['sum of log-probability scores (log_sc + log_ec + log_se + log_ee)'],
            "Relation_3": scores_rel3['sum of log-probability scores (log_sc + log_ec + log_se + log_ee)'],
        }
        
        # Sort the relations by score (descending)
        sorted_log_scores = sorted(log_scores.items(), key=lambda x: x[1], reverse=True)
        
        st.write("RANKING - Sum of log-probability scores (log_sc + log_ec + log_se + log_ee):")
        for rel, score in sorted_log_scores:
            st.write(f"{rel}: {score}")
            
        
        # Display first relation
        st.write("## Relation 1:")

        if cause_text1 is None or effect_text1 is None: 
            st.write("The prediction is not correct for at least one span.")
        else:
            st.markdown(f"**Cause:** {cause_text1}", unsafe_allow_html=True)
            st.markdown(f"**Effect:** {effect_text1}", unsafe_allow_html=True)
            st.markdown(f"**Signal:** {signal_text1}", unsafe_allow_html=True)
            st.markdown(f"<strong>Scores:</strong>", unsafe_allow_html=True)
            st.write(scores_rel1)
            
                

        
        # Display second relation 
        st.write("## Relation 2:")
        if cause_text2 is None or effect_text2 is None: 
            st.write("The prediction is not correct for at least one span.")
        else:
            st.markdown(f"**Cause:** {cause_text2}", unsafe_allow_html=True)
            st.markdown(f"**Effect:** {effect_text2}", unsafe_allow_html=True)
            st.markdown(f"**Signal:** {signal_text2}", unsafe_allow_html=True)
            st.markdown(f"<strong>Scores:</strong>", unsafe_allow_html=True)
            st.write(scores_rel2)

        # Display third relation 
        st.write("## Relation 3:")
        if cause_text3 is None or effect_text3 is None: 
            st.write("The prediction is not correct for at least one span.")
        else:
            st.markdown(f"**Cause:** {cause_text3}", unsafe_allow_html=True)
            st.markdown(f"**Effect:** {effect_text3}", unsafe_allow_html=True)
            st.markdown(f"**Signal:** {signal_text3}", unsafe_allow_html=True)
            st.markdown(f"<strong>Scores:</strong>", unsafe_allow_html=True)
            st.write(scores_rel3)
            
      

    else:
        st.warning("Please enter some text before extracting.")