import gradio as gr
import cohere
import numpy as np
import os

# Initialize Cohere client using API key from environment variable
co = cohere.Client(os.getenv('COHERE_API_KEY'))

def calculate_relevancy(meta_title, h1_heading, first_paragraph, body, keywords, model_type):
    # Select the appropriate model based on the radio button selection
    model = 'embed-english-v3.0' if model_type == 'English' else 'embed-multilingual-v3.0'
    
    # Generate embeddings for each part
    response_meta_title = co.embed(texts=[meta_title], model=model, input_type="search_document")
    meta_title_embedding = response_meta_title.embeddings[0]
    
    response_h1_heading = co.embed(texts=[h1_heading], model=model, input_type="search_document")
    h1_heading_embedding = response_h1_heading.embeddings[0]
    
    response_first_paragraph = co.embed(texts=[first_paragraph], model=model, input_type="search_document")
    first_paragraph_embedding = response_first_paragraph.embeddings[0]
    
    response_body = co.embed(texts=[body], model=model, input_type="search_document")
    body_embedding = response_body.embeddings[0]
    
    # Generate embeddings for each keyword
    keyword_list = [kw.strip() for kw in keywords.split(",")]
    response_keywords = co.embed(texts=keyword_list, model=model, input_type="search_query")
    keyword_embeddings = response_keywords.embeddings
    
    # Calculate relevancy score (cosine similarity) for each part
    relevancy_scores = []
    for kw, kw_embedding in zip(keyword_list, keyword_embeddings):
        scores = [
            kw,
            np.dot(meta_title_embedding, kw_embedding) / (np.linalg.norm(meta_title_embedding) * np.linalg.norm(kw_embedding)),
            np.dot(h1_heading_embedding, kw_embedding) / (np.linalg.norm(h1_heading_embedding) * np.linalg.norm(kw_embedding)),
            np.dot(first_paragraph_embedding, kw_embedding) / (np.linalg.norm(first_paragraph_embedding) * np.linalg.norm(kw_embedding)),
            np.dot(body_embedding, kw_embedding) / (np.linalg.norm(body_embedding) * np.linalg.norm(kw_embedding))
        ]
        relevancy_scores.append(scores)
    return relevancy_scores
js_func = """
function refresh() {
    const url = new URL(window.location);

    if (url.searchParams.get('__theme') !== 'light') {
        url.searchParams.set('__theme', 'light');
        window.location.href = url.href;
    }
}
"""
# Create Gradio interface
def gradio_interface():
    with gr.Blocks(js=js_func) as demo:
        with gr.Row():
            meta_title_input = gr.Textbox(label="Meta Title", placeholder="Enter Meta Title here")
            h1_heading_input = gr.Textbox(label="H1 Heading", placeholder="Enter H1 Heading here")
        first_paragraph_input = gr.Textbox(label="First Paragraph", placeholder="Enter First Paragraph here")
        body_input = gr.Textbox(label="Body", placeholder="Enter Body here", lines=5)
        keywords_input = gr.Textbox(label="Keywords", placeholder="Enter keywords separated by commas")
        
        model_type = gr.Radio(["English", "Multilingual"], label="Language Model", value="English")
        
        relevancy_button = gr.Button("Calculate Relevancy", variant="primary")
        relevancy_output = gr.Dataframe(headers=["Keyword", "Meta Title", "H1 Heading", "First Paragraph", "Body"])
        
        relevancy_button.click(calculate_relevancy, 
                               inputs=[meta_title_input, h1_heading_input, first_paragraph_input, body_input, keywords_input, model_type], 
                               outputs=relevancy_output)
        
    return demo

# Launch Gradio app
if __name__ == "__main__":
    demo = gradio_interface()
    demo.launch()