assets/quarter_sim/generate_data_mp_version_2.py

#!/usr/bin/env python3
"""
Multiprocessing data generation script for laminate simulations.
Loads configuration from YAML file and generates data using multiprocessing.
"""

import yaml
import numpy as np
from pathlib import Path
from itertools import product, combinations_with_replacement, combinations
from math import comb, ceil, factorial
from collections import Counter
import multiprocessing as mp
from functools import partial
import sys
from tqdm import tqdm
from datetime import datetime
import shutil
import json
import matplotlib.pyplot as plt
# Import from lam.py (don't modify original)
#lam_path = Path(__file__).parent 
lam_path = Path(__file__).parent 
sys.path.insert(0, str(lam_path))


# Import lam module (version from config, default "lam.py")
import importlib.util
spec = importlib.util.spec_from_file_location("lam", lam_path / "lam2.py")
lam = importlib.util.module_from_spec(spec)
spec.loader.exec_module(lam)

# Import needed functions and classes
read_dataset_metadata = lam.read_dataset_metadata
load_ud_material_from_files = lam.load_ud_material_from_files
build_full_symmetric_stack = lam.build_full_symmetric_stack
stack_label_from_upper = lam.stack_label_from_upper
Ply = lam.Ply
Laminate = lam.Laminate
T_eps = lam.T_eps
T_sigma = lam.T_sigma
T = lam.T
orthotropic_C_prime = lam.orthotropic_C_prime
COND_MAX = lam.COND_MAX


def parse_discrete_value(value):
    """
    Parse discrete value from YAML.
    Can be:
    - A list: [1, 2, 3]
    - A dict with start, end, interval: {start: 0, end: 90, interval: 10}
    """
    if isinstance(value, list):
        return value
    elif isinstance(value, dict):
        if 'start' in value and 'end' in value and 'interval' in value:
            start = float(value['start'])
            end = float(value['end'])
            interval = float(value['interval'])
            # Generate list from start to end (inclusive) with given interval
            result = []
            current = start
            while current <= end + 1e-9:  # Add small epsilon for floating point
                result.append(current)
                current += interval
            return result
        else:
            raise ValueError(f"Discrete dict must have 'start', 'end', and 'interval' keys")
    else:
        raise ValueError(f"Discrete value must be a list or dict, got {type(value)}")


def parse_continuous_value(value):
    """
    Parse continuous value from YAML.
    Must be a dict with 'num_samples', 'min', 'max'.
    """
    if isinstance(value, dict):
        if 'num_samples' in value and 'min' in value and 'max' in value:
            num_samples = int(value['num_samples'])
            min_val = float(value['min'])
            max_val = float(value['max'])
            return {
                'num_samples': num_samples,
                'min': min_val,
                'max': max_val
            }
        else:
            raise ValueError(f"Continuous value must have 'num_samples', 'min', and 'max' keys")
    else:
        raise ValueError(f"Continuous value must be a dict, got {type(value)}")

def make_sample(material_type, vol_fraction, quarter_angles):
    """
    Build one canonical sample dictionary for the balanced-symmetric quarter-angle workflow.

    Parameters
    ----------
    material_type : str
    vol_fraction : str | float
    quarter_angles : sequence of numbers
        Positive-angle quarter block. Order does not matter.

    Returns
    -------
    dict
        Sample dictionary with the new schema:
        - material_type
        - vol_fraction
        - quarter_layer_count
        - quarter_angles
        - half_layer_count
        - half_angles
        - center_angle
        - full_layer_count
        - full_angles
        - unique_angle_count
        - unique_angle_family
    """
    ordered_quarter_angles = tuple(sorted(float(a) for a in quarter_angles))
    half_angles = tuple(build_half_from_quarter_angles(ordered_quarter_angles))
    full_angles = tuple(build_full_from_quarter_angles(ordered_quarter_angles))
    unique_angle_family = tuple(sorted({abs(float(a)) for a in ordered_quarter_angles}))

    return {
        'material_type': material_type,
        'vol_fraction': vol_fraction,
        'quarter_layer_count': len(ordered_quarter_angles),
        'quarter_angles': list(ordered_quarter_angles),
        'half_layer_count': len(half_angles),
        'half_angles': list(half_angles),
        'center_angle': None,
        'full_layer_count': len(full_angles),
        'full_angles': list(full_angles),
        'unique_angle_count': len(unique_angle_family),
        'unique_angle_family': unique_angle_family,
    }

def parse_unique_k_proportions(config, max_unique_k):
    """
    Parse proportion weights for active feasible family sizes only.

    Important behavior:
    - Keys 1..max_unique_k are required.
    - Keys > max_unique_k are allowed and ignored.
    - Values are non-negative weights (not required to sum to 1).
    """
    proportion_cfg = config.get('unique_angle_k_proportions')

    if not isinstance(proportion_cfg, dict) or not proportion_cfg:
        raise ValueError(
            "Config must define 'unique_angle_k_proportions' "
            "as a non-empty dictionary."
        )

    for key in proportion_cfg.keys():
        try:
            key_int = int(key)
        except (TypeError, ValueError):
            raise ValueError(
                f"Invalid family key '{key}' in "
                "'unique_angle_k_proportions'. "
                "Keys must be integers like 1, 2, 3, ..."
            )
        if key_int < 1:
            raise ValueError(
                f"Invalid family key '{key}'. Family sizes must be >= 1."
            )

    proportions = {}
    for k in range(1, max_unique_k + 1):
        raw_value = proportion_cfg.get(k, proportion_cfg.get(str(k)))
        if raw_value is None:
            raise ValueError(
                f"Missing proportion for active family size k={k} "
                f"in 'unique_angle_k_proportions'."
            )
        value = float(raw_value)
        if value < 0.0:
            raise ValueError(
                f"Proportion weight for k={k} must be >= 0, got {value}."
            )
        proportions[k] = value

    if sum(proportions.values()) <= 0.0:
        raise ValueError(
            "At least one active k in 'unique_angle_k_proportions' must have a positive value."
        )

    return proportions
