main.py

from fastapi import FastAPI, Request, Form
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import yfinance as yf
import pandas as pd
import numpy as np
import json
import plotly.utils
from datetime import datetime
import contextlib

# ==============================================================================
# CONFIGURATION
# ==============================================================================

app = FastAPI(title="ECM Quant AI")

# Mount static files and templates
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

# Hardcoded sector proxies
SECTOR_PROXIES = {
    'SaaS': ['CRM', 'SNOW', 'HUBS', 'NET', 'DDOG'],
    'Fintech': ['SQ', 'PYPL', 'UPST', 'AFRM', 'SOFI'],
    'Biotech': ['XBI', 'IBB', 'MRNA', 'VRTX', 'REGN'],
    'AI_Infra': ['NVDA', 'AMD', 'AVGO', 'MSFT', 'GOOGL']
}

BENCHMARK = '^GSPC'

# ==============================================================================
# FINANCIAL LOGIC
# ==============================================================================

def get_session():
    import requests
    session = requests.Session()
    session.headers.update({
        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
    })
    return session

def fetch_macro_data():
    """Fetches VIX and 10Y Treasury Yield"""
    try:
        # ^VIX: Volatility, ^TNX: 10Y Yield
        macro = yf.download(['^VIX', '^TNX'], period="5d", progress=False)
        if macro.empty: return {'vix': 15.0, 'tnx': 4.0} # Fallback
        
        # Handle MultiIndex columns if present
        if isinstance(macro.columns, pd.MultiIndex):
            vix = macro['Adj Close']['^VIX'].iloc[-1] if '^VIX' in macro['Adj Close'] else 15.0
            tnx = macro['Adj Close']['^TNX'].iloc[-1] if '^TNX' in macro['Adj Close'] else 4.0
        else:
            # Fallback for flat structure/failures
            vix = macro['Adj Close'].iloc[-1] if 'Adj Close' in macro else 15.0
            tnx = 4.0
            
        return {'vix': float(vix), 'tnx': float(tnx)}
    except:
        return {'vix': 18.5, 'tnx': 4.2} # Safe defaults

def fetch_market_data(tickers):
    try:
        all_tickers = tickers + [BENCHMARK]
        # Bulk download
        data = yf.download(all_tickers, period="6mo", progress=False)
        
        if data.empty: raise ValueError("Empty data returned")

        # Handle yfinance 0.2+ MultiIndex columns
        if isinstance(data.columns, pd.MultiIndex):
             prices = data['Adj Close']
        elif 'Adj Close' in data.columns:
            prices = data['Adj Close']
        elif 'Close' in data.columns:
            prices = data['Close']
        else:
            prices = data
        return prices
    except Exception as e:
        print(f"Error fetching data: {e}")
        # Panic Fallback
        dates = pd.date_range(end=datetime.today(), periods=120)
        dummy_data = {}
        for t in tickers + [BENCHMARK]:
            start = 150 if t == BENCHMARK else 100
            returns = np.random.normal(0.001, 0.02, 120)
            dummy_data[t] = start * (1 + returns).cumprod()
        return pd.DataFrame(dummy_data, index=dates)

