import streamlit as st
import torch
from transformers import AutoConfig, AutoTokenizer, AutoModel
from huggingface_hub import login
import re
import copy
from modeling_st2 import ST2ModelV2, SignalDetector
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.1
            self.signal_classification = True
            self.pretrained_signal_detector = False

    args = Args()

    model = ST2ModelV2(args)


    repo_id = "anamargarida/SpanExtractionWithSignalCls_2"  
    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, beam_search=True):
     
    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"][0]
    end_cause_logits = outputs["end_arg0_logits"][0]
    start_effect_logits = outputs["start_arg1_logits"][0]
    end_effect_logits = outputs["end_arg1_logits"][0]
    start_signal_logits = outputs["start_sig_logits"][0]
    end_signal_logits = outputs["end_sig_logits"][0]


    # Set the first and last token logits to a very low value to ignore them
    start_cause_logits[0] = -1e-4
    end_cause_logits[0] = -1e-4
    start_effect_logits[0] = -1e-4
    end_effect_logits[0] = -1e-4
    start_cause_logits[len(inputs["input_ids"][0]) - 1] = -1e-4
    end_cause_logits[len(inputs["input_ids"][0]) - 1] = -1e-4
    start_effect_logits[len(inputs["input_ids"][0]) - 1] = -1e-4
    end_effect_logits[len(inputs["input_ids"][0]) - 1] = -1e-4

    
    # Beam Search for position selection
    if beam_search:
        indices1, indices2, score1, score2, topk_scores = model.beam_search_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,
            topk=5
        )
        start_cause1, end_cause1, start_effect1, end_effect1 = indices1
        start_cause2, end_cause2, start_effect2, end_effect2 = indices2
    else:
        start_cause1 = start_cause_logits.argmax().item()
        end_cause1 = end_cause_logits.argmax().item()
        start_effect1 = start_effect_logits.argmax().item()
        end_effect1 = end_effect_logits.argmax().item()

        start_cause2, end_cause2, start_effect2, end_effect2 = None, None, None, None

    
    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"])

    if has_signal:
        start_signal_logits[0] = -1e-4
        end_signal_logits[0] = -1e-4
    
        start_signal_logits[len(inputs["input_ids"][0]) - 1] = -1e-4
        end_signal_logits[len(inputs["input_ids"][0]) - 1] = -1e-4
       
        start_signal = start_signal_logits.argmax().item()
        end_signal_logits[:start_signal] = -1e4
        end_signal_logits[start_signal + 5:] = -1e4
        end_signal = end_signal_logits.argmax().item()

    if not has_signal:
        start_signal, end_signal = None, None
        
    
    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_cause1, end_cause1, "#FFD700")  # yellow for cause
    effect_text1 = mark_text_by_position(input_text, start_effect1, end_effect1, "#90EE90")  # green for effect
    
    if start_signal is not None and end_signal is not None:
        signal_text = mark_text_by_position(input_text, start_signal, end_signal, "#FF6347")  # red for signal
    else: 
        signal_text = None
        
    if beam_search:
        cause_text2 = mark_text_by_position(input_text, start_cause2, end_cause2, "#FFD700")
        effect_text2 = mark_text_by_position(input_text, start_effect2, end_effect2, "#90EE90")
    else:
        cause_text2 = None
        effect_text2 = None

    if beam_search: 
        start_cause_probs = torch.softmax(start_cause_logits, dim=-1)
        end_cause_probs = torch.softmax(end_cause_logits, dim=-1)
        start_effect_probs = torch.softmax(start_effect_logits, dim=-1)
        end_effect_probs = torch.softmax(end_effect_logits, dim=-1)

        best_start_cause_score = start_cause_probs[start_cause1].item()
        best_end_cause_score = end_cause_probs[end_cause1].item()
        best_start_effect_score = start_effect_probs[start_effect1].item()
        best_end_effect_score = end_effect_probs[end_effect1].item()

        second_start_cause_score = start_cause_probs[start_cause2].item()
        second_end_cause_score = end_cause_probs[end_cause2].item()
        second_start_effect_score = start_effect_probs[start_effect2].item()
        second_end_effect_score = end_effect_probs[end_effect2].item()

        best_scores = {
                "Start Cause Score": round(best_start_cause_score, 4),
                "End Cause Score": round(best_end_cause_score, 4),
                "Start Effect Score": round(best_start_effect_score, 4),
                "End Effect Score": round(best_end_effect_score, 4),
                "Total Best Score [sum of log-probability scores]": round(score1, 4)
            }

        second_best_scores = {
            "Start Cause Score": round(second_start_cause_score, 4),
            "End Cause Score": round(second_end_cause_score, 4),
            "Start Effect Score": round(second_start_effect_score, 4),
            "End Effect Score": round(second_end_effect_score, 4),
            "Total Second Best Score [sum of log-probability scores]": round(score2, 4)
            }

        top5_scores = sorted(topk_scores.items(), key=lambda x: x[1], reverse=True)[:5]
        top5_scores = [(k, round(v, 4)) for k, v in top5_scores]

    
    else:
        best_scores = {}
        second_best_scores = {}
        top5_scores = {}
        
        
    return cause_text1, effect_text1, signal_text, cause_text2, effect_text2, best_scores, second_best_scores, top5_scores 
    

st.title("Causal Relation Extraction")
input_text = st.text_area("Enter your text here:", height=100)
beam_search = st.radio("Enable Position Selector & Beam Search?", ('Yes', 'No')) == 'Yes'


if st.button("Extract"):
    if input_text:
        cause_text1, effect_text1, signal_text, cause_text2, effect_text2, best_scores, second_best_scores, top5_scores = extract_arguments(input_text, tokenizer, model, beam_search=beam_search)

        # 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: The position of the predicted end token comes before the position of the start token.")
        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_text}", unsafe_allow_html=True)

            if beam_search:

                # Display dictionary in Streamlit
                st.markdown(f"<strong>Best Tuple Scores:</strong>", unsafe_allow_html=True)
                st.json(best_scores)

        
                # Display second relation if beam search is enabled
                st.write("## Relation 2:")
                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_text}", unsafe_allow_html=True)
                
                st.markdown(f"<strong>Second best Tuple Scores:</strong>", unsafe_allow_html=True)
                st.json(second_best_scores)

                st.markdown(f"<strong>top5_scores [sum of log-probability scores]:</strong>", unsafe_allow_html=True)
                # Unpack top 5 scores 
                # first, second, third, fourth, fifth = top_5_scores
                st.json(top5_scores)

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