"""
Factor expression parser and calculator module

This module provides tools to parse factor expressions from different frameworks
and calculate factor values based on those expressions.
"""

import numpy as np
import pandas as pd
import re
import ast
from typing import Dict, Any, List, Optional, Tuple, Union
import logging

from ..operators.unified_operators import get_operator_registry, UnifiedOperatorRegistry

logger = logging.getLogger(__name__)


class ExpressionNormalizer:
    """
    Normalizes factor expressions from different frameworks to a unified format.

    Supports expressions from:
    - gpfactor: sub(close)(low) format (C++ style)
    - mas1: Mul(Rank(...), EMA(...)) format (CamelCase)
    - masfactorMiner: mul(-1, cs_rank(ts_mean(...))) format (lowercase with prefixes)
    """

    # Pattern for C++ style function calls: func(arg1)(arg2)
    CPP_CALL_PATTERN = re.compile(r'(\w+)\(([^()]*)\)\(([^()]*)\)')

    # Data field mappings across different frameworks
    FIELD_ALIASES = {
        # Standard OHLCV
        'open': ['open', 'Open', 'OPEN', 'o'],
        'high': ['high', 'High', 'HIGH', 'h'],
        'low': ['low', 'Low', 'LOW', 'l'],
        'close': ['close', 'Close', 'CLOSE', 'c'],
        'volume': ['volume', 'Volume', 'VOLUME', 'vol', 'Vol', 'VOL'],
        'amount': ['amount', 'Amount', 'AMOUNT', 'amt'],
        # Derived fields
        'pct_chg': ['pct_chg', 'pctChg', 'returns', 'ret', 'change'],
        'circ_mv': ['circ_mv', 'circMv', 'market_cap', 'mktcap', 'circ_market_cap'],
        'pre_close': ['pre_close', 'preClose', 'prev_close'],
        'vwap': ['vwap', 'VWAP', 'avg_price'],
    }

    def __init__(self, registry: Optional[UnifiedOperatorRegistry] = None):
        self.registry = registry or get_operator_registry()
        self._build_field_alias_map()

    def _build_field_alias_map(self):
        """Build reverse alias map for fields"""
        self._field_alias_map = {}
        for canonical, aliases in self.FIELD_ALIASES.items():
            for alias in aliases:
                self._field_alias_map[alias.lower()] = canonical

    def normalize_field(self, field: str) -> str:
        """Normalize a data field name to canonical form"""
        return self._field_alias_map.get(field.lower(), field.lower())

    def normalize_operator(self, op_name: str) -> str:
        """Normalize an operator name to canonical form"""
        canonical = self.registry.get_canonical_name(op_name)
        return canonical if canonical else op_name.lower()

    def normalize_expression(self, expression: str, source_framework: str = "auto") -> str:
        """
        Normalize an expression to the unified format.

        Args:
            expression: The original expression string
            source_framework: One of "gpfactor", "mas1", "masfactorMiner", or "auto"

        Returns:
            Normalized expression string
        """
        if source_framework == "auto":
            source_framework = self._detect_framework(expression)

        if source_framework == "gpfactor":
            return self._normalize_gpfactor(expression)
        elif source_framework == "mas1":
            return self._normalize_mas1(expression)
        elif source_framework in ("masfactorMiner", "mas2"):
            return self._normalize_masfactor_miner(expression)
        else:
            # Default: try to normalize operators and fields
            return self._normalize_generic(expression)

    def _detect_framework(self, expression: str) -> str:
        """Auto-detect the source framework based on expression patterns"""
        # gpfactor: uses func(arg1)(arg2) pattern
        if self.CPP_CALL_PATTERN.search(expression):
            return "gpfactor"

        # masfactorMiner: uses ts_ or cs_ prefixes with lowercase
        if re.search(r'\b(ts_|cs_)\w+', expression):
            return "masfactorMiner"

        # mas1: uses CamelCase operators like Mul, Rank, EMA
        if re.search(r'\b(Mul|Div|Add|Sub|Rank|EMA|Mean)\s*\(', expression):
            return "mas1"

        return "generic"

    def _normalize_gpfactor(self, expression: str) -> str:
        """
        Normalize gpfactor expressions.
        Example: sub(close)(low) -> sub(close, low)
        """
        result = expression

        # Convert func(arg1)(arg2) to func(arg1, arg2)
        while True:
            match = self.CPP_CALL_PATTERN.search(result)
            if not match:
                break
            func_name = match.group(1)
            arg1 = match.group(2)
            arg2 = match.group(3)

            # Normalize operator name
            normalized_op = self.normalize_operator(func_name)
            replacement = f"{normalized_op}({arg1}, {arg2})"
            result = result[:match.start()] + replacement + result[match.end():]

        # Normalize remaining operators and fields
        return self._normalize_generic(result)

    def _normalize_mas1(self, expression: str) -> str:
        """
        Normalize mas1 expressions (CamelCase operators).
        Example: Mul(Rank(x), EMA(y, 5)) -> mul(cs_rank(x), ts_ema(y, 5))
        """
        result = expression

        # Find all operator calls and normalize them
        op_pattern = re.compile(r'\b([A-Za-z_][A-Za-z0-9_]*)\s*\(')
        offset = 0

        for match in op_pattern.finditer(expression):
            op_name = match.group(1)
            normalized_op = self.normalize_operator(op_name)

            if normalized_op != op_name:
                start = match.start(1) + offset
                end = match.end(1) + offset
                result = result[:start] + normalized_op + result[end:]
                offset += len(normalized_op) - len(op_name)

        # Normalize field names
        return self._normalize_fields(result)

    def _normalize_masfactor_miner(self, expression: str) -> str:
        """
        Normalize masfactorMiner expressions (already mostly normalized).
        Just ensure field names are canonical.
        """
        return self._normalize_fields(expression)

    def _normalize_generic(self, expression: str) -> str:
        """Generic normalization for operators and fields"""
        result = expression

        # Normalize operators
        op_pattern = re.compile(r'\b([A-Za-z_][A-Za-z0-9_]*)\s*\(')

        # Collect all operators first to avoid issues with overlapping replacements
        operators = [(m.start(1), m.end(1), m.group(1)) for m in op_pattern.finditer(result)]

        # Replace in reverse order to maintain positions
        for start, end, op_name in reversed(operators):
            normalized_op = self.normalize_operator(op_name)
            result = result[:start] + normalized_op + result[end:]

        # Normalize fields
        return self._normalize_fields(result)

    def _normalize_fields(self, expression: str) -> str:
        """Normalize data field names in an expression"""
        result = expression

        # Find all identifiers that are not operators (not followed by '(')
        # and normalize them if they are known fields
        field_pattern = re.compile(r'\b([A-Za-z_][A-Za-z0-9_]*)\b(?!\s*\()')

        offset = 0
        for match in field_pattern.finditer(expression):
            field = match.group(1)
            normalized = self.normalize_field(field)

            if normalized != field:
                start = match.start(1) + offset
                end = match.end(1) + offset
                result = result[:start] + normalized + result[end:]
                offset += len(normalized) - len(field)

        return result


