agent.py

import json
import logging
import asyncio
import time
import hashlib
import random
from typing import List, Dict, Any, Optional
from openai import AsyncOpenAI
import config
from config import NEMOTRON_API_KEY, CLAUDE_API_KEY, GEMINI_FREE_API_KEY, GEMINI_PAID_API_KEY, LLM_API_KEY, DATA_DIR, GEMINI_KEYS

logger = logging.getLogger("AlphaAgent")

# Global state for frontend tracking and rate limiting
CURRENT_MODEL = "None"
USED_LOCAL_FALLBACK = False  # Track if last generation was local fallback

# Rotate multiple Gemini keys to avoid concurrent rate-limiting
_gemini_key_counter = 0
_KEY_COOLDOWNS: Dict[str, float] = {}
_PER_KEY_COOLDOWN_SECS = 120  # 2 minutes key cooldown

def get_next_gemini_key() -> Optional[str]:
    global _gemini_key_counter
    keys = config.GEMINI_KEYS
    if not keys:
        return None
    
    now = time.time()
    # Try to find a key that is not on cooldown
    for _ in range(len(keys)):
        key = keys[_gemini_key_counter % len(keys)]
        _gemini_key_counter += 1
        
        # Check if key is on cooldown
        if now >= _KEY_COOLDOWNS.get(key, 0.0):
            return key
            
    # Return None if all keys are on cooldown
    return None

# Deduplication tracking for local fallback
_GENERATED_HASHES: set = set()

# Expanded WorldQuant data fields covering multiple categories
FUNDAMENTAL_FIELDS = [
    "ebit", "capex", "assets", "debt", "equity", "cash", "cashflow"
]

PRICE_VOLUME_FIELDS = [
    "close", "open", "high", "low", "volume", "vwap", "returns",
    "adv20", "adv60", "sharesout",
]

# Formula templates for local fallback — covering diverse alpha categories
FORMULA_TEMPLATES = [
    # Cross-sectional rank-based (fundamental ratios)
    "-rank({f1} / {f2})",
    "rank({f1} / {f2})",
    "-rank({f1}) + rank({f2})",
    "-rank({f1} - {f2})",
    "rank({f2}) - rank({f1})",
    "-rank(({f1} - {f2}) / {f2})",
    # Time-series with rank
    "-rank(ts_delta({f1}, {d}))",
    "-rank(ts_delta({f1} / {f2}, {d}))",
    "-ts_decay_linear(rank({f1} / {f2}), {d})",
    "ts_decay_linear(rank({f2} / {f1}), {d})",
    "-rank(ts_mean({f1}, {d}) / {f2})",
    "rank(ts_std_dev({f1}, {d}) / ts_mean({f1}, {d}))",
    # Group-neutral variations
    "-group_rank({f1} / {f2}, sector)",
    "group_rank({f2}, subindustry) - group_rank({f1}, subindustry)",
    "-group_rank(ts_delta({f1}, {d}), sector)",
    # Zscore variants
    "-zscore(rank({f1} / {f2}))",
    "zscore({f2}) - zscore({f1})",
    "-zscore(ts_delta({f1} / {f2}, {d}))",
    # Complex multi-field
    "-rank(({f1} + {f2}) / ({f3} + {f4}))",
    "rank({f3}) * rank({f2} / {f1})",
    "-ts_decay_linear(rank({f1} / ({f2} + {f3})), {d})",
    "-rank(ts_mean({f1} / {f2}, {d}) - ts_mean({f1} / {f2}, {d2}))",
    # Momentum / mean-reversion
    "-rank(ts_delta({f1}, 5) - ts_delta({f1}, 20))",
    "rank(ts_mean({f1}, 5) / ts_mean({f1}, 20))",
    "-rank({f1} / ts_mean({f1}, {d}))",
    # Statistical arbitrage (correlation / covariance with price-volume fields)
    "-rank(ts_corr({f1}, close, {d}))",
    "-rank(ts_corr({f1} / {f2}, returns, {d}))",
    "rank(ts_covariance(rank({f1}), rank(returns), {d}))",
    "-group_rank(ts_corr({f1}, returns, {d}), subindustry)",
    "-ts_decay_linear(ts_corr(rank({f1}), rank(close), {d}), 10)",
    # Signed power
    "-signed_power(rank({f1} / {f2}), 2)",
    "signed_power(rank({f2}) - rank({f1}), 1.5)",
    # Tail / clipping
    "-rank(ts_skewness({f1}, {d}))",
    "-rank(ts_kurt({f1}, {d}))",
    # Hybrid
    "-rank({f1} / {f2}) * rank(ts_delta({f3}, {d}))",
    "-ts_decay_linear(rank({f1} / {f2}) + rank({f3} / {f4}), {d})",
]


