import torch
import torch.nn as nn
import gradio as gr

# ---------------------------------------------------------
# Industry-style MLP for Customer Churn Risk
# ---------------------------------------------------------
class ChurnRiskMLP(nn.Module):
    def __init__(self, input_size=20, hidden_size=64, output_size=1, dropout_p=0.5):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=dropout_p)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x1 = self.fc1(x)
        x2 = self.relu(x1)
        x3 = self.dropout(x2)
        x4 = self.fc2(x3)
        out = self.sigmoid(x4)
        return x1, x2, x3, out


# ---------------------------------------------------------
# Inference function for Gradio
# ---------------------------------------------------------
def run_inference(batch_size, hidden_size, dropout_p, mode):
    # Create dummy customer data
    dummy = torch.randn(batch_size, 20)

    # Build model
    model = ChurnRiskMLP(
        input_size=20,
        hidden_size=hidden_size,
        output_size=1,
        dropout_p=dropout_p
    )

    # Set mode
    if mode == "train (dropout ON)":
        model.train()
    else:
        model.eval()

    # Forward pass
    with torch.no_grad():
        fc1, relu, drop, out = model(dummy)

    # Convert tensors to readable text
    result = (
        f"=== Dummy Customer Input (batch={batch_size}) ===\n{dummy}\n\n"
        f"=== After fc1 ===\n{fc1}\n\n"
        f"=== After ReLU ===\n{relu}\n\n"
        f"=== After Dropout ({mode}) ===\n{drop}\n\n"
        f"=== Final Churn Risk Predictions ===\n{out}\n"
    )

    return result


# ---------------------------------------------------------
# Gradio Interface
# ---------------------------------------------------------
with gr.Blocks(title="Customer Churn Risk Explorer") as demo:
    gr.Markdown(
        """
        # 📊 Customer Churn Risk Explorer  
        Investigate how a real industry-style MLP behaves with **dropout**,  
        adjustable **hidden size**, and **batch size**.

        This tool helps you visualize:
        - Layer activations  
        - Dropout effects  
        - Final churn-risk predictions  
        """
    )

    batch = gr.Slider(1, 32, value=8, step=1, label="Batch Size (# of customers)")
    hidden = gr.Slider(8, 256, value=64, step=8, label="Hidden Layer Size")
    dropout = gr.Slider(0.0, 0.9, value=0.5, step=0.1, label="Dropout Probability")
    mode = gr.Radio(["train (dropout ON)", "eval (dropout OFF)"], value="train (dropout ON)", label="Mode")

    output_box = gr.Textbox(label="Model Output", lines=25)

    run_button = gr.Button("Run Model")

    run_button.click(
        fn=run_inference,
        inputs=[batch, hidden, dropout, mode],
        outputs=output_box
    )

demo.launch()