import streamlit as st
import cv2
import numpy as np
from PIL import Image
from io import BytesIO
from cvzone.ClassificationModule import Classifier
import pandas as pd

# Set page configuration
st.set_page_config(page_title="Infrastructure Grading & Facility Verification", page_icon="🏗", layout="wide")

# Initialize session state variables if they don't exist
if 'total_score' not in st.session_state:
    st.session_state.total_score = 0
if 'facility_results' not in st.session_state:
    st.session_state.facility_results = []
if 'deficiencies' not in st.session_state:
    st.session_state.deficiencies = []
if 'verified_facilities' not in st.session_state:
    st.session_state.verified_facilities = set()
if 'facility_status' not in st.session_state:   
    st.session_state.facility_status = {}

# Streamlit app title with subheader
st.title("🏗 Infrastructure Grading & Facility Verification System")

# Create tabs
tab1, tab2, tab3, tab4 = st.tabs(["📚 College Info", "🏢 Facility Verification", "📊 Results", "❓ Help"])

# Define models dictionary
models = {
    "Restroom Model": ('Models/Restroom model/keras_model.h5',
                       'Models/Restroom model/labels.txt'),
    "Dispenser Model": ('Models/Dispenser model/keras_model.h5',
                        'Models/Dispenser model/labels.txt'),
    "Safety Equipment Model": ('Models/safety equipment model/keras_model.h5',
                               'Models/safety equipment model/labels.txt'),
    "Computer Lab Model": ('Models/computer lab model/keras_model.h5',
                           'Models/computer lab model/labels.txt'),
    "Server Room Model": ('Models/Server Room model/keras_model.h5',
                           'Models/Server Room model/labels.txt'),
    "Lab Equipment Model": ('Models/lab equipment model/keras_model.h5',
                            'Models/lab equipment model/labels.txt'),
    "Sports Equipment Model": ('Models/sports equipment model/keras_model.h5',
                               'Models/sports equipment model/labels.txt'),
    "Bicycle Stand Model": ('Models/bicycle stand model/keras_model.h5',
                            'Models/bicycle stand model/labels.txt'),
    "Medical Room Model": ('Models/Medical Room Model/keras_model.h5',
                               'Models/Medical Room Model/labels.txt'),
    "Workshop/Mechanical Lab Model": ('Models/workshop model/keras_model.h5',
                                      'Models/workshop model/labels.txt'),
    "Bus/Transport Model": ('Models/bus transport model/keras_model.h5',
                            'Models/bus transport model/labels.txt'),
    "COVID-19 Protocol Model": ('Models/COVID-19 protocol model/keras_model.h5',
                                'Models/COVID-19 protocol model/labels.txt'),
    "Canteen Model": ('Models/canteen model/keras_model.h5',
                      'Models/canteen model/labels.txt'),
    "CCTV Model": ('Models/CCTV model/keras_model.h5',
                   'Models/CCTV model/labels.txt'),
    "Classroom Model": ('Models/classroom model/keras_model.h5',
                        'Models/classroom model/labels.txt'),
    "Elearning Model": ('Models/elearning model/keras_model.h5',
                        'Models/elearning model/labels.txt'),
    "Faculty Cabin Model": ('Models/faculty cabin model/keras_model.h5',
                            'Models/faculty cabin model/labels.txt'),
    "Fire Extinguisher Model": ('Models/fire extinguisher model/keras_model.h5',
                                'Models/fire extinguisher model/labels.txt'),
    "Generator Model": ('Models/generator model/keras_model.h5',
                        'Models/generator model/labels.txt'),
    "Ground Model": ('Models/ground model (1)/keras_model.h5',
                     'Models/ground model (1)/labels.txt'),
    "Laptop Model": ('Models/laptop model/keras_model.h5',
                     'Models/laptop model/labels.txt'),
    "Library Model": ('Models/library model/keras_model.h5',
                      'Models/library model/labels.txt'),
    "Parking Model": ('Models/parking model/keras_model.h5',
                      'Models/parking model/labels.txt'),
    "Pothole Model": ('Models/pothole model/keras_model.h5',
                      'Models/pothole model/labels.txt'),
    "Seminar Hall Model": ('Models/seminar hall model/keras_model.h5',
                           'Models/seminar hall model/labels.txt'),
    "TPO Model": ('Models/tpo model/keras_model.h5',
                  'Models/tpo model/labels.txt'),
    "Audi Model": ('Models/Audi model/keras_model.h5',
                   'Models/Audi model/labels.txt'),
    "Conference Halls Model": ('Models/conference halls model/keras_model.h5',
                               'Models/conference halls model/labels.txt'),
    "Drawing Halls Model": ('Models/Drawing halls model/keras_model.h5',
                            'Models/Drawing halls model/labels.txt'),
}

