app.py

import argparse
import os
import spaces

import random
import torchvision.transforms as transforms
from monai.transforms import Compose, RandRotate, RandFlip, ToTensor,EnsureType,RandSpatialCrop,Resize,Orientation, CenterSpatialCrop,ScaleIntensityRange
from torchvision.transforms.functional import InterpolationMode
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import gradio as gr
from transformers import StoppingCriteriaList
from minigpt4.common.config import Config
from minigpt4.common.dist_utils import get_rank
from minigpt4.common.registry import registry
from minigpt4.conversation.conversation import Chat, CONV_VISION_Vicuna0, CONV_VISION_LLama2, StoppingCriteriaSub
import pydicom
# imports modules for registration
from minigpt4.datasets.builders import *
from minigpt4.models import *
from minigpt4.processors import *
from minigpt4.runners import *
from minigpt4.tasks import *

import zipfile
import tarfile
import gzip
import nibabel as nib
user_avatar='./icon/robot.png'
# custom_css = """
# .custom-column {
#     max-width: 300px; /* 限制列的最大宽度 */
#     padding: 5px; /* 列内边距，控制整体间距 */
#     box-sizing: border-box; /* 确保 padding 包含在宽度内 */
    
# }
# .center-title {
#     text-align: center !important;
#     flex-grow: 1; /* 占据剩余空间 */
#     display: flex;
#     justify-content: center; /* 居中对齐 */
#     align-items: center; /* 垂直居中 */
#     height: 60px;
# }
# .gradio-container {
#     background-color: #F5F6FA; /* 浅灰色背景 */
# }
# .logo-img {
#     max-width: 60px; /* 调整 Logo 宽度 */
#     height: auto;
#     vertical-align: middle; /* 垂直居中对齐标题 */
#     margin-left: 10px; /* 与标题的间距 */
#     float: left
# }

# .button.svelte-izfbkb.center.boundedheight.flex button {
#     max-height: 98px !important;
#     max-width: 60px; !important;
# }

# upload-button button {
#     background-color: #2E3192; /* 深蓝色按钮 */
#     color: white;
#     font-size: 4px;
#     height: 20px;
#     border-radius: 8px;
#     padding: 12px 24px;
#     font-weight: bold;
# }

# restart_button button {
#     background-color: #D3D3D3; /* 浅灰色次按钮 */
#     color: black;
#     border-radius: 8px;
#     padding: 10px 20px;
# }
# restart_button button:hover {
#     background-color: #B0B0B0; /* 悬停加深 */
# }

# .gradio-input input, .gradio-input label {
#     color: #333333; /* 输入框文字颜色 */
#     border-color: #E0E0E0; /* 输入框边框 */
# }



# .accordion-style {
#     background-color: white;
#     border: 1px solid #E0E0E0;
#     border-radius: 8px;
#     padding: 15px;
#     margin-top: 20px;
#     box-shadow: 0 2px 4px rgba(0,0,0,0.1);
# }

# .message-row img {
#     margin: 0px !important;
# }

# .message.bot.message-bubble-border {
#     background-color: #f0f8ff; /* 浅蓝背景 */
#     border-radius: 10px; /* 圆角 */
#     padding: 15px; /* 内边距 */
#     margin: 10px 0; /* 外边距 */
# }

# .avatar-container img {
#     padding: 0px !important;
# }

# .avatar-image {
#     width: 15px;
#     height: 15px;
#     border-radius: 50%;
#     object-fit: cover;
#     position: relative;
#     top: 0px;
#     left: 0px;
#     margin: 2px;
# }

# """



