#!/usr/bin/env python3
"""
SymQNet Molecular Optimization CLI - CORRECTED UNIVERSAL VERSION

Usage:
    symqnet-molopt --hamiltonian molecule.json --shots 1024 --output results.json

Output summary:
    • symqnet_results: coupling/field estimates, uncertainty stats, rollouts
    • hamiltonian_info / experimental_config / metadata
    • optional performance_analysis, universal_wrapper, validation blocks

Metadata slots:
    Uses the MetadataLayout contract (see ARCHITECTURE_CONTRACT.md) to encode
    qubit/basis/time selections, shot budget (optional), and posterior features.

Shot handling:
    The --shots value drives measurement sampling. If the checkpoint includes
    shots_encoding metadata, the shot budget is normalized as
    log1p(shots) / log1p(1_000_000) and injected into the metadata vector.
"""
import click
import json
import random
import secrets
import torch
import numpy as np
from pathlib import Path
import logging
from typing import Dict, List, Tuple, Optional
import sys
import os
import warnings

# Add current directory to path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Core imports (always available)
from hamiltonian_parser import HamiltonianParser
from measurement_simulator import MeasurementSimulator
from policy_engine import PolicyEngine, InferenceError
from bootstrap_estimator import BootstrapEstimator
from utils import setup_logging, validate_inputs, save_results

#GRACEFUL IMPORT HANDLING FOR Uni COMPONENTS
UNIVERSAL_MODE = False
DEFAULT_TRAINED_QUBITS = 10
try:
    from universal_wrapper import UniversalSymQNetWrapper
    from performance_estimator import PerformanceEstimator, get_performance_warning
    UNIVERSAL_MODE = True
    logger = logging.getLogger(__name__)
    logger.info("🌍 Universal mode available")
except ImportError as e:
    logger = logging.getLogger(__name__)
    logger.warning(f" Universal components not available: {e}")
    logger.warning(f"🔧Falling back to {DEFAULT_TRAINED_QUBITS}-qubit-only mode")
    
    # Fallback implementations
    class PerformanceEstimator:
        def __init__(self, optimal_qubits=10):
            self.optimal_qubits = optimal_qubits
        
        def estimate_performance(self, n_qubits):
            from dataclasses import dataclass
            from enum import Enum

            optimal_qubits = self.optimal_qubits
            
            class PerformanceLevel(Enum):
                OPTIMAL = "optimal"
                POOR = "poor"
            
            @dataclass
            class PerformanceReport:
                performance_factor: float = 1.0 if n_qubits == optimal_qubits else 0.0
                level: PerformanceLevel = (
                    PerformanceLevel.OPTIMAL if n_qubits == optimal_qubits else PerformanceLevel.POOR
                )
                recommendations: List[str] = None
                
                def __post_init__(self):
                    if self.recommendations is None:
                        if n_qubits != optimal_qubits:
                            self.recommendations = [f"Use exactly {optimal_qubits}-qubit systems for this version"]
                        else:
                            self.recommendations = []
            
            return PerformanceReport()
        
        def get_recommended_parameters(self, n_qubits):
            return {'shots': 1024, 'n_rollouts': 5}
    
    def get_performance_warning(n_qubits, optimal_qubits=10):
        if n_qubits != optimal_qubits:
            return f"Only {optimal_qubits}-qubit systems supported in fallback mode"
        return None

# Architecture imports
try:
    from architectures import (
        VariationalAutoencoder,
        FixedSymQNetWithEstimator,
        GraphEmbed,
        TemporalContextualAggregator,
        PolicyValueHead,
        SpinChainEnv,
        MetadataLayout,
    )
except ImportError as e:
    logger.warning(f" Architecture imports failed: {e}")
    # continue anyway so that we can let it fail later with clearer error
    MetadataLayout = None

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────────────────
#  Helper: locate model & VAE weights no matter where the CLI is launched
# ─────────────────────────────────────────────────────────────────────────────

