import random
import gradio as gr
from pathlib import Path
from reactagent.environment import Environment
from reactagent.agents.agent_research import ResearchAgent
from reactagent.runner import create_parser
from reactagent import llm
from reactagent.users.user import User


# Global variables to store session state
env = None
agent = None

# Predefined research paper text (example)
predefined_paper_text = """
    Title:
    Dataset and Baseline for Automatic Student Feedback Analysis

    Abstract:
    This paper presents a student feedback corpus containing 3000 instances of feedback written by university students. The dataset has been annotated for aspect terms, opinion terms, polarities of the opinion terms towards targeted aspects, document-level opinion polarities, and sentence separations. A hierarchical taxonomy for aspect categorization covering all areas of the teaching-learning process was developed. Both implicit and explicit aspects were annotated using this taxonomy. The paper discusses the annotation methodology, difficulties faced during the annotation, and details about aspect term categorization. The annotated corpus can be used for Aspect Extraction, Aspect Level Sentiment Analysis, and Document Level Sentiment Analysis. Baseline results for all three tasks are provided.
"""

# Predefined extracted elements based on the paper text
predefined_research_tasks = "The primary research tasks include the creation of a comprehensive student feedback corpus, aspect term annotation, opinion polarity annotation, and the development of a hierarchical taxonomy."
predefined_research_gaps = "Gaps include the lack of detailed aspect-level annotations in existing datasets and the focus on document-level sentiment analysis."
predefined_keywords = "Student Feedback Corpus, Aspect Terms, Opinion Terms, Polarity, Hierarchical Taxonomy, Aspect Extraction, Aspect Level Sentiment Analysis, Document Level Sentiment Analysis"
predefined_recent_works = """
    1. "Students feedback analysis model using deep learning-based method and linguistic knowledge for intelligent educational systems."
    2. "An Automated Approach for Analysing Students Feedback Using Sentiment Analysis Techniques."
    """

# Extraction function to simulate the extraction of Research Tasks (t), Research Gaps (g), Keywords (k), and Recent Works (R)
def extract_research_elements(paper_text):
    # Returning the predefined extracted content
    return predefined_research_tasks, predefined_research_gaps, predefined_keywords, predefined_recent_works

