from streamlit_extras.colored_header import colored_header
from streamlit_extras.add_vertical_space import add_vertical_space
from PIL import Image
import numpy as np
from transformers import ViTFeatureExtractor, ViTForImageClassification
from sentence_transformers import SentenceTransformer
import streamlit as st
import torch
import matplotlib.pyplot as plt
import logging
import faiss
from typing import List, Dict
from datetime import datetime
from groq import Groq
import os
from functools import lru_cache
import time
from streamlit_card import card
import plotly.graph_objects as go

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class RAGSystem:
    def __init__(self):
        # Load models only when needed
        self._embedding_model = None
        self._vector_store = None
        self._knowledge_base = None

    @property
    def embedding_model(self):
        if self._embedding_model is None:
            self._embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        return self._embedding_model

    @property
    def knowledge_base(self):
        if self._knowledge_base is None:
            self._knowledge_base = self.load_knowledge_base()
        return self._knowledge_base

    @property
    def vector_store(self):
        if self._vector_store is None:
            self._vector_store = self.create_vector_store()
        return self._vector_store

    @staticmethod
    @lru_cache(maxsize=1)  # Cache the knowledge base
    def load_knowledge_base() -> List[Dict]:
        """Load and preprocess knowledge base"""
        kb = {
            "spalling": [
                {
                    "severity": "Critical",
                    "description": "Severe concrete spalling with exposed reinforcement",
                    "repair_method": "Remove deteriorated concrete, clean reinforcement",
                    "immediate_action": "Evacuate area, install support",
                    "prevention": "Regular inspections, waterproofing"
                }
            ],
            "structural_cracks": [
                {
                    "severity": "High",
                    "description": "Active structural cracks >5mm width",
                    "repair_method": "Structural analysis, epoxy injection",
                    "immediate_action": "Install crack monitors",
                    "prevention": "Regular monitoring, load management"
                }
            ],
            "surface_deterioration": [
                {
                    "severity": "Medium",
                    "description": "Surface scaling and deterioration",
                    "repair_method": "Surface preparation, patch repair",
                    "immediate_action": "Document extent, plan repairs",
                    "prevention": "Surface sealers, proper drainage"
                }
            ],
            "corrosion": [
                {
                    "severity": "High",
                    "description": "Corrosion of reinforcement leading to cracks",
                    "repair_method": "Remove rust, apply inhibitors",
                    "immediate_action": "Isolate affected area",
                    "prevention": "Anti-corrosion coatings, proper drainage"
                }
            ],
            "efflorescence": [
                {
                    "severity": "Low",
                    "description": "White powder deposits on concrete surfaces",
                    "repair_method": "Surface cleaning, sealant application",
                    "immediate_action": "Identify moisture source",
                    "prevention": "Improve waterproofing, reduce moisture ingress"
                }
            ],
            "delamination": [
                {
                    "severity": "Medium",
                    "description": "Separation of layers in concrete",
                    "repair_method": "Resurface or replace delaminated sections",
                    "immediate_action": "Inspect bonding layers",
                    "prevention": "Proper curing and bonding agents"
                }
            ],
            "honeycombing": [
                {
                    "severity": "Medium",
                    "description": "Voids in concrete caused by improper compaction",
                    "repair_method": "Grout injection, patch repair",
                    "immediate_action": "Assess structural impact",
                    "prevention": "Proper vibration during pouring"
                }
            ],
            "water_leakage": [
                {
                    "severity": "High",
                    "description": "Water ingress through cracks or joints",
                    "repair_method": "Injection grouting, waterproofing membranes",
                    "immediate_action": "Stop water flow, apply sealants",
                    "prevention": "Drainage systems, joint sealing"
                }
            ],
            "settlement_cracks": [
                {
                    "severity": "High",
                    "description": "Cracks due to uneven foundation settlement",
                    "repair_method": "Foundation underpinning, grouting",
                    "immediate_action": "Monitor movement, stabilize foundation",
                    "prevention": "Soil compaction, proper foundation design"
                }
            ],
            "shrinkage_cracks": [
                {
                    "severity": "Low",
                    "description": "Minor cracks caused by shrinkage during curing",
                    "repair_method": "Sealant application, surface repairs",
                    "immediate_action": "Monitor cracks",
                    "prevention": "Proper curing and moisture control"
                }
            ]
        }

        documents = []
        for category, items in kb.items():
            for item in items:
                doc_text = f"Category: {category}\n"
                for key, value in item.items():
                    doc_text += f"{key}: {value}\n"
                documents.append({"text": doc_text, "metadata": {"category": category}})

        return documents

    def create_vector_store(self):
        """Create FAISS vector store"""
        texts = [doc["text"] for doc in self.knowledge_base]
        embeddings = self.embedding_model.encode(texts)
        dimension = embeddings.shape[1]
        index = faiss.IndexFlatL2(dimension)
        index.add(np.array(embeddings).astype('float32'))
        return index

    @lru_cache(maxsize=32)  # Cache recent query results
    def get_relevant_context(self, query: str, k: int = 2) -> str:
        """Retrieve relevant context based on query"""
        try:
            query_embedding = self.embedding_model.encode([query])
            D, I = self.vector_store.search(np.array(query_embedding).astype('float32'), k)
            context = "\n\n".join([self.knowledge_base[i]["text"] for i in I[0]])
            return context
        except Exception as e:
            logger.error(f"Error retrieving context: {e}")
            return ""

