"""
CHRONOS Quantum Risk Calculation Engine
========================================

Quantitative assessment of quantum computing threats and post-quantum migration strategies.
Implements Monte Carlo simulation, Q-Day probability modeling, and financial analysis.
"""

import json
import csv
import numpy as np
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any
from collections import defaultdict

from scipy.stats import weibull_min
import matplotlib.pyplot as plt
import matplotlib.dates as mdates


class DataSensitivity(str, Enum):
    """Data sensitivity classification."""
    PUBLIC = "public"
    INTERNAL = "internal"
    CONFIDENTIAL = "confidential"
    RESTRICTED = "restricted"


class CryptographyType(str, Enum):
    """Type of cryptographic implementation."""
    RSA = "rsa"
    ECC = "ecc"
    DH = "dh"
    AES = "aes"
    SHA = "sha"
    HYBRID = "hybrid"
    PQC = "pqc"


class DataAssetType(str, Enum):
    """Type of data asset."""
    DATABASE = "database"
    FILE_STORAGE = "file_storage"
    VAULT = "vault"
    COMMUNICATION = "communication"
    IDENTITY = "identity"
    CERTIFICATE = "certificate"
    KEY = "key"


@dataclass
class DataAsset:
    """Represents a data asset with quantum risk profile."""
    
    name: str
    asset_type: DataAssetType
    sensitivity: DataSensitivity
    crypto_type: CryptographyType
    key_size: int
    creation_date: datetime = field(default_factory=datetime.now)
    location: str = "on-premises"
    business_impact: int = 5  # 1-10 scale
    recovery_time_hours: int = 24
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def get_value_score(self) -> float:
        """Calculate data asset value (1-100)."""
        sensitivity_weights = {
            DataSensitivity.PUBLIC: 10,
            DataSensitivity.INTERNAL: 40,
            DataSensitivity.CONFIDENTIAL: 70,
            DataSensitivity.RESTRICTED: 100,
        }
        
        base_value = sensitivity_weights[self.sensitivity]
        impact_multiplier = 1.0 + (self.business_impact / 10.0)
        
        return base_value * impact_multiplier


@dataclass
class QDayProbabilityModel:
    """Quantum computing threat emergence (Q-Day) probability model."""
    
    target_year: int = 2035
    standard_deviation: int = 5
    peak_capability: str = "Brute-force attack on 2048-bit RSA"
    confidence_level: float = 0.68  # 1-sigma
    
    def get_probability_at_year(self, year: int) -> float:
        """
        Calculate probability that Q-Day occurs by given year.
        Uses Weibull distribution centered on target year.
        
        Args:
            year: Target year for Q-Day occurrence
        
        Returns:
            Probability (0.0-1.0)
        """
        # Weibull parameters
        # Shape: 2.0 (bell curve-like)
        # Scale: standard_deviation * sqrt(π/2)
        shape = 2.0
        scale = self.standard_deviation * np.sqrt(np.pi / 2)
        
        # Normalize to target year
        x = (year - self.target_year + scale / 2) / scale
        
        # Weibull CDF
        if x <= 0:
            return 0.0
        
        prob = 1.0 - np.exp(-(x ** shape))
        return min(1.0, max(0.0, prob))
    
    def get_year_at_probability(self, prob: float) -> int:
        """
        Get year when Q-Day probability reaches given level.
        
        Args:
            prob: Probability threshold (0.0-1.0)
        
        Returns:
            Year when probability is reached
        """
        if prob <= 0:
            return 1970
        if prob >= 1:
            return 2100
        
        shape = 2.0
        scale = self.standard_deviation * np.sqrt(np.pi / 2)
        
        # Inverse Weibull CDF
        x = (-np.log(1 - prob)) ** (1 / shape)
        year = int(self.target_year + (x - scale / 2) * scale)
        
        return year


