import gradio as gr
import pandas as pd
import numpy as np
import os
import base64
from together import Together

def extract_medicines(api_key, image):
    """
    Extract medicine names from a prescription image using Together AI's Llama-Vision-Free model
    """
    # Check if API key is provided
    if not api_key:
        return "Please enter your Together API key."
    
    if image is None:
        return "Please upload an image."
    
    try:
        # Initialize Together client with the provided API key
        client = Together(api_key=api_key)
        
        # Convert image to base64
        with open(image, "rb") as img_file:
            img_data = img_file.read()
            b64_img = base64.b64encode(img_data).decode('utf-8')
        
        # Make API call with base64 encoded image
        response = client.chat.completions.create(
            model="meta-llama/Llama-Vision-Free",
            messages=[
                {
                    "role": "system",
                    "content": "You are an expert in identifying medicine names from prescription images."
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Please extract the names of the medicines only."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{b64_img}"
                            }
                        }
                    ]
                }
            ]
        )
        
        # Extract medicine names from response
        medicine_list = response.choices[0].message.content
        return medicine_list
    
    except Exception as e:
        return f"Error: {str(e)}"

def recommend_medicine(api_key, medicine_name, csv_file=None):
    """
    Use Together API to recommend alternative medicines based on input medicine name
    using data from the provided CSV file with specific column structure.
    It will use AI to find similar medicines even if the exact name isn't in the dataset.
    """
    try:
        # If CSV file is provided, use it; otherwise use default
        if csv_file is not None:
            # Read the uploaded CSV
            if isinstance(csv_file, str):  # Path to default CSV
                df = pd.read_csv(csv_file)
            else:  # Uploaded file
                df = pd.read_csv(csv_file.name)
        else:
            # Use the default medicine_dataset.csv in the current directory
            try:
                df = pd.read_csv("medicine_dataset.csv")
            except FileNotFoundError:
                return "Error: Default medicine_dataset.csv not found. Please upload a CSV file."
        
        # Check if medicine is in the dataset
        medicine_exists = medicine_name in df['name'].values
        
        # Create a helpful context about the dataset to send to the LLM
        dataset_overview = f"The dataset contains {len(df)} medicines with columns for name, substitutes, side effects, uses, chemical class, etc."
        
        # Sample of medicine names to give the model context
        sample_names = df['name'].sample(min(20, len(df))).tolist()
        medicine_sample = f"Sample medicines in the dataset: {', '.join(sample_names)}"
        
        # Extract specific medicine data if available
        medicine_data = None
        medicine_info_str = ""
        if medicine_exists:
            medicine_data = df[df['name'] == medicine_name]
            medicine_info_str = medicine_data.to_string(index=False)
        
        # Create system prompt with dataset context
        system_prompt = f"""You are a pharmaceutical expert system that recommends alternative medicines based on a comprehensive medicine dataset. The user has provided the medicine name "{medicine_name}".
DATASET INFORMATION:
{dataset_overview}
{medicine_sample}
The dataset has the following columns:
- name: Medicine name
- substitute0 through substitute4: Potential substitute medicines
- sideEffect0 through sideEffect41: Possible side effects
- use0 through use4: Medical uses
- Chemical Class: The chemical classification
- Habit Forming: Whether the medicine is habit-forming
- Therapeutic Class: The therapeutic classification
- Action Class: How the medicine works
YOUR TASK:
{"The medicine was found in the dataset with the following information:" if medicine_exists else "The medicine was NOT found in the dataset with an exact match. Your task is to:"}
{medicine_info_str if medicine_exists else "1. Identify what kind of medicine this likely is based on its name (e.g., antibiotics, pain relievers, etc.)"}
{'' if medicine_exists else "2. Look for medicines in the sample list that might be similar or serve similar purposes"}
Please recommend alternative medicines for "{medicine_name}" with the following details for each:
1. Name of the alternative medicine
2. Why it's a good alternative (similar chemical composition, therapeutic use, etc.)
3. Potential side effects to be aware of
4. Usage recommendations
5. Similarity to the original medicine (high, medium, low)
Include at least 3-5 alternatives if possible.
IMPORTANT:
- If the medicine name contains strength or formulation (like "500mg" or "Duo"), focus on finding the base medicine first
- Explain why these alternatives might be suitable replacements
- Include appropriate medical disclaimers
- Format your response clearly with headings for each alternative medicine
"""

        # Initialize Together client with the API key
        client = Together(api_key=api_key)
        
        # Make API call
        response = client.chat.completions.create(
            model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free",
            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": f"Please recommend alternatives for {medicine_name} based on the available information."
                }
            ],
            max_tokens=2000,
            temperature=0.7  # Slightly higher temperature for creative recommendations
        )
        
        # Get the raw response
        recommendation_text = response.choices[0].message.content
        
        # Add disclaimer
        final_response = recommendation_text + "\n\n---\n\n**DISCLAIMER:** This information is for educational purposes only. Always consult with a healthcare professional before making any changes to your medication."
        
        return final_response
    
    except Exception as e:
        return f"Error: {str(e)}"