class ExpressionCalculator:
    """
    Calculate factor values from expression strings.

    Supports calculation for both panel data (grouped by code) and single series.
    """

    def __init__(
        self,
        registry: Optional[UnifiedOperatorRegistry] = None,
        normalizer: Optional[ExpressionNormalizer] = None
    ):
        self.registry = registry or get_operator_registry()
        self.normalizer = normalizer or ExpressionNormalizer(self.registry)

    def calculate(
        self,
        expression: str,
        data: pd.DataFrame,
        code_column: str = "code",
        date_column: str = "date",
        normalize: bool = True,
        source_framework: str = "auto"
    ) -> pd.Series:
        """
        Calculate factor values from expression.

        Args:
            expression: Factor expression string
            data: DataFrame with market data
            code_column: Name of code column
            date_column: Name of date column
            normalize: Whether to normalize the expression first
            source_framework: Source framework for normalization

        Returns:
            Series of factor values with MultiIndex (code, date)
        """
        if normalize:
            expression = self.normalizer.normalize_expression(expression, source_framework)

        if not expression:
            raise ValueError("Empty expression")

        # Build namespace with operators
        namespace = self._build_namespace()

        try:
            # Sort data by code and date
            data = data.sort_values([code_column, date_column]).copy()

            # Group by code and apply expression
            results = []
            for code, group in data.groupby(code_column):
                group_result = self._evaluate_expression(expression, group, namespace)
                if group_result is not None and len(group_result) > 0:
                    if isinstance(group_result, pd.Series):
                        group_result = group_result.values
                    result_series = pd.Series(
                        group_result,
                        index=pd.MultiIndex.from_arrays(
                            [group[code_column].values, group[date_column].values],
                            names=[code_column, date_column]
                        )
                    )
                    results.append(result_series)

            if not results:
                raise ValueError("No valid results from expression evaluation")

            final_result = pd.concat(results)
            return final_result.astype(np.float32)

        except Exception as e:
            logger.error(f"Error calculating expression '{expression}': {e}")
            raise

    def calculate_cross_sectional(
        self,
        expression: str,
        data: pd.DataFrame,
        date_column: str = "date",
        normalize: bool = True,
        source_framework: str = "auto"
    ) -> pd.DataFrame:
        """
        Calculate factor values in cross-sectional format (date x code matrix).

        Args:
            expression: Factor expression string
            data: DataFrame in wide format (date index, code columns) or long format
            date_column: Name of date column (for long format)
            normalize: Whether to normalize the expression first
            source_framework: Source framework for normalization

        Returns:
            DataFrame with date index and code columns
        """
        if normalize:
            expression = self.normalizer.normalize_expression(expression, source_framework)

        namespace = self._build_namespace()

        # If data is in long format, pivot to wide format
        if date_column in data.columns:
            # Assume data has code, date, and value columns
            raise NotImplementedError("Long format input not yet supported for cross-sectional calculation")

        # Data is already in wide format (date index, code columns)
        try:
            local_ns = namespace.copy()
            # Add each column as a field
            for col in data.columns:
                local_ns[col.lower()] = data[col]
                local_ns[col] = data[col]

            result = eval(expression, {"__builtins__": {}}, local_ns)
            if isinstance(result, pd.DataFrame):
                return result.astype(np.float32)
            elif isinstance(result, pd.Series):
                return pd.DataFrame(result).astype(np.float32)
            else:
                return pd.DataFrame(result, index=data.index, columns=data.columns).astype(np.float32)

        except Exception as e:
            logger.error(f"Error calculating cross-sectional expression '{expression}': {e}")
            raise

    def _build_namespace(self) -> Dict[str, Any]:
        """Build namespace for expression evaluation"""
        namespace = {
            'np': np,
            'pd': pd,
        }

        # Add all operators
        for name, info in self.registry.get_all_operators().items():
            namespace[name] = info.func
            # Also add aliases
            for alias in info.aliases:
                namespace[alias] = info.func

        return namespace

    def _evaluate_expression(
        self,
        expression: str,
        group: pd.DataFrame,
        namespace: Dict[str, Any]
    ) -> pd.Series:
        """Evaluate expression for a single group (code)"""
        local_ns = namespace.copy()

        # Add field values as Series
        for col in group.columns:
            if col not in ['code', 'date', 'datetime']:
                series_data = group[col].reset_index(drop=True)
                # Add with various case variants
                local_ns[col] = series_data
                local_ns[col.lower()] = series_data
                local_ns[col.upper()] = series_data
                local_ns[col.capitalize()] = series_data

        try:
            result = eval(expression, {"__builtins__": {}}, local_ns)
            if isinstance(result, (int, float)):
                result = pd.Series([result] * len(group), index=group.index)
            elif isinstance(result, np.ndarray):
                result = pd.Series(result, index=group.index)
            elif isinstance(result, pd.Series):
                result.index = group.index
            return result
        except Exception as e:
            logger.warning(f"Expression evaluation failed for group: {e}")
            return pd.Series([np.nan] * len(group), index=group.index)

    def validate_expression(
        self,
        expression: str,
        available_fields: List[str],
        normalize: bool = True
    ) -> Tuple[bool, str]:
        """
        Validate expression syntax and operators.

        Args:
            expression: Factor expression string
            available_fields: List of available field names

        Returns:
            Tuple of (is_valid, error_message)
        """
        if normalize:
            expression = self.normalizer.normalize_expression(expression)

        # Check for empty expression
        if not expression or not expression.strip():
            return False, "Empty expression"

        # Check for Chinese characters
        if any('\u4e00' <= c <= '\u9fff' for c in expression):
            return False, "Expression contains Chinese characters"

        # Check expression length
        if len(expression) < 3:
            return False, "Expression too short"

        # Check for balanced parentheses
        if expression.count('(') != expression.count(')'):
            return False, "Unbalanced parentheses"

        # Extract operator names from expression
        op_pattern = re.compile(r'([A-Za-z_][A-Za-z0-9_]*)\s*\(')
        used_ops = [m.group(1) for m in op_pattern.finditer(expression)]

        if not used_ops:
            # Expression might just be a field name
            if expression.strip() not in [f.lower() for f in available_fields]:
                return False, "No operators found and expression is not a valid field"
            return True, ""

        # Check that all operators are known
        for op in used_ops:
            if not self.registry.has_operator(op):
                # Check if it might be a field name being called incorrectly
                return False, f"Unknown operator: {op}"

        # Check for common syntax errors
        if re.search(r'\(\s*\)', expression):
            return False, "Empty parentheses found"

        if re.search(r',\s*\)', expression):
            return False, "Missing argument before closing parenthesis"

        if re.search(r'\(\s*,', expression):
            return False, "Missing argument after opening parenthesis"

        return True, ""


