import gradio as gr
import neuralfoil as nf
import numpy as np
import json
from pathlib import Path
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend

def load_airfoil_from_dat(filepath):
    """Load airfoil coordinates from a .dat file"""
    with open(filepath, 'r') as f:
        lines = f.readlines()
    
    # Skip the first line (airfoil name)
    coords = []
    for line in lines[1:]:
        line = line.strip()
        if line:
            parts = line.split()
            if len(parts) >= 2:
                try:
                    x, y = float(parts[0]), float(parts[1])
                    coords.append([x, y])
                except ValueError:
                    continue
    
    return np.array(coords)

def format_coordinates(coords_array):
    """Format coordinates array as a string for display"""
    if coords_array is None or len(coords_array) == 0:
        return ""
    
    lines = []
    for x, y in coords_array:
        lines.append(f"{x:.6f}     {y:.6f}")
    return "\n".join(lines)

def parse_coordinates(coords_text):
    """Parse coordinates from text input"""
    lines = coords_text.strip().split('\n')
    coords = []
    
    for line in lines:
        line = line.strip()
        if line:
            parts = line.split()
            if len(parts) >= 2:
                try:
                    x, y = float(parts[0]), float(parts[1])
                    coords.append([x, y])
                except ValueError:
                    continue
    
    return np.array(coords)

def create_pressure_plot(coords, result, alpha, reynolds):
    """Create a plot of the airfoil with pressure distribution"""
    try:
        fig, ax = plt.subplots(figsize=(12, 6))
        
        # Plot airfoil shape
        ax.plot(coords[:, 0], coords[:, 1], 'k-', linewidth=2, label='Airfoil')
        ax.fill(coords[:, 0], coords[:, 1], color='lightgray', alpha=0.3)
        
        # Calculate pressure coefficient from edge velocity
        # Cp = 1 - (ue/vinf)^2
        upper_ue = np.array([result[f'upper_bl_ue/vinf_{i}'][0] for i in range(32) if f'upper_bl_ue/vinf_{i}' in result])
        lower_ue = np.array([result[f'lower_bl_ue/vinf_{i}'][0] for i in range(32) if f'lower_bl_ue/vinf_{i}' in result])
        
        upper_cp_32 = 1 - upper_ue**2
        lower_cp_32 = 1 - lower_ue**2
        
        # Create x positions for 32 points (from leading edge to trailing edge)
        x_bl_upper = np.linspace(0, 1, len(upper_cp_32))
        x_bl_lower = np.linspace(0, 1, len(lower_cp_32))
        
        # Find upper and lower surface points
        # Typically airfoil coords go: TE (top) -> LE -> TE (bottom)
        le_idx = np.argmin(coords[:, 0])  # Leading edge is minimum x
        
        upper_surface = coords[:le_idx+1]  # From TE to LE (top)
        lower_surface = coords[le_idx:]     # From LE to TE (bottom)
        
        # Get high-resolution x-coordinates and interpolate Cp
        if len(upper_surface) > 1:
            x_upper_hires = upper_surface[::-1, 0]  # Reverse to go LE to TE
            y_upper_hires = upper_surface[::-1, 1]
            upper_cp_hires = np.interp(x_upper_hires, x_bl_upper, upper_cp_32)
        else:
            x_upper_hires = np.array([])
            y_upper_hires = np.array([])
            upper_cp_hires = np.array([])
            
        if len(lower_surface) > 1:
            x_lower_hires = lower_surface[:, 0]  # Already LE to TE
            y_lower_hires = lower_surface[:, 1]
            lower_cp_hires = np.interp(x_lower_hires, x_bl_lower, lower_cp_32)
        else:
            x_lower_hires = np.array([])
            y_lower_hires = np.array([])
            lower_cp_hires = np.array([])
        
        # Also get y-coordinates for 32-point data
        if len(upper_surface) > 1:
            y_upper_32 = np.interp(x_bl_upper, upper_surface[::-1, 0], upper_surface[::-1, 1])
        else:
            y_upper_32 = np.zeros_like(x_bl_upper)
            
        if len(lower_surface) > 1:
            y_lower_32 = np.interp(x_bl_lower, lower_surface[:, 0], lower_surface[:, 1])
        else:
            y_lower_32 = np.zeros_like(x_bl_lower)
        
        # Scale factor for pressure lines
        scale = 0.15 * np.max(coords[:, 1] - np.min(coords[:, 1]))
        
        # Calculate pressure line coordinates for high-resolution data
        y_upper_pressure_hires = y_upper_hires - upper_cp_hires * scale
        y_lower_pressure_hires = y_lower_hires + lower_cp_hires * scale
        
        # Calculate pressure line coordinates for 32-point data
        y_upper_pressure_32 = y_upper_32 - upper_cp_32 * scale
        y_lower_pressure_32 = y_lower_32 + lower_cp_32 * scale
        
        # Plot high-resolution pressure distribution
        ax.plot(x_upper_hires, y_upper_pressure_hires, 'b-', linewidth=2, label='Upper Surface Cp (interpolated)', alpha=0.8)
        ax.plot(x_lower_hires, y_lower_pressure_hires, 'r-', linewidth=2, label='Lower Surface Cp (interpolated)', alpha=0.8)
        
        # Plot 32-point data with markers
        ax.plot(x_bl_upper, y_upper_pressure_32, 'bo', markersize=4, label='Upper Surface Cp (32 pts)', alpha=0.6)
        ax.plot(x_bl_lower, y_lower_pressure_32, 'ro', markersize=4, label='Lower Surface Cp (32 pts)', alpha=0.6)
        
        # Fill the area between airfoil surface and pressure line
        ax.fill_between(x_upper_hires, y_upper_hires, y_upper_pressure_hires, color='blue', alpha=0.2)
        ax.fill_between(x_lower_hires, y_lower_hires, y_lower_pressure_hires, color='red', alpha=0.2)
        
        ax.legend(loc='upper right')
        
        ax.set_xlabel('x/c', fontsize=12)
        ax.set_ylabel('y/c', fontsize=12)
        ax.set_title(f'Airfoil with Pressure Distribution (α={alpha}°, Re={reynolds:.1e})', fontsize=14)
        ax.grid(True, alpha=0.3)
        ax.set_aspect('equal')
        ax.axhline(y=0, color='k', linestyle='--', alpha=0.3, linewidth=0.5)
        
        plt.tight_layout()
        return fig
        
    except Exception as e:
        # Return a simple error plot
        fig, ax = plt.subplots(figsize=(12, 6))
        ax.text(0.5, 0.5, f'Error creating plot: {str(e)}', 
                ha='center', va='center', fontsize=12)
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        return fig

