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)
            X_in = np.array([[float(row.get(gw, fill)) for gw in gws]])
            
            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: {e}")

    if st.session_state.test_results:
        res = st.session_state.test_results
        
        # --- RIPRISTINO INFORMAZIONI COORDINATE ---
        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)}**")

        # --- LEGENDA AGGIORNATA CON NUOVI COLORI ---
        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'])

                # 🔵 PUNTO DI TEST (Celeste)
                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)
                
                # 🟠 PUNTO PREDETTO (Arancio)
                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 Errore
                    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)