app.py

import asyncio
import json
import logging
import time
import aiohttp
import pandas as pd
import numpy as np
from aiohttp import web
import websockets
from sklearn.ensemble import RandomForestRegressor

# Configuration
SYMBOL_KRAKEN = "BTC/USD"
PORT = 7860
BROADCAST_RATE = 1.0 
PREDICTION_HORIZON = 100 

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

market_state = {
    "ohlc_history": [],
    "ready": False,
    "model": None,
    "last_training_time": 0
}

connected_clients = set()

# --- Indicator Logic ---
def calculate_indicators(candles):
    if len(candles) < 50: return None
    
    df = pd.DataFrame(candles)
    cols = ['open', 'high', 'low', 'close', 'volume']
    for c in cols: df[c] = df[c].astype(float)

    # EMA 20
    df['ema'] = df['close'].ewm(span=20, adjust=False).mean()

    # Bollinger Bands
    df['sma20'] = df['close'].rolling(window=20).mean()
    df['std'] = df['close'].rolling(window=20).std()
    df['bb_upper'] = df['sma20'] + (df['std'] * 2)
    df['bb_lower'] = df['sma20'] - (df['std'] * 2)

    # RSI
    delta = df['close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    rs = gain / loss
    df['rsi'] = 100 - (100 / (1 + rs))

    # MACD
    k = df['close'].ewm(span=12, adjust=False).mean()
    d = df['close'].ewm(span=26, adjust=False).mean()
    df['macd'] = k - d
    df['macd_signal'] = df['macd'].ewm(span=9, adjust=False).mean()
    df['macd_hist'] = df['macd'] - df['macd_signal']

    # Stochastic
    low_min = df['low'].rolling(window=14).min()
    high_max = df['high'].rolling(window=14).max()
    df['stoch_k'] = 100 * ((df['close'] - low_min) / (high_max - low_min))

    # ATR
    df['tr0'] = abs(df['high'] - df['low'])
    df['tr1'] = abs(df['high'] - df['close'].shift())
    df['tr2'] = abs(df['low'] - df['close'].shift())
    df['tr'] = df[['tr0', 'tr1', 'tr2']].max(axis=1)
    df['atr'] = df['tr'].rolling(window=14).mean()

    # OBV
    df['obv'] = (np.sign(df['close'].diff()) * df['volume']).fillna(0).cumsum()

    # VWAP
    df['tp'] = (df['high'] + df['low'] + df['close']) / 3
    df['vwap'] = (df['tp'] * df['volume']).cumsum() / df['volume'].cumsum()
    
    return df

# --- Machine Learning Logic ---
def train_model(df):
    logging.info("Training ML Model...")
    
    # 1. Prepare Features (X)
    feature_cols = ['close', 'ema', 'bb_upper', 'bb_lower', 'rsi', 'macd', 'stoch_k', 'atr', 'obv', 'vwap']
    
    # Create a clean copy for training data
    data = df.dropna().copy()
    
    # 2. Create Targets (y) - OPTIMIZED to fix Fragmentation Warning
    # Instead of adding columns in a loop, we create a dict and concat once
    future_shifts = {}
    targets = []
    
    for i in range(1, PREDICTION_HORIZON + 1):
        col_name = f'target_{i}'
        future_shifts[col_name] = data['close'].shift(-i)
        targets.append(col_name)
        
    # Concatenate all target columns at once
    target_df = pd.DataFrame(future_shifts, index=data.index)
    data = pd.concat([data, target_df], axis=1)

    # Drop rows where we don't have future data (the last 100 rows)
    data = data.dropna()

    if len(data) < 100:
        logging.warning("Not enough data to train model yet.")
        return None

    X = data[feature_cols].values
    y = data[targets].values

    # Train Random Forest
    model = RandomForestRegressor(n_estimators=50, max_depth=10, n_jobs=-1, random_state=42)
    model.fit(X, y)
    
    logging.info(f"Model Trained on {len(X)} samples.")
    return model

def get_prediction(df, model):
    if model is None: return []
    
    # Get the very last row of data (current market state)
    feature_cols = ['close', 'ema', 'bb_upper', 'bb_lower', 'rsi', 'macd', 'stoch_k', 'atr', 'obv', 'vwap']
    last_row = df.iloc[[-1]][feature_cols]
    
    # Check for NaNs
    if last_row.isnull().values.any(): return []

    # Predict
    prediction = model.predict(last_row.values)[0] 
    
    # Format for frontend
    current_time = int(df.iloc[-1]['time'])
    pred_data = []
    for i, price in enumerate(prediction):
        pred_data.append({
            "time": current_time + ((i + 1) * 60), 
            "value": float(price)
        })
    
    return pred_data

def process_market_data():
    if not market_state['ready'] or not market_state['ohlc_history']: 
        return {"error": "Initializing..."}

    # 1. Calculate DataFrame
    df = calculate_indicators(market_state['ohlc_history'])
    if df is None or len(df) < 50: return {"error": "Not enough data"}

    # 2. Train Model (Periodically)
    # Train initially or every 15 minutes (900 seconds)
    if market_state['model'] is None or (time.time() - market_state['last_training_time'] > 900):
        market_state['model'] = train_model(df)
        market_state['last_training_time'] = time.time()

    # 3. Get Prediction
    predictions = get_prediction(df, market_state['model'])

    # 4. Prepare JSON
    full_data = df.where(pd.notnull(df), None).to_dict('records')
    
    return {
        "data": full_data,
        "prediction": predictions
    }

# --- Frontend ---
HTML_PAGE = f"""
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>{SYMBOL_KRAKEN} AI Predictor</title>
    <script src="https://unpkg.com/lightweight-charts@4.1.1/dist/lightweight-charts.standalone.production.js"></script>
    <style>
        body {{ margin: 0; background: #000; color: #fff; font-family: 'Segoe UI', sans-serif; height: 100vh; display: flex; flex-direction: column; overflow: hidden; }}
        .header {{ height: 32px; background: #0a0a0a; border-bottom: 1px solid #333; display: flex; align-items: center; padding: 0 12px; font-size: 13px; font-weight: 600; justify-content: space-between; }}
        #charts-container {{ flex: 1; display: flex; flex-direction: column; }}
        .chart-row {{ width: 100%; position: relative; border-bottom: 1px solid #222; }}
        #main-chart {{ flex: 4; }}
        #osc-chart {{ flex: 1; min-height: 100px; }}
        .legend {{ position: absolute; top: 8px; left: 10px; z-index: 10; font-size: 11px; color: #aaa; pointer-events: none; text-shadow: 1px 1px 2px #000; }}
        .l-item {{ margin-right: 12px; }}
    </style>
</head>
<body>
    <div class="header">
        <span style="color:#00e676">{SYMBOL_KRAKEN} + Random Forest (Next 100 Candles)</span>
        <span id="clock" style="color:#888">Initializing...</span>
    </div>

    <div id="charts-container">
        <div id="main-chart" class="chart-row">
            <div class="legend">
                <span class="l-item" style="color:#00ff9d">Price</span>
                <span class="l-item" style="color:#bf5af2">AI Prediction</span>
                <span class="l-item" style="color:#2962FF">EMA</span>
            </div>
        </div>
        <div id="osc-chart" class="chart-row">
            <div class="legend">
                <span class="l-item" style="color:#9C27B0">RSI</span>
                <span class="l-item" style="color:#00BCD4">MACD</span>
            </div>
        </div>
    </div>

<script>
    document.addEventListener('DOMContentLoaded', () => {{
        const mainEl = document.getElementById('main-chart');
        const oscEl = document.getElementById('osc-chart');

        const commonOpts = {{
            layout: {{ background: {{ type: 'solid', color: '#000' }}, textColor: '#888' }},
            grid: {{ vertLines: {{ color: '#111' }}, horzLines: {{ color: '#111' }} }},
            timeScale: {{ timeVisible: true, secondsVisible: false, borderColor: '#333' }},
            rightPriceScale: {{ borderColor: '#333' }},
            crosshair: {{ mode: 1 }}
        }};

        const mainChart = LightweightCharts.createChart(mainEl, commonOpts);
        const candles = mainChart.addCandlestickSeries({{ upColor: '#00ff9d', downColor: '#ff3b3b', borderVisible: false }});
        const ema = mainChart.addLineSeries({{ color: '#2962FF', lineWidth: 1 }});
        const predLine = mainChart.addLineSeries({{ color: '#bf5af2', lineWidth: 2, lineStyle: 2, title: 'AI Forecast' }});

        const oscChart = LightweightCharts.createChart(oscEl, commonOpts);
        const rsi = oscChart.addLineSeries({{ color: '#9C27B0', lineWidth: 1 }});
        const macdHist = oscChart.addHistogramSeries({{ priceScaleId: 'macd', color: '#2962FF' }});
        oscChart.priceScale('macd').applyOptions({{ scaleMargins: {{ top: 0.8, bottom: 0 }} }});

        new ResizeObserver(entries => {{
             for (let e of entries) {{
                 if(e.target === mainEl) mainChart.applyOptions({{ width: e.contentRect.width, height: e.contentRect.height }});
                 if(e.target === oscEl) oscChart.applyOptions({{ width: e.contentRect.width, height: e.contentRect.height }});
             }}
        }}).observe(document.body);

        function syncCharts(source, targets) {{
            source.timeScale().subscribeVisibleLogicalRangeChange(range => {{
                targets.forEach(t => t.timeScale().setVisibleLogicalRange(range));
            }});
        }}
        syncCharts(mainChart, [oscChart]);
        syncCharts(oscChart, [mainChart]);

        function connect() {{
            const ws = new WebSocket((location.protocol === 'https:' ? 'wss' : 'ws') + '://' + location.host + '/ws');
            ws.onmessage = (e) => {{
                const payload = JSON.parse(e.data);
                if (!payload.data) return;
                
                const d = payload.data;
                const mapData = (key) => d.map(x => ({{ time: x.time, value: x[key] }})).filter(x => x.value !== null);

                candles.setData(d.map(x => ({{ time: x.time, open: x.open, high: x.high, low: x.low, close: x.close }})));
                ema.setData(mapData('ema'));
                rsi.setData(mapData('rsi'));
                
                if(payload.prediction && payload.prediction.length > 0) {{
                    predLine.setData(payload.prediction);
                }}

                macdHist.setData(d.map(x => ({{ 
                    time: x.time, 
                    value: x.macd_hist || 0, 
                    color: (x.macd_hist||0) >= 0 ? '#26a69a' : '#ef5350' 
                }})));
                
                document.getElementById('clock').innerText = new Date().toISOString().split('T')[1].split('.')[0] + ' UTC';
            }};
            ws.onclose = () => setTimeout(connect, 2000);
        }}
        connect();
    }});
</script>
</body>
</html>
"""

async def kraken_worker():
    global market_state
    try:
        async with aiohttp.ClientSession() as session:
            url = "https://api.kraken.com/0/public/OHLC?pair=XBTUSD&interval=1" 
            async with session.get(url) as response:
                if response.status == 200:
                    data = await response.json()
                    if 'result' in data:
                        for key in data['result']:
                            if key != 'last':
                                raw = data['result'][key]
                                market_state['ohlc_history'] = [
                                    {
                                        'time': int(c[0]), 
                                        'open': float(c[1]), 
                                        'high': float(c[2]), 
                                        'low': float(c[3]), 
                                        'close': float(c[4]), 
                                        'volume': float(c[6])
                                    }
                                    for c in raw[-720:] 
                                ]
                                market_state['ready'] = True
                                break
    except Exception as e:
        logging.error(f"Init Error: {e}")

    # WebSocket Stream
    while True:
        try:
            async with websockets.connect("wss://ws.kraken.com/v2") as ws:
                logging.info("WebSocket Connected")
                await ws.send(json.dumps({"method": "subscribe", "params": {"channel": "trade", "symbol": [SYMBOL_KRAKEN]}}))
                await ws.send(json.dumps({"method": "subscribe", "params": {"channel": "ohlc", "symbol": [SYMBOL_KRAKEN], "interval": 1}}))

                async for message in ws:
                    payload = json.loads(message)
                    channel = payload.get("channel")
                    data = payload.get("data", [])

                    if channel == "trade":
                        for trade in data:
                            try:
                                price = float(trade['price'])
                                vol = float(trade['qty'])
                                current_minute = int(time.time()) // 60 * 60
                                
                                if market_state['ohlc_history']:
                                    last = market_state['ohlc_history'][-1]
                                    if last['time'] == current_minute:
                                        last['close'] = price
                                        last['volume'] += vol
                                        if price > last['high']: last['high'] = price
                                        if price < last['low']: last['low'] = price
                                    elif current_minute > last['time']:
                                        market_state['ohlc_history'].append({
                                            'time': current_minute, 
                                            'open': price, 
                                            'high': price, 
                                            'low': price, 
                                            'close': price, 
                                            'volume': vol
                                        })
                                        if len(market_state['ohlc_history']) > 800: 
                                            market_state['ohlc_history'].pop(0)
                            except: pass
                    
                    elif channel == "ohlc":
                        for c in data:
                            try:
                                t = int(float(c['endtime'])) - 60
                                c_data = {
                                    'time': t, 
                                    'open': float(c['open']), 
                                    'high': float(c['high']), 
                                    'low': float(c['low']), 
                                    'close': float(c['close']), 
                                    'volume': float(c['volume'])
                                }
                                if market_state['ohlc_history']:
                                    if market_state['ohlc_history'][-1]['time'] == t: 
                                        market_state['ohlc_history'][-1] = c_data
                                    elif market_state['ohlc_history'][-1]['time'] < t: 
                                        market_state['ohlc_history'].append(c_data)
                                        if len(market_state['ohlc_history']) > 800: 
                                            market_state['ohlc_history'].pop(0)
                            except: pass
        except Exception as e:
            logging.warning(f"Reconnecting: {e}")
            await asyncio.sleep(2)

async def broadcast_worker():
    while True:
        if connected_clients and market_state['ready']:
            payload = process_market_data()
            if payload and "data" in payload:
                msg = json.dumps(payload)
                for ws in list(connected_clients):
                    try: await ws.send_str(msg)
                    except: pass
        await asyncio.sleep(BROADCAST_RATE)

async def websocket_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    connected_clients.add(ws)
    try:
        async for msg in ws: pass
    finally: connected_clients.remove(ws)
    return ws

async def handle_index(request):
    return web.Response(text=HTML_PAGE, content_type='text/html')

async def main():
    app = web.Application()
    app.router.add_get('/', handle_index)
    app.router.add_get('/ws', websocket_handler)
    asyncio.create_task(kraken_worker())
    asyncio.create_task(broadcast_worker())
    runner = web.AppRunner(app)
    await runner.setup()
    await web.TCPSite(runner, '0.0.0.0', PORT).start()
    print(f"🚀 AI Quant: http://localhost:{PORT}")
    await asyncio.Event().wait()

if __name__ == "__main__":
    try: asyncio.run(main())
    except KeyboardInterrupt: pass