import gradio as gr
from PIL import Image
import re
import cv2
import numpy as np
import torch
from transformers import TrOCRProcessor, VisionEncoderDecoderModel

# =========================
# MODEL INITIALIZATION
# =========================

device = "cuda" if torch.cuda.is_available() else "cpu"

processor = TrOCRProcessor.from_pretrained(
    "microsoft/trocr-base-handwritten"
)

model = VisionEncoderDecoderModel.from_pretrained(
    "microsoft/trocr-base-handwritten"
).to(device)

# =========================
# OCR FUNCTIONS
# =========================

def extract_text_per_line(line_image):

    try:
        if line_image.mode != "RGB":
            line_image = line_image.convert("RGB")

        pixel_values = processor(
            images=line_image,
            return_tensors="pt"
        ).pixel_values.to(device)

        generated_ids = model.generate(
            pixel_values,
            max_length=64,
            num_beams=5,
            early_stopping=True
        )

        return processor.batch_decode(
            generated_ids,
            skip_special_tokens=True
        )[0]

    except Exception as e:
        print(f"OCR Error: {e}")
        return ""


def full_page_ocr(pil_image):

    if pil_image is None:
        return ""

    img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    thresh = cv2.adaptiveThreshold(
        gray,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        21,
        15
    )

    projection = np.sum(thresh, axis=1)

    upper_bounds = []
    lower_bounds = []

    in_line = False
    min_intensity = 1500

    for i, val in enumerate(projection):

        if not in_line and val > min_intensity:
            upper_bounds.append(i)
            in_line = True

        elif in_line and val < min_intensity:
            lower_bounds.append(i)
            in_line = False

    all_text = []

    h, w = img.shape[:2]

    for up, low in zip(upper_bounds, lower_bounds):

        pad = 5

        line_roi = img[
            max(0, up - pad):min(h, low + pad),
            0:w
        ]

        line_pil = Image.fromarray(
            cv2.cvtColor(line_roi, cv2.COLOR_BGR2RGB)
        )

        text = extract_text_per_line(line_pil)

        if text.strip():
            all_text.append(text.strip())

    return " ".join(all_text)

# =========================
# SCORING FUNCTIONS
# =========================

def count_sentences(text):
    return len([
        s for s in re.split(r'[.!?]+', text)
        if s.strip()
    ])


def score_structure(text):

    n = count_sentences(text)

    if n > 0:
        return min(4, max(1, n + 1))

    return 1


def score_punctuation(text):

    sentences = count_sentences(text)

    if sentences == 0:
        return 1

    if text.count('.') >= sentences:
        return 4

    elif text.count('.') >= 1:
        return 3

    return 2


def score_spacing(text):

    words = text.split()

    if len(words) > 5:
        return 4

    return 2


def score_content(text):

    words = len(text.split())

    if words > 30:
        return 4

    if words > 15:
        return 3

    if words > 5:
        return 2

    return 1


def score_readability(text):

    if len(text) > 50:
        return 4

    if len(text) > 25:
        return 3

    return 2


def score_critical(text):

    text_l = text.lower()

    has_reason = any(
        w in text_l
        for w in ["karena", "jadi", "sebab"]
    )

    has_pred = any(
        w in text_l
        for w in ["akan", "mungkin", "bila", "kalau"]
    )

    if has_reason and has_pred:
        return 4

    if has_reason or has_pred:
        return 3

    return 2

# =========================
# FEEDBACK
# =========================

def generate_feedback(scores):

    fb = []

    if scores["readability"] >= 3:
        fb.append(
            "🌟 Hebat! Tulisanmu rapi dan bagus."
        )
    else:
        fb.append(
            "✏️ Coba tulis lebih rapi ya."
        )

    if scores["punctuation"] < 3:
        fb.append(
            "📝 Jangan lupa beri tanda titik."
        )

    if scores["structure"] >= 3:
        fb.append(
            "💡 Ide ceritamu menarik."
        )
    else:
        fb.append(
            "📖 Tambahkan lebih banyak kalimat."
        )

    if scores["critical"] >= 3:
        fb.append(
            "🧠 Kamu sudah berpikir kritis."
        )
    else:
        fb.append(
            "🤔 Coba tambahkan alasan menggunakan kata 'karena'."
        )

    return "\n\n".join(fb)