def send_medicine_to_recommender(api_key, medicine_names, csv_file):
    """
    Takes medicine names extracted from prescription and gets recommendations
    """
    if not medicine_names or medicine_names.startswith("Error") or medicine_names.startswith("Please"):
        return "Please extract valid medicine names first"
    
    # Extract the first medicine name from the list (assuming it's the first line or first item)
    medicine_lines = medicine_names.strip().split('\n')
    if not medicine_lines:
        return "No valid medicine name found in extraction results"
    
    # Get the first medicine name (remove any bullet points or numbers)
    first_medicine = medicine_lines[0]
    # Clean up the medicine name (remove bullets, numbers, etc.)
    first_medicine = first_medicine.lstrip('•-*0123456789. ').strip()
    
    # Check if we have a valid medicine name
    if not first_medicine:
        return "Could not identify a valid medicine name from extraction"
    
    # Call the recommend medicine function with the first extracted medicine
    return recommend_medicine(api_key, first_medicine, csv_file)

def analyze_full_prescription(api_key, medicine_names, csv_file):
    """
    Takes all extracted medicine names and analyzes their interactions and provides comprehensive information
    """
    if not medicine_names or medicine_names.startswith("Error") or medicine_names.startswith("Please"):
        return "Please extract valid medicine names first"
    
    try:
        # Parse the medicine names from the extracted text
        medicine_lines = medicine_names.strip().split('\n')
        cleaned_medicines = []
        
        # Clean up medicine names (remove bullets, numbers, etc.)
        for medicine in medicine_lines:
            cleaned_medicine = medicine.lstrip('•-*0123456789. ').strip()
            if cleaned_medicine:
                cleaned_medicines.append(cleaned_medicine)
        
        if not cleaned_medicines:
            return "No valid medicine names found in extraction"
        
        # Create a prompt for the LLM to analyze the full prescription
        medicines_list = ", ".join(cleaned_medicines)
        
        system_prompt = f"""You are a pharmaceutical expert analyzing a full prescription containing the following medicines: {medicines_list}.
Please provide a comprehensive analysis including:
1. Purpose: The likely medical condition(s) being treated with this combination of medicines
2. Potential interactions: Any known drug interactions between these medicines
3. Side effects: Common side effects to watch for when taking this combination
4. Recommendations: General advice for the patient taking these medicines
5. Questions for the doctor: Important questions the patient should ask their healthcare provider
Base your analysis on pharmacological knowledge about these medicines and their typical uses.
"""

        # Initialize Together client with the API key
        client = Together(api_key=api_key)
        
        # Make API call
        response = client.chat.completions.create(
            model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free",
            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": f"Please analyze this prescription with the following medicines: {medicines_list}"
                }
            ],
            max_tokens=2000,
            temperature=0.3  # Lower temperature for more factual responses
        )
        
        analysis_text = response.choices[0].message.content
        
        # Add disclaimer
        final_response = analysis_text + "\n\n---\n\n**DISCLAIMER:** This analysis is for informational purposes only and should not replace professional medical advice. Always consult with your healthcare provider about your prescription."
        
        return final_response
        
    except Exception as e:
        return f"Error: {str(e)}"

