import time
from streamlit_extras.colored_header import colored_header
from streamlit_extras.add_vertical_space import add_vertical_space
from streamlit_card import card
import plotly.graph_objects as go
import streamlit as st
import torch
from PIL import Image
import numpy as np
from transformers import ViTFeatureExtractor, ViTForImageClassification
from sentence_transformers import SentenceTransformer
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

# 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):
    """Create an interactive confidence chart using Plotly"""
    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_white',
        height=400,
        margin=dict(l=20, r=20, t=40, b=20),
        xaxis=dict(range=[0, 1])
    )
    return fig

def create_defect_card(title, description, severity, repair_method):
    """Create a styled card for defect information"""
    severity_colors = {
        "Critical": "red",
        "High": "orange",
        "Medium": "yellow",
        "Low": "green"
    }
    
    return f"""
    <div style="border: 1px solid #ddd; border-radius: 10px; padding: 15px; margin: 10px 0;">
        <h3 style="color: #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, 'gray')}">
                {severity}
            </span>
        </p>
        <p><strong>Repair Method:</strong> {repair_method}</p>
    </div>
    """

def main():
    st.set_page_config(
        page_title="Smart Construction Defect Analyzer",
        page_icon="🏗️",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Custom CSS
    st.markdown("""
        <style>
        .stApp {
            background-color: #f8f9fa;
        }
        .css-1d391kg {
            padding: 2rem 1rem;
        }
        .stButton>button {
            width: 100%;
        }
        .upload-text {
            text-align: center;
            padding: 2rem;
            border: 2px dashed #ccc;
            border-radius: 10px;
            background-color: #ffffff;
        }
        .info-box {
            background-color: #e9ecef;
            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"
        )
        
        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
                fig = create_plotly_confidence_chart(results)
                st.plotly_chart(fig, use_container_width=True)
                
                # 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:
                        fig = create_plotly_confidence_chart(analysis['results'])
                        st.plotly_chart(fig, use_container_width=True)
        else:
            st.info("No analysis history available")

if __name__ == "__main__":
    main()