def get_fundamentals(tickers):
    metrics = []
    
    # ADVANCED BACKUP DATA (Rule of 40, EV/Rev included)
    BACKUP_DATA = {
        'CRM':  {'marketCap': 280e9, 'forwardPE': 28.5, 'revenueGrowth': 0.11, 'margins': 0.18, 'evToRev': 6.5, 'beta': 1.05},
        'SNOW': {'marketCap': 55e9,  'forwardPE': 45.2, 'revenueGrowth': 0.22, 'margins': -0.05, 'evToRev': 12.0, 'beta': 1.45},
        'HUBS': {'marketCap': 32e9,  'forwardPE': 65.0, 'revenueGrowth': 0.18, 'margins': -0.02, 'evToRev': 9.5, 'beta': 1.25},
        'NET':  {'marketCap': 35e9,  'forwardPE': 85.5, 'revenueGrowth': 0.28, 'margins': -0.03, 'evToRev': 14.2, 'beta': 1.35},
        'DDOG': {'marketCap': 42e9,  'forwardPE': 55.8, 'revenueGrowth': 0.21, 'margins': 0.02, 'evToRev': 11.5, 'beta': 1.40},
        'SQ':   {'marketCap': 45e9,  'forwardPE': 25.0, 'revenueGrowth': 0.12, 'margins': 0.05, 'evToRev': 3.5, 'beta': 1.60},
        'PYPL': {'marketCap': 70e9,  'forwardPE': 14.5, 'revenueGrowth': 0.08, 'margins': 0.16, 'evToRev': 2.8, 'beta': 1.15},
        'NVDA': {'marketCap': 2500e9,'forwardPE': 35.0, 'revenueGrowth': 0.90, 'margins': 0.55, 'evToRev': 25.0, 'beta': 1.70},
    }

    for t in tickers:
        try:
            info = yf.Ticker(t).info
            
            def get_val(key, default=0.0):
                val = info.get(key)
                if val is None and t in BACKUP_DATA:
                    return BACKUP_DATA[t].get(key, default)
                return val if val is not None else default

            m_cap = get_val('marketCap')
            pe = get_val('forwardPE') or get_val('trailingPE')
            growth = get_val('revenueGrowth')
            margins = get_val('profitMargins')
            ev_rev = get_val('enterpriseToRevenue')
            beta = get_val('beta', 1.0)
            
            # Rule of 40: Growth % + Margin % (e.g., 0.20 + 0.10 => 30)
            rule_40 = (growth + margins) * 100

            metrics.append({
                'ticker': t,
                'market_cap': m_cap,
                'pe': pe,
                'ev_rev': ev_rev,
                'rule_40': rule_40,
                'growth': growth,
                'beta': beta
            })
        except Exception:
            # Fallback
            if t in BACKUP_DATA:
                bk = BACKUP_DATA[t]
                rule_40 = (bk['revenueGrowth'] + bk['margins']) * 100
                metrics.append({
                    'ticker': t,
                    'market_cap': bk['marketCap'],
                    'pe': bk['forwardPE'],
                    'ev_rev': bk['evToRev'],
                    'rule_40': rule_40,
                    'growth': bk['revenueGrowth'],
                    'beta': bk['beta']
                })
            else:
                metrics.append({'ticker': t, 'market_cap': 0, 'pe': 0, 'ev_rev': 0, 'rule_40': 0, 'growth': 0, 'beta': 1.0})
                
    return metrics

def calculate_signals(prices_df, sector_tickers):
    signals = {}
    returns = prices_df.pct_change().dropna()
    
    if len(prices_df) < 30: return {}

    # Benchmark returns
    sp500_ret = returns[BENCHMARK] if BENCHMARK in returns.columns else pd.Series(0, index=returns.index)
    momentum_spx = (prices_df[BENCHMARK].iloc[-1] / prices_df[BENCHMARK].iloc[-30] - 1) * 100 if BENCHMARK in prices_df.columns else 0

    # Volatility & Momentum
    volatility = returns.rolling(window=30).std() * np.sqrt(252) * 100
    current_vol = volatility.iloc[-1]
    momentum = (prices_df.iloc[-1] / prices_df.iloc[-30] - 1) * 100
    
    for t in sector_tickers:
        if t in returns.columns:
            cov = returns[t].cov(sp500_ret)
            var = sp500_ret.var()
            beta = cov / var if var != 0 else 1.0
            
            signals[t] = {
                'momentum': momentum.get(t, 0),
                'rel_strength': momentum.get(t, 0) - momentum_spx,
                'volatility': current_vol.get(t, 0),
                'beta': beta
            }
    return signals