def ensure_model_files(model_path: Optional[Path],
                       vae_path: Optional[Path],
                       auto_symlink: bool = True) -> Tuple[Path, Path]:
    """
    Resolve BEST_SYMQNET_MODEL_PPOV2.pth and vae_M10_f.pth with auto-download fallback.
    """
    MODEL_FILE = "BEST_SYMQNET_MODEL_PPOV2.pth"
    VAE_FILE = "vae_M10_f.pth"
    
    # GitHub URLs for auto-download
    MODEL_URL = "https://github.com/YTomar79/symqnet-molopt/raw/main/models/BEST_SYMQNET_MODEL_PPOV2.pth"
    VAE_URL = "https://github.com/YTomar79/symqnet-molopt/raw/main/models/vae_M10_f.pth"

    def _download_file(url: str, filepath: Path) -> Path:
        """Download file if missing."""
        if filepath.exists():
            return filepath
            
        print(f"⬇️ Downloading {filepath.name}...")
        filepath.parent.mkdir(parents=True, exist_ok=True)
        
        try:
            import urllib.request
            urllib.request.urlretrieve(url, filepath)
            print(f" Downloaded {filepath.name}")
            return filepath
        except Exception as e:
            raise click.ClickException(f"Failed to download {filepath.name}: {e}")

    def _resolve(
        user: Optional[Path],
        bundled_dir: Path,
        installed_dir: Path,
        rel_default: Path,
        url: str,
    ) -> Path:
        # 1. User-provided path
        if user is not None:
            up = Path(user).expanduser().resolve()
            if up.exists():
                return up
            raise click.ClickException(f" File not found: {up}")
        
        # 2. Wheel-bundled (next to this module)
        filename = MODEL_FILE if "SYMQNET" in url else VAE_FILE
        bp = bundled_dir / filename
        if bp.exists():
            return bp

        # 3. Installed data_files (sys.prefix/models)
        ip = installed_dir / filename
        if ip.exists():
            return ip
            
        # 4. Local ./models/
        if rel_default.exists():
            return rel_default
            
        # 5. Auto-download as last resort
        return _download_file(url, rel_default)

    # Find wheel-bundled models directory
    try:
        pkg_dir = Path(__file__).parent
        bundled_models = pkg_dir / "models"
    except:
        bundled_models = Path("nonexistent")

    # Local fallback paths
    cwd_models = Path.cwd() / "models"
    rel_model = cwd_models / MODEL_FILE
    rel_vae = cwd_models / VAE_FILE
    installed_models = Path(sys.prefix) / "models"

    try:
        final_model = _resolve(model_path, bundled_models, installed_models, rel_model, MODEL_URL)
        final_vae = _resolve(vae_path, bundled_models, installed_models, rel_vae, VAE_URL)
        return final_model, final_vae
    except Exception as e:
        raise click.ClickException(f" Model resolution failed: {e}")


