import streamlit as st
import pandas as pd
from sentence_transformers import SentenceTransformer
import numpy as np
import faiss

# Load data
df = pd.read_csv("frameworks.csv")

# Embed the scope column
model = SentenceTransformer("all-MiniLM-L6-v2")
scope_embeddings = model.encode(df["scope"].tolist())
index = faiss.IndexFlatL2(scope_embeddings.shape[1])
index.add(np.array(scope_embeddings))

st.set_page_config(page_title="AI Framework/DPS Finder", layout="centered")
st.title("🤖 AI-Powered Framework & DPS Finder")

st.markdown("""Fill in your procurement details and our AI will recommend suitable frameworks or DPS options.""")

with st.form("proc_form"):
    authority_type = st.selectbox("1. What type of contracting authority are you?", ["Local authority", "NHS", "Central government", "Utility"])
    award_method = st.selectbox("2. Would you prefer to direct award or run a mini-competition?", ["Direct award", "Mini-competition", "Either"])
    contract_value = st.number_input("3. Expected contract value (incl. VAT)", min_value=0)
    duration = st.text_input("4. Expected contract duration (incl. extensions)")
    description = st.text_area("5. Briefly describe the type of service needed and main deliverables")
    submit = st.form_submit_button("🔍 Find Frameworks")

if submit:
    if description.strip() == "":
        st.warning("Please enter a service description.")
    else:
        query_vec = model.encode([description])
        D, I = index.search(np.array(query_vec), k=5)
        candidates = df.iloc[I[0]]

        # Filter based on authority, value cap, and award method
        matches = candidates[
            (candidates["authority_types"].str.contains(authority_type, case=False)) &
            (candidates["value_cap"] >= contract_value) &
            (
                (award_method == "Either") |
                ((award_method == "Direct award") & (candidates["direct_award"] == "Yes")) |
                ((award_method == "Mini-competition") & (candidates["direct_award"] == "No"))
            )
        ]

        st.subheader("🔎 Recommended Frameworks")
        if matches.empty:
            st.error("No suitable frameworks found. You may need to consider running a new procurement competition or expanding your criteria.")
        else:
            top_n = min(len(matches), 3)
            for i in range(top_n):
                row = matches.iloc[i]
                st.success(f"✅ {row['name']}")
                st.markdown(f"- **Provider:** {row['provider']}")
                st.markdown(f"- **Scope:** {row['scope']}")
                st.markdown(f"- **Direct Award:** {row['direct_award']}")
                st.markdown(f"- **Value Cap:** £{row['value_cap']:,}")
                st.markdown(f"- [🔗 View Framework]({row['link']})")
                st.markdown("---")