app.py

import os
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import gradio as gr
import requests
from einops import rearrange
from transformers import pipeline

# ============================================================
# CONFIG
# ============================================================
SEQ_LEN = 90
N_FEATURES = 5
PATCH_LEN = 10
D_MODEL = 128
N_HEADS = 4
N_LAYERS = 4
N_CLASSES = 2
DEVICE = "cpu"

MODEL_PATH = "epoch_20 (1).pth"
ALPHA_KEY = os.getenv("vishal")

# ============================================================
# MODEL
# ============================================================
class PatchEmbedding(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(PATCH_LEN * N_FEATURES, D_MODEL)

    def forward(self, x):
        x = rearrange(x, 'b (n p) f -> b n (p f)', p=PATCH_LEN)
        return self.proj(x)

class PatchTST(nn.Module):
    def __init__(self):
        super().__init__()
        self.patch_embed = PatchEmbedding()
        encoder = nn.TransformerEncoderLayer(
            d_model=D_MODEL,
            nhead=N_HEADS,
            batch_first=True
        )
        self.encoder = nn.TransformerEncoder(encoder, num_layers=N_LAYERS)
        self.head = nn.Linear(D_MODEL, N_CLASSES)

    def forward(self, x):
        x = self.patch_embed(x)
        x = self.encoder(x)
        x = x.mean(dim=1)
        return self.head(x)

model = PatchTST()
model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model.eval()

# Small lightweight LLM (safe for HF)
generator = pipeline(
    "text-generation",
    model="google/flan-t5-small",
    device=-1
)

# ============================================================
# ALPHA SYMBOL SEARCH
# ============================================================
def resolve_symbol(query):
    print("Resolving symbol for:", query)

    url = (
        f"https://www.alphavantage.co/query?"
        f"function=SYMBOL_SEARCH&keywords={query}&apikey={ALPHA_KEY}"
    )

    response = requests.get(url)
    data = response.json()

    print("========== SYMBOL SEARCH RAW ==========")
    print(data)
    print("=======================================")

    if "bestMatches" not in data or len(data["bestMatches"]) == 0:
        return None, "No matching symbol found."

    best = data["bestMatches"][0]
    return best["1. symbol"], None

# ============================================================
# FETCH DAILY DATA
# ============================================================
def fetch_daily(symbol):
    print("Fetching daily data for:", symbol)

    url = (
        f"https://www.alphavantage.co/query?"
        f"function=TIME_SERIES_DAILY_ADJUSTED&symbol={symbol}"
        f"&outputsize=compact&apikey={ALPHA_KEY}"
    )

    response = requests.get(url)
    data = response.json()

    print("========== DAILY DATA RAW ==========")
    print(data)
    print("====================================")

    if "Note" in data:
        return None, "Alpha rate limit reached."

    if "Error Message" in data:
        return None, "Invalid symbol or API error."

    if "Time Series (Daily)" not in data:
        return None, "Unexpected Alpha response."

    ts = data["Time Series (Daily)"]
    df = pd.DataFrame.from_dict(ts, orient="index").astype(float)
    df = df.rename(columns={
        "1. open": "Open",
        "2. high": "High",
        "3. low": "Low",
        "4. close": "Close",
        "6. volume": "Volume"
    })

    df = df.sort_index()
    return df[["Open", "High", "Low", "Close", "Volume"]], None

# ============================================================
# PREPROCESS
# ============================================================
def preprocess(df):
    df = df.tail(SEQ_LEN)

    if len(df) < SEQ_LEN:
        return None

    features = df.values
    features = (features - features.mean(axis=0)) / (features.std(axis=0) + 1e-6)

    return torch.tensor(features, dtype=torch.float32).unsqueeze(0)

# ============================================================
# MAIN PIPELINE
# ============================================================
def predict(asset_name):

    if not ALPHA_KEY:
        return 0, "Error", "API key missing."

    symbol, err = resolve_symbol(asset_name)
    if err:
        return 0, "Error", err

    df, err = fetch_daily(symbol)
    if err:
        return 0, "Error", err

    x = preprocess(df)
    if x is None:
        return 0, "Error", "Not enough data."

    with torch.no_grad():
        logits = model(x)
        probs = torch.softmax(logits, dim=1).numpy()[0]

    label = "HIGH Volatility" if np.argmax(probs) == 1 else "LOW Volatility"
    confidence = float(np.max(probs))

    prompt = (
        f"Explain in simple storytelling style: "
        f"The stock {symbol} is currently in {label} "
        f"with confidence {confidence:.2f}."
    )

    explanation = generator(prompt, max_new_tokens=80)[0]["generated_text"]

    return confidence, label, explanation

# ============================================================
# UI
# ============================================================
interface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Enter stock name (e.g. Apple, AAPL, SBI)"),
    outputs=[
        gr.Number(label="Model Confidence"),
        gr.Text(label="Volatility Regime"),
        gr.Textbox(label="Explanation")
    ],
    title="Stock Volatility Regime Detector (Alpha Debug Mode)",
    description="Debug version with full Alpha API logging."
)

interface.launch()