# main.py

import os
import sys
import streamlit as st

from crewai import Task, Crew
from crewai.process import Process

# Update the import to fix deprecation warning
from langchain_community.chat_models import ChatOpenAI

from agents import create_agents
from stream_to_expander import StreamToExpander

# Set the OpenAI model name
os.environ["OPENAI_API_KEY"] = "L5FBmIV3asw5aj21LNq7F_LC3isEA"

OPENAI_MODEL_NAME = "gpt-3.5-turbo"

# Ensure the OpenAI API key is set
if "OPENAI_API_KEY" not in os.environ or os.environ["OPENAI_API_KEY"] == "your_openai_api_key_here":
    st.error("Please set your OpenAI API key as an environment variable.")
    st.stop()

# Setup Manager LLM
manager_llm = ChatOpenAI(model_name=OPENAI_MODEL_NAME, temperature=0)

# Streamlit App
st.set_page_config(page_title="Molecule Design Assistant", page_icon="🧪", layout="wide")

# Add GIF at the heading
st.markdown(
    """
    <div style='text-align: center;'>
        <img src='https://i.giphy.com/media/v1.Y2lkPTc5MGI3NjExaWNobDVjMjJpNTc1a2J6dm5oZW40eTdlbWMxcnN1NHZzOW43NWF1MiZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/VcAt04a901woawqLv2/giphy.gif' alt='GIF Image' width='200'>
    </div>
    """,
    unsafe_allow_html=True,
)

st.markdown("<h1 style='text-align: center;'>🧪 Molecule Design Assistant</h1>", unsafe_allow_html=True)
st.markdown("<h3 style='text-align: center;'>Interact with AI agents to design and optimize molecules.</h3>", unsafe_allow_html=True)

# Initialize session state variables
if "generated" not in st.session_state:
    st.session_state["generated"] = []
if "past" not in st.session_state:
    st.session_state["past"] = []
if "input_text" not in st.session_state:
    st.session_state["input_text"] = ""
if "results" not in st.session_state:
    st.session_state["results"] = []

# Tabs for main app and results table
tab1, tab2 = st.tabs(["Main App", "Results Table"])

with tab1:
    # Create columns for layout
    col1, col2 = st.columns([3, 1])  # Adjust the width ratio as needed

    with col2:
        st.header("User Input")
        # Get user input
        def get_text():
            input_text = st.text_area(
                label="Your Message",
                placeholder="Enter your request here...",
                value=st.session_state["input_text"],
                key="input",
                height=100,
            )
            return input_text

        user_input = get_text()

        submit_button = st.button(label="Submit", key="submit_btn")

        if submit_button:
            if user_input:
                st.session_state["past"].append(user_input)
                st.session_state["input_text"] = ""  # Clear input after submission

    # Create a placeholder for the agent outputs
    with col1:
        st.header("Agent Outputs")
        # Create an expander for agent outputs
        agent_expander = st.expander("💬 Agent Conversations", expanded=False)
        # We can create a placeholder inside the expander
        # output_placeholder = agent_expander.empty()

    # Display the conversation history in the right column
    with col2:
        if st.session_state["generated"]:
            st.subheader("Conversation History")
            for i in range(len(st.session_state["generated"])):
                with st.container():
                    st.markdown(f"**You:** {st.session_state['past'][i]}")
                    st.markdown(f"**Assistant:** {st.session_state['generated'][i]}")
                    st.markdown("---")

    # Function to process user request
    def process_user_request(input_text: str):
        try:
            # Redirect stdout to StreamToExpander
            sys.stdout = StreamToExpander(agent_expander)

            # Create agents with verbose output
            agents = create_agents(verbose=True)

            # Create Crew with Hierarchical Process and Memory
            crew = Crew(
                agents=agents,
                process=Process.hierarchical,
                manager_llm=manager_llm,
                memory=True,
                planning=True,
                verbose=True,
            )

            task = Task(
                description=input_text,
                expected_output="Provide the appropriate output based on the user's request.",
                agent=None,  # In hierarchical process, tasks are not pre-assigned
            )
            crew.tasks = [task]
            result = crew.kickoff()

            # Reset stdout
            sys.stdout = sys.__stdout__

            return result, agents  # Return both result and agents
        except Exception as e:
            st.error(f"An error occurred during processing: {e}")
            # Reset stdout
            sys.stdout = sys.__stdout__
            return None, None

    # Process user input if available
    if "past" in st.session_state and len(st.session_state["past"]) > len(st.session_state["generated"]):
        user_input = st.session_state["past"][-1]
        with st.spinner("🤖 **Agents at work...**"):
            output, agents = process_user_request(user_input)
        if output:
            # Store the output
            st.session_state.generated.append(output.raw if hasattr(output, 'raw') else str(output))
            # Display the final answer
            with col1:
                st.subheader("Final Answer")
                st.markdown(output.raw if hasattr(output, 'raw') else str(output))
            # Save to results table
            st.session_state["results"].append({
                "Input": user_input,
                "Agents Used": ", ".join([agent.role for agent in agents]),
                "Output": output.raw if hasattr(output, 'raw') else str(output),
            })

