import streamlit as st
import pandas as pd
import os
import joblib
import numpy as np
import math
import folium
import json
import base64
from io import BytesIO
from pathlib import Path
from PIL import Image
from streamlit_folium import st_folium

# --- UTILS ---
@st.cache_data
def get_image_base64(img_path):
    img = Image.open(img_path).convert("RGBA")
    w, h = img.size
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("ascii")
    return f"data:image/png;base64,{img_str}", w, h

def calculate_error(z_real, x_real, y_real, z_pred, x_pred, y_pred):
    z_err = abs(z_real - z_pred)
    dist_err = math.sqrt((x_real - x_pred)**2 + (y_real - y_pred)**2)
    return z_err, dist_err

def show_test_inference(cfg):
    st.subheader("🧪 Test Inferenza Offline")
    
    MODEL_DIR = Path("/data/model")
    TEST_SAMPLES_DIR = Path("/data/train/testsamples")
    MAPS_DIR = Path("/data/maps")

    available_models = sorted([m.name for m in MODEL_DIR.glob("model_camp_*.joblib")], reverse=True)
    test_files = sorted([f.name for f in TEST_SAMPLES_DIR.glob("*.csv")])

    if not available_models or not test_files:
        st.warning("Verifica la presenza di modelli e campioni di test.")
        return

    col1, col2 = st.columns(2)
    selected_model = col1.selectbox("🎯 Seleziona Modello:", available_models)
    selected_test = col2.selectbox("📄 Seleziona Fingerprint:", test_files)

    if "test_results" not in st.session_state:
        st.session_state.test_results = None

    if st.button("🚀 ESEGUI TEST DI PRECISIONE", type="primary", use_container_width=True):
        try:
            m_pkg = joblib.load(MODEL_DIR / selected_model)
            delim = cfg.get('paths', {}).get('csv_delimiter', ';')
            df_test = pd.read_csv(TEST_SAMPLES_DIR / selected_test, sep=delim)
            row = df_test.iloc[0]
            
            z_real, x_real, y_real = int(round(float(row['z']))), float(row['x']), float(row['y'])
            gws = m_pkg['gateways_order']
            fill = m_pkg.get('nan_fill', -110.0)

            # --- DEBUG VETTORE INPUT ---
            raw_vals = []
            for gw in gws:
                val = row.get(gw)
                raw_vals.append(float(val) if val is not None and not pd.isna(val) else fill)
            
            X_in = np.array([raw_vals])
            
            # Visualizzazione Debug in UI per l'operatore
            with st.expander("🔍 Analisi Vettore di Input (Fingerprint vs Modello)"):
                debug_df = pd.DataFrame({
                    "Gateway MAC": gws,
                    "RSSI Letto": [row.get(gw, "NON TROVATO") for gw in gws],
                    "RSSI Finale (Input AI)": raw_vals
                })
                st.dataframe(debug_df)
                match_count = np.sum(X_in[0] > fill)
                st.write(f"**Gateway Corrispondenti:** {match_count} su {len(gws)}")
                if match_count == 0:
                    st.error("ERRORE: Il file di test non contiene nessuno dei Gateway usati per l'addestramento!")

            # Predizione
            z_pred = int(m_pkg['floor_clf'].predict(X_in)[0])
            x_pred, y_pred = -1.0, -1.0
            if z_pred in m_pkg['xy_by_floor']:
                xy = m_pkg['xy_by_floor'][z_pred].predict(X_in)[0]
                x_pred, y_pred = float(xy[0]), float(xy[1])

            z_err, dist_err = calculate_error(z_real, x_real, y_real, z_pred, x_pred, y_pred)
            st.session_state.test_results = {
                "z_real": z_real, "x_real": x_real, "y_real": y_real,
                "z_pred": z_pred, "x_pred": x_pred, "y_pred": y_pred,
                "z_err": z_err, "dist_err": dist_err, "model_used": selected_model
            }
        except Exception as e:
            st.error(f"Errore durante l'inferenza: {e}")

    # --- VISUALIZZAZIONE GRAFICA DEI RISULTATI ---
    if st.session_state.test_results:
        res = st.session_state.test_results
        
        c_info1, c_info2 = st.columns(2)
        with c_info1:
            st.info(f"📍 **Test Reale** | Piano: {res['z_real']} | X: **{res['x_real']}** | Y: **{res['y_real']}**")
        with c_info2:
            st.success(f"🔮 **Predizione** | Piano: {res['z_pred']} | X: **{round(res['x_pred'],1)}** | Y: **{round(res['y_pred'],1)}**")

        st.markdown(f"""
        <div style="background-color: #f8f9fa; padding: 12px; border-radius: 8px; border-left: 5px solid #00bcd4; margin: 10px 0;">
            <h5 style="margin:0 0 8px 0;">📍 Legenda Mappa</h5>
            <span style="color: #00bcd4; font-weight: bold;">● PUNTO DI TEST:</span> Posizione reale del rilievo.<br>
            <span style="color: #ff9800; font-weight: bold;">● PUNTO PREDETTO:</span> Posizione calcolata dal modello AI.<br>
            <span style="color: #ffeb3b; font-weight: bold;">--- LINEA GIALLA:</span> Scostamento di <b>{round(res['dist_err'], 2)} px</b>.
        </div>
        """, unsafe_allow_html=True)

        img_filename = f"floor_{res['z_pred']}.png"
        meta_filename = f"meta_{res['z_pred']}.json"
        img_p = MAPS_DIR / img_filename
        meta_p = MAPS_DIR / meta_filename

        if img_p.exists() and meta_p.exists():
            with open(meta_p, "r") as f: meta = json.load(f)
            img_data, w, h = get_image_base64(img_p)
            bounds = [[0, 0], [h, w]]
            
            m = folium.Map(location=[h/2, w/2], crs="Simple", tiles=None, zoom_start=0, attribution_control=False)
            m.fit_bounds(bounds)
            m.options.update({"minZoom": -5, "maxZoom": 5, "maxBounds": bounds, "zoomSnap": 0.25})
            folium.raster_layers.ImageOverlay(image=img_data, bounds=bounds).add_to(m)

            if meta.get("calibrated"):
                def to_px(mx, my):
                    px = (mx * meta["pixel_ratio"]) + meta["origin"][0]
                    py = meta["origin"][1] - (my * meta["pixel_ratio"])
                    return [py, px]

                p_real = to_px(res['x_real'], res['y_real'])
                p_pred = to_px(res['x_pred'], res['y_pred'])

                # Marker Posizione Reale
                folium.CircleMarker(
                    location=p_real, radius=11, color="#00838f", fill=True, 
                    fill_color="#00bcd4", fill_opacity=0.85, tooltip="PUNTO DI TEST (REALE)"
                ).add_to(m)
                
                # Marker Predizione
                if res['x_pred'] != -1.0:
                    folium.CircleMarker(
                        location=p_pred, radius=11, color="#e65100", fill=True, 
                        fill_color="#ff9800", fill_opacity=0.85,
                        tooltip=f"PREDIZIONE (Modello: {res['model_used']})"
                    ).add_to(m)
                    
                    # Linea di errore (Gialla tratteggiata)
                    folium.PolyLine(
                        locations=[p_real, p_pred], color="#ffeb3b", 
                        weight=4, opacity=0.7, dash_array='8'
                    ).add_to(m)

            st_folium(m, height=700, width=None, key=f"test_map_final", use_container_width=True)