def same_angle(a, b, tol=1e-9):
    return abs(float(a) - float(b)) < tol


def is_self_balanced_angle(angle, tol=1e-9):
    """
    Angles that are self-balanced under sign reversal for this workflow.
    """
    a = abs(float(angle)) % 180.0
    return abs(a - 0.0) < tol or abs(a - 90.0) < tol


def balance_partner(angle):
    """
    Canonical balance partner for the quarter-angle workflow.
    Off-axis angles map to their sign-reversed partner.
    Self-balanced angles (0 and 90) map to themselves.
    """
    a = float(angle)
    if is_self_balanced_angle(a):
        a_mod = abs(a) % 180.0
        if same_angle(a_mod, 0.0):
            return 0.0
        if same_angle(a_mod, 90.0):
            return 90.0
        return a_mod
    return -a


def build_half_from_quarter_angles(quarter_angles):
    """
    Build the ordered half-stack from an unordered quarter-angle multiset.
    """
    ordered_quarter = tuple(sorted(float(a) for a in quarter_angles))
    mirrored_quarter = tuple(balance_partner(a) for a in reversed(ordered_quarter))
    return ordered_quarter + mirrored_quarter


def build_full_from_quarter_angles(quarter_angles):
    """
    Build the full balanced-symmetric stack from quarter angles.
    """
    half_angles = build_half_from_quarter_angles(quarter_angles)
    return tuple(build_full_symmetric_stack(half_angles, center_angle=None))


def count_distinct_quarter_angles(quarter_angles):
    """
    Number of distinct quarter-angle families in the quarter multiset.
    """
    return len({float(a) for a in quarter_angles})


def iter_positive_compositions(total, parts):
    """
    Yield all positive compositions of `total` into `parts` strictly positive integers.
    """
    if parts < 1 or total < parts:
        return

    if parts == 1:
        yield (total,)
        return

    for cuts in combinations(range(1, total), parts - 1):
        prev = 0
        comp = []
        for cut in cuts:
            comp.append(cut - prev)
            prev = cut
        comp.append(total - prev)
        yield tuple(comp)


def sample_positive_composition(total, parts, rng):
    """
    Uniformly sample one positive composition of `total` into `parts` parts.
    """
    if parts < 1 or total < parts:
        raise ValueError(f"Cannot compose total={total} into parts={parts}.")

    if parts == 1:
        return (total,)

    cuts = sorted(rng.choice(np.arange(1, total), size=parts - 1, replace=False).tolist())
    prev = 0
    comp = []
    for cut in cuts:
        comp.append(cut - prev)
        prev = cut
    comp.append(total - prev)
    return tuple(comp)
