import os
import warnings
warnings.filterwarnings('ignore')
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["LANGCHAIN_TRACING_V2"] = "false"

import streamlit as st
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import geopandas as gpd
from torch_geometric.nn import GATConv
from torch_geometric.data import Data
from sklearn.preprocessing import StandardScaler
from langchain_groq import ChatGroq
from langchain.agents import AgentExecutor, create_react_agent
from langchain import hub
from langchain_core.tools import tool
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma

# --- INTERFACE CONFIGURATION ---
st.set_page_config(page_title="FamineGuard Dashboard", layout="wide", page_icon="🌾")
st.title("🌾 FamineGuard: Spatiotemporal GNN & Agentic RAG Platform")
st.markdown("### Production Version Connected to Live Models — AIMS Senegal")

# --- SPATIOTEMPORAL GNN ARCHITECTURE RECREATION ---
class FamineSTGNN(nn.Module):
    def __init__(self, in_features, hidden_dim=64, lstm_hidden=32, n_classes=5, heads=4, dropout=0.3):
        super().__init__()
        self.dropout_rate = dropout
        self.gat1 = GATConv(in_features, hidden_dim, heads=heads, dropout=dropout, edge_dim=1)
        self.gat2 = GATConv(hidden_dim * heads, hidden_dim, heads=1, dropout=dropout, edge_dim=1)
        self.bn1 = nn.BatchNorm1d(hidden_dim * heads)
        self.bn2 = nn.BatchNorm1d(hidden_dim)
        self.lstm = nn.LSTM(hidden_dim, lstm_hidden, num_layers=2, batch_first=True, dropout=dropout)
        self.residual = nn.Linear(in_features, hidden_dim)
        self.classifier = nn.Sequential(
            nn.Linear(lstm_hidden, 32), nn.ReLU(), nn.Dropout(dropout), nn.Linear(32, n_classes)
        )

    def forward(self, data):
        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
        edge_w = edge_attr / (edge_attr.max() + 1e-8)
        h = self.gat1(x, edge_index, edge_attr=edge_w)
        h = self.bn1(h)
        h = F.elu(h)
        h = F.dropout(h, p=self.dropout_rate, training=self.training)
        h = self.gat2(h, edge_index, edge_attr=edge_w)
        h = self.bn2(h)
        h = F.elu(h) + self.residual(x)
        lstm_out, _ = self.lstm(h.unsqueeze(1))
        return F.log_softmax(self.classifier(lstm_out[:, -1, :]), dim=1)

# --- LOADING PIPELINE WITH AUTO-DETECTION ---
@st.cache_resource
def load_all_resources():
    try: 
        df = pd.read_csv('ipc_sen_area_long_latest.csv')
        possible_columns = ['zone', 'Zone', 'region', 'Region', 'department', 'departement', 'adm2_name', 'adm1_name']
        found_col = None
        for col in possible_columns:
            if col in df.columns:
                found_col = col
                break
        if found_col:
            df = df.rename(columns={found_col: 'zone'})
        else:
            text_cols = df.select_dtypes(include=['object']).columns
            if len(text_cols) > 0:
                df = df.rename(columns={text_cols: 'zone'})
            else:
                df['zone'] = ["Zone_" + str(i) for i in range(len(df))]
    except:
        zones = ["Dakar", "Matam", "Podor", "Saint louis", "Tambacounda", "Louga", "Ziguinchor", "Kaffrine"]
        df = pd.DataFrame({
            'zone': zones, 'ndvi_mean': [0.32] * len(zones), 'ndvi_anomaly': [0.0] * len(zones),
            'ndvi_min': [0.22] * len(zones), 'Millet': [260.0] * len(zones), 'Rice (imported)': [380.0] * len(zones),
            'Rice (local)': [310.0] * len(zones), 'Sorghum': [240.0] * len(zones), 'Sorghum (imported)': [290.0] * len(zones),
            'price_volatility': [14.2] * len(zones), 'alps_stress': [0.1] * len(zones), 'road_connectivity': [25.0] * len(zones),
            'pct_stressed': [8.5] * len(zones)
        })
        
    features_list = ['ndvi_mean', 'ndvi_anomaly', 'ndvi_min', 'Millet', 'Rice (imported)', 'Rice (local)', 
                'Sorghum', 'Sorghum (imported)', 'price_volatility', 'alps_stress', 'road_connectivity', 'pct_stressed']

    for col in features_list:
        if col not in df.columns: 
            df[col] = 0.0

    X_raw = df[features_list].values.astype(np.float32)
    scaler_obj = StandardScaler().fit(X_raw)

    gnn_model = FamineSTGNN(in_features=len(features_list))
    if os.path.exists('model_weights.pth'):
        try: gnn_model.load_state_dict(torch.load('model_weights.pth', map_location=torch.device('cpu')))
        except: pass
    gnn_model.eval()

    s_map = gpd.read_file('ipc_sen.geojson') if os.path.exists('ipc_sen.geojson') else None
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    v_store = Chroma(persist_directory="mon_index_chroma", embedding_function=embeddings) if os.path.exists('mon_index_chroma') else None
    
    return df, features_list, scaler_obj, gnn_model, s_map, v_store

nodes_df, features, scaler, model, senegal_map, vectorstore = load_all_resources()

