"""Generates cambered airfoil shapes based on the NACA 4-digit series.

For info on NACA airfoils, see https://en.wikipedia.org/wiki/NACA_airfoil
"""

import argparse
import math
import sys
from typing import Tuple

try:
    import vtkmodules
except ImportError:
    print('Generates cambered airfoil shapes based on the NACA 4-digit series.')
    print('Requires vtk module to run. ("pip install vtk")')
    sys.exit(-1)

import vtkmodules.vtkInteractionStyle
import vtkmodules.vtkRenderingOpenGL2
from vtkmodules.vtkCommonCore import (
    vtkPoints
)
from vtkmodules.vtkCommonDataModel import (
    vtkCellArray,
    vtkPolyData,
    vtkPolyLine,
    vtkTriangleStrip
)
from vtkmodules.vtkCommonTransforms import vtkTransform
from vtkmodules.vtkFiltersGeneral import vtkTransformPolyDataFilter
from vtkmodules.vtkFiltersModeling import vtkLinearExtrusionFilter
from vtkmodules.vtkIOGeometry import (
    vtkOBJWriter,
    vtkSTLWriter
)
from vtkmodules.vtkIOXML import vtkXMLPolyDataWriter
from vtkmodules.vtkInteractionStyle import vtkInteractorStyleTrackballCamera
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkCamera,
    vtkPolyDataMapper,
    vtkRenderWindow,
    vtkRenderWindowInteractor,
    vtkRenderer
)


def display_polydata(polydata, title='PolyData'):
    """"""
    mapper = vtkPolyDataMapper()
    mapper.SetInputData(polydata)

    actor = vtkActor()
    actor.SetMapper(mapper)
    actor.GetProperty().SetColor(0.8, 0.8, 0.8)

    ren = vtkRenderer()
    ren.SetBackground(0.2, 0.2, 0.2)
    ren.AddActor(actor)

    renWin = vtkRenderWindow()
    renWin.SetWindowName(title)
    renWin.AddRenderer(ren)
    renWin.SetSize(800, 800)

    iren = vtkRenderWindowInteractor()
    iren.SetRenderWindow(renWin)

    style = vtkInteractorStyleTrackballCamera()
    iren.SetInteractorStyle(style)

    camera = vtkCamera()
    camera.SetPosition(0.5, 0, 10)
    camera.SetFocalPoint(0.5, 0, 0)
    camera.SetViewUp(0, 1, 0)

    ren.SetActiveCamera(camera)
    ren.ResetCamera()
    ren.ResetCameraClippingRange()

    renWin.Render()
    iren.Start()


