import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad

# Panel class to store panel properties
class Panel:
    def __init__(self, xa, ya, xb, yb):
        self.xa, self.ya = xa, ya  # Start point
        self.xb, self.yb = xb, yb  # End point
        self.xc = (xa + xb) / 2    # Control point (midpoint)
        self.yc = (ya + yb) / 2
        self.length = np.sqrt((xb - xa)**2 + (yb - ya)**2)
        self.normal = np.array([-(yb - ya), xb - xa]) / self.length
        self.tangent = np.array([xb - xa, yb - ya]) / self.length
        self.sigma = 0.0  # Source strength (or slope for linear)
        self.mu = 0.0     # Doublet strength (or slope for linear)
        self.vt = 0.0     # Tangential velocity
        self.cp = 0.0     # Pressure coefficient
        self.mu_prev = 0.0  # Previous doublet for unsteady

# Wake vortex class
class WakeVortex:
    def __init__(self, x, y, gamma):
        self.x = x
        self.y = y
        self.gamma = gamma  # Vortex strength

def generate_naca0012(n_panels):
    """Generate NACA0012 airfoil coordinates with n_panels."""
    x = np.linspace(0, 1, n_panels // 2 + 1)
    t = 0.12  # Maximum thickness
    yt = 5 * t * (0.2969 * np.sqrt(x) - 0.1260 * x - 0.3516 * x**2 + 0.2843 * x**3 - 0.1015 * x**4)
    x_upper = x[:-1]
    y_upper = yt[:-1]
    x_lower = x[1:][::-1]
    y_lower = -yt[1:][::-1]
    x = np.concatenate([x_upper, x_lower])
    y = np.concatenate([y_upper, y_lower])
    return x, y

def create_panels(x, y):
    """Create panels from airfoil coordinates."""
    panels = []
    for i in range(len(x) - 1):
        panels.append(Panel(x[i], y[i], x[i+1], y[i+1]))
    return panels

def apply_motion(x, y, t, motion_func):
    """Apply rigid body motion to airfoil coordinates."""
    dx, dy, theta, dx_dt, dy_dt, omega = motion_func(t)
    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)
    x_new = cos_theta * x - sin_theta * y + dx
    y_new = sin_theta * x + cos_theta * y + dy
    return x_new, y_new, np.array([dx_dt, dy_dt]), omega

def source_influence(x, y, panel, method='constant'):
    """Compute velocity induced by a source panel at point (x, y)."""
    epsilon = 1e-10  # Small offset to prevent division by zero
    if method == 'constant':
        def integrand(s):
            xp = panel.xa + s * (panel.xb - panel.xa) / panel.length
            yp = panel.ya + s * (panel.yb - panel.ya) / panel.length
            r = np.sqrt((x - xp)**2 + (y - yp)**2 + epsilon)
            return (1 / (2 * np.pi)) * ((x - xp) * panel.normal[0] + (y - yp) * panel.normal[1]) / r**2
        u, _ = quad(lambda s: integrand(s) * (panel.xb - panel.xa) / panel.length, 0, panel.length)
        v, _ = quad(lambda s: integrand(s) * (panel.yb - panel.ya) / panel.length, 0, panel.length)
        return panel.sigma * u, panel.sigma * v
    else:  # linear
        def integrand(s, weight):
            xp = panel.xa + s * (panel.xb - panel.xa) / panel.length
            yp = panel.ya + s * (panel.yb - panel.ya) / panel.length
            r = np.sqrt((x - xp)**2 + (y - yp)**2 + epsilon)
            return weight * (1 / (2 * np.pi)) * ((x - xp) * panel.normal[0] + (y - yp) * panel.normal[1]) / r**2
        u1, _ = quad(lambda s: integrand(s, s / panel.length), 0, panel.length)
        v1, _ = quad(lambda s: integrand(s, s / panel.length) * (panel.yb - panel.ya) / panel.length, 0, panel.length)
        u2, _ = quad(lambda s: integrand(s, 1 - s / panel.length), 0, panel.length)
        v2, _ = quad(lambda s: integrand(s, 1 - s / panel.length) * (panel.yb - panel.ya) / panel.length, 0, panel.length)
        return panel.sigma * (u1 + u2), panel.sigma * (v1 + v2)

