import streamlit as st
import pandas as pd
import numpy as np
import faiss
import shap
import os
from sklearn.linear_model import LinearRegression
from sentence_transformers import SentenceTransformer
from groq import Groq

st.set_page_config(page_title="⚽ Explainable Match Summaries", page_icon="⚽", layout="centered")

# ── HuggingFace API Token ─────────────────────────────────────
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "") or os.environ.get("Token", "")

# ── Load Data ─────────────────────────────────────────────────
@st.cache_data
def load_data():
    df = pd.read_excel("ucl_stats.xlsx")
    df = df.dropna(subset=["Home Team", "Away Team", "Home Team Goals", "Away Team Goals"])
    df["Home Team Goals"] = pd.to_numeric(df["Home Team Goals"], errors="coerce").fillna(0)
    df["Away Team Goals"] = pd.to_numeric(df["Away Team Goals"], errors="coerce").fillna(0)

    def outcome(row):
        if row["Home Team Goals"] > row["Away Team Goals"]: return "Home Win"
        elif row["Home Team Goals"] < row["Away Team Goals"]: return "Away Win"
        return "Draw"

    df["Outcome"] = df.apply(outcome, axis=1)
    return df

@st.cache_data
def build_documents(df):
    docs = []
    for _, row in df.iterrows():
        try:
            home  = row["Home Team"]
            away  = row["Away Team"]
            hg    = int(row["Home Team Goals"])
            ag    = int(row["Away Team Goals"])
            phase = row.get("Phase", "UCL 2025")
            winner = row.get("Winner", "Unknown")

            # optional stats
            def safe(col): return row[col] if col in df.columns and pd.notna(row.get(col)) else None

            shots_h  = safe("Home Team Total shots attempts")
            shots_a  = safe("Away Team Total shots attempts")
            poss_col = next((c for c in df.columns if "Possession" in c and "Home" in c), None)
            poss_h   = safe(poss_col) if poss_col else None
            corners_h = safe("Home Corners taken")
            corners_a = safe("Away Corners taken")

            parts = [f"{phase}: {home} vs {away} — Final score {hg}-{ag}. Winner: {winner}."]
            if shots_h and shots_a:
                parts.append(f"{home} had {int(shots_h)} shots; {away} had {int(shots_a)} shots.")
            if poss_h:
                parts.append(f"{home} ball possession: {poss_h}%.")
            if corners_h and corners_a:
                parts.append(f"Corners: {home} {int(corners_h)}, {away} {int(corners_a)}.")

            docs.append(" ".join(parts))
        except Exception:
            continue
    return docs

# ── Embedding + FAISS ─────────────────────────────────────────
@st.cache_resource
def build_index(docs):
    model = SentenceTransformer("all-MiniLM-L6-v2")
    embeddings = model.encode(docs, show_progress_bar=False)
    idx = faiss.IndexFlatL2(embeddings.shape[1])
    idx.add(np.array(embeddings))
    return model, idx, embeddings

def retrieve(query, model, index, docs, top_k=3):
    q_emb = model.encode([query])
    _, indices = index.search(np.array(q_emb), top_k)
    return [docs[i] for i in indices[0]]

# ── LLM via HuggingFace API ───────────────────────────────────
def generate_summary(query, evidence, groq_key):
    evidence_text = "\n".join([f"- {e}" for e in evidence])
    try:
        client = Groq(api_key=groq_key)
        response = client.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=[
                {
                    "role": "system",
                    "content": "You are a UEFA Champions League analyst. Generate concise factual match summaries using ONLY the evidence provided. Do NOT invent facts. Keep it under 100 words."
                },
                {
                    "role": "user",
                    "content": f"QUERY: {query}\n\nEVIDENCE:\n{evidence_text}\n\nWrite a concise factual summary:"
                }
            ],
            max_tokens=150,
            temperature=0.2,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"⚠️ API Error: {str(e)}"

