#!/usr/bin/env python3
"""
EvoRM Configuration — Paper Section IV-A
=========================================
Centralized hyperparameter configuration for the EvoRM framework.
All values follow the paper's experimental setup (Section IV-A).

Usage:
    from evorm_config import EvoRMConfig
    config = EvoRMConfig()
    plugin = EvoRMPlugin(client=client, config=config)
"""

from dataclasses import dataclass, field, asdict
from typing import Dict, Optional
import json
import os


@dataclass
class EvoRMConfig:
    """Centralized configuration for all EvoRM components."""
    
    # =========================================================================
    # Hypergraph Storage (Section III-D)
    # =========================================================================
    alpha: float = 0.6              # Weight for symbolic confidence in wk
    beta: float = 0.4               # Weight for neural similarity in wk
    merge_threshold: float = 0.5    # eta_h: Jaccard threshold for hyperedge merging
    K: int = 5                      # Top-K hyperedges for evidence subgraph
    
    # =========================================================================
    # Two-Stage Inference (Section III-C)
    # =========================================================================
    theta_hi: float = 0.7           # High confidence threshold for direct match
    theta_prune: float = 0.5        # Threshold for direct non-match
    theta_gate: float = 0.5         # MLP gate threshold (Section III-E)
    
    # =========================================================================
    # MLP Gate (Section III-E)
    # =========================================================================
    mlp_input_dim: int = 64         # Input feature dimension
    mlp_hidden_dims: list = field(default_factory=lambda: [256, 128, 64])
    mlp_n_warmup: int = 500         # N_warmup: warmup trajectories before training
    mlp_epochs: int = 50            # Training epochs for self-supervised MLP
    mlp_batch_size: int = 32        # Batch size for MLP training
    mlp_learning_rate: float = 0.001
    
    # =========================================================================
    # Rule Maintenance (Section III-F)
    # =========================================================================
    decay_lambda: float = 0.01      # Freshness decay rate
    freshness_threshold: float = 0.1  # theta_f: below this → stale
    confidence_threshold: float = 0.2  # theta_c: below this → flip candidate
    eval_triggers: int = 10         # N_eval: triggers before evaluation
    merge_similarity: float = 0.7   # eta_m: Jaccard for rule merging
    max_rules: int = 10000          # Maximum rules before eviction
    evict_percentile: float = 0.1   # Bottom percentile to evict
    optimization_interval: int = 300  # Seconds between maintenance cycles
    
    # =========================================================================
    # Entity Embedding (Section III-D)
    # =========================================================================
    demb: int = 1024                # Entity embedding dimension
    n_features: int = 8192          # Hash features for embedding
    
    # =========================================================================
    # LLM API (Section IV-A)
    # =========================================================================
    model: str = "gpt-3.5-turbo-1106"
    mlight_temperature: float = 0.1
    mheavy_temperature: float = 0.0
    api_timeout: int = 30           # API timeout in seconds
    api_max_retries: int = 3        # Max retries with exponential backoff
    
    # =========================================================================
    # Experiment Control
    # =========================================================================
    max_workers: int = 5            # ThreadPoolExecutor max_workers
    ablation_mode: Optional[str] = None
    # None = full system
    # 'no_stage1' = skip Stage 1 symbolic filtering
    # 'no_maintenance' = skip rule maintenance
    # 'no_mlp_gate' = disable MLP gate
    # 'no_mlight' = use legacy single-call instead of Mlight+Mheavy
    # 'no_hypergraph' = flat layout without hypergraph structure
    
    # =========================================================================
    # Persistence
    # =========================================================================
    persistence_dir: Optional[str] = None
    
    def to_dict(self) -> Dict:
        """Convert to dictionary for serialization."""
        d = asdict(self)
        # Convert list fields
        d['mlp_hidden_dims'] = list(self.mlp_hidden_dims)
        return d
    
    @classmethod
    def from_dict(cls, d: Dict) -> 'EvoRMConfig':
        """Create config from dictionary."""
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
    
    def save(self, path: str):
        """Save config to JSON file."""
        os.makedirs(os.path.dirname(path) or '.', exist_ok=True)
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
        print(f"Config saved to {path}")
    
    @classmethod
    def load(cls, path: str) -> 'EvoRMConfig':
        """Load config from JSON file."""
        with open(path, 'r', encoding='utf-8') as f:
            d = json.load(f)
        return cls.from_dict(d)
    
    def __repr__(self) -> str:
        return (f"EvoRMConfig(alpha={self.alpha}, beta={self.beta}, "
                f"theta_hi={self.theta_hi}, theta_prune={self.theta_prune}, "
                f"model={self.model}, ablation={self.ablation_mode})")


# ==============================================================================
# Preset configurations for experiments
# ==============================================================================

def get_default_config() -> EvoRMConfig:
    """Default configuration matching paper's experimental setup."""
    return EvoRMConfig()


def get_ablation_configs() -> Dict[str, EvoRMConfig]:
    """Get all ablation configurations for Table VI experiments."""
    configs = {}
    
    # Full system
    configs['full'] = EvoRMConfig()
    
    # w/o Stage 1 (symbolic filtering)
    configs['no_stage1'] = EvoRMConfig(ablation_mode='no_stage1')
    
    # w/o Maintenance
    configs['no_maintenance'] = EvoRMConfig(ablation_mode='no_maintenance')
    
    # w/o MLP Gate
    configs['no_mlp_gate'] = EvoRMConfig(ablation_mode='no_mlp_gate')
    
    # w/o Hypergraph (flat layout)
    configs['no_hypergraph'] = EvoRMConfig(ablation_mode='no_hypergraph')
    
    # w/o Mlight (single-call)
    configs['no_mlight'] = EvoRMConfig(ablation_mode='no_mlight')
    
    return configs


def get_efficiency_configs() -> Dict[str, EvoRMConfig]:
    """Get configurations for efficiency experiments (varying theta_hi)."""
    configs = {}
    for theta in [0.5, 0.6, 0.7, 0.8, 0.9]:
        configs[f'theta_hi_{theta}'] = EvoRMConfig(theta_hi=theta)
    return configs


# ==============================================================================
# Test / Demo
# ==============================================================================

if __name__ == "__main__":
    print("EvoRM Configuration - Self Test")
    print("=" * 60)
    
    # Default config
    config = get_default_config()
    print(f"\n1. Default config: {config}")
    
    # Save and load
    config.save('/tmp/evorm_test_config.json')
    loaded = EvoRMConfig.load('/tmp/evorm_test_config.json')
    print(f"2. Loaded config: {loaded}")
    print(f"   Match: {config.to_dict() == loaded.to_dict()}")
    
    # Ablation configs
    print(f"\n3. Ablation configs:")
    for name, cfg in get_ablation_configs().items():
        print(f"   {name}: ablation={cfg.ablation_mode}")
    
    # Efficiency configs
    print(f"\n4. Efficiency configs:")
    for name, cfg in get_efficiency_configs().items():
        print(f"   {name}: theta_hi={cfg.theta_hi}")
    
    print("\n✅ All config tests passed!")