import streamlit as st
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torch
import json

# Embedded knowledge base
KNOWLEDGE_BASE = {
    "spalling": [
        {
            "severity": "High",
            "description": "Severe concrete spalling with exposed reinforcement",
            "repair_method": ["Remove damaged concrete", "Clean exposed reinforcement", "Apply rust inhibitor", "Apply bonding agent", "Patch with repair mortar"],
            "estimated_cost": "High ($5,000-$10,000)",
            "timeframe": "2-3 weeks",
            "location": "Column/Beam",
            "required_expertise": "Structural Engineer",
            "immediate_action": "Area isolation and temporary support if structural",
            "prevention": "Regular maintenance, waterproofing, proper concrete cover"
        },
        {
            "severity": "Medium",
            "description": "Surface spalling without exposed reinforcement",
            "repair_method": ["Remove loose concrete", "Clean surface", "Apply repair mortar", "Surface treatment"],
            "estimated_cost": "Medium ($2,000-$5,000)",
            "timeframe": "1-2 weeks",
            "location": "Non-structural elements",
            "required_expertise": "Concrete Repair Specialist",
            "immediate_action": "Mark affected areas and monitor",
            "prevention": "Surface coating, regular inspections"
        }
    ],
    "reinforcement_corrosion": [
        {
            "severity": "Critical",
            "description": "Advanced corrosion with significant section loss",
            "repair_method": ["Install temporary support", "Remove damaged concrete", "Replace/supplement reinforcement", "Apply corrosion inhibitor", "Reconstruct concrete cover"],
            "estimated_cost": "Very High ($15,000+)",
            "timeframe": "3-4 weeks",
            "location": "Primary structural elements",
            "required_expertise": "Structural Engineer, Specialist Contractor",
            "immediate_action": "Immediate area evacuation and temporary support",
            "prevention": "Waterproofing, crack sealing, chloride protection"
        }
    ],
    "structural_crack": [
        {
            "severity": "High",
            "description": "Load-bearing element cracks >3mm",
            "repair_method": ["Structural assessment", "Crack injection with epoxy", "External reinforcement if needed", "Monitor crack progression"],
            "estimated_cost": "High ($8,000-$15,000)",
            "timeframe": "2-3 weeks",
            "location": "Load-bearing walls/beams",
            "required_expertise": "Structural Engineer",
            "immediate_action": "Install crack gauges, temporary support",
            "prevention": "proper design, load management, movement joints"
        }
    ],
    "dampness": [
        {
            "severity": "Medium",
            "description": "Persistent moisture penetration",
            "repair_method": ["Identify water source", "Install drainage system", "Apply waterproofing membrane", "Improve ventilation"],
            "estimated_cost": "Medium ($3,000-$7,000)",
            "timeframe": "1-2 weeks",
            "location": "Walls, floors",
            "required_expertise": "Waterproofing Specialist",
            "immediate_action": "Improve ventilation, dehumidification",
            "prevention": "Proper drainage, regular maintenance of water systems"
        }
    ],
    "no_damage": [
        {
            "severity": "Low",
            "description": "No visible structural issues",
            "repair_method": ["Regular inspection", "Preventive maintenance", "Document condition"],
            "estimated_cost": "Low ($500-$1,000)",
            "timeframe": "1-2 days",
            "location": "General structure",
            "required_expertise": "Building Inspector",
            "immediate_action": "Continue regular maintenance",
            "prevention": "Maintain inspection schedule"
        }
    ]
}

DAMAGE_TYPES = {
    0: {'name': 'spalling', 'risk': 'High'},
    1: {'name': 'reinforcement_corrosion', 'risk': 'Critical'},
    2: {'name': 'structural_crack', 'risk': 'High'},
    3: {'name': 'dampness', 'risk': 'Medium'},
    4: {'name': 'no_damage', 'risk': 'Low'}
}

@st.cache_resource
def load_model():
    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

def analyze_damage(image, model, processor):
    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

def main():
    st.set_page_config(
        page_title="Structural Damage Analyzer",
        page_icon="🏗️",
        layout="wide",
        initial_sidebar_state="expanded"
    )

    # Custom CSS
    st.markdown("""
        <style>
        .main {
            padding: 2rem;
        }
        .stProgress > div > div > div > div {
            background-image: linear-gradient(to right, #ff6b6b, #f06595);
        }
        .damage-card {
            padding: 1.5rem;
            border-radius: 0.5rem;
            background: #f8f9fa;
            margin-bottom: 1rem;
        }
        </style>
    """, unsafe_allow_html=True)

    col1, col2, col3 = st.columns([1, 3, 1])
    with col2:
        st.title("🏗️ Structural Damage Analyzer")
        st.markdown("### Upload a photo of structural damage for instant analysis")

    model, processor = load_model()
    
    upload_col1, upload_col2, upload_col3 = st.columns([1, 2, 1])
    with upload_col2:
        uploaded_file = st.file_uploader("Choose an image file", type=['jpg', 'jpeg', 'png'])

    if uploaded_file:
        image = Image.open(uploaded_file)
        
        analysis_col1, analysis_col2 = st.columns(2)
        
        with analysis_col1:
            st.image(image, caption="Uploaded Structure", use_column_width=True)
        
        with analysis_col2:
            with st.spinner("🔍 Analyzing structural damage..."):
                predictions = analyze_damage(image, model, processor)
                
                st.markdown("### 📊 Damage Assessment Results")
                
                for idx, prob in enumerate(predictions):
                    confidence = float(prob) * 100
                    if confidence > 15:
                        damage_type = DAMAGE_TYPES[idx]['name']
                        cases = KNOWLEDGE_BASE[damage_type]
                        
                        st.markdown(f"""
                        <div class="damage-card">
                            <h4>{damage_type.replace('_', ' ').title()} Detected</h4>
                        </div>
                        """, unsafe_allow_html=True)
                        
                        col1, col2 = st.columns([3, 1])
                        with col1:
                            st.progress(confidence / 100)
                        with col2:
                            st.write(f"Confidence: {confidence:.1f}%")
                        
                        tabs = st.tabs(["Details", "Repair Methods", "Recommendations"])
                        
                        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]:
                            st.markdown("#### Repair Steps")
                            for step in cases[0]['repair_method']:
                                st.markdown(f"- {step}")
                            st.markdown(f"""
                            - **Estimated Cost:** {cases[0]['estimated_cost']}
                            - **Timeframe:** {cases[0]['timeframe']}
                            """)
                        
                        with tabs[2]:
                            st.markdown("#### Immediate Actions")
                            st.warning(cases[0]['immediate_action'])
                            
                            st.markdown("#### Prevention Measures")
                            st.info(cases[0]['prevention'])

    # Footer
    st.markdown("---")
    st.markdown("""
        <div style='text-align: center'>
            <p>🏗️ Structural Damage Analysis Tool | Built with Streamlit</p>
        </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()