import gradio as gr
from transformers import pipeline
from langchain import PromptTemplate
from langchain.chains import LLMChain
from langchain_google_genai import ChatGoogleGenerativeAI
import os

api_key = os.environ.get('GOOGLE_API_KEY')
if api_key is None:
    raise ValueError("No API key found. Please set the 'GOOGLE_API_KEY' environment variable.")

os.environ['GOOGLE_API_KEY'] = api_key

# Initialize the OCR pipeline
ocr_pipe = pipeline("image-to-text", model="jinhybr/OCR-Donut-CORD")

# Define detailed descriptions
descriptions = {
    ("Academic", "Task 1"): """You are evaluating an Academic Task 1. The candidate is expected to describe, summarize, or explain a visual representation (such as a graph, chart, table, or diagram). They should write a minimum of 150 words. Key features to look for include:

1. **Introduction:** A brief introduction of what the visual representation is about.
2. **Overview:** A clear and concise summary of the main trends, differences, or stages without going into too much detail.
3. **Detailed Description:** Description of the main features in detail, highlighting significant trends, making comparisons where relevant, and mentioning important data points or changes.
4. **Logical Structure:** The report should be organized logically, ensuring it flows coherently from one point to the next.
5. **Formal Tone:** The writing should maintain a formal and objective tone, avoiding personal opinions or casual language.
6. **Accurate Data Reporting:** Accurate reporting of data without inaccuracies or overgeneralizations.
7. **Use of Linking Words:** Appropriate use of linking words and phrases to connect ideas and sentences (e.g., however, therefore, in contrast, similarly).""",
    
    ("General", "Task 1"): """You are evaluating a General Training Task 1. The candidate is expected to write a letter in response to a given situation. They should write a minimum of 150 words. Key features to look for include:

1. **Introduction:** A clear statement of the purpose of the letter in the opening sentence.
2. **Body:** Inclusion of relevant details depending on the task (e.g., explaining a problem, extending an invitation). If writing a complaint, they should explain the problem and possible solutions. If writing an invitation, include necessary details like date, time, and location.
3. **Conclusion:** A polite closing statement.
4. **Appropriate Tone:** Matching the tone of the letter to the given situation (formal for official letters, semi-formal for professional acquaintances, and informal for friends or family).
5. **Structure and Organization:** Organized into clear paragraphs, each addressing a specific point or aspect of the task.
6. **Polite and Courteous Language:** Polite and courteous language, especially in formal and semi-formal letters.
7. **Use of Linking Words:** Appropriate use of linking words and phrases to ensure coherence (e.g., firstly, moreover, consequently).""",

    ("Academic", "Task 2"): """You are evaluating an Academic Task 2. The candidate is expected to write an essay in response to a point of view, argument, or problem. They should write a minimum of 250 words. Key features to look for include:

1. **Introduction:** An introduction that paraphrases the task statement and states the thesis or main argument.
2. **Body Paragraphs:** Development of the argument in 2-3 body paragraphs, each focusing on a single main idea, supported by examples and evidence.
3. **Clear Position:** A clear and consistent position on the topic throughout the essay. If discussing both sides of an issue, making it clear which side they agree with (if any).
4. **Logical Progression:** Logical organization and coherence from one point to the next.
5. **Formal Tone:** Maintenance of a formal and academic tone throughout.
6. **Use of Linking Words:** Appropriate use of linking words and phrases to connect ideas and paragraphs (e.g., on the other hand, in addition, furthermore, in conclusion).
7. **Conclusion:** A conclusion that summarizes the main points and restates the position or provides a final thought.
8. **Complex Sentence Structures:** Demonstration of a variety of complex sentence structures.
9. **Vocabulary Range:** Use of a wide range of vocabulary accurately and appropriately.""",

    ("General", "Task 2"): """You are evaluating a General Training Task 2. The candidate is expected to write an essay in response to a point of view, argument, or problem. They should write a minimum of 250 words. Key features to look for include:

1. **Introduction:** An introduction that paraphrases the task statement and states the thesis or main argument.
2. **Body Paragraphs:** Development of the argument in 2-3 body paragraphs, each focusing on a single main idea, supported by examples and evidence.
3. **Clear Position:** A clear and consistent position on the topic throughout the essay. If discussing both sides of an issue, making it clear which side they agree with (if any).
4. **Logical Progression:** Logical organization and coherence from one point to the next.
5. **Formal Tone:** Maintenance of a formal and academic tone throughout.
6. **Use of Linking Words:** Appropriate use of linking words and phrases to connect ideas and paragraphs (e.g., on the other hand, in addition, furthermore, in conclusion).
7. **Conclusion:** A conclusion that summarizes the main points and restates the position or provides a final thought.
8. **Complex Sentence Structures:** Demonstration of a variety of complex sentence structures.
9. **Vocabulary Range:** Use of a wide range of vocabulary accurately and appropriately."""
}

# Define the initial prompt for the LLM
initial_prompt = '''
You are an IELTS writing examiner. Below, you will receive a series of details pertaining to a writing task. Please evaluate the writing based on the information provided:

{description}

Please follow these steps for your evaluation:

1. Provide a score on the IELTS scale (0-9 points).
2. Leave a space.
3. Detail the weaknesses and mistakes in the writing.

Evaluate the writing based on the following four criteria:
- Task Response
- Coherence and Cohesion
- Lexical Resource
- Grammatical Range and Accuracy

Respond in this exact format:

Task type: [task type]

Task number: [task number]

score: [score]

Weaknesses and Mistakes:

[Weaknesses and Mistakes]

Notes:

- The text may have been obtained via OCR, which could result in some errors.
- Disregard any text enclosed in <>. They separate different parts of the text.
- If the Task type is empty, try to identify the task type from the question. If you cannot determine the task type, mention that the task type is unclear.
- If the specific question is not present in the task content, mention that the question does not exist in the task response.


Task type: {task_type}

Task number: {task_number}

Question: {question}

Task content: {content}
'''

# Initialize the LLM
llm_model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.7, top_p=0.85)

# Define the prompt template
prompt = PromptTemplate(input_variables=['task_type', 'task_number', 'question', 'content', 'description'], template=initial_prompt)

# Define the LLM chain
chain = LLMChain(
    llm=llm_model,
    prompt=prompt,
)

def evaluate(task_type, task_number, question, image):
    # Process the image to extract text
    text_content = ocr_pipe(image)
    content = text_content[0]['generated_text']

    # Select the appropriate description based on user input
    description = descriptions.get((task_type, task_number), "")

    # Run the chain
    result = chain.run({
        'task_type': task_type,
        'task_number': task_number,
        'question': question,
        'content': content,
        'description': description
    })

    return result

# Create the Gradio interface
inputs = [
    gr.Dropdown(choices=["Academic", "General"], label="Test Type", value="Academic"),
    gr.Dropdown(choices=["Task 1", "Task 2"], label="Task Number", value="Task 1"),
    gr.Textbox(label="Question", value=""),
    gr.Image(type="pil", label="Upload Image")
]

outputs = gr.Markdown(label="Result")

gr.Interface(fn=evaluate, inputs=inputs, outputs=outputs, title="IELTS Writing Evaluation").launch(share=True)