custom_css="""
/* ========================================
   全局配色方案 - 医学主题
======================================== */

/* 主容器和整体背景 */
body, .gradio-container {
    background: linear-gradient(135deg, #f0f8ff 0%, #e6f3ff 100%) !important;
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
    color: #2c3e50 !important;
}

/* 强制 bot 消息占满一行，消除侧边空隙 */
.bot.svelte-1nr59td.message {
    width: 100% !important;
    max-width: 100% !important;
    display: flex !important;
    justify-content: flex-start !important; /* 确保整体靠左 */
}

/* 调整内部文字包装层，取消宽度限制 */
.bot.svelte-1nr59td.message .message-wrap {
    # width: 100% !important;
    width: fit-content !important;
    max-width: 85% !important;
    flex-grow: 1 !important;
}

/* 如果文字在一个气泡/容器里，去掉它的 margin */
.bot.svelte-1nr59td.message .md {
    max-width: 100% !important;
}


/* 1. 取消全宽限制，允许容器根据内容收缩 */
.bot.svelte-1nr59td.message .panel-full-width {
    width: fit-content !important;
    min-width: unset !important;
    max-width: 85% !important; /* 限制最大宽度，防止长文本溢出 */
}

/* 2. 针对消息内容容器，确保其宽度也是自适应的 */
.bot.svelte-1nr59td.message [data-testid="bot"] {
    width: fit-content !important;
    display: inline-block !important; /* 或者 flex */
}

/* 3. 确保文字排版不会因为收缩而显得拥挤 */
.bot.svelte-1nr59td.message .message-content {
    padding-right: 15px !important; /* 给右侧留一点呼吸空间 */
}


/* 关键：调整消息内容的包装器，让其填满左侧 */
.bot.svelte-1nr59td.message > .message-wrap {
    margin-left: 0 !important;
    padding-left: 0 !important;
}

# .bot.svelte-1nr59td.message,
# .md.svelte-1hf8a14.chatbot.prose {
#     width: 100% !important;
#     height: 100% !important;
#     display: flex !important;
#     align-items: stretch !important;
#     justify-content: stretch !important;
#     box-sizing: border-box !important;
#     padding: 15px !important;
#     background-color: #f0f8ff !important; /* 或已定义的背景色 */
#     border-radius: 10px !important;
#     min-height: 60px !important;
# }

# .md.svelte-1hf8a14.chatbot.prose {
#     flex: 1 1 auto !important;
#     width: 100% !important;
#     height: 100% !important;
#     display: flex !important;
#     align-items: center !important;
#     word-break: break-word !important;
#     box-sizing: border-box !important;
# }

/* 2. 移除右侧多余滚动条并美化 */
/* 尝试隐藏外层导致双滚动条的溢出 */
div.gradio-container {
    overflow: visible !important;
}

/* 针对 Gradio Chatbot 的主容器 */
.svelte-1nr59td.chatbot {
    border: none !important; /* 移除边框有时能解决视觉上的滚动条错觉 */
}

/* 隐藏特定的滚动条，只保留主对话区域的滚动功能 */
.svelte-1nr59td.chatbot .wrapper::-webkit-scrollbar {
    width: 6px; /* 设置一个很细的滚动条，或者设为 0 完全隐藏 */
}



/* 标题样式 */
h1, h2, h3, h4, h5, h6 {
    color: #2E3192 !important;
    font-weight: 600 !important;
}

/* 自定义列样式 */
.custom-column {
    max-width: 200px !important;
    width: 200px !important;
    padding: 15px !important;
    box-sizing: border-box !important;
    background: rgba(255, 255, 255, 0.9) !important;
    border-radius: 12px !important;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.08) !important;
    margin: 10px !important;
    border: 1px solid #e0e0e0 !important;
    flex-shrink: 0 !important;
}



/* 针对所有可能的列选择器 */
div[data-testid="column"]:first-child,
.gradio-row .gradio-column:first-child,
.gr-row .gr-column:first-child {
    max-width: 200px !important;
    width: 200px !important;
    flex: 0 0 200px !important;
}

div#component-7.block.custom-avatar.svelte-12cmxck {
    height: 510px !important;
}


/* ========================================
   聊天消息样式 - 医学主题
======================================== */


/* 优化脑肿瘤诊断网站配色方案
   - 主色调采用医学相关的蓝绿色，传递专业、安心、清晰的视觉体验
   - 突出操作按钮，采用协调的高饱和色，焦点内容保持简洁柔和
   - 明确消息区分，加强医生/机器人和用户的消息对比度
*/

/* 全局背景与字体主色 */
body, .main-background {
    background: linear-gradient(135deg, #f2fcfc 0%, #eaf3fb 100%) !important;
    # color: #22324a !important;
}

/* 主要按钮颜色（医学蓝绿色，操作性高） */
.btn-primary, 
.submit-button button, 
button#component-20.lg.secondary.svelte-cmf5ev {
    background: linear-gradient(90deg, #42a5f5 0%, #26c6da 100%) !important;
    color: #fff !important;
    border: none !important;
    border-radius: 8px !important;
    font-weight: bold !important;
    box-shadow: 0 4px 18px rgba(38, 198, 218, 0.14) !important;
}

button:hover, .btn-primary:hover, .submit-button button:hover {
    background: linear-gradient(90deg, #1e88e5 0%, #00bcd4 100%) !important;
    box-shadow: 0 6px 24px rgba(33, 150, 243, 0.22) !important;
}

/* 次要按钮、文件上传/重启等采用辅助绿色或暖色 */
button.upload-button, .btn-secondary, button#component-20.lg.primary.upload-button.svelte-cmf5ev {
    background: linear-gradient(90deg, #43e97b 0%, #38f9d7 100%) !important;
    color: #fff !important;
}

button.upload-button:hover, .btn-secondary:hover,
button#component-20.lg.primary.upload-button.svelte-cmf5ev:hover {
    background: linear-gradient(90deg, #0bab64 0%, #3bb78f 100%) !important;
}

/* 诊断相关内容强调色：高亮医学橙/红色，仅用于提示/错误 */
.diagnosis-highlights, .warning, .error {
    color: #f74c4c !important;
    background: #ffecec !important;
    border-left: 4px solid #f74c4c !important;
    border-radius: 6px !important;
    padding: 8px 14px !important;
}

/* 聊天区优化——医生/机器人消息（宁静青绿） */
/* 
   修复：等待 model 回复时，先出现橙色背景再变为绿色——
   通过增加 div.message.bot 的权重，确保无论任何加载/刷新顺序下都优先使用绿色背景
   建议配合 JS/模板侧确认 message.bot 标记的及时性
*/
div.message.bot, div.message.bot.pending, div.message.bot.incomplete {
    background: linear-gradient(135deg, #b9ecfb 0%, #b7e7c7 100%) !important;
    # /* color: #198b6c !important; */
    border-left: 5px solid #26c6da !important;
    font-size: 19px !important;
    padding: 14px 20px !important;
    border-radius: 18px 18px 18px 8px !important;
    margin: 10px 0 !important;
    box-shadow: 0 3px 12px rgba(38, 198, 218, 0.18) !important;
    width: fit-content !important;
    word-break: break-word !important;
    display: inline-block !important;
    line-height: 1.6 !important;
    color: #22324a !important;
}

/* 用户消息（健康橙-暖色调，易辨识且亲切） */
div.message:not(.bot) {
    background: linear-gradient(135deg, #b9ecfb 0%, #b7e7c7 100%) !important;
    # color: #ff9800 !important;
    # border-right: 5px solid #ff9800 !important;
    font-size: 19px !important;
    padding: 14px 20px !important;
    border-radius: 18px 18px 8px 18px !important;
    margin-left: auto !important;
    margin: 10px 0 !important;
    # box-shadow: 0 3px 12px rgba(255, 152, 0, 0.15) !important;
    width: fit-content !important;
    min-width: 60px !important;
    max-width: 85% !important;
    word-break: break-word !important;
    display: inline-block !important;
    line-height: 1.6 !important;
}

/* 重要：防止权重冲突，避免 bot 消息被 !:not(.bot) 规则选中 */
div.message.bot {
    color: #22324a !important;
}



/* 输入框及主要交互区域高亮 */
input, textarea, .gradio-input, .gr-input {
    background: #f7fcfd !important;
    color: #22324a !important;
    border: 1.5px solid #26c6da !important;
    border-radius: 8px !important;
    box-shadow: 0 2px 7px rgba(38, 198, 218, 0.07) !important;
    padding: 10px 14px !important;
    font-size: 16px !important;
}

input:focus, textarea:focus, .gradio-input:focus, .gr-input:focus {
    outline: 2px solid #42a5f5 !important;
    border-color: #42a5f5 !important;
}



.header, .navbar, .logo {
    background: linear-gradient(90deg, #26c6da, #42a5f5) !important;
    color: #fff !important;
    border-bottom: 3px solid #fff !important;
}



/* 消息边框样式 */
.message-bubble-border {
    border: none !important;                /* 移除默认边框 */
    background: transparent !important;     /* 仅保留文字颜色, 无背景颜色 */
}


/* Chat box 背景样式 */
div#component-13.block.svelte-90oupt.padded {
    # background: linear-gradient(135deg, #f0f8ff 0%, #e6f3ff 100%) !important;
    border-radius: 12px !important;
    padding: 20px !important;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08) !important;
}

/* 针对 svelte 类名的备用方案 */
.svelte-1pjfiar {
    background-color: inherit !important;
}




button.svelte-izfbkb.center.boundedheight.flex {
    max-height: 110px !important;
    # max-width: 60px !important;
    content: "请拖拽文件到此区域或点击上传您想要的文件！"; /* 自定义文字 */
}



.gradio-file::before {
    content: "请拖拽文件到此区域或点击上传您想要的文件！";
    display: block;
    color: #2E7D32;
    font-size: 14px;
    font-weight: 500;
    margin: 10px 0;
}


/* 针对输入框的特定类名 */
.scroll-hide.svelte-1f354aw {
    width: 100% !important;
    height: 40px !important;
    resize: none !important;
    max-height: 40px !important;
    min-height: 40px !important;
    overflow: hidden !important;
}

.scroll-hide.svelte-1f354aw textarea {
    height: 40px !important;
    max-height: 40px !important;
    min-height: 40px !important;
    resize: none !important;
    overflow: hidden !important;
}

/* Textbox 和按钮的包装容器 */
.textbox-with-button {
    position: relative !important;
}

/* 提交按钮样式 - 绝对定位在输入框内 */
.submit-button {
    position: absolute !important;
    top: 50% !important;
    right: 8px !important;
    transform: translateY(-50%) !important;
    z-index: 10 !important;
    width: auto !important;
    height: auto !important;
}

/* 针对特定 ID 和类名的按钮 */
#component-18.lg.primary.submit-button.svelte-cmf5ev {
    position: absolute !important;
    top: 50% !important;
    right: 8px !important;
    transform: translateY(-50%) !important;
    z-index: 10 !important;
    width: auto !important;
    height: auto !important;
}


.submit-button button {
    background-color: #2E3192 !important;
    color: white !important;
    border-radius: 6px !important;
    height: 32px !important;
    padding: 6px 12px !important;
    font-weight: bold !important;
    border: none !important;
    transition: background-color 0.3s ease !important;
    font-size: 14px !important;
    min-width: 60px !important;
}



/* 针对特定按钮的样式 */
button#component-10.lg.primary.submit-button.svelte-cmf5ev {
    background-color: white !important; 
    color: #2E3192 !important;
    border-radius: 6px !important;
    height: 32px !important;
    padding: 6px 12px !important;
    font-weight: bold !important;
    border: none !important;
    transition: background-color 0.3s ease !important;
    font-size: 14px !important;
    min-width: 60px !important;
    line-height: 1 !important;
}

.submit-button button:hover {
    background-color: #1A1F71 !important;
}

#component-18.lg.primary.submit-button.svelte-cmf5ev button:hover {
    background-color: #1A1F71 !important;
}

.submit-button button:disabled {
    background-color: #D3D3D3 !important;
    color: #666666 !important;
    cursor: not-allowed !important;
}

#component-18.lg.primary.submit-button.svelte-cmf5ev button:disabled {
    background-color: #D3D3D3 !important;
    color: #666666 !important;
    cursor: not-allowed !important;
}

/* Upload & Start Chat 按钮样式 (component-6) */
button#component-20.lg.primary.upload-button.svelte-cmf5ev {
    background: linear-gradient(135deg, #B0B0B0 0%, #7D7D7D 100%) !important;
    color: white !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 12px 24px !important;
    font-size: 14px !important;
    font-weight: bold !important;
    height: 45px !important;
    min-width: 100px !important;
    box-shadow: 0 4px 15px rgba(76, 175, 80, 0.3) !important;
    transition: all 0.3s ease !important;
    cursor: pointer !important;
    text-transform: uppercase !important;
    letter-spacing: 0.5px !important;
}

button#component-20.lg.primary.upload-button.svelte-cmf5ev:hover {
    background: linear-gradient(135deg, #7D7D7D 0%, #B0B0B0 100%) !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(76, 175, 80, 0.4) !important;
}

button#component-20.lg.primary.upload-button.svelte-cmf5ev:active {
    transform: translateY(0) !important;
    box-shadow: 0 2px 10px rgba(76, 175, 80, 0.3) !important;
}

button#component-20.lg.primary.upload-button.svelte-cmf5ev:disabled {
    background: #E0E0E0 !important;
    color: #9E9E9E !important;
    cursor: not-allowed !important;
    transform: none !important;
    box-shadow: none !important;
}



/* Restart 按钮样式 (component-7) */
button#component-21.lg.secondary.svelte-cmf5ev {
    background: linear-gradient(135deg, #FF7043 0%, #FF5722 100%) !important;
    color: white !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 10px 20px !important;
    font-size: 14px !important;
    font-weight: bold !important;
    height: 45px !important;
    min-width: 100px !important;
    box-shadow: 0 4px 15px rgba(255, 112, 67, 0.3) !important;
    transition: all 0.3s ease !important;
    cursor: pointer !important;
    text-transform: uppercase !important;
    letter-spacing: 0.5px !important;
}


#component-21 button:hover {
    background: linear-gradient(135deg, #FF5722 0%, #FF7043 100%) !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(255, 112, 67, 0.4) !important;
}

#component-21 button:active {
    transform: translateY(0) !important;
    box-shadow: 0 2px 10px rgba(255, 112, 67, 0.3) !important;
}

/* Textbox 样式调整 - 为按钮留出空间 */
.textbox-with-button .scroll-hide.svelte-1f354aw {
    padding-right: 80px !important; /* 为按钮留出空间 */
}

.textbox-with-button .scroll-hide.svelte-1f354aw textarea {
    padding-right: 80px !important; /* 为按钮留出空间 */
}

.avatar-image.svelte-1pijsyv {
    width: 80px !important;
    height: 80px !important;
    border-radius: 50% !important; /* 保持为圆形 */
    border: 4px solid #cde3f9 !important; /* 可自定义边框颜色与粗细 */
    box-sizing: border-box !important;
    transition: width 0.2s, height 0.2s, border-width 0.2s; /* 平滑变化(可选) */
}

/* 允许父元素动态调整图片和边框大小，增加适应性，如果图片尺寸变化则边框同步变化 */
.avatar-container.svelte-1e1jlin .avatar-image.svelte-1pijsyv {
    width: 100% !important;
    height: 100% !important;
    max-width: 120px !important;   /* 限制最大尺寸，可根据需求调整 */
    max-height: 120px !important;
}

/* 可选: 父容器控制整体 avatar 尺寸，图片和边框会自适应填充 */
.avatar-container.svelte-1e1jlin {
    width: 45px !important;
    height: 45px !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
}
/* 取消头像的圆形边框，使图片为常规矩形且无边框 */
.avatar-image.svelte-1pijsyv {
    width: 90% !important;
    height: 90% !important;
    border-radius: 0 !important;
    border: none !important;
    box-sizing: border-box !important;
    transition: width 0.2s, height 0.2s, border-width 0.2s;
}

.avatar-container.svelte-1e1jlin div.avatar-image.svelte-1pijsyv {
    width: 55px !important;
    height: 55px !important;
}



/* 聊天框背景样式 */
div.bubble-wrap.svelte-1pjfiar {
    background: linear-gradient(135deg, #e6f3ff 0%, #f0f8ff 100%) !important;
    border-radius: 12px !important;
    padding: 20px !important;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08) !important;
    margin: 10px 0 !important;

    height: 520px !important;
    # overflow-y: auto !important;
}

div.wrapper.svelte-nab2ao {
    background: linear-gradient(135deg, #e6f3ff 0%, #f0f8ff 100%) !important;
    # border: 2px solid #e3f2fd !important;
    height: 510px !important;
}

div#component-6.block.custom-avatar.svelte-90oupt {
    height: 500px !important;
}

block#component-7.block.custom-avatar.svelte-12cmxck {
    height: 510px !important;
}


div.flex-wrap.bot.svelte-1e1jlin {
    border: none !important;
    /* 或者若需完全透明边框也可以如下： */
    /* border: 1px solid transparent !important; */
    box-shadow: none !important;
}


# /* 针对文件上传组件的固定高度与滚动条设置，及使文件上传中始终保持自定义边框样式 */
# #custom-file-upload,
# #custom-file-upload:focus,
# #custom-file-upload:active,
# #custom-file-upload[aria-busy="true"],
# div.custom-file-upload.block.svelte-90oupt,
# div.custom-file-upload.block.svelte-90oupt:focus,
# div.custom-file-upload.block.svelte-90oupt:active,
# div.custom-file-upload.block.svelte-90oupt[aria-busy="true"] {
#     max-height: 120px !important;  /* 限制最大高度 */
#     overflow-y: auto !important;   /* 内容超出时显示垂直滚动条 */
#     border-style: dashed !important; /* 始终保持虚线边框 */
#     border-width: 2px !important;
#     border-color: #90caf9 !important;
#     box-shadow: none !important;
# }


/* ============================ */
/* 聊天消息主气泡美化-消除双背景分层 */
/* ============================ */

/* 强制_message整体_单一渐变背景&无白色内层 */
div.message.bot, div.message.user, 
div.message.bot.pending, div.message.bot.incomplete,
div.message.user.pending, div.message.user.incomplete {
    /* 直接应用柔和青绿色医学渐变 */
    background: linear-gradient(135deg, #b9ecfb 0%, #b7e7c7 100%) !important;
    /* 移除 padding 产生的白色卡层 */
    padding: 18px 22px !important;
    border-radius: 15px !important;
    margin: 12px 0 !important;
    /* 阴影更柔和，避免高对比分割 */
    box-shadow: 0 2px 12px rgba(41, 141, 150, 0.08) !important;
    border: none !important;
}

/* 彻底去掉任何“内层”白色/默认背景：泡泡内容区、Markdown区 */
div.message.bot > div, 
div.message.user > div, 
div.message.bot .markdown, 
div.message.user .markdown,
div.message.bot .bubble-content, 
div.message.user .bubble-content
{
    background: transparent !important;
    box-shadow: none !important;
    border: none !important;
}

/* 用户消息气泡同一模式，但偏浅蓝些区分身份 */
div.message.user, div.message.user.pending, div.message.user.incomplete {
    background: linear-gradient(135deg, #e3f0fa 0%, #c8e2f7 100%) !important;
}

/* 去除可能附着在底下的多余盒子背景 */
div.svelte-1e1jlin, div.bubble-wrap, div.bubble-content, div.markdown-body {
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
}

/* 可选：调整文字清晰度、字号与高亮内容 */
div.message.bot, div.message.user {
    color: #243e3e !important;
    font-size: 1.08em !important;
    line-height: 1.55 !important;
    font-family: 'Segoe UI', 'Arial', 'PingFang SC', 'Microsoft Yahei', sans-serif !important;
}

div.message.bot strong, div.message.user strong {
    color: #1769aa !important;
    font-weight: 600 !important;
}

/* ✅ 精简视觉分界，让渐变自然延展到边缘 */
div.message.bot, div.message.user {
    border: none !important;
    /* 去除双边框 */
    outline: none !important;
}

/* 可选：让头像和泡泡贴得更紧凑，不留“白带” */
.avatar-container, .avatar-image {
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
    margin: 0 !important;
    padding: 0 !important;
}



# /* 聊天框滚动条样式 */
# div.bubble-wrap.svelte-1pjfiar::-webkit-scrollbar {
#     width: 6px !important;
# }

# div.bubble-wrap.svelte-1pjfiar::-webkit-scrollbar-track {
#     background: #f1f8ff !important;
#     border-radius: 3px !important;
# }

# div.bubble-wrap.svelte-1pjfiar::-webkit-scrollbar-thumb {
#     background: linear-gradient(135deg, #4CAF50 0%, #45a049 100%) !important;
#     border-radius: 3px !important;
# }

# div.bubble-wrap.svelte-1pjfiar::-webkit-scrollbar-thumb:hover {
#     background: linear-gradient(135deg, #45a049 0%, #4CAF50 100%) !important;
# }


.fixed-height-textbox textarea {
    resize: none !important;
    height: 40px !important;
    min-height: 40px !important;
    max-height: 40px !important;
}


/* 
   === 等待消息（pending/思考中）气泡样式：确保气泡出现在bot头像右侧 ===
   重点：当 role="status"（即 bot pending 消息）时，应靠右对齐，并紧贴 bot 头像
   兼容用户待发送气泡靠左、bot pending 靠右的“对话气泡”体验
*/

/* 基础样式：pending 状态的消息气泡 */
div.message.pending.svelte-1gpwetz {
    display: flex !important;
    align-items: center !important;
    min-height: 48px !important;
    box-shadow: none !important;
    border-radius: 14px !important;
    margin: 10px 0 !important;
    padding: 10px 20px !important;
    position: relative;
    /* 默认左对齐，但后面有更高优先级的覆盖 */
    justify-content: flex-start !important;
    background: #f0f7ff !important;
}

/* user（右侧/人类）等待消息：靠右 */
div.message.user.pending.svelte-1gpwetz,
div.message[role="user"].pending.svelte-1gpwetz {
    justify-content: flex-end !important;
    margin-left: auto !important;
    margin-right: 0 !important;
    background: #e3f2fd !important;
}

/* bot（左侧/机器人）等待消息：紧贴bot头像+显示bot头像 */
/* 
   假设 pending message 时 DOM 结构应包含和正常 bot 消息一致的头像区域（avatar-container/avatar-image），
   则只需调整 margin & flex 排布，确保pending气泡紧贴头像而不靠最左边。增加头像在pending状态下可见的条件样式。
*/
/* 
如何显示avatar的bot头像？
1. 确保 Gradio/前端 DOM 结构中， pending 状态的 bot 消息 (role="status"或role="assistant") DOM 内也包含 
   <div class="avatar-container"><img class="avatar-image" ...></div> 结构，和正常 bot 消息一致。
2. 样式上，只需保证 pending 泡泡与头像并排显示，不要用 display: none 隐藏头像容器。
3. 可通过如下样式让 pending bot 消息的 bubble 紧贴头像并确保头像显示。
*/

div.message.bot.pending.svelte-1gpwetz,
div.message[role="status"].pending.svelte-1gpwetz,
div.message[role="assistant"].pending.svelte-1gpwetz {
    display: flex !important;
    align-items: center !important;
    /* 让flex内第一个元素是avatar-container，第二个是bubble，则bubble就会紧贴avatar展示 */
    /* 如 Gradio DOM 正常包含头像，此样式即可显示bot pending时的头像 */
    justify-content: flex-start !important;
    margin-left: 20 !important;  /* 头像默认有宽度+外边距，无需再留额外边距 */
    margin-right: auto !important;
    background: #f0f8ff !important;
    max-width: 88% !important;

    
}



/* 移动端适配 */
@media (max-width: 600px) {
    div.message.pending.svelte-1gpwetz {
        padding: 8px 8px !important;
        min-height: 42px !important;
    }
}

/* 再确保优先级覆盖，实锤靠右 */
div.message.bot.pending.svelte-1gpwetz,
div.message[role="status"].pending.svelte-1gpwetz,
div.message[role="assistant"].pending.svelte-1gpwetz {
    justify-content: flex-end !important;
}

/* 兼容极端情况下 pending 消息内容溢出，保持 bubble 不挤压头像 */
div.message.bot.pending.svelte-1gpwetz,
div.message[role="status"].pending.svelte-1gpwetz,
div.message[role="assistant"].pending.svelte-1gpwetz {
    max-width: 88%; /* 留一点空间避免紧贴屏幕右侧，如需更靠近头像可微调 */
}

/* 缩减上传按钮高度（适当减少高度和内侧间距）*/
button.svelte-1x5qevo.center.boundedheight.flex[aria-label="file upload"],
button.svelte-1x5qevo.center.boundedheight.flex,
button.svelte-1x5qevo.center.boundedheight.flex > .wrap.svelte-12ioyct {
    min-height: 100px !important;
    height: 100px !important;         /* 原可能较高，改为较低高度 */
    padding-top: 4px !important;
    padding-bottom: 4px !important;
    box-sizing: border-box !important;
}
button.svelte-1x5qevo.center.boundedheight.flex {
    height: 44px !important;         /* 按钮整体高度缩减 */
    max-height: 48px !important;
}

button.svelte-1x5qevo.center.boundedheight.flex svg {
    max-height: 26px !important;
    max-width: 26px !important;
}

button.svelte-1x5qevo.center.boundedheight.flex span,
button.svelte-1x5qevo.center.boundedheight.flex .wrap.svelte-12ioyct {
    line-height: 1.2 !important;
    font-size: 1rem !important;
}

/* 缩小上传按钮内的input高度和间距 */
button.svelte-1x5qevo.center.boundedheight.flex input[type="file"] {
    min-height: 40px !important;
    height: 40px !important;
    padding: 0 !important;
    margin: 0 !important;
}

/* 如果 needed: 缩减“将文件拖放到此处 - 或 - 点击上传”的内容垂直边距 */
button.svelte-1x5qevo.center.boundedheight.flex .wrap.svelte-12ioyct {
    padding-top: 0 !important;
    padding-bottom: 0 !important;
}

/* 兼容边框圆角调整 */
button.svelte-1x5qevo.center.boundedheight.flex {
    border-radius: 9px !important;
}

.upload-fixed-height .svelte-1x5qevo.center.boundedheight.flex,
.upload-fixed-height button.svelte-1x5qevo.center.boundedheight.flex {
    min-height: 130px !important;
    max-height: 130px !important;
    height: 130px !important;
    box-sizing: border-box !important;
}

.upload-fixed-height {
    min-height: 110px !important;
    max-height: 110px !important;
    height: 110px !important;
    overflow-y: auto !important;
    box-sizing: border-box !important;
    /* 若想四个齿轮按钮/Drag区域完全等高则必须加这个保证外框也一致 */
}

.progress-text.svelte-au1olv.meta-text {
    font-size: 1.35rem !important;
    font-weight: 600 !important;
}


/* --- 解决气泡不贴合文字 --- */

/* 核心：取消强制全宽类名 */
.panel-full-width.svelte-1nr59td {
    width: fit-content !important;
    max-width: 90% !important;
    min-width: unset !important;
}

/* 确保内部消息内容也是自适应宽度 */
.bot.svelte-1nr59td.message [data-testid="bot"],
.user.svelte-1nr59td.message [data-testid="user"] {
    width: fit-content !important;
    display: inline-block !important;
}

/* 消息行容器：确保消息靠边对齐而不是撑开 */
.message-row.svelte-1nr59td {
    width: 100% !important;
    display: flex !important;
}

/* --- 解决多余滚动条 --- */

/* 1. 针对你自定义的 block 容器，隐藏它的滚动条，只保留内部滚动 */
.block.custom-avatar {
    overflow: hidden !important; 
}

# /* 2. 针对 Gradio 的内部包装器，确保它是唯一的滚动源 */
# .svelte-1wizwbi, .bubble-wrap.svelte-kpz1 {
#     overflow-y: auto !important;
#     height: 100% !important;
# }

div.svelte-1wizwbi {
    # overflow-y: auto !important;
    height: 100% !important;
    max-height: 600px !important;
}

/* 3. 彻底隐藏浏览器默认滚动条的视觉效果（可选，保持界面清爽） */
.block.custom-avatar::-webkit-scrollbar {
    display: none;
}



/* 1. 核心消息容器：取消全宽，改为根据内容收缩 */
.bot.svelte-1nr59td.message {
    width: fit-content !important;
    max-width: 85% !important; /* 限制最大宽度，防止长文本顶到头 */
    display: block !important;  /* 确保它不会默认占满整行 */
}

/* 2. 内部包装层：确保它们也是自适应宽度 */
.bot.svelte-1nr59td.message .message-wrap,
.bot.svelte-1nr59td.message .panel-full-width,
.bot.svelte-1nr59td.message [data-testid="bot"] {
    width: fit-content !important;
    min-width: unset !important;
    max-width: 100% !important;
}

/* 3. 消除 Markdown 内容区域的默认最小宽度 */
.bot.svelte-1nr59td.message .prose {
    width: fit-content !important;
    min-width: unset !important;
}

/* 4. 调整气泡内边距（Padding），让文字和边缘更紧凑 */
div.message.bot {
    padding: 8px 14px !important; /* 原来是 14px 20px，缩小后更紧凑 */
    line-height: 1.4 !important;  /* 稍微收缩行高 */
}



/* 1. 用户消息核心容器：改为自适应宽度并靠右 */
.user.svelte-1nr59td.message {
    width: fit-content !important;
    max-width: 85% !important;
    margin-left: auto !important; /* 关键：推向右侧 */
    margin-right: 0 !important;
    display: block !important;
}

/* 2. 内部层级处理：消除强制全宽 */
.user.svelte-1nr59td.message .message-wrap,
.user.svelte-1nr59td.message .panel-full-width,
.user.svelte-1nr59td.message [data-testid="user"] {
    width: fit-content !important;
    min-width: unset !important;
    max-width: 100% !important;
    margin-left: auto !important; /* 确保内部块也靠右 */
}

/* 3. 针对 Markdown 和气泡外观的微调 */
div.message.user {
    padding: 8px 14px !important; /* 紧凑内边距 */
    line-height: 1.4 !important;
    border-radius: 18px 18px 8px 18px !important; /* 保持用户气泡特有的圆角 */
    text-align: left !important; /* 气泡内部文字依然左对齐，方便阅读 */
}

/* 4. 消除可能导致撑开的 prose (Markdown) 最小宽度 */
.user.svelte-1nr59td.message .prose {
    width: fit-content !important;
    min-width: unset !important;
}
"""


