#!/usr/bin/env python3
"""
Plot vane forces vs deflection angle from Phase 2 CFD results.
"""

import csv
import os
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CSV_FILE = os.path.join(BASE_DIR, "vane_forces.csv")

def main():
    # Read CSV
    angles, Fx, Fy, Fz = [], [], [], []
    Mx, My, Mz = [], [], []

    with open(CSV_FILE, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            angles.append(float(row['angle']))
            Fx.append(float(row['Fx']))
            Fy.append(float(row['Fy']))
            Fz.append(float(row['Fz']))
            Mx.append(float(row['Mx']))
            My.append(float(row['My']))
            Mz.append(float(row['Mz']))

    angles = np.array(angles)
    Fx = np.array(Fx)
    Fy = np.array(Fy)
    Fz = np.array(Fz)
    Mx = np.array(Mx)
    My = np.array(My)
    Mz = np.array(Mz)

    # Lateral force magnitude (X-Y plane)
    F_lateral = np.sqrt(Fx**2 + Fy**2)

    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle('Phase 2 CFD: Control Vane Forces vs Deflection Angle\n'
                 f'Duct ID=40mm, Vane 38x18mm, V_inlet=30 m/s',
                 fontsize=13, fontweight='bold')

    # Plot 1: Lateral force (Fx) vs angle
    ax = axes[0, 0]
    ax.plot(angles, Fx * 1000, 'bo-', linewidth=2, markersize=8, label='Fx (lateral)')
    ax.set_xlabel('Vane Deflection Angle [deg]')
    ax.set_ylabel('Force [mN]')
    ax.set_title('Lateral Force (Fx) vs Angle')
    ax.grid(True, alpha=0.3)
    ax.legend()

    # Plot 2: Axial force (Fz) vs angle
    ax = axes[0, 1]
    ax.plot(angles, Fz * 1000, 'rs-', linewidth=2, markersize=8, label='Fz (axial/drag)')
    ax.set_xlabel('Vane Deflection Angle [deg]')
    ax.set_ylabel('Force [mN]')
    ax.set_title('Axial Force (Fz) vs Angle')
    ax.grid(True, alpha=0.3)
    ax.legend()

    # Plot 3: All forces
    ax = axes[1, 0]
    ax.plot(angles, Fx * 1000, 'bo-', linewidth=2, markersize=6, label='Fx')
    ax.plot(angles, Fy * 1000, 'g^-', linewidth=2, markersize=6, label='Fy')
    ax.plot(angles, Fz * 1000, 'rs-', linewidth=2, markersize=6, label='Fz')
    ax.plot(angles, F_lateral * 1000, 'kd--', linewidth=2, markersize=6, label='|F_lateral|')
    ax.set_xlabel('Vane Deflection Angle [deg]')
    ax.set_ylabel('Force [mN]')
    ax.set_title('All Force Components')
    ax.grid(True, alpha=0.3)
    ax.legend()

    # Plot 4: Moments
    ax = axes[1, 1]
    ax.plot(angles, Mx * 1000, 'bo-', linewidth=2, markersize=6, label='Mx')
    ax.plot(angles, My * 1000, 'g^-', linewidth=2, markersize=6, label='My')
    ax.plot(angles, Mz * 1000, 'rs-', linewidth=2, markersize=6, label='Mz')
    ax.set_xlabel('Vane Deflection Angle [deg]')
    ax.set_ylabel('Moment [mN-m]')
    ax.set_title('Moment Components')
    ax.grid(True, alpha=0.3)
    ax.legend()

    plt.tight_layout()
    plot_file = os.path.join(BASE_DIR, "vane_forces_plot.png")
    plt.savefig(plot_file, dpi=150, bbox_inches='tight')
    print(f"Plot saved to {plot_file}")

    # Also print a summary
    print("\nSummary:")
    print(f"{'Angle':>6} {'Fx[mN]':>10} {'Fy[mN]':>10} {'Fz[mN]':>10} {'F_lat[mN]':>10}")
    print("-" * 50)
    for i in range(len(angles)):
        print(f"{angles[i]:>6.0f} {Fx[i]*1000:>10.3f} {Fy[i]*1000:>10.3f} "
              f"{Fz[i]*1000:>10.3f} {F_lateral[i]*1000:>10.3f}")


if __name__ == '__main__':
    main()