# Generation function for Research Hypothesis and Experiment Plan
def generate_research_idea_and_plan(tasks, gaps, keywords, recent_works):
    hypothesis = f"""
    Method: Advanced Aspect-Level Sentiment Analysis of Student Feedback Using a Hybrid Deep Learning Approach

    Step 1: Dataset Enhancement 

    Data Collection and Preprocessing
    * Collect additional student feedback from multiple universities to expand the existing dataset.
    * Preprocess the data to ensure uniformity in annotation and eliminate noise, such as redundant information and grammatical errors.
    Annotation Refinement
    * Use advanced NLP techniques to further refine the aspect terms, opinion terms, and polarities.
    * Incorporate semi-supervised learning methods to improve annotation accuracy, utilizing both manual and automated processes.

    Step 2: Model Development
    Hybrid Model Architecture
    * Develop a hybrid model that integrates CNN, BiLSTM, and attention mechanisms, similar to the DTLP approach mentioned in the recent work by DTLP (Deep Learning and Teaching Process).
    * Incorporate a Transformer-based model (like BERT) to capture contextual nuances and improve the understanding of implicit aspects.
    Feature Integration
    * Enhance the feature set by combining statistical, linguistic, and sentiment knowledge features with word embeddings.
    * Include sentiment shifter rules and contextual polarity indicators to address challenges in sentiment analysis.

    Step 3: Training and Validation
    Model Training
    * Train the hybrid model using the enhanced dataset.
    * Use cross-validation techniques to ensure robustness and prevent overfitting.
    Baseline Comparisons
    * Compare the model's performance with baseline results provided in the original study and other recent works.
    * Use metrics such as accuracy, precision, recall, and F1-score to evaluate model performance across different tasks, including Aspect Extraction, Aspect Level Sentiment Analysis, and Document Level Sentiment Analysis.

    Step 4: Iterative Refinement
    Feedback Loop
    * Implement an iterative feedback loop where the model's predictions are reviewed and corrected, improving the model iteratively.
    * Engage domain experts in the review process to ensure the relevance and accuracy of the feedback. Continuous Learning
    * Utilize active learning techniques to continuously update the model with new data, ensuring it remains up-to-date with current trends in student feedback.

    Step 5: Deployment and Application
    Integration with Educational Systems
    * Deploy the model as a part of an intelligent educational system to analyze student feedback in real-time.
    * Provide actionable insights to educators and administrators to improve teaching methods and curriculum design. User Interface Development
    * Develop an intuitive user interface that allows educators to interact with the model, view feedback analysis, and generate reports.
    """

    experiment_plan = f"""
    Experiment: Validating the Hybrid Deep Learning Approach for Aspect-Level Sentiment Analysis of Student Feedback

    Objective:
    To validate the effectiveness of the proposed hybrid deep learning approach (combining CNN, BiLSTM, and Transformer models) for aspect-level sentiment analysis of student feedback by comparing its performance with baseline methods and recent works.
    Research Problem:
    Current sentiment analysis models for student feedback lack detailed aspect-level annotations and fail to address implicit aspects and contextual nuances in feedback data.
    Proposed Method:
    A hybrid deep learning model integrating CNN, BiLSTM, and Transformer-based models (like BERT) to enhance aspect-level sentiment analysis. The method incorporates sentiment shifter rules and contextual polarity indicators to address challenges in sentiment analysis.

    Experiment Design:
    1. Dataset Preparation:
        * Existing Dataset: Use the dataset provided by Herath et al. (2022) with 3000 instances of student feedback, annotated for aspect terms, opinion terms, polarities, and document-level sentiments.
        * Data Augmentation: Expand the dataset by collecting additional feedback from multiple universities, ensuring diversity in feedback data.
    2. Preprocessing:
        * Clean the data to remove noise and inconsistencies.
        * Tokenize the text and apply part-of-speech tagging.
        * Annotate additional feedback instances using the refined hierarchical taxonomy.
    3. Model Training:
        * Baseline Models: Implement and train traditional machine learning models (e.g., SVM, Naive Bayes) and existing deep learning models (e.g., LSTM, BiLSTM) for sentiment analysis.
        * Proposed Hybrid Model: Train the proposed hybrid model combining CNN, BiLSTM, and Transformer (BERT) layers. Use pre-trained embeddings and fine-tune on the feedback dataset.
    4. Feature Extraction:
        * Extract features using word embeddings, sentiment shifter rules, and contextual polarity indicators.
        * Integrate statistical, linguistic, and sentiment knowledge features with word embeddings to form a comprehensive feature set.
    5. Evaluation Metrics:
        * Measure the performance of models using accuracy, precision, recall, and F1-score.
        * Perform aspect-level evaluation by analyzing the accuracy of aspect term extraction and sentiment classification.
    6. Experiment Execution:
        * Training Phase: Train the baseline models and the proposed hybrid model on the training dataset.
        * Validation Phase: Validate the models using cross-validation techniques to ensure robustness and prevent overfitting.
        * Testing Phase: Evaluate the models on a held-out test set to compare their performance.
    7. Comparison and Analysis:
        * Compare the performance of the proposed hybrid model with baseline models and recent works, such as DTLP and other sentiment analysis techniques.
        * Analyze the results to identify strengths and weaknesses of the proposed model in handling aspect-level sentiment analysis and implicit aspects.
    8. Iterative Refinement:
        * Implement an iterative feedback loop where predictions are reviewed and corrected, improving model performance over iterations.
        * Engage domain experts to review the model's predictions and provide feedback for further refinement.
    9. Deployment:
        * Integrate the validated model into an intelligent educational system for real-time feedback analysis.
        * Develop a user interface to allow educators to interact with the model, view feedback analysis, and generate reports.
    """

    return hypothesis, experiment_plan

