"""
Synthetic DGP Atlas Extended (v0.2) - 500 DGPs for Hierarchical Bayesian Calibration

Generates 10 families × 50 DGPs = 500 total synthetic datasets for learning
accurate calibration models.

Families:
1. F1_clean_var: Clean VAR (TRUE NEGATIVES)
2. F2_structural_breaks: Structural breaks
3. F3_irregular_sampling: MCAR, MAR, MNAR-seasonal
4. F4_high_persistence: High persistence (eigenvalue ~0.95)
5. F5_latent_confounders: Latent confounders
6. F6_seasonality: Annual, quarterly, monthly cycles (NEW)
7. F7_nonlinear: Quadratic, sigmoid, threshold (NEW)
8. F8_non_gaussian: Heavy tails, skewed, bimodal (NEW)
9. F9_mixed_violations: 2-3 violations combined (NEW)
10. F10_extreme_cases: Edge cases for stress testing (NEW)

Output:
- Location: /Volumes/X10 Pro/autocause_resources/data/synthetic_dgp_atlas_v02/
- Per DGP: {family}/{dgp_id}_data.csv + {dgp_id}_metadata.json
- Metadata includes ground truth risks (0-1 scale) for calibration

Author: Marco Rueda
Date: December 1, 2024
"""

import numpy as np
import pandas as pd
from pathlib import Path
import json
from datetime import datetime, timedelta
from typing import Dict, Tuple
from dataclasses import dataclass, asdict
from scipy import stats
from sklearn.preprocessing import StandardScaler


