import streamlit as st
import pandas as pd
import numpy as np
import torch
import openai
import os
import io
import re
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt

# --- Groq API setup ---
openai.api_key = os.getenv("GROQ_API_KEY")
openai.api_base = "https://api.groq.com/openai/v1"
GROQ_MODEL = "llama3-8b-8192"

# --- Load Excel ---
@st.cache_data
def load_excel(file):
    xl = pd.read_excel(file, sheet_name=None)
    all_data = pd.concat(xl.values(), ignore_index=True)
    return all_data

# --- Chunk using regex instead of nltk ---
def chunk_data(df, max_tokens=100):
    text = "\n".join(df.astype(str).apply(lambda row: " | ".join(row), axis=1))
    # Split on period, newline, or semicolon
    sentences = re.split(r'(?<=[.;])\s+|\n+', text)
    chunks, current_chunk, current_len = [], [], 0

    for sent in sentences:
        tokens = sent.split()
        if current_len + len(tokens) > max_tokens:
            chunks.append(" ".join(current_chunk))
            current_chunk, current_len = [], 0
        current_chunk.append(sent)
        current_len += len(tokens)

    if current_chunk:
        chunks.append(" ".join(current_chunk))

    return chunks

# --- Embed chunks ---
@st.cache_resource
def embed_chunks(chunks):
    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    embeddings = model.encode(chunks)
    return embeddings, model

# --- Query chunks ---
def query_embedding(user_query, chunks, embeddings, model):
    query_vec = model.encode([user_query])
    similarities = cosine_similarity(query_vec, embeddings)[0]
    top_indices = similarities.argsort()[::-1][:5]
    top_chunks = "\n\n".join([chunks[i] for i in top_indices])
    return top_chunks

# --- Generate Estimate from Groq ---
def generate_estimate(context, user_input):
    prompt = f"""You are a construction estimator in Pakistan. Use the following schedule of rates:

{context}

Estimate a full itemized construction BOQ for:
{user_input}

Include all relevant items for a complete house: excavation, foundation, RCC, masonry, plastering, flooring, doors, windows, paint, distemper, fans, lights, wiring, DBs, plumbing, sanitary fittings, water supply, cupboards, wardrobes, gate, etc.

Provide output in a markdown table with columns: Item No, Description, Qty, Unit, Rate, Amount in Rs.
"""
    response = openai.ChatCompletion.create(
        model=GROQ_MODEL,
        messages=[{"role": "user", "content": prompt}]
    )
    return response['choices'][0]['message']['content']

# --- Quantity Calculator ---
def calculate_quantities(rooms, area, baths, car_porch, living):
    return {
        "Total Area (sqft)": area,
        "No. of Rooms": rooms,
        "No. of Bathrooms": baths,
        "Living Rooms": living,
        "Car Porch Area (est.)": car_porch * 200
    }

# --- Floor Plan Sketch ---
def draw_floor_plan(rooms, baths, living, car_porch, area):
    total_spaces = rooms + baths + living + car_porch
    cols = int(np.ceil(np.sqrt(total_spaces)))
    rows = int(np.ceil(total_spaces / cols))

    fig, ax = plt.subplots(figsize=(10, 8))
    scale = np.sqrt(area) / 10
    width, height = scale, scale * 0.75

    labels = (["Room"] * rooms + ["Bath"] * baths +
              ["Living"] * living + ["Car Porch"] * car_porch)

    for i, label in enumerate(labels):
        row = i // cols
        col = i % cols
        x = col * width
        y = (rows - 1 - row) * height
        ax.add_patch(plt.Rectangle((x, y), width, height, edgecolor='black', facecolor='lightblue'))
        ax.text(x + width / 2, y + height / 2, label, ha='center', va='center', fontsize=8)

    ax.set_xlim(0, cols * width)
    ax.set_ylim(0, rows * height)
    ax.set_aspect('equal')
    ax.set_title(f"Tentative Floor Plan (Scale: 1 unit = {int(scale)} sqft)")
    ax.axis('off')

    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    return buf

# --- Main App ---
def main():
    st.set_page_config(page_title="Construction Estimator", layout="centered")
    st.title("🧱 Construction Estimator (RAG + LLaMA 3 + Sketch)")

    excel_file = st.file_uploader("Upload Schedule of Rates (.xlsx or .xlsm)", type=["xlsx", "xlsm"])
    if excel_file:
        df = load_excel(excel_file)
        st.success("Excel file loaded successfully.")
        chunks = chunk_data(df)
        embeddings, model = embed_chunks(chunks)

        st.subheader("🏗️ Enter Project Details")
        rooms = st.number_input("Number of Rooms", min_value=1, value=3)
        area = st.number_input("Total Covered Area (sqft)", min_value=100, value=1200)
        baths = st.number_input("Number of Washrooms", min_value=1, value=2)
        living = st.number_input("Number of Living Rooms", min_value=0, value=1)
        car_porch = st.number_input("Number of Car Porches", min_value=0, value=1)

        if st.button("Generate Estimate"):
            quantities = calculate_quantities(rooms, area, baths, car_porch, living)
            user_query = f"Estimate cost for {rooms} rooms, {baths} bathrooms, {living} living rooms, total area {area} sqft, and {car_porch} car porch(es)."
            context = query_embedding(user_query, chunks, embeddings, model)
            response = generate_estimate(context, user_query)

            st.subheader("📊 Input Quantities")
            st.json(quantities)

            st.subheader("💸 Estimated Construction Cost (BOQ Style)")
            st.markdown(response)

            buf = draw_floor_plan(rooms, baths, living, car_porch, area)
            st.subheader("🏠 Tentative Floor Plan Sketch")
            st.image(buf, caption="Auto-generated Line Plan", use_column_width=True)
            st.download_button("📥 Download Sketch", buf, file_name="floor_plan.png", mime="image/png")

if __name__ == "__main__":
    main()