# <div class="message pending svelte-1gpwetz" role="status" aria-label="Loading response" aria-live="polite" style="border-radius: var(--radius-xxl);"><span class="sr-only">Loading content</span> <div class="dot-flashing svelte-1gpwetz"></div>
# 	&nbsp;
# 	<div class="dot-flashing svelte-1gpwetz"></div>
# 	&nbsp;
# 	<div class="dot-flashing svelte-1gpwetz"></div></div> 

# h1 {
#     color: #2E3192; /* 标题颜色 */
# }
# p {
#     color: #666666; /* 副标题颜色 */
# }

# 
# .custom-avatar .avatar {
        # width: 40px;
        # height: 40px;
    # }
# 

custom_theme = gr.themes.Default(
    primary_hue="blue",  # 主色调为蓝色
    secondary_hue="gray",  # 次色调为灰色
    neutral_hue="gray",  # 中性色调
    text_size="lg",  # 文字大小适中
    radius_size="md",  # 圆角适中
    font=["Arial", "sans-serif"],  # 使用清晰字体
    spacing_size="md"  # 间距适中
)


def parse_args():
    parser = argparse.ArgumentParser(description="Demo")
    parser.add_argument("--cfg-path", default='./eval_configs/minigpt4_eval.yaml', help="path to configuration file.")
    parser.add_argument("--gpu-id", type=int, default=0, help="specify the gpu to load the model.")
    parser.add_argument(
        "--options",
        nargs="+",
        help="override some settings in the used config, the key-value pair "
        "in xxx=yyy format will be merged into config file (deprecate), "
        "change to --cfg-options instead.",
    )
    args = parser.parse_args()
    return args