class ImageAnalyzer:
    def __init__(self, model_name="microsoft/swin-base-patch4-window7-224-in22k"):
        self.device = "cpu"
        self.defect_classes = ["spalling", "structural_cracks", "surface_deterioration"]
        self.model_name = model_name
        self._model = None
        self._feature_extractor = None

    @property
    def model(self):
        if self._model is None:
            self._model = self._load_model()
        return self._model

    @property
    def feature_extractor(self):
        if self._feature_extractor is None:
            self._feature_extractor = self._load_feature_extractor()
        return self._feature_extractor

    def _load_feature_extractor(self):
        """Load the appropriate feature extractor based on model type"""
        try:
            if "swin" in self.model_name:
                from transformers import AutoFeatureExtractor
                return AutoFeatureExtractor.from_pretrained(self.model_name)
            elif "convnext" in self.model_name:
                from transformers import ConvNextFeatureExtractor
                return ConvNextFeatureExtractor.from_pretrained(self.model_name)
            else:
                from transformers import ViTFeatureExtractor
                return ViTFeatureExtractor.from_pretrained(self.model_name)
        except Exception as e:
            logger.error(f"Feature extractor initialization error: {e}")
            return None

    def _load_model(self):
        try:
            if "swin" in self.model_name:
                from transformers import SwinForImageClassification
                model = SwinForImageClassification.from_pretrained(
                    self.model_name,
                    num_labels=len(self.defect_classes),
                    ignore_mismatched_sizes=True
                )
            elif "convnext" in self.model_name:
                from transformers import ConvNextForImageClassification
                model = ConvNextForImageClassification.from_pretrained(
                    self.model_name,
                    num_labels=len(self.defect_classes),
                    ignore_mismatched_sizes=True
                )
            else:
                from transformers import ViTForImageClassification
                model = ViTForImageClassification.from_pretrained(
                    self.model_name,
                    num_labels=len(self.defect_classes),
                    ignore_mismatched_sizes=True
                )

            model = model.to(self.device)
            
            # Reinitialize the classifier layer
            with torch.no_grad():
                if hasattr(model, 'classifier'):
                    in_features = model.classifier.in_features
                    model.classifier = torch.nn.Linear(in_features, len(self.defect_classes))
                elif hasattr(model, 'head'):
                    in_features = model.head.in_features
                    model.head = torch.nn.Linear(in_features, len(self.defect_classes))
                
            return model
        except Exception as e:
            logger.error(f"Model initialization error: {e}")
            return None

    def preprocess_image(self, image_bytes):
        """Preprocess image for model input"""
        return _cached_preprocess_image(image_bytes, self.model_name)

    def analyze_image(self, image):
        """Analyze image for defects"""
        try:
            if self.model is None:
                raise ValueError("Model not properly initialized")

            inputs = self.feature_extractor(
                images=image,
                return_tensors="pt"
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            with torch.no_grad():
                outputs = self.model(**inputs)
            
            probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
            
            confidence_threshold = 0.3
            results = {
                self.defect_classes[i]: float(probs[i]) 
                for i in range(len(self.defect_classes))
                if float(probs[i]) > confidence_threshold
            }
            
            if not results:
                max_idx = torch.argmax(probs)
                results = {self.defect_classes[int(max_idx)]: float(probs[max_idx])}
            
            return results
            
        except Exception as e:
            logger.error(f"Analysis error: {str(e)}")
            return None

@st.cache_data
def _cached_preprocess_image(image_bytes, model_name):
    """Cached version of image preprocessing"""
    try:
        image = Image.open(image_bytes)
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Adjust size based on model requirements
        if "convnext" in model_name:
            width, height = 384, 384
        else:
            width, height = 224, 224
            
        image = image.resize((width, height), Image.Resampling.LANCZOS)
        return image
    except Exception as e:
        logger.error(f"Image preprocessing error: {e}")
        return None

@st.cache_data      
def get_groq_response(query: str, context: str) -> str:
    """Get response from Groq LLM with caching"""
    try:
        if not os.getenv("GROQ_API_KEY"):
            return "Error: Groq API key not configured"

        client = Groq(api_key=os.getenv("GROQ_API_KEY"))
        
        prompt = f"""Based on the following context about construction defects, answer the question.
        Context: {context}
        Question: {query}
        Provide a detailed answer based on the given context."""

        response = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You are a construction defect analysis expert."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            model="llama-3.3-70b-versatile",
            temperature=0.7,
        )
        return response.choices[0].message.content
    except Exception as e:
        logger.error(f"Groq API error: {e}", exc_info=True)
        return f"Error: Unable to get response from AI model. Exception: {str(e)}"