# Function to handle input and image upload
def input_and_upload(label, min_value, default_value, model_name, model_path, labels_path):
    col1, col2 = st.columns(2)
    with col1:
        if isinstance(default_value, int):
            user_value = st.number_input(f"Number of {label}", min_value=min_value, value=int(default_value), step=1)
        else:
            user_value = st.number_input(f"Number of {label}", min_value=float(min_value), value=float(default_value), step=0.1)
    with col2:
        uploaded_files = st.file_uploader(f"Upload images for {label} (max {int(user_value)})", 
                                        type=["jpg", "jpeg", "png"], 
                                        accept_multiple_files=True)

    # Limit the number of processed files to user_value
    uploaded_files = uploaded_files[:int(user_value)] if uploaded_files else []

    score = 0
    if uploaded_files:
        score = classify_image(model_path, labels_path, uploaded_files, model_name)
        
        # Update the facility status immediately after classification
        if label not in st.session_state.verified_facilities and score > 0:
            st.session_state.verified_facilities.add(label)
            st.session_state.total_score = sum(
                data["Score"] 
                for data in st.session_state.facility_status.values()
            )

    return user_value, score


# Function to classify image
def classify_image(model_path, labels_path, uploaded_files, model_name):
    try:
        classifier = Classifier(model_path, labels_path)
        class_names = open(labels_path).read().splitlines()
    except FileNotFoundError:
        st.error(f"Model files for {model_name} not found. Please check the model path.")
        return 0
    except Exception as e:
        st.error(f"Error loading model {model_name}: {str(e)}")
        return 0
        
    score = 0

    # Create a grid layout for images
    cols = st.columns(3)  # Adjust the number of columns as needed

    for idx, uploaded_file in enumerate(uploaded_files):
        try:
            img = Image.open(BytesIO(uploaded_file.read()))
            img_array = np.array(img)

            # Convert to BGR if the image is RGB
            if img_array.shape[-1] == 3:
                img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

            # Resize for prediction (but don't display this resized version)
            img_resized = cv2.resize(img_array, (224, 224))

            prediction = classifier.getPrediction(img_resized)
            class_id = prediction[1]
            class_name = class_names[class_id] if class_id < len(class_names) else "Unknown"

            # Display image and results in a compact format
            with cols[idx % 3]:
                st.image(img, use_column_width=True)
                if "no" in class_name.lower() or "not" in class_name.lower():
                    st.warning(f"{model_name} not verified.")
                else:
                    st.success(f"{model_name} is verified.")
                    score += 10

        except Exception as e:
            st.error(f"An error occurred: {e}")

    return score

# Grade calculation function
def calculate_grade(total_score, max_score):
    percentage = (total_score / max_score) * 100
    if percentage >= 80:
        return "A"
    elif percentage >= 60:
        return "B"
    elif percentage >= 40:
        return "C"
    else:
        return "D"