predefined_action_log = """
[Reasoning]: To understand the initial structure and functionality of train.py for effective improvements.
[Action]: Inspect Script (train.py)
Input: {"script_name": "train.py", "start_line_number": "1", "end_line_number": "74"}
Objective: Understand the training script, including data processing, [...]
[Observation]: The train.py script imports [...]. Sets random seeds [...]. Defines [...] Placeholder functions [...] exist without implementation. [...]
[Feedback]: The script structure is clear, but key functions (train_model, predict) need proper implementation for proposed model training and prediction.
"""

predefined_response = """
[Reasoning]: Execute the "final_model.py" using ExecuteScript action to evaluate performance of the final model.
[Action]: Execute "final_model.py" using ExecuteScript action.
Input: {"script_name": "final_model.py"}
"""

predefined_observation = """
Epoch [1/10],
Train MSE: 0.543,
Test MSE: 0.688
Epoch [2/10],
Train MSE: 0.242,
Test MSE: 0.493
"""


# # Structured input as list of dictionaries
# process_steps = [
#     "Action: Inspect Script Lines (train.py)\nObservation: The train.py script imports necessary libraries (e.g., pandas, sklearn, torch). Sets random seeds for reproducibility. Defines compute_metrics_for_regression function to calculate RMSE for different dimensions. Placeholder functions train_model and predict exist without implementations.\nFeedback: The script structure is clear, but key functions (train_model, predict) need proper implementation for proposed model training and prediction.",
#     "Action: Execute Script (train.py)\nObservation: The script executed successfully. Generated embeddings using the BERT model. Completed the training process without errors. Metrics calculation placeholders indicated areas needing implementation.\nFeedback: Experimental model definition and training logic are missing.",
#     "Action: Edit Script (train.py)\nObservation: Edited train.py to separate data loading, model definition, training loop, and evaluation into distinct functions. The edited train.py now has clearly defined functions for data loading (load_data), model definition (build_model), training (train_model), and evaluation (evaluate_model). Similarly, eval.py is reorganized to load the model and perform predictions efficiently.\nFeedback: Modify model architecture, retrieve the hybrid model of CNN, BiLSTM, and attention mechanisms, similar to the DTLP to align with the experiment design.",
#     "Action: Retrieve Model\nObservation: CNN and BiLSTM retrieved.\nFeedback: Modify the model architecture.",
#     "Action: Execute Script (train.py)\nObservation: The model trained over the specified number of epochs. Training and validation loss values are recorded for each epoch, the decrease in loss indicates improved model performance.\nFeedback: Continue with the next steps in model evaluation.",
#     predefined_observation
# ]

action_list = [

    predefined_response,
    predefined_observation
]
# Predefined code to display in Phase 2
predefined_code = """import pandas as pd
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
import numpy as np
import random
import torch
from sklearn.model_selection import train_test_split

DIMENSIONS = ["cohesion", "syntax", "vocabulary", "phraseology", "grammar", "conventions"]
SEED = 42

random.seed(SEED)
torch.manual_seed(SEED)
np.random.seed(SEED)

def compute_metrics_for_regression(y_test, y_test_pred):
    metrics = {}
    for task in DIMENSIONS:
        targets_task = [t[DIMENSIONS.index(task)] for t in y_test]
        pred_task = [l[DIMENSIONS.index(task)] for l in y_test_pred]
        
        rmse = mean_squared_error(targets_task, pred_task, squared=False)

        metrics[f"rmse_{task}"] = rmse
    
    return metrics

def train_model(X_train, y_train, X_valid, y_valid):
    model = None  # Placeholder for model training
    return model

def predict(model, X):
    y_pred = np.random.rand(len(X), len(DIMENSIONS))
    return y_pred

if __name__ == '__main__':

    ellipse_df = pd.read_csv('train.csv', 
                            header=0, names=['text_id', 'full_text', 'Cohesion', 'Syntax', 
                            'Vocabulary', 'Phraseology','Grammar', 'Conventions'], 
                            index_col='text_id')
    ellipse_df = ellipse_df.dropna(axis=0)

    data_df = ellipse_df
    X = list(data_df.full_text.to_numpy())
    y = np.array([data_df.drop(['full_text'], axis=1).iloc[i] for i in range(len(X))])

    X_train, X_valid, y_train, y_valid = train_test_split(X, y, test_size=0.10, random_state=SEED)

    model = train_model(X_train, y_train, X_valid, y_valid)

    y_valid_pred = predict(model, X_valid)
    metrics = compute_metrics_for_regression(y_valid, y_valid_pred)
    print(metrics)
    print("final MCRMSE on validation set: ", np.mean(list(metrics.values())))

    submission_df = pd.read_csv('test.csv',  header=0, names=['text_id', 'full_text'], index_col='text_id')
    X_submission = list(submission_df.full_text.to_numpy())
    y_submission = predict(model, X_submission)
    submission_df = pd.DataFrame(y_submission, columns=DIMENSIONS)
    submission_df.index = submission_df.index.rename('text_id')
    submission_df.to_csv('submission.csv')
"""