def validate_checkpoint_metadata(model_path: Path) -> Dict[str, any]:
    """Validate checkpoint metadata schema and return normalized metadata."""
    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
    if not isinstance(checkpoint, dict):
        raise click.ClickException("Model checkpoint must be a dictionary-based state bundle.")
    if MetadataLayout is None:
        raise click.ClickException(
            "Architecture metadata tools unavailable; cannot validate checkpoint metadata."
        )

    required_keys = {
        "model_state_dict",
        "meta_dim",
        "shots_encoding",
        "n_qubits",
        "M_evo",
        "rollout_steps",
    }
    missing = sorted(required_keys - checkpoint.keys())
    if missing:
        raise click.ClickException(f"Checkpoint missing required metadata keys: {missing}")

    shots_encoding = checkpoint["shots_encoding"]
    if shots_encoding is not None and not isinstance(shots_encoding, dict):
        raise click.ClickException("Checkpoint shots_encoding must be a dict or None.")
    if isinstance(shots_encoding, dict) and "type" not in shots_encoding:
        raise click.ClickException("Checkpoint shots_encoding dict must include a 'type' field.")

    metadata_layout = MetadataLayout.from_problem(
        n_qubits=int(checkpoint["n_qubits"]),
        M_evo=int(checkpoint["M_evo"]),
    )
    meta_dim = int(checkpoint["meta_dim"])
    if meta_dim != metadata_layout.meta_dim:
        raise click.ClickException(
            "Checkpoint meta_dim mismatch: "
            f"{meta_dim} (checkpoint) vs {metadata_layout.meta_dim} (expected)."
        )

    checkpoint_format = checkpoint.get("checkpoint_format")
    checkpoint_version = checkpoint.get("checkpoint_version")
    if (checkpoint_format is None) != (checkpoint_version is None):
        raise click.ClickException(
            "Checkpoint must include both checkpoint_format and checkpoint_version, or neither."
        )

    return {
        "meta_dim": meta_dim,
        "shots_encoding": shots_encoding,
        "n_qubits": int(checkpoint["n_qubits"]),
        "M_evo": int(checkpoint["M_evo"]),
        "rollout_steps": int(checkpoint["rollout_steps"]),
        "checkpoint_format": checkpoint_format,
        "checkpoint_version": checkpoint_version,
    }


def find_hamiltonian_file(hamiltonian_path: Path) -> Path:
    """Find Hamiltonian file in examples or user directories"""
    
    # If absolute path or relative path that exists, use as-is
    if hamiltonian_path.is_absolute() or hamiltonian_path.exists():
        return hamiltonian_path
    
    # Check user directory first
    user_path = Path("user_hamiltonians") / hamiltonian_path
    if user_path.exists():
        logger.info(f"Found in user directory: {user_path}")
        return user_path
    
    # Check examples directory
    examples_path = Path("examples") / hamiltonian_path
    if examples_path.exists():
        logger.info(f"Found in examples directory: {examples_path}")
        return examples_path
    
    # Not found
    raise ValueError(
        f"Hamiltonian file not found: {hamiltonian_path}\n"
        f"Searched in:\n"
        f"  • Current directory\n"
        f"  • user_hamiltonians/\n"
        f"  • examples/\n\n"
        f"Use 'symqnet-add {hamiltonian_path}' to add your file to the system."
    )

def validate_hamiltonian_universal(hamiltonian_path: Path, optimal_qubits: int) -> Dict[str, any]:
    """Validate Hamiltonian with universal support or fallback"""
    
    try:
        with open(hamiltonian_path, 'r') as f:
            hamiltonian_data = json.load(f)
        
        n_qubits = hamiltonian_data.get('n_qubits', 0)
        
        if UNIVERSAL_MODE:
            # Universal validation - any qubit count ≥2
            if n_qubits < 2:
                raise ValueError(f"Minimum 2 qubits required, got {n_qubits}")
            
            # Performance guidance
            if n_qubits > 25:
                logger.warning(f"Large system ({n_qubits} qubits) may have very long runtime")
            
            logger.info(f" Validated: {n_qubits}-qubit Hamiltonian (Universal mode)")
            
        else:
            # Fallback mode - only trained qubits
            if n_qubits != optimal_qubits:
                raise ValueError(
                    f" FALLBACK MODE: Only {optimal_qubits}-qubit systems supported.\n"
                    f"   Your Hamiltonian: {n_qubits} qubits\n"
                    f"   Required: exactly {optimal_qubits} qubits\n\n"
                    f"💡 To enable universal support:\n"
                    f"   • Install universal components (universal_wrapper.py, performance_estimator.py)\n"
                    f"   • Or use a {optimal_qubits}-qubit molecular representation"
                )
            
            logger.info(f" Validated: {n_qubits}-qubit Hamiltonian (Fallback mode)")
        
        return hamiltonian_data
        
    except json.JSONDecodeError:
        raise ValueError(f"Invalid JSON file: {hamiltonian_path}")
    except FileNotFoundError:
        raise ValueError(f"Hamiltonian file not found: {hamiltonian_path}")