def doublet_influence(x, y, panel, method='constant'):
    """Compute velocity induced by a doublet panel at point (x, y)."""
    epsilon = 1e-10
    if method == 'constant':
        def integrand(s):
            xp = panel.xa + s * (panel.xb - panel.xa) / panel.length
            yp = panel.ya + s * (panel.yb - panel.ya) / panel.length
            r = np.sqrt((x - xp)**2 + (y - yp)**2 + epsilon)
            return (1 / (2 * np.pi)) * ((y - yp) * panel.normal[0] - (x - xp) * panel.normal[1]) / r**2
        u, _ = quad(lambda s: integrand(s) * (panel.xb - panel.xa) / panel.length, 0, panel.length)
        v, _ = quad(lambda s: integrand(s) * (panel.yb - panel.ya) / panel.length, 0, panel.length)
        return panel.mu * u, panel.mu * v
    else:  # linear
        def integrand(s, weight):
            xp = panel.xa + s * (panel.xb - panel.xa) / panel.length
            yp = panel.ya + s * (panel.yb - panel.ya) / panel.length
            r = np.sqrt((x - xp)**2 + (y - yp)**2 + epsilon)
            return weight * (1 / (2 * np.pi)) * ((y - yp) * panel.normal[0] - (x - xp) * panel.normal[1]) / r**2
        u1, _ = quad(lambda s: integrand(s, s / panel.length), 0, panel.length)
        v1, _ = quad(lambda s: integrand(s, s / panel.length) * (panel.yb - panel.ya) / panel.length, 0, panel.length)
        u2, _ = quad(lambda s: integrand(s, 1 - s / panel.length), 0, panel.length)
        v2, _ = quad(lambda s: integrand(s, 1 - s / panel.length) * (panel.yb - panel.ya) / panel.length, 0, panel.length)
        return panel.mu * (u1 + u2), panel.mu * (v1 + v2)

def vortex_influence(x, y, vortex):
    """Compute velocity induced by a point vortex at point (x, y)."""
    epsilon = 1e-10
    r = np.sqrt((x - vortex.x)**2 + (y - vortex.y)**2 + epsilon)
    u = -vortex.gamma / (2 * np.pi * r**2) * (y - vortex.y)
    v = vortex.gamma / (2 * np.pi * r**2) * (x - vortex.x)
    return u, v

def build_matrix(panels, wake_vortices, freestream, body_vel, omega, unsteady=True, method='constant'):
    """Build influence coefficient matrix."""
    n = len(panels)
    A = np.zeros((n + 1, n + 1))
    b = np.zeros(n + 1)
    
    for i, panel_i in enumerate(panels):
        for j, panel_j in enumerate(panels):
            u_d, v_d = doublet_influence(panel_i.xc, panel_i.yc, panel_j, method)
            u_s, v_s = source_influence(panel_i.xc, panel_i.yc, panel_j, method)
            A[i, j] = (u_d + u_s) * panel_i.normal[0] + (v_d + v_s) * panel_i.normal[1]
        
        if unsteady and wake_vortices:
            for vortex in wake_vortices:
                u_v, v_v = vortex_influence(panel_i.xc, panel_i.yc, vortex)
                b[i] -= (u_v * panel_i.normal[0] + v_v * panel_i.normal[1])
        
        rel_vel = freestream - (body_vel + omega * np.array([-panel_i.yc, panel_i.xc]))
        b[i] -= rel_vel[0] * panel_i.normal[0] + rel_vel[1] * panel_i.normal[1]
    
    # Kutta condition: Equal tangential velocities at trailing edge panels
    for j, panel_j in enumerate(panels):
        u_d1, v_d1 = doublet_influence(panels[0].xc, panels[0].yc, panel_j, method)
        u_s1, v_s1 = source_influence(panels[0].xc, panels[0].yc, panel_j, method)
        u_d2, v_d2 = doublet_influence(panels[-1].xc, panels[-1].yc, panel_j, method)
        u_s2, v_s2 = source_influence(panels[-1].xc, panels[-1].yc, panel_j, method)
        A[n, j] = ((u_d1 + u_s1) * panels[0].tangent[0] + (v_d1 + v_s1) * panels[0].tangent[1]) - \
                  ((u_d2 + u_s2) * panels[-1].tangent[0] + (v_d2 + v_s2) * panels[-1].tangent[1])
    
    if unsteady and wake_vortices:
        u_v1, v_v1 = 0, 0
        u_v2, v_v2 = 0, 0
        for vortex in wake_vortices:
            u_v1 += vortex_influence(panels[0].xc, panels[0].yc, vortex)[0]
            v_v1 += vortex_influence(panels[0].xc, panels[0].yc, vortex)[1]
            u_v2 += vortex_influence(panels[-1].xc, panels[-1].yc, vortex)[0]
            v_v2 += vortex_influence(panels[-1].xc, panels[-1].yc, vortex)[1]
        b[n] = -((u_v1 * panels[0].tangent[0] + v_v1 * panels[0].tangent[1]) - \
                 (u_v2 * panels[-1].tangent[0] + v_v2 * panels[-1].tangent[1]))
    
    return A, b