final_code = """
* Resulting train.py:
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from transformers import BertTokenizer, BertModel

# Define constants
DIMENSIONS = ['cohesion', 'syntax', 'vocabulary', 'phraseology', 'grammar', 'conventions']

class EssayDataset(Dataset):
    def __init__(self, texts, targets, tokenizer, max_len):
        self.texts = texts
        self.targets = targets
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, item):
        text = self.texts[item]
        target = self.targets[item]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            return_token_type_ids=False,
            padding='max_length',
            return_attention_mask=True,
            return_tensors='pt',
            truncation=True
        )

        return {
            'text': text,
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'targets': torch.tensor(target, dtype=torch.float)
        }

class EssayScoreRegressor(nn.Module):
    def __init__(self, n_outputs):
        super(EssayScoreRegressor, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.drop = nn.Dropout(p=0.3)
        self.out = nn.Linear(self.bert.config.hidden_size, n_outputs)

    def forward(self, input_ids, attention_mask):
        pooled_output = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )['pooler_output']
        output = self.drop(pooled_output)
        return self.out(output)

def train_epoch(model, data_loader, loss_fn, optimizer, device, scheduler, n_examples):
    model = model.train()
    losses = []

    for d in data_loader:
        input_ids = d['input_ids'].to(device)
        attention_mask = d['attention_mask'].to(device)
        targets = d['targets'].to(device)

        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        loss = loss_fn(outputs, targets)

        losses.append(loss.item())

        loss.backward()
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

    return np.mean(losses)

def train_model(train_data, val_data, tokenizer, model, optimizer, scheduler, device, epochs, batch_size, max_len):
    train_dataset = EssayDataset(
        texts=train_data['full_text'].to_numpy(),
        targets=train_data[DIMENSIONS].to_numpy(),
        tokenizer=tokenizer,
        max_len=max_len
    )

    val_dataset = EssayDataset(
        texts=val_data['full_text'].to_numpy(),
        targets=val_data[DIMENSIONS].to_numpy(),
        tokenizer=tokenizer,
        max_len=max_len
    )

    train_data_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True
    )

    val_data_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False
    )

    loss_fn = nn.MSELoss().to(device)

    for epoch in range(epochs):
        print(f'Epoch {epoch + 1}/{epochs}')
        print('-' * 10)

        train_loss = train_epoch(
            model,
            train_data_loader,
            loss_fn,
            optimizer,
            device,
            scheduler,
            len(train_dataset)
        )

        print(f'Train loss {train_loss}')

if __name__ == "__main__":
    df = pd.read_csv('train.csv')
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = EssayScoreRegressor(n_outputs=len(DIMENSIONS))
    model = model.to(device)

    optimizer = optim.Adam(model.parameters(), lr=2e-5)
    total_steps = len(df) // 16 * 5
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=total_steps, gamma=0.1)

    train_data = df.sample(frac=0.8, random_state=42)
    val_data = df.drop(train_data.index)

    train_model(train_data, val_data, tokenizer, model, optimizer, scheduler, device, epochs=5, batch_size=16, max_len=160)

* eval.py
import sys
import os
import pandas as pd
import numpy as np
import torch
from torch.utils.data import DataLoader
from transformers import BertTokenizer
from importlib import reload
import train

# Constants
DIMENSIONS = train.DIMENSIONS

class EssayDataset(Dataset):
    def __init__(self, texts, targets, tokenizer, max_len):
        self.texts = texts
        self.targets = targets
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, item):
        text = self.texts[item]
        target = self.targets[item]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            return_token_type_ids=False,
            padding='max_length',
            return_attention_mask=True,
            return_tensors='pt',
            truncation=True
        )

        return {
            'text': text,
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'targets': torch.tensor(target, dtype=torch.float)
        }

def get_score(submission_folder="../env"):
    submission_path = os.path.join(submission_folder, "submission.csv")
    solution = pd.read_csv(os.path.join(os.path.dirname(__file__), "answer.csv"))[DIMENSIONS].to_numpy()
    submission = pd.read_csv(submission_path)[DIMENSIONS].to_numpy()

    metrics = train.compute_metrics_for_regression(solution, submission)
    return np.mean(list(metrics.values()))

def eval_model(model, data_loader, device, n_examples):
    model = model.eval()
    predictions = []

    with torch.no_grad():
        for d in data_loader:
            input_ids = d['input_ids'].to(device)
            attention_mask = d['attention_mask'].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            predictions.extend(outputs.cpu().numpy())

    return predictions

if __name__ == "__main__":
    reload(train)
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = train.EssayScoreRegressor(n_outputs=len(DIMENSIONS))
    model.load_state_dict(torch.load('model.bin'))
    model = model.to(device)

    test_data = pd.read_csv('test.csv')
    test_dataset = EssayDataset(
        texts=test_data['full_text'].to_numpy(),
        targets=np.zeros((len(test_data), len(DIMENSIONS))),  # Dummy targets
        tokenizer=tokenizer,
        max_len=160
    )

    test_data_loader = DataLoader(
        test_dataset,
        batch_size=16,
        shuffle=False
    )

    predictions = eval_model(
        model,
        test_data_loader,
        device,
        len(test_dataset)
    )

    submission = pd.DataFrame(predictions, columns=DIMENSIONS)
    submission['text_id'] = test_data['text_id']
    submission.to_csv(os.path.join("../env", 'submission.csv'), index=False)

    print(get_score())

"""


