"""
香蕉叶病害智能识别系统
基于 ViT 视觉变换器的香蕉叶片病害分类应用
"""

import os

import streamlit as st
import plotly.graph_objects as go
from PIL import Image

# 国内 HuggingFace 镜像加速
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

from transformers import ViTForImageClassification, ViTImageProcessor, pipeline  # noqa: E402

# ─── Disease Label Mapping ──────────────────────────────────────────────────
# 模型输出 LABEL_0 ~ LABEL_6，映射为实际病害名称
DISEASE_INFO = {
    "LABEL_0": {
        "name": "细菌性枯萎病",
        "name_en": "Banana Bacterial Wilt (BBW)",
        "color": "#e74c3c",
        "description": "由细菌引起，叶片从边缘开始变黄枯萎，最终整株死亡。",
        "advice": "及时清除病株，避免刀具交叉感染，使用无菌工具操作，加强田间排水。",
    },
    "LABEL_1": {
        "name": "黑条斑病",
        "name_en": "Black Sigatoka",
        "color": "#8e44ad",
        "description": "真菌性病害，叶片出现黑色条纹，严重时叶片大面积枯死。",
        "advice": "选用抗病品种，合理密植保持通风，及时清除病叶，适时喷施杀菌剂。",
    },
    "LABEL_2": {
        "name": "苞片花叶病毒病",
        "name_en": "Bract Mosaic Virus (BBMV)",
        "color": "#e67e22",
        "description": "病毒性病害，苞片和叶片出现花叶状斑纹，影响果实品质。",
        "advice": "使用无病种苗，防治传毒蚜虫，及时拔除病株，避免机械传播。",
    },
    "LABEL_3": {
        "name": "健康叶片",
        "name_en": "Healthy",
        "color": "#27ae60",
        "description": "叶片状态良好，无明显病害症状，颜色鲜绿。",
        "advice": "继续保持良好的田间管理，合理施肥灌溉，定期巡检。",
    },
    "LABEL_4": {
        "name": "花叶病毒病",
        "name_en": "Banana Mosaic Virus",
        "color": "#f39c12",
        "description": "病毒性病害，叶片出现黄绿相间的花叶症状，植株矮化。",
        "advice": "选用无毒组培苗，控制蚜虫等传毒媒介，拔除病株，工具消毒。",
    },
    "LABEL_5": {
        "name": "巴拿马病（枯萎病）",
        "name_en": "Panama Disease (Fusarium Wilt)",
        "color": "#c0392b",
        "description": "由尖孢镰刀菌引起，维管束变褐，叶片从下部开始黄化枯萎。",
        "advice": "种植抗病品种，严格检疫，病区轮作或休耕，避免土壤传播。",
    },
    "LABEL_6": {
        "name": "黄条斑病",
        "name_en": "Yellow Sigatoka",
        "color": "#d4ac0d",
        "description": "真菌性病害，叶片出现黄色条纹，比黑条斑病症状较轻。",
        "advice": "保持田间通风透光，清除病残叶，适时用药防治，合理施肥增强抗性。",
    },
}


# ─── Page Config ────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="香蕉叶病害识别",
    page_icon="🍌",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown("""
<style>
html, body, [class*="css"] {
    font-family: "PingFang SC", "Microsoft YaHei", "Noto Sans SC", sans-serif;
}
</style>
""", unsafe_allow_html=True)


# ─── Model Loading ──────────────────────────────────────────────────────────
@st.cache_resource
def load_model():
    """加载 HuggingFace 模型（首次运行自动下载，约 343MB）"""
    model_name = "Dmitry43243242/banana-disease-leaf-model"
    # 模型的 preprocessor_config.json 中 image_processor_type 为旧版 ViTFeatureExtractor，
    # AutoImageProcessor 无法识别，直接用具体的 ViT 类加载
    processor = ViTImageProcessor.from_pretrained(model_name)
    model = ViTForImageClassification.from_pretrained(model_name)
    classifier = pipeline(
        "image-classification", model=model, image_processor=processor
    )
    return classifier


# ─── Sidebar ────────────────────────────────────────────────────────────────
with st.sidebar:
    st.header("🍌 香蕉叶病害识别")
    st.caption("上传叶片照片，AI 帮你识别病害")
    st.divider()

    st.subheader("📷 上传图片")
    uploaded_file = st.file_uploader(
        "选择香蕉叶片图片",
        type=["jpg", "jpeg", "png", "webp"],
        help="支持 JPG、PNG、WebP 格式",
    )

    st.divider()
    st.subheader("ℹ️ 使用说明")
    st.markdown("""
1. 上传一张香蕉叶片照片
2. 系统自动识别病害类型
3. 查看各类别置信度
4. 获取防治建议

**Tips:**
- 图片越清晰，识别越准确
- 尽量拍摄完整叶片
- 病变部位清晰可见
""")

    st.divider()
    st.subheader("📋 可识别病害")
    for label_key, info in DISEASE_INFO.items():
        st.markdown(
            f"<span style='color:{info['color']}'>●</span> **{info['name']}** ({info['name_en']})",
            unsafe_allow_html=True,
        )


