import gradio as gr
import random
import plotly.graph_objects as go

# We use a simple dict for state to keep Hugging Face happy
init_state = {
    "balance": 10000.0,
    "shares": 0,
    "avg_buy_price": 0.0,
    "stock_price": 100.0,
    "price_history": [100.0],
    "current_mode": "Sideways ➡️",
    "tick_counter": 0
}

def tick(state):
    state["tick_counter"] += 1
    
    # Change market modes
    if state["tick_counter"] % 15 == 0:
        state["current_mode"] = random.choice(["Bull 📈", "Bear 📉", "Sideways ➡️"])
    
    # Percentage calculation
    base_move = random.uniform(-0.03, 0.03)
    if state["current_mode"] == "Bull 📈":
        base_move += random.uniform(0.005, 0.02)
    elif state["current_mode"] == "Bear 📉":
        base_move -= random.uniform(0.005, 0.02)
        
    if random.random() < 0.02: # Black swan
        base_move += random.choice([-0.08, 0.08])
        
    state["stock_price"] = max(1.0, round(state["stock_price"] * (1 + base_move), 2))
    state["price_history"].append(state["stock_price"])
    if len(state["price_history"]) > 30:
        state["price_history"].pop(0)
        
    return state, *get_ui_updates(state)

def buy_stock(state):
    if state["balance"] >= state["stock_price"]:
        shares_to_buy = min(10, int(state["balance"] // state["stock_price"]))
        if shares_to_buy > 0:
            total_cost = shares_to_buy * state["stock_price"]
            total_shares = state["shares"] + shares_to_buy
            state["avg_buy_price"] = ((state["shares"] * state["avg_buy_price"]) + total_cost) / total_shares
            state["shares"] = total_shares
            state["balance"] -= total_cost
    return state, *get_ui_updates(state)

def sell_stock(state):
    if state["shares"] > 0:
        shares_to_sell = min(10, state["shares"])
        state["balance"] += shares_to_sell * state["stock_price"]
        state["shares"] -= shares_to_sell
        if state["shares"] == 0:
            state["avg_buy_price"] = 0.0
    return state, *get_ui_updates(state)

def get_ui_updates(state):
    pnl = (state["shares"] * state["stock_price"]) - (state["shares"] * state["avg_buy_price"])
    pnl_str = f"+${pnl:,.2f}" if pnl >= 0 else f"-${abs(pnl):,.2f}"
    net_worth = state["balance"] + (state["shares"] * state["stock_price"])
    
    fig = go.Figure()
    fig.add_trace(go.Scatter(y=state["price_history"], mode='lines+markers', line=dict(color='#2196F3', width=3)))
    fig.update_layout(template="plotly_dark", margin=dict(l=10, r=10, t=10, b=10), height=300)
    
    return (
        f"${state['balance']:,.2f}",
        f"${state['stock_price']:,.2f}",
        str(state["shares"]),
        f"${state['avg_buy_price']:,.2f}",
        pnl_str,
        f"${net_worth:,.2f}",
        f"Market: {state['current_mode']}",
        fig
    )

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📈 StockSim")
    state = gr.State(default=init_state.copy())
    
    with gr.Row():
        with gr.Column(scale=2):
            chart = gr.Plot()
            regime_display = gr.Markdown("Market: Sideways ➡️")
        with gr.Column(scale=1):
            net_worth_box = gr.Textbox(label="Net Worth")
            balance_box = gr.Textbox(label="Cash")
            shares_box = gr.Textbox(label="Shares")
            avg_box = gr.Textbox(label="Avg Cost")
            pnl_box = gr.Textbox(label="P&L")

    with gr.Row():
        btn_buy = gr.Button("BUY 10", variant="primary")
        btn_sell = gr.Button("SELL 10", variant="secondary")

    outputs = [balance_box, gr.Textbox(visible=False), shares_box, avg_box, pnl_box, net_worth_box, regime_display, chart]
    
    btn_buy.click(buy_stock, inputs=[state], outputs=[state] + outputs)
    btn_sell.click(sell_stock, inputs=[state], outputs=[state] + outputs)
    
    # Timer updates state explicitly
    timer = gr.Timer(1.0)
    timer.tick(tick, inputs=[state], outputs=[state] + outputs)

demo.launch()