def generate_advisory(signals, macro, fundamentals, last_private_price):
    """
    The Brain: Generates the Executive Commentary and Pricing Advice
    """
    if not signals: return {}

    avg_mom = np.mean([v['momentum'] for v in signals.values()])
    avg_vol = np.mean([v['volatility'] for v in signals.values()])
    avg_ev_rev = np.mean([f['ev_rev'] for f in fundamentals if f['ev_rev'] > 0])
    
    # 1. PRICING LOGIC
    # Revised Logic: Relative Valuation Scaling
    base_price = 30.0 # Anchor (Simplified assumption)
    
    # NEW: Percentage-based Momentum Premium/Discount
    if avg_mom > 5: 
        base_price *= 1.12  # +12% Premium for Hot Sector
    elif avg_mom < -5:
        base_price *= 0.88  # -12% Discount for Cold Sector

    if avg_vol > 40: base_price *= 0.95 # -5% for High Volatility
    if macro['vix'] > 20: base_price *= 0.90 # -10% for Macro Fear
    if macro['tnx'] > 4.5: base_price *= 0.95 # -5% for Rates
    
    low_px = base_price * 0.93 # +/- 7% Range
    high_px = base_price * 1.07
    
    # 2. MARKET WINDOW LOGIC
    window_status = "OPEN"
    window_color = "#4ade80" # Green
    
    if macro['vix'] > 25:
        window_status = "CLOSED"
        window_color = "#ef4444" # Red
        advice_text = "Market volatility (VIX > 25) indicates a closed issuance window. Severe price dislocation risk."
    elif avg_mom < -5 or macro['vix'] > 20:
        window_status = "CAUTION"
        window_color = "#facc15" # Yellow/Orange
        advice_text = "Headwinds present. Buy-side demand is highly selective. Recommend widening the price talk."
    else:
        advice_text = "Constructive backdrop. Strong peer momentum (Avg Beta < 1.0) supports a premium valuation relative to the sector curve."

    # 3. DOWN-ROUND DETECTOR
    down_round_alert = False
    if last_private_price:
        try:
            lpp = float(last_private_price)
            if lpp > high_px:
                advice_text += f"<br><br><b>⚠️ DOWN-ROUND RISK:</b> Implied range (${low_px:.2f}-${high_px:.2f}) is below last private mark (${lpp:.2f}). Expect significant cap table friction."
        except:
            pass

    # 4. YIELD SENSITIVITY (New)
    yield_impact = "Neutral"
    if macro['tnx'] > 4.5:
        yield_impact = "Negative (High Rates)"
    elif macro['tnx'] < 3.5:
        yield_impact = "Positive (Low Rates)"

    # 5. RISK CHECKLIST GENERATION
    risk_matrix = []
    
    # Risk 1: Volatility
    req_vix = "🟢" if macro['vix'] < 20 else ("🔴" if macro['vix'] > 25 else "🟡")
    risk_matrix.append({
        "factor": "Market Volatility (VIX)",
        "status": f"{req_vix} {macro['vix']:.2f}",
        "impact": "Constructive" if macro['vix'] < 20 else "Dislocated"
    })
    
    # Risk 2: Valuation Gap (Down-Round)
    if last_private_price:
        try:
            lpp = float(last_private_price)
            gap_pct = ((high_px - lpp) / lpp) * 100
            
            if gap_pct < 0:
                icon = "🔴"
                impact_text = "Down-Round Risk"
            elif gap_pct < 15:
                icon = "🟡"
                impact_text = "Flat/Small Uplift"
            else:
                icon = "🟢"
                impact_text = "Accretive IPO"
                
            risk_matrix.append({
                "factor": "Valuation Step-Up",
                "status": f"{icon} {gap_pct:+.1f}%",
                "impact": impact_text
            })
        except:
            pass
            
    # Risk 3: Yield Environment
    risk_matrix.append({
        "factor": "Yield Environment (10Y)",
        "status": f"{'🔴' if macro['tnx'] > 4.2 else '🟢'} {macro['tnx']:.2f}%",
        "impact": "Valuation Drag" if macro['tnx'] > 4.2 else "Supportive"
    })
    
    # Risk 4: Sector Health (Rule of 40)
    avg_r40 = np.mean([f['rule_40'] for f in fundamentals]) if fundamentals else 0
    risk_matrix.append({
        "factor": "Sector Health (Rule of 40)",
        "status": f"{'🟢' if avg_r40 >= 40 else '🟡'} {avg_r40:.0f}",
        "impact": "Premium Multiples" if avg_r40 >= 40 else "Discount Applied"
    })

    # 6. FORMATTING (Restore V1 Style + Yield Note)
    final_text = (
        f"<b>Market Conditions:</b> {window_status} ({macro['vix']:.1f} VIX). "
        f"{advice_text}<br><br>"
        f"<b>Strategic Recommendation:</b> Based on current implied volatility and peer multiples (Avg {avg_ev_rev:.1f}x EV/Rev), "
        f"we recommend an initial pricing range of <b>${low_px:.2f} - ${high_px:.2f}</b> per share."
    )
    if macro['tnx'] > 4.2:
        final_text += " <br><i>Note: Valuation pressured by rising 10yr yields (>4.2%).</i>"

    return {
        'commentary': final_text,
        'sentiment': window_status,
        'low': round(low_px, 2),
        'high': round(high_px, 2),
        'color': window_color,
        'avg_ev_rev': round(avg_ev_rev, 1),
        'risk_matrix': risk_matrix
    }

# ==============================================================================
# IPO STRUCTURING LOGIC (Professional)
# ==============================================================================

