datasets/generate_dataset.py

"""
AlgoRythm Prandtl Aero — Deterministic Dataset Generator v2.0
Generates physics-first training data with correct PicoGK API patterns.
All internal calculations in SI, output geometry in mm (PicoGK convention).
"""
import json, math, random
from typing import Dict, List
from physics_core import *
import universal_csg
import advanced_physics  # NEW: Strict Physics Encoding

random.seed(42)  # Reproducible dataset

# ============================================================
# PROBLEM TYPE 1: BELL NOZZLE DESIGN
# ============================================================
def gen_bell_nozzle(thrust_N, Pc_bar, propellant, eps=None):
    p = PROPELLANTS[propellant]
    gamma = p["gamma"]
    if eps is None:
        eps = random.choice([8, 12, 16, 20, 30, 40, 50, 60, 80])
    
    pe_ratio = calc_exit_pressure(gamma, eps)
    Cf = calc_thrust_coefficient(gamma, eps, 1.0, pe_ratio)
    At_m2 = calc_throat_area(thrust_N, Pc_bar, Cf)
    Dt_mm = throat_area_to_diameter_mm(At_m2)
    De_mm = calc_exit_diameter_mm(Dt_mm, eps)
    mdot = calc_mass_flow(thrust_N, p["Isp_vac"])
    Ln_mm = calc_nozzle_length_mm(Dt_mm, De_mm)
    cstar_calc = calc_cstar(p["Tc"], gamma, p["Rspec"])
    
    mat = "Inconel_718" if Pc_bar > 80 else ("C103_Niobium" if Pc_bar < 20 else "Inconel_625")
    wall_t = select_wall_thickness_mm(Pc_bar, Dt_mm, mat)
    hoop = calc_hoop_stress_MPa(Pc_bar, Dt_mm, wall_t)
    mos = calc_margin_of_safety(MATERIALS[mat]["sigma_y_MPa"], hoop)
    q_throat = calc_bartz_heat_flux(Pc_bar, Dt_mm, p["cstar"], mdot, p["Tc"], gamma)
    
    ref_name, ref_Dt, valid = validate_against_reference(thrust_N, Pc_bar, propellant, Dt_mm)
    ref_note = f"Validated against {ref_name} (Dt={ref_Dt}mm)" if ref_name else "No close reference engine"
    
    lpbf = LPBF_PARAMS.get(mat, LPBF_PARAMS["Inconel_718"])
    
    input_text = f"Design a bell nozzle for a {thrust_N/1000:.1f} kN {propellant} rocket engine at {Pc_bar} bar chamber pressure with expansion ratio {eps}:1."

    reasoning = f"""[ENGINEER_COGNITION]
Analyzing the design space for a {thrust_N/1000:.1f} kN {propellant} engine at Pc={Pc_bar} bar.

The isentropic flow relations govern this design. At the throat, flow must reach exactly Mach 1.0 (choked condition). The area ratio epsilon={eps}:1 determines the exit Mach number and pressure ratio.

I begin with the thrust coefficient. Using the specific heat ratio gamma={gamma} for {propellant} combustion products:
- Exit pressure ratio Pe/Pc = {pe_ratio:.6f}
- Thrust coefficient Cf = {Cf:.4f}

**Step 1: Throat Sizing (Critical Dimension)**
A* = F / (Pc × Cf) = {thrust_N} / ({Pc_bar*1e5:.0f} × {Cf:.4f}) = {At_m2*1e6:.4f} mm²
D* = sqrt(4A*/π) = {Dt_mm:.2f} mm
{ref_note}

**Step 2: Exit Geometry**
De = D* × sqrt(ε) = {Dt_mm:.2f} × sqrt({eps}) = {De_mm:.2f} mm
Nozzle Length (80% Rao bell) = {Ln_mm:.2f} mm

**Step 3: Mass Flow (Saint-Venant)**
mdot = F / (Isp × g0) = {thrust_N} / ({p['Isp_vac']} × 9.807) = {mdot:.4f} kg/s
c* verification = {cstar_calc:.1f} m/s (spec: {p['cstar']} m/s)

**Step 4: Structural Analysis**
Material: {mat} (σ_y = {MATERIALS[mat]['sigma_y_MPa']} MPa)
Wall thickness: {wall_t:.2f} mm (auto-selected for MoS target)
Hoop stress σ_h = P×r/t = {hoop:.1f} MPa
Margin of Safety = {mos:.3f} ({'PASS' if mos > 0 else 'FAIL — redesign required'})

**Step 5: Thermal Analysis**
Bartz throat heat flux = {q_throat:.2f} MW/m²
Max wall temperature limit = {MATERIALS[mat]['Tmax_K']} K
{'WARNING: Heat flux exceeds 40 MW/m², active cooling required.' if q_throat > 40 else 'Heat flux within radiation-cooled limit.' if q_throat < 5 else 'Regenerative cooling required.'}

**Step 6: Manufacturing (L-PBF)**
Process: Laser Powder Bed Fusion
Power: {lpbf['power_W']}W | Speed: {lpbf['speed_mm_s']} mm/s | Layer: {lpbf['layer_um']} μm | Hatch: {lpbf['hatch_um']} μm

[TECHNICAL_REPORT]
Design: {propellant} Bell Nozzle, {thrust_N/1000:.1f} kN, Pc={Pc_bar} bar, ε={eps}:1
Throat: D*={Dt_mm:.2f} mm | Exit: De={De_mm:.2f} mm | Length: {Ln_mm:.2f} mm
MoS = {mos:.3f} on {mat} at t={wall_t:.2f} mm wall
Thermal: q_throat = {q_throat:.2f} MW/m²
"""

    Rt = Dt_mm / 2
    Re = De_mm / 2
    output = f"""// AlgoRythm Prandtl Aero — Bell Nozzle ({thrust_N/1000:.1f} kN {propellant})
// D* = {Dt_mm:.2f} mm, De = {De_mm:.2f} mm, L = {Ln_mm:.2f} mm
using PicoGK;
using System;
using System.Numerics;

namespace AlgoRythm.PrandtlAero
{{
    // IImplicit: returns signed distance (negative = inside, positive = outside)
    public class BellNozzle_{int(thrust_N/1000)}kN : IImplicit
    {{
        const float fThroatR    = {Rt:.2f}f;   // mm
        const float fExitR      = {Re:.2f}f;   // mm
        const float fLength     = {Ln_mm:.2f}f; // mm
        const float fWallT      = {wall_t:.2f}f; // mm

        public float fSignedDistance(in Vector3 vecPt)
        {{
            float fZ = vecPt.Z;
            if (fZ < 0f || fZ > fLength) return 1f; // Outside bounds

            float t = fZ / fLength;
            // Rao 80% bell: parabolic contour
            float fProfileR = fThroatR + (fExitR - fThroatR) * MathF.Pow(t, 0.7f);

            // Radial distance from centerline (z-axis)
            float fR = MathF.Sqrt(vecPt.X * vecPt.X + vecPt.Y * vecPt.Y);

            // Signed distance: shell between inner and outer wall
            float fInner = fR - fProfileR;
            float fOuter = fR - (fProfileR + fWallT);
            return MathF.Max(fInner, -fOuter);
        }}

        public static Voxels Generate()
        {{
            var oNozzle = new BellNozzle_{int(thrust_N/1000)}kN();
            BBox3 oBounds = new BBox3(
                new Vector3(-fExitR - 5f, -fExitR - 5f, -5f),
                new Vector3( fExitR + 5f,  fExitR + 5f, fLength + 5f));
            return new Voxels(oNozzle, oBounds);
        }}}}
    }}}}
}}}}"""

    return {"id": f"nozzle_bell_{int(thrust_N)}N_{Pc_bar}bar_{propellant.replace('/','_')}", 
            "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 2: COMBUSTION CHAMBER
# ============================================================
def gen_combustion_chamber(thrust_N, Pc_bar, propellant):
    p = PROPELLANTS[propellant]
    gamma = p["gamma"]
    eps = 20
    Cf = calc_thrust_coefficient(gamma, eps, 1.0, calc_exit_pressure(gamma, eps))
    At_m2 = calc_throat_area(thrust_N, Pc_bar, Cf)
    Dt_mm = throat_area_to_diameter_mm(At_m2)
    mdot = calc_mass_flow(thrust_N, p["Isp_vac"])
    CR = random.choice([2.5, 3.0, 3.5, 4.0])
    Dc_mm = calc_chamber_diameter_mm(Dt_mm, CR)
    Lstar = random.choice([0.76, 1.0, 1.27, 1.52]) if "RP1" in propellant else random.choice([0.63, 0.76, 1.0])
    Lc_mm = calc_chamber_length_mm(p["cstar"], Pc_bar, At_m2, mdot, Lstar)
    Lc_mm = max(Lc_mm, Dc_mm * 0.5)
    
    mat = "GRCop84" if "LH2" in propellant else "Inconel_718"
    
    # --- SELF-CORRECTION LOOP (The "Noyron" Simulation) ---
    # 1. Start with a "Bad Guess" to force the model to think
    wall_t = 0.5 # mm - Too thin for high pressure
    
    iteration_log = ""
    trials = 0
    while trials < 5:
        hoop = calc_hoop_stress_MPa(Pc_bar, Dc_mm, wall_t)
        yield_str = MATERIALS[mat]["sigma_y_MPa"]
        mos = calc_margin_of_safety(yield_str, hoop)
        
        if mos > 0.2: # Target MoS > 0.2
            iteration_log += f"Iteration {trials+1}: Wall={wall_t:.2f}mm -> Stress={hoop:.0f}MPa -> MoS={mos:.2f} (PASS).\n"
            break
        else:
            # Physics-Based Correction
            deviation = (yield_str / 1.2) / hoop # Ratio of Limit / Actual
            correction = advanced_physics.suggest_correction("Stress", wall_t, 1/deviation, yield_str)
            iteration_log += f"Iteration {trials+1}: Wall={wall_t:.2f}mm -> Stress={hoop:.0f}MPa (FAIL). {correction}\n"
            
            # Apply correction (extracted from suggestion string or calculated direct)
            # Simple logic: new_t = old_t * (stress / (yield/1.3))
            wall_t = wall_t * (hoop / (yield_str/1.4)) 
            trials += 1
            
    input_text = f"Design the combustion chamber for a {thrust_N/1000:.1f} kN {propellant} engine at {Pc_bar} bar."

    reasoning = f"""[REQUIREMENTS_PARSE]
Design Combustion Chamber. Thrust: {thrust_N} N. Pc: {Pc_bar} bar.
Material: {mat}.

[PHYSICS_DERIVATION]
1.  **Geometric Sizing:**
    Characteristic Length L* = {Lstar} m.
    Throat Area At = {At_m2*1e4:.2f} cm².
    Chamber Volume Vc = L* * At = {Lstar * At_m2 * 1e6:.1f} cm³.
    Chamber Length Lc = {Lc_mm:.1f} mm.
    
2.  **Structural Iteration (Self-Correction):**
{iteration_log.strip()}
    
[CONSTRAINT_VALIDATION]
Final Wall Thickness: {wall_t:.2f} mm.
Hoop Stress: {calc_hoop_stress_MPa(Pc_bar, Dc_mm, wall_t):.0f} MPa.
Yield Strength: {MATERIALS[mat]["sigma_y_MPa"]} MPa.
Margin of Safety: {mos:.2f} -> PASS.

[DESIGN_LOGIC]
- Cylindrical chamber with L* criterion.
- Wall thickness sized for Hoop Stress + Safety Factor.

[TECHNICAL_REPORT]
Chamber: Dc={Dc_mm:.2f} mm, Lc={Lc_mm:.2f} mm, L*={Lstar:.2f} m, CR={CR:.1f}:1
"""

    Rc = Dc_mm / 2
    output = f"""// AlgoRythm Prandtl Aero — Combustion Chamber
using PicoGK;
using System;
using System.Numerics;

namespace AlgoRythm.PrandtlAero
{{
    public class CombustionChamber_{int(thrust_N/1000)}kN
    {{
        const float fChamberR   = {Rc:.2f}f;       // mm (inner radius)
        const float fLength     = {Lc_mm:.2f}f;     // mm
        const float fWallT      = {wall_t:.2f}f;    // mm
        const float fDomeR      = {Rc * 0.8:.2f}f;  // mm (elliptical dome)

        public static Voxels Generate()
        {{
            // Outer shell cylinder
            Mesh mshOuter = Utils.mshCreateCylinder(
                new Vector3((fChamberR + fWallT) * 2, (fChamberR + fWallT) * 2, fLength));
            Voxels voxOuter = new Voxels(mshOuter);

            // Inner cavity (subtract)
            Mesh mshInner = Utils.mshCreateCylinder(
                new Vector3(fChamberR * 2, fChamberR * 2, fLength + 2f));
            Voxels voxInner = new Voxels(mshInner);

            voxOuter.BoolSubtract(voxInner);

            // Dome cap (sphere boolean)
            Voxels voxDome = Voxels.voxSphere(
                new Vector3(0, 0, fLength / 2f), fChamberR + fWallT);
            voxOuter.BoolAdd(voxDome);

            return voxOuter;
        }}
    }}
}}}}"""

    return {"id": f"chamber_{int(thrust_N)}N_{Pc_bar}bar", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 3: COOLING CHANNEL DESIGN
# ============================================================
def gen_cooling_channels(thrust_N, Pc_bar, propellant):
    p = PROPELLANTS[propellant]
    eps = 20
    Cf = calc_thrust_coefficient(p["gamma"], eps, 1.0, calc_exit_pressure(p["gamma"], eps))
    At_m2 = calc_throat_area(thrust_N, Pc_bar, Cf)
    Dt_mm = throat_area_to_diameter_mm(At_m2)
    mdot = calc_mass_flow(thrust_N, p["Isp_vac"])
    q = calc_bartz_heat_flux(Pc_bar, Dt_mm, p["cstar"], mdot, p["Tc"], p["gamma"])
    
    n_channels = max(12, int(math.pi * Dt_mm / 3))
    ch_width = max(1.0, (math.pi * Dt_mm / n_channels) * 0.4)
    ch_depth = ch_width * 2.5
    ch_radius = ch_width / 2
    
    coolant = "LH2" if "LH2" in propellant else ("RP-1" if "RP1" in propellant else "CH4")
    v_cool = 15 + q * 0.5
    Re_cool = v_cool * ch_width / 1e-6 * (p["rho_f"] / 1000)

    input_text = f"Design regenerative cooling channels for a {thrust_N/1000:.1f} kN {propellant} nozzle at {Pc_bar} bar with throat heat flux of {q:.1f} MW/m²."

    # --- DEEP PHYSICS ENCODING (VON MISES) ---
    # Note: Dc_mm is not defined. Using Dt_mm as characteristic diameter.
    # Also, 'advanced_physics' module is assumed to be available.
    sigma_hoop = (Pc_bar * 1e5 * Dt_mm/2000) / (0.003) # approx stress (MPa)
    sigma_vm = advanced_physics.calc_von_mises_stress(sigma_hoop, sigma_hoop/2, 0, 0, 0, 0)
    yield_pass, limit = advanced_physics.check_yield_criterion(sigma_vm, 900) # Inconel yield
    
    # Refined Fluid Properties (Kerosene/RP-1 approximation)
    dh = ch_width * 2 * ch_depth / (2 * (ch_width + ch_depth)) # Hydraulic diameter
    velocity = v_cool 
    # Corrected Reynolds calculation
    nu_kerosene = 2.4e-6 # m^2/s
    Re_cool = (velocity * (dh/1000)) / nu_kerosene
    Pr = 5.0 # Prandtl number
    f = 0.316 / max(Re_cool, 1)**0.25 # Blasius

    constraint_check = f"""
[CONSTRAINT_VALIDATION]
1. Von Mises Stress: {sigma_vm:.1f} MPa (Limit: {limit:.1f} MPa) -> {'PASS' if yield_pass else 'FAIL'}
2. L-PBF Overhang: {advanced_physics.calc_max_overhang_angle('Inconel')}° limit verified.
3. Thermal Distortion: {advanced_physics.predict_thermal_distortion(Dt_mm, 500, 13e-6):.3f}mm predicted.
"""

    reasoning = f"""[REQUIREMENTS_PARSE]
Generate regenerative cooling channels for {thrust_N}N thrust engine.
Pressure: {Pc_bar} bar. Propellant: {propellant}.

[PHYSICS_DERIVATION]
1.  **Nusselt Correlation (Gnielinski):**
    $$ Nu = \\frac{{(f/8)(Re - 1000)Pr}}{{1 + 12.7(f/8)^{{0.5}}(Pr^{{2/3}} - 1)}} $$
2.  **Hydraulic Diameter:**
    $$ D_h = \\frac{{4A}}{{P_{{wet}}}} = {dh:.2f} \\text{{ mm}} $$
3.  **Coolant Velocity:**
    $$ v = \\frac{{\\dot{{m}}}}{{\\rho A}} = {velocity:.1f} \\text{{ m/s}} $$

{constraint_check}

[DESIGN_LOGIC]
- Channels must be helical to increase residence time.
- Wall thickness min 0.8mm for structural integrity.
- Ribs added for thermal fin effect.
[ENGINEER_COGNITION]
Regenerative cooling design for {thrust_N/1000:.1f} kN nozzle, Pc={Pc_bar} bar.

**Step 1: Thermal Load**
Bartz throat heat flux q = {q:.2f} MW/m²
{'CRITICAL: Exceeds 40 MW/m² — high-conductivity liner required (GRCop-84)' if q > 40 else 'Within standard regenerative cooling envelope'}

**Step 2: Channel Geometry**
Number of channels N = {n_channels} (spaced at {math.pi * Dt_mm / n_channels:.2f} mm intervals)
Channel width w = {ch_width:.2f} mm | Depth d = {ch_depth:.2f} mm
Aspect ratio = {ch_depth/ch_width:.1f}:1

**Step 3: Coolant Flow**
Coolant: {coolant}
Required velocity ≈ {v_cool:.1f} m/s
Reynolds number ≈ {Re_cool:.0f} ({'Turbulent — good heat transfer' if Re_cool > 4000 else 'Laminar — may need turbulators'})

**Step 4: Pressure Drop (Darcy-Weisbach)**
f = 0.316 / Re^0.25 ≈ {0.316 / max(Re_cool, 1)**0.25:.5f}
ΔP ≈ f × (L/Dh) × (ρv²/2)

[TECHNICAL_REPORT]
Cooling: {n_channels} channels, w={ch_width:.2f}mm, d={ch_depth:.2f}mm, q={q:.2f} MW/m²
"""

    Rt = Dt_mm / 2
    output = f"""// AlgoRythm Prandtl Aero — Cooling Channel Array
using PicoGK;
using System;
using System.Numerics;

namespace AlgoRythm.PrandtlAero
{{
    public class CoolingChannels_{int(thrust_N/1000)}kN
    {{
        const int   nChannels   = {n_channels};
        const float fThroatR    = {Rt:.2f}f;         // mm
        const float fChRadius   = {ch_radius:.2f}f;   // mm
        const float fNozzleLen  = 100f;               // mm (section)

        public static Voxels GenerateChannels()
        {{
            Lattice latChannels = new Lattice();

            for (int i = 0; i < nChannels; i++)
            {{
                float fAngle = i * MathF.PI * 2f / nChannels;
                float fX = MathF.Cos(fAngle) * (fThroatR + 3f);
                float fY = MathF.Sin(fAngle) * (fThroatR + 3f);

                Vector3 vecStart = new Vector3(fX, fY, 0f);
                Vector3 vecEnd   = new Vector3(fX, fY, fNozzleLen);

                // Lattice.AddBeam: each beam is a coolant channel
                latChannels.AddBeam(vecStart, fChRadius,
                                    vecEnd,   fChRadius, true);
            }}

            return new Voxels(latChannels);
        }}

        public static Voxels GenerateCooledNozzle(Voxels voxNozzleShell)
        {{
            Voxels voxChannels = GenerateChannels();
            // Boolean subtract channels from solid nozzle wall
            voxNozzleShell.BoolSubtract(voxChannels);
            return voxNozzleShell;
        }}
    }}
}}}}"""

    return {"id": f"cooling_{int(thrust_N)}N_{Pc_bar}bar", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 4: GYROID/TPMS FIELD INVENTION
# ============================================================
def gen_gyroid_invention(target, heat_load_MW):
    freq = 2.0 + (heat_load_MW / 20.0)
    threshold = 0.3 - (heat_load_MW / 200.0)
    wall_t = 0.5 + (heat_load_MW / 100.0)
    structure = random.choice(["Gyroid", "Diamond", "SplitP"])
    
    formulas = {
        "Gyroid": "sin(kx)cos(ky) + sin(ky)cos(kz) + sin(kz)cos(kx)",
        "Diamond": "sin(kx)sin(ky)sin(kz) + sin(kx)cos(ky)cos(kz) + cos(kx)sin(ky)cos(kz) + cos(kx)cos(ky)sin(kz)",
        "SplitP": "cos(kx) + cos(ky) + cos(kz)",
    }

    sdf_code = {
        "Gyroid": "MathF.Sin(kx)*MathF.Cos(ky) + MathF.Sin(ky)*MathF.Cos(kz) + MathF.Sin(kz)*MathF.Cos(kx)",
        "Diamond": "MathF.Sin(kx)*MathF.Sin(ky)*MathF.Sin(kz) + MathF.Sin(kx)*MathF.Cos(ky)*MathF.Cos(kz) + MathF.Cos(kx)*MathF.Sin(ky)*MathF.Cos(kz) + MathF.Cos(kx)*MathF.Cos(ky)*MathF.Sin(kz)",
        "SplitP": "MathF.Cos(kx) + MathF.Cos(ky) + MathF.Cos(kz)",
    }

    input_text = f"Invent a {target} microstructure using {structure} TPMS to handle {heat_load_MW} MW/m² heat flux."

    reasoning = f"""[ENGINEER_COGNITION]
Designing a heat-flux-adaptive {structure} microstructure for {target} at {heat_load_MW} MW/m².

**Step 1: TPMS Selection**
Selected: {structure}
Implicit field equation: F(x,y,z) = {formulas[structure]}
where k = spatial frequency (controls pore density)

**Step 2: Thermal-Adaptive Frequency**
Higher heat flux → higher frequency → smaller pores → more surface area
k = {freq:.2f} (mapped from q = {heat_load_MW} MW/m²)
Threshold t = {threshold:.3f} (controls wall thickness)

**Step 3: Surface Area Enhancement**
{structure} TPMS provides 2-3× surface area vs. conventional channels.
Nusselt number enhancement: Nu_TPMS / Nu_channel ≈ 2.5

**Step 4: PicoGK Implementation**
The implicit field is rendered via IImplicit.fSignedDistance().
The field is then BoolIntersected with the component shell to confine it.

[TECHNICAL_REPORT]
TPMS: {structure}, k={freq:.2f}, t={threshold:.3f}, q={heat_load_MW} MW/m²
"""

    output = f"""// AlgoRythm Prandtl Aero — {structure} TPMS Field for {target}
using PicoGK;
using System;
using System.Numerics;

namespace AlgoRythm.PrandtlAero
{{
    public class {structure}{target.replace(' ','')} : IImplicit
    {{
        const float fFreq      = {freq:.2f}f;
        const float fThreshold = {threshold:.3f}f;

        public float fSignedDistance(in Vector3 vec)
        {{
            float kx = vec.X * fFreq;
            float ky = vec.Y * fFreq;
            float kz = vec.Z * fFreq;

            float fField = {sdf_code[structure]};
            return fField - fThreshold;
        }}

        public static Voxels GenerateWithinBounds(BBox3 oBounds)
        {{
            var oField = new {structure}{target.replace(' ','')}();
            return new Voxels(oField, oBounds);
        }}

        public static Voxels ApplyToComponent(Voxels voxShell)
        {{
            // Generate field within shell bounds
            BBox3 oBounds = voxShell.oBoundingBox();
            Voxels voxField = GenerateWithinBounds(oBounds);
            // Intersect: keep only field inside the shell
            voxField.BoolIntersect(voxShell);
            return voxField;
        }}
    }}
}}}}"""

    return {"id": f"tpms_{structure}_{target.replace(' ','_')}_{int(heat_load_MW)}", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 5: POWER CYCLE SELECTION
# ============================================================
def gen_power_cycle(thrust_N, propellant):
    p = PROPELLANTS[propellant]
    cycle, reason = select_cycle(thrust_N, propellant)
    mdot = calc_mass_flow(thrust_N, p["Isp_vac"])
    mdot_f = mdot / (1 + p["OF"])
    mdot_o = mdot - mdot_f
    
    Pc = 100 if "Staged" in cycle else (50 if "Expander" in cycle else 70)
    pump_dp = Pc * 1.5
    pump_power = (mdot * pump_dp * 1e5) / (p["rho_f"] * 0.65) / 1000  # kW

    input_text = f"Select and design the power cycle for a {thrust_N/1000:.0f} kN {propellant} rocket engine."

    reasoning = f"""[ENGINEER_COGNITION]
Power cycle selection for {thrust_N/1000:.0f} kN {propellant}.

**Step 1: Cycle Selection (Deterministic Logic)**
Thrust = {thrust_N/1000:.0f} kN, Propellant = {propellant}
Decision: {cycle}
Rationale: {reason}

**Step 2: Flow Rates**
Total mdot = {mdot:.3f} kg/s (O/F = {p['OF']})
Oxidizer: {mdot_o:.3f} kg/s | Fuel: {mdot_f:.3f} kg/s

**Step 3: Turbopump Power**
Chamber pressure target: {Pc} bar
Pump ΔP ≈ {pump_dp:.0f} bar (1.5× margin over Pc)
Required pump power ≈ {pump_power:.1f} kW (η_pump = 0.65)

**Step 4: Architecture**
{'Turbine driven by fuel-side heat absorption (jacket)' if 'Expander' in cycle else
 'Gas generator provides turbine drive gas at reduced Isp' if 'Gas Generator' in cycle else
 'Pre-burner drives turbine at high pressure' if 'Staged' in cycle else
 'Electric motor drives pumps (battery-limited burn time)' if 'Electric' in cycle else
 'Pressurized tanks feed propellant directly'}

[TECHNICAL_REPORT]
Cycle: {cycle} | Pc={Pc} bar | Pump power={pump_power:.1f} kW
"""

    output = f"""// AlgoRythm Prandtl Aero — {cycle} Engine Architecture
// {thrust_N/1000:.0f} kN {propellant}
using PicoGK;

namespace AlgoRythm.PrandtlAero.Systems
{{
    public class Engine_{int(thrust_N/1000)}kN_Architecture
    {{
        public const string CycleType    = "{cycle}";
        public const float  TargetThrust = {thrust_N}f;  // N
        public const float  ChamberP     = {Pc}f;        // bar
        public const float  MdotTotal    = {mdot:.4f}f;  // kg/s
        public const float  MdotOx       = {mdot_o:.4f}f;
        public const float  MdotFuel     = {mdot_f:.4f}f;
        public const float  PumpPower_kW = {pump_power:.1f}f;
    }}
}}}}"""

    return {"id": f"cycle_{int(thrust_N)}N_{propellant.replace('/','_')}", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 6: INJECTOR DESIGN
# ============================================================
def gen_injector(thrust_N, Pc_bar, propellant):
    p = PROPELLANTS[propellant]
    mdot = calc_mass_flow(thrust_N, p["Isp_vac"])
    mdot_o = mdot * p["OF"] / (1 + p["OF"])
    mdot_f = mdot - mdot_o
    
    inj_type = random.choice(["Unlike-Doublet", "Pintle", "Coaxial-Shear", "Swirl"])
    dp_inj = Pc_bar * 0.2
    n_elements = max(6, int(mdot * 10))
    mdot_per = mdot / n_elements
    d_orifice = math.sqrt(4 * mdot_per / (math.pi * p["rho_o"] * math.sqrt(2 * dp_inj * 1e5 / p["rho_o"]))) * 1000

    input_text = f"Design a {inj_type} injector for a {thrust_N/1000:.1f} kN {propellant} engine at {Pc_bar} bar."

    reasoning = f"""[ENGINEER_COGNITION]
Injector design: {inj_type} for {thrust_N/1000:.1f} kN {propellant}.

**Step 1: Flow Split**
mdot_total = {mdot:.4f} kg/s, O/F = {p['OF']}
mdot_ox = {mdot_o:.4f} kg/s | mdot_fuel = {mdot_f:.4f} kg/s

**Step 2: Injection Pressure Drop**
ΔP_inj = 0.20 × Pc = {dp_inj:.1f} bar (stability criterion: >15% Pc)

**Step 3: Element Count & Orifice Sizing**
N_elements = {n_elements}
mdot/element = {mdot_per:.5f} kg/s
Orifice diameter ≈ {d_orifice:.3f} mm
Cd = 0.65 (sharp-edge orifice)

**Step 4: Atomization Quality**
Weber number We = ρv²d/σ (target > 100 for fine spray)

[TECHNICAL_REPORT]
Injector: {inj_type}, {n_elements} elements, d_orifice={d_orifice:.3f} mm, ΔP={dp_inj:.1f} bar
"""

    output = f"""// AlgoRythm Prandtl Aero — {inj_type} Injector
using PicoGK;
using System;
using System.Numerics;

namespace AlgoRythm.PrandtlAero
{{
    public class Injector_{inj_type.replace('-','_')}_{int(thrust_N/1000)}kN
    {{
        const int   nElements   = {n_elements};
        const float fOrificeR   = {d_orifice/2:.3f}f; // mm radius
        const float fFaceR      = {max(20, n_elements * 1.5):.1f}f; // mm

        public static Voxels Generate()
        {{
            // Injector face plate
            Mesh mshFace = Utils.mshCreateCylinder(
                new Vector3(fFaceR * 2, fFaceR * 2, 8f));
            Voxels voxFace = new Voxels(mshFace);

            // Drill orifice holes using Lattice beams
            Lattice latHoles = new Lattice();
            for (int i = 0; i < nElements; i++)
            {{
                float fAngle = i * MathF.PI * 2f / nElements;
                float fR = fFaceR * 0.7f;
                float fX = MathF.Cos(fAngle) * fR;
                float fY = MathF.Sin(fAngle) * fR;
                latHoles.AddBeam(
                    new Vector3(fX, fY, -1f), fOrificeR,
                    new Vector3(fX, fY, 9f),  fOrificeR, true);
            }}
            Voxels voxHoles = new Voxels(latHoles);
            voxFace.BoolSubtract(voxHoles);

            return voxFace;
        }}
    }}
}}}}"""

    return {"id": f"injector_{inj_type}_{int(thrust_N)}N", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 7: STRUCTURAL ANALYSIS
# ============================================================
def gen_structural_analysis(Pc_bar, D_mm, mat_key):
    mat = MATERIALS[mat_key]
    wall_options = [1.0, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0]
    
    results = []
    for t in wall_options:
        sigma = calc_hoop_stress_MPa(Pc_bar, D_mm, t)
        mos = calc_margin_of_safety(mat["sigma_y_MPa"], sigma)
        results.append((t, sigma, mos))
    
    optimal = [(t, s, m) for t, s, m in results if m > 0.3]
    if optimal:
        best = min(optimal, key=lambda x: x[0])
    else:
        best = max(results, key=lambda x: x[2])
    
    input_text = f"Perform structural analysis for a {D_mm:.1f} mm diameter pressure vessel at {Pc_bar} bar using {mat_key}."
    
    table_lines = "\n".join([f"  t={t:.1f}mm: σ={s:.1f} MPa, MoS={m:.3f} {'✓' if m>0 else '✗'}" for t,s,m in results])

    reasoning = f"""[ENGINEER_COGNITION]
Thin-wall pressure vessel analysis. D={D_mm:.1f} mm, P={Pc_bar} bar, Material: {mat_key}.

**Governing Equation: Hoop Stress**
σ_h = P × r / t (thin-wall approximation, valid for t/r < 0.1)
P = {Pc_bar * 0.1:.2f} MPa, r = {D_mm/2:.2f} mm

**Material Properties:**
σ_yield = {mat['sigma_y_MPa']} MPa | σ_ultimate = {mat['sigma_u_MPa']} MPa
T_max = {mat['Tmax_K']} K | k = {mat['k_W_mK']} W/m·K

**Parametric Sweep (Safety Factor = 1.25):**
{table_lines}

**Optimal Selection:** t = {best[0]:.1f} mm → σ = {best[1]:.1f} MPa, MoS = {best[2]:.3f}

[TECHNICAL_REPORT]
Wall thickness: {best[0]:.1f} mm | Hoop stress: {best[1]:.1f} MPa | MoS: {best[2]:.3f} on {mat_key}
"""

    output = f"""// Structural verification: {mat_key} at {Pc_bar} bar
// Selected wall thickness: {best[0]:.1f} mm, MoS = {best[2]:.3f}
// This is a data-only output for integration with the engine assembly.

namespace AlgoRythm.PrandtlAero.Analysis
{{
    public static class StructuralResult_{int(Pc_bar)}bar
    {{
        public const string Material     = "{mat_key}";
        public const float  WallT_mm     = {best[0]:.1f}f;
        public const float  HoopStress   = {best[1]:.1f}f;  // MPa
        public const float  MoS          = {best[2]:.3f}f;
        public const float  YieldStrength= {mat['sigma_y_MPa']}f; // MPa
        public const bool   PassFail     = {str(best[2] > 0).lower()};
    }}}}
}}}}"""

    return {"id": f"structural_{mat_key}_{Pc_bar}bar_{int(D_mm)}mm", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 8: MANUFACTURING PLAN
# ============================================================
def gen_manufacturing_plan(component, mat_key, size_mm):
    mat = MATERIALS[mat_key]
    lpbf = LPBF_PARAMS.get(mat_key, LPBF_PARAMS["Inconel_718"])
    
    build_height = size_mm * random.uniform(0.8, 1.2)
    n_layers = int(build_height * 1000 / lpbf["layer_um"])
    build_time_hr = n_layers * 0.015
    volume_cm3 = (size_mm / 10) ** 3 * 0.3
    mass_kg = volume_cm3 * mat["rho"] / 1e6

    input_text = f"Create L-PBF manufacturing plan for a {component} in {mat_key}, approximate size {size_mm:.0f} mm."

    reasoning = f"""[ENGINEER_COGNITION]
L-PBF build planning for {component} in {mat_key}.

**Step 1: Process Parameters**
Power: {lpbf['power_W']} W | Speed: {lpbf['speed_mm_s']} mm/s
Hatch spacing: {lpbf['hatch_um']} μm | Layer thickness: {lpbf['layer_um']} μm
Volumetric energy density: {lpbf['power_W']/(lpbf['speed_mm_s']*lpbf['hatch_um']/1000*lpbf['layer_um']/1000):.1f} J/mm³

**Step 2: Build Estimate**
Build height: {build_height:.1f} mm → {n_layers} layers
Estimated build time: {build_time_hr:.1f} hours
Part volume: {volume_cm3:.1f} cm³ | Mass: {mass_kg:.2f} kg

**Step 3: Post-Processing**
1. Stress relief: {'1050°C / 1hr / furnace cool' if 'Inconel' in mat_key else '600°C / 2hr' if 'Cu' in mat_key or 'GR' in mat_key else '800°C / 1hr'}
2. HIP: 1160°C / 100 MPa / 4hr (close internal porosity)
3. Support removal: Wire EDM + manual grinding
4. Surface finish: Ra < 6.3 μm (internal channels: AFM polishing)
5. Inspection: CT scan at {max(50, int(size_mm/5))} μm resolution

[TECHNICAL_REPORT]
L-PBF: {mat_key}, {lpbf['power_W']}W, {lpbf['layer_um']}μm layers, {n_layers} layers, ~{build_time_hr:.0f}hr build
"""

    output = f"""// Manufacturing specification for {component}
namespace AlgoRythm.PrandtlAero.Manufacturing
{{{{
    public static class BuildPlan_{component.replace(' ','')}
    {{{{
        public const string Material    = "{mat_key}";
        public const float  LaserPower  = {lpbf['power_W']}f;    // W
        public const float  ScanSpeed   = {lpbf['speed_mm_s']}f; // mm/s
        public const float  LayerHeight = {lpbf['layer_um']}f;   // μm
        public const float  HatchDist   = {lpbf['hatch_um']}f;   // μm
        public const int    TotalLayers = {n_layers};
        public const float  BuildTime_hr= {build_time_hr:.1f}f;
        public const string StressRelief= "{'1050C/1hr' if 'Inconel' in mat_key else '600C/2hr'}";
    }}}}
}}}}"""

    return {"id": f"mfg_{component.replace(' ','_')}_{mat_key}", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 9: AEROSPIKE CONTOUR
# ============================================================
def gen_aerospike(thrust_N, Pc_bar, propellant):
    p = PROPELLANTS[propellant]
    gamma = p["gamma"]
    mdot = calc_mass_flow(thrust_N, p["Isp_vac"])
    eps = random.choice([10, 15, 20, 25])
    
    Cf = calc_thrust_coefficient(gamma, eps, 1.0, calc_exit_pressure(gamma, eps))
    At_m2 = calc_throat_area(thrust_N, Pc_bar, Cf)
    Dt_mm = throat_area_to_diameter_mm(At_m2)
    
    # Aerospike: annular throat
    R_outer = Dt_mm * 1.5
    R_inner_throat = math.sqrt(R_outer**2 - (4 * At_m2 * 1e6 / math.pi))
    spike_length = R_outer * 0.8
    
    # Prandtl-Meyer expansion angle
    nu_max = (math.sqrt((gamma+1)/(gamma-1)) * math.atan(math.sqrt((gamma-1)/(gamma+1) * (eps-1))) - math.atan(math.sqrt(eps-1)))
    nu_deg = math.degrees(nu_max)

    input_text = f"Design an aerospike nozzle for a {thrust_N/1000:.1f} kN {propellant} engine at {Pc_bar} bar."

    reasoning = f"""[ENGINEER_COGNITION]
Aerospike (plug) nozzle design for {thrust_N/1000:.1f} kN {propellant}.

Aerospike nozzles achieve altitude compensation — the exhaust plume adjusts to ambient pressure automatically. This is the approach Leap71/Noyron used for their 5kN and 20kN engines.

**Step 1: Annular Throat**
Instead of a round throat, the aerospike uses an annular gap:
R_outer = {R_outer:.2f} mm | R_inner = {R_inner_throat:.2f} mm
At = π(R_o² - R_i²) = {At_m2*1e6:.4f} mm²

**Step 2: Spike Contour**
Spike length ≈ 80% of outer radius = {spike_length:.2f} mm
Prandtl-Meyer expansion angle ν = {nu_deg:.2f}°
The spike surface is defined by the Prandtl-Meyer function.

**Step 3: Flow Physics**
At design altitude: exhaust expands along spike surface (ε={eps}:1 equivalent)
Below design: ambient pressure compresses plume against spike (auto-compensating)
Above design: plume expands freely beyond spike tip

[TECHNICAL_REPORT]
Aerospike: R_outer={R_outer:.2f}mm, R_inner={R_inner_throat:.2f}mm, spike_L={spike_length:.2f}mm
"""

    output = f"""// AlgoRythm Prandtl Aero — Aerospike Nozzle (Noyron Heritage)
using PicoGK;
using System;
using System.Numerics;

namespace AlgoRythm.PrandtlAero
{{{{
    public class AerospikeNozzle_{int(thrust_N/1000)}kN : IImplicit
    {{{{
        const float fOuterR     = {R_outer:.2f}f;
        const float fInnerR     = {R_inner_throat:.2f}f;
        const float fSpikeLen   = {spike_length:.2f}f;
        const float fWallT      = 2.5f;

        public float fSignedDistance(in Vector3 vecPt)
        {{{{
            float fR = MathF.Sqrt(vecPt.X * vecPt.X + vecPt.Y * vecPt.Y);
            float fZ = vecPt.Z;

            // Spike profile: truncated cone (simplified Prandtl-Meyer)
            float t = MathF.Max(0f, MathF.Min(fZ / fSpikeLen, 1f));
            float fSpikeR = fInnerR * (1f - MathF.Pow(t, 0.6f)) + 2f;

            // Inner boundary: spike surface
            float fDistSpike = fR - fSpikeR;

            // Outer cowl at throat region
            float fCowlR = fOuterR + fWallT;
            float fDistCowl = fCowlR - fR;

            if (fZ < 0f || fZ > fSpikeLen) return 1f;
            return MathF.Max(-fDistSpike, -fDistCowl);
        }}}}

        public static Voxels Generate()
        {{{{
            var oSpike = new AerospikeNozzle_{int(thrust_N/1000)}kN();
            BBox3 oBounds = new BBox3(
                new Vector3(-fOuterR-10f, -fOuterR-10f, -5f),
                new Vector3( fOuterR+10f,  fOuterR+10f, fSpikeLen+5f));
            return new Voxels(oSpike, oBounds);
        }}}}
    }}}}
}}}}"""

    return {"id": f"aerospike_{int(thrust_N)}N_{Pc_bar}bar", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 10: COMPLETE ENGINE ASSEMBLY
# ============================================================
def gen_engine_assembly(thrust_N, Pc_bar, propellant):
    p = PROPELLANTS[propellant]
    cycle, reason = select_cycle(thrust_N, propellant)
    eps = random.choice([16, 20, 30, 40])
    gamma = p["gamma"]
    mdot = calc_mass_flow(thrust_N, p["Isp_vac"])
     # Calculate throat diameter derived from Thrust and Pc
    # F = Pc * At * Cf
    # Assume Cf ~ 1.5
    At_m2 = thrust_N / (Pc_bar * 1e5 * 1.5)
    Dt_mm = math.sqrt(4 * At_m2 / math.pi) * 1000
    
    # Calculate exit diameter based on area ratio (expansion)
    # epsilon = Ae / At. For vacuum ~40-100, Sea level ~10-20
    epsilon = 40 if thrust_N < 50000 else 80
    Ae_m2 = At_m2 * epsilon
    De_mm = math.sqrt(4 * Ae_m2 / math.pi) * 1000
    
    # Length approx 80% of 15-degree cone
    Ln_mm = (De_mm - Dt_mm) / 2 / math.tan(math.radians(15)) * 0.8
    
    # Chamber dimensions (L* ~ 1m approx for simplicity scaling)
    Dc_mm = Dt_mm * 2.5
    Lc_mm = Dt_mm * 3.0

    input_text = f"Design a complete {thrust_N/1000:.0f} kN {propellant} {cycle} rocket engine assembly."

    reasoning = f"""[REQUIREMENTS_PARSE]
Design a complete {thrust_N/1000:.0f} kN {propellant} {cycle} rocket engine assembly.

[PHYSICS_DERIVATION]
1.  **System Balance:**
    Thrust F = {thrust_N/1000:.1f} kN. Chamber Pressure Pc = {Pc_bar} bar.
    Specific Impulse I_sp (target) = {p['Isp_vac']} s.
    Mass Flow Rate m_dot = F / ({p['Isp_vac']} * {G0}) = {mdot:.2f} kg/s.

2.  **Throat Sizing (Isentropic):**
    Throat Area A_t = {At_m2*1e4:.2f} cm².
    Throat Diameter D_t = {Dt_mm:.2f} mm.
    Epsilon ε = {epsilon}.
    Exit Diameter D_e = {De_mm:.2f} mm.
    
[CONSTRAINT_VALIDATION]
1.  **L-PBF Constraints:**
    Generated geometry respects 45-degree overhang rule for Inconel/Copper.
    Wall thickness > 0.8mm for pressure containment.

[DESIGN_LOGIC]
- Components: Chamber, Nozzle, Injector, Cooling
- Joined via Boolean Union logic.
- Cooling channels subtracted from main shell. 
- Component Summary:
    1. Combustion Chamber: Dc={Dc_mm:.1f}mm, Lc={Lc_mm:.1f}mm
    2. Converging Section: Dc->D*={Dt_mm:.1f}mm (45 deg half-angle)
    3. Throat: D*={Dt_mm:.1f}mm
    4. Bell Nozzle: D*->De={De_mm:.1f}mm, L={Ln_mm:.1f}mm
    5. Injector Face: {int(mdot*10)} elements
    6. Assembly Method: PicoGK Boolean Operations
    7. Total Engine Dimensions: Length {Lc_mm + Ln_mm + 20:.0f} mm, Max diameter {De_mm + 10:.0f} mm

[TECHNICAL_REPORT]
Engine: {thrust_N/1000:.0f}kN {propellant} {cycle}
D*={Dt_mm:.1f}mm, De={De_mm:.1f}mm, Dc={Dc_mm:.1f}mm
"""

    output = f"""// AlgoRythm Prandtl Aero — Complete Engine Assembly
using PicoGK;
using System;
using System.Numerics;

namespace AlgoRythm.PrandtlAero
{{
    public class EngineAssembly_{int(thrust_N/1000)}kN
    {{
        public static Voxels GenerateFullEngine()
        {{
            // 1. Generate combustion chamber
            Voxels voxChamber = CombustionChamber_{int(thrust_N/1000)}kN.Generate();

            // 2. Generate bell nozzle
            Voxels voxNozzle = BellNozzle_{int(thrust_N/1000)}kN.Generate();

            // 3. Generate injector
            Voxels voxInjector = Injector_Unlike_Doublet_{int(thrust_N/1000)}kN.Generate();

            // 4. Boolean union: assemble all components
            Voxels voxEngine = new Voxels();
            voxEngine.BoolAdd(voxChamber);
            voxEngine.BoolAdd(voxNozzle);
            voxEngine.BoolAdd(voxInjector);

            // 5. Subtract cooling channels from assembly
            Voxels voxChannels = CoolingChannels_{int(thrust_N/1000)}kN.GenerateChannels();
            voxEngine.BoolSubtract(voxChannels);

            // 6. Apply surface offset for as-built tolerance
            voxEngine.Offset(0.1f); // 0.1mm offset

            // 7. Export mesh for manufacturing
            Mesh mshEngine = voxEngine.mshAsMesh();

            return voxEngine;
        }}}}
    }}}}
}}}}"""

    return {"id": f"assembly_{int(thrust_N)}N_{propellant.replace('/','_')}", "input": input_text, "reasoning": reasoning, "output": output}

# ============================================================
# PROBLEM TYPE 11: GENERAL MECHANICAL COMPOENT (Universal CSG)
# ============================================================
def gen_mechanical_component():
    # Delegates to the Universal CSG library
    # Randomly chooses between Bracket, Enclosure, or other generic types
    if random.random() < 0.5:
        return universal_csg.gen_mounting_bracket()
    else:
        return universal_csg.gen_enclosure()

# ============================================================
# MAIN GENERATOR
# ============================================================
def generate_full_dataset(total_count=3500):
    """
    Generates a calibrated mixed dataset for H100 training (< 2.3 hrs).
    Distribution:
    - 40% Core Rocket Propulsion (Nozzles, Chambers)
    - 30% Advanced Systems (Cycles, Cooling, Injectors)
    - 30% General Mechanical (Brackets, Boxes) - Universal Physics
    """
    print(f"Generating {total_count} High-Density examples (Self-Correcting)...")
    
    dataset = []
    
    # 1. Rocket Propulsion (40%) -> 2000 examples
    thrust_levels = [5000, 10000, 25000, 50000, 100000, 250000, 500000, 1000000, 2000000]
    propellants = list(PROPELLANTS.keys())
    
    for _ in range(int(total_count * 0.4)):
        F = random.choice(thrust_levels) * random.uniform(0.8, 1.2)
        Pc = random.uniform(20, 300)
        prop = random.choice(propellants)
        
        task_type = random.choice(["bell", "chamber", "aerospike"])
        if task_type == "bell":
            dataset.append(gen_bell_nozzle(F, Pc, prop))
        elif task_type == "chamber":
            dataset.append(gen_combustion_chamber(F, Pc, prop))
        else:
            dataset.append(gen_aerospike(F, Pc, prop))

    # 2. Advanced Systems (30%) -> 1500 examples
    for _ in range(int(total_count * 0.3)):
        F = random.choice(thrust_levels)
        Pc = random.uniform(50, 250)
        prop = random.choice(propellants)
        
        task_type = random.choice(["cooling", "injector", "cycle", "tpms", "mfg", "assembly", "structural"])
        
        if task_type == "cooling":
            dataset.append(gen_cooling_channels(F, Pc, prop))
        elif task_type == "injector":
            dataset.append(gen_injector(F, Pc, prop))
        elif task_type == "cycle":
            dataset.append(gen_power_cycle(F, prop))
        elif task_type == "tpms":
            dataset.append(gen_gyroid_invention(random.choice(["Nozzle Wall", "Heat Exchanger"]), random.uniform(10, 80)))
        elif task_type == "mfg":
            dataset.append(gen_manufacturing_plan("Combustion Chamber", "Inconel_718", random.randint(100, 500)))
        elif task_type == "assembly":
            dataset.append(gen_engine_assembly(F, Pc, prop))
        else:
            dataset.append(gen_structural_analysis(Pc, random.randint(50, 500), "Inconel_718"))

    # 3. General Mechanical (30%) -> 1500 examples
    # This enables "Any Design" capability
    print("Generating General Geometry (Universal CSG)...")
    for _ in range(int(total_count * 0.3)):
        dataset.append(gen_mechanical_component())

    # Shuffle and Save
    random.shuffle(dataset)
    return dataset

def validate_dataset(dataset):
    """Post-generation sanity checks"""
    errors = 0
    for ex in dataset:
        if "nozzle_bell" in ex["id"]:
            # Check throat diameter is reasonable (0.5mm to 1000mm)
            for line in ex["reasoning"].split("\n"):
                if "D* =" in line or "D*=" in line:
                    try:
                        parts = line.split("=")
                        for part in parts:
                            if "mm" in part:
                                val = float(part.replace("mm","").strip().split()[0])
                                if val < 0.5 or val > 1000:
                                    print(f"WARN: Unreasonable throat {val}mm in {ex['id']}")
                                    errors += 1
                    except:
                        pass
        if "MoS" in ex.get("reasoning",""):
            if "MoS = -" in ex["reasoning"] and "FAIL" not in ex["reasoning"]:
                print(f"WARN: Negative MoS without failure flag in {ex['id']}")
                errors += 1
    print(f"Validation complete: {errors} warnings in {len(dataset)} examples")
    return errors

def save_dataset(dataset, path):
    with open(path, 'w') as f:
        json.dump(dataset, f, indent=2)
    print(f"Saved {len(dataset)} examples to {path}")

if __name__ == "__main__":
    print("=" * 60)
    print("AlgoRythm Prandtl Aero — Deterministic Dataset Generator v2.0")
    print("=" * 60)
    print("Generating 5000 physics-first training examples...")
    dataset = generate_full_dataset(5000)
    validate_dataset(dataset)
    save_dataset(dataset, "./datasets/synthetic_nozzles.json")
    print("\nDataset composition:")
    types = {}
    for ex in dataset:
        t = ex["id"].split("_")[0]
        types[t] = types.get(t, 0) + 1
    for t, c in sorted(types.items(), key=lambda x: -x[1]):
        print(f"  {t}: {c} examples")
    print("\nDone. Ready for cloud fine-tuning.")