import os
import streamlit as st
from groq import Groq
import matplotlib.pyplot as plt
from fpdf import FPDF
import io

# Set your API key
os.environ["GROQ_API_KEY"] = "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"

# Initialize Groq client
client = Groq(api_key=os.environ["GROQ_API_KEY"])

# App title
st.title("Cutting Plan Generator")

# Input for blank plates
st.header("Blank Plates Details")
num_blank_types = st.number_input("Number of Blank Plate Types", min_value=1, step=1, key="num_blank_types")
blank_data = []
for i in range(num_blank_types):
    with st.expander(f"Blank Plate {i+1}"):
        blank_width = st.number_input(f"Width (mm) for Blank {i+1}", min_value=0.0, step=0.1, key=f"blank_width_{i}")
        blank_length = st.number_input(f"Length (mm) for Blank {i+1}", min_value=0.0, step=0.1, key=f"blank_length_{i}")
        num_blank_plates = st.number_input(f"Number of Plates for Blank {i+1}", min_value=1, step=1, key=f"num_blank_plates_{i}")
        blank_data.append({"width": blank_width, "length": blank_length, "count": num_blank_plates})

# Input for pieces
st.header("Pieces to be Cut Details")
num_piece_types = st.number_input("Number of Piece Types", min_value=1, step=1, key="num_piece_types")
piece_data = []
for i in range(num_piece_types):
    with st.expander(f"Piece {i+1}"):
        piece_width = st.number_input(f"Width (mm) for Piece {i+1}", min_value=0.0, step=0.1, key=f"piece_width_{i}")
        piece_length = st.number_input(f"Length (mm) for Piece {i+1}", min_value=0.0, step=0.1, key=f"piece_length_{i}")
        num_pieces = st.number_input(f"Number of Pieces for Piece {i+1}", min_value=1, step=1, key=f"num_pieces_{i}")
        piece_data.append({"width": piece_width, "length": piece_length, "count": num_pieces})

# Saw cut size and trimming size
st.header("Other Specifications")
saw_cut_size = st.number_input("Saw Cut Size (mm)", min_value=0.0, step=0.1, key="saw_cut_size")
trimming_size = st.number_input("Trimming Size (mm)", min_value=0.0, step=0.1, key="trimming_size")

# Priority of blank sizes
priority = st.selectbox("Priority of Blank Sizes", options=["Smallest First", "Largest First"], key="priority")

# Generate Cutting Plan and Sketch
if st.button("Generate Cutting Plan"):
    total_cut_area = 0
    blank_areas = []
    for blank in blank_data:
        blank_area = (blank["width"] - 2 * trimming_size) * (blank["length"] - 2 * trimming_size) * blank["count"]
        blank_areas.append(blank_area)

    for piece in piece_data:
        total_cut_area += piece["width"] * piece["length"] * piece["count"]

    total_blank_area = sum(blank_areas)
    unused_area = total_blank_area - total_cut_area
    usable_area_percentage = (total_cut_area / total_blank_area) * 100

    st.subheader("Results")
    st.write(f"Total Usable Surface Area: {total_blank_area:.2f} mm²")
    st.write(f"Total Cut Square Area: {total_cut_area:.2f} mm²")
    st.write(f"Total Non-used Square Area: {unused_area:.2f} mm²")
    st.write(f"Percentage of Usable Area: {usable_area_percentage:.2f}%")

    # Generate sketch
    fig, ax = plt.subplots(figsize=(8, 6))
    for i, blank in enumerate(blank_data):
        ax.add_patch(plt.Rectangle((0, i * (blank["length"] + 10)), blank["width"], blank["length"], edgecolor="blue", facecolor="none", label=f"Blank {i+1}" if i == 0 else ""))
        for j, piece in enumerate(piece_data):
            ax.add_patch(plt.Rectangle((j * (piece["width"] + 5), i * (blank["length"] + 10) + trimming_size), piece["width"], piece["length"], edgecolor="red", facecolor="pink", label=f"Piece {j+1}" if i == 0 else ""))

    ax.set_xlim(0, max(blank["width"] for blank in blank_data) + 50)
    ax.set_ylim(0, (len(blank_data) * max(blank["length"] for blank in blank_data)) + 100)
    ax.set_title("Cutting Plan Sketch")
    ax.legend(loc="upper right")
    st.pyplot(fig)

    # Save as PDF
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    pdf.cell(200, 10, txt="Cutting Plan Report", ln=True, align="C")
    pdf.cell(200, 10, txt=f"Total Usable Surface Area: {total_blank_area:.2f} mm²", ln=True)
    pdf.cell(200, 10, txt=f"Total Cut Square Area: {total_cut_area:.2f} mm²", ln=True)
    pdf.cell(200, 10, txt=f"Total Non-used Square Area: {unused_area:.2f} mm²", ln=True)
    pdf.cell(200, 10, txt=f"Percentage of Usable Area: {usable_area_percentage:.2f}%", ln=True)

    pdf_output = io.BytesIO()
    pdf.output(pdf_output)
    st.download_button(
        label="Download Cutting Plan as PDF",
        data=pdf_output.getvalue(),
        file_name="cutting_plan.pdf",
        mime="application/pdf",
    )