# ─── Main Content ───────────────────────────────────────────────────────────
st.title("🍌 香蕉叶病害智能识别系统")
st.caption("基于 ViT 视觉变换器 · 7 类香蕉叶病害自动分类")

# 加载模型
with st.spinner("正在加载模型（首次运行需下载约 343MB）..."):
    classifier = load_model()

if uploaded_file is not None:
    # 显示上传的图片
    image = Image.open(uploaded_file).convert("RGB")

    col_img, col_result = st.columns([1, 1])

    with col_img:
        st.subheader("🖼️ 上传图片")
        st.image(image, use_container_width=True)

    # 推理
    with st.spinner("正在识别病害..."):
        results = classifier(image)

    # 解析结果
    top_result = results[0]
    top_label = top_result["label"]
    top_score = top_result["score"]
    disease = DISEASE_INFO.get(top_label, {
        "name": top_label,
        "name_en": top_label,
        "color": "#95a5a6",
        "description": "未知病害类型",
        "advice": "建议咨询专业植保人员",
    })

    with col_result:
        st.subheader("🔍 识别结果")
        is_healthy = top_label == "LABEL_3"
        if is_healthy:
            st.success(f"**{disease['name']}** — 置信度 {top_score:.1%}")
        else:
            st.warning(f"**{disease['name']}** — 置信度 {top_score:.1%}")

        st.markdown(f"**英文名**: {disease['name_en']}")
        st.markdown(f"**症状描述**: {disease['description']}")

    # ─── Confidence Bar Chart ────────────────────────────────────────────────
    st.subheader("📊 各类别置信度")

    labels = [r["label"] for r in results]
    scores = [r["score"] for r in results]
    display_names = [
        f"{DISEASE_INFO.get(l, {}).get('name', l)}" for l in labels
    ]
    colors = [
        DISEASE_INFO.get(l, {}).get("color", "#95a5a6") for l in labels
    ]

    fig = go.Figure()
    fig.add_trace(go.Bar(
        x=scores,
        y=display_names,
        orientation="h",
        marker=dict(
            color=colors,
            opacity=0.85,
            line=dict(color="rgba(0,0,0,0.08)", width=1),
        ),
        text=[f"{s:.1%}" for s in scores],
        textposition="outside",
        textfont=dict(color="#5a5a5a", size=11),
    ))
    fig.update_layout(
        xaxis=dict(
            title="置信度",
            range=[0, 1.15],
            tickformat=".0%",
            color="#5a5a5a",
            gridcolor="rgba(0,0,0,0.06)",
        ),
        yaxis=dict(color="#2c2c2c", autorange="reversed"),
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        font=dict(color="#2c2c2c", size=11),
        margin=dict(t=10, b=30, l=130, r=60),
        height=320,
        showlegend=False,
    )
    st.plotly_chart(fig, use_container_width=True)

    # ─── Pie Chart ──────────────────────────────────────────────────────────
    col_pie, col_advice = st.columns([1, 1])

    with col_pie:
        st.subheader("📈 置信度分布")
        fig_pie = go.Figure()
        fig_pie.add_trace(go.Pie(
            labels=display_names,
            values=scores,
            marker=dict(colors=colors),
            textinfo="label+percent",
            textfont=dict(size=10),
            hole=0.45,
        ))
        fig_pie.update_layout(
            paper_bgcolor="rgba(0,0,0,0)",
            font=dict(color="#2c2c2c", size=10),
            margin=dict(t=10, b=10, l=10, r=10),
            height=300,
            showlegend=False,
        )
        st.plotly_chart(fig_pie, use_container_width=True)

    # ─── Disease Advice ─────────────────────────────────────────────────────
    with col_advice:
        st.subheader("💡 防治建议")
        if is_healthy:
            st.success(f"叶片状态健康，继续保持良好的田间管理！")
        else:
            st.error(f"检测到 **{disease['name']}** 病害")
        st.info(disease["advice"])

        # 如果置信度不高，给出提示
        if top_score < 0.6:
            st.warning(
                "⚠️ 置信度较低，建议：\n"
                "- 确保图片清晰且病变部位可见\n"
                "- 尝试拍摄完整叶片\n"
                "- 可咨询专业植保人员确认"
            )

else:
    # 无图片上传时的占位内容
    st.info("👈 请在左侧上传一张香蕉叶片图片开始识别")

    # 展示示例区域
    st.subheader("📖 支持识别的病害类型")
    cols = st.columns(4)
    for i, (label_key, info) in enumerate(DISEASE_INFO.items()):
        with cols[i % 4]:
            with st.container(border=True):
                st.markdown(
                    f"<div style='color:{info['color']}; font-size:24px; font-weight:bold;'>●</div>",
                    unsafe_allow_html=True,
                )
                st.markdown(f"**{info['name']}**")
                st.caption(info["name_en"])

st.divider()
st.caption("香蕉叶病害智能识别系统 · 基于 ViT (Vision Transformer) · 仅供参考，确诊请咨询专业植保人员")