def run_optimization_universal(hamiltonian_data, model_path, vae_path, device, shots, n_rollouts, max_steps,
                               warn_performance=True, trained_qubits: int = DEFAULT_TRAINED_QUBITS):
    """
    Run optimization with working parameter extraction
    Uses proven rollout logic for both universal and fallback modes
    """
    
    original_qubits = hamiltonian_data['n_qubits']
    
    if UNIVERSAL_MODE:
        logger.info("Using Universal SymQNet with parameter extraction")
        
        # Step 1: Normalize Hamiltonian for trained-qubit processing
        universal_wrapper = UniversalSymQNetWrapper(
            trained_model_path=model_path,
            trained_vae_path=vae_path,
            device=device
        )

        universal_wrapper._validate_normalized_qubits(hamiltonian_data)
        
        # Get normalized hamiltonian
        normalized_hamiltonian = universal_wrapper._normalize_hamiltonian(hamiltonian_data)
        if normalized_hamiltonian.get("n_qubits") != universal_wrapper.trained_qubits:
            raise ValueError(
                "Normalized Hamiltonian qubit count mismatch: "
                f"expected {universal_wrapper.trained_qubits}, got {normalized_hamiltonian.get('n_qubits')}."
            )
        logger.info(
            f"🔄 Normalized {original_qubits}-qubit → "
            f"{universal_wrapper.trained_qubits}-qubit system"
        )
        
        # Step 2: Run rollouts through the wrapper to ensure consistent metadata layout
        try:
            normalized_results = universal_wrapper._run_optimization(
                normalized_hamiltonian, shots, n_rollouts, max_steps
            )
        except InferenceError as e:
            logger.error(f"Aborting optimization due to inference failure: {e}")
            logger.error("MAE will be invalid if inference is broken.")
            return {
                'symqnet_results': None,
                'rollout_results': [],
                'error': {
                    'type': type(e).__name__,
                    'message': str(e),
                },
                'aborted': True,
                'universal_metadata': {
                    'original_qubits': original_qubits,
                    'normalized_to': universal_wrapper.trained_qubits,
                    'normalization_applied': True,
                }
            }
        
        # Step 3: Denormalize results back to original system
        logger.info(
            f"🔄 Denormalizing results: {universal_wrapper.trained_qubits}-qubit → "
            f"{original_qubits}-qubit"
        )
        
        final_results = universal_wrapper._denormalize_results(normalized_results, original_qubits)
        
        # Add universal metadata
        final_results['universal_metadata'] = {
            'original_qubits': original_qubits,
            'normalized_to': universal_wrapper.trained_qubits,
            'expected_performance': universal_wrapper._calculate_performance_factor(original_qubits),
            'normalization_applied': True,
            'optimal_at': universal_wrapper.trained_qubits,
            'fixed_parameter_extraction': True
        }
        
        return final_results
    
    else:
        # Fallback to original implementation (this already works)
        logger.info(f"🔧 Using fallback {trained_qubits}-qubit implementation")
        
        # Initialize components directly
        policy = PolicyEngine(model_path, vae_path, device, shots=shots)
        simulator = MeasurementSimulator(hamiltonian_data, shots, device)
        estimator = BootstrapEstimator()
        
        # Run rollouts
        rollout_results = []
        try:
            for i in range(n_rollouts):
                logger.info(f"  Rollout {i+1}/{n_rollouts}")
                policy.reset()
                
                measurements = []
                parameter_estimates = []
                current_measurement = simulator.get_initial_measurement()
                
                for step in range(max_steps):
                    action_info = policy.get_action(current_measurement)
                    measurement_result = simulator.execute_measurement(
                        qubit_indices=action_info['qubits'],
                        pauli_operators=action_info['operators'],
                        evolution_time=action_info['time']
                    )
                    
                    measurements.append(measurement_result)
                    param_estimate = policy.get_parameter_estimate()
                    parameter_estimates.append(param_estimate)
                    measurement_mask = measurement_result['measurement_mask']
                    current_measurement[measurement_mask] = (
                        measurement_result['expectation_values'][measurement_mask]
                    )
                    
                    if step > 5 and policy.has_converged(parameter_estimates):
                        break
                
                rollout_results.append({
                    'rollout_id': i,
                    'final_estimate': parameter_estimates[-1] if parameter_estimates else None,
                    'convergence_step': step
                })
        except InferenceError as e:
            logger.error(f"Aborting optimization due to inference failure: {e}")
            logger.error("MAE will be invalid if inference is broken.")
            return {
                'symqnet_results': None,
                'rollout_results': rollout_results,
                'error': {
                    'type': type(e).__name__,
                    'message': str(e),
                },
                'aborted': True,
                'fallback_mode': True
            }
        
        # Bootstrap analysis
        logger.info(" Computing confidence intervals...")
        bootstrap_results = estimator.compute_intervals(rollout_results)
        
        return {
            'symqnet_results': bootstrap_results,
            'rollout_results': rollout_results,
            'fallback_mode': True
        }