def solve_system(panels, wake_vortices, freestream, body_vel, omega, unsteady=True, method='constant'):
    """Solve for source and doublet strengths."""
    A, b = build_matrix(panels, wake_vortices, freestream, body_vel, omega, unsteady, method)
    try:
        strengths = np.linalg.solve(A, b)
    except np.linalg.LinAlgError:
        print("Warning: Matrix A is singular or ill-conditioned.")
        strengths = np.linalg.lstsq(A, b, rcond=None)[0]
    for i, panel in enumerate(panels):
        panel.sigma = strengths[i]
        panel.mu = strengths[i]
    if unsteady:
        return strengths[-1]  # Wake vortex strength
    return 0

def compute_velocity(panels, wake_vortices, freestream, body_vel, omega, x, y, unsteady=True, method='constant'):
    """Compute velocity at point (x, y)."""
    u, v = freestream - (body_vel + omega * np.array([-y, x]))
    for panel in panels:
        u_s, v_s = source_influence(x, y, panel, method)
        u_d, v_d = doublet_influence(x, y, panel, method)
        u += panel.sigma * u_s + panel.mu * u_d
        v += panel.sigma * v_s + panel.mu * v_d
    if unsteady:
        for vortex in wake_vortices:
            u_v, v_v = vortex_influence(x, y, vortex)
            u += u_v
            v += v_v
    return u, v

def compute_pressure(panels, freestream, body_vel, omega, V_inf, dt=0, unsteady=True, method='constant'):
    """Compute pressure coefficients."""
    for panel in panels:
        u, v = compute_velocity(panels, [], freestream, body_vel, omega, panel.xc, panel.yc, unsteady=False, method=method)
        panel.vt = u * panel.tangent[0] + v * panel.tangent[1]
        if unsteady:
            dmu_dt = (panel.mu - panel.mu_prev) / dt if hasattr(panel, 'mu_prev') else 0
            panel.cp = 1 - (panel.vt / V_inf)**2 - 2 * dmu_dt / V_inf
        else:
            panel.cp = 1 - (panel.vt / V_inf)**2
        panel.mu_prev = panel.mu

def compute_forces_moments(panels, V_inf, rho=1.225, c=1.0, unsteady=True):
    """Compute lift, drag, and moment coefficients."""
    cl, cd, cm = 0, 0, 0
    for panel in panels:
        fx = -rho * panel.length * panel.cp * panel.normal[0]
        fy = -rho * panel.length * panel.cp * panel.normal[1]
        cl += fy / (0.5 * rho * V_inf**2 * c)
        cd += fx / (0.5 * rho * V_inf**2 * c)
        cm += -fx * (panel.xc - 0.25) - fy * panel.yc / (0.5 * rho * V_inf**2 * c**2)
    return cl, cd, cm

def update_wake(panels, wake_vortices, freestream, body_vel, omega, gamma_wake, dt, V_inf, method='constant'):
    """Update wake vortex positions and add new vortex."""
    te_x = (panels[0].xa + panels[-1].xa) / 2
    te_y = (panels[0].ya + panels[-1].ya) / 2
    wake_vortices.append(WakeVortex(te_x + V_inf * dt / 3, te_y, gamma_wake))
    
    for vortex in wake_vortices:
        u, v = compute_velocity(panels, wake_vortices, freestream, body_vel, omega, vortex.x, vortex.y, method=method)
        vortex.x += u * dt
        vortex.y += v * dt

