import streamlit as st
import subprocess
import os
import time

st.title('Swin-Unet Model Testing')

st.write('Click the button below to test the model on the test set.')

# Initialize session state for test status
if 'test_completed' not in st.session_state:
    st.session_state.test_completed = False

if st.button('Test Model'):
    st.session_state.test_completed = False
    st.write('Testing in progress...')
    
    # Clear previous overlays
    overlays_dir = os.path.join('overlays')
    if os.path.exists(overlays_dir):
        for file in os.listdir(overlays_dir):
            os.remove(os.path.join(overlays_dir, file))
    
    # Define the command to run the test script with correct paths
    command = [
        'python', 'src/test_isic.py',
        '--dataset', 'ISIC',
        '--cfg', 'src/configs/swin_tiny_patch4_window7_224_lite.yaml',
        '--is_savenii',
        '--root_path', 'src/datasets/test_npz',
        '--output_dir', 'src/outputs',
        '--max_epochs', '150',
        '--base_lr', '0.05',
        '--img_size', '224',
        '--batch_size', '24',
        '--list_dir', 'src/lists/ISIC'
    ]
    # Run the command and capture output
    result = subprocess.run(command, capture_output=True, text=True)
    
    # Display only the metrics from the output
    output_lines = result.stdout.split('\n')
    metrics_section = False
    metrics = []
    
    for line in output_lines:
        if "Test Results:" in line:
            metrics_section = True
            continue
        if metrics_section and line.strip():
            metrics.append(line.strip())
    
    if metrics:
        st.subheader('Test Results:')
        for metric in metrics:
            st.text(metric)
    
    # Filter out progress bar output from stderr
    error_lines = [line for line in result.stderr.split('\n') 
                  if line.strip() and not any(x in line for x in ['it/s', '?it/s', 's/it'])]
    if error_lines:
        st.subheader('Errors:')
        st.code('\n'.join(error_lines))
    
    st.session_state.test_completed = True

# Display overlays only after test is completed
if st.session_state.test_completed:
    overlays_dir = os.path.join('overlays')
    if os.path.exists(overlays_dir):
        overlay_files = [f for f in os.listdir(overlays_dir) if f.endswith('.png')]
        if overlay_files:
            st.subheader('Generated Overlays:')
            cols = st.columns(3)
            for i, overlay_file in enumerate(overlay_files[:3]):  # Show only first 3 overlays
                with cols[i]:
                    st.image(os.path.join(overlays_dir, overlay_file), caption=overlay_file)