def setup_seeds(config):
    seed = config.run_cfg.seed + get_rank()

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    cudnn.benchmark = False
    cudnn.deterministic = True


# ========================================
#             Model Initialization
# ========================================

conv_dict = {'pretrain_vicuna0': CONV_VISION_Vicuna0,
             'pretrain_llama2': CONV_VISION_LLama2}

print('Initializing Chat')
args = parse_args()
cfg = Config(args)

model_config = cfg.model_cfg
model_config.device_8bit = args.gpu_id
model_cls = registry.get_model_class(model_config.arch)


# breakpoint() 
# breakpoint()

model = model_cls.from_config(model_config).to('cuda')


import copy

# cfg_model_2 = cfg.model_cfg
# cfg_model_2['ckpt'] = '/home/ynwang/brain_llm/MiniGPT-4/eval_configs/minigpt4_eval_report.yaml'
model_config.ckpt = './checkpoint_1.pth'
model_2 = model_cls.from_config(model_config).to('cuda')

# model_2 = copy.deepcopy(model.to('cpu')).to('cuda:{}'.format(args.gpu_id+1))
# model.to('cuda:{}'.format(args.gpu_id))

# model_2 = model
CONV_VISION = conv_dict[model_config.model_type]

# vis_processor_cfg = cfg.datasets_cfg.radiopediaclass.vis_processor.train
vis_processor = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.RandomResizedCrop(
                    224,
                    scale=(0.5, 1.0),
                    interpolation=InterpolationMode.BICUBIC,
                ),
                # self.normalize,
            ]
        )

stop_words_ids = [[835], [2277, 29937]]
stop_words_ids = [torch.tensor(ids) for ids in stop_words_ids]
stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])

# breakpoint()
# model_2.load_state_dict(torch.load('/home/ynwang/brain_llm/MiniGPT-4/minigpt4/ckpt/20250611193/checkpoint_1.pth'),strict=False)
chat = Chat([model,model_2], vis_processor, device='cpu', stopping_criteria=stopping_criteria)
print('Initialization Finished')


# ========================================
#             Gradio Setting
# ========================================

def handle_upload(files, chat_state, img_list):
    # files 是 gr.File 返回的文件对象列表
    img_list = [file.name for file in files]  # 获取每个文件的路径
    # 提取文件名（不包含路径）
    file_names = [os.path.basename(file_path) for file_path in img_list]
    chat_state = f"Uploaded {len(file_names)} files: {', '.join(file_names)}"
    return chat_state, img_list, [["System", f"Uploaded files: {', '.join(file_names)}"]]



# def gradio_reset(chat_state, img_list):
#     # breakpoint()
#     if chat_state is not None:
#         chat_state.messages = []
#     if img_list is not None:
#         img_list = []
        