# =========================
# MAIN EVALUATION
# =========================

def evaluate(image):

    text = full_page_ocr(image)

    if not text:
        return (
            "⚠️ Gambar tidak terbaca. "
            "Pastikan foto jelas dan terang."
        )

    scores = {
        "structure": score_structure(text),
        "punctuation": score_punctuation(text),
        "spacing": score_spacing(text),
        "content": score_content(text),
        "readability": score_readability(text),
        "critical": score_critical(text)
    }

    total = sum(scores.values())

    final_score = (total / 24) * 100

    feedback = generate_feedback(scores)

    result = f"""
🎉 HORE! INI HASIL PENILAIANMU 🎉

📊 NILAI AKHIR: {final_score:.1f}/100

🔍 Detail Nilai:
- Struktur Cerita: {scores['structure']}
- Tanda Baca: {scores['punctuation']}
- Kerapian Spasi: {scores['spacing']}
- Isi Cerita: {scores['content']}
- Kerapian Tulisan: {scores['readability']}
- Berpikir Kritis: {scores['critical']}

💬 Feedback:
{feedback}
"""

    return result

# =========================
# THEME
# =========================

colorful_theme = gr.themes.Soft(
    primary_hue="purple",
    secondary_hue="blue",
    neutral_hue="slate"
)

custom_css = """
.gradio-container {
    background: linear-gradient(
        135deg,
        #f5f7fa 0%,
        #c3cfe2 100%
    );
}
"""

# =========================
# DOWNLOAD FUNCTION
# =========================

def prepare_download(result):

    if "HORE" not in result:
        return gr.update(visible=False)

    filename = "hasil_penilaian_GMUSH.txt"

    with open(filename, "w", encoding="utf-8") as f:
        f.write(result)

    return gr.update(
        value=filename,
        visible=True
    )

# =========================
# UI
# =========================

with gr.Blocks() as demo:

    gr.HTML(
        """
        <h1 style='color:black; text-align:center;'>
        🌈 GMUSH
        </h1>
        """
    )

    # =========================
    # INSTRUCTIONS
    # =========================

    with gr.Group():

        gr.Markdown("""
### 🚀 Cara Menggunakan Aplikasi GMUSH:

1. 📝 Lihat gambar tugas di sebelah kiri.
2. ✍️ Tulis jawabanmu di kertas.
3. 📸 Foto tulisan tanganmu lalu upload.
4. ✨ Klik tombol **MULAI PENILAIAN!**
5. ⏳ Tunggu sekitar **5 menit yaa** 😊  
   Robot GMUSH sedang membaca tulisanmu dengan teliti.

---

## ✏️ Yuk Jawab Pertanyaannya!

1. Apa yang terjadi pada gambar?
2. Mengapa hal itu bisa terjadi?
3. Apa yang akan terjadi selanjutnya?
        """)

    # =========================
    # MAIN LAYOUT
    # =========================

    with gr.Row():

        with gr.Column(scale=1):

            gr.Markdown("### 📖 Tugas")

            gr.Image(
                "https://images.unsplash.com/photo-1596464716127-f2a82984de30?w=600",
                label="Perhatikan gambar berikut"
            )

        with gr.Column(scale=2):

            image_input = gr.Image(
                type="pil",
                label="📷 Upload Foto Tulisanmu"
            )

            btn = gr.Button(
                "✨ MULAI PENILAIAN! (Tunggu ±5 Menit 😊)",
                variant="primary"
            )

    # =========================
    # OUTPUT
    # =========================

    result_output = gr.Textbox(
        label="📊 HASIL",
        lines=12
    )

    download_btn = gr.DownloadButton(
        "📥 SIMPAN HASILKU!",
        visible=False
    )

    # =========================
    # BUTTON ACTIONS
    # =========================

    btn.click(
        evaluate,
        inputs=image_input,
        outputs=result_output
    ).then(
        prepare_download,
        inputs=result_output,
        outputs=download_btn
    )

# =========================
# RUN APP
# =========================

if __name__ == "__main__":

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=colorful_theme,
        css=custom_css
    )