import streamlit as st
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torch
import time
import io
import base64

# Cache the model globally
MODEL = None
PROCESSOR = None

# Embedded knowledge base
# Knowledge base
KNOWLEDGE_BASE = {
    "spalling": [
        {
            "severity": "Critical",
            "description": "Severe concrete spalling with exposed reinforcement and section loss",
            "repair_method": ["Install temporary support", "Remove deteriorated concrete", "Clean and treat reinforcement", "Apply corrosion inhibitor", "Apply bonding agent", "High-strength repair mortar"],
            "estimated_cost": "Very High ($15,000+)",
            "timeframe": "3-4 weeks",
            "location": "Primary structural elements",
            "required_expertise": "Structural Engineer + Specialist Contractor",
            "immediate_action": "Evacuate area, install temporary support, prevent access",
            "prevention": "Regular inspections, waterproofing, chloride protection"
        },
        {
            "severity": "High",
            "description": "Surface spalling with visible reinforcement",
            "repair_method": ["Remove damaged concrete", "Treat reinforcement", "Apply repair mortar", "Surface treatment"],
            "estimated_cost": "High ($8,000-$15,000)",
            "timeframe": "2-3 weeks",
            "location": "Structural elements",
            "required_expertise": "Structural Engineer",
            "immediate_action": "Area isolation, temporary support assessment",
            "prevention": "Protective coatings, drainage improvement"
        }
    ],
    "reinforcement_corrosion": [
        {
            "severity": "Critical",
            "description": "Severe corrosion with >30% section loss",
            "repair_method": ["Structural support", "Remove concrete", "Replace reinforcement", "Corrosion protection", "Concrete repair"],
            "estimated_cost": "Critical ($20,000+)",
            "timeframe": "4-6 weeks",
            "location": "Load-bearing elements",
            "required_expertise": "Senior Structural Engineer",
            "immediate_action": "Immediate evacuation, emergency shoring",
            "prevention": "Waterproofing, cathodic protection"
        }
    ],
    "structural_crack": [
        {
            "severity": "High",
            "description": "Cracks >5mm in structural elements",
            "repair_method": ["Structural analysis", "Epoxy injection", "Carbon fiber reinforcement", "Crack monitoring"],
            "estimated_cost": "High ($10,000-$20,000)",
            "timeframe": "2-4 weeks",
            "location": "Primary structure",
            "required_expertise": "Structural Engineer",
            "immediate_action": "Install crack monitors, load restriction",
            "prevention": "Load management, joint maintenance"
        }
    ],
    "dampness": [
        {
            "severity": "Medium",
            "description": "Active water penetration with efflorescence",
            "repair_method": ["Water source identification", "Drainage improvement", "Waterproof membrane", "Ventilation"],
            "estimated_cost": "Medium ($5,000-$10,000)",
            "timeframe": "1-2 weeks",
            "location": "Various",
            "required_expertise": "Waterproofing Specialist",
            "immediate_action": "Dehumidification, efflorescence cleaning",
            "prevention": "Proper drainage, vapor barriers"
        }
    ],
    "no_damage": [
        {
            "severity": "Low",
            "description": "No significant structural issues",
            "repair_method": ["Regular inspection", "Preventive maintenance"],
            "estimated_cost": "Low ($500-$2,000)",
            "timeframe": "1-2 days",
            "location": "General",
            "required_expertise": "Building Inspector",
            "immediate_action": "Continue monitoring",
            "prevention": "Regular maintenance schedule"
        }
    ]
}

DAMAGE_TYPES = {
    0: {'name': 'spalling', 'risk': 'High', 'color': '#ff4d4d'},
    1: {'name': 'reinforcement_corrosion', 'risk': 'Critical', 'color': '#800000'},
    2: {'name': 'structural_crack', 'risk': 'High', 'color': '#ff6b6b'},
    3: {'name': 'dampness', 'risk': 'Medium', 'color': '#4dabf7'},
    4: {'name': 'no_damage', 'risk': 'Low', 'color': '#40c057'}
}

def init_session_state():
    if 'history' not in st.session_state:
        st.session_state.history = []
    if 'dark_mode' not in st.session_state:
        st.session_state.dark_mode = False

@st.cache_resource
def load_model():
    try:
        model = ViTForImageClassification.from_pretrained(
            "google/vit-base-patch16-224",
            num_labels=len(DAMAGE_TYPES),
            ignore_mismatched_sizes=True
        )
        processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
        return model, processor
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None, None

def analyze_damage(image, model, processor):
    try:
        image = image.convert('RGB')
        inputs = processor(images=image, return_tensors="pt")
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
        return probs
    except Exception as e:
        st.error(f"Error analyzing image: {str(e)}")
        return None

