import gradio as gr
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import plotly.graph_objects as go

torch.manual_seed(42)
np.random.seed(42)

# Intercepts encode baseline preferences when no rivals are encountered.
# Luke pushes asparagus & milk; Rey pushes steak & milk.
B_true = np.array([
# intercepts: asp milk appl steak
[ 0.6, 0.3, 0.0, -0.2],
# Luke effect:
[ 0.8, 0.5, -0.2, -0.6],
# Rey effect:
[-0.6, 0.5, -0.3, 0.9],
])

FOODS = ['Asparagus', 'Milk', 'Apples', 'Steak']
COLORS = ['#00CC00', '#0066FF', '#FF0000', '#FF9900']

def softmax(z):
    z = np.atleast_2d(z)
    z = z - np.max(z, axis=1, keepdims=True)
    return np.exp(z) / np.sum(np.exp(z), axis=1, keepdims=True)

def generate_logs(n=100):
    """Simulate Kylo's ~100 weeks of logs."""
    luke = np.random.randint(0, 6, n)
    rey = np.random.randint(0, 6, n)
    X = np.column_stack([np.ones(n), luke, rey])
    probs = softmax(X @ B_true)
    choices = np.array([np.random.choice(4, p=p) for p in probs])
    return X, choices

class SoftmaxModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(3, 4, bias=False)

    def forward(self, x):
        return self.linear(x)

def train_model(X, y):
    """Fit multinomial logistic regression in PyTorch."""
    model = SoftmaxModel()
    opt = optim.Adam(model.parameters(), lr=0.5)
    loss_fn = nn.CrossEntropyLoss()
    X_t = torch.tensor(X, dtype=torch.float32)
    y_t = torch.tensor(y, dtype=torch.long)

    for _ in range(500):
        opt.zero_grad()
        loss_fn(model(X_t), y_t).backward()
        opt.step()

    return model.linear.weight.detach().numpy().T

# Train model once on startup
X_train, y_train = generate_logs(100)
B_learned = train_model(X_train, y_train)


def make_simplex(luke, rey):
    """Show predicted probabilities on tetrahedron for given encounters."""
    # Predict using learned model
    x = np.array([[1, luke, rey]])
    probs = softmax(x @ B_learned)[0]

    # Tetrahedron vertices
    v = np.array([
        [1, 0, -1/np.sqrt(2)],
        [-1, 0, -1/np.sqrt(2)],
        [0, 1, 1/np.sqrt(2)],
        [0, -1, 1/np.sqrt(2)],
    ])

    # Predicted point location
    pt = probs @ v

    fig = go.Figure()

    # Draw tetrahedron edges
    for i, j in [(0,1), (0,2), (0,3), (1,2), (1,3), (2,3)]:
        fig.add_trace(go.Scatter3d(
            x=[v[i,0], v[j,0]], y=[v[i,1], v[j,1]], z=[v[i,2], v[j,2]],
            mode='lines', line=dict(color='gray', width=2),
            showlegend=False, hoverinfo='skip'
        ))

    # Draw predicted point
    fig.add_trace(go.Scatter3d(
        x=[pt[0]], y=[pt[1]], z=[pt[2]],
        mode='markers', marker=dict(size=12, color='white', symbol='diamond'),
        name='Prediction'
    ))

    # Vertex labels with probabilities
    for i, (food, color) in enumerate(zip(FOODS, COLORS)):
        label = f"{food}\n({probs[i]:.1%})"
        fig.add_trace(go.Scatter3d(
            x=[v[i,0]*1.3], y=[v[i,1]*1.3], z=[v[i,2]*1.3],
            mode='text', text=[label], textfont=dict(size=24, color=color),
            showlegend=False, hoverinfo='skip'
        ))

    fig.update_layout(
        title=f"Luke: {luke}, Rey: {rey}",
        scene=dict(
            xaxis=dict(visible=False),
            yaxis=dict(visible=False),
            zaxis=dict(visible=False),
            bgcolor='#1a1a1a'
        ),
        paper_bgcolor='#1a1a1a',
        font=dict(color='white'),
        margin=dict(l=0, r=0, t=40, b=0),
        height=500
    )
    return fig

def predict(luke, rey):
    x = np.array([[1, luke, rey]])
    probs = softmax(x @ B_learned)[0]
    fig = make_simplex(int(luke), int(rey))

    text = f"**Predicted probabilities for {int(luke)} Luke, {int(rey)} Rey encounters:**\n\n"
    for food, p, color in zip(FOODS, probs, COLORS):
        text += f"- {food}: {p:.1%}\n"

    return fig, text

# ============================================================================
# GRADIO APP
# ============================================================================
with gr.Blocks(title="Kylo's Cravings") as demo:
    gr.Markdown("# Kylo Ren's Craving Predictor")
    gr.Markdown("Model trained on ~100 weeks of logs. Adjust sliders to predict cravings.")

    with gr.Row():
        luke_slider = gr.Slider(0, 10, value=0, step=1, label="Luke encounters this week")
        rey_slider = gr.Slider(0, 10, value=0, step=1, label="Rey encounters this week")

    plot = gr.Plot()
    output = gr.Markdown()

    luke_slider.change(predict, [luke_slider, rey_slider], [plot, output])
    rey_slider.change(predict, [luke_slider, rey_slider], [plot, output])

    demo.load(predict, [luke_slider, rey_slider], [plot, output])

    gr.Markdown(f"""---
**Learned B:**
```
{np.array2string(B_learned, precision=2)}
```
""")

demo.launch()