app.py

import streamlit as st
import torch
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd
import numpy as np
from PIL import Image

# =========================
# 1) LIBRERIA OUTPUT: possibili nuovi usi
#   Ognuno ha un "mini embedding" e un'immagine di esempio
# =========================
output_libreria = {
    "ArmadioAnta": {
        "embedding": [1.0, 0.0],   # fittizio
        "descrizione": "Riutilizzo come anta armadio",
        "img": "armadio.jpg"       # se ce l'hai, o lascialo vuoto
    },
    "DockRigenerato": {
        "embedding": [0.5, 0.5],
        "descrizione": "Base dock con componenti EoL",
        "img": "dock.jpg"
    },
    "OggettoVetrina": {
        "embedding": [0.0, 1.0],
        "descrizione": "Pannello estetico o vetrina",
        "img": "vetrina.jpg"
    }
}
# In un caso reale potresti mettere molte voci, con embeddings più ricchi
# e parametri più precisi

# =========================
# 2) MODELLO AI: un MLP
#    Input: [volume, area, lunghezza, spessore, usura, (materiale cod?),  + output_embedding(2)]
#    Output: punteggio compatibilità (float)
# =========================
class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(8, 16)
        self.fc2 = nn.Linear(16, 8)
        self.fc3 = nn.Linear(8, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = torch.sigmoid(self.fc3(x))  # punteggio tra 0 e 1
        return x

# =========================
# Salviamo il modello in session_state per poterlo usare da più sezioni
# =========================
if "model" not in st.session_state:
    st.session_state["model"] = SimpleMLP()
if "trained" not in st.session_state:
    st.session_state["trained"] = False

# Ottimizziamo i param in session
optimizer = torch.optim.Adam(st.session_state["model"].parameters(), lr=1e-3)

# ============== FUNZIONE DI TRAIN FAKE ================
def train_fake_model(epochs=20, batch_size=32):
    """
    Genera un dataset fittizio di EoL + output => score
    e allena un MLP in session_state["model"].
    """
    # Generiamo un dataset random con 1000 esempi
    # EoL -> (vol, area, lunghezza, spessore, usura, materialeCod)
    # Output -> 2D embedding
    # Score -> random 0..1 ma correlato un po' a "logica" fittizia
    n = 1000
    rng = np.random.default_rng(42)
    vol = rng.normal(50,15,n)
    area = rng.normal(20,5,n)
    lung = rng.normal(100,30,n)
    spess = rng.uniform(0.5,5,n)
    usura = rng.uniform(0,1,n)
    materiale = rng.integers(0,3,n).astype(float)

    # Scegli output a caso
    out_keys = list(output_libreria.keys())
    output_idx = rng.integers(0,len(out_keys),n)
    # costruiamo embedding 2D
    out_emb = []
    for i in output_idx:
        key = out_keys[i]
        out_emb.append(output_libreria[key]["embedding"])
    out_emb = np.array(out_emb)

    # Score fittizio
    # Esempio: se volume < 200 e usura < 0.5 => punteggio + 
    # + un offset se output = "ArmadioAnta" e spess < 3 => punteggio alto
    base_score = 0.5 * np.ones(n)
    base_score -= (usura * 0.3)
    base_score -= ((vol - 200)/400)
    # se spess < 3 e out=Armadio => +0.2
    # se out=DockRigenerato => +0.1 random
    for i in range(n):
        if out_keys[output_idx[i]]=="ArmadioAnta" and spess[i]<3:
            base_score[i]+=0.2
        if out_keys[output_idx[i]]=="DockRigenerato":
            base_score[i]+=rng.uniform(0,0.1)
    # bounding punteggio
    base_score = np.clip(base_score,0,1)

    X = np.column_stack([vol,area,lung,spess,usura,materiale,out_emb])
    y = base_score.reshape(-1,1)

    # Torch
    X_t = torch.tensor(X, dtype=torch.float32)
    y_t = torch.tensor(y, dtype=torch.float32)

    dataset = torch.utils.data.TensorDataset(X_t, y_t)
    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

    model = st.session_state["model"]
    # reinit param
    def weights_init(m):
        if isinstance(m, nn.Linear):
            nn.init.xavier_uniform_(m.weight)
            nn.init.zeros_(m.bias)
    model.apply(weights_init)

    lossf = nn.MSELoss()
    for ep in range(epochs):
        totloss=0
        for xb,yb in loader:
            optimizer.zero_grad()
            pred = model(xb)
            loss=lossf(pred,yb)
            loss.backward()
            optimizer.step()
            totloss+=loss.item()
    st.session_state["trained"] = True
    return totloss

# ============== UI Streamlit ================
st.title("WEEKO – Flusso: EoL + Output-libreria + AI + Overlay Estetico")

# TAB: Train AI + Avanzamento
tab1, tab2, tab3 = st.tabs(["1) Train AI (finto)","2) Valuta EoL + Output","3) Overlay Estetico"])

with tab1:
    st.subheader("Fase 1: Train fittizio del modello AI")
    st.write("Generiamo un dataset casuale simulando molte componenti EoL e possibili output.")
    if st.button("🏋️ Esegui training fittizio"):
        final_loss = train_fake_model(epochs=20)
        st.success(f"Training completato, loss={final_loss:.4f}")
    if st.session_state["trained"]:
        st.info("Modello AI pronto all'uso.")
    else:
        st.warning("Non hai ancora addestrato il modello fittizio, i punteggi saranno casuali.")

with tab2:
    st.subheader("Fase 2: Inserisci EoL e scegli un Output")

    # Carico eventuale foto EoL
    file_eol = st.file_uploader("Foto EoL (opzionale)", type=["jpg","png"], key="eolFile2")
    if file_eol:
        st.image(file_eol, caption="Componente EoL", use_column_width=True)

    # Dati EoL
    vol_in = st.number_input("Volume (cm³)", 0.0, 1000.0, 100.0)
    area_in = st.number_input("Area (cm²)", 0.0, 500.0, 30.0)
    lung_in = st.number_input("Lunghezza (mm)", 0.0, 2000.0, 120.0)
    spess_in = st.number_input("Spessore (mm)", 0.5, 10.0, 2.0)
    usura_in = st.slider("Usura (0=nuovo,1=molto usurato)", 0.0,1.0,0.3)
    mat_in = st.selectbox("Materiale", ["Plastica","Metallo","PCB"])
    mat_code = {"Plastica":0.0,"Metallo":1.0,"PCB":2.0}[mat_in]

    # Scegli un output
    scelte = list(output_libreria.keys())
    out_sel = st.selectbox("Scegli l'output-libreria", scelte)

    if st.button("Calcola punteggio AI"):
        model = st.session_state["model"]
        embed = output_libreria[out_sel]["embedding"]  # 2D
        x_np = np.array([[vol_in, area_in, lung_in, spess_in, usura_in, mat_code, embed[0], embed[1]]], dtype=np.float32)
        x_t = torch.from_numpy(x_np)
        with torch.no_grad():
            pred = model(x_t).item()
        st.info(f"Punteggio di compatibilità con '{out_sel}': {pred:.3f}")
        if pred>0.7:
            st.success("Ottima compatibilità! Procedi con la rigenerazione.")
        elif pred>0.4:
            st.warning("Compatibilità media: valuta modifiche/manuale rifinitura.")
        else:
            st.error("Compatibilità bassa, meglio cambiare output o scartare.")

with tab3:
    st.subheader("Fase 3: Overlay Estetico")
    st.write("Carica la foto EoL e la foto dell'Output finale per un effetto wow. (Puramente estetico)")

    col1, col2 = st.columns(2)
    with col1:
        file_eol_3 = st.file_uploader("Foto EoL", type=["jpg","png"], key="uplEol3")
    with col2:
        file_out_3 = st.file_uploader("Foto Oggetto Finale (es. Armadio?)", type=["jpg","png"], key="uplOut3")

    alpha_val = st.slider("Trasparenza EoL", 0.0,1.0,0.5)

    if file_eol_3 and file_out_3:
        eol_img = Image.open(file_eol_3).convert("RGBA")
        out_img = Image.open(file_out_3).convert("RGBA")
        # ridimensioniamo la EoL all'immagine dell'oggetto
        eol_img = eol_img.resize(out_img.size)

        blend = Image.blend(out_img, eol_img, alpha_val)
        st.image(blend, caption="Overlay EoL + Oggetto Finale", use_column_width=True)
    else:
        st.info("Attendi di caricare entrambe le immagini per vedere l'overlay.")