def generate_parameter_space(config):
    """
    Generate all parameter combinations for the balanced-symmetric quarter-angle workflow.

    Workflow:
    - material_types: exhaustive pool
    - vol_fractions: exhaustive supplied VF cases
    - candidate_angles: finite pool of positive quarter-angle candidates
    - quarter_layer_counts: allowed quarter-block sizes
    - total_samples: requested final sample count
    - unique_angle_k_proportions:
        relative weights for each unique-angle family size k in final dataset

    Returns:
    - List of sample dicts
    - Total number of samples
    - Count breakdown by k
    - Allocation report dictionary
    """
    workflow = config.get('sampling_workflow', 'family_percentage_exhaustive_mat_vf')
    if workflow != 'family_percentage_exhaustive_mat_vf':
        raise ValueError(
            "This version of generate_parameter_space only supports "
            "sampling_workflow='family_percentage_exhaustive_mat_vf'."
        )

    # ------------------------------------------------------------------
    # Parse exhaustive outer-loop variables
    # ------------------------------------------------------------------
    material_types = config['material_types']

    vf_config = config.get('vol_fractions', {})
    if isinstance(vf_config, dict) and 'values' in vf_config:
        vol_fractions = parse_discrete_value(vf_config['values'])
    else:
        vol_fractions = parse_discrete_value(vf_config)

    if not vol_fractions:
        raise ValueError("vol_fractions cannot be empty.")

    # ------------------------------------------------------------------
    # Parse candidate quarter-angle pool
    # ------------------------------------------------------------------
    angles_config = config.get('candidate_angles', {})
    if isinstance(angles_config, dict) and 'type' in angles_config:
        if angles_config['type'] != 'discrete':
            raise ValueError(
                "In the balanced-symmetric quarter-angle workflow, "
                "candidate_angles must define a finite angle pool."
            )
        candidate_angles = parse_discrete_value(angles_config['values'])
    else:
        candidate_angles = parse_discrete_value(angles_config)

    if not candidate_angles:
        raise ValueError("candidate_angles cannot be empty.")

    candidate_angles = sorted(float(a) for a in candidate_angles)

    if len(set(candidate_angles)) != len(candidate_angles):
        raise ValueError(
            "candidate_angles must contain unique values for the quarter-angle workflow."
        )

    # ------------------------------------------------------------------
    # Parse allowed quarter-block sizes
    # ------------------------------------------------------------------
    raw_quarter_layer_counts = parse_discrete_value(config.get('quarter_layer_counts', [1]))
    quarter_layer_counts = sorted(int(n) for n in raw_quarter_layer_counts)

    if not quarter_layer_counts:
        raise ValueError("quarter_layer_counts cannot be empty.")

    if any(n < 1 for n in quarter_layer_counts):
        raise ValueError("All quarter_layer_counts must be >= 1.")

    max_quarter_layers = max(quarter_layer_counts)

    # ------------------------------------------------------------------
    # Determine maximum unique-angle family size
    # ------------------------------------------------------------------
    proportion_cfg = config.get('unique_angle_k_proportions', {})
    if not isinstance(proportion_cfg, dict) or not proportion_cfg:
        raise ValueError(
            "Config must define 'unique_angle_k_proportions' "
            "as a non-empty dictionary."
        )

    try:
        requested_k_values = sorted(int(k) for k in proportion_cfg.keys())
    except (TypeError, ValueError):
        raise ValueError(
            "'unique_angle_k_proportions' keys must be integers like 1, 2, 3, ..."
        )

    if not requested_k_values or requested_k_values[0] < 1:
        raise ValueError(
            "'unique_angle_k_proportions' must start from k=1."
        )

    max_unique_k_requested = max(requested_k_values)
    max_unique_k_allowed = min(max_quarter_layers, len(candidate_angles))

    active_max_unique_k = min(max_unique_k_requested, max_unique_k_allowed)
    ignored_k_values = [k for k in requested_k_values if k > active_max_unique_k]

    if active_max_unique_k < 1:
        raise ValueError(
            "No feasible unique-angle family sizes are available. "
            "Check candidate_angles and quarter_layer_counts."
        )

    proportion_by_k = parse_unique_k_proportions(config, active_max_unique_k)

    total_samples_requested = config.get('total_samples')
    if total_samples_requested is None:
        raise ValueError(
            "Config must define 'total_samples' for workflow 'family_percentage_exhaustive_mat_vf'."
        )
    total_samples_requested = int(total_samples_requested)
    if total_samples_requested < 1:
        raise ValueError("total_samples must be >= 1.")

    # ------------------------------------------------------------------
    # Reproducibility
    # ------------------------------------------------------------------
    random_seed = config.get('random_seed')
    rng = np.random.default_rng(random_seed)

    # ------------------------------------------------------------------
    # Capacity by k
    # ------------------------------------------------------------------
    mat_vf_multiplier = len(material_types) * len(vol_fractions)

    capacity_by_k = {k: 0 for k in range(1, active_max_unique_k + 1)}
    capacity_by_k_and_q = {k: {} for k in range(1, active_max_unique_k + 1)}

    for k in range(1, active_max_unique_k + 1):
        for q in quarter_layer_counts:
            if k > q:
                continue

            quarter_multiset_count = comb(len(candidate_angles), k) * comb(q - 1, k - 1)
            capacity_qk = mat_vf_multiplier * quarter_multiset_count

            if capacity_qk <= 0:
                continue

            capacity_by_k[k] += capacity_qk
            capacity_by_k_and_q[k][q] = capacity_qk

    total_capacity = sum(capacity_by_k.values())
    normalized_total = sum(proportion_by_k.values())
    normalized_proportion_by_k = {
        k: (proportion_by_k[k] / normalized_total) for k in range(1, active_max_unique_k + 1)
    }

    # Initial requested targets by k via largest remainder method
    raw_targets = {
        k: total_samples_requested * normalized_proportion_by_k[k]
        for k in range(1, active_max_unique_k + 1)
    }
    requested_target_by_k = {k: int(raw_targets[k]) for k in raw_targets}
    fractional_order = sorted(
        range(1, active_max_unique_k + 1),
        key=lambda k: (raw_targets[k] - requested_target_by_k[k], normalized_proportion_by_k[k], -k),
        reverse=True,
    )
    remainder = total_samples_requested - sum(requested_target_by_k.values())
    for k in fractional_order[:remainder]:
        requested_target_by_k[k] += 1

    # Cap by per-k capacity
    allocated_by_k = {
        k: min(requested_target_by_k[k], capacity_by_k[k])
        for k in range(1, active_max_unique_k + 1)
    }

    # Redistribute shortfall to k with remaining capacity
    shortfall = total_samples_requested - sum(allocated_by_k.values())
    while shortfall > 0:
        eligible = [k for k in range(1, active_max_unique_k + 1) if allocated_by_k[k] < capacity_by_k[k]]
        if not eligible:
            break

        weight_sum = sum(normalized_proportion_by_k[k] for k in eligible)
        if weight_sum <= 0:
            per_weight = {k: 1.0 / len(eligible) for k in eligible}
        else:
            per_weight = {k: normalized_proportion_by_k[k] / weight_sum for k in eligible}

        raw_extra = {k: shortfall * per_weight[k] for k in eligible}
        base_extra = {}
        used = 0
        for k in eligible:
            room = capacity_by_k[k] - allocated_by_k[k]
            add = min(room, int(raw_extra[k]))
            base_extra[k] = add
            used += add

        for k, add in base_extra.items():
            allocated_by_k[k] += add
        shortfall -= used
        if shortfall <= 0:
            break

        eligible_after_base = [k for k in eligible if allocated_by_k[k] < capacity_by_k[k]]
        if not eligible_after_base:
            break

        fractional_sorted = sorted(
            eligible_after_base,
            key=lambda k: (raw_extra[k] - int(raw_extra[k]), per_weight[k], -k),
            reverse=True,
        )
        if used == 0:
            fractional_sorted = sorted(
                eligible_after_base,
                key=lambda k: (capacity_by_k[k] - allocated_by_k[k], per_weight[k], -k),
                reverse=True,
            )

        progress = False
        for k in fractional_sorted:
            if shortfall <= 0:
                break
            if allocated_by_k[k] >= capacity_by_k[k]:
                continue
            allocated_by_k[k] += 1
            shortfall -= 1
            progress = True

        if not progress:
            break

    # ------------------------------------------------------------------
    # Build final sample list exactly per-k allocations
    # ------------------------------------------------------------------
    sample_pool_enumeration_threshold = 1_000_000
    samples = []
    count_breakdown_by_k = {k: 0 for k in range(1, active_max_unique_k + 1)}

    def generate_all_candidates_for_k(k):
        for mat_type in material_types:
            for vf in vol_fractions:
                for quarter_layer_count in quarter_layer_counts:
                    if quarter_layer_count < k:
                        continue

                    for quarter_angles in combinations_with_replacement(candidate_angles, quarter_layer_count):
                        if count_distinct_quarter_angles(quarter_angles) != k:
                            continue

                        yield make_sample(
                            material_type=mat_type,
                            vol_fraction=vf,
                            quarter_angles=quarter_angles,
                        )

    for k in range(1, active_max_unique_k + 1):
        target_k = allocated_by_k[k]
        if target_k <= 0:
            continue

        capacity_k = capacity_by_k[k]
        exhaustive_for_k = target_k >= capacity_k

        if exhaustive_for_k or capacity_k <= sample_pool_enumeration_threshold:
            candidate_samples = list(generate_all_candidates_for_k(k))
            if exhaustive_for_k:
                selected_samples = candidate_samples
            else:
                selected_idx = rng.choice(len(candidate_samples), size=target_k, replace=False)
                selected_samples = [candidate_samples[i] for i in np.sort(selected_idx)]
            samples.extend(selected_samples)
            count_breakdown_by_k[k] += len(selected_samples)
            continue

        # Large-space random unique sampling for this k
        seen_keys = set()
        material_indices = np.arange(len(material_types))
        vf_indices = np.arange(len(vol_fractions))
        angle_indices = np.arange(len(candidate_angles))

        q_choices = np.array(sorted(capacity_by_k_and_q[k].keys()), dtype=int)
        q_weights = np.array([capacity_by_k_and_q[k][q] for q in q_choices], dtype=float)
        q_weights = q_weights / q_weights.sum()

        max_attempts = max(50_000, target_k * 300)
        attempts = 0

        while len(seen_keys) < target_k and attempts < max_attempts:
            mat_idx = int(rng.choice(material_indices))
            vf_idx = int(rng.choice(vf_indices))
            quarter_layer_count = int(rng.choice(q_choices, p=q_weights))

            family_idx_tuple = tuple(sorted(rng.choice(angle_indices, size=k, replace=False)))
            family_angles = tuple(candidate_angles[i] for i in family_idx_tuple)
            family_counts = sample_positive_composition(quarter_layer_count, k, rng)

            quarter_angles = []
            for angle, count in zip(family_angles, family_counts):
                quarter_angles.extend([angle] * count)

            sample = make_sample(
                material_type=material_types[mat_idx],
                vol_fraction=vol_fractions[vf_idx],
                quarter_angles=quarter_angles,
            )

            key = (
                sample['material_type'],
                format_vol_fraction(sample['vol_fraction']),
                tuple(sample['quarter_angles']),
            )

            if key not in seen_keys:
                seen_keys.add(key)
                samples.append(sample)
                count_breakdown_by_k[k] += 1

            attempts += 1

        if count_breakdown_by_k[k] < target_k:
            raise RuntimeError(
                f"Could not reach allocated target for k={k}. "
                f"Requested={target_k:,}, generated={count_breakdown_by_k[k]:,}, "
                f"capacity={capacity_k:,}. Try lowering total_samples or adjusting proportions."
            )

    samples.sort(
        key=lambda s: (
            s['material_type'],
            format_vol_fraction(s['vol_fraction']),
            s['unique_angle_count'],
            s['quarter_layer_count'],
            tuple(s['quarter_angles']),
            tuple(s['full_angles'])
        )
    )

    total_samples = len(samples)
    allocation_report = {
        'requested_total_samples': total_samples_requested,
        'total_capacity': total_capacity,
        'global_shortfall': max(0, total_samples_requested - total_capacity),
        'ignored_k_values': ignored_k_values,
        'active_max_unique_k': active_max_unique_k,
        'quarter_layer_counts': quarter_layer_counts,
        'normalized_proportion_by_k': normalized_proportion_by_k,
        'requested_target_by_k': requested_target_by_k,
        'capacity_by_k': capacity_by_k,
        'allocated_by_k': allocated_by_k,
        'generated_by_k': count_breakdown_by_k.copy(),
    }
    return samples, total_samples, count_breakdown_by_k, allocation_report


