app.py

# ==============================================================================
# SYSTEM: SADIM AI RADAR 77M - ADVANCED TACTICAL EXPERT INTERFACE
# VERSION: 4.2.0 (CLASSIC TERMINAL UI, STABLE CONE SEARCH, PARALLAX)
# ==============================================================================

import streamlit as st
import streamlit.components.v1 as components
import pandas as pd
import numpy as np
import math
import duckdb
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
import os
import gc

# ------------------------------------------------------------------------------
# 1. GLOBAL PAGE CONFIGURATION
# ------------------------------------------------------------------------------
st.set_page_config(
    page_title="SADIM AI RADAR",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# ------------------------------------------------------------------------------
# 2. ADVANCED CSS STYLING (CLASSIC ACADEMIC TERMINAL UI)
# ------------------------------------------------------------------------------
st.markdown("""
    <style>
    /* Full Page Reset */
    html, body, [data-testid="stAppViewContainer"], [data-testid="stMainViewContainer"], .main, .block-container {
        padding: 0 !important;
        margin: 0 !important;
        width: 100vw !important;
        height: 100vh !important;
        overflow: hidden !important;
        background-color: #000000 !important;
    }

    /* Hide Native UI */
    [data-testid="stHeader"], footer { display: none !important; }
    
    /* THE ACADEMIC SEARCH BAR BLOCK - Sharp Edges, Minimalist */
    div[data-testid="stHorizontalBlock"] {
        position: fixed !important;
        bottom: 30px !important;
        left: 30px !important;
        z-index: 100000 !important;
        background-color: transparent !important;
        width: fit-content !important; 
        display: flex !important;
        flex-direction: row !important;
        gap: 0px !important; /* Zero gap for seamless attachment */
        align-items: center !important;
        box-shadow: 5px 5px 15px rgba(0,0,0,0.9) !important;
    }
    
    div[data-testid="column"] {
        width: auto !important;
        flex: none !important;
        min-width: 0 !important;
        padding: 0 !important;
    }

    /* Input Field - Classic Rectangle */
    div[data-testid="stTextInput"] {
        margin-bottom: 0 !important;
        width: 190px !important; 
    }
    div[data-testid="stTextInput"] > div:first-child { display: none !important; } 
    div[data-testid="stTextInput"] input {
        background-color: #d8d8d8 !important; /* Classic terminal grey */
        color: #000000 !important; 
        border: 1px solid #555555 !important;
        border-right: none !important; /* Open right edge to fuse with button */
        font-family: 'Courier New', Courier, monospace !important; 
        font-size: 13px !important; 
        font-weight: bold !important;
        padding: 4px 10px !important; 
        border-radius: 0px !important; /* Sharp corners */
        height: 34px !important; 
        outline: none !important;
        box-shadow: none !important;
    }
    
    /* Scan Button - Classic Square/Rectangle */
    div[data-testid="stButton"] {
        margin-bottom: 0 !important;
        width: 75px !important;
    }
    div[data-testid="stButton"] button {
        background-color: #d8d8d8 !important; 
        color: #000000 !important; 
        border: 1px solid #555555 !important;
        font-family: 'Arial', sans-serif !important; 
        font-size: 11px !important; 
        font-weight: 900 !important; 
        padding: 0 !important;
        border-radius: 0px !important; /* Sharp corners */
        text-transform: uppercase !important; 
        height: 34px !important; /* Exact match to input height */
        line-height: 34px !important;
        width: 100% !important;
        transition: background-color 0.1s;
        display: block !important;
        cursor: crosshair !important;
    }
    div[data-testid="stButton"] button:hover { background-color: #bbbbbb !important; }
    div[data-testid="stButton"] button:active { background-color: #999999 !important; border-color:#000; }

    /* Aladin Viewport */
    iframe { 
        position: fixed !important;
        top: 0 !important;
        left: -1px !important; 
        width: calc(100vw + 2px) !important;
        height: 100vh !important; 
        border: none !important;
        margin: 0 !important; 
        padding: 0 !important;
    }
    </style>
    """, unsafe_allow_html=True)


# ------------------------------------------------------------------------------
# 3. AI ARCHITECTURE (RESIDUAL BLOCKS)
# ------------------------------------------------------------------------------
class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super(ResidualBlock, self).__init__()
        self.ln = nn.LayerNorm(dim)
        self.fc = nn.Sequential(
            nn.Linear(dim, dim),
            nn.GELU(),
            nn.Linear(dim, dim)
        )
    def forward(self, x):
        return x + self.fc(self.ln(x))

class SADIM_V2_54M_Expert(nn.Module):
    def __init__(self, input_dim=12, latent_dim=64):
        super(SADIM_V2_54M_Expert, self).__init__()
        h = 3584 
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, h), nn.LayerNorm(h), nn.GELU(),
            nn.Linear(h, h), nn.LayerNorm(h), nn.GELU(),
            ResidualBlock(h), ResidualBlock(h), ResidualBlock(h)
        )
        self.fc_mu = nn.Linear(h, latent_dim)
        self.fc_logvar = nn.Linear(h, latent_dim)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, h), nn.GELU(),
            nn.Linear(h, h), nn.LayerNorm(h), nn.GELU(),
            ResidualBlock(h), ResidualBlock(h),
            nn.Linear(h, input_dim)
        )
    def forward(self, x):
        hidden = self.encoder(x)
        mu = self.fc_mu(hidden)
        logvar = torch.clamp(self.fc_logvar(hidden), -10, 10)
        std = torch.exp(0.5 * logvar)
        z = mu + torch.randn_like(std) * std
        return self.decoder(z), mu, logvar