@dataclass
class QuantumThreatAssessment:
    """Assessment of quantum threat to a specific asset."""
    
    asset_name: str
    crypto_type: CryptographyType
    key_size: int
    threat_level: str  # "low", "medium", "high", "critical"
    harvest_now_risk: float  # 0.0-1.0
    future_risk_year_2030: float
    future_risk_year_2035: float
    future_risk_year_2040: float
    qday_survival_probability: float  # Probability asset survives until QDay
    recommended_migration: str
    urgency_score: float  # 0.0-100.0
    
    def to_dict(self) -> Dict:
        """Convert to dictionary."""
        return {
            "asset_name": self.asset_name,
            "crypto_type": self.crypto_type.value,
            "key_size": self.key_size,
            "threat_level": self.threat_level,
            "harvest_now_risk": round(self.harvest_now_risk, 4),
            "future_risk_2030": round(self.future_risk_year_2030, 4),
            "future_risk_2035": round(self.future_risk_year_2035, 4),
            "future_risk_2040": round(self.future_risk_year_2040, 4),
            "qday_survival_probability": round(self.qday_survival_probability, 4),
            "recommended_migration": self.recommended_migration,
            "urgency_score": round(self.urgency_score, 2),
        }


@dataclass
class DataValueDecay:
    """Data value decay over time based on sensitivity and lifetime."""
    
    asset_name: str
    initial_value: float
    sensitivity: DataSensitivity
    creation_date: datetime
    retention_years: int
    decay_rate: float = 0.15  # 15% per year
    
    def get_value_at_year(self, year: int) -> float:
        """
        Calculate data value at specific year.
        Uses exponential decay: V(t) = V0 * e^(-λt)
        
        Args:
            year: Target year
        
        Returns:
            Value at that year
        """
        if year < self.creation_date.year:
            return self.initial_value
        
        years_elapsed = year - self.creation_date.year
        retention_end = self.creation_date.year + self.retention_years
        
        # After retention period, value is zero
        if year > retention_end:
            return 0.0
        
        # Exponential decay
        lambda_param = self.decay_rate
        value = self.initial_value * np.exp(-lambda_param * years_elapsed)
        
        return max(0.0, value)
    
    def get_value_curve(self, start_year: int, end_year: int, step: int = 1) -> Dict[int, float]:
        """
        Get value decay curve over time range.
        
        Args:
            start_year: Start year
            end_year: End year
            step: Year step size
        
        Returns:
            Dictionary of year -> value
        """
        curve = {}
        for year in range(start_year, end_year + 1, step):
            curve[year] = self.get_value_at_year(year)
        return curve


@dataclass
class MigrationScenario:
    """Post-quantum cryptography migration scenario."""
    
    name: str
    migration_year: int
    pqc_algorithm: str  # ML-KEM, ML-DSA, SLH-DSA
    re_encrypt_existing: bool = True
    cost_per_asset: float = 5000.0  # Per asset migration cost
    disruption_hours: float = 4.0  # Average disruption per asset
    success_probability: float = 0.95  # Probability of successful migration
    testing_months: int = 6
    rollout_months: int = 12
    
    def get_total_cost(self, num_assets: int) -> float:
        """Calculate total migration cost."""
        return self.cost_per_asset * num_assets
    
    def get_completion_date(self) -> datetime:
        """Get estimated completion date."""
        start = datetime(self.migration_year, 1, 1)
        months = self.testing_months + self.rollout_months
        return start + timedelta(days=months * 30)


@dataclass
class RiskMetrics:
    """Aggregated risk metrics for assessment."""
    
    total_quantum_risk: float  # 0-100
    harvest_now_assets: int
    critical_threat_assets: int
    avg_urgency_score: float
    timeline_to_resilience: int  # Years until quantum-safe
    estimated_migration_cost: float
    resources_needed: int  # FTE equivalents
    
    def to_dict(self) -> Dict:
        """Convert to dictionary."""
        return asdict(self)


