import streamlit as st
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsRegressor

# ---------------------------
# Page config
# ---------------------------
st.set_page_config(
    page_title="Neural Tuner – Car Mod Performance Estimator",
    page_icon="🔥",
    layout="wide"
)

# ---------------------------
# SESSION STATE INITIALIZATION (Fixes Yellow Warning)
# ---------------------------
def init_session_state():
    # Only set these defaults if they don't exist yet
    defaults = {
        "engine_disp": 2.0,
        "engine_layout": "I6",
        "base_hp": 240,
        "weight_kg": 1500,
        "weight_reduction": 0,
        "intake": "Stock",
        "headers": "Stock",
        "exhaust": "Stock",
        "exhaust_dia": 60,
        "cam": "Stock",
        "intake_manifold": "Stock",
        "induction": "None",
        "boost_psi": 0,
        "turbo_size": "N/A",
        "intercooler": "None",
        "meth": False,
        "tune": "None",
        "fuel": "93",
        "altitude": 200,
        "traction_mode": "Spirited"
    }
    
    for key, value in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value

# Run init immediately
init_session_state()

# ---------------------------
# Global CSS (Fixes "Cutting Out" / Rounds)
# ---------------------------
st.markdown(
    """
    <style>
    html, body, [data-testid="stAppViewContainer"] {
        background: radial-gradient(circle at top, #020617 0, #020617 35%, #020617 40%, #000000 100%) !important;
        color: #e5e7eb;
        font-family: system-ui, -apple-system, BlinkMacSystemFont, "SF Pro Text",
                     "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
    }
    [data-testid="stHeader"] { background: transparent; }
    [data-testid="stToolbar"] { display: none; }

    /* Fix: Removed overflow:hidden to stop cutting off shadows */
    .hero {
        border-radius: 24px;
        padding: 24px 28px;
        background: radial-gradient(circle at top left, rgba(96,165,250,0.35), transparent 55%),
                    radial-gradient(circle at bottom right, rgba(236,72,153,0.35), transparent 55%),
                    rgba(15,23,42,0.94);
        box-shadow: 0 20px 50px rgba(0,0,0,0.5); 
        border: 1px solid rgba(148,163,184,0.35);
        position: relative;
        margin-bottom: 20px;
    }
    
    .hero-title {
        font-size: 1.9rem;
        font-weight: 700;
        background: linear-gradient(90deg, #f97316, #facc15, #22c55e, #38bdf8, #a855f7, #f97316);
        background-size: 400% 100%;
        -webkit-background-clip: text;
        color: transparent;
        animation: moveGradient 9s ease infinite;
    }
    .hero-sub {
        color: #9ca3af;
        font-size: 0.95rem;
    }
    .hero-pill {
        display: inline-flex;
        align-items: center;
        gap: 6px;
        padding: 3px 11px;
        border-radius: 999px;
        background: rgba(15,118,110,0.2);
        border: 1px solid rgba(34,197,94,0.6);
        font-size: 0.75rem;
        color: #bbf7d0;
        margin-right: 8px;
    }
    @keyframes moveGradient {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }

    /* Fix: Removed overflow:hidden and added margin for shadow breathing room */
    .glass {
        background: radial-gradient(circle at top left, rgba(148,163,184,0.24), transparent 55%),
                    rgba(15,23,42,0.96);
        border-radius: 20px;
        padding: 20px;
        border: 1px solid rgba(148,163,184,0.4);
        box-shadow: 0 15px 35px rgba(0,0,0,0.6);
        margin: 5px; 
    }

    .chip {
        display:inline-block;
        padding:4px 10px;
        margin:3px 4px 3px 0;
        border-radius:999px;
        font-size:0.78rem;
        background:rgba(59,130,246,0.14);
        border:1px solid rgba(59,130,246,0.45);
        color:#bfdbfe;
    }

    .meter-label {
        font-size: 0.85rem;
        color: #9ca3af;
        margin-bottom: 3px;
    }
    .muted {
        font-size: 0.8rem;
        color: #6b7280;
    }

    /* Tabs styling */
    button[data-baseweb="tab"] {
        background: transparent !important;
        border-radius: 999px !important;
        padding: 0.5rem 1rem !important;
        margin-right: 0.4rem;
        color: #9ca3af !important;
    }
    button[data-baseweb="tab"][aria-selected="true"] {
        background: rgba(59,130,246,0.2) !important;
        color: #e5e7eb !important;
        box-shadow: 0 0 0 1px rgba(59,130,246,0.7);
    }

    /* Rev limiter gauge - Only this keeps overflow hidden for the semi-circle */
    .rev-wrap {
        display:flex;
        justify-content:center;
        align-items:center;
        margin-top:4px;
        margin-bottom:4px;
    }
    .rev-gauge {
        position:relative;
        width:170px;
        height:90px;
        border-radius:170px 170px 14px 14px;
        background: radial-gradient(circle at 50% 120%, #991b1b, transparent 60%),
                    radial-gradient(circle at 0% 0%, rgba(59,130,246,0.4), transparent 60%),
                    #020617;
        overflow:hidden; 
        border:1px solid rgba(148,163,184,0.5);
        box-shadow: 0 12px 30px rgba(0,0,0,0.8);
    }
    .rev-arc {
        position:absolute;
        inset:10px 10px auto 10px;
        height:62px;
        border-radius:999px;
        border-top:6px solid rgba(249,250,251,0.16);
        border-left:6px solid rgba(34,197,94,0.75);
        border-right:6px solid rgba(220,38,38,0.9);
        border-bottom:none;
    }
    .rev-needle {
        --base-angle: -40deg;
        position:absolute;
        left:50%;
        bottom:8px;
        width:3px;
        height:68px;
        background: linear-gradient(to top, #fecaca, #f97316);
        border-radius:999px;
        transform-origin:50% 90%;
        animation: revBounce 0.6s ease-in-out infinite alternate;
    }
    .rev-center {
        position:absolute;
        left:50%;
        bottom:4px;
        transform:translateX(-50%);
        width:22px;
        height:22px;
        border-radius:999px;
        background:#020617;
        border:2px solid #9ca3af;
        display:flex;
        align-items:center;
        justify-content:center;
        font-size:0.55rem;
        color:#9ca3af;
    }
    .rev-label {
        position:absolute;
        right:8px;
        bottom:6px;
        font-size:0.7rem;
        color:#fecaca;
    }
    @keyframes revBounce {
        from { transform: rotate(var(--base-angle)); }
        to   { transform: rotate(calc(var(--base-angle) + 14deg)); }
    }
    </style>
    """,
    unsafe_allow_html=True
)