@st.cache_resource
def initialize_deep_ai():
    device = torch.device("cpu")
    model = SADIM_V2_54M_Expert().to(device)
    try:
        hf_token = os.environ.get("HF_TOKEN")
        model_path = hf_hub_download(repo_id="KilmaAI/Sadim-77M", filename="pytorch_model.bin", token=hf_token)
        model.load_state_dict(torch.load(model_path, map_location=device))
    except Exception:
        pass 
    model.eval()
    return model, device

model, device = initialize_deep_ai()


# ------------------------------------------------------------------------------
# 4. STELLAR DATA ENGINE (CONE SEARCH & LOCAL KINEMATICS)
# ------------------------------------------------------------------------------
def execute_stellar_scan(source_id):
    token = os.environ.get("HF_TOKEN")
    if not token:
        return {"error": "SYSTEM_FATAL: HF_TOKEN NOT FOUND IN SECRETS"}

    try:
        db = duckdb.connect(database=':memory:')
        db.execute("INSTALL httpfs; LOAD httpfs;")
        db.execute(f"CREATE SECRET IF NOT EXISTS tac_sec (TYPE HUGGINGFACE, TOKEN '{token}');")
        
        # 4.1. Index Lookup
        index_path = "hf://datasets/samfatnassi/sadim-gaia-index/master_index.parquet"
        index_query = f"SELECT filename FROM '{index_path}' WHERE source_id = {source_id} LIMIT 1"
        index_df = db.query(index_query).df()
        
        if index_df.empty:
            db.close()
            return {"error": "TARGET_ID NOT FOUND IN MASTER INDEX"}
            
        exact_file_path = index_df.iloc[0]['filename']
        if not exact_file_path.startswith("hf://"):
            exact_file_path = "hf://" + exact_file_path
            
        # 4.2. Target Precision Coordinate Extraction
        coord_query = f"SELECT ra, dec FROM '{exact_file_path}' WHERE source_id = {source_id} LIMIT 1"
        coord_df = db.query(coord_query).df()
        if coord_df.empty:
            db.close()
            return {"error": "COORDINATES EXTRACTION FAILED"}
            
        t_ra = coord_df.iloc[0]['ra']
        t_dec = coord_df.iloc[0]['dec']

        # 4.3. Cone Search (Radius 0.05 degrees)
        cone_query = f"""
            SELECT source_id, ra, dec, pmra, pmdec, d_pc, abs_m, bp_rp, l, b, x, y, z
            FROM '{exact_file_path}' 
            WHERE ra BETWEEN {t_ra - 0.05} AND {t_ra + 0.05}
              AND dec BETWEEN {t_dec - 0.05} AND {t_dec + 0.05}
            LIMIT 50
        """
        neighborhood_df = db.query(cone_query).df()
        db.close()
        
    except Exception as e:
        return {"error": f"IO_FAILURE: {str(e)}"}

    if neighborhood_df.empty:
        return {"error": "DATASET RETURNED EMPTY BATCH"}

    # 4.4. Extract Target Data
    target_row = neighborhood_df[neighborhood_df['source_id'] == int(source_id)]
    if target_row.empty:
        return {"error": "TARGET LOST IN CONE SEARCH"}
    row = target_row.iloc[0]

    # Kinematics & Parallax Math
    dist_pc = row['d_pc'] if pd.notna(row['d_pc']) else 1000.0
    dist_ly = dist_pc * 3.26156
    parallax_mas = 1000.0 / dist_pc if dist_pc > 0 else 0.0
    pm_total = math.sqrt(row['pmra']**2 + row['pmdec']**2)
    v_tangential = 4.74 * 1e-3 * pm_total * dist_pc 

    # 4.5. Local Kinematic Anomaly Detection
    neighborhood_df['pm_tot'] = np.sqrt(neighborhood_df['pmra']**2 + neighborhood_df['pmdec']**2)
    local_pm_mean = neighborhood_df['pm_tot'].mean()
    local_pm_std = neighborhood_df['pm_tot'].std()
    
    is_local_outlier = False
    if len(neighborhood_df) > 5 and local_pm_std > 0:
        if pm_total > (local_pm_mean + (3 * local_pm_std)):
            is_local_outlier = True

    # 4.6. AI Inference
    try:
        raw_vector = np.array([
            row['ra'], row['dec'], row['l'], row['b'], 
            row['pmra'], row['pmdec'], row['d_pc'], 
            row['x'], row['y'], row['z'], 
            row['abs_m'], row['bp_rp']
        ], dtype=np.float32)
        
        norm_v = raw_vector.copy()
        norm_v[0] /= 360.0;  norm_v[1] /= 90.0;  norm_v[2] /= 360.0; norm_v[3] /= 90.0
        norm_v[4] /= 500.0;  norm_v[5] /= 500.0;  norm_v[6] /= 10000.0
        norm_v[7] /= 15000.0; norm_v[8] /= 15000.0; norm_v[9] /= 15000.0
        
        input_tensor = torch.tensor(np.nan_to_num(norm_v)).unsqueeze(0).to(device)
        
        with torch.no_grad():
            reconstructed, mu, logvar = model(input_tensor)
            mse_val = torch.mean((reconstructed - input_tensor)**2).item()
            kl_val = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()).item()
            
        hybrid_total = mse_val + (0.1 * kl_val)
    except Exception:
        mse_val, kl_val, hybrid_total = 0.0, 0.0, 0.0

    # 4.7. The Filter Hierarchy (Physics -> Local -> AI)
    final_status = "NOMINAL (NORMAL)"
    final_color = "#888888"

    if v_tangential > 600.0:
        final_status = "PHYSICAL ANOMALY: HYPERVELOCITY"
        final_color = "#ff3333" # Red
    elif is_local_outlier:
        final_status = "KINEMATIC OUTLIER (LOCAL)"
        final_color = "#aa33ff" # Purple
    elif hybrid_total > 5.0:
        final_status = "AI GLOBAL ANOMALY"
        final_color = "#00ffff" # Cyan

    # 4.8. Neighbors Export for JS Map
    neighbors_list = []
    for _, n_row in neighborhood_df.iterrows():
        if int(n_row['source_id']) != int(source_id):
            neighbors_list.append(f"A.source({n_row['ra']}, {n_row['dec']})")
    js_neighbors_array = ",\n".join(neighbors_list)

    result_dict = {
        'id': source_id, 'ra': row['ra'], 'dec': row['dec'], 
        'dist': round(dist_ly, 1), 'plx': round(parallax_mas, 3), 'vel': round(v_tangential, 2), 
        'mse': round(mse_val, 6), 'kl': round(kl_val, 6), 
        'status': final_status, 'color': final_color,
        'neighbors_js': js_neighbors_array,
        'neighbor_count': len(neighbors_list)
    }
    
    # 4.9. Aggressive Memory Sweep
    del neighborhood_df
    del index_df
    del coord_df
    gc.collect()
    
    return result_dict


