#!/usr/bin/env python3
"""
Write AI scores into the Excel file and highlight modified cells with yellow background.
"""

import json
import pandas as pd
from openpyxl import load_workbook
from openpyxl.styles import PatternFill, Font, Alignment, Border, Side
from copy import copy

EXCEL_SRC = "/home/ubuntu/upload/MemelyAlphaStockRanking（副本）（副本）.xlsx"
SCORES_JSON = "/home/ubuntu/memely-alpha-stock-ranking/research_scores.json"
EXCEL_OUT = "/home/ubuntu/memely-alpha-stock-ranking/MemelyAlphaStockRanking_scored.xlsx"

# Load scores
with open(SCORES_JSON, 'r') as f:
    scores = json.load(f)

print(f"Loaded scores for {len(scores)} stocks")

# Load workbook
wb = load_workbook(EXCEL_SRC)

# Yellow fill for modified cells
yellow_fill = PatternFill(start_color="FFFF00", end_color="FFFF00", fill_type="solid")
header_fill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid")
header_font = Font(bold=True, color="FFFFFF", size=11)
score_font = Font(size=10)
thin_border = Border(
    left=Side(style='thin'),
    right=Side(style='thin'),
    top=Side(style='thin'),
    bottom=Side(style='thin')
)

# ========== Modify "stock master list" sheet ==========
ws = wb['stock master list']

# Find the last column with data
max_col = ws.max_column

# Add new score columns
new_cols = {
    max_col + 1: 'AI Overall Score',
    max_col + 2: 'AI Momentum Score',
    max_col + 3: 'AI Theme Score',
    max_col + 4: 'AI Risk Score',
    max_col + 5: 'AI Social Buzz Score',
    max_col + 6: 'AI Brief Analysis',
}

# Write headers
for col_offset, header_name in new_cols.items():
    cell = ws.cell(row=1, column=col_offset, value=header_name)
    cell.fill = header_fill
    cell.font = header_font
    cell.alignment = Alignment(horizontal='center', wrap_text=True)
    cell.border = thin_border