def get_custom_css():
    return """
    <style>
    .main {
        padding: 2rem;
    }
    .stProgress > div > div > div > div {
        background-image: linear-gradient(to right, var(--progress-color, #ff6b6b), var(--progress-color-end, #f06595));
    }
    .damage-card {
        padding: 1.5rem;
        border-radius: 0.5rem;
        background: var(--card-bg, #f8f9fa);
        margin-bottom: 1rem;
        border: 1px solid var(--border-color, #dee2e6);
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    .damage-header {
        font-size: 1.25rem;
        font-weight: bold;
        margin-bottom: 1rem;
        color: var(--text-color, #212529);
    }
    .dark-mode {
        background-color: #1a1a1a;
        color: #ffffff;
    }
    .dark-mode .damage-card {
        background: #2d2d2d;
        border-color: #404040;
    }
    .tooltip {
        position: relative;
        display: inline-block;
        border-bottom: 1px dotted #ccc;
    }
    .tooltip .tooltiptext {
        visibility: hidden;
        background-color: #555;
        color: #fff;
        text-align: center;
        border-radius: 6px;
        padding: 5px;
        position: absolute;
        z-index: 1;
        bottom: 125%;
        left: 50%;
        margin-left: -60px;
        opacity: 0;
        transition: opacity 0.3s;
    }
    </style>
    """

def display_header():
    st.markdown(
        """
        <div style='text-align: center; padding: 1rem;'>
            <h1>🏗️ Smart Construction Defect Analyzer</h1>
            <p style='font-size: 1.2rem;'>Advanced AI-powered Construction Defect tool</p>
        </div>
        """,
        unsafe_allow_html=True
    )

def main():
    init_session_state()
    st.set_page_config(
        page_title="Smart Construction Defect Analyzer",
        page_icon="🏗️",
        layout="wide",
        initial_sidebar_state="expanded"
    )

    st.markdown(get_custom_css(), unsafe_allow_html=True)
    
    # Sidebar
    with st.sidebar:
        st.markdown("### ⚙️ Settings")
        st.session_state.dark_mode = st.toggle("Dark Mode", st.session_state.dark_mode)
        st.markdown("### 📖 Analysis History")
        if st.session_state.history:
            for item in st.session_state.history[-5:]:
                st.markdown(f"- {item}")
    
    display_header()

    # Load model
    global MODEL, PROCESSOR
    if MODEL is None or PROCESSOR is None:
        with st.spinner("Loading AI model..."):
            MODEL, PROCESSOR = load_model()
            if MODEL is None:
                st.error("Failed to load model. Please refresh the page.")
                return

    # File upload with drag and drop
    uploaded_file = st.file_uploader(
        "Drag and drop or click to upload an image",
        type=['jpg', 'jpeg', 'png'],
        help="Supported formats: JPG, JPEG, PNG"
    )

    if uploaded_file:
        try:
            # Image display and analysis
            image = Image.open(uploaded_file)
            col1, col2 = st.columns([1, 1])
            
            with col1:
                st.image(image, caption="Uploaded Structure", use_container_width =True)
            
            with col2:
                with st.spinner("🔍 Analyzing damage..."):
                    start_time = time.time()
                    predictions = analyze_damage(image, MODEL, PROCESSOR)
                    analysis_time = time.time() - start_time
                    
                    if predictions is not None:
                        st.markdown("### 📊 Analysis Results")
                        st.markdown(f"*Analysis completed in {analysis_time:.2f} seconds*")
                        
                        detected = False
                        for idx, prob in enumerate(predictions):
                            confidence = float(prob) * 100
                            if confidence > 15:
                                detected = True
                                damage_type = DAMAGE_TYPES[idx]['name']
                                cases = KNOWLEDGE_BASE[damage_type]
                                
                                with st.expander(f"{damage_type.replace('_', ' ').title()} - {confidence:.1f}%", expanded=True):
                                    # Progress bar with custom color
                                    st.markdown(
                                        f"""
                                        <style>
                                        .stProgress > div > div > div > div {{
                                            background-color: {DAMAGE_TYPES[idx]['color']} !important;
                                        }}
                                        </style>
                                        """,
                                        unsafe_allow_html=True
                                    )
                                    st.progress(confidence / 100)
                                    
                                    tabs = st.tabs(["📋 Details", "🔧 Repairs", "⚠️ Actions"])
                                    
                                    with tabs[0]:
                                        for case in cases:
                                            st.markdown(f"""
                                            - **Severity:** {case['severity']}
                                            - **Description:** {case['description']}
                                            - **Location:** {case['location']}
                                            - **Required Expertise:** {case['required_expertise']}
                                            """)
                                    
                                    with tabs[1]:
                                        for step in cases[0]['repair_method']:
                                            st.markdown(f"✓ {step}")
                                        st.info(f"**Estimated Cost:** {cases[0]['estimated_cost']}")
                                        st.info(f"**Timeframe:** {cases[0]['timeframe']}")
                                    
                                    with tabs[2]:
                                        st.warning("**Immediate Actions Required:**")
                                        st.markdown(cases[0]['immediate_action'])
                                        st.success("**Prevention Measures:**")
                                        st.markdown(cases[0]['prevention'])
                        
                        if not detected:
                            st.info("No significant Construction Defect detected. Regular maintenance recommended.")
                        
                        # Add to history
                        st.session_state.history.append(f"Analyzed image: {uploaded_file.name}")

        except Exception as e:
            st.error(f"Error processing image: {str(e)}")
            st.info("Please try uploading a different image.")

    # Footer
    st.markdown("---")
    st.markdown(
        """
        <div style='text-align: center'>
            <p>🏗️ Smart Construction Defect Analyzer | Built with Streamlit & Transformers</p>
            <p style='font-size: 0.8rem;'>For professional use only. Always consult with a structural engineer.</p>
        </div>
        """,
        unsafe_allow_html=True
    )

if __name__ == "__main__":
    main()