def format_vol_fraction(vf):
    """
    Format volume fraction as a 4-decimal string for file naming.
    Examples:
        0.1       -> "0.1000"
        "0.1"     -> "0.1000"
        "0.1000"  -> "0.1000"
    """
    try:
        return f"{float(vf):.4f}"
    except (TypeError, ValueError):
        return str(vf)

def save_sample_plot(plot_data, meta, save_path):
    """
    Save one 7-panel plot for a single generated sample.

    Expected plot_data structure:
        plot_data = {
            "11": {
                "x": np.ndarray,
                "stress": np.ndarray,
                "lateral": np.ndarray,
                "eps33": np.ndarray,
            },
            "22": {
                "x": np.ndarray,
                "stress": np.ndarray,
                "lateral": np.ndarray,
                "eps33": np.ndarray,
            },
            "12": {
                "x": np.ndarray,
                "stress": np.ndarray,
            },
        }

    Expected meta structure:
        meta = {
            "material_type": str,
            "vol_fraction": str,
            "quarter_angles": list,
        }
    """
    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)

    fig, axes = plt.subplots(3, 3, figsize=(16, 16))
    ax = axes

    # ---------------- Row 1: stress ----------------
    if "11" in plot_data:
        ax[0, 0].plot(plot_data["11"]["x"], plot_data["11"]["stress"], marker="o")
    ax[0, 0].set_title("Mode 11 - Stress")
    ax[0, 0].set_xlabel("Strain")
    ax[0, 0].set_ylabel("Stress (MPa)")
    ax[0, 0].grid(True, alpha=0.3)

    if "22" in plot_data:
        ax[0, 1].plot(plot_data["22"]["x"], plot_data["22"]["stress"], marker="o")
    ax[0, 1].set_title("Mode 22 - Stress")
    ax[0, 1].set_xlabel("Strain")
    ax[0, 1].set_ylabel("Stress (MPa)")
    ax[0, 1].grid(True, alpha=0.3)

    if "12" in plot_data:
        ax[0, 2].plot(plot_data["12"]["x"], plot_data["12"]["stress"], marker="o")
    ax[0, 2].set_title("Mode 12 - Stress")
    ax[0, 2].set_xlabel("Strain")
    ax[0, 2].set_ylabel("Stress (MPa)")
    ax[0, 2].grid(True, alpha=0.3)

    # ---------------- Row 2: lateral ----------------
    if "11" in plot_data:
        ax[1, 0].plot(plot_data["11"]["x"], plot_data["11"]["lateral"], marker="o")
    ax[1, 0].set_title("Mode 11 - Lateral")
    ax[1, 0].set_xlabel("Strain")
    ax[1, 0].set_ylabel("Lateral Strain")
    ax[1, 0].grid(True, alpha=0.3)

    if "22" in plot_data:
        ax[1, 1].plot(plot_data["22"]["x"], plot_data["22"]["lateral"], marker="o")
    ax[1, 1].set_title("Mode 22 - Lateral")
    ax[1, 1].set_xlabel("Strain")
    ax[1, 1].set_ylabel("Lateral Strain")
    ax[1, 1].grid(True, alpha=0.3)

    ax[1, 2].axis("off")

    # ---------------- Row 3: eps_33 ----------------
    if "11" in plot_data:
        ax[2, 0].plot(plot_data["11"]["x"], plot_data["11"]["eps33"], marker="o")
    ax[2, 0].set_title("Mode 11 - eps_33")
    ax[2, 0].set_xlabel("Strain")
    ax[2, 0].set_ylabel("eps_33")
    ax[2, 0].grid(True, alpha=0.3)

    if "22" in plot_data:
        ax[2, 1].plot(plot_data["22"]["x"], plot_data["22"]["eps33"], marker="o")
    ax[2, 1].set_title("Mode 22 - eps_33")
    ax[2, 1].set_xlabel("Strain")
    ax[2, 1].set_ylabel("eps_33")
    ax[2, 1].grid(True, alpha=0.3)

    ax[2, 2].axis("off")

    material_type = meta.get("material_type", "Unknown")
    vol_fraction = meta.get("vol_fraction", "Unknown")
    quarter_angles = meta.get("quarter_angles", [])

    fig.suptitle(
        f"Material: {material_type}, VF={vol_fraction}, Q-angles={quarter_angles}",
        fontsize=12
    )

    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fig.savefig(save_path, dpi=200, bbox_inches="tight")
    plt.close(fig)