#     return gr.update(value=[]), gr.update(value=None, interactive=True), gr.update(placeholder='Please upload your image first', interactive=False),gr.update(value="Upload & Start Chat", interactive=True), gr.update(interactive=False), chat_state, img_list
# chatbot, image, text_input, upload_button, submit_button, chat_state, img_list

def encode_img(img_list, modalities):
    instruction, LR_image_list, HR_image_list, modalities = chat.encode_img(img_list, modalities)
    return instruction, LR_image_list, HR_image_list, modalities

def upload_img(chatbot,t1c_ax_file, t1_file, t2_file, fla_file, text_input, chat_state, model_type):
    # 允许5个文件输入为None，自动跳过为None的对象
    # modalities=['t1c ax','t1 ax','t2 ax','t2f ax','t1c sag']
    # 将五个文件list合并成一个新的list
    modalities = []
    files = []
    for file_list in [t1c_ax_file, t1_file, t2_file, fla_file]:
        if file_list is not None:
            # breakpoint()
            # 检查 file_list[0] 是否为压缩文件
            new_file_list = []
            for file in file_list:
                is_compressed = file.endswith('.zip') or file.endswith('.tar') or file.endswith('.tar.gz') or file.endswith('.tgz')
                if is_compressed:

                    modalities.append(file.name.split('/')[-1].split('.')[0])
                    if file.endswith('.zip'):
                        extract_dir = file.name.replace('.zip', '_extracted')
                        with zipfile.ZipFile(file.name, 'r') as zip_ref:
                            zip_ref.extractall(extract_dir)
                        new_file_list = [
                            os.path.join(extract_dir, sub_file)
                            for sub_file in zip_ref.namelist()
                            if not sub_file.endswith('/') 
                            and not sub_file.startswith('__MACOSX/')
                            and not sub_file.endswith('.DS_Store')
                            and os.path.isfile(os.path.join(extract_dir, sub_file))
                            and 'dcm' in sub_file.lower()  # 排除目录条目
                        ]

                        # new_file_list.extend(new_file_list)
                    elif file.endswith('.tar'):
                        extract_dir = file.name.replace('.tar', '_extracted')
                        with tarfile.open(file.name, 'r') as tar_ref:
                            tar_ref.extractall(extract_dir)
                        # 不能在with语句结束后再使用tar_ref对象；
                        # tar_ref.getnames()需要在with内部调用并保存，否则with结束后文件被关闭
                            tar_names = tar_ref.getnames()
                            new_file_list = [
                                os.path.join(extract_dir, sub_file)
                                for sub_file in tar_names
                                if not sub_file.endswith('/') 
                                and not sub_file.startswith('__MACOSX/')
                                and not sub_file.endswith('.DS_Store')
                                and os.path.isfile(os.path.join(extract_dir, sub_file))
                                and 'dcm' in sub_file.lower()  # 排除目录条目
                            ]
                            # breakpoint()
                        # file_list = new_file_list
                    elif file.endswith('.tar.gz'):
                        extract_dir = file.name.replace('.tar.gz', '_extracted')
                        with tarfile.open(file.name, 'r') as tar_ref:
                            tar_ref.extractall(extract_dir)
                        new_file_list = [
                            os.path.join(extract_dir, sub_file)
                            for sub_file in tar_ref.namelist()
                            if not sub_file.endswith('/') 
                            and not sub_file.startswith('__MACOSX/')
                            and not sub_file.endswith('.DS_Store')
                            and os.path.isfile(os.path.join(extract_dir, sub_file))
                            and 'dcm' in sub_file.lower()  # 排除目录条目
                        ]
                        # new_file_list.extend(new_file_list)


                    elif file.endswith('.tgz'):
                        extract_dir = file.name.replace('.tgz', '_extracted')
                        with tarfile.open(file.name, 'r') as tar_ref:
                            tar_ref.extractall(extract_dir)
                            new_file_list = [
                                os.path.join(extract_dir, sub_file)
                                for sub_file in tar_ref.namelist()
                                if not sub_file.endswith('/') 
                                and not sub_file.startswith('__MACOSX/')
                                and not sub_file.endswith('.DS_Store')
                                and os.path.isfile(os.path.join(extract_dir, sub_file))
                                and 'dcm' in sub_file.lower()  # 排除目录条目
                            ]

                    dcm_files = []
                    dcm_orders=[]
                    for dcm_file_path in new_file_list:
                        # print(dcm_file_path)
                        dcm_data = pydicom.dcmread(dcm_file_path,force=True)
                        if np.max(dcm_data.pixel_array) > 0:
                            img = dcm_data.pixel_array
                            img = ((img - np.min(img)) / (np.max(img) - np.min(img) + 1e-8) * 255).astype(np.uint8)
                            dcm_files.append(img)
                            dcm_orders.append(int(dcm_data.InstanceNumber))
                    if len(dcm_files) > 0:
                        sorted_indices = np.argsort(dcm_orders)
                        dcm_files = [dcm_files[i] for i in sorted_indices]
                        # dcm_orders = [dcm_orders[i] for i in sorted_indices]
                        dcm_files = np.stack(dcm_files, axis=0)
                        print('dcm_files.shape',dcm_files.shape)
                        # np.save('dcm_files.npy',dcm_files)


                        files.append(dcm_files)

                elif file.endswith('.nii.gz'):
                    modalities.append(file.name.split('/')[-1].split('.')[0])

                    files.append(nib.load(file).get_fdata().astype(np.uint8))
                
                elif file.endswith('npy'):
                    modalities.append(file.name.split('/')[-1].split('.')[0])
                    files.append(np.load(file))
                
            #             os.path.join(extract_dir, file)
            #             for file in zip_ref.namelist()
            #             if not file.endswith('/') 
            #             and not file.startswith('__MACOSX/')
            #             and not file.endswith('.DS_Store')
            #             and os.path.isfile(os.path.join(extract_dir, file))
            #             and 'dcm' in file.lower()  # 排除目录条目及非影像相关
            #         ]
            #         file_list = new_file_list

            #         # breakpoint()
            #     elif filename.endswith('.tar'):
            #         extract_dir = file_list[0].name.replace('.tar', '_extracted')
            #         with tarfile.open(file_list[0].name, 'r') as tar_ref:
            #             tar_ref.extractall(extract_dir)
            #         new_file_list = [
            #             os.path.join(extract_dir, file)
            #             for file in tar_ref.namelist()
            #             if not file.endswith('/') 
            #             and not file.startswith('__MACOSX/')
            #             and not file.endswith('.DS_Store')
            #             and os.path.isfile(os.path.join(extract_dir, file)) 
            #             and 'dcm' in file.lower()  # 排除目录条目
            #         ]
            #         file_list = new_file_list  # 
            #     elif filename.endswith('.tar.gz'):
            #         extract_dir = file_list[0].name.replace('.tar.gz', '_extracted')
            #         with tarfile.open(file_list[0].name, 'r') as tar_ref:
            #             tar_ref.extractall(extract_dir)
            #         new_file_list = [
            #             os.path.join(extract_dir, file)
            #             for file in tar_ref.namelist()
            #             if not file.endswith('/') 
            #             and not file.startswith('__MACOSX/')
            #             and not file.endswith('.DS_Store')
            #             and os.path.isfile(os.path.join(extract_dir, file)) 
            #             and 'dcm' in file.lower()  # 排除目录条目
            #         ]
            #         file_list = new_file_list  # 
            # # breakpoint()
            # if '.dcm' in file_list[0].lower():
            #     dcm_files = []
            #     dcm_orders=[]
            #     for dcm_file in file_list:
            #         dcm_data = pydicom.dcmread(dcm_file)
                    
            #         if np.max(dcm_data.pixel_array) > 0:
            #             dcm_files.append(dcm_data.pixel_array)
            #             dcm_orders.append(int(dcm_data.InstanceNumber))

            #         # breakpoint()
            #     # 将 list 中的 np.array 沿第0维拼接成一个ndarray
            #     # 根据dcm_orders的顺序对dcm_files和dcm_orders同时排序（升序）
            #     if len(dcm_files) > 0:
            #         sorted_indices = np.argsort(dcm_orders)
            #         print(sorted_indices)
            #         print(dcm_orders)
            #         dcm_files = [dcm_files[i] for i in sorted_indices]
            #         dcm_orders = [dcm_orders[i] for i in sorted_indices]
            #     dcm_files = np.stack(dcm_files, axis=0)
            #     print(dcm_files.shape)
            #     files.append(dcm_files)
            #     modalities.append(modality)
            # else:
            #     files.extend(file_list)
            #     for file in file_list:
            #         modalities.append(modality)
    # img_name_list = [os.path.basename(f.name) for f in valid_files]
    # print([os.path.abspath(f.name) for f in valid_files])
    if not files:
        return None, None, gr.update(interactive=True), gr.update(interactive=False), chat_state, None

    # chat_state = CONV_VISION.copy()

    # 只传递非None的文件对象和文件名

    img_list = []
    llm_message = chat.upload_img(files, chat_state, img_list)
    
    instruction, LR_image_list, HR_image_list, modalities = encode_img(img_list, modalities)
    # instruction, LR_image_list, HR_image_list, modalities = chat.encode_img(img_list, modalities)
    return chatbot, chat_state, instruction, LR_image_list, HR_image_list, modalities, gr.update(interactive=False),gr.update(interactive=False)