class ExpressionParser:
    """
    Parse factor expressions into an AST for analysis and transformation.
    """

    def __init__(self, registry: Optional[UnifiedOperatorRegistry] = None):
        self.registry = registry or get_operator_registry()
        self.normalizer = ExpressionNormalizer(self.registry)

    def parse(self, expression: str, normalize: bool = True) -> Dict[str, Any]:
        """
        Parse an expression into a structured representation.

        Returns a dictionary with:
        - 'type': 'call', 'field', or 'constant'
        - 'operator': operator name (for 'call' type)
        - 'args': list of parsed arguments (for 'call' type)
        - 'value': field name or constant value
        """
        if normalize:
            expression = self.normalizer.normalize_expression(expression)

        expression = expression.strip()

        # Try to parse as a number
        try:
            val = float(expression)
            return {'type': 'constant', 'value': val}
        except ValueError:
            pass

        # Check if it's a simple field reference
        if re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', expression):
            return {'type': 'field', 'value': self.normalizer.normalize_field(expression)}

        # Parse as function call
        return self._parse_call(expression)

    def _parse_call(self, expression: str) -> Dict[str, Any]:
        """Parse a function call expression"""
        # Find the outermost function call
        match = re.match(r'^([A-Za-z_][A-Za-z0-9_]*)\s*\((.*)\)$', expression.strip(), re.DOTALL)
        if not match:
            raise ValueError(f"Invalid expression: {expression}")

        func_name = match.group(1)
        args_str = match.group(2)

        # Parse arguments
        args = self._split_args(args_str)
        parsed_args = [self.parse(arg, normalize=False) for arg in args]

        return {
            'type': 'call',
            'operator': self.normalizer.normalize_operator(func_name),
            'args': parsed_args
        }

    def _split_args(self, args_str: str) -> List[str]:
        """Split argument string respecting parentheses and quotes"""
        args = []
        current = []
        depth = 0

        for char in args_str:
            if char == ',' and depth == 0:
                args.append(''.join(current).strip())
                current = []
            else:
                if char == '(':
                    depth += 1
                elif char == ')':
                    depth -= 1
                current.append(char)

        if current:
            args.append(''.join(current).strip())

        return [a for a in args if a]

    def extract_operators(self, expression: str) -> List[str]:
        """Extract all operators used in an expression"""
        parsed = self.parse(expression)
        return self._collect_operators(parsed)

    def _collect_operators(self, parsed: Dict[str, Any]) -> List[str]:
        """Recursively collect all operators from a parsed expression"""
        operators = []
        if parsed['type'] == 'call':
            operators.append(parsed['operator'])
            for arg in parsed['args']:
                operators.extend(self._collect_operators(arg))
        return operators

    def extract_fields(self, expression: str) -> List[str]:
        """Extract all data fields used in an expression"""
        parsed = self.parse(expression)
        return list(set(self._collect_fields(parsed)))

    def _collect_fields(self, parsed: Dict[str, Any]) -> List[str]:
        """Recursively collect all fields from a parsed expression"""
        fields = []
        if parsed['type'] == 'field':
            fields.append(parsed['value'])
        elif parsed['type'] == 'call':
            for arg in parsed['args']:
                fields.extend(self._collect_fields(arg))
        return fields

    def get_expression_depth(self, expression: str) -> int:
        """Calculate the nesting depth of an expression"""
        parsed = self.parse(expression)
        return self._calculate_depth(parsed)

    def _calculate_depth(self, parsed: Dict[str, Any]) -> int:
        """Calculate depth of a parsed expression"""
        if parsed['type'] in ('field', 'constant'):
            return 0
        elif parsed['type'] == 'call':
            if not parsed['args']:
                return 1
            return 1 + max(self._calculate_depth(arg) for arg in parsed['args'])
        return 0

    def to_canonical_string(self, expression: str) -> str:
        """Convert an expression to canonical string form"""
        parsed = self.parse(expression)
        return self._to_string(parsed)

    def _to_string(self, parsed: Dict[str, Any]) -> str:
        """Convert parsed expression back to string"""
        if parsed['type'] == 'constant':
            return str(parsed['value'])
        elif parsed['type'] == 'field':
            return parsed['value']
        elif parsed['type'] == 'call':
            args_str = ', '.join(self._to_string(arg) for arg in parsed['args'])
            return f"{parsed['operator']}({args_str})"
        return ""