def print_performance_info(n_qubits: int, performance_estimator: PerformanceEstimator, optimal_qubits: int):
    """Print performance information and recommendations"""
    
    print("\n" + "="*60)
    if UNIVERSAL_MODE:
        print(" UNIVERSAL SYMQNET PERFORMANCE ANALYSIS")
    else:
        print(" FALLBACK MODE ANALYSIS")

    
    report = performance_estimator.estimate_performance(n_qubits)
    print(f" Optimal training size: {optimal_qubits} qubits")
    
    

def print_summary(results: Dict, n_qubits: int, performance_factor: float,
                  requested_shots: int, shots_used: int, optimal_qubits: int):
    """Print a formatted summary of results with performance context."""
    

    print(" SYMQNET MOLECULAR OPTIMIZATION RESULTS")

    
    print(f" System: {n_qubits} qubits")
    print(f" Optimal training size: {optimal_qubits} qubits")
    print(f" Shots: requested={requested_shots}, used={shots_used}")

    

    
    # Extract results from nested structure if needed
    if 'symqnet_results' in results:
        symqnet_results = results['symqnet_results']
    else:
        symqnet_results = results
    
    if 'coupling_parameters' in symqnet_results:
        coupling_params = symqnet_results['coupling_parameters']
        if coupling_params and len(coupling_params) > 0:
            if isinstance(coupling_params[0], tuple):
                # Tuple format: (mean, ci_low, ci_high)
                coupling_count = len(coupling_params)
                print(f"\n COUPLING PARAMETERS ({coupling_count} estimated):")
                for i, (mean, ci_low, ci_high) in enumerate(coupling_params):
                    uncertainty = (ci_high - ci_low) / 2
                    print(f"  J_{i}: {mean:8.6f} ± {uncertainty:.6f} [{ci_low:.6f}, {ci_high:.6f}]")
            else:
                # Dict format
                coupling_count = len(coupling_params)
                print(f"\n COUPLING PARAMETERS ({coupling_count} estimated):")
                for param in coupling_params:
                    i = param.get('index', 0)
                    mean = param.get('mean', 0)
                    uncertainty = param.get('uncertainty', 0)
                    ci = param.get('confidence_interval', [0, 0])
                    print(f"  J_{i}: {mean:8.6f} ± {uncertainty:.6f} [{ci[0]:.6f}, {ci[1]:.6f}]")
    
    if 'field_parameters' in symqnet_results:
        field_params = symqnet_results['field_parameters']
        if field_params and len(field_params) > 0:
            if isinstance(field_params[0], tuple):
                # Tuple format: (mean, ci_low, ci_high)
                field_count = len(field_params)
                print(f"\n FIELD PARAMETERS ({field_count} estimated):")
                for i, (mean, ci_low, ci_high) in enumerate(field_params):
                    uncertainty = (ci_high - ci_low) / 2
                    print(f"  h_{i}: {mean:8.6f} ± {uncertainty:.6f} [{ci_low:.6f}, {ci_high:.6f}]")
            else:
                # Dict format
                field_count = len(field_params)
                print(f"\n FIELD PARAMETERS ({field_count} estimated):")
                for param in field_params:
                    i = param.get('index', 0)
                    mean = param.get('mean', 0)
                    uncertainty = param.get('uncertainty', 0)
                    ci = param.get('confidence_interval', [0, 0])
                    print(f"  h_{i}: {mean:8.6f} ± {uncertainty:.6f} [{ci[0]:.6f}, {ci[1]:.6f}]")
    
    if 'total_uncertainty' in symqnet_results:
        print(f"\n📏 Total Parameter Uncertainty: {symqnet_results['total_uncertainty']:.6f}")
    
    if 'n_rollouts' in symqnet_results:
        print(f"🔄 Rollouts Completed: {symqnet_results['n_rollouts']}")
    