def plot_geometry(panels, wake_vortices=None, t=None, filename=None):
    """Plot airfoil panels and wake vortices."""
    plt.figure()
    for panel in panels:
        plt.plot([panel.xa, panel.xb], [panel.ya, panel.yb], 'b-', lw=1)
    if wake_vortices:
        wake_x = [vortex.x for vortex in wake_vortices]
        wake_y = [vortex.y for vortex in wake_vortices]
        plt.scatter(wake_x, wake_y, c='r', marker='o', s=20, label='Wake Vortices')
    plt.xlabel('x')
    plt.ylabel('y')
    if t is not None:
        plt.title(f'Geometry at t = {t:.2f} s')
    plt.axis('equal')
    plt.legend()
    if filename:
        plt.savefig(filename)
        plt.close()
    else:
        plt.show()

def wagner_function(s):
    """Wagner function for lift response to step change in angle of attack."""
    return 1 - 0.165 * np.exp(-0.0455 * s) - 0.335 * np.exp(-0.3 * s)

def run_simulation(x, y, freestream, V_inf, dt, t_max, motion_func, unsteady=True, method='constant', plot_steps=False):
    """Run unsteady simulation with user-defined motion."""
    wake_vortices = []
    cl_history, cd_history, cm_history = [], [], []
    time = np.arange(0, t_max, dt)
    
    for i, t in enumerate(time):
        x_new, y_new, body_vel, omega = apply_motion(x, y, t, motion_func)
        panels = create_panels(x_new, y_new)
        gamma_wake = solve_system(panels, wake_vortices, freestream, body_vel, omega, unsteady, method)
        compute_pressure(panels, freestream, body_vel, omega, V_inf, dt, unsteady, method)
        cl, cd, cm = compute_forces_moments(panels, V_inf, unsteady)
        cl_history.append(cl)
        cd_history.append(cd)
        cm_history.append(cm)
        if plot_steps and i % 10 == 0:  # Plot every 10 steps
            plot_geometry(panels, wake_vortices, t, f'geometry_t_{t:.2f}.png')
        if unsteady:
            update_wake(panels, wake_vortices, freestream, body_vel, omega, gamma_wake, dt, V_inf, method)
    
    return time, cl_history, cd_history, cm_history

# Test Case 1: Wagner Function Validation
def test_wagner(method='constant'):
    """Validate against Wagner function for step change in angle of attack."""
    n_panels = 100
    alpha = np.deg2rad(5)
    V_inf = 1.0
    dt = 0.1
    t_max = 20
    
    def motion_func(t):
        return 0, 0, alpha, 0, 0, 0
    
    x, y = generate_naca0012(n_panels)
    freestream = np.array([V_inf * np.cos(alpha), V_inf * np.sin(alpha)])
    time, cl_history, _, _ = run_simulation(x, y, freestream, V_inf, dt, t_max, motion_func, method=method, plot_steps=True)
    s = 2 * V_inf * time
    cl_theory = 2 * np.pi * np.sin(alpha) * wagner_function(s)
    
    plt.figure()
    plt.plot(s, cl_history, label=f'Panel Method ({method})')
    plt.plot(s, cl_theory, '--', label='Wagner Function')
    plt.xlabel('Non-dimensional time (2Vt/c)')
    plt.ylabel('Cl')
    plt.legend()
    plt.title(f'Wagner Function Validation ({method})')
    plt.show()

# Test Case 2: Cl vs. Angle of Attack
def test_cl_alpha(method='constant'):
    """Compute Cl vs. angle of attack for steady-state."""
    n_panels = 100
    V_inf = 1.0
    alphas = np.deg2rad(np.linspace(-10, 10, 21))
    cl_values = []
    
    for alpha in alphas:
        print(alpha)
        x, y = generate_naca0012(n_panels)
        panels = create_panels(x, y)
        freestream = np.array([V_inf * np.cos(alpha), V_inf * np.sin(alpha)])
        solve_system(panels, [], freestream, np.array([0, 0]), 0, unsteady=False, method=method)
        compute_pressure(panels, freestream, np.array([0, 0]), 0, V_inf, unsteady=False, method=method)
        cl, _, _ = compute_forces_moments(panels, V_inf, unsteady=False)
        cl_values.append(cl)
    
    plt.figure()
    plt.plot(np.rad2deg(alphas), cl_values, label=f'Panel Method ({method})')
    plt.plot(np.rad2deg(alphas), 2 * np.pi * alphas, '--', label='Theory (2π sin(α))')
    plt.xlabel('Angle of Attack (deg)')
    plt.ylabel('Cl')
    plt.legend()
    plt.title(f'Cl vs. Angle of Attack ({method})')
    plt.show()