def run_neuralfoil_prediction_api(coordinates, alpha, reynolds, model_size):
    """Run NeuralFoil prediction from numpy array/list - for API usage"""
    try:
        # Convert to numpy array if it's a list
        if isinstance(coordinates, list):
            coords = np.array(coordinates)
        else:
            coords = coordinates
        
        if len(coords) < 3:
            return {"error": "Invalid coordinates. Please provide at least 3 coordinate pairs."}
        
        # Run NeuralFoil analysis directly from coordinates
        result = nf.get_aero_from_coordinates(
            coordinates=coords,
            alpha=alpha,
            Re=reynolds,
            model_size=model_size
        )
        
        # Convert result to a serializable dictionary
        output = {}
        
        # Check if result is a dictionary or has attributes
        if isinstance(result, dict):
            # Result is already a dictionary - extract scalar values
            for key, val in result.items():
                if isinstance(val, np.ndarray):
                    output[key] = float(val[0]) if val.size == 1 else val.tolist()
                else:
                    output[key] = val
        else:
            # Result has attributes (older API style)
            # Standard outputs
            if hasattr(result, 'CL'):
                output['CL'] = float(result.CL) if not np.isnan(result.CL) else None
            if hasattr(result, 'CD'):
                output['CD'] = float(result.CD) if not np.isnan(result.CD) else None
            if hasattr(result, 'CM'):
                output['CM'] = float(result.CM) if not np.isnan(result.CM) else None
            
            # Transition locations (if available)
            if hasattr(result, 'Top_Xtr'):
                output['Top_Xtr'] = float(result.Top_Xtr) if not np.isnan(result.Top_Xtr) else None
            if hasattr(result, 'Bot_Xtr'):
                output['Bot_Xtr'] = float(result.Bot_Xtr) if not np.isnan(result.Bot_Xtr) else None
            
            # Confidence metric
            if hasattr(result, 'analysis_confidence'):
                output['analysis_confidence'] = float(result.analysis_confidence) if not np.isnan(result.analysis_confidence) else None
            
            # Include all other attributes
            for attr in dir(result):
                if not attr.startswith('_') and attr not in output:
                    val = getattr(result, attr)
                    if isinstance(val, (int, float, str, bool)):
                        output[attr] = val
                    elif isinstance(val, np.ndarray):
                        output[attr] = val.tolist()
        
        # Calculate pressure coefficients from edge velocity
        # Cp = 1 - (ue/vinf)^2
        upper_ue = np.array([result[f'upper_bl_ue/vinf_{i}'][0] for i in range(32) if f'upper_bl_ue/vinf_{i}' in result])
        lower_ue = np.array([result[f'lower_bl_ue/vinf_{i}'][0] for i in range(32) if f'lower_bl_ue/vinf_{i}' in result])
        
        upper_cp_32 = 1 - upper_ue**2
        lower_cp_32 = 1 - lower_ue**2
        
        # Create x positions for the 32 boundary layer stations
        x_bl_upper = np.linspace(0, 1, len(upper_cp_32))
        x_bl_lower = np.linspace(0, 1, len(lower_cp_32))
        
        # Split airfoil coordinates into upper and lower surfaces
        le_idx = np.argmin(coords[:, 0])  # Leading edge is minimum x
        upper_surface = coords[:le_idx+1]  # From TE to LE (top)
        lower_surface = coords[le_idx:]     # From LE to TE (bottom)
        
        # Get x-coordinates for upper and lower surfaces (in 0-1 range)
        if len(upper_surface) > 1:
            x_upper_coords = upper_surface[::-1, 0]  # Reverse to go LE to TE
            # Interpolate Cp from 32 points to airfoil coordinate resolution
            upper_cp_interp = np.interp(x_upper_coords, x_bl_upper, upper_cp_32)
        else:
            x_upper_coords = np.array([])
            upper_cp_interp = np.array([])
        
        if len(lower_surface) > 1:
            x_lower_coords = lower_surface[:, 0]  # Already LE to TE
            # Interpolate Cp from 32 points to airfoil coordinate resolution
            lower_cp_interp = np.interp(x_lower_coords, x_bl_lower, lower_cp_32)
        else:
            x_lower_coords = np.array([])
            lower_cp_interp = np.array([])
        
        # Add pressure coefficient arrays (both 32-point and interpolated)
        output['pressure_coefficients'] = {
            'upper_surface_cp': upper_cp_interp.tolist(),
            'lower_surface_cp': lower_cp_interp.tolist(),
            'x_upper': x_upper_coords.tolist(),
            'x_lower': x_lower_coords.tolist(),
            'upper_surface_cp_32': upper_cp_32.tolist(),
            'lower_surface_cp_32': lower_cp_32.tolist(),
            'x_upper_32': x_bl_upper.tolist(),
            'x_lower_32': x_bl_lower.tolist()
        }
        
        # Add input parameters for reference
        output['input_parameters'] = {
            'alpha_deg': alpha,
            'reynolds_number': reynolds,
            'model_size': model_size,
            'num_coordinates': len(coords)
        }
        
        return output
        
    except Exception as e:
        return {"error": str(e)}