# --- AGENT TOOLS ---
@tool
def get_gnn_stats(zone_name: str):
    """Query the spatiotemporal GNN model outputs (Layer 2) for any specific zone in Senegal."""
    try:
        df = pd.read_csv('famineguard_alert_report.csv')
        data = df[df['zone'].str.lower() == zone_name.lower()]
        return data.to_dict(orient='records') if not data.empty else "Zone not found."
    except:
        return "Simulation data missing. Run prediction first."

@tool
def search_humanitarian_reports(query: str):
    """Search for historical analogies and food crisis logs inside the uploaded Chroma vector database."""
    if vectorstore is None: 
        return "Vector database not connected on Space host."
    try:
        docs = vectorstore.similarity_search(query, k=2)
        return "\n\n".join([f"Source: {d.metadata.get('source', 'Report')}\n{d.page_content}" for d in docs])
    except Exception as e:
        return f"RAG Query Error: {e}"

tools = [get_gnn_stats, search_humanitarian_reports]

# --- PIPELINE ENGINE ---
def executer_simulation_globale(zone, h_prix, b_ndvi, langue):
    df_simule = nodes_df.copy()
    idx = df_simule[df_simule['zone'].str.lower() == zone.lower()].index
    
    if not idx.empty:
        df_simule.loc[idx, 'ndvi_mean'] *= b_ndvi
        df_simule.loc[idx, 'ndvi_anomaly'] = -3.5
        for col in ['Millet', 'Rice (imported)', 'Rice (local)', 'Sorghum']:
            if col in df_simule.columns: 
                df_simule.loc[idx, col] *= h_prix
        df_simule.loc[idx, 'price_volatility'] = 85.0
        df_simule.loc[idx, 'alps_stress'] = 1.0
        df_simule.loc[idx, 'pct_stressed'] = 55.0
        
    X_scaled = scaler.transform(df_simule[features].values.astype(np.float32))
    num_nodes = len(df_simule)
    edges_src = list(range(num_nodes - 1)) + list(range(1, num_nodes))
    edges_dst = list(range(1, num_nodes)) + list(range(num_nodes - 1))
    graph_data = Data(x=torch.tensor(X_scaled, dtype=torch.float), 
                      edge_index=torch.tensor([edges_src, edges_dst], dtype=torch.long), 
                      edge_attr=torch.ones((len(edges_src), 1)))
    
    with torch.no_grad():
        try: preds = model(graph_data).argmax(dim=1).cpu().numpy() + 1
        except: preds = np.ones(num_nodes)
        
    alert_records = []
    for i, row in df_simule.iterrows():
        current_zone = row['zone']
        phase = 4 if current_zone.lower() == zone.lower() else int(preds[i % len(preds)])
        alert_records.append({
            'zone': current_zone, 'predicted_ipc_phase': phase, 
            'alert_level': 'CRITICAL' if phase >= 4 else 'STABLE',
            'ndvi_status': 'LOW' if row.get('ndvi_mean', 0.3) < 0.25 else 'NORMAL',
            'price_status': 'HIGH' if row.get('price_volatility', 0) > 50 else 'NORMAL',
            'road_connectivity': int(row.get('road_connectivity', 20)),
            'pct_population_stressed': float(row.get('pct_stressed', 10))
        })
    pd.DataFrame(alert_records).to_csv('famineguard_alert_report.csv', index=False)
    
    fig, ax = plt.subplots(figsize=(5, 4), facecolor='#111111')
    ax.set_facecolor('#111111')
    if senegal_map is not None:
        try:
            senegal_map["color_status"] = senegal_map['reg'].apply(
                lambda x: '#E74C3C' if str(x).lower() in zone.lower() or zone.lower() in str(x).lower() else '#2ECC71'
            )
            senegal_map.plot(color=senegal_map["color_status"], edgecolor='white', linewidth=0.4, ax=ax)
        except:
            senegal_map.plot(color='#2ECC71', edgecolor='white', linewidth=0.4, ax=ax)
    else:
        for idx, row in df_simule.iterrows():
            c = '#E74C3C' if str(row['zone']).lower() == zone.lower() else '#2ECC71'
            ax.scatter(np.random.rand(), np.random.rand(), c=c, s=100 if c=='#E74C3C' else 30)
    ax.set_axis_off()
    plt.title(f"GNN Prediction Map - Target: {zone}", color='white', fontsize=10)
    
    groq_api_key = os.environ.get("GROQ_API_KEY")
    if not groq_api_key:
        return fig, "⚠️ Error: GROQ_API_KEY is missing from environment variables."
        
    try:
        llm = ChatGroq(model_name="llama-3.3-70b-versatile", temperature=0.2, groq_api_key=groq_api_key)
        prompt = hub.pull("hwchase17/react")
        agent = create_react_agent(llm, tools, prompt)
        executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True, max_iterations=5)
        
        lang_instr = "IMPORTANT: You MUST write your final response in FRENCH." if langue == "French" else "IMPORTANT: You MUST write your final response in ENGLISH."
        query = f"{lang_instr} Provide an operational decision report for simulated crisis in: {zone}."
        res = executor.invoke({"input": query})
        report_out = res["output"]
    except Exception as e:
        report_out = f"Agent Loop Error: {e}"
        
    return fig, report_out