# Test Case 3: Sinusoidal Pitching Motion
def test_sinusoidal_pitch(method='constant'):
    """Simulate sinusoidal pitching motion."""
    n_panels = 100
    V_inf = 1.0
    dt = 0.05
    t_max = 10
    omega = 0.5
    alpha_mean = np.deg2rad(5)
    alpha_amplitude = np.deg2rad(2)
    
    def motion_func(t):
        alpha = alpha_mean + alpha_amplitude * np.sin(omega * t)
        omega_body = alpha_amplitude * omega * np.cos(omega * t)
        return 0, 0, alpha, 0, 0, omega_body
    
    x, y = generate_naca0012(n_panels)
    freestream = np.array([V_inf, 0])
    time, cl_history, _, _ = run_simulation(x, y, freestream, V_inf, dt, t_max, motion_func, method=method, plot_steps=True)
    
    plt.figure()
    plt.plot(time, cl_history, label='Cl')
    plt.xlabel('Time (s)')
    plt.ylabel('Cl')
    plt.title(f'Sinusoidal Pitching Motion ({method})')
    plt.legend()
    plt.show()

# Test Case 4: Plunging Motion
def test_plunging(method='constant'):
    """Simulate plunging motion."""
    n_panels = 100
    V_inf = 1.0
    dt = 0.05
    t_max = 10
    h_amplitude = 0.1
    omega = 0.5
    
    def motion_func(t):
        h = h_amplitude * np.sin(omega * t)
        h_dot = h_amplitude * omega * np.cos(omega * t)
        return 0, h, 0, 0, h_dot, 0
    
    x, y = generate_naca0012(n_panels)
    freestream = np.array([V_inf, 0])
    time, cl_history, _, _ = run_simulation(x, y, freestream, V_inf, dt, t_max, motion_func, method=method, plot_steps=True)
    
    plt.figure()
    plt.plot(time, cl_history, label='Cl')
    plt.xlabel('Time (s)')
    plt.ylabel('Cl')
    plt.title(f'Plunging Motion ({method})')
    plt.legend()
    plt.show()

# Test Case 5: Convergence Test
def test_convergence(method='constant'):
    """Test convergence with increasing number of panels."""
    V_inf = 1.0
    alpha = np.deg2rad(5)
    dt = 0.1
    t_max = 5
    panel_counts = [50, 100, 200]
    cl_histories = []
    
    def motion_func(t):
        return 0, 0, alpha, 0, 0, 0
    
    x_base, y_base = generate_naca0012(max(panel_counts))
    freestream = np.array([V_inf * np.cos(alpha), V_inf * np.sin(alpha)])
    
    for n_panels in panel_counts:
        idx = np.linspace(0, len(x_base)-1, n_panels, dtype=int)
        x = x_base[idx]
        y = y_base[idx]
        time, cl_history, _, _ = run_simulation(x, y, freestream, V_inf, dt, t_max, motion_func, method=method)
        cl_histories.append(cl_history)
    
    plt.figure()
    for i, n in enumerate(panel_counts):
        plt.plot(time, cl_histories[i], label=f'N = {n}')
    plt.xlabel('Time (s)')
    plt.ylabel('Cl')
    plt.title(f'Convergence Test ({method})')
    plt.legend()
    plt.show()

if __name__ == "__main__":
    #for method in ['constant', 'linear']:
    for method in ['linear']:
        print('method', method)
        test_cl_alpha(method)
        #plt.show()

        #test_wagner(method)
        #test_sinusoidal_pitch(method)
        #test_plunging(method)
        #test_convergence(method)