# Custom CSS for styling the application
custom_css = """
:root {
    --primary-color: #3498db;
    --secondary-color: #2ecc71;
    --accent-color: #e74c3c;
    --background-color: #f9f9f9;
    --card-bg: #ffffff;
    --text-color: #333333;
    --light-text: #777777;
    --border-radius: 10px;
    --shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    --hover-shadow: 0 10px 15px rgba(0, 0, 0, 0.15);
}

body {
    font-family: 'Roboto', sans-serif;
    background-color: var(--background-color);
    color: var(--text-color);
}

/* Header Styling */
.app-header {
    background: linear-gradient(135deg, var(--primary-color), var(--secondary-color));
    padding: 20px;
    border-radius: var(--border-radius);
    color: white;
    margin-bottom: 20px;
    box-shadow: var(--shadow);
    text-align: center;
}

.app-header h1 {
    margin: 0;
    font-size: 2.5rem;
    font-weight: 700;
}

.app-header p {
    margin: 10px 0 0;
    font-size: 1.1rem;
    opacity: 0.9;
}

/* General Card Styling */
.card {
    background-color: var(--card-bg);
    border-radius: var(--border-radius);
    padding: 20px;
    margin-bottom: 20px;
    box-shadow: var(--shadow);
    transition: box-shadow 0.3s ease;
}

.card:hover {
    box-shadow: var(--hover-shadow);
}

.card-header {
    border-bottom: 1px solid #eee;
    padding-bottom: 10px;
    margin-bottom: 15px;
    font-weight: 600;
    color: var(--primary-color);
}

/* Button Styling */
.primary-btn button {
    background: linear-gradient(to right, var(--primary-color), #2980b9) !important;
    border: none !important;
    color: white !important;
    padding: 10px 20px !important;
    border-radius: var(--border-radius) !important;
    font-weight: 600 !important;
    transition: transform 0.2s, box-shadow 0.2s !important;
    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1) !important;
}

.primary-btn button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.15) !important;
}

.secondary-btn button {
    background: linear-gradient(to right, var(--secondary-color), #27ae60) !important;
    border: none !important;
    color: white !important;
    padding: 10px 20px !important;
    border-radius: var(--border-radius) !important;
    font-weight: 600 !important;
    transition: transform 0.2s, box-shadow 0.2s !important;
    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1) !important;
}

.secondary-btn button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.15) !important;
}

.accent-btn button {
    background: linear-gradient(to right, var(--accent-color), #c0392b) !important;
    border: none !important;
    color: white !important;
    padding: 10px 20px !important;
    border-radius: var(--border-radius) !important;
    font-weight: 600 !important;
    transition: transform 0.2s, box-shadow 0.2s !important;
    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1) !important;
}

.accent-btn button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.15) !important;
}

/* Input Styling */
.custom-input input, .custom-input textarea {
    border-radius: var(--border-radius) !important;
    border: 1px solid #ddd !important;
    padding: 12px !important;
    transition: border-color 0.3s, box-shadow 0.3s !important;
}

.custom-input input:focus, .custom-input textarea:focus {
    border-color: var(--primary-color) !important;
    box-shadow: 0 0 0 2px rgba(52, 152, 219, 0.2) !important;
}

/* Tab Styling */
.custom-tabs .tabs {
    border-bottom: 2px solid #eee !important;
    margin-top: -100px;
}

.custom-tabs .tab-nav button {
    border: none !important;
    background: transparent !important;
    padding: 10px 20px !important;
    margin: 0 !important;
    color: var(--light-text) !important;
    font-weight: 600 !important;
    transition: color 0.3s !important;
}

.custom-tabs .tab-nav button[data-selected="true"] {
    color: var(--primary-color) !important;
    border-bottom: 3px solid var(--primary-color) !important;
}

/* File Upload Styling */
.custom-file-upload {
    border: 2px dashed #ddd !important;
    border-radius: var(--border-radius) !important;
    padding: 20px !important;
    text-align: center !important;
    transition: border-color 0.3s !important;
}

.custom-file-upload:hover {
    border-color: var(--primary-color) !important;
}

/* Output Styling */
.output-area {
    background-color: #f5f7fa !important;
    border-radius: var(--border-radius) !important;
    padding: 15px !important;
    border-left: 4px solid var(--primary-color) !important;
}

/* Info Box Styling */
.info-box {
    background-color: rgba(52, 152, 219, 0.1) !important;
    border-left: 4px solid var(--primary-color) !important;
    padding: 15px !important;
    border-radius: var(--border-radius) !important;
    margin: 20px 0 !important;
}

/* Animation for loading state */
@keyframes pulse {
    0% { opacity: 1; }
    50% { opacity: 0.6; }
    100% { opacity: 1; }
}

.loading {
    animation: pulse 1.5s infinite ease-in-out;
}

/* Responsive improvements */
@media (max-width: 768px) {
    .app-header h1 {
        font-size: 2rem;
    }
    
    .card {
        padding: 15px;
    }
}

/* Custom markdown styling */
.custom-markdown h1, .custom-markdown h2, .custom-markdown h3 {
    color: var(--primary-color);
}

.custom-markdown a {
    color: var(--secondary-color);
    text-decoration: none;
}

.custom-markdown a:hover {
    text-decoration: underline;
}

.custom-markdown blockquote {
    border-left: 4px solid var(--secondary-color);
    padding-left: 15px;
    color: var(--light-text);
}

.custom-markdown code {
    background-color: #f0f0f0;
    padding: 2px 4px;
    border-radius: 4px;
}

/* Disclaimer styling */
.disclaimer {
    background-color: rgba(231, 76, 60, 0.1);
    border-left: 4px solid var(--accent-color);
    padding: 10px 15px;
    border-radius: var(--border-radius);
    margin-top: 20px;
    font-size: 0.9rem;
}

/* Icon styling */
.icon {
    vertical-align: middle;
    margin-right: 8px;
}

/* Feature box styling */
.feature-box {
    display: flex;
    background-color: var(--card-bg);
    border-radius: var(--border-radius);
    padding: 15px;
    margin: 10px 0;
    box-shadow: var(--shadow);
}

.feature-icon {
    width: 50px;
    height: 50px;
    background-color: rgba(52, 152, 219, 0.2);
    border-radius: 50%;
    display: flex;
    align-items: center;
    justify-content: center;
    margin-right: 15px;
}

.feature-content {
    flex: 1;
}

.feature-content h3 {
    margin-top: 0;
    color: var(--primary-color);
}

/* API key styling */
.api-key-box {
    background-color: rgba(46, 204, 113, 0.1);
    border-radius: var(--border-radius);
    padding: 15px;
    margin-bottom: 20px;
    border: 1px solid rgba(46, 204, 113, 0.3);
}
"""