class AlphaAgent:
    def __init__(self):
        # We initialize clients dynamically per request to support fallbacks
        pass

    async def _call_llm_with_fallback(self, prompt: str, task_type: str) -> Optional[str]:
        global CURRENT_MODEL, USED_LOCAL_FALLBACK
        
        now = time.time()
        
        # Determine fallback chain based on task type.
        if task_type == "refine":
            chain = [
                {"name": "Claude 3.5 Sonnet", "model": "anthropic/claude-3.5-sonnet", "key": CLAUDE_API_KEY},
                {"name": "Gemini 3.5 Flash", "model": "google/gemini-3.5-flash", "key": None},
                {"name": "Gemini 2.5 Flash", "model": "google/gemini-2.5-flash", "key": None},
                {"name": "Gemini 2.5 Pro", "model": "google/gemini-2.5-pro", "key": None},
                {"name": "Gemini 3.1 Pro Preview", "model": "google/gemini-3.1-pro-preview", "key": None},
                {"name": "Gemini 3.1 Flash-Lite", "model": "google/gemini-3.1-flash-lite", "key": None},
                {"name": "Gemini 2.5 Flash-Lite", "model": "google/gemini-2.5-flash-lite", "key": None},
                {"name": "Nemotron 3 Ultra", "model": "nvidia/nemotron-3-ultra-550b-a55b", "key": NEMOTRON_API_KEY}
            ]
        else:
            chain = [
                {"name": "Gemini 3.5 Flash", "model": "google/gemini-3.5-flash", "key": None},
                {"name": "Gemini 2.5 Flash", "model": "google/gemini-2.5-flash", "key": None},
                {"name": "Gemini 2.5 Pro", "model": "google/gemini-2.5-pro", "key": None},
                {"name": "Gemini 3.1 Pro Preview", "model": "google/gemini-3.1-pro-preview", "key": None},
                {"name": "Gemini 3.1 Flash-Lite", "model": "google/gemini-3.1-flash-lite", "key": None},
                {"name": "Gemini 2.5 Flash-Lite", "model": "google/gemini-2.5-flash-lite", "key": None},
                {"name": "Nemotron 3 Ultra", "model": "nvidia/nemotron-3-ultra-550b-a55b", "key": NEMOTRON_API_KEY},
                {"name": "Claude 3.5 Sonnet", "model": "anthropic/claude-3.5-sonnet", "key": CLAUDE_API_KEY}
            ]
            
        for attempt in chain:
            model_name_key = attempt["name"]
            is_gemini = "gemini" in model_name_key.lower() or "gemini" in attempt["model"].lower()
            
            if is_gemini:
                key = get_next_gemini_key()
            else:
                key = attempt["key"] or LLM_API_KEY
                
            if not key or key.startswith("your_") or key == "DUMMY_KEY":
                logger.debug(f"Skipping {attempt['name']} due to missing or placeholder API key.")
                continue
            
            # Key-level cooldown check
            key_cooldown_until = _KEY_COOLDOWNS.get(key, 0.0)
            if now < key_cooldown_until:
                remaining = int(key_cooldown_until - now)
                logger.info(f"Skipping {attempt['name']} because its key (prefix: {key[:6]}...) is cooling down for {remaining}s. Trying next...")
                continue
                
            CURRENT_MODEL = attempt['name']
            
            # Route Anthropic API keys to their native endpoint
            if key.startswith("sk-ant"):
                logger.info(f"Attempting native Anthropic call with {CURRENT_MODEL} for task: {task_type} (Key prefix: {key[:6]}...)")
                try:
                    from anthropic import AsyncAnthropic
                    client = AsyncAnthropic(api_key=key)
                    model_id = "claude-3-5-sonnet-20241022" if "claude" in attempt["model"] else attempt["model"]
                    
                    response = await asyncio.wait_for(
                        client.messages.create(
                            model=model_id,
                            max_tokens=4000,
                            temperature=0.7 if task_type == "generate" else 0.3,
                            messages=[{"role": "user", "content": prompt}]
                        ),
                        timeout=45.0
                    )
                    USED_LOCAL_FALLBACK = False
                    return response.content[0].text
                except Exception as e:
                    logger.warning(f"Error with native Anthropic {CURRENT_MODEL}: {e}. Falling back...")
                    continue

            # Route Google Gemini API keys to native OpenAI-compatible endpoint
            elif key.startswith("AIzaSy") or key.startswith("AQ."):
                base_url = "https://generativelanguage.googleapis.com/v1beta/openai/"
                model_id = attempt["model"].replace("google/", "")
                logger.info(f"Attempting native Gemini API call with {CURRENT_MODEL} using model {model_id} at generativelanguage.googleapis.com for task: {task_type} (Key prefix: {key[:6]}...)")
            # Route Nvidia API keys to their native endpoint
            elif key.startswith("nvapi-"):
                base_url = "https://integrate.api.nvidia.com/v1"
                model_id = "nvidia/nemotron-3-ultra-550b-a55b"
                logger.info(f"Attempting native LLM call with {CURRENT_MODEL} at integrate.api.nvidia.com for task: {task_type} (Key prefix: {key[:6]}...)")
            else:
                base_url = "https://openrouter.ai/api/v1"
                model_id = attempt["model"]
                logger.info(f"Attempting OpenRouter LLM call with {CURRENT_MODEL} for task: {task_type} (Key prefix: {key[:6]}...)")
            
            client = AsyncOpenAI(
                base_url=base_url,
                api_key=key
            )
            
            try:
                kwargs = {
                    "model": model_id,
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": 0.7 if task_type == "generate" else 0.3,
                }
                
                # Only use response_format if not using Nvidia build API (which does not support it for all models)
                if not key.startswith("nvapi-"):
                    kwargs["response_format"] = {"type": "json_object"}
                    
                response = await asyncio.wait_for(client.chat.completions.create(**kwargs), timeout=45.0)
                USED_LOCAL_FALLBACK = False
                return response.choices[0].message.content
                
            except asyncio.TimeoutError:
                logger.warning(f"Timeout Error with {CURRENT_MODEL}. Falling back...")
            except Exception as e:
                err_str = str(e).lower()
                if "429" in err_str or "rate limit" in err_str or "resource_exhausted" in err_str:
                    logger.warning(f"Rate Limit Error (429) with {CURRENT_MODEL}: {e}. Cooling down key (prefix: {key[:6]}...) for {_PER_KEY_COOLDOWN_SECS}s.")
                    _KEY_COOLDOWNS[key] = time.time() + _PER_KEY_COOLDOWN_SECS
                else:
                    logger.warning(f"Error with {CURRENT_MODEL}: {e}. Falling back...")
                    
        logger.error("All LLM fallbacks exhausted.")
        USED_LOCAL_FALLBACK = True
        return None

    async def generate_mutations(self, seed_alpha: str, data_fields: List[str], successful_examples: Optional[List[str]] = None) -> List[Dict[str, Any]]:
        """
        Generates variations of a seed alpha using an LLM, falling back to a rule-based generator if needed.
        """
        logger.info(f"Generating mutations for seed alpha: {seed_alpha}")
        
        cache_file = DATA_DIR / "data_fields_cache.json"
        categorized_prompt_str = ""
        fallback_flat_fields = []
        
        if cache_file.exists():
            try:
                with open(cache_file, "r") as f:
                    cached_data = json.load(f)
                if isinstance(cached_data, dict) and cached_data:
                    # Choose a subset of categories (up to 4)
                    available_categories = list(cached_data.keys())
                    selected_categories = random.sample(
                        available_categories, 
                        k=min(4, len(available_categories))
                    )
                    
                    prompt_lines = []
                    for cat in selected_categories:
                        prompt_lines.append(f"Category: {cat}")
                        fields_in_cat = cached_data[cat]
                        # Sample 5-7 fields per category
                        sampled_fields = random.sample(
                            fields_in_cat,
                            k=min(6, len(fields_in_cat))
                        )
                        for field in sampled_fields:
                            fid = field.get("id")
                            fdesc = field.get("desc", "No description available.")
                            ftype = field.get("type", "MATRIX")
                            prompt_lines.append(f"  - {fid} (Type: {ftype}): {fdesc}")
                            fallback_flat_fields.append(fid)
                    categorized_prompt_str = "\n".join(prompt_lines)
            except Exception as e:
                logger.warning(f"Error loading/sampling cached data fields: {e}")
                
        if not categorized_prompt_str:
            categorized_prompt_str = f"Here are some available data fields: {', '.join(data_fields)}"
            fallback_flat_fields = data_fields
            
        examples_str = ""
        if successful_examples:
            examples_str = (
                "Here are some examples of successful, high-Sharpe alphas from our previous simulations. "
                "Draw inspiration from their mathematical structures and combinations:\n" +
                "\n".join(f"- {ex}" for ex in successful_examples) + "\n\n"
            )
            
        prompt = (
            "You are an elite quantitative researcher at WorldQuant. Your task is to generate new trading alphas.\n"
            f"Here is a proven seed alpha expression: {seed_alpha}\n\n"
            f"{examples_str}"
            "Here is a subset of available data fields from our vast library, grouped by category. "
            "Each field includes its ID, data type (MATRIX or VECTOR), and description:\n"
            f"{categorized_prompt_str}\n\n"
            "Standard price-volume and basic fields you can ALWAYS use in your formulas (case-sensitive, must be lowercase):\n"
            "- close, open, high, low, volume, vwap, returns, cap\n"
            "Standard group fields you can use in group operators (case-sensitive, must be lowercase):\n"
            "- subindustry, industry, sector, market\n\n"
            "Generate 5 to 10 logical variations of the seed alpha. Use these WorldQuant operators:\n"
            "- rank(x), zscore(x), group_rank(x, group), ts_delta(x, days), ts_mean(x, days)\n"
            "- ts_std_dev(x, days), ts_decay_linear(x, days), signed_power(x, exp)\n"
            "- ts_skewness(x, days), ts_kurt(x, days), group_neutralize(x, group)\n"
            "- ts_corr(x, y, days), ts_covariance(x, y, days), ts_rank(x, days), ts_scale(x, days)\n\n"
            "Valid neutralization groups: SUBINDUSTRY, INDUSTRY, SECTOR, MARKET, NONE.\n"
            "Valid region: USA.\n"
            "Valid universes: TOP3000, TOP2000, TOP1000, TOP500, TOP200, TOPSP500.\n"
            "Valid delays: 0, 1.\n"
            "We strongly encourage you to vary the settings (neutralization, delay, decay, universe) to produce diverse alphas that run in different neutralization groups, delays, decays, and universes, which reduces correlation!\n\n"
            "Ensure each alpha is dimensionless (wrap raw fields in rank() or zscore()).\n"
            "Make diverse variations: change fields, operators, time horizons, and structures.\n"
            "Do NOT use universes outside this list.\n\n"
            "WARNING: Do NOT guess or use field names like 'sales', 'income', 'net_income', 'eps', 'revenue', 'market_cap', or 'total_assets'. Only use the fields explicitly listed in the library above or the standard price-volume/basic fields (close, open, high, low, volume, vwap, returns, cap).\n\n"
            "WARNING: Do NOT use scientific notation like '1e-9' or '1e-5' to avoid division by zero or as values. WorldQuant Brain does NOT support scientific notation. Always write float values as full decimals, e.g. use '0.000000001' or '1.0' instead of '1e-9'.\n\n"
            "Output strictly a JSON object containing a key 'mutations' which is a list of dictionaries. "
            "Each dictionary MUST have the following keys:\n"
            "- 'expression': The alpha formula string.\n"
            "- 'settings': A dictionary with keys 'delay' (0 or 1), 'decay' (e.g. 5 to 30), 'truncation' (0.01), 'neutralization' (e.g. 'SUBINDUSTRY'), 'region' (e.g. 'USA'), 'universe' (e.g. 'TOP3000', 'TOP2000', 'TOP1000', 'TOP500', 'TOP200', 'TOPSP500'). All setting text values must be strictly uppercase.\n"
            "Do not include any markdown formatting, explanations, or code blocks. Output pure JSON."
        )
        
        content = await self._call_llm_with_fallback(prompt, "generate")
        mutations = []
        
        if content:
            try:
                clean_content = content.strip()
                if clean_content.startswith("```"):
                    lines = clean_content.splitlines()
                    if lines[0].startswith("```"):
                        lines = lines[1:]
                    if lines and lines[-1].startswith("```"):
                        lines = lines[:-1]
                    clean_content = "\n".join(lines).strip()
                
                mutations_data = json.loads(clean_content)
                mutations = mutations_data.get("mutations", [])
                if mutations:
                    logger.info(f"LLM successfully generated {len(mutations)} mutations.")
            except (json.JSONDecodeError, ValueError):
                # Robust fallback: try to extract JSON object from within text/markdown
                import re
                json_match = re.search(r'\{[\s\S]*"mutations"[\s\S]*\}', content)
                if json_match:
                    try:
                        mutations_data = json.loads(json_match.group())
                        mutations = mutations_data.get("mutations", [])
                        if mutations:
                            logger.info(f"LLM JSON extracted via regex fallback: {len(mutations)} mutations.")
                    except Exception:
                        logger.warning("Regex JSON extraction also failed. Falling back to local generation.")
                else:
                    logger.warning("No JSON found in LLM response. Falling back to local generation.")
            except Exception as e:
                logger.exception("Failed to parse LLM response for mutations. Falling back to local generation.")
        
        if not mutations:
            mutations = self._generate_local_fallback_mutations(fallback_flat_fields)
            
        return mutations

    def _generate_local_fallback_mutations(self, data_fields: List[str]) -> List[Dict[str, Any]]:
        """
        Expanded local rule-based fallback generator with 30+ templates and deduplication.
        """
        logger.info("Using local rule-based fallback generator for mutations.")
        
        # Prevent hash exhaustion after long runs
        if len(_GENERATED_HASHES) > 3000:
            logger.info(f"Clearing dedup hash set ({len(_GENERATED_HASHES)} entries) to prevent generation exhaustion.")
            _GENERATED_HASHES.clear()
        
        # Combine provided fields with our expanded set for more diversity
        all_fields = list(set(data_fields + FUNDAMENTAL_FIELDS))
        random.shuffle(all_fields)
        
        mutations = []
        attempts = 0
        max_attempts = 80  # Try up to 80 combinations to find 10 unique ones
        target_count = 10
        
        days_options = [5, 10, 20, 40, 60]
        decay_options = [5, 10, 15, 20]
        neutralization_options = ["SUBINDUSTRY", "INDUSTRY", "SECTOR", "MARKET", "NONE"]
        
        while len(mutations) < target_count and attempts < max_attempts:
            attempts += 1
            
            # Pick random fields and parameters
            f1, f2 = random.sample(all_fields, 2)
            f3, f4 = random.sample(all_fields, 2)
            d = random.choice(days_options)
            d2 = random.choice([x for x in days_options if x != d])  # Different lookback
            
            template = random.choice(FORMULA_TEMPLATES)
            expression = template.format(f1=f1, f2=f2, f3=f3, f4=f4, d=d, d2=d2)
            
            # Deduplication via hash
            expr_hash = hashlib.md5(expression.encode()).hexdigest()
            if expr_hash in _GENERATED_HASHES:
                continue
            _GENERATED_HASHES.add(expr_hash)
            
            decay = random.choice(decay_options)
            neut = random.choice(neutralization_options)
            
            # Varied local settings
            region = "USA"
            universe = random.choice(["TOP3000", "TOP2000", "TOP1000", "TOP500", "TOP200", "TOPSP500"])
            delay = random.choice([0, 1])
            
            mutations.append({
                "expression": expression,
                "settings": {
                    "delay": delay,
                    "decay": decay,
                    "truncation": 0.01,
                    "neutralization": neut,
                    "region": region,
                    "universe": universe
                }
            })
        
        logger.info(f"Successfully generated {len(mutations)} local fallback mutated alphas (from {attempts} attempts, {len(_GENERATED_HASHES)} total unique hashes).")
        return mutations

    async def refine_error(self, alpha_dict: Dict[str, Any], error_log: str) -> Optional[Dict[str, Any]]:
        """
        Self-corrects a failed alpha based on the execution error log or failed metrics.
        """
        logger.info(f"Refining alpha due to error: {alpha_dict.get('expression')}")
        prompt = (
            "You are an elite quantitative researcher and debugging expert.\n"
            f"The following alpha expression failed during simulation:\n"
            f"Expression: {alpha_dict.get('expression')}\n"
            f"Settings: {json.dumps(alpha_dict.get('settings', {}))}\n"
            f"Error/Status Log: {error_log}\n\n"
            "Please apply self-correction to fix the issue based on the following rules:\n"
            "1. If the error is a unit mismatch (e.g., 'Incompatible unit... expected \"Unit\", found \"Unit\"'), "
            "wrap the offending variables in `rank()` or `zscore()` to strip their dimensionality.\n"
            "2. If the error is 'Turnover > 70%' or high turnover, increase the `ts_decay_linear` parameter or add a `trade_when` volatility gate to freeze weights.\n"
            "3. If the error is a weight concentration limit, reduce the truncation to 0.01.\n"
            "4. If the error is 'FAILED_CORRELATION' (meaning it overlaps too heavily with existing alphas), mathematically orthogonalize the signal. Do this by either subtracting a market capitalization rank component, e.g. `(original_expression) - rank(cap)`, or changing the neutralization group from `Market` to `Subindustry`.\n"
            "5. If Sharpe is too low, try changing the data fields or adding momentum factors.\n"
            "6. If Fitness is too low, try adjusting decay or adding ts_decay_linear.\n"
            "7. Crucial: Do NOT guess or use variable names. If a variable name in the failed expression (like 'sales', 'income', 'net_income', 'eps', 'revenue', 'market_cap') was rejected, you MUST replace it with a valid variable name. Valid price-volume fields are: close, open, high, low, volume, vwap, returns, and cap (case-sensitive, must be lowercase). Valid fundamental fields are ebit, capex, assets, debt, equity, cash, cashflow.\n"
            "8. Crucial: Do NOT use scientific notation like '1e-9' or '1e-5' to avoid division by zero or as values. WorldQuant Brain does NOT support scientific notation. Always write float values as full decimals, e.g. use '0.000000001' or '1.0' instead of '1e-9'.\n"
            "9. For any other syntax or logical errors, use your expertise to resolve them.\n\n"
            "Return the corrected alpha strictly in a JSON object with keys 'expression' and 'settings' "
            "(delay, decay, truncation, neutralization, region, universe). Ensure settings text values are strictly uppercase (e.g., neutralization='SUBINDUSTRY', region='USA', universe='TOP3000').\n"
            "Output pure JSON only, no markdown formatting or explanations."
        )
        
        content = await self._call_llm_with_fallback(prompt, "refine")
        refined_alpha = None
        
        if content:
            try:
                clean_content = content.strip()
                if clean_content.startswith("```"):
                    lines = clean_content.splitlines()
                    if lines[0].startswith("```"):
                        lines = lines[1:]
                    if lines and lines[-1].startswith("```"):
                        lines = lines[:-1]
                    clean_content = "\n".join(lines).strip()
                refined_alpha = json.loads(clean_content)
            except (json.JSONDecodeError, ValueError):
                import re
                json_match = re.search(r'\{[\s\S]*"expression"[\s\S]*\}', content)
                if json_match:
                    try:
                        refined_alpha = json.loads(json_match.group())
                    except Exception:
                        logger.warning("Regex JSON extraction also failed for refinement. Falling back to local corrector.")
                else:
                    logger.warning("No JSON found in LLM refinement response. Falling back to local corrector.")
            except Exception as e:
                logger.exception("Failed to parse LLM response for alpha refinement. Falling back to local corrector.")
                
        if not refined_alpha or not isinstance(refined_alpha, dict) or 'expression' not in refined_alpha:
            logger.info("Using local rule-based fallback corrector for refining error.")
            expression = alpha_dict.get("expression", "")
            settings = dict(alpha_dict.get("settings", {}))
            
            err_lower = error_log.lower()
            if "unit mismatch" in err_lower or "incompatible unit" in err_lower:
                if not (expression.startswith("rank(") or expression.startswith("-rank(")):
                    expression = f"rank({expression})"
            elif "turnover" in err_lower:
                settings["decay"] = settings.get("decay", 10) + 10
                if "ts_decay_linear" not in expression:
                    expression = f"ts_decay_linear({expression}, 20)"
            elif "concentration" in err_lower or "weight" in err_lower:
                settings["truncation"] = 0.01
            elif "correlation" in err_lower or "failed_correlation" in err_lower:
                settings["neutralization"] = "SUBINDUSTRY"
                if not expression.endswith(" - rank(cap)"):
                    expression = f"({expression}) - rank(cap)"
            elif "sharpe" in err_lower or "fitness" in err_lower:
                # Try wrapping in time-series decay for stability
                if "ts_decay_linear" not in expression:
                    expression = f"ts_decay_linear({expression}, 10)"
                settings["decay"] = min(settings.get("decay", 10) + 5, 30)
            else:
                if not (expression.startswith("rank(") or expression.startswith("-rank(")):
                    expression = f"rank({expression})"
                    
            refined_alpha = {
                "expression": expression,
                "settings": {
                    "delay": settings.get("delay", 1),
                    "decay": settings.get("decay", 10),
                    "truncation": settings.get("truncation", 0.01),
                    "neutralization": settings.get("neutralization", "SUBINDUSTRY"),
                    "region": settings.get("region", "USA"),
                    "universe": settings.get("universe", "TOP3000")
                }
            }
            logger.info(f"Successfully refined alpha using local rules to: {refined_alpha['expression']}")
            
        return refined_alpha

    async def orthogonalize_alpha(self, alpha_dict: Dict[str, Any], correlation_details: str) -> Optional[Dict[str, Any]]:
        """
        Takes an alpha that failed the correlation/overlap checks and prompts the LLM to orthogonalize it,
        decoupling it from standard/existing factors.
        """
        logger.info(f"Attempting to orthogonalize alpha: {alpha_dict.get('expression')}")
        
        prompt = (
            "You are an elite quantitative researcher at WorldQuant. Your task is to orthogonalize a trading alpha "
            "that has failed correlation checks due to overlapping too heavily with existing alphas.\n\n"
            f"Failed Alpha Expression: {alpha_dict.get('expression')}\n"
            f"Failed Settings: {json.dumps(alpha_dict.get('settings'))}\n"
            f"Platform Rejection Reason / Correlation Details: {correlation_details}\n\n"
            "To successfully orthogonalize this alpha signal, you must modify its structure or settings. Apply one or more of these techniques:\n"
            "1. Subtract a rank of market capitalization, e.g., wrapping the expression as `(original_expression) - rank(cap)` to reduce size exposure.\n"
            "2. Neutralize the expression against group hierarchies. If it is currently neutralized against MARKET, change it to SUBINDUSTRY, INDUSTRY, or SECTOR in the settings.\n"
            "3. Mathematically combine the expression with an orthogonal factor like a long-term momentum index or price-volume indicator (e.g. subtracting rank(close) or using ts_decay_linear).\n"
            "4. Change the time horizon parameters (decay, delay, or time-series operators duration).\n\n"
            "WARNING: Do NOT use scientific notation like '1e-9' or '1e-5' to avoid division by zero or as values. WorldQuant Brain does NOT support scientific notation. Always write float values as full decimals, e.g. use '0.000000001' or '1.0' instead of '1e-9'.\n\n"
            "Return the orthogonalized alpha strictly in a JSON object with keys 'expression' and 'settings' "
            "(delay, decay, truncation, neutralization, region, universe). Ensure settings text values are strictly uppercase (e.g., neutralization='SUBINDUSTRY', region='USA', universe='TOP3000').\n"
            "Output pure JSON only, no markdown formatting or explanations."
        )
        
        content = await self._call_llm_with_fallback(prompt, "refine")
        orthogonal_alpha = None
        
        if content:
            try:
                clean_content = content.strip()
                if clean_content.startswith("```"):
                    lines = clean_content.splitlines()
                    if lines[0].startswith("```"):
                        lines = lines[1:]
                    if lines and lines[-1].startswith("```"):
                        lines = lines[:-1]
                    clean_content = "\n".join(lines).strip()
                orthogonal_alpha = json.loads(clean_content)
            except (json.JSONDecodeError, ValueError):
                import re
                json_match = re.search(r'\{[\s\S]*"expression"[\s\S]*\}', content)
                if json_match:
                    try:
                        orthogonal_alpha = json.loads(json_match.group())
                    except Exception:
                        logger.warning("Regex JSON extraction also failed for orthogonalization. Falling back to local rules.")
                else:
                    logger.warning("No JSON found in LLM orthogonalization response. Falling back to local rules.")
            except Exception as e:
                logger.exception("Failed to parse LLM response for alpha orthogonalization. Falling back to local rules.")
                
        if not orthogonal_alpha or not isinstance(orthogonal_alpha, dict) or 'expression' not in orthogonal_alpha:
            logger.info("Using local rule-based fallback for orthogonalization.")
            expression = alpha_dict.get("expression", "")
            settings = dict(alpha_dict.get("settings", {}))
            
            # Local rules: apply subindustry neutralization and size deduction
            settings["neutralization"] = "SUBINDUSTRY"
            if not expression.endswith(" - rank(cap)"):
                # Subtract market cap rank to reduce size exposure
                expression = f"({expression}) - rank(cap)"
                
            orthogonal_alpha = {
                "expression": expression,
                "settings": {
                    "delay": settings.get("delay", 1),
                    "decay": settings.get("decay", 10),
                    "truncation": settings.get("truncation", 0.01),
                    "neutralization": settings.get("neutralization", "SUBINDUSTRY"),
                    "region": settings.get("region", "USA"),
                    "universe": settings.get("universe", "TOP3000")
                }
            }
            logger.info(f"Successfully orthogonalized alpha locally: {orthogonal_alpha['expression']}")
            
        return orthogonal_alpha