class NACA4:
    """Generates cambered airfoil shapes based on the NACA 4-digit series.

    For info on NACA airfoils, see https://en.wikipedia.org/wiki/NACA_airfoil
    """

    def __init__(self):
        pass

    def profile_2d(self, naca4_code: str, n: int = 40) -> Tuple[list, list, list, list]:
        """Generates 2d profiles for upper and lower sides.

        Inputs:
            naca4_code: (string) NACA series 4 code, e.g., 2412.

        Returns 4 lists for x_upper, y_upper, x_lower, y_lower coords, respectively.
        Each list includes first and last points (0, 0), (1, 0).
        """
        # Validate input code
        if len(naca4_code) != 4:
            print('Invalid NACA4 code: string length must be 4 not {}'.format(len(naca4_code)))
            return None

        # Check for digits
        if not naca4_code.isdecimal():
            print('Invalid NACA4 code: some char(s) not decimal digits')
            return None

        m = float(naca4_code[0]) / 100.0  # max camber
        p = float(naca4_code[1]) / 10.0  # position
        t = float(naca4_code[2:4]) / 100.0  # max thickness
        return self.profile_2d_params(m, p, t, n)

    def profile_2d_params(self,
                          max_camber: float,
                          position: float,
                          max_thickness: float,
                          n: int = 40) -> Tuple[list, list, list, list]:
        """Generates 2d profiles for upper and lower sides.

        Best math is at http://airfoiltools.com/airfoil/naca4digit
        """
        m = max_camber
        p = position
        t = max_thickness

        print('max camber m = {}'.format(m))
        print('position p = {}'.format(p))
        print('max thickness t = {}'.format(t))

        x_upper = [None] * (n + 2)
        y_upper = [None] * (n + 2)
        x_lower = [None] * (n + 2)
        y_lower = [None] * (n + 2)

        # First point
        x_upper[0] = 0.0
        y_upper[0] = 0.0
        x_lower[0] = 0.0
        y_lower[0] = 0.0

        for i in range(1, n + 1):
            beta = math.pi * float(i) / (n + 2)
            x = (1 - math.cos(beta)) / 2.0

            if x < p:
                yc = m * (2 * p * x - x * x) / (p * p)
                grad = 2 * m * (p - x) / (p * p)
            else:
                yc = m * (1 - 2 * p + 2 * p * x - x * x) / ((1 - p) * (1 - p))
                grad = 2 * m * (p - x) / ((1 - p) * (1 - p))
            # Note: for open profile use -0.1015 for last coefficient
            yt = 5 * t * (0.2969 * math.sqrt(x) - 0.1260 * x - 0.3516 * x *
                          x + 0.2843 * x * x * x - 0.1036 * x * x * x * x)
            theta = math.atan(grad)
            xu = x - yt * math.sin(theta)
            yu = yc + yt * math.cos(theta)

            # print(x, yc, grad, yt, xu, yu)
            x_upper[i] = xu
            y_upper[i] = yu

            xl = x + yt * math.sin(theta)
            yl = yc - yt * math.cos(theta)

            x_lower[i] = xl
            y_lower[i] = yl

        # Last point
        x_upper[n + 1] = 1.0
        y_upper[n + 1] = 0.0
        x_lower[n + 1] = 1.0
        y_lower[n + 1] = 0.0

        return (x_upper, y_upper, x_lower, y_lower)

    def merge_profiles(self, x_upper, y_upper, x_lower, y_lower):
        """Combines upper and lower profiles into one."""
        # Reverse lower lists and clip off first and last entries
        x_clip_lower = x_lower[-2:0:-1]
        y_clip_lower = y_lower[-2:0:-1]
        return x_upper + x_clip_lower, y_upper + y_clip_lower

    def extrude_3d(self, profile_2d, zmin, zmax, rotate_deg=None):
        """Generates vtkPolyData by extruding the 2D profile."""
        points = vtkPoints()
        x_coords, y_coords = profile_2d
        num_points = len(x_coords)
        # print('profile_2d num_points', num_points)
        for i in range(num_points):
            point_id = points.InsertNextPoint(x_coords[i], y_coords[i], zmin)

        polydata = vtkPolyData()
        lines = vtkCellArray()
        line = vtkPolyLine()
        line.GetPointIds().SetNumberOfIds(num_points + 1)
        for i in range(0, num_points):
            line.GetPointIds().SetId(i, i)
        line.GetPointIds().SetId(num_points, 0)  # close the loop
        lines.InsertNextCell(line)

        polydata.SetPoints(points)
        polydata.SetLines(lines)
        # print(polydata)

        # Apply rotation (angle of attack) if specified
        if rotate_deg is not None:
            xform = vtkTransform()
            xform.Translate(0.5, 0.0, 0.0)
            xform.RotateZ(rotate_deg)
            xform.Translate(-0.5, 0.0, 0.0)
            # print(xform)

            filter = vtkTransformPolyDataFilter()
            filter.SetInputData(polydata)
            filter.SetTransform(xform)
            filter.Update()
            vtk_data = filter.GetOutput()
        else:
            vtk_data = polydata
        # display_polydata(polydata)

        # Extrude in Z direction
        extruder = vtkLinearExtrusionFilter()
        extruder.SetInputData(vtk_data)
        extruder.SetExtrusionTypeToNormalExtrusion()
        extruder.SetVector(0, 0, zmax - zmin)
        extruder.CappingOff()
        extruder.Update()
        extrusion = extruder.GetOutput()
        # print(extrusion)
        # display_polydata(extrusion)

        # Add end cap faces (tri strips)
        strips = extrusion.GetStrips()
        for face in range(2):
            # print('Checkpoint number of ids:', num_points)
            offset = face * num_points
            strip = vtkTriangleStrip()
            strip.GetPointIds().SetNumberOfIds(num_points)
            # print('Setting i {} id {}'.format(0, offset))
            strip.GetPointIds().SetId(0, offset)
            iupper = 1 + offset
            ilower = num_points + offset - 1
            for i in range(1, num_points - 2, 2):
                # print('Setting i {} id {}'.format(i, iupper))
                strip.GetPointIds().SetId(i, iupper)
                iupper += 1
                # print('Setting i {} id {}'.format(i+1, ilower))
                strip.GetPointIds().SetId(i + 1, ilower)
                ilower -= 1
            # print('Setting i {} id {}'.format(num_points-1, iupper))
            strip.GetPointIds().SetId(num_points - 1, iupper)
            strips.InsertNextCell(strip)

        return extrusion