def worker_function(sample, curve_dir, out_dir, config):
    """
    Worker function for multiprocessing.
    Generates one output file for a given parameter combination.
    """
    try:
        mat_type = sample['material_type']
        vf = sample['vol_fraction']
        quarter_angles = sample['quarter_angles']
        half_angles = sample['half_angles']
        center_angle = sample['center_angle']
        full_angles = sample['full_angles']

        vf_str = format_vol_fraction(vf)
        prefix = f"{mat_type}_{vf_str}"

        # Load material (this reads from files)
        curve_dir_path = Path(curve_dir)

        # Temporarily modify CURVE_DIR in lam module
        original_curve_dir = lam.CURVE_DIR
        lam.CURVE_DIR = curve_dir_path

        try:
            metadata_result = read_dataset_metadata(prefix)

            if not isinstance(metadata_result, tuple) or len(metadata_result) != 3:
                lam.CURVE_DIR = original_curve_dir
                return False, (
                    f"read_dataset_metadata returned unexpected format for {prefix}: "
                    f"{type(metadata_result)}, length: "
                    f"{len(metadata_result) if hasattr(metadata_result, '__len__') else 'N/A'}"
                )

            vf_meta, centers_meta, n_fibers = metadata_result
            mat = load_ud_material_from_files(prefix)

        except Exception as e:
            lam.CURVE_DIR = original_curve_dir
            import traceback
            return False, (
                f"Failed to load material for {prefix}: "
                f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
            )
        finally:
            lam.CURVE_DIR = original_curve_dir

        
        # Use quarter angles for the file name, but keep the full stack label for metadata.
        quarter_label_file = "_".join(
            str(int(round(float(a))))
            for a in quarter_angles
        )
        stack_label_file = quarter_label_file
        
        stack_label_human = stack_label_from_upper(full_angles)

        combined_blocks = {}
        plot_data = {}

        save_sample_plots = config.get('save_sample_plots', False)
        plot_out_dir_path = None

        if save_sample_plots:
            plot_output_directory = config.get('plot_output_directory')
            if plot_output_directory:
                plot_out_dir_path = Path(plot_output_directory)
            else:
                plot_out_dir_path = Path(str(out_dir) + "_plots")

            plot_out_dir_path.mkdir(parents=True, exist_ok=True)

        # Run simulations for all three modes
        mode_errors = []
        for mode in ("11", "22", "12"):
            try:
                result = run_simulation_with_mat(prefix, full_angles, mode, mat)

                # Verify we got the expected number of return values
                if len(result) != 8:
                    raise ValueError(f"Expected 8 return values, got {len(result)}")

                ex, sx, ey, gxy, ezz, g23, g13, e11 = result

            except Exception as e:
                # Skip this mode if it fails
                import traceback
                tb_str = traceback.format_exc()
                error_detail = f"{type(e).__name__}: {str(e)}"

                if "not enough values to unpack" in str(e):
                    error_detail += f"\nTraceback:\n{tb_str}"

                mode_errors.append(f"Mode {mode}: {error_detail}")
                continue

            # Process output (configurable number of points)
            num_output_points = config.get('num_output_points', 10)
            x_out = np.linspace(ex[0], ex[-1], num_output_points)
            sx_out_MPa = np.interp(x_out, ex, sx / 1e6)

            if mode == "11":
                strain_label = "eps_11"
                stress_label = "sig_11"
                lateral_label = "eps_22"
                lateral_series = ey
                e33_label = "eps_33"
                e33_series = ezz
            elif mode == "22":
                strain_label = "eps_22"
                stress_label = "sig_22"
                lateral_label = "eps_11"
                lateral_series = e11
                e33_label = "eps_33"
                e33_series = ezz
            elif mode == "12":
                strain_label = "eps_12"
                stress_label = "sig_12"
                lateral_label = None
                lateral_series = None
                e33_label = None
                e33_series = None
            else:
                raise ValueError(f"Unknown mode {mode}")

            if lateral_series is not None:
                lateral_out = np.interp(x_out, ex, lateral_series)
            else:
                lateral_out = None

            if e33_series is not None:
                e33_out = np.interp(x_out, ex, e33_series)
            else:
                e33_out = None

            if lateral_label is None:
                header_line = f"{strain_label:<8}  {stress_label:<8}"
            else:
                header_line = (
                    f"{strain_label:<8}  {stress_label:<8}  "
                    f"{lateral_label:<8}  {e33_label:<8}"
                )

            rows = []
            if lateral_label is None:
                for eps_val, sig_val in zip(x_out, sx_out_MPa):
                    line = f"{eps_val:8.6f}  {sig_val:8.3f}"
                    rows.append(line)
            else:
                for eps_val, sig_val, lat_val, e33_val in zip(
                    x_out, sx_out_MPa, lateral_out, e33_out
                ):
                    line = (
                        f"{eps_val:8.6f}  {sig_val:8.3f}  "
                        f"{lat_val:8.6f}  {e33_val:8.6f}"
                    )
                    rows.append(line)

            combined_blocks[mode] = (header_line, rows)
            if mode in ("11", "22"):
                plot_data[mode] = {
                    "x": np.array(x_out, copy=True),
                    "stress": np.array(sx_out_MPa, copy=True),
                    "lateral": np.array(lateral_out, copy=True),
                    "eps33": np.array(e33_out, copy=True),
                }
            elif mode == "12":
                plot_data[mode] = {
                    "x": np.array(x_out, copy=True),
                    "stress": np.array(sx_out_MPa, copy=True),
                }
        # Write output file
        out_dir_path = Path(out_dir)
        out_dir_path.mkdir(parents=True, exist_ok=True)
        combined_file = out_dir_path / (
            f"{mat_type}_{vf_str}_{stack_label_file}.txt"
        )

        # Only write file if at least one mode succeeded
        if not combined_blocks:
            error_msg = "All modes failed for this sample"
            if mode_errors:
                error_msg += f" - Errors: {'; '.join(mode_errors)}"
            return False, error_msg

        with open(combined_file, "w") as fc:
            # Write 11, then 22, then 12 in order
            for m in ("11", "22", "12"):
                if m in combined_blocks:
                    header_line, rows = combined_blocks[m]
                    fc.write(header_line + "\n")
                    for line in rows:
                        fc.write(line + "\n")
                    fc.write("\n")

            # One common metadata block at the end
            fc.write(f"volume fraction= {vf_meta:.6f}\n")
            fc.write(f"material type= {mat_type}\n")
            fc.write("loading modes= 11, 22, 12\n")
            fc.write(f"stacking sequence= {stack_label_human}\n")
            fc.write(f"number of fibers= {n_fibers}\n")
            fc.write(f"fiber_centers_YZ={centers_meta}\n")

        if save_sample_plots and plot_data:
            plot_meta = {
                "material_type": mat_type,
                "vol_fraction": vf_str,
                "quarter_angles": quarter_angles,
            }

            plot_file = plot_out_dir_path / (
                f"{mat_type}_{vf_str}_{stack_label_file}.png"
            )

            save_sample_plot(plot_data, plot_meta, plot_file)

        return True, None

    except Exception as e:
        return False, str(e)