def run_neuralfoil_prediction(coords_text, alpha, reynolds, model_size):
    """Run NeuralFoil prediction and return full JSON output and plot - for UI usage"""
    try:
        # Parse coordinates
        coords = parse_coordinates(coords_text)
        
        if len(coords) < 3:
            error_fig, ax = plt.subplots(figsize=(12, 6))
            ax.text(0.5, 0.5, 'Invalid coordinates. Please provide at least 3 coordinate pairs.', 
                   ha='center', va='center', fontsize=12)
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            return json.dumps({"error": "Invalid coordinates. Please provide at least 3 coordinate pairs."}, indent=2), error_fig
        
        # Run NeuralFoil analysis directly from coordinates
        result = nf.get_aero_from_coordinates(
            coordinates=coords,
            alpha=alpha,
            Re=reynolds,
            model_size=model_size
        )
        # print(result)
        
        # Convert result to a serializable dictionary
        output = {}
        
        # Standard outputs
        if hasattr(result, 'CL'):
            output['CL'] = float(result.CL) if not np.isnan(result.CL) else None
        if hasattr(result, 'CD'):
            output['CD'] = float(result.CD) if not np.isnan(result.CD) else None
        if hasattr(result, 'CM'):
            output['CM'] = float(result.CM) if not np.isnan(result.CM) else None
        
        # Transition locations (if available)
        if hasattr(result, 'Top_Xtr'):
            output['Top_Xtr'] = float(result.Top_Xtr) if not np.isnan(result.Top_Xtr) else None
        if hasattr(result, 'Bot_Xtr'):
            output['Bot_Xtr'] = float(result.Bot_Xtr) if not np.isnan(result.Bot_Xtr) else None
        
        # Confidence metric
        if hasattr(result, 'analysis_confidence'):
            output['analysis_confidence'] = float(result.analysis_confidence) if not np.isnan(result.analysis_confidence) else None
        
        # Include all other attributes
        for attr in dir(result):
            if not attr.startswith('_') and attr not in output:
                val = getattr(result, attr)
                if isinstance(val, (int, float, str, bool)):
                    output[attr] = val
                elif isinstance(val, np.ndarray):
                    output[attr] = val.tolist()
        
        # Calculate pressure coefficients from edge velocity
        # Cp = 1 - (ue/vinf)^2
        upper_ue = np.array([result[f'upper_bl_ue/vinf_{i}'][0] for i in range(32) if f'upper_bl_ue/vinf_{i}' in result])
        lower_ue = np.array([result[f'lower_bl_ue/vinf_{i}'][0] for i in range(32) if f'lower_bl_ue/vinf_{i}' in result])
        
        upper_cp_32 = 1 - upper_ue**2
        lower_cp_32 = 1 - lower_ue**2
        
        # Create x positions for the 32 boundary layer stations
        x_bl_upper = np.linspace(0, 1, len(upper_cp_32))
        x_bl_lower = np.linspace(0, 1, len(lower_cp_32))
        
        # Split airfoil coordinates into upper and lower surfaces
        le_idx = np.argmin(coords[:, 0])  # Leading edge is minimum x
        upper_surface = coords[:le_idx+1]  # From TE to LE (top)
        lower_surface = coords[le_idx:]     # From LE to TE (bottom)
        
        # Get x-coordinates for upper and lower surfaces (in 0-1 range)
        if len(upper_surface) > 1:
            x_upper_coords = upper_surface[::-1, 0]  # Reverse to go LE to TE
            # Interpolate Cp from 32 points to airfoil coordinate resolution
            upper_cp_interp = np.interp(x_upper_coords, x_bl_upper, upper_cp_32)
        else:
            x_upper_coords = np.array([])
            upper_cp_interp = np.array([])
        
        if len(lower_surface) > 1:
            x_lower_coords = lower_surface[:, 0]  # Already LE to TE
            # Interpolate Cp from 32 points to airfoil coordinate resolution
            lower_cp_interp = np.interp(x_lower_coords, x_bl_lower, lower_cp_32)
        else:
            x_lower_coords = np.array([])
            lower_cp_interp = np.array([])
        
        # Add pressure coefficient arrays (both 32-point and interpolated)
        output['pressure_coefficients'] = {
            'upper_surface_cp': upper_cp_interp.tolist(),
            'lower_surface_cp': lower_cp_interp.tolist(),
            'x_upper': x_upper_coords.tolist(),
            'x_lower': x_lower_coords.tolist(),
            'upper_surface_cp_32': upper_cp_32.tolist(),
            'lower_surface_cp_32': lower_cp_32.tolist(),
            'x_upper_32': x_bl_upper.tolist(),
            'x_lower_32': x_bl_lower.tolist()
        }
        
        # Add input parameters for reference
        output['input_parameters'] = {
            'alpha_deg': alpha,
            'reynolds_number': reynolds,
            'model_size': model_size,
            'num_coordinates': len(coords)
        }

        # Use the API function to get results
        output = run_neuralfoil_prediction_api(coords, alpha, reynolds, model_size)
        
        if "error" in output:
            error_fig, ax = plt.subplots(figsize=(12, 6))
            ax.text(0.5, 0.5, output["error"], ha='center', va='center', fontsize=12)
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            return json.dumps(output, indent=2), error_fig
        
        # Get result back for plotting
        result = nf.get_aero_from_coordinates(
            coordinates=coords,
            alpha=alpha,
            Re=reynolds,
            model_size=model_size
        )
        
        # Create the pressure plot
        fig = create_pressure_plot(coords, result, alpha, reynolds)
        
        return json.dumps(output, indent=2), fig
        
    except Exception as e:
        error_fig, ax = plt.subplots(figsize=(12, 6))
        ax.text(0.5, 0.5, f'Error: {str(e)}', ha='center', va='center', fontsize=12)
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        return json.dumps({"error": str(e)}, indent=2), error_fig