with tab2:
    st.header("Results Table")
    if st.session_state["results"]:
        import pandas as pd
        df = pd.DataFrame(st.session_state["results"])
        # Convert Output to string to avoid serialization issues
        df['Output'] = df['Output'].astype(str)
        st.table(df)
    else:
        st.write("No results yet.")

# Custom CSS for better UI
st.markdown(
    """
    <style>
    div.stButton > button:first-child {
        width: 100%;
        background-color: #2196F3 !important;  /* Changed to blue */
        color: white !important;
        border: none;
        padding: 10px 20px;
        text-align: center;
        font-size: 16px;
    }
    .stTextArea textarea {
        height: 100px !important;
    }
    </style>
    """,
    unsafe_allow_html=True,
)


# agents.py

from crewai import Agent
from tools.text_to_molecule_tool import TextToMoleculeTool
from tools.molecule_to_text_tool import MoleculeToTextTool
from tools.valid_molecule_tool import ValidMoleculeTool
from tools.molecule_optimizer_tool import MoleculeOptimizerTool
from tools.visualization_tool import VisualizationTool
from tools.rdkit_tools import CalculateLogPTool, CalculateMolWtTool, LipinskiRuleOfFiveTool, SimilarityMapTool

def create_agents(verbose=False):
    # Agent 1: Chemistry Expert with RDKit tools
    chemistry_expert = Agent(
        role="Chemistry Expert",
        goal="Provide chemistry knowledge and visualize molecules",
        backstory="""You are an experienced chemist with expertise in chemical analysis and visualization.
        Always include the outputs from any tools you use in your final answer.
        When you generate a molecule visualization, mention that and ensure it will be displayed to the user.
        Present calculated properties results clearly in the final answer.
        """,
        tools=[
            ValidMoleculeTool(),
            VisualizationTool(),
            CalculateLogPTool(),
            CalculateMolWtTool(),
            LipinskiRuleOfFiveTool(),
            SimilarityMapTool(),
        ],
        
        allow_delegation=True,
        verbose=verbose,
    )

    # Agent 2: Molecule Generator and Descriptor
    molecule_generator = Agent(
        role="Molecule Generator",
        goal="Generate molecules from descriptions and describe molecules",
        backstory="""You specialize in converting text descriptions to molecules and vice versa.""",
        tools=[TextToMoleculeTool(), MoleculeToTextTool()],
        allow_delegation=True,
        verbose=verbose,
    )

    # Agent 3: Molecule Editor
    molecule_editor = Agent(
        role="Molecule Editor",
        goal="Suggest molecule edits and alternatives",
        backstory="""You are skilled in suggesting modifications to molecular structures.""",
        tools=[MoleculeOptimizerTool()],
        allow_delegation=True,
        verbose=verbose,
    )

    # Additional agents (Agent 4 and Agent 5) can be added similarly when their tools are implemented

    return [chemistry_expert, molecule_generator, molecule_editor]



# tools/visualization_tool.py

from crewai_tools import BaseTool
from pydantic import Field
from rdkit import Chem
from rdkit.Chem import Draw
import streamlit as st
import base64
from io import BytesIO

class VisualizationTool(BaseTool):
    name: str = Field(default="Molecule Visualization Tool")
    description: str = Field(default="Generates an image of the molecule from a SMILES string")

    def _run(self, smiles: str) -> str:
        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return "Invalid molecule for visualization."

            # Generate the image
            img = Draw.MolToImage(mol, size=(300, 300))

            # Convert image to base64 string
            buffered = BytesIO()
            img.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()

            # Return an HTML image tag
            html_img = f'<img src="data:image/png;base64,{img_str}" alt="Molecule Image"/>'

            # Since Streamlit cannot render HTML in markdown, we use components
            st.markdown(html_img, unsafe_allow_html=True)

            return "Molecule visualization displayed."
        except Exception as e:
            print(f"Error in molecule visualization: {str(e)}")
            return f"Error in molecule visualization: {str(e)}"

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"