# ── SHAP Explainability ───────────────────────────────────────
def compute_shap(df, home_team, away_team):
    try:
        shap_df = pd.DataFrame({
            "goals_scored":   pd.to_numeric(df["Home Team Goals"], errors="coerce").fillna(0),
            "goals_conceded": pd.to_numeric(df["Away Team Goals"], errors="coerce").fillna(0),
        })
        shap_df["goal_difference"] = shap_df["goals_scored"] - shap_df["goals_conceded"]
        shap_df["is_win"]          = (shap_df["goal_difference"] > 0).astype(int)
        shap_df = shap_df.dropna()

        X = shap_df[["goals_scored", "goals_conceded", "goal_difference", "is_win"]]
        y = shap_df["goal_difference"]

        model = LinearRegression()
        model.fit(X, y)

        explainer   = shap.LinearExplainer(model, X)
        shap_values = explainer.shap_values(X)
        mean_shap   = np.abs(shap_values).mean(axis=0)

        return dict(zip(X.columns, mean_shap))
    except Exception:
        return {}

# ── UI ────────────────────────────────────────────────────────
st.title("⚽ Explainable UCL Match Summaries")
st.markdown(
    "RAG + **Mistral-7B** + **SHAP** explainability on the self-curated "
    "UCL 2025 dataset (189 matches). Grounded summaries — no hallucination."
)
st.divider()

# Load
with st.spinner("Loading UCL 2025 dataset..."):
    df   = load_data()
    docs = build_documents(df)

with st.spinner("Building FAISS index with Sentence-BERT..."):
    emb_model, faiss_index, _ = build_index(docs)

st.success(f"✅ {len(docs)} match records indexed | {len(df['Home Team'].unique())} teams")

# Token input
st.subheader("🔑 Groq API Key")
if GROQ_API_KEY:
    st.info("✅ Groq API key loaded from Space secrets (GROQ_API_KEY)")
    token = GROQ_API_KEY
else:
    token = st.text_input(
        "Enter your Groq API key (free at console.groq.com):",
        type="password",
        placeholder="gsk_..."
    )

st.divider()

# Query input
st.subheader("🔍 Ask About a Match")

teams = sorted(set(df["Home Team"].dropna()) | set(df["Away Team"].dropna()))
col1, col2 = st.columns(2)
with col1:
    team1 = st.selectbox("Team 1", teams, index=teams.index("Real Madrid") if "Real Madrid" in teams else 0)
with col2:
    team2 = st.selectbox("Team 2", teams, index=teams.index("Liverpool") if "Liverpool" in teams else 1)

query_type = st.selectbox("Query type", [
    "Match summary",
    "Who won?",
    "Goals and shots analysis",
    "Possession and corners breakdown",
])

query_map = {
    "Match summary":                f"Summarize the match between {team1} and {team2}",
    "Who won?":                     f"Who won the match between {team1} and {team2} and by how many goals?",
    "Goals and shots analysis":     f"Analyze the goals and shots for {team1} vs {team2}",
    "Possession and corners breakdown": f"Describe the possession and corners for {team1} vs {team2}",
}

custom_query = st.text_input("Or type your own query:", placeholder=f"e.g. How did {team1} perform against {team2}?")
final_query  = custom_query if custom_query.strip() else query_map[query_type]

if st.button("🚀 Generate Summary", type="primary"):
    if not token:
        st.error("❌ Please enter your Groq API key above!")
    else:
        with st.spinner("🔍 Retrieving evidence from FAISS..."):
            evidence = retrieve(final_query, emb_model, faiss_index, docs, top_k=3)

        with st.spinner("🤖 Generating summary with Mistral-7B..."):
            summary = generate_summary(final_query, evidence, token)

        st.divider()
        st.subheader("📋 Generated Summary")
        st.success(summary)

        st.subheader("📚 Retrieved Evidence (RAG)")
        for i, ev in enumerate(evidence, 1):
            st.info(f"**Evidence {i}:** {ev}")

        # SHAP
        st.subheader("🔍 SHAP Feature Importance")
        shap_scores = compute_shap(df, team1, team2)
        if shap_scores:
            shap_df_display = pd.DataFrame({
                "Feature":    list(shap_scores.keys()),
                "SHAP Value": [round(v, 4) for v in shap_scores.values()]
            }).sort_values("SHAP Value", ascending=False)

            st.bar_chart(shap_df_display.set_index("Feature"))
            st.caption("SHAP values show which features most influenced the match outcome prediction.")
        else:
            st.warning("SHAP computation unavailable for this query.")

st.divider()
st.markdown(
    "Built by **Bharath Kesav R** · "
    "[GitHub](https://github.com/bk1210) · "
    "[Portfolio](https://bk1210.github.io/portfolio) · "
    "Model: Mistral-7B-Instruct via HuggingFace API · RAG: Sentence-BERT + FAISS"
)