# ---------------------------
# Data & model
# ---------------------------
def generate_dataset(n=400):
    np.random.seed(42)
    data = []
    for _ in range(n):
        engine = np.random.choice([1.6, 2.0, 2.5, 3.0, 3.5, 5.0])
        cyl = np.random.choice([4, 6, 8])
        base_hp = int(engine * cyl * np.random.uniform(18, 22))

        intake = np.random.choice([0, 1, 2])
        exhaust = np.random.choice([0, 1, 2])
        induction = np.random.choice([0, 1, 2])
        fuel = np.random.choice([0, 1, 2, 3])
        tune = np.random.choice([0, 1, 2])
        altitude = np.random.uniform(0, 2000)

        hp_gain = (
            intake * np.random.uniform(3, 10) +
            exhaust * np.random.uniform(5, 20) +
            induction * np.random.uniform(25, 100) +
            tune * np.random.uniform(10, 35) +
            fuel * np.random.uniform(2, 8) -
            altitude * 0.01 +
            np.random.uniform(-3, 3)
        )

        data.append([engine, cyl, base_hp, intake, exhaust,
                     induction, fuel, tune, altitude, hp_gain])

    columns = [
        "engine", "cyl", "base_hp", "intake", "exhaust",
        "induction", "fuel", "tune", "altitude", "hp_gain"
    ]
    return pd.DataFrame(data, columns=columns)

df = generate_dataset()
X = df.drop("hp_gain", axis=1)
y = df["hp_gain"]

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

model = KNeighborsRegressor(n_neighbors=5, weights="distance")
model.fit(X_scaled, y)

