#!/usr/bin/env python3
"""
Phase 2 CFD: Control vane parametric sweep.
Creates a 3D duct with a deflected vane at various angles,
runs simpleFoam, and extracts forces.

Duct: cylinder R=20mm, L=137mm (approximated as O-grid in blockMesh)
Vane: span=38mm (diameter, so half-span=19mm from center), chord=18mm, thickness=1.5mm
Vane center at Z=36mm from inlet (midpoint of Z=34..38mm in real geometry)
"""

import os
import sys
import subprocess
import shutil
import math
import csv
import struct

# ── Configuration ──────────────────────────────────────────────
ANGLES = [0, 5, 10, 15, 20, 25, 30]  # degrees
DUCT_R = 0.020       # 20mm radius
DUCT_L = 0.137       # 137mm length
VANE_SPAN = 0.038    # 38mm (full span across duct)
VANE_CHORD = 0.018   # 18mm
VANE_THICK = 0.0015  # 1.5mm
VANE_Z_CENTER = 0.036  # 36mm from inlet
INLET_VELOCITY = 30.0  # m/s

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
OF_SOURCE = "source /usr/lib/openfoam/openfoam2406/etc/bashrc"


def run_of(cmd, case_dir, log_name=None):
    """Run an OpenFOAM command in case_dir."""
    full_cmd = f"{OF_SOURCE} && cd {case_dir} && {cmd}"
    if log_name:
        full_cmd += f" > {log_name} 2>&1"
    result = subprocess.run(full_cmd, shell=True, executable="/bin/bash",
                          capture_output=(log_name is None), text=True, timeout=600)
    return result


def write_stl_binary(filename, triangles):
    """Write binary STL. triangles = list of (normal, v1, v2, v3) tuples."""
    with open(filename, 'wb') as f:
        f.write(b'\0' * 80)  # header
        f.write(struct.pack('<I', len(triangles)))
        for normal, v1, v2, v3 in triangles:
            for coord in normal:
                f.write(struct.pack('<f', coord))
            for v in (v1, v2, v3):
                for coord in v:
                    f.write(struct.pack('<f', coord))
            f.write(struct.pack('<H', 0))  # attribute byte count


def generate_vane_stl(angle_deg, filename):
    """Generate STL for a flat vane at given deflection angle.
    Vane is centered at origin in X-Y, centered at VANE_Z_CENTER in Z.
    Deflection rotates around Y-axis (tilts in X-Z plane).
    """
    angle_rad = math.radians(angle_deg)
    cos_a = math.cos(angle_rad)
    sin_a = math.sin(angle_rad)

    half_chord = VANE_CHORD / 2
    half_span = VANE_SPAN / 2
    half_thick = VANE_THICK / 2

    # Vane corners in local frame (chord along X, span along Y, thickness along Z)
    # Then rotate around Y by angle (deflection in X-Z plane)
    corners_local = [
        (-half_chord, -half_span, -half_thick),  # 0
        ( half_chord, -half_span, -half_thick),  # 1
        ( half_chord,  half_span, -half_thick),  # 2
        (-half_chord,  half_span, -half_thick),  # 3
        (-half_chord, -half_span,  half_thick),  # 4
        ( half_chord, -half_span,  half_thick),  # 5
        ( half_chord,  half_span,  half_thick),  # 6
        (-half_chord,  half_span,  half_thick),  # 7
    ]

    def rotate_y(x, y, z):
        """Rotate around Y axis, then translate to vane Z position."""
        # The vane chord is along X, flow is along Z
        # Deflection rotates the chord plane: X' = X*cos - Z*sin, Z' = X*sin + Z*cos
        xr = x * cos_a - z * sin_a
        zr = x * sin_a + z * cos_a + VANE_Z_CENTER
        return (xr, y, zr)

    corners = [rotate_y(*c) for c in corners_local]

    # 6 faces, 2 triangles each = 12 triangles
    faces = [
        # bottom (Z-)
        (0, 1, 2), (0, 2, 3),
        # top (Z+)
        (4, 6, 5), (4, 7, 6),
        # front (Y+)
        (3, 2, 6), (3, 6, 7),
        # back (Y-)
        (0, 5, 1), (0, 4, 5),
        # left (X-)
        (0, 3, 7), (0, 7, 4),
        # right (X+)
        (1, 5, 6), (1, 6, 2),
    ]

    def cross(a, b):
        return (a[1]*b[2]-a[2]*b[1], a[2]*b[0]-a[0]*b[2], a[0]*b[1]-a[1]*b[0])

    def sub(a, b):
        return (a[0]-b[0], a[1]-b[1], a[2]-b[2])

    def normalize(v):
        mag = math.sqrt(sum(x*x for x in v))
        if mag < 1e-12:
            return (0, 0, 1)
        return tuple(x/mag for x in v)

    triangles = []
    for i0, i1, i2 in faces:
        v0, v1, v2 = corners[i0], corners[i1], corners[i2]
        n = normalize(cross(sub(v1, v0), sub(v2, v0)))
        triangles.append((n, v0, v1, v2))

    write_stl_binary(filename, triangles)