class QuantumRiskCalculator:
    """
    Comprehensive quantum risk assessment and calculation engine.
    
    Performs Monte Carlo simulations, Q-Day probability modeling, and
    cost-benefit analysis for post-quantum migration strategies.
    """
    
    def __init__(self, verbose: bool = False):
        """
        Initialize risk calculator.
        
        Args:
            verbose: Enable verbose logging
        """
        self.verbose = verbose
        self.assets: List[DataAsset] = []
        self.qday_model = QDayProbabilityModel()
        self.assessments: List[QuantumThreatAssessment] = []
        self.scenarios: List[MigrationScenario] = []
    
    def add_asset(self, asset: DataAsset) -> None:
        """Add data asset to assessment."""
        self.assets.append(asset)
        if self.verbose:
            print(f"Added asset: {asset.name}")
    
    def add_assets_batch(self, assets: List[DataAsset]) -> None:
        """Add multiple assets."""
        self.assets.extend(assets)
        if self.verbose:
            print(f"Added {len(assets)} assets")
    
    def assess_quantum_threat(self, asset: DataAsset) -> QuantumThreatAssessment:
        """
        Assess quantum threat to a specific asset.
        
        Args:
            asset: Data asset to assess
        
        Returns:
            QuantumThreatAssessment with threat analysis
        """
        # Determine threat level based on crypto type and key size
        threat_level = self._calculate_threat_level(asset)
        
        # Calculate "harvest now" risk
        # Adversaries harvest encrypted data now, decrypt later with quantum computer
        harvest_now_risk = self._calculate_harvest_now_risk(asset)
        
        # Calculate future risks at key years
        risk_2030 = self.qday_model.get_probability_at_year(2030) * harvest_now_risk
        risk_2035 = self.qday_model.get_probability_at_year(2035) * harvest_now_risk
        risk_2040 = self.qday_model.get_probability_at_year(2040) * harvest_now_risk
        
        # Calculate Q-Day survival probability
        # Probability that data is still valuable when Q-Day occurs
        data_lifetime = asset.metadata.get("data_lifetime", 10)
        qday_year = self.qday_model.get_year_at_probability(0.5)  # 50% probability
        years_until_qday = max(0, qday_year - datetime.now().year)
        
        # If data expires before Q-Day, lower risk
        qday_survival = 1.0 if data_lifetime > years_until_qday else 0.3
        
        # Calculate urgency score (0-100)
        urgency = self._calculate_urgency_score(
            threat_level, harvest_now_risk, risk_2030, qday_survival
        )
        
        # Recommend migration path
        migration_rec = self._recommend_migration(asset, threat_level)
        
        assessment = QuantumThreatAssessment(
            asset_name=asset.name,
            crypto_type=asset.crypto_type,
            key_size=asset.key_size,
            threat_level=threat_level,
            harvest_now_risk=harvest_now_risk,
            future_risk_year_2030=risk_2030,
            future_risk_year_2035=risk_2035,
            future_risk_year_2040=risk_2040,
            qday_survival_probability=qday_survival,
            recommended_migration=migration_rec,
            urgency_score=urgency,
        )
        
        self.assessments.append(assessment)
        return assessment
    
    def assess_all_assets(self) -> List[QuantumThreatAssessment]:
        """Assess quantum threat to all added assets."""
        self.assessments = []
        for asset in self.assets:
            self.assess_quantum_threat(asset)
        return self.assessments
    
    def _calculate_threat_level(self, asset: DataAsset) -> str:
        """Determine threat level based on cryptography and key size."""
        if asset.crypto_type == CryptographyType.PQC:
            return "low"
        
        if asset.crypto_type == CryptographyType.AES:
            # AES-256 is relatively safe due to large key space
            return "low" if asset.key_size >= 256 else "medium"
        
        if asset.crypto_type == CryptographyType.RSA:
            if asset.key_size < 2048:
                return "critical"
            elif asset.key_size < 3072:
                return "high"
            elif asset.key_size < 4096:
                return "medium"
            else:
                return "low"
        
        if asset.crypto_type == CryptographyType.ECC:
            if asset.key_size < 256:
                return "critical"
            elif asset.key_size < 384:
                return "high"
            else:
                return "medium"
        
        if asset.crypto_type == CryptographyType.DH:
            if asset.key_size < 2048:
                return "critical"
            elif asset.key_size < 3072:
                return "high"
            else:
                return "medium"
        
        return "medium"  # Default for unknown types
    
    def _calculate_harvest_now_risk(self, asset: DataAsset) -> float:
        """
        Calculate harvest-now-decrypt-later risk.
        
        Returns:
            Risk score 0.0-1.0
        """
        base_risk = {
            DataSensitivity.PUBLIC: 0.0,
            DataSensitivity.INTERNAL: 0.3,
            DataSensitivity.CONFIDENTIAL: 0.7,
            DataSensitivity.RESTRICTED: 1.0,
        }[asset.sensitivity]
        
        # Adjust for data lifetime
        data_lifetime = asset.metadata.get("data_lifetime", 10)
        lifetime_factor = min(1.0, data_lifetime / 20)  # Normalize to 20 years
        
        # Adjust for key size
        key_factor = 1.0
        if asset.crypto_type == CryptographyType.RSA:
            if asset.key_size >= 4096:
                key_factor = 0.8
            elif asset.key_size >= 2048:
                key_factor = 0.9
        elif asset.crypto_type == CryptographyType.ECC:
            if asset.key_size >= 384:
                key_factor = 0.7
        
        return base_risk * lifetime_factor * key_factor
    
    def _calculate_urgency_score(
        self,
        threat_level: str,
        harvest_risk: float,
        risk_2030: float,
        qday_survival: float,
    ) -> float:
        """Calculate urgency score (0-100)."""
        threat_weights = {
            "critical": 100,
            "high": 75,
            "medium": 50,
            "low": 25,
        }
        
        base_score = threat_weights.get(threat_level, 50)
        harvest_component = harvest_risk * 30  # Up to 30 points
        risk_component = risk_2030 * 25  # Up to 25 points
        survival_component = qday_survival * 15  # Up to 15 points
        
        urgency = base_score + harvest_component + risk_component + survival_component
        return min(100.0, max(0.0, urgency))
    
    def _recommend_migration(self, asset: DataAsset, threat_level: str) -> str:
        """Recommend migration strategy."""
        if threat_level == "critical":
            return "Immediate: Migrate to post-quantum cryptography (ML-KEM/ML-DSA)"
        elif threat_level == "high":
            return "Urgent: Plan PQC migration for 2024-2025"
        elif threat_level == "medium":
            return "Plan: Begin PQC evaluation and pilot programs"
        else:
            return "Monitor: Continue using current cryptography with periodic reviews"
    
    def monte_carlo_simulation(
        self, num_simulations: int = 10000, years: int = 20
    ) -> Dict[int, float]:
        """
        Run Monte Carlo simulation of quantum threat evolution.
        
        Args:
            num_simulations: Number of simulation runs
            years: Simulation period in years
        
        Returns:
            Dictionary of year -> average risk
        """
        if self.verbose:
            print(f"Running {num_simulations} Monte Carlo simulations over {years} years...")
        
        current_year = datetime.now().year
        year_risks = defaultdict(list)
        
        # Run simulations
        for sim in range(num_simulations):
            # Randomly vary Q-Day target year
            qday_variation = np.random.normal(0, 2)  # ±2 year variation
            qday_year = self.qday_model.target_year + qday_variation
            
            for year_offset in range(years):
                year = current_year + year_offset
                
                # Probability Q-Day has occurred by this year
                if year >= qday_year:
                    qday_occurred = 1.0
                else:
                    # Weibull probability
                    qday_occurred = self.qday_model.get_probability_at_year(year)
                
                # Calculate aggregate risk for this year
                year_risk = 0.0
                if self.assessments:
                    for assessment in self.assessments:
                        risk_at_year = (
                            assessment.harvest_now_risk * qday_occurred * 100
                        )
                        year_risk += risk_at_year
                    
                    year_risk /= len(self.assessments)
                
                year_risks[year_offset].append(year_risk)
        
        # Calculate averages
        result = {}
        for year_offset, risks in year_risks.items():
            result[current_year + year_offset] = float(np.mean(risks))
        
        if self.verbose:
            print(f"Monte Carlo simulation complete")
        
        return result
    
    def calculate_migration_roi(
        self, scenario: MigrationScenario
    ) -> Dict[str, float]:
        """
        Calculate ROI for a migration scenario.
        
        Args:
            scenario: Migration scenario to evaluate
        
        Returns:
            ROI metrics dictionary
        """
        num_assets = len(self.assets)
        
        # Calculate costs
        migration_cost = scenario.get_total_cost(num_assets)
        total_disruption_hours = scenario.disruption_hours * num_assets
        disruption_cost = total_disruption_hours * 150  # $150/hour labor
        
        total_cost = migration_cost + disruption_cost
        
        # Calculate benefits (avoided quantum breach risk)
        quantum_breach_prob = 0.0
        quantum_breach_impact = 0.0
        
        for asset in self.assets:
            # If not migrated, calculate breach probability and impact
            if asset.creation_date.year < scenario.migration_year:
                value = asset.get_value_score()
                impact = asset.business_impact * 10000  # $10k per impact point
                
                # Probability of breach if not migrated
                breach_prob = self.qday_model.get_probability_at_year(
                    scenario.migration_year + 5
                )
                
                quantum_breach_prob += breach_prob * value
                quantum_breach_impact += breach_prob * impact
        
        # Expected loss without mitigation
        expected_loss = quantum_breach_impact
        
        # Calculate ROI
        net_benefit = expected_loss - total_cost
        roi_percent = (net_benefit / total_cost * 100) if total_cost > 0 else 0
        
        return {
            "migration_cost": migration_cost,
            "disruption_cost": disruption_cost,
            "total_cost": total_cost,
            "expected_loss_avoided": expected_loss,
            "net_benefit": net_benefit,
            "roi_percentage": roi_percent,
            "payback_years": total_cost / (expected_loss / 10) if expected_loss > 0 else 999,
            "success_probability": scenario.success_probability,
        }
    
    def get_risk_metrics(self) -> RiskMetrics:
        """
        Calculate aggregated risk metrics.
        
        Returns:
            RiskMetrics with overall assessment
        """
        if not self.assessments:
            self.assess_all_assets()
        
        # Count critical assets
        critical_count = sum(
            1 for a in self.assessments if a.threat_level == "critical"
        )
        
        # Count harvest-now vulnerable assets
        harvest_count = sum(
            1 for a in self.assessments if a.harvest_now_risk > 0.5
        )
        
        # Calculate average urgency
        avg_urgency = (
            np.mean([a.urgency_score for a in self.assessments])
            if self.assessments
            else 0
        )
        
        # Calculate total quantum risk (0-100 scale)
        total_risk = min(100.0, avg_urgency * 1.2)
        
        # Estimate timeline to resilience (migration)
        # Assume 2 years for full PQC deployment if started immediately
        timeline = 2 + (len(self.assets) // 50)  # Add time for large estates
        
        # Estimate migration cost ($5k per asset average)
        migration_cost = len(self.assets) * 5000
        
        # Estimate FTE requirements (1 FTE per 30 assets)
        fte_required = max(1, len(self.assets) // 30)
        
        return RiskMetrics(
            total_quantum_risk=total_risk,
            harvest_now_assets=harvest_count,
            critical_threat_assets=critical_count,
            avg_urgency_score=avg_urgency,
            timeline_to_resilience=timeline,
            estimated_migration_cost=migration_cost,
            resources_needed=fte_required,
        )
    
    def calculate_data_value_decay(
        self, asset: DataAsset, retention_years: int = 10
    ) -> DataValueDecay:
        """
        Calculate data value decay for an asset.
        
        Args:
            asset: Data asset
            retention_years: Expected data retention period
        
        Returns:
            DataValueDecay model
        """
        decay = DataValueDecay(
            asset_name=asset.name,
            initial_value=asset.get_value_score() * 1000,  # Convert to monetary value
            sensitivity=asset.sensitivity,
            creation_date=asset.creation_date,
            retention_years=retention_years,
            decay_rate=0.15,  # 15% annual decay
        )
        
        return decay
    
    def export_risk_report_json(self, filepath: str) -> None:
        """
        Export risk assessment to JSON format.
        
        Args:
            filepath: Output file path
        """
        if not self.assessments:
            self.assess_all_assets()
        
        metrics = self.get_risk_metrics()
        
        report = {
            "timestamp": datetime.now().isoformat(),
            "summary": asdict(metrics),
            "assessments": [a.to_dict() for a in self.assessments],
            "qday_model": {
                "target_year": self.qday_model.target_year,
                "standard_deviation": self.qday_model.standard_deviation,
                "probability_2030": self.qday_model.get_probability_at_year(2030),
                "probability_2035": self.qday_model.get_probability_at_year(2035),
                "probability_2040": self.qday_model.get_probability_at_year(2040),
            },
        }
        
        with open(filepath, "w") as f:
            json.dump(report, f, indent=2)
        
        if self.verbose:
            print(f"Risk report exported to {filepath}")
    
    def export_risk_report_csv(self, filepath: str) -> None:
        """
        Export risk assessment to CSV format.
        
        Args:
            filepath: Output file path
        """
        if not self.assessments:
            self.assess_all_assets()
        
        with open(filepath, "w", newline="") as f:
            fieldnames = [
                "asset_name",
                "crypto_type",
                "key_size",
                "threat_level",
                "harvest_now_risk",
                "risk_2030",
                "risk_2035",
                "risk_2040",
                "qday_survival",
                "urgency_score",
                "recommendation",
            ]
            
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            
            for assessment in self.assessments:
                writer.writerow({
                    "asset_name": assessment.asset_name,
                    "crypto_type": assessment.crypto_type.value,
                    "key_size": assessment.key_size,
                    "threat_level": assessment.threat_level,
                    "harvest_now_risk": round(assessment.harvest_now_risk, 4),
                    "risk_2030": round(assessment.future_risk_year_2030, 4),
                    "risk_2035": round(assessment.future_risk_year_2035, 4),
                    "risk_2040": round(assessment.future_risk_year_2040, 4),
                    "qday_survival": round(assessment.qday_survival_probability, 4),
                    "urgency_score": round(assessment.urgency_score, 2),
                    "recommendation": assessment.recommended_migration,
                })
        
        if self.verbose:
            print(f"Risk report exported to {filepath}")
    
    def plot_risk_timeline(
        self,
        filepath: Optional[str] = None,
        show: bool = False,
        years: int = 20,
    ) -> Optional[plt.Figure]:
        """
        Create risk timeline visualization.
        
        Args:
            filepath: Save to file if provided
            show: Display plot
            years: Timeline in years
        
        Returns:
            Matplotlib figure if not saving
        """
        # Run Monte Carlo simulation
        mc_results = self.monte_carlo_simulation(num_simulations=1000, years=years)
        
        # Get Q-Day probability curve
        current_year = datetime.now().year
        qday_probs = {}
        for year_offset in range(years):
            year = current_year + year_offset
            qday_probs[year] = self.qday_model.get_probability_at_year(year)
        
        # Create figure with subplots
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10))
        
        # Plot 1: Quantum Risk Timeline
        years_list = sorted(mc_results.keys())
        risks = [mc_results[year] for year in years_list]
        
        ax1.fill_between(years_list, 0, risks, alpha=0.3, color="red", label="Quantum Risk")
        ax1.plot(years_list, risks, color="red", linewidth=2, marker="o", markersize=4)
        ax1.set_xlabel("Year", fontsize=12)
        ax1.set_ylabel("Risk Score (0-100)", fontsize=12)
        ax1.set_title("Quantum Computing Threat Evolution", fontsize=14, fontweight="bold")
        ax1.grid(True, alpha=0.3)
        ax1.legend(fontsize=10)
        ax1.set_ylim(0, 100)
        
        # Plot 2: Q-Day Probability
        qday_years = sorted(qday_probs.keys())
        qday_values = [qday_probs[year] * 100 for year in qday_years]
        
        ax2.fill_between(qday_years, 0, qday_values, alpha=0.3, color="orange")
        ax2.plot(qday_years, qday_values, color="orange", linewidth=2, marker="s", markersize=4)
        ax2.axvline(
            x=self.qday_model.target_year,
            color="red",
            linestyle="--",
            linewidth=2,
            label=f"Target Q-Day: {self.qday_model.target_year}±{self.qday_model.standard_deviation}y"
        )
        ax2.set_xlabel("Year", fontsize=12)
        ax2.set_ylabel("Probability (%)", fontsize=12)
        ax2.set_title("Quantum Computer Emergence Probability", fontsize=14, fontweight="bold")
        ax2.grid(True, alpha=0.3)
        ax2.legend(fontsize=10)
        ax2.set_ylim(0, 100)
        
        plt.tight_layout()
        
        if filepath:
            plt.savefig(filepath, dpi=300, bbox_inches="tight")
            if self.verbose:
                print(f"Timeline visualization saved to {filepath}")
        
        if show:
            plt.show()
        
        return fig if not filepath else None
    
    def plot_threat_distribution(
        self,
        filepath: Optional[str] = None,
        show: bool = False,
    ) -> Optional[plt.Figure]:
        """
        Create threat distribution visualization.
        
        Args:
            filepath: Save to file if provided
            show: Display plot
        
        Returns:
            Matplotlib figure if not saving
        """
        if not self.assessments:
            self.assess_all_assets()
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(14, 10))
        
        # Plot 1: Threat Level Distribution
        threat_counts = defaultdict(int)
        for a in self.assessments:
            threat_counts[a.threat_level] += 1
        
        levels = ["critical", "high", "medium", "low"]
        counts = [threat_counts.get(level, 0) for level in levels]
        colors = ["red", "orange", "yellow", "green"]
        
        ax1.bar(levels, counts, color=colors, alpha=0.7, edgecolor="black")
        ax1.set_title("Threat Level Distribution", fontsize=12, fontweight="bold")
        ax1.set_ylabel("Number of Assets", fontsize=10)
        ax1.grid(True, alpha=0.3, axis="y")
        
        # Plot 2: Urgency Score Distribution
        urgencies = [a.urgency_score for a in self.assessments]
        ax2.hist(urgencies, bins=20, color="purple", alpha=0.7, edgecolor="black")
        ax2.axvline(np.mean(urgencies), color="red", linestyle="--", linewidth=2, label=f"Mean: {np.mean(urgencies):.1f}")
        ax2.set_title("Urgency Score Distribution", fontsize=12, fontweight="bold")
        ax2.set_xlabel("Urgency Score (0-100)", fontsize=10)
        ax2.set_ylabel("Number of Assets", fontsize=10)
        ax2.legend()
        ax2.grid(True, alpha=0.3, axis="y")
        
        # Plot 3: Harvest-Now Risk vs Crypto Type
        crypto_risks = defaultdict(list)
        for a in self.assessments:
            crypto_risks[a.crypto_type.value].append(a.harvest_now_risk)
        
        crypto_types = list(crypto_risks.keys())
        avg_risks = [np.mean(crypto_risks[ct]) for ct in crypto_types]
        
        ax3.barh(crypto_types, avg_risks, color="blue", alpha=0.7, edgecolor="black")
        ax3.set_title("Average Harvest-Now Risk by Crypto Type", fontsize=12, fontweight="bold")
        ax3.set_xlabel("Harvest-Now Risk (0-1)", fontsize=10)
        ax3.grid(True, alpha=0.3, axis="x")
        
        # Plot 4: Key Size vs Urgency
        key_sizes = [a.key_size for a in self.assessments]
        urgencies = [a.urgency_score for a in self.assessments]
        colors_scatter = [
            "red" if a.threat_level == "critical"
            else "orange" if a.threat_level == "high"
            else "yellow" if a.threat_level == "medium"
            else "green"
            for a in self.assessments
        ]
        
        ax4.scatter(key_sizes, urgencies, c=colors_scatter, s=100, alpha=0.6, edgecolor="black")
        ax4.set_title("Key Size vs Urgency Score", fontsize=12, fontweight="bold")
        ax4.set_xlabel("Key Size (bits)", fontsize=10)
        ax4.set_ylabel("Urgency Score (0-100)", fontsize=10)
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if filepath:
            plt.savefig(filepath, dpi=300, bbox_inches="tight")
            if self.verbose:
                print(f"Threat distribution visualization saved to {filepath}")
        
        if show:
            plt.show()
        
        return fig if not filepath else None
    
    def plot_migration_roi_comparison(
        self,
        scenarios: List[MigrationScenario],
        filepath: Optional[str] = None,
        show: bool = False,
    ) -> Optional[plt.Figure]:
        """
        Compare ROI across migration scenarios.
        
        Args:
            scenarios: List of migration scenarios
            filepath: Save to file if provided
            show: Display plot
        
        Returns:
            Matplotlib figure if not saving
        """
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(14, 10))
        
        scenario_names = [s.name for s in scenarios]
        roi_results = [self.calculate_migration_roi(s) for s in scenarios]
        
        # Plot 1: Total Cost
        costs = [r["total_cost"] / 1000 for r in roi_results]  # In thousands
        ax1.bar(scenario_names, costs, color="steelblue", alpha=0.7, edgecolor="black")
        ax1.set_title("Migration Cost by Scenario", fontsize=12, fontweight="bold")
        ax1.set_ylabel("Cost ($1000s)", fontsize=10)
        ax1.grid(True, alpha=0.3, axis="y")
        
        # Plot 2: Expected Benefit
        benefits = [r["expected_loss_avoided"] / 1000 for r in roi_results]
        ax2.bar(scenario_names, benefits, color="green", alpha=0.7, edgecolor="black")
        ax2.set_title("Expected Loss Avoided", fontsize=12, fontweight="bold")
        ax2.set_ylabel("Benefit ($1000s)", fontsize=10)
        ax2.grid(True, alpha=0.3, axis="y")
        
        # Plot 3: ROI Percentage
        roi_percents = [r["roi_percentage"] for r in roi_results]
        colors_roi = ["green" if r > 0 else "red" for r in roi_percents]
        ax3.bar(scenario_names, roi_percents, color=colors_roi, alpha=0.7, edgecolor="black")
        ax3.axhline(y=0, color="black", linestyle="-", linewidth=0.8)
        ax3.set_title("Return on Investment (%)", fontsize=12, fontweight="bold")
        ax3.set_ylabel("ROI (%)", fontsize=10)
        ax3.grid(True, alpha=0.3, axis="y")
        
        # Plot 4: Payback Period
        paybacks = [min(r["payback_years"], 20) for r in roi_results]  # Cap at 20 years
        ax4.bar(scenario_names, paybacks, color="purple", alpha=0.7, edgecolor="black")
        ax4.set_title("Payback Period", fontsize=12, fontweight="bold")
        ax4.set_ylabel("Years", fontsize=10)
        ax4.grid(True, alpha=0.3, axis="y")
        
        plt.tight_layout()
        
        if filepath:
            plt.savefig(filepath, dpi=300, bbox_inches="tight")
            if self.verbose:
                print(f"ROI comparison visualization saved to {filepath}")
        
        if show:
            plt.show()
        
        return fig if not filepath else None