# Custom HTML components
html_header = """
<div class="app-header">
    <h1>💊 Medicine Assistant</h1>
    <p>AI-powered tools for medicine analysis, alternatives, and prescription insights</p>
</div>
"""

html_api_key_section = """
<div class="api-key-box">
    <h3>🔑 API Access</h3>
    <p>Your Together API key is securely used for this session only and is not stored.</p>
</div>
"""

html_about_section = """
<div class="card">
    <div class="card-header">About Medicine Assistant</div>
    <p>This advanced application combines AI-powered tools to help you understand your prescriptions better:</p>
    
    <div class="feature-box">
        <div class="feature-icon">📋</div>
        <div class="feature-content">
            <h3>Prescription Extraction</h3>
            <p>Upload an image of your prescription and the AI will identify medicine names using advanced computer vision.</p>
        </div>
    </div>
    
    <div class="feature-box">
        <div class="feature-icon">🔄</div>
        <div class="feature-content">
            <h3>Alternative Medicine Finder</h3>
            <p>Get detailed information about alternative medications that might serve similar purposes.</p>
        </div>
    </div>
    
    <div class="feature-box">
        <div class="feature-icon">🔍</div>
        <div class="feature-content">
            <h3>Prescription Analysis</h3>
            <p>Analyze entire prescriptions for potential interactions, uses, and important information.</p>
        </div>
    </div>
    
    <div class="disclaimer">
        <strong>Important:</strong> This application is for informational purposes only. Always consult with a healthcare professional before making any changes to your medication regimen.
    </div>
</div>
"""

html_how_to_use_extractor = """
<div class="info-box">
    <h3>How to use the Prescription Extractor:</h3>
    <ol>
        <li>Enter your Together API key above</li>
        <li>Upload a clear image of your prescription</li>
        <li>Click "Extract Medicines" to identify medicines in the image</li>
        <li>After extraction, you can:
            <ul>
                <li>Get alternative recommendations for the first medicine</li>
                <li>Analyze all medicines for potential interactions and insights</li>
            </ul>
        </li>
    </ol>
    <p><strong>Tip:</strong> For best results, ensure your prescription image is clear, well-lit, and shows all medicine names clearly.</p>
    <p><strong>Tip 2:</strong> Please wait for few seconds for results.</p>
</div>
"""

html_how_to_use_recommender = """
<div class="info-box">
    <h3>How to use the Alternative Recommender:</h3>
    <ol>
        <li>Enter your Together API key (same key used across the application)</li>
        <li>Type the name of a medicine you'd like to find alternatives for</li>
        <li>Click "Get Recommendations" to see detailed information about possible alternatives</li>
    </ol>
    <p><strong>Features:</strong></p>
    <ul>
        <li>Finds alternatives even if the exact medicine isn't in the database</li>
        <li>Analyzes medicine names to determine likely purpose and composition</li>
        <li>Provides detailed information about substitutes, side effects, and usage</li>
    </ul>
</div>
"""