def create_plotly_confidence_chart(results, chart_id):
    """Create an interactive confidence chart using Plotly with unique ID"""
    colors = {
        'light': {
            'bg': 'white',
            'text': 'black',
            'grid': '#eee'
        },
        'dark': {
            'bg': '#262730',
            'text': 'white',
            'grid': '#333'
        }
    }
    
    theme = 'dark' if st.get_option('theme.base') == 'dark' else 'light'
    
    fig = go.Figure(data=[
        go.Bar(
            x=list(results.values()),
            y=list(results.keys()),
            orientation='h',
            marker_color='rgb(26, 118, 255)',
            text=[f'{v:.1%}' for v in results.values()],
            textposition='auto',
        )
    ])
    
    fig.update_layout(
        title='Defect Detection Confidence Levels',
        xaxis_title='Confidence',
        yaxis_title='Defect Type',
        template='plotly_dark' if theme == 'dark' else 'plotly_white',
        height=400,
        margin=dict(l=20, r=20, t=40, b=20),
        xaxis=dict(range=[0, 1]),
        plot_bgcolor=colors[theme]['bg'],
        paper_bgcolor=colors[theme]['bg'],
        font=dict(color=colors[theme]['text'])
    )
    
    return fig

def create_defect_card(title, description, severity, repair_method):
    """Create a styled card for defect information with theme support"""
    severity_colors = {
        "Critical": "#ff4444",
        "High": "#ffa000",
        "Medium": "#ffeb3b",
        "Low": "#4caf50"
    }
    
    # Get current theme
    is_dark = st.get_option('theme.base') == 'dark'
    
    bg_color = '#1e1e1e' if is_dark else '#ffffff'
    text_color = '#ffffff' if is_dark else '#000000'
    border_color = '#333333' if is_dark else '#dddddd'
    
    return f"""
    <div style="border: 1px solid {border_color}; border-radius: 10px; padding: 15px; margin: 10px 0; background-color: {bg_color}; color: {text_color};">
        <h3 style="color: {'#00a0dc' if is_dark else '#1f77b4'}; margin: 0 0 10px 0;">{title}</h3>
        <p><strong>Description:</strong> {description}</p>
        <p><strong>Severity:</strong> 
            <span style="color: {severity_colors.get(severity, '#808080')}">
                {severity}
            </span>
        </p>
        <p><strong>Repair Method:</strong> {repair_method}</p>
    </div>
    """