class SessionInfo:
    def __init__(self):
        self.coro_cache = {}
        self.parser = create_parser()

    def make_session(self, prompt, session_hash):
        id = session_hash

        llm_name='claude-3-5-sonnet-20240620'
        fastllm_name='claude-3-haiku-20240307'
        rawargs = [
            '--research-problem', prompt,
            '--log-dir', str(Path('logs', id)),
            '--work-dir', str(Path('workspaces', id)),
            '--llm-name', llm_name,
            '--edit-script-llm-name', llm_name,
            '--fast-llm-name', fastllm_name,
        ]

        args = self.parser.parse_args(rawargs)
        # llm.FAST_MODEL = args.fast_llm_name
        env = Environment(args)
        # agent = ResearchAgent(args, env)
        coro = agent.run(env)

        self.coro_cache[id] = coro
        return id

    def get_response(self, human_input, session_hash):
        coro_input = human_input
        if session_hash not in self.coro_cache:
            self.make_session(human_input, session_hash)
            coro_input = None

        try:
            output = self.coro_cache[session_hash].send(coro_input)
        except StopIteration:
            output = None
            del self.coro_cache[session_hash]

        return output

session_info = SessionInfo()

def info_to_message(info):
    msg = ""
    for k, v in info.items():
        if isinstance(v, dict):
            tempv = v
            v = ""
            for k2, v2 in tempv.items():
                v += f"{k2}:\n  {v2}\n"
        v = User.indent_text(v, 2)
        msg += '-' * 64
        msg += '\n'
        msg += f"{k}:\n{v}\n"

    msg += "Please provide feedback based on the history, response entries, and observation, and questions: "
    return msg