# ---------------------------
# PRESET CALLBACKS
# ---------------------------
def preset_daily():
    s = st.session_state
    s.engine_disp = 1.6
    s.engine_layout = "I4"
    s.base_hp = 120
    s.weight_kg = 1350
    s.weight_reduction = 0
    s.intake = "Stock"
    s.headers = "Stock"
    s.exhaust = "Stock"
    s.exhaust_dia = 50
    s.cam = "Stock"
    s.intake_manifold = "Stock"
    s.induction = "None"
    s.boost_psi = 0
    s.turbo_size = "N/A"
    s.intercooler = "None"
    s.meth = False
    s.tune = "None"
    s.fuel = "87"
    s.altitude = 200
    s.traction_mode = "Daily"

def preset_street():
    s = st.session_state
    s.engine_disp = 2.0
    s.engine_layout = "I4"
    s.base_hp = 190
    s.weight_kg = 1400
    s.weight_reduction = 5
    s.intake = "Cold Air"
    s.headers = "Shorty"
    s.exhaust = "Cat-back"
    s.exhaust_dia = 60
    s.cam = "Stage 1 – Road"
    s.intake_manifold = "High-flow"
    s.induction = "Turbo"
    s.boost_psi = 12
    s.turbo_size = "Small 45-55mm"
    s.intercooler = "Air-to-Air"
    s.meth = False
    s.tune = "Stage 1"
    s.fuel = "93"
    s.altitude = 150
    s.traction_mode = "Spirited"

def preset_drag():
    s = st.session_state
    s.engine_disp = 5.0
    s.engine_layout = "V8"
    s.base_hp = 420
    s.weight_kg = 1500
    s.weight_reduction = 18
    s.intake = "Performance"
    s.headers = "Long Tube"
    s.exhaust = "Straight Pipe"
    s.exhaust_dia = 90
    s.cam = "Stage 2 – Aggressive"
    s.intake_manifold = "High-flow"
    s.induction = "Twin-Turbo"
    s.boost_psi = 24
    s.turbo_size = "Big 66-75mm"
    s.intercooler = "Front-mount High-Flow"
    s.meth = True
    s.tune = "Kill Mode"
    s.fuel = "E85 / Race Blend"
    s.altitude = 100
    s.traction_mode = "Track / Drag"

def preset_max():
    s = st.session_state
    s.engine_disp = 6.0
    s.engine_layout = "V12"
    s.base_hp = 650
    s.weight_kg = 1400
    s.weight_reduction = 25
    s.intake = "Performance"
    s.headers = "Long Tube"
    s.exhaust = "Straight Pipe"
    s.exhaust_dia = 100
    s.cam = "Stage 3 – Race"
    s.intake_manifold = "Individual throttle bodies"
    s.induction = "Twincharged"
    s.boost_psi = 30
    s.turbo_size = "XL 76mm+"
    s.intercooler = "Front-mount High-Flow"
    s.meth = True
    s.tune = "Kill Mode"
    s.fuel = "E85 / Race Blend"
    s.altitude = 0
    s.traction_mode = "Track / Drag"

# ---------------------------
# HERO SECTION
# ---------------------------
st.markdown(
    """
    <div class="hero">
        <div style="display:flex;justify-content:space-between;align-items:center;gap:14px;">
            <div>
                <div class="hero-pill">⚙️ KNN model + synthetic dyno data</div>
                <div class="hero-pill">🎛️ Live tuning playground</div>
                <div class="hero-title">Neural Tuner – Car Mod Performance Lab</div>
                <p class="hero-sub">
                    Mash intakes, boost, tunes and fuel in real time. 
                    Watch horsepower climb and the rev limiter dance. Built to make car nerds go feral. 🐺
                </p>
            </div>
            <div style="
                width:170px;height:110px;
                border-radius:22px;
                background:conic-gradient(from 220deg,
                    #22c55e, #38bdf8, #a855f7, #f97316, #facc15, #22c55e);
                padding:2px;
            ">
                <div style="
                    width:100%;height:100%;
                    border-radius:19px;
                    background:radial-gradient(circle at 30% 0%, rgba(248,250,252,0.2), transparent 55%),
                               #020617;">
                    <div style="display:flex;flex-direction:column;justify-content:center;align-items:center;height:100%;">
                        <div style="font-size:0.78rem;color:#9ca3af;">Virtual</div>
                        <div style="font-size:1.6rem;font-weight:700;color:#e5e7eb;">Dyno</div>
                        <div style="font-size:0.75rem;color:#22c55e;">No hardware, all chaos</div>
                    </div>
                </div>
            </div>
        </div>
    </div>
    """,
    unsafe_allow_html=True
)