# tools/valid_molecule_tool.py

from crewai_tools import BaseTool
from pydantic import Field
from rdkit import Chem

class ValidMoleculeTool(BaseTool):
    name: str = Field(default="Valid Molecule Checker")
    description: str = Field(default="Checks if a SMILES string is a valid molecule")

    def _run(self, smiles: str) -> str:
        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return "Invalid molecule"
            return "Valid molecule"
        except Exception as e:
            print(f"Error validating molecule: {str(e)}")
            return f"Error validating molecule: {str(e)}"

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"


# tools/text_to_molecule_tool.py

from crewai_tools import BaseTool
from transformers import T5Tokenizer, T5ForConditionalGeneration
from pydantic import Field
import torch

class TextToMoleculeTool(BaseTool):
    name: str = Field(default="Text to Molecule Converter")
    description: str = Field(
        default="Converts text descriptions to SMILES molecular representations using MolT5 model"
    )

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.tokenizer = T5Tokenizer.from_pretrained(
            "laituan245/molt5-large-caption2smiles", model_max_length=512
        )
        self.model = T5ForConditionalGeneration.from_pretrained(
            "laituan245/molt5-large-caption2smiles"
        )
        if torch.cuda.is_available():
            self.model = self.model.to("cuda")

    def _run(self, text: str) -> str:
        try:
            input_ids = self.tokenizer(text, return_tensors="pt").input_ids
            if torch.cuda.is_available():
                input_ids = input_ids.to("cuda")

            outputs = self.model.generate(
                input_ids, num_beams=5, max_length=512, temperature=0.7, do_sample=True
            )
            smiles = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

            return smiles
        except Exception as e:
            print(f"Error in text to molecule conversion: {str(e)}")
            return f"Error in text to molecule conversion: {str(e)}"

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"


# tools/rdkit_tools.py

from crewai_tools import BaseTool
from pydantic import Field
from rdkit import Chem
from rdkit.Chem import Descriptors, Draw, rdMolDescriptors, AllChem
from rdkit.Chem.Draw import SimilarityMaps
import streamlit as st
import base64
from io import BytesIO

class CalculateLogPTool(BaseTool):
    name: str = Field(default="Calculate LogP")
    description: str = Field(default="Calculates the LogP of the compound.")

    def _run(self, smiles: str) -> str:
        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return "Invalid SMILES string"
            logp = Descriptors.MolLogP(mol)
            return f"LogP: {logp}"
        except Exception as e:
            print(f"Error calculating LogP: {str(e)}")
            return f"Error calculating LogP: {str(e)}"

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"

class CalculateMolWtTool(BaseTool):
    name: str = Field(default="Calculate Molecular Weight")
    description: str = Field(default="Calculates the molecular weight of the compound.")

    def _run(self, smiles: str) -> str:
        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return "Invalid SMILES string"
            mol_wt = Descriptors.ExactMolWt(mol)
            return f"Molecular Weight: {mol_wt}"
        except Exception as e:
            print(f"Error calculating Molecular Weight: {str(e)}")
            return f"Error calculating Molecular Weight: {str(e)}"

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"

class LipinskiRuleOfFiveTool(BaseTool):
    name: str = Field(default="Lipinski Rule of Five")
    description: str = Field(default="Evaluates Lipinski's Rule of Five for druglikeness.")

    def _run(self, smiles: str) -> str:
        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return "Invalid SMILES string"

            MW = Descriptors.MolWt(mol)
            HBA = rdMolDescriptors.CalcNumHBA(mol)
            HBD = rdMolDescriptors.CalcNumHBD(mol)
            LogP = Descriptors.MolLogP(mol)

            conditions = [MW <= 500, HBA <= 10, HBD <= 5, LogP <= 5]
            violations = 4 - conditions.count(True)

            result = f"Lipinski Rule of Five:\n"
            result += f" - Molecular Weight: {MW} Da\n"
            result += f" - Hydrogen Bond Acceptors: {HBA}\n"
            result += f" - Hydrogen Bond Donors: {HBD}\n"
            result += f" - LogP: {LogP}\n"
            result += f" - Violations: {violations}\n"
            result += "Compound is likely to be orally active." if violations <= 1 else "Compound is unlikely to be orally active."

            return result
        except Exception as e:
            print(f"Error evaluating Lipinski's Rule: {str(e)}")
            return f"Error evaluating Lipinski's Rule: {str(e)}"

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"