# t1c_ax_file, t1_file, t2_file, fla_file, t1c_file, text_input, chat_state, model_type
def gradio_ask(user_message, age,gender,chatbot, chat_state,model_type):
    # breakpoint()
    if model_type == 'Diagnosis':
        if age != '' and gender != '':
            user_message = f'Please make a diagnosis for this {age} years old {gender} patient based on the MRI files, providing MRI report.'
        else:
            user_message = f'Please make a diagnosis for this patient based on the MRI files, providing MRI report.'
    elif model_type == 'Diagnosis and Confidence':
        if age != '' and gender != '':
            user_message = f'Please make a diagnosis for this {age} years old {gender} patient based on the MRI files, providing MRI report and your confidence level.'
        else:
            user_message = f'Please make a diagnosis for this patient based on the MRI files, providing MRI report and your confidence level.'
    # breakpoint()
    chat_state = CONV_VISION.copy()
    # chat.ask(user_message, chat_state)
    if len(user_message) == 0:
        return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
    chat.ask(user_message, chat_state)
    chatbot = chatbot + [{'role': 'user', 'content': user_message}]
    return '', chatbot, chat_state,gr.update(interactive=False)


def button_wait():
    return gr.update(interactive=False),gr.update(interactive=False),gr.update(interactive=False)


@spaces.GPU(duration=180)
def gradio_answer(chatbot, chat_state,instruction, lr_img_list, hr_img_list, modalities,model_type,age='',gender=''):
    print('instruction',instruction)
    try:
        llm_message = chat.answer(conv=chat_state,
                              instruction=instruction,
                              lr_img_list=lr_img_list.to('cuda'),
                              hr_img_list=hr_img_list.to('cuda'),
                              model_type=model_type,
                              modalities=modalities,
                              age=age,
                              gender=gender,
                              max_new_tokens=300,
                              max_length=2000)
        chatbot =chatbot + [{'role': 'assistant', 'content': llm_message}]
    except Exception as e:
        print(f'Error in gradio_answer: {e}')
        chatbot =chatbot + [{'role': 'assistant', 'content': 'The data files type is not supported, please check the file shape, we only accept data files with shape (slice, W, H) or (slice, W, H, 3)'}]
    return chatbot, chat_state, lr_img_list, hr_img_list,gr.update(interactive=True),gr.update(interactive=True),gr.update(interactive=True)


title = """<h1 align="center">BrainVLM</h1>"""
description = """<h3>This is the demo of BrainVLM. Upload patient MRI files and start chatting</h3>"""
article = """
<p style="display: inline-block; margin-right: 2px; vertical-align: top;">
    <a href='https://github.com/Vision-CAIR/MiniGPT-4'><img src='https://img.shields.io/badge/Github-Code-blue' style="height: 20px; object-fit: contain;"></a>
</p>
"""
# initial_messages = [
#     # (f"This is BrainVLM, a large vision-language model for brain image understanding. Please upload patient MRI files to start the chat!😊",f"Question template: Please make a diagnosis for this patient based on the MRI files, providing MRI report and your confidence level.")
#     [None,f"This is BrainVLM, a large vision-language model for brain tumor image understanding. Please upload patient MRI files to start the chat!😊\n Using Pipeline: 1. Please upload patient MRI files via the upload module! 🏥 2. Input age and gender information of the patient(optional). 💬\n 3. Click the diagnosis button to get the diagnosis result! ✨ "],
#     # (f"This is user, try to upload patient MRI files to start the chat!😊",f"Using Pipeline: 1. Please upload patient MRI files via the upload module! 🏥 2. Click the Start Chat button to start the chat! 💬\n 3. Input your question and click the Send button to get the answer! ✨ \n Question template: Please make a diagnosis for this patient based on the MRI files, providing MRI report.")

# ]
# [user_message, None]
initial_messages = [
    gr.ChatMessage(role="assistant", content=(
        "This is BrainVLM, a large vision-language model for brain tumor image understanding. "
        "Please upload patient MRI files to start the chat!\n"
        "Using Pipeline: 1. Rename your file with modality and view name, e.g. 'T1c_ax', 'T1_cor', 'T2_sag', 'T2f_ax', 'T1c_sag', etc. \n"
        "2. Please upload patient MRI files via the upload module! \n"
        "3. Input age and gender information of the patient (optional).\n"
        "4. Click the diagnosis button to get the diagnosis result!\n"
        "5. Click the restart button to process new files!\n"
        "Attention: BrainVLM was trained based on 5 MRI sequences (T1c, T1, T2, FLAIR and another view of T1c).\n Full modality diagnosis (based on 5 sequences) is recommended. "
        "We accept missing modality diagnosis, however, the diagnosis result may not be very accurate!\n"
        "Support File type: .nii.gz, .npy and zip file with dcm slice files."
    ))
]