st.markdown("")

# ---------------------------
# Layout
# ---------------------------
left, right = st.columns([1.15, 1])

# ========= LEFT SIDE =========
with left:
    st.markdown('<div class="glass">', unsafe_allow_html=True)
    st.subheader("🎛️ Tune Your Build")

    # Preset buttons
    bcol1, bcol2, bcol3, bcol4 = st.columns(4)
    bcol1.button("🚙 Daily Driver", on_click=preset_daily)
    bcol2.button("🌃 Street Fun", on_click=preset_street)
    bcol3.button("🏁 Drag Strip", on_click=preset_drag)
    bcol4.button("🔥 Max Attack", on_click=preset_max)

    st.markdown("")

    tabs = st.tabs(["Core Specs", "Bolt-Ons", "Boost & Cooling", "ECU & Fuel"])

    # NOTE: I removed 'index=...' from these selectboxes.
    # They now just rely on the key being present in st.session_state (handled by init_session_state)
    
    # --- Core specs ---
    with tabs[0]:
        col1, col2 = st.columns(2)
        engine_disp_options = [0.8,1.0,1.2,1.4,1.6,1.8,2.0,2.2,2.4,2.5,
                               2.8,3.0,3.2,3.5,4.0,4.4,5.0,6.0,7.0,8.0]
        
        # We manually update value if it's missing (failsafe), though init handles it.
        if "engine_disp" not in st.session_state: st.session_state.engine_disp = 2.0
        
        engine_disp = col1.selectbox(
            "Engine Displacement (L)",
            engine_disp_options,
            key="engine_disp"
        )
        engine_layout = col2.selectbox(
            "Engine Layout",
            ["I3","I4","I6","V6","V8","V10","V12"],
            key="engine_layout"
        )

        layout_to_cyl = {"I3":3,"I4":4,"I6":6,"V6":6,"V8":8,"V10":10,"V12":12}
        cyl = layout_to_cyl[engine_layout]

        base_hp = st.number_input(
            "Base Horsepower (stock dyno)",
            min_value=60, max_value=3000, 
            key="base_hp"
        )

        col3, col4 = st.columns(2)
        weight_kg = col3.number_input(
            "Vehicle Weight (kg)",
            min_value=700, max_value=4000,
            key="weight_kg"
        )
        weight_reduction = col4.slider(
            "Weight Reduction (%)",
            0, 40,
            key="weight_reduction"
        )

    # --- Bolt-ons ---
    with tabs[1]:
        col1, col2, col3 = st.columns(3)
        intake = col1.selectbox("Intake", ["Stock", "Cold Air", "Performance"], key="intake")
        headers = col2.selectbox("Headers", ["Stock", "Shorty", "Long Tube"], key="headers")
        exhaust = col3.selectbox("Exhaust", ["Stock", "Cat-back", "Straight Pipe"], key="exhaust")
        exhaust_dia = st.slider("Exhaust Diameter (mm)", 40, 120, key="exhaust_dia")

        cam = st.selectbox(
            "Cam Profile",
            ["Stock", "Stage 1 – Road", "Stage 2 – Aggressive", "Stage 3 – Race"],
            key="cam"
        )
        intake_manifold = st.selectbox(
            "Intake Manifold",
            ["Stock", "High-flow", "Individual throttle bodies"],
            key="intake_manifold"
        )

    # --- Boost & cooling ---
    with tabs[2]:
        induction = st.selectbox(
            "Forced Induction Setup",
            ["None", "Turbo", "Twin-Turbo", "Supercharger", "Twincharged"],
            key="induction"
        )
        boost_psi = st.slider("Target Boost (psi)", 0, 40, key="boost_psi")
        turbo_size = st.selectbox(
            "Turbo Size",
            ["N/A", "Small 45-55mm", "Medium 56-65mm", "Big 66-75mm", "XL 76mm+"],
            key="turbo_size"
        )
        intercooler = st.selectbox(
            "Intercooler Type",
            ["None", "Air-to-Air", "Air-to-Water", "Front-mount High-Flow"],
            key="intercooler"
        )
        meth = st.checkbox("Methanol Injection Kit", key="meth")

    # --- ECU & fuel ---
    with tabs[3]:
        tune = st.selectbox(
            "ECU Tune Level",
            ["None", "Mild Street", "Stage 1", "Stage 2", "Kill Mode"],
            key="tune"
        )
        fuel = st.selectbox(
            "Fuel Type",
            ["87", "91", "93", "E85 / Race Blend"],
            key="fuel"
        )
        altitude = st.slider("Altitude (meters)", 0, 3500, key="altitude")
        traction_mode = st.radio(
            "Traction Mode Vibe",
            ["Daily", "Spirited", "Track / Drag"],
            horizontal=True,
            key="traction_mode"
        )

    st.markdown(
        '<p class="muted" style="margin-top:8px;">Numbers are synthetic; this is a dyno-inspired playground, not a tuning bible. 🧪</p>',
        unsafe_allow_html=True
    )
    st.markdown('</div>', unsafe_allow_html=True)