def convert_to_serializable(obj):
    """Convert numpy types to Python native types for JSON serialization."""
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {k: convert_to_serializable(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_to_serializable(item) for item in obj]
    return obj


@dataclass
class DGPMetadata:
    """Metadata for each generated DGP."""
    dgp_id: str
    family: str
    description: str
    n_variables: int
    n_samples: int
    true_graph: Dict
    nonstationarity_risk: float
    irregularity_risk: float
    persistence_risk: float
    confounding_risk: float
    generation_params: Dict
    random_seed: int
    generation_timestamp: str


class SyntheticAtlasExtended:
    """
    Generate extended synthetic atlas for v0.2 calibration.
    
    Each family has specific violation patterns that map to ground truth risks.
    """
    
    def __init__(self, output_dir: str, random_seed: int = 42):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.base_rng = np.random.RandomState(random_seed)
        self.generated_count = 0
        
    def generate_all(self, n_per_family: int = 50):
        """Generate all DGPs for all families."""
        families = self._define_families()
        
        print(f"\n{'='*70}")
        print(f"GENERATING EXTENDED SYNTHETIC ATLAS (v0.2)")
        print(f"{'='*70}")
        print(f"Total DGPs: {len(families) * n_per_family}")
        print(f"Output: {self.output_dir}")
        print(f"{'='*70}\n")
        
        for family_name, family_config in families.items():
            print(f"\n[{family_name}] {family_config['description']}")
            print(f"  Expected risks: {family_config['expected_risks']}")
            print(f"  Generating {n_per_family} DGPs...")
            
            family_dir = self.output_dir / family_name
            family_dir.mkdir(parents=True, exist_ok=True)
            
            for i in range(n_per_family):
                # Create unique seed for this DGP
                dgp_seed = self.base_rng.randint(0, 1000000)
                dgp_rng = np.random.RandomState(dgp_seed)
                
                dgp_id = f"{family_name}_{i:03d}"
                
                try:
                    self._generate_single_dgp(
                        dgp_id, family_name, family_config, family_dir, dgp_rng, dgp_seed
                    )
                    self.generated_count += 1
                    
                    if (i + 1) % 10 == 0:
                        print(f"    Progress: {i+1}/{n_per_family}")
                
                except Exception as e:
                    print(f"    ERROR generating {dgp_id}: {e}")
                    continue
        
        print(f"\n{'='*70}")
        print(f"✓ GENERATION COMPLETE")
        print(f"{'='*70}")
        print(f"Total generated: {self.generated_count} DGPs")
        print(f"Location: {self.output_dir}")
        print(f"{'='*70}\n")
    
    def _define_families(self) -> Dict:
        """Define 10 DGP families with expected risk levels."""
        return {
            "F1_clean_var": {
                "description": "Clean VAR, no violations (TRUE NEGATIVES)",
                "generator": self._generate_clean_var,
                "expected_risks": {
                    "nonstationarity": 0.05,
                    "irregularity": 0.05,
                    "persistence": 0.25,
                    "confounding": 0.05
                }
            },
            "F2_structural_breaks": {
                "description": "VAR with 1-3 structural breaks",
                "generator": self._generate_structural_breaks,
                "expected_risks": {
                    "nonstationarity": 0.85,
                    "irregularity": 0.10,
                    "persistence": 0.35,
                    "confounding": 0.10
                }
            },
            "F3_irregular_sampling": {
                "description": "Irregular sampling (MCAR, MAR, MNAR-seasonal)",
                "generator": self._generate_irregular,
                "expected_risks": {
                    "nonstationarity": 0.15,
                    "irregularity": 0.75,
                    "persistence": 0.30,
                    "confounding": 0.10
                }
            },
            "F4_high_persistence": {
                "description": "High persistence (eigenvalue ~0.95)",
                "generator": self._generate_high_persistence,
                "expected_risks": {
                    "nonstationarity": 0.20,
                    "irregularity": 0.10,
                    "persistence": 0.90,
                    "confounding": 0.10
                }
            },
            "F5_latent_confounders": {
                "description": "Latent confounders L -> {X, Y}",
                "generator": self._generate_latent_confounders,
                "expected_risks": {
                    "nonstationarity": 0.20,
                    "irregularity": 0.10,
                    "persistence": 0.35,
                    "confounding": 0.85
                }
            },
            "F6_seasonality": {
                "description": "Strong seasonal patterns (annual, quarterly, monthly)",
                "generator": self._generate_seasonality,
                "expected_risks": {
                    "nonstationarity": 0.70,  # Seasonal = non-stationary
                    "irregularity": 0.10,
                    "persistence": 0.45,
                    "confounding": 0.15
                }
            },
            "F7_nonlinear": {
                "description": "Nonlinear dynamics (quadratic, sigmoid, threshold)",
                "generator": self._generate_nonlinear,
                "expected_risks": {
                    "nonstationarity": 0.25,
                    "irregularity": 0.10,
                    "persistence": 0.35,
                    "confounding": 0.20
                }
            },
            "F8_non_gaussian": {
                "description": "Non-Gaussian distributions (heavy tails, skewed)",
                "generator": self._generate_non_gaussian,
                "expected_risks": {
                    "nonstationarity": 0.15,
                    "irregularity": 0.10,
                    "persistence": 0.30,
                    "confounding": 0.15
                }
            },
            "F9_mixed_violations": {
                "description": "Multiple violations combined (realistic scenarios)",
                "generator": self._generate_mixed,
                "expected_risks": {
                    "nonstationarity": 0.60,
                    "irregularity": 0.50,
                    "persistence": 0.55,
                    "confounding": 0.50
                }
            },
            "F10_extreme_cases": {
                "description": "Edge cases for stress testing",
                "generator": self._generate_extreme,
                "expected_risks": {
                    "nonstationarity": 0.50,
                    "irregularity": 0.45,
                    "persistence": 0.55,
                    "confounding": 0.45
                }
            }
        }
    
    def _generate_single_dgp(self, dgp_id: str, family_name: str, family_config: Dict, 
                            output_dir: Path, rng: np.random.RandomState, seed: int):
        """Generate a single DGP with data + metadata."""
        # Call family-specific generator
        data, true_graph, true_risks, gen_params = family_config["generator"](rng)
        
        # Ensure data has datetime index
        if not isinstance(data.index, pd.DatetimeIndex):
            start_date = datetime(2020, 1, 1)
            data.index = pd.date_range(start=start_date, periods=len(data), freq='D')
        
        # Save data
        data.to_csv(output_dir / f"{dgp_id}_data.csv")
        
        # Create metadata
        metadata = DGPMetadata(
            dgp_id=dgp_id,
            family=family_name,
            description=family_config["description"],
            n_variables=len(data.columns),
            n_samples=len(data),
            true_graph=true_graph,
            nonstationarity_risk=true_risks["nonstationarity"],
            irregularity_risk=true_risks["irregularity"],
            persistence_risk=true_risks["persistence"],
            confounding_risk=true_risks["confounding"],
            generation_params=gen_params,
            random_seed=seed,
            generation_timestamp=datetime.now().isoformat()
        )
        
        # Save metadata (convert numpy types to native Python)
        metadata_dict = convert_to_serializable(asdict(metadata))
        with open(output_dir / f"{dgp_id}_metadata.json", 'w') as f:
            json.dump(metadata_dict, f, indent=2)
    
    # ==================== FAMILY GENERATORS ====================
    
    def _generate_clean_var(self, rng: np.random.RandomState) -> Tuple:
        """
        F1: Clean VAR process with no violations.
        
        This is our TRUE NEGATIVE family - should score low on all risks.
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_samples = rng.choice([500, 750, 1000])
        
        # Generate stable VAR process
        # A matrix: moderate connectivity, stable eigenvalues
        A = self._generate_stable_var_matrix(n_vars, max_eigenvalue=0.7, rng=rng)
        
        # Generate clean time series
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        for t in range(1, n_samples):
            X[t] = A @ X[t-1] + rng.randn(n_vars) * 0.1
        
        # Create DataFrame
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        # True causal graph (from A matrix)
        true_graph = self._extract_graph_from_var(A, threshold=0.1)
        
        # Ground truth risks (all low)
        true_risks = {
            "nonstationarity": 0.05,
            "irregularity": 0.05,
            "persistence": 0.25,  # Some persistence is natural
            "confounding": 0.05
        }
        
        gen_params = {
            "n_vars": int(n_vars),
            "n_samples": int(n_samples),
            "max_eigenvalue": 0.7,
            "noise_std": 0.1
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_structural_breaks(self, rng: np.random.RandomState) -> Tuple:
        """
        F2: VAR with structural breaks.
        
        High NonstationarityRisk due to parameter changes.
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_samples = rng.choice([500, 750, 1000])
        n_breaks = rng.choice([1, 2, 3])
        
        # Determine break points
        break_points = sorted(rng.choice(range(100, n_samples-100), size=n_breaks, replace=False))
        
        # Generate different VAR matrices for each regime
        regimes = []
        A_prev = self._generate_stable_var_matrix(n_vars, max_eigenvalue=0.7, rng=rng)
        
        for i in range(n_breaks + 1):
            if i == 0:
                regimes.append(A_prev)
            else:
                # Create new matrix with some parameter shifts
                A_new = A_prev + rng.randn(n_vars, n_vars) * 0.3
                # Re-stabilize
                A_new = self._stabilize_var_matrix(A_new)
                regimes.append(A_new)
                A_prev = A_new
        
        # Generate time series with breaks
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        current_regime = 0
        for t in range(1, n_samples):
            # Check if we hit a break point
            if current_regime < len(break_points) and t >= break_points[current_regime]:
                current_regime += 1
            
            X[t] = regimes[current_regime] @ X[t-1] + rng.randn(n_vars) * 0.1
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        # True graph (average across regimes)
        true_graph = self._extract_graph_from_var(regimes[0], threshold=0.1)
        
        # Ground truth risks
        # Nonstationarity depends on break magnitude
        break_magnitude = np.mean([np.linalg.norm(regimes[i+1] - regimes[i]) 
                                   for i in range(len(regimes)-1)])
        
        true_risks = {
            "nonstationarity": min(0.95, 0.5 + break_magnitude * 0.5),
            "irregularity": 0.10,
            "persistence": 0.35,
            "confounding": 0.10
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_samples": n_samples,
            "n_breaks": n_breaks,
            "break_points": break_points,
            "break_magnitude": float(break_magnitude)
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_irregular(self, rng: np.random.RandomState) -> Tuple:
        """
        F3: Irregular sampling (MCAR, MAR, MNAR-seasonal).
        
        High IrregularityRisk due to missing data patterns.
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_samples = rng.choice([500, 750, 1000])
        
        # Generate base VAR process
        A = self._generate_stable_var_matrix(n_vars, max_eigenvalue=0.7, rng=rng)
        
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        for t in range(1, n_samples):
            X[t] = A @ X[t-1] + rng.randn(n_vars) * 0.1
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        # Add datetime index
        start_date = datetime(2020, 1, 1)
        data.index = pd.date_range(start=start_date, periods=len(data), freq='D')
        
        # Introduce missing data
        missingness_type = rng.choice(['MCAR', 'MAR', 'MNAR-seasonal'])
        missing_rate = rng.uniform(0.15, 0.35)
        
        if missingness_type == 'MCAR':
            # Completely random
            mask = rng.rand(n_samples, n_vars) < missing_rate
            data[mask] = np.nan
            irregularity_risk = 0.5 + missing_rate
        
        elif missingness_type == 'MAR':
            # Missing depends on observed values
            for i in range(n_vars):
                # Higher values more likely to be missing
                prob = 1 / (1 + np.exp(-data[f"X{i}"].values))
                mask = rng.rand(n_samples) < (prob * missing_rate * 2)
                data.loc[mask, f"X{i}"] = np.nan
            irregularity_risk = 0.6 + missing_rate
        
        else:  # MNAR-seasonal
            # Missing more in winter (seasonal pattern)
            day_of_year = data.index.dayofyear
            winter_months = (day_of_year < 90) | (day_of_year > 300)  # DJF
            
            for i in range(n_vars):
                mask = winter_months & (rng.rand(n_samples) < missing_rate * 2.5)
                data.loc[mask, f"X{i}"] = np.nan
            irregularity_risk = 0.75 + missing_rate * 0.5
        
        true_graph = self._extract_graph_from_var(A, threshold=0.1)
        
        true_risks = {
            "nonstationarity": 0.15,
            "irregularity": min(0.95, irregularity_risk),
            "persistence": 0.30,
            "confounding": 0.10
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_samples": n_samples,
            "missingness_type": missingness_type,
            "missing_rate": float(missing_rate)
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_high_persistence(self, rng: np.random.RandomState) -> Tuple:
        """
        F4: High persistence (eigenvalue ~0.95).
        
        High PersistenceRisk due to long memory.
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_samples = rng.choice([500, 750, 1000])
        
        # Generate VAR with high eigenvalue
        target_eigenvalue = rng.uniform(0.92, 0.98)
        A = self._generate_stable_var_matrix(n_vars, max_eigenvalue=target_eigenvalue, rng=rng)
        
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        for t in range(1, n_samples):
            X[t] = A @ X[t-1] + rng.randn(n_vars) * 0.05  # Lower noise for persistence
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        true_graph = self._extract_graph_from_var(A, threshold=0.1)
        
        # Persistence risk increases with eigenvalue
        persistence_risk = (target_eigenvalue - 0.7) / 0.3  # Scale 0.7-1.0 -> 0-1
        
        true_risks = {
            "nonstationarity": 0.20,
            "irregularity": 0.10,
            "persistence": min(0.95, persistence_risk),
            "confounding": 0.10
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_samples": n_samples,
            "target_eigenvalue": float(target_eigenvalue),
            "actual_eigenvalue": float(np.max(np.abs(np.linalg.eigvals(A))))
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_latent_confounders(self, rng: np.random.RandomState) -> Tuple:
        """
        F5: Latent confounders L -> {X, Y}.
        
        High ConfoundingRisk due to unobserved common causes.
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_latent = rng.choice([1, 2])
        n_samples = rng.choice([500, 750, 1000])
        
        # Generate latent variables
        L = np.zeros((n_samples, n_latent))
        L[0] = rng.randn(n_latent) * 0.1
        
        # Latent VAR process
        A_latent = self._generate_stable_var_matrix(n_latent, max_eigenvalue=0.7, rng=rng)
        
        for t in range(1, n_samples):
            L[t] = A_latent @ L[t-1] + rng.randn(n_latent) * 0.1
        
        # Generate observed variables influenced by latent
        X = np.zeros((n_samples, n_vars))
        
        # Some variables depend on latent confounders
        latent_effects = rng.randn(n_vars, n_latent) * 0.8
        
        # Also some direct causal effects between observed
        A_obs = self._generate_stable_var_matrix(n_vars, max_eigenvalue=0.5, rng=rng)
        
        X[0] = rng.randn(n_vars) * 0.1
        for t in range(1, n_samples):
            X[t] = A_obs @ X[t-1] + latent_effects @ L[t] + rng.randn(n_vars) * 0.1
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        # True graph (observed only, latent not included)
        true_graph = self._extract_graph_from_var(A_obs, threshold=0.1)
        
        # Confounding risk depends on strength of latent effects
        latent_strength = np.mean(np.abs(latent_effects))
        
        true_risks = {
            "nonstationarity": 0.20,
            "irregularity": 0.10,
            "persistence": 0.35,
            "confounding": min(0.95, 0.5 + latent_strength)
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_latent": n_latent,
            "n_samples": n_samples,
            "latent_strength": float(latent_strength)
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_seasonality(self, rng: np.random.RandomState) -> Tuple:
        """
        F6: Strong seasonal patterns.
        
        High NonstationarityRisk due to seasonal non-stationarity.
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_samples = rng.choice([500, 750, 1000])
        
        # Generate base VAR
        A = self._generate_stable_var_matrix(n_vars, max_eigenvalue=0.5, rng=rng)
        
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        # Add seasonal components
        t = np.arange(n_samples)
        seasonal_periods = [365, 365/4, 365/12]  # Annual, quarterly, monthly
        seasonal_strength = rng.uniform(0.5, 1.5)
        
        seasonality = np.zeros((n_samples, n_vars))
        for i in range(n_vars):
            # Each variable gets different seasonal pattern
            period = rng.choice(seasonal_periods)
            amplitude = rng.uniform(0.3, 0.8) * seasonal_strength
            phase = rng.uniform(0, 2*np.pi)
            seasonality[:, i] = amplitude * np.sin(2*np.pi*t/period + phase)
        
        for t_idx in range(1, n_samples):
            X[t_idx] = A @ X[t_idx-1] + seasonality[t_idx] + rng.randn(n_vars) * 0.1
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        true_graph = self._extract_graph_from_var(A, threshold=0.1)
        
        # Seasonality -> high nonstationarity risk
        true_risks = {
            "nonstationarity": min(0.95, 0.5 + seasonal_strength * 0.3),
            "irregularity": 0.10,
            "persistence": 0.45,
            "confounding": 0.15
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_samples": n_samples,
            "seasonal_strength": float(seasonal_strength),
            "seasonal_periods": seasonal_periods
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_nonlinear(self, rng: np.random.RandomState) -> Tuple:
        """
        F7: Nonlinear dynamics.
        
        Tests nonlinearity detection (not a core risk in minimal v0.2).
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_samples = rng.choice([500, 750, 1000])
        
        # Generate with nonlinear transformation
        A = self._generate_stable_var_matrix(n_vars, max_eigenvalue=0.6, rng=rng)
        
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        nonlinearity_type = rng.choice(['quadratic', 'sigmoid', 'threshold'])
        
        for t in range(1, n_samples):
            linear_part = A @ X[t-1]
            
            if nonlinearity_type == 'quadratic':
                nonlinear_part = 0.3 * (linear_part ** 2)
            elif nonlinearity_type == 'sigmoid':
                nonlinear_part = 0.5 * np.tanh(2 * linear_part)
            else:  # threshold
                nonlinear_part = np.where(linear_part > 0, linear_part * 1.2, linear_part * 0.8)
            
            X[t] = linear_part + nonlinear_part + rng.randn(n_vars) * 0.1
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        true_graph = self._extract_graph_from_var(A, threshold=0.1)
        
        true_risks = {
            "nonstationarity": 0.25,
            "irregularity": 0.10,
            "persistence": 0.35,
            "confounding": 0.20
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_samples": n_samples,
            "nonlinearity_type": nonlinearity_type
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_non_gaussian(self, rng: np.random.RandomState) -> Tuple:
        """
        F8: Non-Gaussian distributions.
        
        Tests gaussianity detection (auxiliary diagnostic).
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_samples = rng.choice([500, 750, 1000])
        
        A = self._generate_stable_var_matrix(n_vars, max_eigenvalue=0.6, rng=rng)
        
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        # Different noise distributions per variable
        noise_types = rng.choice(['student_t', 'laplace', 'gamma'], size=n_vars)
        
        for t in range(1, n_samples):
            noise = np.zeros(n_vars)
            for i in range(n_vars):
                if noise_types[i] == 'student_t':
                    noise[i] = rng.standard_t(df=3) * 0.1
                elif noise_types[i] == 'laplace':
                    noise[i] = rng.laplace(0, 0.07)
                else:  # gamma
                    noise[i] = (rng.gamma(2, 0.05) - 0.1)
            
            X[t] = A @ X[t-1] + noise
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        true_graph = self._extract_graph_from_var(A, threshold=0.1)
        
        true_risks = {
            "nonstationarity": 0.15,
            "irregularity": 0.10,
            "persistence": 0.30,
            "confounding": 0.15
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_samples": n_samples,
            "noise_types": noise_types.tolist()
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_mixed(self, rng: np.random.RandomState) -> Tuple:
        """
        F9: Multiple violations combined (realistic scenarios).
        
        Moderate-to-high risks on multiple dimensions.
        """
        n_vars = rng.choice([5, 6, 7, 8])
        n_samples = rng.choice([500, 750, 1000])
        
        # Combine: structural breaks + persistence + some irregularity
        violations = rng.choice(['breaks+persist', 'breaks+irregular', 'persist+seasonal', 
                                'confound+irregular', 'all'])
        
        if 'breaks' in violations or violations == 'all':
            # Add 1-2 breaks
            n_breaks = rng.choice([1, 2])
            break_points = sorted(rng.choice(range(100, n_samples-100), size=n_breaks, replace=False))
            nonstationarity_risk = 0.65
        else:
            n_breaks = 0
            break_points = []
            nonstationarity_risk = 0.20
        
        if 'persist' in violations or violations == 'all':
            max_eigenvalue = 0.90
            persistence_risk = 0.75
        else:
            max_eigenvalue = 0.65
            persistence_risk = 0.35
        
        if 'irregular' in violations or violations == 'all':
            missing_rate = 0.20
            irregularity_risk = 0.60
        else:
            missing_rate = 0
            irregularity_risk = 0.10
        
        if 'confound' in violations or violations == 'all':
            n_latent = 1
            confounding_risk = 0.70
        else:
            n_latent = 0
            confounding_risk = 0.15
        
        # Generate with these violations
        A = self._generate_stable_var_matrix(n_vars, max_eigenvalue=max_eigenvalue, rng=rng)
        
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        current_regime = 0
        for t in range(1, n_samples):
            if current_regime < len(break_points) and t >= break_points[current_regime]:
                # Apply parameter shift
                A = A + rng.randn(n_vars, n_vars) * 0.2
                A = self._stabilize_var_matrix(A)
                current_regime += 1
            
            X[t] = A @ X[t-1] + rng.randn(n_vars) * 0.1
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        # Add missing data if needed
        if missing_rate > 0:
            mask = rng.rand(n_samples, n_vars) < missing_rate
            data[mask] = np.nan
        
        true_graph = self._extract_graph_from_var(A, threshold=0.1)
        
        true_risks = {
            "nonstationarity": nonstationarity_risk,
            "irregularity": irregularity_risk,
            "persistence": persistence_risk,
            "confounding": confounding_risk
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_samples": n_samples,
            "violations": violations,
            "n_breaks": n_breaks,
            "max_eigenvalue": float(max_eigenvalue),
            "missing_rate": float(missing_rate)
        }
        
        return data, true_graph, true_risks, gen_params
    
    def _generate_extreme(self, rng: np.random.RandomState) -> Tuple:
        """
        F10: Edge cases for stress testing.
        
        Tests robustness of diagnostics.
        """
        n_vars = rng.choice([3, 4, 10, 12])  # Very small or large
        n_samples = rng.choice([200, 300, 1500, 2000])  # Very short or long
        
        # Random extreme case
        case_type = rng.choice(['very_short', 'very_sparse', 'near_unit_root', 'high_dim'])
        
        if case_type == 'very_short':
            n_samples = 200
            n_vars = 6
        elif case_type == 'very_sparse':
            # Lots of missing data
            missing_rate = 0.40
        elif case_type == 'near_unit_root':
            max_eigenvalue = 0.99
        else:  # high_dim
            n_vars = 12
            n_samples = 500
        
        # Generate base process
        A = self._generate_stable_var_matrix(n_vars, max_eigenvalue=rng.uniform(0.6, 0.9), rng=rng)
        
        X = np.zeros((n_samples, n_vars))
        X[0] = rng.randn(n_vars) * 0.1
        
        for t in range(1, n_samples):
            X[t] = A @ X[t-1] + rng.randn(n_vars) * 0.15
        
        data = pd.DataFrame(X, columns=[f"X{i}" for i in range(n_vars)])
        
        if case_type == 'very_sparse':
            mask = rng.rand(n_samples, n_vars) < missing_rate
            data[mask] = np.nan
        
        true_graph = self._extract_graph_from_var(A, threshold=0.1)
        
        # Risks depend on case type
        true_risks = {
            "nonstationarity": rng.uniform(0.3, 0.7),
            "irregularity": 0.7 if case_type == 'very_sparse' else 0.2,
            "persistence": 0.9 if case_type == 'near_unit_root' else 0.4,
            "confounding": rng.uniform(0.2, 0.6)
        }
        
        gen_params = {
            "n_vars": n_vars,
            "n_samples": n_samples,
            "case_type": case_type
        }
        
        return data, true_graph, true_risks, gen_params
    
    # ==================== UTILITY FUNCTIONS ====================
    
    def _generate_stable_var_matrix(self, n_vars: int, max_eigenvalue: float, 
                                    rng: np.random.RandomState) -> np.ndarray:
        """Generate a stable VAR(1) matrix with specified max eigenvalue."""
        # Generate random matrix
        A = rng.randn(n_vars, n_vars) * 0.3
        
        # Make sparse (50% zeros)
        mask = rng.rand(n_vars, n_vars) > 0.5
        A[mask] = 0
        
        # Scale to achieve desired max eigenvalue
        eigenvalues = np.linalg.eigvals(A)
        max_eig = np.max(np.abs(eigenvalues))
        
        if max_eig > 0:
            A = A * (max_eigenvalue / max_eig)
        
        return A
    
    def _stabilize_var_matrix(self, A: np.ndarray) -> np.ndarray:
        """Ensure VAR matrix is stable (max eigenvalue < 1)."""
        eigenvalues = np.linalg.eigvals(A)
        max_eig = np.max(np.abs(eigenvalues))
        
        if max_eig >= 1.0:
            A = A * (0.95 / max_eig)
        
        return A
    
    def _extract_graph_from_var(self, A: np.ndarray, threshold: float = 0.1) -> Dict:
        """Extract causal graph from VAR matrix."""
        n_vars = A.shape[0]
        graph = {}
        
        for i in range(n_vars):
            parents = []
            for j in range(n_vars):
                if abs(A[i, j]) > threshold:
                    parents.append(f"X{j}")
            graph[f"X{i}"] = parents
        
        return graph


def main():
    """Generate full extended atlas."""
    atlas = SyntheticAtlasExtended(
        output_dir="/Volumes/X10 Pro/autocause_resources/data/synthetic_dgp_atlas_v02",
        random_seed=42
    )
    
    atlas.generate_all(n_per_family=50)


if __name__ == "__main__":
    main()