def calculate_ipo_structure(implied_price, discount_pct, greenshoe_active, existing_shares_m, target_raise_m=250.0):
    """
    Calculates final deal structure based on banking levers.
    Target Capital Raise is now dynamic (default $250M)
    """
    target_raise = float(target_raise_m)
    
    # 1. Apply Discount
    discount_factor = (100 - discount_pct) / 100
    final_price = implied_price * discount_factor
    
    if final_price <= 0: final_price = 1.0 # Safety
    
    # 2. Calculate Issuance
    new_shares_m = target_raise / final_price
    
    # 3. Greenshoe Adjustment (Standard 15% Over-allotment)
    greenshoe_shares_m = 0.0
    if greenshoe_active:
        greenshoe_shares_m = new_shares_m * 0.15
        new_shares_m += greenshoe_shares_m
        target_raise += (greenshoe_shares_m * final_price)
    
    # 4. Dilution & Ownership
    total_shares_m = existing_shares_m + new_shares_m
    dilution_pct = (new_shares_m / total_shares_m) * 100
    
    # Ownership Split
    ownership = {
        "Existing Shareholders": round(existing_shares_m, 2),
        "New Public Investors": round(new_shares_m, 2)
    }
    
    return {
        "final_price": round(final_price, 2),
        "capital_raised": round(target_raise, 2),
        "new_shares": round(new_shares_m, 2),
        "total_shares": round(total_shares_m, 2),
        "dilution": round(dilution_pct, 1),
        "ownership": ownership
    }

# ==============================================================================
# ROUTES
# ==============================================================================

@app.get("/", response_class=HTMLResponse)
async def read_root(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.get("/health")
async def health_check():
    return {"status": "ok"}

@app.post("/analyze")
async def analyze(request: Request, 
                  query: str = Form(...), 
                  last_private: str = Form(None),
                  ipo_discount: float = Form(15.0),  # Default 15%
                  greenshoe: bool = Form(False),     # Default Off
                  primary_shares: float = Form(100.0), # Default 100M shares
                  target_raise: float = Form(250.0)): # Default $250M
    
    # 1. Determine Sector
    sector_key = 'SaaS'
    q_lower = query.lower()
    if 'fintech' in q_lower: sector_key = 'Fintech'
    elif 'bio' in q_lower: sector_key = 'Biotech'
    elif 'ai' in q_lower: sector_key = 'AI_Infra'
    
    target_tickers = SECTOR_PROXIES.get(sector_key, SECTOR_PROXIES['SaaS'])
    
    # 2. Fetch Data (Market + Macro)
    prices = fetch_market_data(target_tickers)
    macro = fetch_macro_data()
    
    if prices.empty:
        return JSONResponse(status_code=500, content={"error": "Failed to fetch market data"})
        
    # 3. Core Calculations
    signals = calculate_signals(prices, target_tickers)
    fundamentals = get_fundamentals(target_tickers)
    
    # 4. The Advisor Engine
    # Get base advisory first to get the 'High' implied price
    advisory = generate_advisory(signals, macro, fundamentals, last_private)
    
    # 5. IPO Structuring (The Pro Layer)
    # We use the midpoint of the implied range as the base for discounting
    implied_midpoint = (advisory['low'] + advisory['high']) / 2
    structure = calculate_ipo_structure(implied_midpoint, ipo_discount, greenshoe, primary_shares, target_raise)
    
    # Update Advisory with Final Price Context
    final_price = structure['final_price']
    
    # Re-Run Down-Round Logic on FINAL PRICE
    down_round_alert = False
    dr_text = ""
    if last_private:
        try:
            lpp = float(last_private)
            if final_price < lpp:
                down_round_alert = True
                diff = ((final_price - lpp) / lpp) * 100
                dr_text = f"🚨 <b>DOWN-ROUND ALERT:</b> Final IPO Price (${final_price}) is {diff:.1f}% below Last Private Round (${lpp})."
        except:
            pass

    # 6. Chart Data
    normalized = prices / prices.iloc[0] * 100
    chart_data = []
    for col in normalized.columns:
        if col in target_tickers or col == BENCHMARK:
            chart_data.append({
                'x': normalized.index.strftime('%Y-%m-%d').tolist(),
                'y': normalized[col].values.tolist(),
                'name': col,
                'type': 'scatter',
                'mode': 'lines'
            })
            
    # 7. Response
    response_data = {
        'sector': sector_key,
        'advisory': advisory,
        'structure': structure, # NEW
        'down_round': {'is_active': down_round_alert, 'text': dr_text}, # NEW
        'macro': macro,
        'metrics': { 
            'avg_momentum': np.mean([s['momentum'] for s in signals.values()]) if signals else 0,
            'avg_beta': np.mean([s['beta'] for s in signals.values()]) if signals else 0,
            'avg_vol': np.mean([s['volatility'] for s in signals.values()]) if signals else 0,
            'avg_ev_rev': advisory['avg_ev_rev'],
            'avg_rule_40': np.mean([f['rule_40'] for f in fundamentals]) if fundamentals else 0
        },
        'chart_json': chart_data,
        'comparables': fundamentals,
        'signals': signals 
    }
    
    return JSONResponse(content=response_data)