def get_theme_specific_styles():
    """Get theme-specific CSS styles"""
    is_dark = st.get_option('theme.base') == 'dark'
    
    if is_dark:
        return """
        <style>
        .stApp {
            background-color: #0e1117;
        }
        .upload-text {
            border: 2px dashed #444;
            background-color: #1e1e1e;
        }
        .info-box {
            background-color: #262730;
            border: 1px solid #333;
        }
        .stAlert {
            background-color: #262730;
            border: 1px solid #333;
        }
        </style>
        """
    else:
        return """
        <style>
        .stApp {
            background-color: #f8f9fa;
        }
        .upload-text {
            border: 2px dashed #ccc;
            background-color: #ffffff;
        }
        .info-box {
            background-color: #e9ecef;
            border: 1px solid #dee2e6;
        }
        </style>
        """

def main():
    st.set_page_config(
        page_title="Smart Construction Defect Analyzer",
        page_icon="🏗️",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Apply theme-specific styles
    st.markdown(get_theme_specific_styles(), unsafe_allow_html=True)
    
    # Base CSS that works for both themes
    st.markdown("""
        <style>
        .css-1d391kg {
            padding: 2rem 1rem;
        }
        .stButton>button {
            width: 100%;
        }
        .upload-text {
            text-align: center;
            padding: 2rem;
            border-radius: 10px;
        }
        .info-box {
            padding: 1rem;
            border-radius: 10px;
            margin: 1rem 0;
        }
        </style>
    """, unsafe_allow_html=True)
    
    # Initialize session state
    if 'analyzer' not in st.session_state:
        st.session_state.analyzer = ImageAnalyzer()
    if 'rag_system' not in st.session_state:
        st.session_state.rag_system = RAGSystem()
    if 'analysis_history' not in st.session_state:
        st.session_state.analysis_history = []
    
    # Sidebar
    with st.sidebar:
        colored_header(
            label="System Controls",
            description="Settings and Information",
            color_name="blue-70"
        )
        
        # Theme selector
        theme = st.selectbox(
            "Choose Theme",
            options=["Light", "Dark"],
            index=1 if st.get_option('theme.base') == 'dark' else 0
        )
        
        if os.getenv("GROQ_API_KEY"):
            st.success("🟢 AI System: Connected")
        else:
            st.error("🔴 AI System: Not configured")
        
        add_vertical_space(2)
        
        with st.expander("ℹ️ About", expanded=True):
            st.write("""
            ### Smart Construction Defect Analyzer
            
            This advanced tool combines computer vision and AI to:
            - Detect construction defects in images
            - Provide detailed repair recommendations
            - Answer technical questions
            - Track analysis history
            """)
        
        with st.expander("🔧 Settings"):
            if st.button("Clear Analysis History"):
                st.session_state.analysis_history = []
                st.cache_data.clear()
                st.success("History cleared!")
            
            confidence_threshold = st.slider(
                "Detection Confidence Threshold",
                min_value=0.0,
                max_value=1.0,
                value=0.3,
                step=0.1
            )
    
    # Main content
    colored_header(
        label="Construction Defect Analyzer",
        description="Upload images and get instant defect analysis",
        color_name="blue-70"
    )
    
    tab1, tab2, tab3 = st.tabs(["📸 Image Analysis", "❓ Ask Expert", "📊 Analysis History"])
    
    with tab1:
        col1, col2 = st.columns([1, 1])
        
        with col1:
            st.markdown('<div class="upload-text">', unsafe_allow_html=True)
            uploaded_file = st.file_uploader(
                "Drop your construction image here",
                type=["jpg", "jpeg", "png"],
                key="image_uploader"
            )
            st.markdown('</div>', unsafe_allow_html=True)
            
            if uploaded_file:
                try:
                    with st.spinner('Processing image...'):
                        processed_image = st.session_state.analyzer.preprocess_image(uploaded_file)
                        if processed_image:
                            st.image(processed_image, caption='Analyzed Image', use_column_width=True)
                            
                            results = st.session_state.analyzer.analyze_image(processed_image)
                            if results:
                                # Store analysis in history
                                st.session_state.analysis_history.append({
                                    'timestamp': datetime.now(),
                                    'results': results,
                                    'image': processed_image
                                })
                except Exception as e:
                    st.error(f"Error: {str(e)}")
        
        with col2:
            if uploaded_file and results:
                st.markdown("### Analysis Results")
                
                # Interactive confidence chart with unique ID
                fig = create_plotly_confidence_chart(results, "main_analysis")
                st.plotly_chart(fig, use_container_width=True, key="main_chart")
                
                # Most critical defect
                most_likely_defect = max(results.items(), key=lambda x: x[1])[0]
                st.info(f"🔍 Primary Defect Detected: {most_likely_defect}")
                
                # Get detailed information about the defect
                context = st.session_state.rag_system.get_relevant_context(most_likely_defect)
                if context:
                    st.markdown("### Defect Details")
                    st.markdown(create_defect_card(
                        most_likely_defect,
                        context.split('\n')[2].split(': ')[1],
                        context.split('\n')[1].split(': ')[1],
                        context.split('\n')[3].split(': ')[1]
                    ), unsafe_allow_html=True)
    
    with tab2:
        st.markdown("### Ask the Construction Expert")
        
        query_placeholder = "Example: What are the best repair methods for structural cracks?"
        user_query = st.text_input("Your Question:", placeholder=query_placeholder)
        
        if user_query:
            with st.spinner('Consulting AI expert...'):
                context = st.session_state.rag_system.get_relevant_context(user_query)
                if context:
                    response = get_groq_response(user_query, context)
                    
                    if not response.startswith("Error"):
                        st.markdown("### Expert Response")
                        st.markdown(response)
                        
                        with st.expander("View Source Information"):
                            st.markdown(context)
                    else:
                        st.error(response)
    
    with tab3:
        if st.session_state.analysis_history:
            for i, analysis in enumerate(reversed(st.session_state.analysis_history)):
                with st.expander(f"Analysis {i+1} - {analysis['timestamp'].strftime('%Y-%m-%d %H:%M')}"):
                    col1, col2 = st.columns([1, 1])
                    with col1:
                        st.image(analysis['image'], caption='Analyzed Image', use_column_width=True)
                    with col2:
                        # Create chart with unique ID for history items
                        fig = create_plotly_confidence_chart(analysis['results'], f"history_{i}")
                        st.plotly_chart(fig, use_container_width=True, key=f"history_chart_{i}")
        else:
            st.info("No analysis history available")

    # Handle theme change
    if theme == "Dark" and st.get_option('theme.base') != 'dark':
        st.markdown("""
            <script>
                var elements = window.parent.document.getElementsByTagName('iframe');
                for (var i = 0; i < elements.length; i++) {
                    if (elements[i].height === '0') {
                        elements[i].remove();
                    }
                }
            </script>
            """, unsafe_allow_html=True)
        st.experimental_set_query_params(theme='dark')
        st.experimental_rerun()
    elif theme == "Light" and st.get_option('theme.base') != 'light':
        st.experimental_set_query_params(theme='light')
        st.experimental_rerun()

if __name__ == "__main__":
    main()