import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import numpy as np

# -------------------------------
# Load model and tokenizer
# -------------------------------
MODEL_NAME = "JacobLinCool/IELTS_essay_scoring_safetensors"

model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model.eval()

# -------------------------------
# Scoring + Feedback function
# -------------------------------
def score_essay(essay: str):
    if not essay.strip():
        return {
            "Task Achievement": 0.0,
            "Coherence & Cohesion": 0.0,
            "Vocabulary": 0.0,
            "Grammar": 0.0,
            "Overall": 0.0,
            "Feedback": "Please write an essay to get a score."
        }

    # Tokenize input
    inputs = tokenizer(
        essay,
        return_tensors="pt",
        truncation=True,
        max_length=512
    )

    # Model inference
    with torch.no_grad():
        outputs = model(**inputs)

    preds = outputs.logits.squeeze().cpu().numpy()

    # -------------------------------
    # Refined piecewise scaling
    # -------------------------------
    scores = preds.copy()
    for i in range(len(scores)):
        if scores[i] < 4.0:       # very weak essays
            scores[i] = scores[i]*1.25 + 0.3
        elif scores[i] < 6.0:     # medium essays
            scores[i] = scores[i]*1.1 + 0.2
        elif scores[i] < 8.0:     # good essays
            scores[i] = scores[i]*1.05 + 0.1
        else:                      # very strong essays
            scores[i] = scores[i]  # keep unchanged

    # -------------------------------
    # Word count bonus
    # -------------------------------
    wc = len(essay.split())
    if wc >= 250:
        scores += 0.3
    elif wc >= 200:
        scores += 0.15

    # -------------------------------
    # Clamp and round to 0.5
    # -------------------------------
    scores = np.clip(scores, 0, 9)
    rounded = np.round(scores * 2) / 2

    ta, cc, vocab, gram, overall_dummy = rounded  # we will compute weighted overall

    # -------------------------------
    # Weighted overall band
    # -------------------------------
    overall = 0.3*ta + 0.25*cc + 0.2*vocab + 0.25*gram
    overall = np.clip(overall, 0, 9)
    overall = np.round(overall * 2) / 2
    rounded[-1] = overall  # replace last element

    labels = ["Task Achievement", "Coherence & Cohesion", "Vocabulary", "Grammar", "Overall"]

    # -------------------------------
    # Dynamic mini feedback
    # -------------------------------
    feedback_lines = []

    if ta < 5:
        feedback_lines.append("Task: Ideas are underdeveloped; give examples and elaboration.")
    if cc < 5:
        feedback_lines.append("Coherence: Improve transitions; use connectors like 'however', 'moreover'.")
    if vocab < 5:
        feedback_lines.append("Vocabulary: Use a wider range of words and avoid repetition.")
    if gram < 5:
        feedback_lines.append("Grammar: Check articles, plurals, tenses, and sentence structures.")
    if overall >= 5:
        feedback_lines.append("Overall: Good attempt! Refine grammar and vocabulary to increase your band.")

    feedback = "\n".join(feedback_lines[:4])  # max 4 lines

    result = {label: float(score) for label, score in zip(labels, rounded)}
    result["Feedback"] = feedback
    return result

# -------------------------------
# Gradio UI
# -------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 📝 Refined Automated IELTS Writing Scorer")
    gr.Markdown(
        "Paste your IELTS Task 2 essay below. "
        "The system estimates band scores with mini feedback."
    )

    essay_input = gr.Textbox(
        lines=12,
        placeholder="Paste your IELTS essay here..."
    )

    score_output = gr.JSON(label="Estimated IELTS Band Scores + Feedback")

    submit_btn = gr.Button("Score Essay")
    submit_btn.click(
        fn=score_essay,
        inputs=essay_input,
        outputs=score_output
    )

    gr.Markdown(
        "⚠️ Note: AI-based estimator; not an official IELTS score."
    )

# -------------------------------
# Launch app
# -------------------------------
demo.launch()