# ========= MODEL & CALCS =========

intake_map = {"Stock":0, "Cold Air":1, "Performance":2}
exhaust_map = {"Stock":0, "Cat-back":1, "Straight Pipe":2}
induction_model_map = {"None":0, "Turbo":1, "Twin-Turbo":1, "Supercharger":2, "Twincharged":2}
fuel_map = {"87":0, "91":1, "93":2, "E85 / Race Blend":3}
tune_base_map = {"None":0, "Mild Street":1, "Stage 1":1, "Stage 2":2, "Kill Mode":2}

input_for_model = np.array([[
    engine_disp,
    cyl,
    base_hp,
    intake_map.get(intake,0),
    exhaust_map.get(exhaust,0),
    induction_model_map.get(induction,0),
    fuel_map.get(fuel,0),
    tune_base_map.get(tune,0),
    altitude
]])

input_scaled = scaler.transform(input_for_model)
pred_base = float(model.predict(input_scaled)[0])

cam_gain_map = {
    "Stock":0.0,
    "Stage 1 – Road":6.0,
    "Stage 2 – Aggressive":12.0,
    "Stage 3 – Race":20.0
}
headers_gain_map = {"Stock":0.0, "Shorty":5.0, "Long Tube":9.0}
intake_man_gain_map = {
    "Stock":0.0,
    "High-flow":5.0,
    "Individual throttle bodies":10.0
}
intercooler_gain_map = {
    "None":0.0,
    "Air-to-Air":4.0,
    "Air-to-Water":6.5,
    "Front-mount High-Flow":9.0
}
turbo_size_map = {
    "N/A":0.0,
    "Small 45-55mm":5.0,
    "Medium 56-65mm":12.0,
    "Big 66-75mm":20.0,
    "XL 76mm+":28.0
}

cam_gain = cam_gain_map.get(cam, 0.0)
headers_gain = headers_gain_map.get(headers, 0.0)
intake_man_gain = intake_man_gain_map.get(intake_manifold, 0.0)
intercooler_gain = intercooler_gain_map.get(intercooler, 0.0)
turbo_size_gain = turbo_size_map.get(turbo_size, 0.0)

if induction in ["Turbo", "Twin-Turbo"]:
    boost_gain = boost_psi * 1.8
elif induction in ["Supercharger", "Twincharged"]:
    boost_gain = boost_psi * 1.4
else:
    boost_gain = 0.0

meth_gain = 12.0 if meth else 0.0
exhaust_dia_gain = max(0.0, (exhaust_dia - 55) * 0.1)

extra_gain = (
    cam_gain +
    headers_gain +
    intake_man_gain +
    intercooler_gain +
    turbo_size_gain +
    boost_gain * 0.9 +
    meth_gain +
    exhaust_dia_gain
)

traction_multiplier = {"Daily":0.9, "Spirited":1.0, "Track / Drag":1.05}[traction_mode]