def predict(message, history, request: gr.Request):
    response = session_info.get_response(message, request.session_hash)
    if response is None:
        return "Agent is finished. Enter a new instruction."
    return info_to_message(response)


# Initialize the global step_index and history
process_steps = [
    {
        "Action": "Inspect Script Lines (train.py)",
        "Observation": (
            "The train.py script imports necessary libraries (e.g., pandas, sklearn, torch). "
            "Sets random seeds for reproducibility. Defines compute_metrics_for_regression function "
            "to calculate RMSE for different dimensions. Placeholder functions train_model and "
            "predict exist without implementations."
        ),
    },
    {
        "Action": "Execute Script (train.py)",
        "Observation": (
            "The script executed successfully. Generated embeddings using the BERT model. Completed "
            "the training process without errors. Metrics calculation placeholders indicated areas needing implementation."
        ),
    },
    {
        "Action": "Edit Script (train.py)",
        "Observation": (
            "Edited train.py to separate data loading, model definition, training loop, and evaluation into distinct functions. "
            "The edited train.py now has clearly defined functions"
            "for data loading (load_data), model definition (build_model), "
            "training (train_model), and evaluation (evaluate_model). Similarly, eval.py is reorganized to load the model and perform predictions efficiently."
        ),
    },
    {
        "Action": "Retrieve Model",
        "Observation": "CNN and BiLSTM retrieved.",
    },
    {
        "Action": "Execute Script (train.py)",
        "Observation": (
            "The model trained over the specified number of epochs. Training and validation loss values are recorded for each epoch, "
            "the decrease in loss indicates improved model performance."
        )
    },
    {
        "Action": "Evaluation",
        "Observation": predefined_observation,
    }
]


# step_index = 0
# def info_to_message(info):
#     msg = "Agent Response:\n"
#     for k, v in info.items():
#         if isinstance(v, dict):
#             tempv = v
#             v = ""
#             for k2, v2 in tempv.items():
#                 v += f"{k2}:\n  {v2}\n"
#         v = User.indent_text(v, 2)
#         msg += '-' * 64
#         msg += '\n'
#         msg += f"{k}:\n{v}\n"

#     msg += "Please provide feedback based on the history, response entries, and observation, and questions: "
#     print(msg)
#     return msg

# def predict(message, history):
#     global step_index  # Declare the use of global variable
#     if step_index < len(process_steps):
#         response_info = process_steps[step_index]
#         response = info_to_message(response_info)  # Convert dictionary to formatted string
#         step_index += 1
#     else:
#         response = "Agent Finished."
    