class SimilarityMapTool(BaseTool):
    name: str = Field(default="Generate Similarity Map")
    description: str = Field(default="Generates a similarity map between two molecules.")

    def _run(self, smiles1: str, smiles2: str) -> str:
        try:
            mol1 = Chem.MolFromSmiles(smiles1)
            mol2 = Chem.MolFromSmiles(smiles2)
            if mol1 is None or mol2 is None:
                return "Invalid SMILES string(s)"

            # Generate the similarity map
            fig = SimilarityMaps.GetSimilarityMapForFingerprint(mol1, mol2, SimilarityMaps.GetMorganFingerprint)

            # Convert figure to image
            buf = BytesIO()
            fig.savefig(buf, format='png')
            buf.seek(0)
            img_str = base64.b64encode(buf.getvalue()).decode()

            # Display image in Streamlit
            html_img = f'<img src="data:image/png;base64,{img_str}" alt="Similarity Map"/>'
            st.markdown(html_img, unsafe_allow_html=True)

            return "Similarity map displayed."
        except Exception as e:
            print(f"Error generating similarity map: {str(e)}")
            return f"Error generating similarity map: {str(e)}"

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"


# tools/molecule_to_text_tool.py

from crewai_tools import BaseTool
from transformers import T5Tokenizer, T5ForConditionalGeneration
from pydantic import Field
import torch

class MoleculeToTextTool(BaseTool):
    name: str = Field(default="Molecule to Text Converter")
    description: str = Field(
        default="Converts SMILES representations to text descriptions using MolT5 model"
    )

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.tokenizer = T5Tokenizer.from_pretrained(
            "laituan245/molt5-large-smiles2caption", model_max_length=512
        )
        self.model = T5ForConditionalGeneration.from_pretrained(
            "laituan245/molt5-large-smiles2caption"
        )

        if torch.cuda.is_available():
            self.model = self.model.to("cuda")

    def _run(self, smiles: str) -> str:
        try:
            input_ids = self.tokenizer(smiles, return_tensors="pt").input_ids
            if torch.cuda.is_available():
                input_ids = input_ids.to("cuda")

            outputs = self.model.generate(
                input_ids, num_beams=5, max_length=512, temperature=0.7, do_sample=True
            )
            description = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            return description
        except Exception as e:
            print(f"Error in molecule to text conversion: {str(e)}")
            return f"Error in molecule to text conversion: {str(e)}"

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"


# tools/molecule_optimizer_tool.py

from crewai_tools import BaseTool
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from pydantic import Field
from rdkit import Chem
import torch

class MoleculeOptimizerTool(BaseTool):
    name: str = Field(default="Molecule Optimizer")
    description: str = Field(
        default="Suggests optimizations for molecular properties using MolGen model"
    )

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.tokenizer = AutoTokenizer.from_pretrained("zjunlp/MolGen-large")
        self.model = AutoModelForSeq2SeqLM.from_pretrained("zjunlp/MolGen-large")

        if torch.cuda.is_available():
            self.model = self.model.to("cuda")

    def _run(self, smiles: str) -> list:
        try:
            sf_input = self.tokenizer(smiles, return_tensors="pt")
            if torch.cuda.is_available():
                sf_input = {k: v.to("cuda") for k, v in sf_input.items()}

            molecules = self.model.generate(
                input_ids=sf_input["input_ids"],
                num_beams=5,
                num_return_sequences=3,
                max_length=512,
                temperature=0.7,
                do_sample=True,
            )

            suggestions = []
            for mol in molecules:
                smiles_output = (
                    self.tokenizer.decode(
                        mol, skip_special_tokens=True, clean_up_tokenization_spaces=True
                    )
                    .replace(" ", "")
                    .strip()
                )
                if Chem.MolFromSmiles(smiles_output) is not None:
                    suggestions.append(smiles_output)

            if not suggestions:
                suggestions = ["No valid optimizations found"]

            return suggestions
        except Exception as e:
            print(f"Error in optimization: {str(e)}")
            return [f"Error in optimization: {str(e)}"]

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"



the idea it to get the output in a proper format 

now i am getting the text as the final result and if needed molecule visualization 

rather than that 

text final output 

if their is molecule output that molecule visualization down the final output and beside it all the properties calcualted and validity checker results as well if valid tick mark 

how to do this