def get_recommended_params_for_system(n_qubits: int, 
                                     user_shots: int, 
                                     user_rollouts: int,
                                     performance_estimator: PerformanceEstimator) -> Dict[str, int]:
    """Get recommended parameters based on system size and user preferences"""
    
    # Get performance-based recommendations
    recommended = performance_estimator.get_recommended_parameters(n_qubits)
    
    # Respect user choices but warn if they seem too low
    final_shots = max(user_shots, int(recommended['shots'] * 0.8))
    final_rollouts = max(user_rollouts, int(recommended['n_rollouts'] * 0.8))
    
    if user_shots < recommended['shots'] or user_rollouts < recommended['n_rollouts']:
        logger.warning(f"Parameters may be too low for {n_qubits}-qubit system. "
                      f"Recommended: shots={recommended['shots']}, rollouts={recommended['n_rollouts']}")
    
    return {
        'shots': final_shots,
        'n_rollouts': final_rollouts,
        'recommended_shots': recommended['shots'],
        'recommended_rollouts': recommended['n_rollouts']
    }

@click.command()
@click.option('--hamiltonian', '-h', 
              type=click.Path(path_type=Path),
              required=True,
              help='Path to molecular Hamiltonian JSON file')
@click.option('--shots', '-s', 
              type=int, 
              default=1024,
              help='Number of measurement shots per observable (default: 1024)')
@click.option('--output', '-o', 
              type=click.Path(path_type=Path),
              required=True,
              help='Output JSON file for estimates and uncertainties')
@click.option('--model-path', '-m',
              type=click.Path(path_type=Path),
              default=None,
              help='Path to trained SymQNet model (default: models/BEST_SYMQNET_MODEL_PPOV2.pth). '
                   'Checkpoint must include meta_dim and shots_encoding metadata.')
@click.option('--vae-path', '-v',
              type=click.Path(path_type=Path),
              default=None,
              help='Path to pre-trained VAE (default: models/vae_M10_f.pth)')
@click.option('--max-steps', '-t',
              type=int,
              default=50,
              help='Maximum measurement steps per rollout (default: 50)')
@click.option('--n-rollouts', '-r',
              type=int,
              default=10,
              help='Number of policy rollouts for averaging (default: 10)')
@click.option('--confidence', '-c',
              type=float,
              default=0.95,
              help='Confidence level for uncertainty intervals (default: 0.95)')
@click.option('--device', '-d',
              type=click.Choice(['cpu', 'cuda', 'auto']),
              default='auto',
              help='Compute device (default: auto)')
@click.option('--seed', 
              type=int,
              default=None,
              help='Random seed for reproducibility (default: None for randomized runs)')
@click.option('--verbose', '-V',
              is_flag=True,
              help='Enable verbose logging')
