app.py

import streamlit as st
from groq import Groq
import json
import os
import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objects as go
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize Groq client
try:
    client = Groq(api_key=os.getenv("GROQ_API_KEY"))
except Exception as e:
    st.error(f"Failed to initialize Groq client: {e}")

# App title
st.title("🏠 AI Architect: Home Planner")
st.markdown("Get an easy-to-understand house plan with 2D/3D visualization!")

# --- USER INPUTS ---
st.header("📏 Step 1: Plot Details")
col1, col2 = st.columns(2)
with col1:
    plot_length = st.number_input("Plot Length (ft)", min_value=20, max_value=500, value=50)
with col2:
    plot_width = st.number_input("Plot Width (ft)", min_value=20, max_value=500, value=40)

st.header("🚪 Step 2: Room Requirements")
master_bed = st.number_input("Master Bedrooms (12x14 ft min)", min_value=0, max_value=5, value=1)
bedrooms = st.number_input("Standard Bedrooms (10x12 ft min)", min_value=0, max_value=10, value=2)
kitchens = st.number_input("Kitchens (8x10 ft min)", min_value=1, max_value=3, value=1)
study_rooms = st.number_input("Study Rooms (8x8 ft min)", min_value=0, max_value=3, value=0)
lounges = st.number_input("Lounges (12x16 ft min)", min_value=1, max_value=3, value=1)

st.header("🚽 Step 3: Washrooms")
attached_washrooms = st.number_input("Attached Washrooms (6x8 ft min)", min_value=0, max_value=5, value=1)
common_washrooms = st.number_input("Common Washrooms (5x7 ft min)", min_value=0, max_value=3, value=1)

st.header("🌳 Step 4: Extras")
parking = st.checkbox("Parking Space (12x20 ft min) 🚗")
lawn = st.checkbox("Lawn/Garden 🌳")

# --- AI GENERATION ---
if st.button("Generate Professional Plan"):
    prompt = f"""
    Design a house plan for a {plot_length}x{plot_width} ft plot following architectural best practices.
    Return response in STRICT JSON format with these EXACT keys:
    {{
        "floor_plan": {{
            "rooms": [
                {{
                    "name": "string",
                    "x": number,
                    "y": number,
                    "width": number,
                    "height": number
                }}
            ],
            "washrooms": [
                {{
                    "name": "string",
                    "x": number,
                    "y": number,
                    "width": number,
                    "height": number
                }}
            ]
        }},
        "visualization_tips": ["string"],
        "material_suggestions": ["string"],
        "warnings": ["string"]
    }}

    Requirements:
    1. All dimensions must be numbers (not arrays)
    2. Positions must be x,y coordinates as separate numbers
    3. Room names should be clear (e.g., "Master Bedroom")
    4. Include at least {master_bed} master bedroom(s) and {bedrooms} standard bedroom(s)
    """

    try:
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama3-70b-8192",
            response_format={"type": "json_object"},
            temperature=0.3  # More deterministic output
        )
        
        # Parse and validate response
        plan = json.loads(response.choices[0].message.content)
        
        # Validate required fields
        required_keys = ["floor_plan", "visualization_tips", "material_suggestions", "warnings"]
        for key in required_keys:
            if key not in plan:
                raise ValueError(f"Missing required key: {key}")
        
        if "rooms" not in plan["floor_plan"]:
            raise ValueError("Missing rooms in floor plan")
        
        st.success("✅ Professionally Designed Plan Generated!")
        
        # --- 2D Visualization ---
        st.header("📐 2D Floor Plan")
        
        fig, ax = plt.subplots(figsize=(10, 8))
        colors = plt.cm.tab20.colors
        
        # Plot rooms
        for i, room in enumerate(plan['floor_plan']['rooms']):
            rect = plt.Rectangle(
                (room['x'], room['y']), 
                room['width'], 
                room['height'],
                linewidth=2,
                edgecolor='black',
                facecolor=colors[i % len(colors)])
            ax.add_patch(rect)
            
            # Add room label
            center_x = room['x'] + room['width']/2
            center_y = room['y'] + room['height']/2
            ax.text(center_x, center_y, 
                   f"{room.get('name', 'Room')}\n{room['width']}x{room['height']}ft",
                   ha='center', va='center', fontsize=8)
        
        # Plot washrooms
        for washroom in plan['floor_plan'].get('washrooms', []):
            rect = plt.Rectangle(
                (washroom['x'], washroom['y']), 
                washroom['width'], 
                washroom['height'],
                linewidth=2,
                edgecolor='black',
                facecolor='lightblue',
                hatch='/')
            ax.add_patch(rect)
            ax.text(washroom['x'] + washroom['width']/2,
                   washroom['y'] + washroom['height']/2,
                   "Bath",
                   ha='center', va='center', fontsize=8)
        
        # Set plot limits and labels
        ax.set_xlim(0, plot_length)
        ax.set_ylim(0, plot_width)
        ax.set_aspect('equal')
        ax.set_xlabel('Length (ft)')
        ax.set_ylabel('Width (ft)')
        ax.set_title('2D Floor Plan')
        st.pyplot(fig)
        
        # --- 3D Visualization ---
        st.header("✨ 3D Visualization")
        
        fig_3d = go.Figure()
        
        # Add rooms as 3D boxes
        for i, room in enumerate(plan['floor_plan']['rooms']):
            fig_3d.add_trace(go.Mesh3d(
                x=[room['x'], room['x'], room['x']+room['width'], room['x']+room['width'], room['x']],
                y=[room['y'], room['y']+room['height'], room['y']+room['height'], room['y'], room['y']],
                z=[0, 0, 0, 0, 0],
                opacity=0.7,
                color=f'rgb({colors[i][0]*255},{colors[i][1]*255},{colors[i][2]*255})',
                name=room.get('name', 'Room'),
                showlegend=True
            ))
            fig_3d.add_trace(go.Mesh3d(
                x=[room['x'], room['x'], room['x']+room['width'], room['x']+room['width'], room['x']],
                y=[room['y'], room['y']+room['height'], room['y']+room['height'], room['y'], room['y']],
                z=[10, 10, 10, 10, 10],
                opacity=0.7,
                color=f'rgb({colors[i][0]*255},{colors[i][1]*255},{colors[i][2]*255})',
                showlegend=False
            ))
        
        fig_3d.update_layout(
            scene=dict(
                xaxis=dict(title='Length (ft)'),
                yaxis=dict(title='Width (ft)'),
                zaxis=dict(title='Height (ft)'),
                aspectmode='manual',
                aspectratio=dict(x=1, y=plot_width/plot_length, z=0.2)
            ),
            margin=dict(r=20, l=10, b=10, t=30),
            legend=dict(x=1, y=1)
        )
        st.plotly_chart(fig_3d)
        
        # --- Additional Information ---
        st.header("💡 Visualization Tips")
        for tip in plan['visualization_tips']:
            st.markdown(f"- {tip}")
        
        st.header("🏗️ Material Suggestions")
        for material in plan['material_suggestions']:
            st.markdown(f"- {material}")
        
        st.header("⚠️ Warnings")
        for warning in plan['warnings']:
            st.markdown(f"- ⚠️ {warning}")
        
    except json.JSONDecodeError:
        st.error("Failed to parse the AI response. Please try again.")
    except Exception as e:
        st.error(f"Error: {str(e)}. Please adjust your inputs and try again.")

# Footer
st.markdown("---")
st.caption("✅ Easy-to-understand visualizations | Powered by Groq API")