def write_blockmesh_dict(case_dir):
    """Write blockMeshDict for a 3D rectangular duct approximating the cylinder.
    Using a simple box: X=[-R,R], Y=[-R,R], Z=[0,L]
    This is a square duct approximation - good enough for force trends.
    """
    R = DUCT_R * 1000  # convert to mm (scale 0.001 in blockMeshDict)
    L = DUCT_L * 1000

    content = f"""FoamFile
{{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      blockMeshDict;
}}

scale   0.001;

vertices
(
    ({-R}  {-R}   0)    // 0
    ( {R}  {-R}   0)    // 1
    ( {R}   {R}   0)    // 2
    ({-R}   {R}   0)    // 3
    ({-R}  {-R} {L})    // 4
    ( {R}  {-R} {L})    // 5
    ( {R}   {R} {L})    // 6
    ({-R}   {R} {L})    // 7
);

blocks
(
    hex (0 1 2 3 4 5 6 7) (20 20 70) simpleGrading (1 1 1)
);

boundary
(
    inlet
    {{
        type patch;
        faces ( (0 1 2 3) );
    }}

    outlet
    {{
        type patch;
        faces ( (4 5 6 7) );
    }}

    ductWall
    {{
        type wall;
        faces
        (
            (0 1 5 4)
            (1 2 6 5)
            (2 3 7 6)
            (3 0 4 7)
        );
    }}
);
"""
    os.makedirs(os.path.join(case_dir, "system"), exist_ok=True)
    with open(os.path.join(case_dir, "system", "blockMeshDict"), 'w') as f:
        f.write(content)