@click.option('--no-performance-warnings',
              is_flag=True,
              help='Disable performance degradation warnings')
@click.option('--show-performance-analysis',
              is_flag=True,
              help='Show detailed performance analysis')
def main(hamiltonian: Path, shots: int, output: Path, model_path: Optional[Path], 
         vae_path: Optional[Path], max_steps: int, n_rollouts: int, confidence: float,
         device: str, seed: Optional[int], verbose: bool, no_performance_warnings: bool,
         show_performance_analysis: bool):
    """
    SymQNet Molecular Optimization CLI
    
     support (any qubit count) with WORKING parameter extraction
     Fallback to trained-qubit-only mode if universal components unavailable
     Checkpoint metadata expectations: meta_dim and optional shots_encoding

    Metadata slots:
        See ARCHITECTURE_CONTRACT.md for the full metadata slot layout.

    Shot handling:
        If the checkpoint declares shots_encoding, the shot budget is injected into the
        metadata vector using log1p(shots) / log1p(1_000_000). Otherwise the shots slot is 0.

    Migration note:
        Older checkpoints without model_state_dict/meta_dim/shots_encoding/n_qubits/M_evo/rollout_steps
        must be re-exported with the current checkpoint_format/version fields.

    Examples:
        symqnet-molopt --hamiltonian H2O_10q.json --output results.json
        symqnet-molopt --hamiltonian molecule.json --output results.json --shots 2048
    """
    
    # Setup logging first
    if verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    setup_logging(verbose)
    
    trained_qubits = DEFAULT_TRAINED_QUBITS

    # Display mode information
    if UNIVERSAL_MODE:
        logger.info(" Universal SymQNet mode enabled ")

    else:
        logger.info(" Fallback mode active")
    
    #  FIXED: Handle model file paths gracefully
    try:
        model_path, vae_path = ensure_model_files(model_path, vae_path)
        logger.info(f" Using model: {model_path}")
        logger.info(f" Using VAE: {vae_path}")
        checkpoint_metadata = validate_checkpoint_metadata(model_path)
        shots_encoding = checkpoint_metadata["shots_encoding"]
        shots_encoding_type = shots_encoding.get("type") if isinstance(shots_encoding, dict) else None
        logger.info(
            " Checkpoint metadata: meta_dim=%s, shots_encoding=%s, n_qubits=%s, M_evo=%s, rollout_steps=%s",
            checkpoint_metadata["meta_dim"],
            shots_encoding_type or "none",
            checkpoint_metadata["n_qubits"],
            checkpoint_metadata["M_evo"],
            checkpoint_metadata["rollout_steps"],
        )
        trained_qubits = int(checkpoint_metadata["n_qubits"])
        if shots_encoding_type is None:
            logger.info(" Shots encoding disabled; model metadata will not include shot count conditioning.")
        else:
            logger.info(" Shots encoding enabled; model expects shot count conditioning via metadata.")
        if not UNIVERSAL_MODE:
            logger.info(f"  Supports exactly {trained_qubits}-qubit systems only")
    except click.ClickException:
        raise  # Re-raise click exceptions as-is
    except Exception as e:
        raise click.ClickException(f" Model setup failed: {e}")
    
    # Find hamiltonian file early
    try:
        hamiltonian_path = find_hamiltonian_file(hamiltonian)
    except ValueError as e:
        raise click.ClickException(str(e))
    
    # Validate with mode-appropriate constraints
    try:
        hamiltonian_data = validate_hamiltonian_universal(hamiltonian_path, optimal_qubits=trained_qubits)
        n_qubits = hamiltonian_data['n_qubits']
        
    except ValueError as e:
        raise click.ClickException(str(e))
    
    # Set device
    if device == 'auto':
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    device = torch.device(device)
    logger.info(f"Using device: {device}")
    
    # Set random seeds
    if seed is None:
        seed = secrets.randbits(32)
        logger.info(f"No seed provided. Generated random seed: {seed}")

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    
    # Initialize performance estimator
    performance_estimator = PerformanceEstimator(optimal_qubits=trained_qubits)
    
    # Performance analysis and warnings
    if not no_performance_warnings:
        warning = get_performance_warning(n_qubits, optimal_qubits=trained_qubits)
        if warning:
            logger.warning(warning)
    
    if show_performance_analysis:
        print_performance_info(n_qubits, performance_estimator, optimal_qubits=trained_qubits)
    
    # Get recommended parameters based on system size
    requested_shots = shots
    param_recommendations = get_recommended_params_for_system(
        n_qubits, shots, n_rollouts, performance_estimator
    )
    
    # Use recommended parameters if significantly different
    if param_recommendations['shots'] > shots:
        logger.info(f"Increasing shots: {shots} → {param_recommendations['shots']} (recommended for {n_qubits} qubits)")
        shots = param_recommendations['shots']
    
    if param_recommendations['n_rollouts'] > n_rollouts:
        logger.info(f"Increasing rollouts: {n_rollouts} → {param_recommendations['n_rollouts']} (recommended for {n_qubits} qubits)")
        n_rollouts = param_recommendations['n_rollouts']
    shots_used = shots
    logger.info(f"Shots summary: requested={requested_shots}, used={shots_used}")
    
    try:
        # Validate inputs
        validate_inputs(hamiltonian_path, shots, confidence, max_steps, n_rollouts)
        
        # 1. Parse Hamiltonian
        logger.info("🔍 Parsing molecular Hamiltonian...")
        parser = HamiltonianParser()
        hamiltonian_data = parser.load_hamiltonian(hamiltonian_path, warn_performance=(not no_performance_warnings))
        logger.info(f"Loaded {hamiltonian_data['n_qubits']}-qubit Hamiltonian "
                   f"with {len(hamiltonian_data['pauli_terms'])} terms")
        
        # 2. Run Parameter Estimation 
        performance_report = performance_estimator.estimate_performance(n_qubits)
        logger.info(f" Expected performance: {performance_report.performance_factor:.1%} of optimal")
        
        logger.info(f" Running parameter estimation...")
        logger.info(f" Configuration: {shots} shots, {n_rollouts} rollouts, {max_steps} max steps")
        
        final_results = run_optimization_universal(
            hamiltonian_data=hamiltonian_data,
            model_path=model_path,
            vae_path=vae_path,
            device=device,
            shots=shots,
            n_rollouts=n_rollouts,
            max_steps=max_steps,
            warn_performance=(not no_performance_warnings),
            trained_qubits=trained_qubits
        )
        
        # 3. Save Results
        logger.info(f" Saving results to {output}")
        save_results(
            results=final_results,
            hamiltonian_data=hamiltonian_data,
            config={
                'shots': shots,
                'requested_shots': requested_shots,
                'shots_used': shots_used,
                'max_steps': max_steps,
                'n_rollouts': n_rollouts,
                'confidence': confidence,
                'seed': seed,
                'mode': 'universal_fixed' if UNIVERSAL_MODE else 'fallback',
                'performance_metadata': {
                    'expected_performance': performance_report.performance_factor,
                    'performance_level': performance_report.level.value,
                    'optimal_qubits': trained_qubits,
                    'universal_mode': UNIVERSAL_MODE,
                    'parameter_extraction_fixed': True
                }
            },
            output_path=output
        )
        
        # Print summary with performance context
        print_summary(final_results, n_qubits, performance_report.performance_factor,
                      requested_shots, shots_used, optimal_qubits=trained_qubits)
        
        # Final performance note
        success_msg = " Universal molecular optimization completed successfully!" if UNIVERSAL_MODE else " Molecular optimization completed successfully!"
        logger.info(success_msg)
        
    except Exception as e:
        logger.error(f" Error: {e}")
        if verbose:
            import traceback
            traceback.print_exc()
        raise click.ClickException(str(e))

if __name__ == '__main__':
    main()