# Write scores for each stock
for row in range(2, ws.max_row + 1):
    symbol = ws.cell(row=row, column=1).value
    if symbol and symbol in scores:
        s = scores[symbol]
        
        # Overall Score
        cell = ws.cell(row=row, column=max_col + 1, value=s.get('overall_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        # Momentum Score
        cell = ws.cell(row=row, column=max_col + 2, value=s.get('momentum_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        # Theme Score
        cell = ws.cell(row=row, column=max_col + 3, value=s.get('theme_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        # Risk Score
        cell = ws.cell(row=row, column=max_col + 4, value=s.get('risk_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        # Social Buzz Score
        cell = ws.cell(row=row, column=max_col + 5, value=s.get('social_buzz_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        # Brief Analysis
        cell = ws.cell(row=row, column=max_col + 6, value=s.get('brief_analysis', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(wrap_text=True)
        cell.border = thin_border

# Set column widths
ws.column_dimensions[ws.cell(row=1, column=max_col + 1).column_letter].width = 15
ws.column_dimensions[ws.cell(row=1, column=max_col + 2).column_letter].width = 15
ws.column_dimensions[ws.cell(row=1, column=max_col + 3).column_letter].width = 15
ws.column_dimensions[ws.cell(row=1, column=max_col + 4).column_letter].width = 15
ws.column_dimensions[ws.cell(row=1, column=max_col + 5).column_letter].width = 18
ws.column_dimensions[ws.cell(row=1, column=max_col + 6).column_letter].width = 60

print(f"Updated 'stock master list' sheet with scores")

# ========== Also modify Sheet4 with scores ==========
ws4 = wb['Sheet4']
max_col4 = ws4.max_column

new_cols4 = {
    max_col4 + 1: 'AI Overall Score',
    max_col4 + 2: 'AI Momentum Score',
    max_col4 + 3: 'AI Theme Score',
    max_col4 + 4: 'AI Risk Score',
    max_col4 + 5: 'AI Social Buzz Score',
    max_col4 + 6: 'AI Brief Analysis',
}

# Write headers for Sheet4
for col_offset, header_name in new_cols4.items():
    cell = ws4.cell(row=1, column=col_offset, value=header_name)
    cell.fill = header_fill
    cell.font = header_font
    cell.alignment = Alignment(horizontal='center', wrap_text=True)
    cell.border = thin_border

# Write scores for Sheet4
for row in range(2, ws4.max_row + 1):
    symbol = ws4.cell(row=row, column=1).value
    if symbol and symbol in scores:
        s = scores[symbol]
        
        cell = ws4.cell(row=row, column=max_col4 + 1, value=s.get('overall_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        cell = ws4.cell(row=row, column=max_col4 + 2, value=s.get('momentum_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        cell = ws4.cell(row=row, column=max_col4 + 3, value=s.get('theme_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        cell = ws4.cell(row=row, column=max_col4 + 4, value=s.get('risk_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        cell = ws4.cell(row=row, column=max_col4 + 5, value=s.get('social_buzz_score', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(horizontal='center')
        cell.border = thin_border
        
        cell = ws4.cell(row=row, column=max_col4 + 6, value=s.get('brief_analysis', ''))
        cell.fill = yellow_fill
        cell.font = score_font
        cell.alignment = Alignment(wrap_text=True)
        cell.border = thin_border

ws4.column_dimensions[ws4.cell(row=1, column=max_col4 + 1).column_letter].width = 15
ws4.column_dimensions[ws4.cell(row=1, column=max_col4 + 2).column_letter].width = 15
ws4.column_dimensions[ws4.cell(row=1, column=max_col4 + 3).column_letter].width = 15
ws4.column_dimensions[ws4.cell(row=1, column=max_col4 + 4).column_letter].width = 15
ws4.column_dimensions[ws4.cell(row=1, column=max_col4 + 5).column_letter].width = 18
ws4.column_dimensions[ws4.cell(row=1, column=max_col4 + 6).column_letter].width = 60

print(f"Updated 'Sheet4' with scores")

# ========== Create a new summary sheet ==========
ws_summary = wb.create_sheet("AI Score Summary", 0)

# Headers
summary_headers = ['Rank', 'Symbol', 'Theme', 'Sector', 'Overall Score', 'Momentum Score', 
                    'Theme Score', 'Risk Score', 'Social Buzz Score', 'Perf YTD', 'Brief Analysis']

for col, header in enumerate(summary_headers, 1):
    cell = ws_summary.cell(row=1, column=col, value=header)
    cell.fill = header_fill
    cell.font = header_font
    cell.alignment = Alignment(horizontal='center', wrap_text=True)
    cell.border = thin_border

# Sort stocks by overall score
sorted_stocks = sorted(scores.items(), key=lambda x: x[1].get('overall_score', 0), reverse=True)

# Read original data for theme/sector info
df = pd.read_excel(EXCEL_SRC, sheet_name='stock master list')
df4 = pd.read_excel(EXCEL_SRC, sheet_name='Sheet4')

symbol_info = {}
for _, row in df.iterrows():
    sym = row['Symbol']
    if pd.notna(sym):
        symbol_info[sym] = {
            'theme': row.get('Theme', ''),
            'perf_ytd': row.get('performance Year to Date', None)
        }

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

# Color coding for scores
from openpyxl.styles import PatternFill

def get_score_fill(score):
    if score >= 70:
        return PatternFill(start_color="92D050", end_color="92D050", fill_type="solid")  # Green
    elif score >= 50:
        return PatternFill(start_color="FFFF00", end_color="FFFF00", fill_type="solid")  # Yellow
    elif score >= 30:
        return PatternFill(start_color="FFC000", end_color="FFC000", fill_type="solid")  # Orange
    else:
        return PatternFill(start_color="FF6B6B", end_color="FF6B6B", fill_type="solid")  # Red

for rank, (symbol, s) in enumerate(sorted_stocks, 1):
    row = rank + 1
    info = symbol_info.get(symbol, {})
    
    ws_summary.cell(row=row, column=1, value=rank).border = thin_border
    ws_summary.cell(row=row, column=1).alignment = Alignment(horizontal='center')
    
    ws_summary.cell(row=row, column=2, value=symbol).border = thin_border
    ws_summary.cell(row=row, column=2).font = Font(bold=True)
    
    theme = info.get('theme', '')
    ws_summary.cell(row=row, column=3, value=theme if pd.notna(theme) else '').border = thin_border
    
    sector = sector_info.get(symbol, '')
    ws_summary.cell(row=row, column=4, value=sector if pd.notna(sector) else '').border = thin_border
    
    # Overall Score with color
    overall = s.get('overall_score', 0)
    cell = ws_summary.cell(row=row, column=5, value=overall)
    cell.fill = get_score_fill(overall)
    cell.font = Font(bold=True, size=11)
    cell.alignment = Alignment(horizontal='center')
    cell.border = thin_border
    
    # Momentum Score
    momentum = s.get('momentum_score', 0)
    cell = ws_summary.cell(row=row, column=6, value=momentum)
    cell.fill = get_score_fill(momentum)
    cell.alignment = Alignment(horizontal='center')
    cell.border = thin_border
    
    # Theme Score
    theme_sc = s.get('theme_score', 0)
    cell = ws_summary.cell(row=row, column=7, value=theme_sc)
    cell.fill = get_score_fill(theme_sc)
    cell.alignment = Alignment(horizontal='center')
    cell.border = thin_border
    
    # Risk Score (higher = more risky, use inverse coloring)
    risk = s.get('risk_score', 0)
    cell = ws_summary.cell(row=row, column=8, value=risk)
    # For risk, high is bad
    if risk >= 70:
        cell.fill = PatternFill(start_color="FF6B6B", end_color="FF6B6B", fill_type="solid")
    elif risk >= 50:
        cell.fill = PatternFill(start_color="FFC000", end_color="FFC000", fill_type="solid")
    else:
        cell.fill = PatternFill(start_color="92D050", end_color="92D050", fill_type="solid")
    cell.alignment = Alignment(horizontal='center')
    cell.border = thin_border
    
    # Social Buzz Score
    buzz = s.get('social_buzz_score', 0)
    cell = ws_summary.cell(row=row, column=9, value=buzz)
    cell.fill = get_score_fill(buzz)
    cell.alignment = Alignment(horizontal='center')
    cell.border = thin_border
    
    # YTD Performance
    ytd = info.get('perf_ytd', None)
    if pd.notna(ytd):
        cell = ws_summary.cell(row=row, column=10, value=ytd)
        cell.number_format = '0.00%'
    else:
        cell = ws_summary.cell(row=row, column=10, value='N/A')
    cell.alignment = Alignment(horizontal='center')
    cell.border = thin_border
    
    # Brief Analysis
    cell = ws_summary.cell(row=row, column=11, value=s.get('brief_analysis', ''))
    cell.alignment = Alignment(wrap_text=True)
    cell.border = thin_border

# Set column widths for summary
col_widths = [6, 10, 20, 25, 14, 16, 14, 12, 18, 12, 60]
for i, w in enumerate(col_widths, 1):
    ws_summary.column_dimensions[ws_summary.cell(row=1, column=i).column_letter].width = w

# Freeze panes
ws_summary.freeze_panes = 'A2'

print(f"Created 'AI Score Summary' sheet with {len(sorted_stocks)} stocks ranked")

# Save
wb.save(EXCEL_OUT)
print(f"\nSaved scored Excel to: {EXCEL_OUT}")
print("All modified/new cells are highlighted with yellow background")