def load_example(example_name):
    """Load an example airfoil"""
    example_files = {
        "NACA 4412": "examples/naca4412.dat",
        "Clark Y": "examples/clarky.dat",
        "RAE 2822": "examples/rae2822.dat"
    }
    
    filepath = example_files.get(example_name)
    if filepath and Path(filepath).exists():
        coords = load_airfoil_from_dat(filepath)
        return format_coordinates(coords)
    return ""

# Load default airfoil (RAE 2822)
default_coords = format_coordinates(load_airfoil_from_dat("examples/rae2822.dat"))

# Create Gradio interface
with gr.Blocks(title="NeuralFoil Airfoil Predictor") as demo:
    gr.Markdown("# NeuralFoil Airfoil Predictor")
    gr.Markdown("""
    This app uses [NeuralFoil](https://github.com/peterdsharpe/NeuralFoil) to predict airfoil aerodynamics.
    
    Provide airfoil coordinates (x, y pairs, one per line) and operating conditions to get predictions for CL, CD, CM, and more.
    """)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Airfoil Coordinates")
            gr.Markdown("Enter x,y coordinate pairs (one per line). Coordinates should trace the airfoil from trailing edge, over the top, to leading edge, then back along the bottom.")
            
            coords_input = gr.Textbox(
                label="Airfoil Coordinates",
                value=default_coords,
                lines=15,
                max_lines=30,
                placeholder="x     y\n1.0   0.0\n0.95  0.01\n..."
            )
            
            gr.Markdown("### Load Example")
            example_buttons = gr.Radio(
                choices=["NACA 4412", "Clark Y", "RAE 2822"],
                label="Example Airfoils",
                value="RAE 2822"
            )
            
            load_btn = gr.Button("Load Example")
            
            gr.Markdown("### Operating Conditions")
            alpha_input = gr.Slider(
                minimum=-10,
                maximum=20,
                value=5.0,
                step=0.5,
                label="Angle of Attack α [deg]"
            )
            
            reynolds_input = gr.Number(
                value=1e6,
                label="Reynolds Number Re [-]"
            )
            
            model_size_input = gr.Dropdown(
                choices=["xxsmall", "xsmall", "small", "medium", "large", "xlarge", "xxlarge", "xxxlarge"],
                value="large",
                label="Model Size"
            )
            
            predict_btn = gr.Button("Run Prediction", variant="primary")
        
        with gr.Column():
            gr.Markdown("### Airfoil with Pressure Distribution")
            output_plot = gr.Plot(label="Pressure Distribution")
            
            gr.Markdown("### Full NeuralFoil Output (JSON)")
            output_json = gr.Textbox(
                label="Prediction Results",
                lines=20,
                max_lines=30,
                placeholder="Results will appear here..."
            )
    
    # Event handlers
    load_btn.click(
        fn=load_example,
        inputs=[example_buttons],
        outputs=[coords_input]
    )
    
    predict_btn.click(
        fn=run_neuralfoil_prediction,
        inputs=[coords_input, alpha_input, reynolds_input, model_size_input],
        outputs=[output_json, output_plot]
    )
    
    gr.Markdown("""
    ---
    ### About
    
    **NeuralFoil** is a neural-network-based surrogate for XFoil that predicts airfoil aerodynamics much faster than traditional CFD.
    
    **Citation**: If you use NeuralFoil, please cite the [GitHub repository](https://github.com/peterdsharpe/NeuralFoil) and Peter Sharpe's PhD thesis.
    """)

# Create API endpoint that accepts numpy arrays
api = gr.Interface(
    fn=run_neuralfoil_prediction_api,
    inputs=[
        gr.JSON(label="Coordinates (2D array: [[x1,y1], [x2,y2], ...])"),
        gr.Number(label="Angle of Attack α [deg]"),
        gr.Number(label="Reynolds Number Re [-]"),
        gr.Dropdown(choices=["xxsmall", "xsmall", "small", "medium", "large", "xlarge", "xxlarge", "xxxlarge"], label="Model Size")
    ],
    outputs=gr.JSON(label="Prediction Results"),
    title="NeuralFoil API",
    description="API endpoint for programmatic access. Input coordinates as a 2D array."
)

# Combine both interfaces in tabs
app = gr.TabbedInterface(
    [demo, api],
    ["Interactive UI", "API"]
)

if __name__ == "__main__":
    app.launch()