if __name__ == '__main__':
    description = 'Generate airfoil geometry (obj/stl/vtk) from NACA 4-digit series.'
    epilog = 'See https://en.wikipedia.org/wiki/NACA_airfoil#Four-digit_series for airfoil info'
    parser = argparse.ArgumentParser(description=description, epilog=epilog)

    # Required arguments
    parser.add_argument('naca4_code', type=str,
                        help='naca4_code: (string) NACA series 4 code, e.g., 2412.')

    # Optional arguments
    parser.add_argument('-a', '--angle_of_attack', type=int, default=None,
                        help='rotate to angle of attack, integer between -90 and +90')
    parser.add_argument('-d', '--display', action='store_true', default=True,
                        help='Display geometry')
    parser.add_argument('-n', '--points_per_side', type=int, default=40,
                        help='number of sample points for each side (upper, lower')
    parser.add_argument('-x', '--no_output', action='store_true',
                        help='do NOT write geometry to OBJ/STL/VTP file')
    parser.add_argument('-z', '--zmin_zmax', type=int, nargs=2, default=[-1, 1],
                        help="min and max z values for extruding 2D profile to 3D surface")

    args = parser.parse_args()
    print(args)

    # Get the number
    naca4_code = sys.argv[1]
    print('Using NACA code {}'.format(args.naca4_code))

    naca4 = NACA4()
    profiles = naca4.profile_2d(args.naca4_code, n=args.points_per_side)
    if profiles is None:
        print('{} is not a valid NACA4 code'.format(args.naca4_code))
        sys.exit(-1)
    x_upper, y_upper, x_lower, y_lower = profiles
    # for i in range(len(x_upper)):
    #     print(x_upper[i], y_upper[i])

    merged_profile = naca4.merge_profiles(*profiles)
    deg = None if args.angle_of_attack is None else -args.angle_of_attack
    polydata = naca4.extrude_3d(merged_profile, *args.zmin_zmax, rotate_deg=deg)
    if not args.no_output:
        # Format filename from naca4 code and angle of attack (if any)
        suffix = ''
        if args.angle_of_attack is not None:
            sign = '-' if args.angle_of_attack < 0.0 else '+'
            suffix = '{0}{1:02d}'.format(sign, args.angle_of_attack)

        # Write STL file
        filename = 'naca{}{}.stl'.format(args.naca4_code, suffix)
        stl_writer = vtkSTLWriter()
        stl_writer.SetFileName(filename)
        stl_writer.SetInputData(polydata)
        stl_writer.Write()
        print('Wrote {}'.format(filename))

        # Write OBJ file
        filename = 'naca{}{}.obj'.format(args.naca4_code, suffix)
        obj_writer = vtkOBJWriter()
        obj_writer.SetFileName(filename)
        obj_writer.SetInputData(polydata)
        obj_writer.Write()
        print('Wrote {}'.format(filename))

        # Write VTK (.vtp) file
        filename = 'naca{}{}.vtp'.format(args.naca4_code, suffix)
        vtp_writer = vtkXMLPolyDataWriter()
        vtp_writer.SetFileName(filename)
        vtp_writer.SetInputData(polydata)
        vtp_writer.SetDataModeToAscii()
        vtp_writer.Write()
        print('Wrote {}'.format(filename))

    if args.display:
        # print(polydata)
        display_polydata(polydata, 'NACA4 Airfoil')