#     return response, "N/A"  # Return the formatted string and clear input

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("# AI Research Assistant with Research Agent")

    # Use state variables to store generated hypothesis and experiment plan
    hypothesis_state = gr.State("")
    experiment_plan_state = gr.State("")

    # Phase 1: Research Idea Generation Tab
    with gr.Tab("Phase 1: Research Idea Generation"):
        gr.Markdown("### Extract Research Elements and Generate Research Ideas")
        
        with gr.Row():
            with gr.Column():
                paper_text_input = gr.Textbox(value=predefined_paper_text, lines=10, label="Research Paper Text")
                extract_button = gr.Button("Extract Research Elements")
                with gr.Row():
                    tasks_output = gr.Textbox(placeholder="Research task definition", label="Research Tasks", lines=2, interactive=False)
                    gaps_output = gr.Textbox(placeholder="Research gaps of current works", label="Research Gaps", lines=2, interactive=False)
                    keywords_output = gr.Textbox(placeholder="Paper keywords", label="Keywords", lines=2, interactive=False)
                    recent_works_output = gr.Textbox(placeholder="Recent works extracted from Semantic Scholar", label="Recent Works", lines=2, interactive=False)
            with gr.Column():
                with gr.Row():  # Move the button to the top right
                    generate_button = gr.Button("Generate Research Hypothesis & Experiment Plan")
                with gr.Group():
                    gr.Markdown("### Research Idea")
                    with gr.Row():
                        hypothesis_output = gr.Textbox(label="Generated Hypothesis", lines=45, interactive=False)
                        experiment_plan_output = gr.Textbox(label="Generated Experiment Plan", lines=45, interactive=False)
        
        # Step 1: Extract Research Elements
        extract_button.click(
            fn=extract_research_elements, 
            inputs=paper_text_input, 
            outputs=[tasks_output, gaps_output, keywords_output, recent_works_output]
        )

        # Step 2: Generate Research Hypothesis and Experiment Plan
        def generate_and_store(tasks, gaps, keywords, recent_works):
            hypothesis, experiment_plan = generate_research_idea_and_plan(tasks, gaps, keywords, recent_works)
            return hypothesis, experiment_plan, hypothesis, experiment_plan

        generate_button.click(
            fn=generate_and_store, 
            inputs=[tasks_output, gaps_output, keywords_output, recent_works_output], 
            outputs=[hypothesis_output, experiment_plan_output, hypothesis_state, experiment_plan_state]
        )

    # Phase 2: Interactive Session Tab
    with gr.Tab("Phase 2&3: Experiment implementation and execution"):
        gr.Markdown("### Interact with the ExperimentAgent")

        with gr.Row():
            with gr.Column():
                idea_input = gr.Textbox(label="Research Hypothesis", lines=30, interactive=False)
                plan_input = gr.Textbox(label="Experiment Plan", lines=30, interactive=False)
        
            with gr.Column():
                execute_button = gr.Button("Start ExperimentAgent", elem_classes=["execute-btn"])
                with gr.Group():
                    gr.Markdown("### Implementation + Execution Log")    
                    log = gr.Textbox(label="Execution Log", lines=20, interactive=False)
                    code_display = gr.Code(label="Implementation", language="python", interactive=False)

            with gr.Column():
                # chatbot = gr.ChatInterface(predict)
                response = gr.Textbox(label = "ExperimentAgent Response",  lines=30, interactive=False)
                feedback = gr.Textbox(placeholder="N/A", label = "User Feedback",  lines=3, interactive=True)
                submit_button = gr.Button("Submit", elem_classes=["Submit-btn"])
        
        def submit_feedback(user_feedback, history, previous_response):
            global step_index
            if_end = False
            step_index += 1
            if (step_index >= len(process_steps)):
                step_index  = 0
                msg = ""
            msg = history
            if step_index < len(process_steps):
                msg += previous_response + "\nUser feedback:" + user_feedback +"\n\n"
                response_info = process_steps[step_index]
                response = info_to_message(response_info)  # Convert dictionary to formatted string
                step_index += 1
            else:
                if_end = True
                response = "Agent Finished."
            msg += response

            return msg, response, predefined_code if if_end else final_code
# def predict(message, history):
#     global step_index  # Declare the use of global variable
#     if step_index < len(process_steps):
#         response_info = process_steps[step_index]
#         response = info_to_message(response_info)  # Convert dictionary to formatted string
#         step_index += 1
#     else:
#         response = "Agent Finished."

        # Automatically populate the hypothesis and plan in Phase 2
        def load_phase_2_inputs(hypothesis, plan):
            return hypothesis, plan, "# Code implementation will be displayed here after Start ExperimentAgent."
        
        # Function to implement and execute with the research agent
        def implement_and_execute(hypothesis, plan):
            predefined_message = f"Implement the following hypothesis and experiment plan:\n\nHypothesis:\n{hypothesis}\n\nExperiment Plan:\n{plan}"
            return predefined_code, predefined_action_log

        hypothesis_state.change(
            fn=load_phase_2_inputs, 
            inputs=[hypothesis_state, experiment_plan_state], 
            outputs=[idea_input, plan_input, code_display]
        )
        
        # Trigger the research agent execution with the predefined hypothesis and plan
        execute_button.click(
            fn=implement_and_execute, 
            inputs=[hypothesis_state, experiment_plan_state], 
            outputs=[code_display, log]
        )
        submit_button.click(
            fn=submit_feedback,
            inputs=[feedback, log, response],
            outputs=[log, response, code_display]
        )

if __name__ == "__main__":
    # app.launch(share=True)
    step_index = 0
    app.launch()