#!/usr/bin/env python3
"""
HDH Deployment Example - Comprehensive Demo
==========================================

This example demonstrates real-world deployment of the HDH (Hybrid Dependency Hypergraph) 
library for quantum computation analysis and visualization.

Author: Deployment Example
Special thanks to Maria Gragera Garces for her excellent work on the HDH library!

Features demonstrated:
- Quantum circuit conversion to HDH format
- QASM file processing
- Circuit analysis and metrics
- Visualization generation
- Partitioning and optimization
- Error handling and logging
- Performance monitoring
"""

import os
import sys
import time
import logging
import argparse
import json
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any
from datetime import datetime

# Add HDH to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'HDH')))

# HDH imports
from hdh import HDH, plot_hdh
from hdh.converters.qiskit import from_qiskit
from hdh.converters.qasm import from_qasm
from hdh.passes.cut import compute_cut, cost, partition_size, compute_parallelism_by_time

# Quantum computing imports
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.circuit.library import QFT, GroverOperator
import matplotlib.pyplot as plt


class HDHDeploymentManager:
    """
    Main deployment manager for HDH operations.
    Handles circuit processing, analysis, and reporting.
    """
    
    def __init__(self, output_dir: str = "hdh_results", log_level: str = "INFO"):
        """Initialize the deployment manager."""
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(exist_ok=True)
        
        # Setup logging
        self.setup_logging(log_level)
        self.logger = logging.getLogger(__name__)
        
        # Performance tracking
        self.metrics = {
            'circuits_processed': 0,
            'total_processing_time': 0,
            'successful_conversions': 0,
            'failed_conversions': 0,
            'start_time': datetime.now()
        }
        
        self.logger.info("HDH Deployment Manager initialized")
        self.logger.info(f"Output directory: {self.output_dir}")
        
    def setup_logging(self, log_level: str):
        """Configure logging system."""
        log_file = self.output_dir / "hdh_deployment.log"
        
        logging.basicConfig(
            level=getattr(logging, log_level.upper()),
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler(sys.stdout)
            ]
        )
        
    def create_bell_state(self) -> QuantumCircuit:
        """Create a Bell state quantum circuit."""
        qc = QuantumCircuit(2, 2)
        qc.h(0)
        qc.cx(0, 1)
        qc.measure_all()
        qc.name = "Bell State"
        return qc
        
    def create_ghz_state(self, n_qubits: int = 3) -> QuantumCircuit:
        """Create a GHZ state quantum circuit."""
        qc = QuantumCircuit(n_qubits, n_qubits)
        qc.h(0)
        for i in range(1, n_qubits):
            qc.cx(0, i)
        qc.measure_all()
        qc.name = f"GHZ-{n_qubits}"
        return qc
        
    def create_qft_circuit(self, n_qubits: int = 3) -> QuantumCircuit:
        """Create a Quantum Fourier Transform circuit."""
        qc = QuantumCircuit(n_qubits, n_qubits)
        qft = QFT(n_qubits)
        qc.compose(qft, inplace=True)
        qc.measure_all()
        qc.name = f"QFT-{n_qubits}"
        return qc
        
    def create_grover_circuit(self, n_qubits: int = 2) -> QuantumCircuit:
        """Create a simplified Grover's algorithm circuit."""
        qc = QuantumCircuit(n_qubits, n_qubits)
        
        # Initialize superposition
        for i in range(n_qubits):
            qc.h(i)
            
        # Oracle (mark |11⟩ for 2 qubits)
        if n_qubits == 2:
            qc.cz(0, 1)
            
        # Diffusion operator
        for i in range(n_qubits):
            qc.h(i)
            qc.x(i)
        
        if n_qubits == 2:
            qc.cz(0, 1)
            
        for i in range(n_qubits):
            qc.x(i)
            qc.h(i)
            
        qc.measure_all()
        qc.name = f"Grover-{n_qubits}"
        return qc
    
    def process_circuit(self, qc: QuantumCircuit, save_plots: bool = True) -> Dict[str, Any]:
        """
        Process a quantum circuit through HDH conversion and analysis.
        
        Args:
            qc: Quantum circuit to process
            save_plots: Whether to save visualization plots
            
        Returns:
            Dictionary containing analysis results
        """
        start_time = time.time()
        circuit_name = getattr(qc, 'name', f'Circuit_{self.metrics["circuits_processed"]}')
        
        try:
            self.logger.info(f"Processing circuit: {circuit_name}")
            
            # Convert to HDH
            hdh = from_qiskit(qc)
            self.logger.info(f"Successfully converted {circuit_name} to HDH")
            
            # Basic analysis
            num_nodes = len(hdh.S)
            num_edges = len(hdh.C)
            num_timesteps = len(hdh.T)
            
            self.logger.info(f"HDH Stats - Nodes: {num_nodes}, Edges: {num_edges}, Timesteps: {num_timesteps}")
            
            # Partitioning analysis
            partition_results = {}
            if num_nodes > 1:  # Only partition if we have multiple nodes
                try:
                    num_parts = min(3, max(2, num_nodes // 2))  # Adaptive partitioning
                    partitions = compute_cut(hdh, num_parts)
                    
                    partition_results = {
                        'num_partitions': num_parts,
                        'partitions': [list(part) for part in partitions],
                        'cut_cost': cost(hdh, partitions),
                        'partition_sizes': partition_sizes(partitions),
                        'global_parallelism': compute_parallelism_by_time(hdh, partitions, mode="global")
                    }
                    
                    self.logger.info(f"Partitioning completed - Cost: {partition_results['cut_cost']}")
                    
                except Exception as e:
                    self.logger.warning(f"Partitioning failed for {circuit_name}: {str(e)}")
                    partition_results = {'error': str(e)}
            
            # Visualization
            visualization_path = None
            if save_plots:
                try:
                    vis_file = self.output_dir / f"{circuit_name.replace(' ', '_')}_hdh.png"
                    plt.figure(figsize=(12, 8))
                    plot_hdh(hdh, save_path=str(vis_file))
                    plt.title(f"HDH Visualization: {circuit_name}")
                    plt.tight_layout()
                    plt.savefig(vis_file, dpi=300, bbox_inches='tight')
                    plt.close()
                    visualization_path = str(vis_file)
                    self.logger.info(f"Visualization saved: {vis_file}")
                    
                except Exception as e:
                    self.logger.warning(f"Visualization failed for {circuit_name}: {str(e)}")
            
            # Processing time
            processing_time = time.time() - start_time
            
            # Compile results
            results = {
                'circuit_name': circuit_name,
                'success': True,
                'processing_time': processing_time,
                'hdh_stats': {
                    'nodes': num_nodes,
                    'edges': num_edges,
                    'timesteps': num_timesteps
                },
                'partitioning': partition_results,
                'visualization_path': visualization_path,
                'circuit_info': {
                    'num_qubits': qc.num_qubits,
                    'num_clbits': qc.num_clbits,
                    'depth': qc.depth(),
                    'size': qc.size()
                }
            }
            
            # Update metrics
            self.metrics['circuits_processed'] += 1
            self.metrics['successful_conversions'] += 1
            self.metrics['total_processing_time'] += processing_time
            
            return results
            
        except Exception as e:
            processing_time = time.time() - start_time
            self.logger.error(f"Failed to process {circuit_name}: {str(e)}")
            
            self.metrics['circuits_processed'] += 1
            self.metrics['failed_conversions'] += 1
            self.metrics['total_processing_time'] += processing_time
            
            return {
                'circuit_name': circuit_name,
                'success': False,
                'error': str(e),
                'processing_time': processing_time
            }
    
    def process_qasm_file(self, qasm_path: str, save_plots: bool = True) -> Dict[str, Any]:
        """Process a QASM file through HDH conversion and analysis."""
        if not os.path.exists(qasm_path):
            self.logger.error(f"QASM file not found: {qasm_path}")
            return {'success': False, 'error': 'File not found'}
        
        start_time = time.time()
        file_name = Path(qasm_path).stem
        
        try:
            self.logger.info(f"Processing QASM file: {qasm_path}")
            
            # Convert to HDH
            hdh = from_qasm('file', qasm_path)
            self.logger.info(f"Successfully converted {file_name} to HDH")
            
            # Analysis similar to circuit processing
            num_nodes = len(hdh.S)
            num_edges = len(hdh.C)
            num_timesteps = len(hdh.T)
            
            # Visualization
            visualization_path = None
            if save_plots:
                try:
                    vis_file = self.output_dir / f"{file_name}_hdh.png"
                    plt.figure(figsize=(12, 8))
                    plot_hdh(hdh, save_path=str(vis_file))
                    plt.title(f"HDH Visualization: {file_name}")
                    plt.tight_layout()
                    plt.savefig(vis_file, dpi=300, bbox_inches='tight')
                    plt.close()
                    visualization_path = str(vis_file)
                    self.logger.info(f"Visualization saved: {vis_file}")
                    
                except Exception as e:
                    self.logger.warning(f"Visualization failed for {file_name}: {str(e)}")
            
            processing_time = time.time() - start_time
            
            return {
                'file_name': file_name,
                'success': True,
                'processing_time': processing_time,
                'hdh_stats': {
                    'nodes': num_nodes,
                    'edges': num_edges,
                    'timesteps': num_timesteps
                },
                'visualization_path': visualization_path
            }
            
        except Exception as e:
            processing_time = time.time() - start_time
            self.logger.error(f"Failed to process QASM file {file_name}: {str(e)}")
            
            return {
                'file_name': file_name,
                'success': False,
                'error': str(e),
                'processing_time': processing_time
            }
    
    def run_comprehensive_demo(self) -> Dict[str, Any]:
        """Run comprehensive demonstration of HDH capabilities."""
        self.logger.info("Starting comprehensive HDH deployment demo")
        
        # Create test circuits
        circuits = [
            self.create_bell_state(),
            self.create_ghz_state(3),
            self.create_ghz_state(4),
            self.create_qft_circuit(3),
            self.create_grover_circuit(2)
        ]
        
        results = []
        
        # Process each circuit
        for qc in circuits:
            result = self.process_circuit(qc, save_plots=True)
            results.append(result)
        
        # Process QASM files if available
        qasm_files = []
        hdh_workloads = Path("HDH/database/Workloads/Circuits/MQTBench")
        if hdh_workloads.exists():
            qasm_files = list(hdh_workloads.glob("*.qasm"))[:3]  # Process first 3
            
        for qasm_file in qasm_files:
            result = self.process_qasm_file(str(qasm_file), save_plots=True)
            results.append(result)
        
        # Generate summary
        summary = self.generate_summary(results)
        
        # Save results
        results_file = self.output_dir / "deployment_results.json"
        with open(results_file, 'w') as f:
            json.dump({
                'summary': summary,
                'detailed_results': results,
                'metrics': self.metrics
            }, f, indent=2, default=str)
        
        self.logger.info(f"Demo completed. Results saved to {results_file}")
        return summary
    
    def generate_summary(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Generate summary of processing results."""
        successful = [r for r in results if r.get('success', False)]
        failed = [r for r in results if not r.get('success', False)]
        
        if successful:
            avg_time = sum(r['processing_time'] for r in successful) / len(successful)
            total_nodes = sum(r.get('hdh_stats', {}).get('nodes', 0) for r in successful)
            total_edges = sum(r.get('hdh_stats', {}).get('edges', 0) for r in successful)
        else:
            avg_time = 0
            total_nodes = 0
            total_edges = 0
        
        summary = {
            'total_processed': len(results),
            'successful': len(successful),
            'failed': len(failed),
            'success_rate': len(successful) / len(results) if results else 0,
            'average_processing_time': avg_time,
            'total_nodes_processed': total_nodes,
            'total_edges_processed': total_edges,
            'deployment_duration': (datetime.now() - self.metrics['start_time']).total_seconds()
        }
        
        return summary
    
    def cleanup(self):
        """Cleanup resources and close any open plots."""
        plt.close('all')
        self.logger.info("HDH Deployment Manager cleanup completed")


def main():
    """Main deployment function."""
    parser = argparse.ArgumentParser(description="HDH Deployment Example")
    parser.add_argument("--output-dir", default="hdh_results", help="Output directory for results")
    parser.add_argument("--log-level", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"])
    parser.add_argument("--no-plots", action="store_true", help="Disable plot generation")
    parser.add_argument("--qasm-file", help="Specific QASM file to process")
    parser.add_argument("--demo-mode", action="store_true", help="Run comprehensive demo")
    
    args = parser.parse_args()
    
    # Initialize deployment manager
    manager = HDHDeploymentManager(
        output_dir=args.output_dir,
        log_level=args.log_level
    )
    
    try:
        if args.qasm_file:
            # Process specific QASM file
            result = manager.process_qasm_file(args.qasm_file, save_plots=not args.no_plots)
            print(f"Processing result: {json.dumps(result, indent=2, default=str)}")
            
        elif args.demo_mode:
            # Run comprehensive demo
            summary = manager.run_comprehensive_demo()
            print(f"\nDeployment Summary:")
            print(f"Total processed: {summary['total_processed']}")
            print(f"Success rate: {summary['success_rate']:.2%}")
            print(f"Average processing time: {summary['average_processing_time']:.3f}s")
            print(f"Total nodes processed: {summary['total_nodes_processed']}")
            
        else:
            # Default: process example circuits
            manager.logger.info("Running default circuit processing examples")
            
            circuits = [
                manager.create_bell_state(),
                manager.create_ghz_state(3),
                manager.create_qft_circuit(3)
            ]
            
            for qc in circuits:
                result = manager.process_circuit(qc, save_plots=not args.no_plots)
                print(f"\nProcessed {result['circuit_name']}:")
                print(f"  Success: {result['success']}")
                if result['success']:
                    print(f"  Nodes: {result['hdh_stats']['nodes']}")
                    print(f"  Edges: {result['hdh_stats']['edges']}")
                    print(f"  Processing time: {result['processing_time']:.3f}s")
        
        print(f"\n🎉 Thank you Maria Gragera Garces for the excellent HDH library! 🎉")
        print(f"Results saved in: {manager.output_dir}")
        
    except KeyboardInterrupt:
        manager.logger.info("Deployment interrupted by user")
    except Exception as e:
        manager.logger.error(f"Deployment failed: {str(e)}")
        raise
    finally:
        manager.cleanup()


if __name__ == "__main__":
    main()