# ------------------------------------------------------------------------------
# 5. SESSION MANAGEMENT & INTERFACE BUILDER
# ------------------------------------------------------------------------------
if 'session' not in st.session_state:
    st.session_state.session = {
        'data': None, 'is_active': False, 'has_error': False, 
        'error_text': "", 'ra': 269.4485, 'dec': 4.7394
    }

search_col_1, search_col_2 = st.columns([2, 1])
with search_col_1:
    target_sid = st.text_input("ID_SEARCH", placeholder="ENTER SOURCE_ID", label_visibility="collapsed")
with search_col_2:
    if st.button("SCAN"):
        if target_sid and target_sid.isdigit():
            scan_res = execute_stellar_scan(target_sid)
            if "error" in scan_res:
                st.session_state.session.update({'data': None, 'is_active': True, 'has_error': True, 'error_text': scan_res['error']})
            else:
                st.session_state.session.update({'data': scan_res, 'ra': scan_res['ra'], 'dec': scan_res['dec'], 'is_active': True, 'has_error': False})

popup_html_component = ""
if st.session_state.session['is_active']:
    state = st.session_state.session
    if state['has_error']:
        popup_html_component = f"""
        <div id="pop-win" class="tactical-popup" style="border-color:#ff3333; border-radius:0px;">
            <div class="pop-header"><span style="color:#ff3333">! ALERT_FAILURE</span><button onclick="document.getElementById('pop-win').style.display='none'" class="pop-close">×</button></div>
            <div style="color:#ff3333;font-size:11px;text-align:center;padding:10px;font-weight:bold">{state['error_text']}</div>
        </div>"""
    elif state['data']:
        sd = state['data']
        popup_html_component = f"""
        <div id="pop-win" class="tactical-popup" style="border-radius:0px;">
            <div class="pop-header"><span>SADIM 77M // SCAN_REPORT</span><button onclick="document.getElementById('pop-win').style.display='none'" class="pop-close">×</button></div>
            <div class="pop-section">
                <div class="pop-row"><span>SOURCE_ID</span><b>{sd['id']}</b></div>
                <div class="pop-row"><span>DIST_LY</span><b>{sd['dist']}</b></div>
                <div class="pop-row"><span>PARALLAX</span><b>{sd['plx']} mas</b></div>
                <div class="pop-row"><span>VEL_KMS (TAN)</span><b style="color:{'#ff3333' if sd['vel']>600 else '#ffffff'}">{sd['vel']}</b></div>
            </div>
            <div class="pop-section">
                <div class="pop-row"><span>LOCAL NEIGHBORS</span><b>{sd['neighbor_count']} STARS</b></div>
            </div>
            <div class="pop-section-ai">
                <div class="pop-row"><span>MSE_DEV</span><b>{sd['mse']}</b></div>
                <div class="pop-row"><span>KL_SURP</span><b>{sd['kl']}</b></div>
            </div>
            <div class="pop-status" style="color:{sd['color']}; border-top: 1px solid {sd['color']}44;">STATUS: [{sd['status']}]</div>
        </div>"""