def run_simulation_with_mat(prefix, full_angles, mode, mat):
    """
    Wrapper to run simulation with material object.
    This works around the issue that run_uniaxial_test_from_files_5d uses mat
    but doesn't receive it as parameter.
    """
    
    import numpy as np
    
    tply = 0.05
    
    plies = [Ply(angle_deg, tply, mat) for angle_deg in full_angles]
    laminate = Laminate(plies)
    
    if mode == "11":
        main_index = 0
        eps_max = 0.10
    elif mode == "22":
        main_index = 1
        eps_max = 0.10
    elif mode == "12":
        main_index = 5
        eps_max = 0.20
    else:
        raise ValueError("mode must be '11', '22' or '12'")
    
    main_steps = np.linspace(0.0, eps_max, 1500)
    
    ex_hist, sx_hist = [], []
    ey_hist, gxy_hist = [], []
    ezz_hist, g23_hist, g13_hist = [], [], []
    e11_hist = []
    
    ex_prev = 0.0
    ey_prev = 0.0
    gxy_prev = 0.0
    ezz_prev = 0.0
    g23_prev = 0.0
    g13_prev = 0.0
    s1_prev = 0.0
    
    for i in range(1, len(main_steps)):
        main_target = main_steps[i]
        
        if main_index == 0:
            main_prev = ex_prev
        elif main_index == 1:
            main_prev = ey_prev
        else:
            main_prev = gxy_prev
        
        dmain = main_target - main_prev
        
        Ceff_prev = laminate.effective_C_from_previous_strains(
            ex_prev, ey_prev, ezz_prev, g23_prev, g13_prev, gxy_prev
        )
        
        cond = np.linalg.cond(Ceff_prev)
        if not np.isfinite(cond) or cond > COND_MAX:
            raise np.linalg.LinAlgError(
                f"Effective C is ill-conditioned (cond={cond:.3e}) at step {i}"
            )
        
        e_j = np.zeros(6)
        e_j[main_index] = 1.0
        
        try:
            S_col = np.linalg.solve(Ceff_prev, e_j)
        except np.linalg.LinAlgError as err:
            raise np.linalg.LinAlgError(
                f"Failed to solve for compliance column at step {i}: {err}"
            )
        
        Sjj = S_col[main_index]
        if abs(Sjj) < 1e-20:
            raise ZeroDivisionError(
                f"Sjj is zero or too small at step {i} (Sjj={Sjj:.3e})."
            )
        
        ds1 = dmain / Sjj
        de_vec = S_col * ds1
        de_vec[main_index] = dmain
        
        de1, de2, de3, de4, de5, de6 = de_vec
        de4 = 0.0
        de5 = 0.0
        
        s1 = s1_prev + ds1
        ex = ex_prev + de1
        ey = ey_prev + de2
        ezz = ezz_prev + de3
        g23 = g23_prev + de4
        g13 = g13_prev + de5
        gxy = gxy_prev + de6
        
        laminate.update_fiber_angles_incremental(de1, de2, de6)
        
        ex_prev, ey_prev, ezz_prev = ex, ey, ezz
        g23_prev, g13_prev, gxy_prev = g23, g13, gxy
        s1_prev = s1
        
        if main_index == 0:
            main_strain = ex
        elif main_index == 1:
            main_strain = ey
        else:
            main_strain = 0.5 * gxy
        
        ex_hist.append(main_strain)
        sx_hist.append(s1)
        ey_hist.append(ey)
        gxy_hist.append(gxy)
        ezz_hist.append(ezz)
        g23_hist.append(g23)
        g13_hist.append(g13)
        e11_hist.append(ex)
    
    # Ensure all arrays are created properly
    result = (
        np.array(ex_hist),
        np.array(sx_hist),
        np.array(ey_hist),
        np.array(gxy_hist),
        np.array(ezz_hist),
        np.array(g23_hist),
        np.array(g13_hist),
        np.array(e11_hist)
    )
    
    # Verify we're returning 8 values
    if len(result) != 8:
        raise ValueError(f"Internal error: expected 8 return values, got {len(result)}")
    
    return result