with tab1:
    st.markdown("### College Information")
    st.markdown("Enter the basic information about the college below:")

    # Create columns for a more compact layout
    col1, col2 = st.columns(2)

    with col1:
        num_divisions = st.number_input("Number of Divisions", min_value=1, value=1, help="Total number of divisions across all years")
        num_courses = st.number_input("Number of Courses", min_value=1, value=1, help="Total number of distinct courses offered")

    with col2:
        num_students = st.number_input("Total Students", min_value=1, value=100, help="Total number of students enrolled")
        course_duration = st.number_input("Course Duration (years)", min_value=1, value=4, help="Average duration of courses in years")

    # Calculations
    classroom_requirement = num_divisions * course_duration * 0.5

    # Lab logic based on student intake
    if num_students <= 600:
        first_year_labs = 4
    else:
        first_year_labs = 4 + (num_students - 600) // 150

    if num_students <= 180 * num_courses:
        labs_other_years = 2 * num_courses * (course_duration - 1)
    else:
        extra_students_per_course = (num_students - 180 * num_courses) // 50
        labs_other_years = 2 * num_courses * (course_duration - 1) + extra_students_per_course

    total_labs = first_year_labs + labs_other_years

    # Other facility requirements
    workshop_requirement = 1 + (num_students - 600) // 600 if num_students > 600 else 1
    cad_centre_requirement = 1 + (num_students - 600) // 600 if num_students > 600 else 1
    computer_centre_requirement = 1 + (num_students - 600) // 600 if num_students > 600 else 1
    seminar_hall_requirement = 1
    library_requirement = 1
    language_lab_requirement = 1
    pc_requirement = max(20, num_students // 10)

    st.markdown("### 📊 Calculated Facility Requirements")
    
    # Create three columns for a more compact display of requirements
    col1, col2, col3 = st.columns(3)

    with col1:
        st.metric("Classrooms", f"{classroom_requirement:.1f}")
        st.metric("Total Labs", f"{total_labs}")
        st.metric("Workshops", f"{workshop_requirement}")

    with col2:
        st.metric("CAD Centres", f"{cad_centre_requirement}")
        st.metric("Computer Centres", f"{computer_centre_requirement}")
        st.metric("Seminar Halls", f"{seminar_hall_requirement}")

    with col3:
        st.metric("Libraries", f"{library_requirement}")
        st.metric("Language Labs", f"{language_lab_requirement}")
        st.metric("PC/Laptops", f"{pc_requirement}")

    st.info("ℹ These calculations are based on standard educational infrastructure guidelines.")

# Input fields and image upload for each facility
facilities = [
    ("Classrooms", classroom_requirement, "Classroom Model"),
    ("Computer Labs", computer_centre_requirement, "Computer Lab Model"), #Done
    ("Workshops", workshop_requirement, "Workshop/Mechanical Lab Model"), # Done
    ("Drawing Halls", cad_centre_requirement, "Drawing Halls Model"), # Done
    ("Seminar Halls", seminar_hall_requirement, "Seminar Hall Model"), # Done
    ("Conference Halls", 1, "Conference Halls Model"), # done
    ("Auditorium", 1, "Audi Model"),
    ("Faculty Cabins", 1, "Faculty Cabin Model"),
    ("Security/CCTV", 1, "CCTV Model"), 
    ("House Keeping", 1, "Safety Equipment Model"), # to be done / store rooms can be done but not added yet
    ("Restrooms", 1, "Restroom Model"), # Done
    ("Canteens", 1, "Canteen Model"), 
    ("Server Room", 1, "Server Room Model"), # Done
    ("First Aid/Medical Room", 1, "Medical Room Model"), # done
    ("Gym/Sports Facilities", 1, "Sports Equipment Model"), #  done
    ("Language Labs", language_lab_requirement, "Computer Lab Model"), #done
    ("Water Coolers/Dispensers", 1, "Dispenser Model"), # done
    ("Generators", 1, "Generator Model"),
    ("TPO Office", 1, "TPO Model"),
    # Other supporting facilities
    ("Libraries", library_requirement, "Library Model"),
    ("PCs/Laptops", pc_requirement, "Laptop Model"),
    ("Bicycle Stands", 1, "Bicycle Stand Model"),
    ("Bus/Transport Facilities", 1, "Bus/Transport Model"),
    ("COVID-19 Protocol Measures", 1, "COVID-19 Protocol Model"),
    ("E-learning Facilities", 1, "Elearning Model"),
    ("Fire Extinguishers", 1, "Fire Extinguisher Model"),
    ("Grounds/Playgrounds", 1, "Ground Model"),
    ("Parking Areas", 1, "Parking Model"),
    ("Road Condition (Potholes)", 1, "Pothole Model"),
]

with tab2:
    st.markdown("### Facility Verification")
    st.markdown("Upload images for each facility to verify their existence and compliance.")

    # Reset scores and results when starting verification
    if st.button("🔄 Reset Verification"):
        # Reset all session state variables
        st.session_state.total_score = 0
        st.session_state.facility_results = []
        st.session_state.deficiencies = []
        st.session_state.verified_facilities = set()
        st.session_state.facility_status = {}
        # Add a rerun to refresh the page
        st.rerun()

    # Create subtabs for different facility categories
    facility_tabs = st.tabs(["🏫 Academic", "🔬 Labs & Workshops", "📚 Support Facilities", "🏥 Essential Services"])
    
    # Define facility categories
    academic_facilities = [
        f for f in facilities 
        if any(name in f[0] for name in ["Classroom", "Faculty", "Drawing", "Seminar", "Conference", "Auditorium"])
    ]
    
    lab_facilities = [
        f for f in facilities 
        if any(name in f[0] for name in ["Computer Lab", "Workshop", "Language Lab", "Server"])
    ]
    
    support_facilities = [
        f for f in facilities 
        if any(name in f[0] for name in ["Library", "TPO", "E-learning", "Sport", "Ground"])
    ]
    
    essential_facilities = [
        f for f in facilities 
        if any(name in f[0] for name in ["Restroom", "Security", "CCTV", "Medical", "Water", "Fire", "Generator"])
    ]

    # Process facilities in each tab
    with facility_tabs[0]:  # Academic Facilities
        for facility, requirement, model_name in academic_facilities:
            st.markdown(f"#### {facility}")
            user_value, score = input_and_upload(facility, 0, requirement, model_name, *models[model_name])
            
            # Update facility status
            st.session_state.facility_status[facility] = {
                "Verified": int(user_value),
                "Required": int(requirement),
                "Score": score,
                "Status": "✅ Met" if user_value >= requirement and score > 0 else "❌ Not Met"
            }
            
            # Only add score if not already verified
            if facility not in st.session_state.verified_facilities and score > 0:
                st.session_state.total_score += score
                st.session_state.verified_facilities.add(facility)
            
            if user_value < requirement:
                st.session_state.deficiencies.append(f"{facility} requires {int(requirement - user_value)} more.")
            st.markdown("---")

    with facility_tabs[1]:  # Labs & Workshops
        for facility, requirement, model_name in lab_facilities:
            st.markdown(f"#### {facility}")
            user_value, score = input_and_upload(facility, 0, requirement, model_name, *models[model_name])
            
            # Update facility status
            st.session_state.facility_status[facility] = {
                "Verified": int(user_value),
                "Required": int(requirement),
                "Score": score,
                "Status": "✅ Met" if user_value >= requirement and score > 0 else "❌ Not Met"
            }
            
            # Only add score if not already verified
            if facility not in st.session_state.verified_facilities and score > 0:
                st.session_state.total_score += score
                st.session_state.verified_facilities.add(facility)
            
            if user_value < requirement:
                st.session_state.deficiencies.append(f"{facility} requires {int(requirement - user_value)} more.")
            st.markdown("---")

    with facility_tabs[2]:  # Support Facilities
        for facility, requirement, model_name in support_facilities:
            st.markdown(f"#### {facility}")
            user_value, score = input_and_upload(facility, 0, requirement, model_name, *models[model_name])
            
            # Update facility status
            st.session_state.facility_status[facility] = {
                "Verified": int(user_value),
                "Required": int(requirement),
                "Score": score,
                "Status": "✅ Met" if user_value >= requirement and score > 0 else "❌ Not Met"
            }
            
            # Only add score if not already verified
            if facility not in st.session_state.verified_facilities and score > 0:
                st.session_state.total_score += score
                st.session_state.verified_facilities.add(facility)
            
            if user_value < requirement:
                st.session_state.deficiencies.append(f"{facility} requires {int(requirement - user_value)} more.")
            st.markdown("---")

    with facility_tabs[3]:  # Essential Services
        for facility, requirement, model_name in essential_facilities:
            st.markdown(f"#### {facility}")
            user_value, score = input_and_upload(facility, 0, requirement, model_name, *models[model_name])
            
            # Update facility status
            st.session_state.facility_status[facility] = {
                "Verified": int(user_value),
                "Required": int(requirement),
                "Score": score,
                "Status": "✅ Met" if user_value >= requirement and score > 0 else "❌ Not Met"
            }
            
            # Only add score if not already verified
            if facility not in st.session_state.verified_facilities and score > 0:
                st.session_state.total_score += score
                st.session_state.verified_facilities.add(facility)
            
            if user_value < requirement:
                st.session_state.deficiencies.append(f"{facility} requires {int(requirement - user_value)} more.")
            st.markdown("---")

with tab3:
    st.markdown("### 📊 Results Summary")
    
    # Calculate max possible score and verification progress
    max_score = len(facilities) * 10
    verified_count = len(st.session_state.verified_facilities)
    verification_progress = (verified_count / len(facilities)) * 100
    
    # Calculate grade based on total score
    current_grade = calculate_grade(st.session_state.total_score, max_score)
    
    # Create three columns for key metrics
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.markdown("#### Grade")
        st.markdown(f"<h2 style='text-align: center; color: {'green' if current_grade == 'A' else 'orange' if current_grade == 'B' else 'red'};'>{current_grade}</h2>", unsafe_allow_html=True)
    
    with col2:
        st.markdown("#### Score")
        st.markdown(f"<h2 style='text-align: center;'>{st.session_state.total_score}/{max_score}</h2>", unsafe_allow_html=True)
        st.progress(st.session_state.total_score / max_score)
    
    with col3:
        st.markdown("#### Verification Progress")
        st.markdown(f"<h2 style='text-align: center;'>{verification_progress:.1f}%</h2>", unsafe_allow_html=True)
        st.progress(verification_progress / 100)

    # Display facility status metrics
    if st.session_state.facility_status:
        st.markdown("### 📋 Facility Status Overview")
        
        # Calculate metrics
        total_facilities = len(st.session_state.facility_status)
        compliant_facilities = sum(1 for data in st.session_state.facility_status.values() if data["Status"] == "✅ Met")
        non_compliant_facilities = total_facilities - compliant_facilities
        
        # Display metrics in columns
        metric_col1, metric_col2, metric_col3 = st.columns(3)
        with metric_col1:
            st.metric("Total Facilities", total_facilities)
        with metric_col2:
            st.metric("Compliant", compliant_facilities, delta=f"{(compliant_facilities/total_facilities)*100:.1f}%")
        with metric_col3:
            st.metric("Non-Compliant", non_compliant_facilities, delta=f"-{(non_compliant_facilities/total_facilities)*100:.1f}%")

        # Create tabs for different views of the results
        results_tab1, results_tab2 = st.tabs(["📊 Summary Table", "⚠️ Deficiencies"])
        
        with results_tab1:
            # Create DataFrame from facility_status
            summary_data = [
                {
                    "Facility": facility,
                    "Verified": data["Verified"],
                    "Required": data["Required"],
                    "Score": data["Score"],
                    "Status": data["Status"]
                }
                for facility, data in st.session_state.facility_status.items()
            ]
            
            summary_df = pd.DataFrame(summary_data)
            summary_df = summary_df[["Facility", "Verified", "Required", "Score", "Status"]]
            
            # Add color coding to the dataframe
            st.dataframe(
                summary_df.style
                .set_properties(**{
                    'background-color': 'white',
                    'color': 'black',
                    'border-color': 'lightgrey',
                    'text-align': 'center'
                })
                .apply(lambda x: ['background-color: #e6ffe6' if v == '✅ Met' else 'background-color: #ffe6e6' for v in x], subset=['Status'])
                .format({'Score': '{:.0f}'})
            )

        with results_tab2:
            # Display deficiencies with better formatting
            current_deficiencies = [
                (facility, data['Required'] - data['Verified'], data['Score'])
                for facility, data in st.session_state.facility_status.items()
                if data['Verified'] < data['Required'] or data['Score'] == 0
            ]
            
            if current_deficiencies:
                st.markdown("#### 🚨 Areas Needing Improvement")
                for facility, shortage, score in current_deficiencies:
                    if score == 0:
                        st.error(f"**{facility}**: Needs verification with proper images")
                    else:
                        st.error(f"**{facility}**: Requires {shortage} more unit{'s' if shortage > 1 else ''}")
                
                # Add recommendations
                st.markdown("#### 💡 Recommendations")
                st.info("""
                To improve your grade:
                1. Focus on addressing critical deficiencies first
                2. Prioritize essential facilities
                3. Document improvements with clear photographs
                4. Ensure all verifications are complete
                """)
            else:
                st.success("🎉 Congratulations! All required facilities are met and verified.")

    else:
        # Show empty state
        st.info("👋 No facilities have been verified yet. Please complete the verification process in the Facility Verification tab.")
        st.markdown("""
        #### Getting Started:
        1. Go to the Facility Verification tab
        2. Upload images for each facility
        3. Complete the verification process
        4. Return here to view your results
        """)

with tab4:
    st.markdown("### How to Use This System")
    st.write("""
    1. **College Information Tab**
        - Enter basic college details
        - View calculated facility requirements
        
    2. **Facility Verification Tab**
        - Navigate through facility categories
        - Upload images for verification
        - Get instant verification results
        
    3. **Results Tab**
        - View overall grade and score
        - Check facility verification summary
        - Review any deficiencies
    """)

    st.markdown("### Grading Criteria")
    st.write("""
    **Grade A:** ≥ 80% (Excellent infrastructure and verification)
    **Grade B:** ≥ 60% (Good infrastructure with minor improvements needed)
    **Grade C:** ≥ 40% (Basic infrastructure present but needs significant improvements)
    **Grade D:** < 40% (Major infrastructure improvements required)
    """)

    st.markdown("### Note")
    st.info("This system uses machine learning models to verify facilities from uploaded images. Ensure that the images clearly show the relevant facility for accurate verification.")

    st.markdown("### Disclaimer")
    st.warning("This tool is for assessment purposes only. The final evaluation and accreditation of educational institutions should be conducted by authorized bodies following official guidelines and on-site inspections.")

# Footer
st.markdown("---")
st.markdown("Developed for educational infrastructure assessment purposes. For support or inquiries, please contact the system administrator.")