#!/usr/bin/env python3
"""
Score all stocks using a specified AI model.
Usage: python3 score_multi_model.py <model_name> <output_json>
"""

import pandas as pd
import json
import os
import sys
import time
from openai import OpenAI

model_name = sys.argv[1]
output_json = sys.argv[2]

client = OpenAI()

EXCEL_PATH = "/home/ubuntu/upload/MemelyAlphaStockRanking（副本）（副本）.xlsx"

# Read the master list
df = pd.read_excel(EXCEL_PATH, sheet_name='stock master list')
df4 = pd.read_excel(EXCEL_PATH, sheet_name='Sheet4')

sector_map = {}
for _, row in df4.iterrows():
    sym = row.get('Symbol')
    if pd.notna(sym):
        sector_map[sym] = {
            'sector': row.get('Sector', ''),
            'theme': row.get('Theme', ''),
        }

# Build stock info
stocks_info = []
seen = set()
for _, row in df.iterrows():
    sym = row['Symbol']
    if pd.isna(sym) or sym in seen:
        continue
    seen.add(sym)
    info = {
        'symbol': sym,
        'theme': str(row.get('Theme', '')),
        'perf_7d': row.get('performance past 7 days', None),
        'perf_ytd': row.get('performance Year to Date', None),
        'perf_specific': row.get('performance on specific dates', None),
        'first_call_kol': str(row.get('first call X kol', '')),
        'top_contributor': str(row.get('top contributor', '')),
    }
    if sym in sector_map:
        info['sector'] = sector_map[sym].get('sector', '')
        info['theme_s4'] = sector_map[sym].get('theme', '')
    stocks_info.append(info)

print(f"[{model_name}] Total stocks to score: {len(stocks_info)}")

def score_batch(batch_stocks, model):
    stocks_text = ""
    for s in batch_stocks:
        perf_7d = f"{s['perf_7d']:.2%}" if pd.notna(s.get('perf_7d')) else "N/A"
        perf_ytd = f"{s['perf_ytd']:.2%}" if pd.notna(s.get('perf_ytd')) else "N/A"
        perf_spec = f"{s['perf_specific']:.2%}" if pd.notna(s.get('perf_specific')) else "N/A"
        sector = s.get('sector', s.get('theme', 'N/A'))
        theme = s.get('theme_s4', s.get('theme', 'N/A'))
        stocks_text += f"""
---
Symbol: {s['symbol']}
Sector: {sector}
Theme: {theme}
Performance (Past 7 Days): {perf_7d}
Performance (Year to Date): {perf_ytd}
Performance (Specific Date): {perf_spec}
First Call KOL: {s.get('first_call_kol', 'N/A')}
Top Contributor: {s.get('top_contributor', 'N/A')}
"""

    prompt = f"""You are a professional stock/crypto analyst. Evaluate each of the following stocks/assets and provide a comprehensive score.

For EACH stock, provide:
1. **overall_score** (1-100): Overall investment attractiveness score
2. **momentum_score** (1-100): Based on recent price performance and momentum
3. **theme_score** (1-100): How strong/relevant is the thematic play (e.g., AI, Defense, Aerospace, Crypto, etc.)
4. **risk_score** (1-100): Risk level (100 = highest risk)
5. **social_buzz_score** (1-100): Social media attention and KOL backing strength
6. **brief_analysis**: 2-3 sentence analysis explaining the scores

Consider these factors:
- YTD performance indicates momentum strength
- Thematic relevance to current market trends (AI, Defense, Aerospace, Crypto, Quantum Computing are hot in 2025-2026)
- Small/micro cap stocks with strong themes get higher theme scores
- Stocks with notable KOL backing get higher social buzz scores
- Higher volatility = higher risk score

Here are the stocks to evaluate:
{stocks_text}

Return ONLY a valid JSON array with objects for each stock. Each object must have: symbol, overall_score, momentum_score, theme_score, risk_score, social_buzz_score, brief_analysis.
Do not include any text outside the JSON array."""

    for attempt in range(3):
        try:
            response = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": "You are a professional financial analyst. Always respond with valid JSON only."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.3,
                max_tokens=8000
            )
            content = response.choices[0].message.content.strip()
            if content.startswith("```"):
                content = content.split("```")[1]
                if content.startswith("json"):
                    content = content[4:]
            results = json.loads(content)
            return results
        except Exception as e:
            print(f"  Attempt {attempt+1} error: {e}")
            time.sleep(3)
    return None

# Process in batches
batch_size = 15
scored = {}
total_batches = (len(stocks_info) + batch_size - 1) // batch_size

for i in range(0, len(stocks_info), batch_size):
    batch = stocks_info[i:i+batch_size]
    batch_num = i // batch_size + 1
    symbols_in_batch = [s['symbol'] for s in batch]
    print(f"[{model_name}] Batch {batch_num}/{total_batches}: {symbols_in_batch}")
    
    results = score_batch(batch, model_name)
    
    if results:
        for r in results:
            sym = r.get('symbol', '')
            if sym:
                scored[sym] = r
                print(f"  ✓ {sym}: Overall={r.get('overall_score')}")
    else:
        print(f"  ✗ Batch {batch_num} failed after retries")
    
    time.sleep(1)

with open(output_json, 'w') as f:
    json.dump(scored, f, indent=2, ensure_ascii=False)

print(f"\n[{model_name}] COMPLETE: {len(scored)}/{len(stocks_info)} scored")
print(f"Saved to: {output_json}")