pred_total_gain = (pred_base + extra_gain) * traction_multiplier
pred_total_gain = max(pred_total_gain, -5.0)
new_hp = max(base_hp + pred_total_gain, 40.0)

effective_weight = weight_kg * (1 - weight_reduction / 100.0)
hp_per_ton = new_hp / (effective_weight / 1000.0)

hype_raw = np.clip(pred_total_gain / 120 * 100, 0, 100)
hype_level = int(hype_raw)

if hype_level < 20:
    verdict = "Sleeper grocery getter 🚙"
elif hype_level < 40:
    verdict = "Respectable street build 🌃"
elif hype_level < 70:
    verdict = "Serious weekend weapon ⚔️"
else:
    verdict = "Full send, tyres cry for mercy 🏁"

# compute base angle for rev gauge (-90deg to +90deg)
angle_deg = -90 + (hype_level / 100.0) * 180

# ========= RIGHT SIDE =========
with right:
    st.markdown('<div class="glass">', unsafe_allow_html=True)
    st.subheader("💥 Build Outcome")

    m1, m2, m3 = st.columns(3)
    m1.metric("Estimated HP Gain", f"{pred_total_gain:.1f} HP")
    m2.metric("New Output", f"{new_hp:.1f} HP")
    m3.metric("HP per Ton", f"{hp_per_ton:.1f}")

    # Animated rev limiter gauge
    st.markdown('<div class="meter-label">Rev Limiter Vibes</div>', unsafe_allow_html=True)
    st.markdown(
        f"""
        <div class="rev-wrap">
          <div class="rev-gauge">
            <div class="rev-arc"></div>
            <div class="rev-needle" style="--base-angle:{angle_deg}deg;"></div>
            <div class="rev-center">RPM</div>
            <div class="rev-label">REDLINE</div>
          </div>
        </div>
        """,
        unsafe_allow_html=True
    )

    st.markdown('<div class="meter-label">Hype Meter (relative craziness of this build)</div>', unsafe_allow_html=True)
    st.progress(hype_level)
    st.markdown(f"**Verdict:** {verdict}")

    st.markdown("")
    st.bar_chart(
        pd.DataFrame(
            {"Horsepower": [base_hp, new_hp]},
            index=["Stock", "Tuned"]
        )
    )

    st.markdown(
        f"""
        <div style="margin-top:6px;">
            <span class="chip">{engine_disp}L {engine_layout}</span>
            <span class="chip">Weight: {weight_kg} kg ▸ −{weight_reduction}%</span>
            <span class="chip">Intake: {intake}</span>
            <span class="chip">Headers: {headers}</span>
            <span class="chip">Exhaust: {exhaust} ({exhaust_dia}mm)</span><br/>
            <span class="chip">Induction: {induction} @ {boost_psi} psi</span>
            <span class="chip">IC: {intercooler}</span>
            <span class="chip">Tune: {tune}</span>
            <span class="chip">Fuel: {fuel}</span>
        </div>
        """,
        unsafe_allow_html=True
    )

    st.markdown("---")

    breakdown = pd.DataFrame({
        "Component": [
            "Model base (from dataset)",
            "Cam profile",
            "Headers",
            "Intake manifold",
            "Intercooler",
            "Turbo size",
            "Boost (net)",
            "Meth kit",
            "Exhaust diameter tweak"
        ],
        "Approx HP": [
            pred_base,
            cam_gain,
            headers_gain,
            intake_man_gain,
            intercooler_gain,
            turbo_size_gain,
            boost_gain * 0.9,
            meth_gain,
            exhaust_dia_gain
        ]
    })

    st.caption("🔬 Where is that extra power coming from?")
    st.dataframe(breakdown, use_container_width=True, height=260)

    st.markdown(
        '<p class="muted" style="margin-top:8px;">Model: distance-weighted KNN on synthetic dyno-style data + heuristic math for advanced mods.</p>',
        unsafe_allow_html=True
    )
    st.markdown('</div>', unsafe_allow_html=True)

# ========= FOOTER =========
st.markdown(
    """
    <div style="text-align:center;margin-top:10px;" class="muted">
        Presets + rev limiter + neon dyno… if this doesn’t make people go mad, we can push it even further. 🔧🔥
    </div>
    """,
    unsafe_allow_html=True
)