def write_snappy_dict(case_dir, stl_file):
    """Write snappyHexMeshDict to refine around vane and snap to it."""
    stl_name = os.path.basename(stl_file)
    content = f"""FoamFile
{{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      snappyHexMeshDict;
}}

castellatedMesh true;
snap            true;
addLayers       false;

geometry
{{
    vane.stl
    {{
        type triSurfaceMesh;
        name vane;
    }}
}}

castellatedMeshControls
{{
    maxLocalCells   50000;
    maxGlobalCells  200000;
    minRefinementCells 10;
    maxLoadUnbalance 0.10;
    nCellsBetweenLevels 3;
    resolveFeatureAngle 30;

    features
    (
    );

    refinementSurfaces
    {{
        vane
        {{
            level (2 2);
            patchInfo
            {{
                type wall;
            }}
        }}
    }}

    refinementRegions
    {{
    }}

    locationInMesh (0 0 0.001);  // point inside the duct (not inside vane)
    allowFreeStandingZoneFaces true;
}}

snapControls
{{
    nSmoothPatch    3;
    tolerance       2.0;
    nSolveIter      100;
    nRelaxIter      5;
    nFeatureSnapIter 10;
    implicitFeatureSnap false;
    explicitFeatureSnap true;
    multiRegionFeatureSnap false;
}}

addLayersControls
{{
    relativeSizes   true;
    layers          {{}}
    expansionRatio  1.0;
    finalLayerThickness 0.3;
    minThickness    0.1;
    nGrow           0;
    featureAngle    60;
    nRelaxIter      3;
    nSmoothSurfaceNormals 1;
    nSmoothNormals  3;
    nSmoothThickness 10;
    maxFaceThicknessRatio 0.5;
    maxThicknessToMedialRatio 0.3;
    minMedialAxisAngle 90;
    nBufferCellsNoExtrude 0;
    nLayerIter      50;
}}

meshQualityControls
{{
    maxNonOrtho     65;
    maxBoundarySkewness 20;
    maxInternalSkewness 4;
    maxConcave      80;
    minVol          1e-13;
    minTetQuality   -1e30;
    minArea         -1;
    minTwist        0.02;
    minDeterminant  0.001;
    minFaceWeight   0.02;
    minVolRatio     0.01;
    minTriangleTwist -1;
    nSmoothScale    4;
    errorReduction  0.75;
    relaxed
    {{
        maxNonOrtho     75;
    }}
}}

debug           0;
mergeTolerance  1e-6;
"""
    with open(os.path.join(case_dir, "system", "snappyHexMeshDict"), 'w') as f:
        f.write(content)


def write_decompose_dict(case_dir):
    """Write decomposeParDict (single processor)."""
    content = """FoamFile
{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      decomposeParDict;
}

numberOfSubdomains 1;
method          simple;
simpleCoeffs
{
    n           (1 1 1);
    delta       0.001;
}
"""
    with open(os.path.join(case_dir, "system", "decomposeParDict"), 'w') as f:
        f.write(content)


def write_control_dict(case_dir, end_time=500):
    """Write controlDict with force function objects on vane patch."""
    content = f"""FoamFile
{{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      controlDict;
}}

application     simpleFoam;
startFrom       startTime;
startTime       0;
stopAt          endTime;
endTime         {end_time};
deltaT          1;

writeControl    timeStep;
writeInterval   100;
purgeWrite      2;

writeFormat     ascii;
writePrecision  8;
writeCompression off;
timeFormat      general;
timePrecision   6;
runTimeModifiable true;

functions
{{
    vaneForces
    {{
        type            forces;
        libs            (forces);
        writeControl    timeStep;
        writeInterval   1;
        log             true;
        patches         (vane);
        rho             rhoInf;
        rhoInf          1.225;
        CofR            (0 0 {VANE_Z_CENTER});
    }}
}}
"""
    with open(os.path.join(case_dir, "system", "controlDict"), 'w') as f:
        f.write(content)


def write_fv_schemes(case_dir):
    content = """FoamFile
{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      fvSchemes;
}

ddtSchemes
{
    default         steadyState;
}

gradSchemes
{
    default         Gauss linear;
    grad(p)         Gauss linear;
    grad(U)         cellLimited Gauss linear 1;
}

divSchemes
{
    default         none;
    div(phi,U)      bounded Gauss linearUpwind grad(U);
    div(phi,k)      bounded Gauss upwind;
    div(phi,omega)  bounded Gauss upwind;
    div((nuEff*dev2(T(grad(U))))) Gauss linear;
}

laplacianSchemes
{
    default         Gauss linear corrected;
}

interpolationSchemes
{
    default         linear;
}

snGradSchemes
{
    default         corrected;
}

wallDist
{
    method          meshWave;
}
"""
    with open(os.path.join(case_dir, "system", "fvSchemes"), 'w') as f:
        f.write(content)