# ------------------------------------------------------------------------------
# 6. RADAR VISUALIZER ENGINE (TARGET + CONE SEARCH MAPPING)
# ------------------------------------------------------------------------------
js_neighbors = st.session_state.session['data']['neighbors_js'] if st.session_state.session['data'] else ""

radar_logic_html = f"""
<!DOCTYPE html>
<html>
<head>
    <link rel="stylesheet" href="https://aladin.u-strasbg.fr/AladinLite/api/v2/latest/aladin.min.css" />
    <style>
        body, html {{ margin:0; padding:0; height:100vh; width:100vw; background: black; overflow:hidden; }}
        #radar-canvas {{ width:100%; height:100%; position:absolute; top:0; left:0; }}

        .tactical-popup {{
            position: absolute; top: 25px; right: 25px; z-index: 9999;
            background: rgba(10, 10, 10, 0.95); border: 1px solid #555;
            padding: 15px; font-family: 'Courier New', Courier, monospace; width: 280px;
            color: white; box-shadow: 0 10px 50px rgba(0,0,0,1);
        }}
        .pop-header {{
            display: flex; justify-content: space-between; align-items: center;
            border-bottom: 1px dashed #555; margin-bottom: 12px;
            font-size: 10px; font-weight: bold; letter-spacing: 1px; color: #aaa;
        }}
        .pop-close {{
            background: none; border: none; color: #777; cursor: pointer;
            font-size: 24px; line-height: 1; padding: 0;
        }}
        .pop-close:hover {{ color: white; }}
        .pop-row {{ display: flex; justify-content: space-between; font-size: 10px; margin-bottom: 6px; }}
        .pop-row span {{ color: #777; }}
        .pop-section {{ margin-bottom: 10px; border-bottom: 1px solid #333; padding-bottom: 8px; }}
        .pop-section-ai {{ margin-bottom: 10px; }}
        .pop-status {{
            text-align: center; font-weight: bold; margin-top: 10px;
            padding-top: 12px; font-size: 11px; letter-spacing: 1px;
        }}

        .radar-crosshair {{
            position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%);
            width: 12px; height: 12px; border: 1px solid rgba(255,255,255,0.3); border-radius: 50%;
            z-index: 9998; pointer-events: none;
        }}
        .radar-crosshair::before, .radar-crosshair::after {{
            content: ''; position: absolute; background: rgba(255,255,255,0.7);
        }}
        .radar-crosshair::before {{ top: 50%; left: -6px; right: -6px; height: 1px; }}
        .radar-crosshair::after {{ left: 50%; top: -6px; bottom: -6px; width: 1px; }}

        .aladin-logo, .aladin-searchControl, .aladin-layers-control {{ display: none !important; }}
    </style>
</head>
<body>
    <div id="radar-canvas"></div>
    <div class="radar-crosshair"></div>
    {popup_html_component}

    <script src="https://code.jquery.com/jquery-1.12.1.min.js"></script>
    <script src="https://aladin.u-strasbg.fr/AladinLite/api/v2/latest/aladin.min.js"></script>
    <script>
        var scanActive = {'true' if st.session_state.session['data'] else 'false'};
        var t_ra = {st.session_state.session['ra']};
        var t_dec = {st.session_state.session['dec']};

        var tactical_radar = A.aladin('#radar-canvas', {{ 
            survey: "P/DSS2/color", 
            fov: scanActive ? 2.0 : 0.2, 
            target: t_ra + " " + t_dec, 
            showReticle: false,
            showLayersControl: true,
            showZoomControl: true 
        }});

        if(scanActive){{
            // Layer 1: The Background Local Neighborhood (Blue dots)
            var neighborsLayer = A.catalog({{name:'NEIGHBORS', color:'#4488ff', shape:'circle', sourceSize:6}});
            tactical_radar.addCatalog(neighborsLayer);
            neighborsLayer.addSources([
                {js_neighbors}
            ]);

            // Layer 2: The Main Target (Red/White Cross)
            var targetLayer = A.catalog({{name:'LOCK_ON', color:'#ffffff', shape:'plus', sourceSize:14}});
            tactical_radar.addCatalog(targetLayer);
            targetLayer.addSources([A.source(t_ra, t_dec)]);
            
            // Zoom-in Animation to see the neighborhood
            setTimeout(function(){{
                tactical_radar.setFoV(0.08); 
            }}, 500);
        }}
    </script>
</body>
</html>
"""

components.html(radar_logic_html, height=1000)

# ==============================================================================
# END OF CODE 
# ==============================================================================