# ===============================
# Imports
# ===============================
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM


# ===============================
# Optional CPU optimizations
# ===============================
torch.set_num_threads(4)  # adjust based on CPU cores


# ===============================
# Load Model & Tokenizer (ONCE)
# ===============================
print("Loading Aksara v1 model on CPU...")

tokenizer = AutoTokenizer.from_pretrained(
    "cropinailab/aksara_v1",
    use_fast=True,
)

model = AutoModelForCausalLM.from_pretrained(
    "cropinailab/aksara_v1",
    torch_dtype=torch.float32,  # CPU-safe
)

model.eval()

print("Model loaded successfully!")


# ===============================
# Core Generation Function
# ===============================
def generate_agri_response(plant, disease):
    prompt = f"""
You are an agricultural expert specializing in plant pathology, crop nutrition, and safe farm management.
Your job is to provide accurate, scientifically correct, and legally safe advice.

Plant: {plant}
Issue: {disease}

Your response MUST follow this structure clearly and must be 100% accurate:

### 1. About the Disease
- Explain what the disease is and identify the correct pathogen type (fungus, bacteria, virus, pest, oomycete, etc.)
- Describe how it spreads (only scientifically correct modes of spread)
- Avoid any incorrect or exaggerated claims

### 2. Symptoms
- Describe accurate symptoms on each relevant plant part:
  - Leaves
  - Stems
  - Roots
  - Fruit (only if the plant actually produces edible fruit)
  - Tubers/roots if the crop is root-based

### 3. Safe & Legal Treatment Options
- Copper-based fungicides
- Mancozeb
- Chlorothalonil
- Biological controls (Trichoderma, Bacillus, Pseudomonas)
- Cultural practices (airflow, sanitation, moisture control)
- NO dosages or banned chemicals

### 4. Prevention
- Resistant varieties
- Crop rotation
- Spacing & airflow
- Drip irrigation
- Field hygiene

### 5. Nutrient Requirements
Explain roles of:
N, P, K, Ca, Mg, S, Fe, Zn, B, Mn, Cu, Mo

### 6. Fertilizer Recommendations (No Dosages)
- Chemical
- Organic
- Biofertilizers

### 7. Additional Good Practices
- Irrigation
- Drainage
- Tool sanitation
"""

    inputs = tokenizer(prompt, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=600,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.15,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )

    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Remove prompt echo safely
    if full_output.startswith(prompt):
        response = full_output[len(prompt):].strip()
    else:
        response = full_output.replace(prompt, "").strip()

    return response


# ===============================
# Gradio Interface
# ===============================
with gr.Blocks(title="🌱 Agricultural Disease Advisor") as demo:
    gr.Markdown(
        """
        # 🌱 Agricultural Disease Advisor
        **CPU-based AI assistant for plant disease management**
        
        Enter the crop and disease name to receive **safe, scientific, and legal agricultural guidance**.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            plant_input = gr.Textbox(
                label="🌾 Plant / Crop Name",
                placeholder="e.g., Potato, Tomato, Rice",
            )

            disease_input = gr.Textbox(
                label="🦠 Disease / Issue",
                placeholder="e.g., Late Blight, Leaf Curl Virus",
            )

            generate_btn = gr.Button("🔍 Generate Advice", variant="primary")

        with gr.Column(scale=2):
            output_text = gr.Markdown(
                label="📋 Agricultural Guidance",
            )

    generate_btn.click(
        fn=generate_agri_response,
        inputs=[plant_input, disease_input],
        outputs=output_text,
    )

    gr.Markdown(
        """
        ---
        ⚠️ **Disclaimer:**  
        This tool provides general agricultural guidance only.  
        Always consult local agricultural extension services for field-level decisions.
        """
    )


# ===============================
# Launch
# ===============================
if __name__ == "__main__":
    demo.launch()