def write_fv_solution(case_dir):
    content = """FoamFile
{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      fvSolution;
}

solvers
{
    p
    {
        solver          GAMG;
        smoother        GaussSeidel;
        tolerance       1e-06;
        relTol          0.1;
    }

    "(U|k|omega)"
    {
        solver          smoothSolver;
        smoother        symGaussSeidel;
        tolerance       1e-06;
        relTol          0.1;
    }
}

SIMPLE
{
    nNonOrthogonalCorrectors 1;
    consistent      yes;

    residualControl
    {
        p               1e-4;
        U               1e-4;
        "(k|omega)"     1e-4;
    }
}

relaxationFactors
{
    fields
    {
        p               0.3;
    }
    equations
    {
        U               0.7;
        k               0.7;
        omega           0.7;
    }
}
"""
    with open(os.path.join(case_dir, "system", "fvSolution"), 'w') as f:
        f.write(content)


def write_initial_conditions(case_dir):
    """Write 0/ boundary conditions for 3D case."""
    zero_dir = os.path.join(case_dir, "0")
    os.makedirs(zero_dir, exist_ok=True)

    # U
    with open(os.path.join(zero_dir, "U"), 'w') as f:
        f.write(f"""FoamFile
{{
    version     2.0;
    format      ascii;
    class       volVectorField;
    object      U;
}}

dimensions      [0 1 -1 0 0 0 0];
internalField   uniform (0 0 {INLET_VELOCITY});

boundaryField
{{
    inlet
    {{
        type            fixedValue;
        value           uniform (0 0 {INLET_VELOCITY});
    }}
    outlet
    {{
        type            zeroGradient;
    }}
    ductWall
    {{
        type            noSlip;
    }}
    vane
    {{
        type            noSlip;
    }}
}}
""")

    # p
    with open(os.path.join(zero_dir, "p"), 'w') as f:
        f.write("""FoamFile
{
    version     2.0;
    format      ascii;
    class       volScalarField;
    object      p;
}

dimensions      [0 2 -2 0 0 0 0];
internalField   uniform 0;

boundaryField
{
    inlet
    {
        type            zeroGradient;
    }
    outlet
    {
        type            fixedValue;
        value           uniform 0;
    }
    ductWall
    {
        type            zeroGradient;
    }
    vane
    {
        type            zeroGradient;
    }
}
""")

    # k
    # k = 1.5 * (U * I)^2, I=5%, U=30 => k = 1.5*(1.5)^2 = 3.375
    k_val = 1.5 * (INLET_VELOCITY * 0.05) ** 2
    with open(os.path.join(zero_dir, "k"), 'w') as f:
        f.write(f"""FoamFile
{{
    version     2.0;
    format      ascii;
    class       volScalarField;
    object      k;
}}

dimensions      [0 2 -2 0 0 0 0];
internalField   uniform {k_val};

boundaryField
{{
    inlet
    {{
        type            fixedValue;
        value           uniform {k_val};
    }}
    outlet
    {{
        type            zeroGradient;
    }}
    ductWall
    {{
        type            kqRWallFunction;
        value           uniform {k_val};
    }}
    vane
    {{
        type            kqRWallFunction;
        value           uniform {k_val};
    }}
}}
""")

    # omega
    # omega = k^0.5 / (Cmu^0.25 * l), l ~ 0.07*D, D=0.04m => l=0.0028
    # omega = sqrt(3.375) / (0.5477 * 0.0028) ~ 1198
    omega_val = k_val**0.5 / (0.09**0.25 * 0.07 * 2 * DUCT_R)
    with open(os.path.join(zero_dir, "omega"), 'w') as f:
        f.write(f"""FoamFile
{{
    version     2.0;
    format      ascii;
    class       volScalarField;
    object      omega;
}}

dimensions      [0 0 -1 0 0 0 0];
internalField   uniform {omega_val:.1f};

boundaryField
{{
    inlet
    {{
        type            fixedValue;
        value           uniform {omega_val:.1f};
    }}
    outlet
    {{
        type            zeroGradient;
    }}
    ductWall
    {{
        type            omegaWallFunction;
        value           uniform {omega_val:.1f};
    }}
    vane
    {{
        type            omegaWallFunction;
        value           uniform {omega_val:.1f};
    }}
}}
""")

    # nut
    with open(os.path.join(zero_dir, "nut"), 'w') as f:
        f.write("""FoamFile
{
    version     2.0;
    format      ascii;
    class       volScalarField;
    object      nut;
}

dimensions      [0 2 -1 0 0 0 0];
internalField   uniform 0;

boundaryField
{
    inlet
    {
        type            calculated;
        value           uniform 0;
    }
    outlet
    {
        type            calculated;
        value           uniform 0;
    }
    ductWall
    {
        type            nutkWallFunction;
        value           uniform 0;
    }
    vane
    {
        type            nutkWallFunction;
        value           uniform 0;
    }
}
""")