def main():
    if len(sys.argv) < 2:
        print("Usage: python generate_data_mp.py <config.yaml>")
        sys.exit(1)
    
    config_path = Path(sys.argv[1])
    if not config_path.exists():
        print(f"Error: Config file not found: {config_path}")
        sys.exit(1)
    
    # Load configuration
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    # Get paths
    curve_dir = Path(config.get('input_directory', 'shahriar_modified_2025_12/RVE_Datasets'))
    out_dir = Path(config.get('output_directory', 'shahriar_modified_2025_12/Output_directory'))
    num_processes = config.get('num_processes', mp.cpu_count())
    
    # Generate parameter space
    print("Generating parameter space...")
    samples, total_samples, count_breakdown_by_k, allocation_report = generate_parameter_space(config)
    
    # New workflow summary
    sampling_mode = "Proportion-Based Unique-Angle Allocation"

    vf_config = config.get('vol_fractions', {})
    angles_config = config.get('candidate_angles', {})
    proportion_cfg = config.get('unique_angle_k_proportions', {})
    quarter_layer_counts = allocation_report.get(
        'quarter_layer_counts',
        config.get('quarter_layer_counts', [])
    )
    
    # Parse candidate angle count for display
    if isinstance(angles_config, dict) and 'values' in angles_config:
        display_candidate_angles = parse_discrete_value(angles_config['values'])
    else:
        display_candidate_angles = parse_discrete_value(angles_config)
    
    # Display values should come from the actual quarter-angle workflow
    max_full_layers = 4 * max(quarter_layer_counts) if quarter_layer_counts else 0
    requested_k_values = sorted(int(k) for k in proportion_cfg.keys()) if proportion_cfg else [1]
    max_unique_k_requested = max(requested_k_values) if requested_k_values else 1
    
    # These must come from the actual generated allocation report,
    # not be recomputed using obsolete full_layer_counts logic
    active_max_unique_k = allocation_report['active_max_unique_k']
    ignored_k_values = allocation_report['ignored_k_values']

    print("\n" + "="*70)
    print("GENERATION SUMMARY")
    print("="*70)
    print(f"Sampling mode: {sampling_mode}")
    print(f"Requested total samples: {allocation_report['requested_total_samples']:,}")
    print(f"Total number of samples (computed before generation): {total_samples:,}")
    print(f"Global feasible capacity: {allocation_report['total_capacity']:,}")
    if allocation_report['global_shortfall'] > 0:
        print(f"Unfillable shortfall (capacity limit): {allocation_report['global_shortfall']:,}")
    print(f"Number of processes: {num_processes}")
    print(f"Number of output points per curve: {config.get('num_output_points', 10)}")
    print(f"Input directory: {curve_dir}")
    print(f"Output directory: {out_dir}")
    print(f"Maximum full layers: {max_full_layers}")
    print(f"Requested maximum unique-angle family size: {max_unique_k_requested}")
    print(f"Feasible maximum unique-angle family size: {active_max_unique_k}")
    print(f"Candidate angle pool size: {len(display_candidate_angles)}")
    if ignored_k_values:
        print(f"Ignored family sizes (not feasible with current candidate angle pool): {ignored_k_values}")
    print("Unique-angle proportions and allocation details:")
    for k in range(1, active_max_unique_k + 1):
        raw_val = proportion_cfg[str(k)] if str(k) in proportion_cfg else proportion_cfg[k]
        norm_prop = allocation_report['normalized_proportion_by_k'][k]
        requested_target = allocation_report['requested_target_by_k'][k]
        capacity = allocation_report['capacity_by_k'][k]
        allocated = allocation_report['allocated_by_k'][k]
        generated = count_breakdown_by_k.get(k, 0)
        print(
            f"  k={k}: weight={float(raw_val):.6g}, normalized={norm_prop:.2%}, "
            f"requested={requested_target:,}, capacity={capacity:,}, "
            f"allocated={allocated:,}, generated={generated:,}"
        )

    print("Sample count breakdown by active family size:")
    for k in range(1, active_max_unique_k + 1):
        print(f"  k={k}: {count_breakdown_by_k.get(k, 0):,}")
    
    print(f"Total samples from breakdown: {sum(count_breakdown_by_k.values()):,}")
    print(f"Total allocated by plan: {sum(allocation_report['allocated_by_k'].values()):,}")
    print(f"Total number of samples: {total_samples:,}")
    
    print("="*70)
    
    # Ask for confirmation (skip if running in non-interactive mode)
    if sys.stdin.isatty():
        response = input("\nProceed with generation? (yes/no): ").strip().lower()
        if response not in ['yes', 'y']:
            print("Generation cancelled.")
            sys.exit(0)
    else:
        print("\nNon-interactive mode: proceeding with generation...")
    
    # Create output directory if it doesn't exist
    out_dir.mkdir(parents=True, exist_ok=True)
    
    # Copy config file to output directory for reference
    config_copy_path = out_dir / "generation_config.yaml"
    shutil.copy2(config_path, config_copy_path)
    print(f"Config file copied to: {config_copy_path}")
    
    # Record start time
    start_time = datetime.now()
    start_timestamp = start_time.isoformat()
    
    # Create worker function with fixed arguments
    worker = partial(worker_function, 
                     curve_dir=str(curve_dir),
                     out_dir=str(out_dir),
                     config=config)
    
    # Run multiprocessing
    print(f"\nStarting generation with {num_processes} processes...\n")
    
    completed = 0
    failed = 0
    errors_list = []
    
    with mp.Pool(processes=num_processes) as pool:
        # Use tqdm for progress bar
        results = pool.imap(worker, samples)
        
        with tqdm(total=total_samples, desc="Generating samples", unit="sample") as pbar:
            for i, (success, error) in enumerate(results, 1):
                if success:
                    completed += 1
                else:
                    failed += 1
                    if error:
                        errors_list.append(f"Sample {i}: {error}")
                
                # Update progress bar with current stats
                pbar.set_postfix({
                    'Completed': completed,
                    'Failed': failed,
                    'Success %': f"{100*completed/i:.1f}%"
                })
                pbar.update(1)
    
    # Record end time
    end_time = datetime.now()
    end_timestamp = end_time.isoformat()
    duration = end_time - start_time
    duration_seconds = duration.total_seconds()
    duration_hours = duration_seconds / 3600
    duration_minutes = (duration_seconds % 3600) / 60
    
    # Format duration string
    if duration_seconds < 60:
        duration_str = f"{duration_seconds:.2f} seconds"
    elif duration_seconds < 3600:
        duration_str = f"{int(duration_minutes)} minutes {duration_seconds % 60:.2f} seconds"
    else:
        duration_str = f"{int(duration_hours)} hours {int(duration_minutes)} minutes {duration_seconds % 60:.2f} seconds"
    
    # Calculate statistics
    success_rate = 100 * completed / total_samples if total_samples > 0 else 0
    avg_time_per_sample = duration_seconds / total_samples if total_samples > 0 else 0
    
    print("\n" + "="*70)
    print("GENERATION COMPLETE")
    print("="*70)
    print(f"Total samples: {total_samples}")
    print(f"Successfully completed: {completed}")
    print(f"Failed: {failed}")
    print(f"Success rate: {success_rate:.2f}%")
    print(f"Total duration: {duration_str}")
    print(f"Average time per sample: {avg_time_per_sample:.2f} seconds")
    print(f"Output directory: {out_dir}")
    
    if errors_list:
        print(f"\nFirst {min(10, len(errors_list))} errors:")
        for err in errors_list[:10]:
            print(f"  - {err}")
        if len(errors_list) > 10:
            print(f"  ... and {len(errors_list) - 10} more errors")
    
    print("="*70)
    
    # Create summary file
    summary = {
        'generation_info': {
            'start_time': start_timestamp,
            'end_time': end_timestamp,
            'duration_seconds': duration_seconds,
            'duration_formatted': duration_str,
            'sampling_mode': sampling_mode,
            'requested_max_unique_k': max_unique_k_requested,
            'active_max_unique_k': allocation_report['active_max_unique_k'],
            'ignored_k_values': allocation_report['ignored_k_values'],
        },
        'statistics': {
            'total_samples': total_samples,
            'successfully_completed': completed,
            'failed': failed,
            'success_rate_percent': round(success_rate, 2),
            'average_time_per_sample_seconds': round(avg_time_per_sample, 2),
        },
        'configuration': {
            'num_processes': num_processes,
            'num_output_points': config.get('num_output_points', 10),
            'input_directory': str(curve_dir),
            'output_directory': str(out_dir),
            'sampling_workflow': config.get('sampling_workflow'),
            'material_types': config.get('material_types', []),
            'vol_fractions': config.get('vol_fractions', {}),
            'candidate_angles': config.get('candidate_angles', {}),
            'quarter_layer_counts': allocation_report.get(
                'quarter_layer_counts',
                config.get('quarter_layer_counts', [])
            ),
            'max_full_layers': max_full_layers,
            'total_samples': config.get('total_samples'),
            'unique_angle_k_proportions': config.get('unique_angle_k_proportions', {}),
            'random_seed': config.get('random_seed'),
        },
        'allocation_report': allocation_report,
        'errors': {
            'total_errors': len(errors_list),
            'error_samples': errors_list[:50] if len(errors_list) <= 50 else errors_list[:50] + [f"... and {len(errors_list) - 50} more errors"]
        }
    }
    
    # Save summary as both YAML and JSON for flexibility
    summary_yaml_path = out_dir / "generation_summary.yaml"
    summary_json_path = out_dir / "generation_summary.json"
    
    with open(summary_yaml_path, 'w') as f:
        yaml.dump(summary, f, default_flow_style=False, sort_keys=False)
    
    with open(summary_json_path, 'w') as f:
        json.dump(summary, f, indent=2, sort_keys=False)
    
    print(f"\nSummary files saved:")
    print(f"  - {summary_yaml_path}")
    print(f"  - {summary_json_path}")
    print(f"  - {config_copy_path}")


if __name__ == "__main__":
    main()