# Create Gradio interface with tabs for all functionalities and custom styling
with gr.Blocks(css=custom_css, title="Medicine Assistant") as app:
    gr.HTML(html_header)
    
    # API key input (shared between tabs)
    with gr.Row():
        with gr.Column():
            gr.HTML(html_api_key_section)
            api_key_input = gr.Textbox(
                label="Together API Key", 
                placeholder="Enter your Together API key here...",
                type="password",
                elem_classes=["custom-input"]
            )
    
    # Create a file input for CSV that can be shared between tabs
    with gr.Row():
        with gr.Column():
            csv_file_input = gr.File(
                label="Upload Medicine CSV (Optional)", 
                file_types=[".csv"],
                type="filepath",
                elem_classes=["custom-file-upload"]
            )
            gr.Markdown("If no CSV is uploaded, the app will use the default 'medicine_dataset.csv' file.")
    
    with gr.Tabs(elem_classes=["custom-tabs"]) as tabs:
        with gr.Tab("Prescription Medicine Extractor"):
            gr.HTML("""<div class="card">
                <div class="card-header">Prescription Medicine Extractor</div>
                <p>Upload a prescription image to extract medicine names using Together AI's Llama-Vision-Free model.</p>
            </div>""")
            
            with gr.Row():
                with gr.Column(scale=1):
                    image_input = gr.Image(
                        type="filepath", 
                        label="Upload Prescription Image",
                        elem_classes=["custom-file-upload"]
                    )
                    extract_btn = gr.Button("Extract Medicines", elem_classes=["primary-btn"])
                    
                with gr.Column(scale=2):
                    extracted_output = gr.Textbox(
                        label="Extracted Medicines", 
                        lines=10,
                        elem_classes=["custom-input", "output-area"]
                    )
            
            with gr.Row():
                with gr.Column(scale=1):
                    with gr.Row():
                        recommend_from_extract_btn = gr.Button(
                            "Get Recommendations (Tab) for First Medicine", 
                            elem_classes=["secondary-btn"]
                        )
                    with gr.Row():
                        analyze_full_btn = gr.Button(
                            "Analyze Full Prescription (Full Analysis Tab)", 
                            elem_classes=["accent-btn"]
                        )
                
                with gr.Column(scale=2):
                    output_tabs = gr.Tabs(elem_classes=["custom-tabs"])
                    with output_tabs:
                        with gr.Tab("Recommendations"):
                            recommendation_from_extract_output = gr.Markdown(elem_classes=["custom-markdown"])
                        with gr.Tab("Full Analysis"):
                            full_analysis_output = gr.Markdown(elem_classes=["custom-markdown"])
            
            gr.HTML(html_how_to_use_extractor)
            
            # Connect the buttons to functions
            extract_btn.click(
                fn=extract_medicines, 
                inputs=[api_key_input, image_input], 
                outputs=extracted_output
            )
            
            recommend_from_extract_btn.click(
                fn=send_medicine_to_recommender,
                inputs=[api_key_input, extracted_output, csv_file_input],
                outputs=recommendation_from_extract_output
            )
            
            analyze_full_btn.click(
                fn=analyze_full_prescription,
                inputs=[api_key_input, extracted_output, csv_file_input],
                outputs=full_analysis_output
            )
            
        with gr.Tab("Medicine Alternative Recommender"):
            gr.HTML("""<div class="card">
                <div class="card-header">Medicine Alternative Recommender</div>
                <p>This tool recommends alternative medicines based on an input medicine name using the Together API.</p>
            </div>""")
            
            with gr.Row():
                with gr.Column():
                    medicine_name = gr.Textbox(
                        label="Medicine Name", 
                        placeholder="Enter a medicine name (e.g., Augmentin 625 Duo)",
                        elem_classes=["custom-input"]
                    )
                    submit_btn = gr.Button("Get Recommendations", elem_classes=["primary-btn"])
                    gr.HTML(html_how_to_use_recommender)
                
                with gr.Column():
                    recommendation_output = gr.Markdown(elem_classes=["custom-markdown"])
            
            submit_btn.click(
                recommend_medicine,
                inputs=[api_key_input, medicine_name, csv_file_input],
                outputs=recommendation_output
            )
    
    gr.HTML(html_about_section)

# Launch the app
if __name__ == "__main__":
    app.launch()