def write_turbulence_properties(case_dir):
    const_dir = os.path.join(case_dir, "constant")
    os.makedirs(const_dir, exist_ok=True)

    with open(os.path.join(const_dir, "turbulenceProperties"), 'w') as f:
        f.write("""FoamFile
{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      turbulenceProperties;
}

simulationType  RAS;

RAS
{
    RASModel        kOmegaSST;
    turbulence      on;
    printCoeffs     on;
}
""")

    with open(os.path.join(const_dir, "transportProperties"), 'w') as f:
        f.write("""FoamFile
{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      transportProperties;
}

transportModel  Newtonian;
nu              [0 2 -1 0 0 0 0] 1.5e-05;
""")


def setup_case(angle_deg, case_dir):
    """Set up a complete OpenFOAM case for a given vane angle."""
    if os.path.exists(case_dir):
        shutil.rmtree(case_dir)
    os.makedirs(case_dir)

    # Generate vane STL
    stl_dir = os.path.join(case_dir, "constant", "triSurface")
    os.makedirs(stl_dir, exist_ok=True)
    stl_file = os.path.join(stl_dir, "vane.stl")
    generate_vane_stl(angle_deg, stl_file)

    # Write all case files
    write_blockmesh_dict(case_dir)
    write_snappy_dict(case_dir, stl_file)
    write_control_dict(case_dir)
    write_fv_schemes(case_dir)
    write_fv_solution(case_dir)
    write_decompose_dict(case_dir)
    write_initial_conditions(case_dir)
    write_turbulence_properties(case_dir)


def run_case(angle_deg, case_dir):
    """Run blockMesh + snappyHexMesh + simpleFoam for a case."""
    print(f"\n{'='*60}")
    print(f"  Running angle = {angle_deg} deg")
    print(f"{'='*60}")

    # blockMesh
    print("  blockMesh...", flush=True)
    r = run_of("blockMesh", case_dir, "log.blockMesh")
    if r.returncode != 0:
        print(f"  ERROR: blockMesh failed for angle {angle_deg}")
        return False

    # snappyHexMesh
    print("  snappyHexMesh...", flush=True)
    r = run_of("snappyHexMesh -overwrite", case_dir, "log.snappyHexMesh")
    if r.returncode != 0:
        print(f"  ERROR: snappyHexMesh failed for angle {angle_deg}")
        return False

    # Check that vane patch exists
    r = run_of("checkMesh", case_dir, "log.checkMesh")

    # simpleFoam
    print("  simpleFoam...", flush=True)
    r = run_of("simpleFoam", case_dir, "log.simpleFoam")
    if r.returncode != 0:
        print(f"  WARNING: simpleFoam returned non-zero for angle {angle_deg} (may have converged early)")

    return True