with gr.Blocks(theme=custom_theme,css=custom_css) as demo:
    # gr.Markdown(title)
    with gr.Row():
        with gr.Column():
            gr.Markdown(
                """
                <div style="width: 100%; display: flex; flex-direction: column; align-items: center; justify-content: center; position: relative;">
                    <div style="width: 100%; display: flex; align-items: center; justify-content: center; position: relative;">
                        <h1 style="
                            margin: 0 auto;
                            font-size: 2.3em;
                            line-height:1;
                            display: inline-block;
                            text-align: center;
                            width: auto;
                            flex-shrink: 0;
                        ">BrainVLM</h1>
                    </div>
                </div>
                """,
                elem_classes=["center-title"]
            )
                
        # print(os.getcwd())

        # breakpoint()
        # logo_path = os.path.join(os.getcwd(), "./icon/brain.png")
        # gr.Markdown(f'![Logo](/file={os.path.join(os.getcwd(), "./icon/HKU.jpg")})', elem_classes=["logo-img"])  # Logo 图片
        # html_content = f'<img src="/file={logo_path}" alt="预定义图片" style="max-width: 35px;margin-left: 10px;">'
        # gr.HTML(html_content)
    # gr.Markdown(description)
    # gr.Markdown(article)

    # with gr.Row():/
        # with gr.Column(scale=0.3, elem_classes=["custom-column"]):
        #     image = gr.File(label="Upload .nii.gz file",file_count="multiple", file_types=[".nii.gz"])
            
        #     upload_button = gr.Button(value="Start Chat", interactive=True, variant="primary",elem_classes=["upload-button"])
        #     clear = gr.Button("Restart")
            
        #     model_type = gr.Dropdown(
        #         choices=["Diagnosis", "Diagnosis and Confidence"],
        #         value="Diagnosis",
        #         label="Task Type",
        #         interactive=True
        #     )
    
    
    # ############################ 原始版本
    # with gr.Row():
    #     # First row - Chat section (80% height)
    #     # with gr.Column(scale=1):
    #     chat_state = gr.State()
    #     img_list = gr.State()
    #     with gr.Accordion("Chat box", open=True):
    #         chatbot = gr.Chatbot(initial_messages,label='BrainVLM',type='messages',avatar_images=["./icon/user.png","./icon/robot.png"],elem_classes="custom-avatar",avatar_position="top", height=600) # Increased height
    #         with gr.Row():
    #             with gr.Column(scale=1):
    #                 text_input = gr.Textbox(label='User', placeholder='Please upload patient nii.gz files first', interactive=False, container=False, lines=1, max_lines=1, elem_classes=["textbox-with-button"])
    #                 submit_button = gr.Button("Send", variant="primary", interactive=False, elem_classes=["submit-button"])
    #     instruction,LR_image_list,HR_image_list=gr.State(),gr.State(),gr.State()
    
    
    # with gr.Row():
        # First row - Chat section (80% height)
        # with gr.Column(scale=1):
    chat_state = gr.State()
    img_list = gr.State()


    with gr.Row():
        chatbot = gr.Chatbot(initial_messages,label='BrainVLM',avatar_images=[".\\icon\\user.png",".\\icon\\robot.png"],elem_classes="custom-avatar",height=500) # Increased height
    # with gr.Row():
    #     with gr.Column(scale=1):
    #         text_input = gr.Textbox(label='User', placeholder='Please upload patient nii.gz files first', interactive=False, container=False, lines=1, max_lines=1, elem_classes=["textbox-with-button"])
    #         submit_button = gr.Button("Send", variant="primary", interactive=False, elem_classes=["submit-button"])
    instruction,LR_image_list,HR_image_list,modalities=gr.State(),gr.State(),gr.State(),gr.State()

    # Second row - File upload section (20% height)
    with gr.Row():
        # with gr.Column(scale=1):
        # 增加自定义 CSS 类 upload-fixed-height，固定上传组件高度并滚动条
        with gr.Row(elem_classes=["upload-column"]):
            with gr.Column():
                # 方法二：通过自定义 Gradio 组件以支持“文件夹上传”（即 <input type="file" webkitdirectory>）。
                # 实现方式为嵌入自定义 HTML + JS + Gradio 的 gr.HTML() 和 gr.File() 配合。
                # 下面实现一个自定义上传按钮，可以让用户直接选取文件夹，所有内部文件自动上传：
                # 注意：当前 Gradio 没有直接内置支持 webkitdirectory 参数，需用 JS 处理。


                # 这个输入组件仍用于支持常规多文件上传
                t1c_ax_file = gr.File(
                    label="T1c MRI files", 
                    file_count="multiple",
                    file_types=[".nii", ".gz", ".nii.gz",'.npy','.zip','.tar'],
                    # file_types=[".nii", ".gz", ".nii.gz", ".dcm", ".npy"], 
                    interactive=True,
                    elem_id="custom-file-upload-t1c",
                    elem_classes=["upload-fixed-height"],
                )

                # 需要配合 gradio 回调，在 folder_upload_callback.click 时读取 window.folderUploadFiles → decode_files_from_js 并更新流程
                # 你可以在 gradio 接口绑定如下回调 (需放在 app/Gradio Blocks 构建后):
                # folder_upload_callback.click(fn=decode_files_from_js, inputs=None, outputs=YOUR_IMG_LIST_GRVAR)

                t1_file=gr.File(
                    label="T1 MRI files",
                    file_count="multiple",
                    file_types=[".nii", ".gz", ".nii.gz",'.npy','.zip','.tar'],
                    interactive=True,
                    elem_id="custom-file-upload-t1",
                    elem_classes=["upload-fixed-height"]
                )
        with gr.Row():
            with gr.Column():
                t2_file=gr.File(
                    label="T2 MRI files",
                    file_count="multiple",
                    # file_types=[".nii", ".gz", ".nii.gz",'.npy','.dcm'],
                    file_types=[".nii", ".gz", ".nii.gz",'.npy','.zip','.tar'],
                    interactive=True,
                    elem_id="custom-file-upload-t2",
                    elem_classes=["upload-fixed-height"]
                )

                fla_file=gr.File(
                    label="FLAIR MRI files",
                    file_count="multiple",
                    # file_types=[".nii", ".gz", ".nii.gz",'.npy','.dcm'],
                    file_types=[".nii", ".gz", ".nii.gz",'.npy','.zip','.tar'],
                    interactive=True,
                    elem_id="custom-file-upload-flair",
                    elem_classes=["upload-fixed-height"]
                )
        # with gr.Row():
        #     with gr.Column():
        #         t1c_cor_file=gr.File(
        #             label="T1c cor files",
        #             file_count="multiple",
        #             file_types=[".nii", ".gz", ".nii.gz",'.npy','.dcm'],
        #             interactive=True,
        #             elem_id="custom-file-upload-t1c-cor",
        #             elem_classes=["upload-fixed-height"]
        #         )

        #         t1c_sag_file=gr.File(
        #             label="T1c sag MRI files",
        #             file_count="multiple",
        #             file_types=[".nii", ".gz", ".nii.gz",'.npy','.dcm'],
        #             interactive=True,
        #             elem_id="custom-file-upload-t1c-sag",
        #             elem_classes=["upload-fixed-height"]
        #         )
        # 让 upload_button 以及下方所有与聊天相关的控件初始为 disabled/不可交互。
        # 通过 image.change 事件，用户选择了文件后，自动激活 start chat 按钮以及相关控件
        with gr.Column():
            Confidence_Diagnosis = gr.Button("Full modality Diagnosis",interactive=False,variant="primary",elem_classes=["upload-button"])
            Diagnosis = gr.Button(
                value="Missing Modality Diagnosis",
                interactive=False,
                variant="primary",
                elem_classes=["diagnosis-button"]
            )

            Restart = gr.Button(
                "Restart",
                interactive=False,
                variant="secondary",
                elem_classes=["restart-button"]
            )
        # with gr.Row():
        with gr.Column():
            # with gr.Column(scale=1):
            age_input = gr.Textbox(
                label="Age",
                placeholder="Enter patient age (optional)",
                interactive=True,
                lines=1,
                max_lines=1,
                elem_classes=["fixed-height-textbox"]
            )
            gender_input = gr.Textbox(
                label="Gender",
                placeholder="Enter patient gender (optional)",
                interactive=True,
                lines=1,
                max_lines=1,
                elem_classes=["fixed-height-textbox"]
            )
            # To ensure the 'fixed-height-textbox' class disables resizing and fixes height,
            # add the following CSS to your css config area:
            #

            # 注意：gr.Dropdown 组件可以直接在回调函数中使用，不需要转换为 State
            # 在回调函数的 inputs 列表中添加 model_type 即可获取用户选择的值
        
        # breakpoint()

        # with gr.Column(scale=2.7):
        #     chat_state = gr.State()
        #     img_list = gr.State()
        #     with gr.Accordion("Chat box", open=True):
        #         chatbot = gr.Chatbot(initial_messages,label='BrainVLM',type='messages',avatar_images=["./icon/brain.png","./icon/robot.png"],elem_classes="custom-avatar",avatar_position="top")
        #         with gr.Row():
        #             with gr.Column(scale=1):
        #                 text_input = gr.Textbox(label='User', placeholder='Please upload patient nii.gz files first', interactive=False, container=False, lines=1, max_lines=1, elem_classes=["textbox-with-button"])
        #                 submit_button = gr.Button("Send", variant="primary", interactive=False, elem_classes=["submit-button"])
        #     instruction,LR_image_list,HR_image_list=gr.State(),gr.State(),gr.State()
    # print('model_type.value',model_type.value)
    # breakpoint()
    # if image is not None:
    #     upload_button.interactive = True
    # 检测用户完成 image 上传：使用 Gradio 的 `.change` 事件，在文件有变动时使 upload_button、text_input、submit_button 变为可交互
    

    
    
    def on_image_change(files):
        # files: list of gr.File objects or None
        if files is not None and len(files) > 0:
            # 已选文件，允许按钮交互，提示可以输入
            return (
                gr.update(interactive=True),    # Diagnosis
                gr.update(interactive=True)    # Confidence_Diagnosis
                # gr.update(interactive=False,placeholder="Click Start Chat button"),  # text_input
                # gr.update(interactive=False)     # submit_button
            )
        else:
            # 未选文件，禁用按钮和输入框
            return (
                gr.update(interactive=False),
                gr.update(interactive=False)
                # gr.update(interactive=False, placeholder="Please upload patient nii.gz files first"),
                # gr.update(interactive=False)
            )

    # 当检测到 image 上传新文件后，先激活 Diagnosis/Confidence_Diagnosis 两按钮，再自动触发 Diagnosis.click (即 upload_img)
    def on_image_and_upload(files, age_val, gender_val, chat_state_val, model_type_val):
        # 激活按钮
        diagnosis_update, confidence_update = on_image_change(files)
        # 如果有文件，则自动调用 upload_img，相当于点击 Diagnosis 按钮
        if files is not None and len(files) > 0:
            # 自动触发 upload_img 的逻辑
            # breakpoint()
            chat_state, instruction,LR_image_list,HR_image_list,modalities=upload_img(files, age_val, chat_state_val, model_type_val)
            # 返回按钮状态更新 + upload_img 所有返回值(注意顺序对齐 outputs)
            # Diagnosis,Confidence_Diagnosis 为按钮状态，upload_img返回的顺序需与对应outputs对齐
            return (
                diagnosis_update,             # Diagnosis按钮状态
                confidence_update,            # Confidence_Diagnosis按钮状态
                chat_state,
                instruction,
                LR_image_list,
                HR_image_list,
                modalities               # 展开 upload_img 返回值
            )
        else:
            # 如果未上传文件，只返回按钮状态，其余输出为None
            n_upload_outputs = 8  # Diagnosis, Confidence_Diagnosis, image, age_input, gender_input, Diagnosis, Confidence_Diagnosis, chat_state, instruction, LR_image_list, HR_image_list
            return (diagnosis_update, confidence_update) + (None,) * 8  # 8个None作为其他output
    model_type=gr.State("Diagnosis")
    # 需要定义outputs序列，前2个为按钮，后面是upload_img返回的输出及所需参数
    # t1c_ax_file.change(
    #     on_image_and_upload,
    #     inputs=[t1c_ax_file, age_input, chat_state, model_type],
    #     outputs=[
    #         Diagnosis, Confidence_Diagnosis,
    #         # upload_img 所有返回（需对齐调用 upload_img 的 signature和互动）
    #         chat_state, instruction, LR_image_list, HR_image_list, modalities
    #     ],
    #     queue=False
    # )

    # 实时检测五个gr.File组件是否都上传了文件，并据此动态设置Confidence Diagnosis按钮
    def all_files_uploaded_modified(t1c_ax_file, t1_file, t2_file, fla_file):
        # 判断前四个必填都不为None且都有文件
        required_files = [t1c_ax_file, t1_file, t2_file, fla_file]
        required_all_uploaded = all(f is not None and (not hasattr(f, "__len__") or len(f) > 0) for f in required_files)
        return required_all_uploaded

    def update_confidence_interactive(t1c_ax_file, t1_file, t2_file, fla_file):
        # 返回的是gr.update(interactive=True/False)
        if all_files_uploaded_modified(t1c_ax_file, t1_file, t2_file, fla_file):
            return gr.update(interactive=True)
        else:
            return gr.update(interactive=False)

    # 实时绑定change事件到五个文件Gradio组件上，只要有变化就检测状态
    t1c_ax_file.change(
        fn=update_confidence_interactive,
        inputs=[t1c_ax_file, t1_file, t2_file, fla_file],
        outputs=Confidence_Diagnosis,
        queue=False,
        show_progress=False
    )
    t1_file.change(
        fn=update_confidence_interactive,
        inputs=[t1c_ax_file, t1_file, t2_file, fla_file],
        outputs=Confidence_Diagnosis,
        queue=False,
        show_progress=False
    )
    t2_file.change(
        fn=update_confidence_interactive,
        inputs=[t1c_ax_file, t1_file, t2_file, fla_file],
        outputs=Confidence_Diagnosis,
        queue=False,
        show_progress=False
    )
    fla_file.change(
        fn=update_confidence_interactive,
        inputs=[t1c_ax_file, t1_file, t2_file, fla_file],
        outputs=Confidence_Diagnosis,
        queue=False,
        show_progress=False
    )
    # t1c_file.change(
    #     fn=update_confidence_interactive,
    #     inputs=[t1c_ax_file, t1_file, t2_file, fla_file, t1c_file],
    #     outputs=Confidence_Diagnosis,
    #     queue=False,
    #     show_progress=False
    # )

  
    # 只要有任意一个 gr.File 上传文件，Diagnosis 按钮设为 interactive=True，否则为False
    def any_file_uploaded(t1c_ax_file, t1_file, t2_file, fla_file):
        files = [t1c_ax_file, t1_file, t2_file, fla_file]
        for f in files:
            # gr.File 有些gradio版本可能返回None或者[]，都视为“未上传”
            if f is not None:
                # 如果f为list（multiple），检查有内容；如果是单File对象/单项，也算
                if hasattr(f, "__len__"):
                    if len(f) > 0:
                        return True
                else:
                    return True
        return False

    def update_diagnosis_interactive(t1c_ax_file, t1_file, t2_file, fla_file):
        if any_file_uploaded(t1c_ax_file, t1_file, t2_file, fla_file):
            return gr.update(interactive=True)
        else:
            return gr.update(interactive=False)

    # 绑定change事件到五个文件组件
    t1c_ax_file.change(
        fn=update_diagnosis_interactive,
        inputs=[t1c_ax_file, t1_file, t2_file, fla_file],
        outputs=Diagnosis,
        queue=False,
        show_progress=False
    )
    t1_file.change(
        fn=update_diagnosis_interactive,
        inputs=[t1c_ax_file, t1_file, t2_file, fla_file],
        outputs=Diagnosis,
        queue=False,
        show_progress=False
    )
    t2_file.change(
        fn=update_diagnosis_interactive,
        inputs=[t1c_ax_file, t1_file, t2_file, fla_file],
        outputs=Diagnosis,
        queue=False,
        show_progress=False
    )
    fla_file.change(
        fn=update_diagnosis_interactive,
        inputs=[t1c_ax_file, t1_file, t2_file, fla_file],
        outputs=Diagnosis,
        queue=False,
        show_progress=False
    )
    
    # t1c_file.change(
    #     fn=update_diagnosis_interactive,
    #     inputs=[t1c_ax_file, t1_file, t2_file, fla_file, t1c_file],
    #     outputs=Diagnosis,
    #     queue=False,
    #     show_progress=False
    # )

    text_input=gr.State(' ')
    text_input_confidence=gr.State(' ')
    # Diagnosis.click(gradio_ask, [text_input, age_input, gender_input, chatbot, chat_state,gr.State('Diagnosis')], [text_input, chatbot, chat_state]).then(
    #     gradio_answer, [chatbot, chat_state,instruction, LR_image_list,HR_image_list, modalities, gr.State('Diagnosis')], [chatbot, chat_state, LR_image_list,HR_image_list]
    # )
    # Confidence_Diagnosis.click(upload_img, [t1c_ax_file, t1_file, t2_file, fla_file,t1c_cor_file,t1c_sag_file, age_input, chat_state, model_type], [chat_state, instruction, LR_image_list, HR_image_list, modalities,Confidence_Diagnosis,Diagnosis]).then(gradio_ask, [text_input_confidence, age_input, gender_input, chatbot, chat_state,gr.State('Diagnosis and Confidence')], [text_input_confidence, chatbot, chat_state,Confidence_Diagnosis]).then(
    #     gradio_answer, [chatbot, chat_state,instruction, LR_image_list,HR_image_list, modalities, gr.State('Diagnosis and Confidence')], [chatbot, chat_state, LR_image_list,HR_image_list,Confidence_Diagnosis,Diagnosis,Restart]
    # )
    # Confidence_Diagnosis.click(button_wait,[], [Confidence_Diagnosis,Diagnosis,Restart]).then(upload_img, [t1c_ax_file, t1_file, t2_file, fla_file, age_input, chat_state, model_type], [chat_state, instruction, LR_image_list, HR_image_list, modalities,Confidence_Diagnosis,Diagnosis]).then(gradio_ask, [text_input_confidence, age_input, gender_input, chatbot, chat_state,gr.State('Diagnosis and Confidence')], [text_input_confidence, chatbot, chat_state,Confidence_Diagnosis]).then(
    #     gradio_answer, [chatbot, chat_state,instruction, LR_image_list,HR_image_list, modalities, gr.State('Diagnosis and Confidence')], [chatbot, chat_state, LR_image_list,HR_image_list,Confidence_Diagnosis,Diagnosis,Restart]
    # )

    Confidence_Diagnosis.click(button_wait,[], [Confidence_Diagnosis,Diagnosis,Restart]).then(gradio_ask, [text_input_confidence, age_input, gender_input, chatbot, chat_state,gr.State('Diagnosis and Confidence')], [text_input_confidence, chatbot, chat_state,Confidence_Diagnosis]).then(upload_img, [chatbot,t1c_ax_file, t1_file, t2_file, fla_file, age_input, chat_state, model_type], [chatbot,chat_state, instruction, LR_image_list, HR_image_list, modalities,Confidence_Diagnosis,Diagnosis]).then(
        gradio_answer, [chatbot, chat_state,instruction, LR_image_list,HR_image_list, modalities, gr.State('Diagnosis and Confidence'),age_input, gender_input], [chatbot, chat_state, LR_image_list,HR_image_list,Confidence_Diagnosis,Diagnosis,Restart]
    )

    Diagnosis.click(button_wait,[], [Diagnosis,gr.State('Diagnosis and Confidence'),Restart]).then(gradio_ask, [text_input, age_input, gender_input, chatbot, chat_state,gr.State('Diagnosis and Confidence')], [text_input, chatbot, chat_state,Diagnosis]).then(upload_img, [chatbot,t1c_ax_file, t1_file, t2_file, fla_file, age_input, chat_state, model_type], [chatbot,chat_state, instruction, LR_image_list, HR_image_list, modalities,Diagnosis,Confidence_Diagnosis]).then(
        gradio_answer, [chatbot, chat_state,instruction, LR_image_list,HR_image_list, modalities, gr.State('Diagnosis and Confidence'),age_input, gender_input], [chatbot, chat_state, LR_image_list,HR_image_list,Diagnosis,gr.State('Diagnosis and Confidence'),Restart]
    )


    # Diagnosis.click(button_wait,[], [Diagnosis,gr.State('Diagnosis and Confidence'),Restart]).then(upload_img, [t1c_ax_file, t1_file, t2_file, fla_file, age_input, chat_state, model_type], [chat_state, instruction, LR_image_list, HR_image_list, modalities,Diagnosis,Confidence_Diagnosis]).then(gradio_ask, [text_input, age_input, gender_input, chatbot, chat_state,gr.State('Diagnosis and Confidence')], [text_input, chatbot, chat_state,Diagnosis]).then(
    #     gradio_answer, [chatbot, chat_state,instruction, LR_image_list,HR_image_list, modalities, gr.State('Diagnosis')], [chatbot, chat_state, LR_image_list,HR_image_list,Diagnosis,gr.State('Diagnosis and Confidence'),Restart]
    # )

    # Diagnosis.click(upload_img, [t1c_ax_file, t1_file, t2_file, fla_file,t1c_cor_file,t1c_sag_file, age_input, chat_state, model_type], [chat_state, instruction, LR_image_list, HR_image_list, modalities,Diagnosis,Confidence_Diagnosis]).then(gradio_ask, [text_input, age_input, gender_input, chatbot, chat_state,gr.State('Diagnosis')], [text_input, chatbot, chat_state,Diagnosis]).then(
    #     gradio_answer, [chatbot, chat_state,instruction, LR_image_list,HR_image_list, modalities, gr.State('Diagnosis and Confidence')], [chatbot, chat_state, LR_image_list,HR_image_list,Diagnosis,Confidence_Diagnosis,Restart]
    # )
    # Confidence_Diagnosis.click(gradio_ask, [text_input_confidence, age_input, gender_input, chatbot, chat_state,gr.State('Diagnosis and Confidence')], [text_input_confidence, chatbot, chat_state,Confidence_Diagnosis]).then(
    #     gradio_answer, [chatbot, chat_state,instruction, LR_image_list,HR_image_list, modalities, gr.State('Diagnosis and Confidence')], [chatbot, chat_state, LR_image_list,HR_image_list,Confidence_Diagnosis]
    # )   
    
    def gradio_reset():
        # breakpoint()
        # 清空 /tmp/gradio/ 目录下的所有文件            
        import os
        import shutil

        gradio_tmp_dir = "/tmp/gradio/"
        if os.path.exists(gradio_tmp_dir):
            for filename in os.listdir(gradio_tmp_dir):
                file_path = os.path.join(gradio_tmp_dir, filename)
                try:
                    if os.path.isfile(file_path) or os.path.islink(file_path):
                        os.unlink(file_path)
                    elif os.path.isdir(file_path):
                        shutil.rmtree(file_path)
                except Exception as e:
                    print(f'Failed to delete {file_path}. Reason: {e}')
        return gr.update(value=None, interactive=True), gr.update(value=[]), gr.update(value=[])
    

    def gradio_reset_all():
        # 清空 /tmp/gradio/ 目录下的所有文件
        import os
        import shutil

        gradio_tmp_dir = "/tmp/gradio/"
        if os.path.exists(gradio_tmp_dir):
            for filename in os.listdir(gradio_tmp_dir):
                file_path = os.path.join(gradio_tmp_dir, filename)
                try:
                    if os.path.isfile(file_path) or os.path.islink(file_path):
                        os.unlink(file_path)
                    elif os.path.isdir(file_path):
                        shutil.rmtree(file_path)
                except Exception as e:
                    print(f'Failed to delete {file_path}. Reason: {e}')
        return gr.update(value=[]), gr.update(value=[]), gr.update(value=[]), gr.update(value=[]), gr.update(value=[]), gr.update(value=[]), gr.update(value=[]), gr.update(value=[]), gr.update(value=[]), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False)

    # t1c_ax_file.clear(gradio_reset,[], [img_list,LR_image_list,HR_image_list])
    # 插入清除按钮和重置逻辑
    

    # 定义一个重置函数，重置所有state和控件到初始状态
    Restart.click(gradio_reset_all,[], [t1c_ax_file, t1_file, t2_file, fla_file, chat_state, instruction, LR_image_list, HR_image_list, modalities, Diagnosis, Confidence_Diagnosis,Restart])

demo.launch(share=True,allowed_paths=['./icon'])