def extract_forces(case_dir):
    """Extract final forces from the forces function object output."""
    forces_dir = os.path.join(case_dir, "postProcessing", "vaneForces")

    if not os.path.exists(forces_dir):
        print(f"  WARNING: No forces directory found at {forces_dir}")
        return None

    # Find the latest time directory
    times = sorted([d for d in os.listdir(forces_dir) if d.replace('.', '').isdigit()],
                   key=float)
    if not times:
        print("  WARNING: No time directories in forces output")
        return None

    latest_time = times[-1]
    force_file = os.path.join(forces_dir, latest_time, "force.dat")
    moment_file = os.path.join(forces_dir, latest_time, "moment.dat")

    if not os.path.exists(force_file):
        # Try older format
        force_file = os.path.join(forces_dir, latest_time, "forces.dat")
        if not os.path.exists(force_file):
            print(f"  WARNING: No force data file found")
            # List what's there
            for t in times:
                td = os.path.join(forces_dir, t)
                print(f"    {t}/: {os.listdir(td)}")
            return None

    # Parse the last line of force data
    # OpenFOAM 2406 format: force.dat has columns:
    # Time (pressure_x pressure_y pressure_z) (viscous_x viscous_y viscous_z) (porous_x porous_y porous_z)
    forces = {'Fx': 0, 'Fy': 0, 'Fz': 0, 'Mx': 0, 'My': 0, 'Mz': 0}

    try:
        with open(force_file, 'r') as f:
            lines = [l.strip() for l in f.readlines() if l.strip() and not l.strip().startswith('#')]
        if lines:
            last_line = lines[-1]
            # Format: time total_x total_y total_z pressure_x ... viscous_x ...
            import re
            nums = re.findall(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?', last_line)
            if len(nums) >= 4:
                # nums[0]=time, nums[1..3]=total force
                forces['Fx'] = float(nums[1])
                forces['Fy'] = float(nums[2])
                forces['Fz'] = float(nums[3])
    except Exception as e:
        print(f"  WARNING: Error parsing force file: {e}")

    if os.path.exists(moment_file):
        try:
            with open(moment_file, 'r') as f:
                lines = [l.strip() for l in f.readlines() if l.strip() and not l.strip().startswith('#')]
            if lines:
                last_line = lines[-1]
                import re
                nums = re.findall(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?', last_line)
                if len(nums) >= 4:
                    # nums[0]=time, nums[1..3]=total moment
                    forces['Mx'] = float(nums[1])
                    forces['My'] = float(nums[2])
                    forces['Mz'] = float(nums[3])
        except Exception as e:
            print(f"  WARNING: Error parsing moment file: {e}")

    return forces


def main():
    results = []

    for angle in ANGLES:
        case_dir = os.path.join(BASE_DIR, f"angle_{angle:02d}")
        print(f"\n>>> Setting up case for angle={angle} deg at {case_dir}")

        setup_case(angle, case_dir)
        success = run_case(angle, case_dir)

        if success:
            forces = extract_forces(case_dir)
            if forces:
                row = {'angle': angle, **forces}
                results.append(row)
                print(f"  Forces: Fx={forces['Fx']:.6f} Fy={forces['Fy']:.6f} Fz={forces['Fz']:.6f}")
                print(f"  Moments: Mx={forces['Mx']:.6f} My={forces['My']:.6f} Mz={forces['Mz']:.6f}")
            else:
                print(f"  Could not extract forces for angle {angle}")
                results.append({'angle': angle, 'Fx': 0, 'Fy': 0, 'Fz': 0,
                               'Mx': 0, 'My': 0, 'Mz': 0})
        else:
            print(f"  Simulation failed for angle {angle}")
            results.append({'angle': angle, 'Fx': 0, 'Fy': 0, 'Fz': 0,
                           'Mx': 0, 'My': 0, 'Mz': 0})

    # Write CSV
    csv_file = os.path.join(BASE_DIR, "vane_forces.csv")
    with open(csv_file, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=['angle', 'Fx', 'Fy', 'Fz', 'Mx', 'My', 'Mz'])
        writer.writeheader()
        writer.writerows(results)

    print(f"\n{'='*60}")
    print(f"Results saved to {csv_file}")
    print(f"{'='*60}")

    # Print summary table
    print(f"\n{'Angle':>6} {'Fx':>12} {'Fy':>12} {'Fz':>12} {'Mx':>12} {'My':>12} {'Mz':>12}")
    print("-" * 84)
    for r in results:
        print(f"{r['angle']:>6} {r['Fx']:>12.6f} {r['Fy']:>12.6f} {r['Fz']:>12.6f} "
              f"{r['Mx']:>12.6f} {r['My']:>12.6f} {r['Mz']:>12.6f}")


if __name__ == '__main__':
    main()