"""
The wing.py module contains class definitions for and various components
we add to an airfoil (spars, stringers, and ribs).

Classes:
    Airfoil: instantiated with class method to provide coordinates to heirs.
    Spar: inherits from Airfoil.
    Stringer: also inherits from Airfoil.

Functions:
    plot_geom(airfoil): generates a 2D plot of the airfoil & any components.
"""

import logging
import numpy as np
from math import sin, cos, atan
import bisect as bi
import matplotlib.pyplot as plt

import creator.base as base
import resources.materials as mt


class Airfoil(base.Component):
    """This class represents a single NACA airfoil.

    The coordinates are saved as two np.arrays
    for the x- and z-coordinates. The coordinates start at
    the leading edge, travel over the airfoil's upper edge,
    then loop back to the leading edge via the lower edge.

    This method was chosen for easier future exports
    to 3D CAD packages like SolidWorks, which can import such
    geometry as coordinates written in a CSV file.
    """
    def __init__(self,
                 parent,
                 name,
                 chord=68,
                 semi_span=150,
                 material=mt.aluminium):
        super().__init__(parent, name)
        parent.wing = self
        if chord > 20:
            self.chord = chord
        else:
            self.chord = 20
            logging.debug('Chord too small, using minimum value of 20.')
        parent
        self.semi_span = semi_span
        self.material = material
        self.spars = []
        self.stringers = []

    def add_naca(self, naca_num=2412):
        """Generate surface geometry for a NACA airfoil.

        The nested functions perform the required steps to generate geometry,
        and can be called to solve the geometry y-coordinate for any 'x' input.
        Equation coefficients were retrieved from Wikipedia.org.

        Parameters:
        naca_num: 4-digit NACA wing

        Return:
        None
        """
        self.naca_num = naca_num
        # Variables extracted from naca_num argument passed to the function
        m = int(str(naca_num)[0]) / 100
        p = int(str(naca_num)[1]) / 10
        t = int(str(naca_num)[2:]) / 100
        # x-coordinate of maximum camber
        p_c = p * self.chord

        def get_camber(x):
            """
            Returns camber z-coordinate from 1 'x' along the airfoil chord.
            """
            z_c = float()
            if 0 <= x < p_c:
                z_c = (m / (p**2)) * (2 * p * (x / self.chord) -
                                      (x / self.chord)**2)
            elif p_c <= x <= self.chord:
                z_c = (m /
                       ((1 - p)**2)) * ((1 - 2 * p) + 2 * p *
                                        (x / self.chord) - (x / self.chord)**2)
            return (z_c * self.chord)

        def get_thickness(x):
            """Return thickness from 1 'x' along the airfoil chord."""
            x = 0 if x < 0 else x
            z_t = 5 * t * self.chord * (+0.2969 *
                                        (x / self.chord)**0.5 - 0.1260 *
                                        (x / self.chord)**1 - 0.3516 *
                                        (x / self.chord)**2 + 0.2843 *
                                        (x / self.chord)**3 - 0.1015 *
                                        (x / self.chord)**4)
            return z_t

        def get_theta(x):
            dz_c = float()
            if 0 <= x < p_c:
                dz_c = ((2 * m) / p**2) * (p - x / self.chord)
            elif p_c <= x <= self.chord:
                dz_c = (2 * m) / ((1 - p)**2) * (p - x / self.chord)

            theta = atan(dz_c)
            return theta

        def get_coord_u(x):
            x = x - get_thickness(x) * sin(get_theta(x))
            z = get_camber(x) + get_thickness(x) * cos(get_theta(x))
            return (x, z)

        def get_coord_l(x):
            x = x + get_thickness(x) * sin(get_theta(x))
            z = get_camber(x) - get_thickness(x) * cos(get_theta(x))
            return (x, z)

        # Densify x-coordinates 10 times for first 1/4 chord length
        x_chord_25_percent = round(self.chord / 4)
        x_chord = [i / 10 for i in range(x_chord_25_percent * 10)]
        x_chord.extend(i for i in range(x_chord_25_percent, self.chord + 1))
        # Generate our airfoil skin geometry from previous sub-functions
        self.x_c = np.array([])
        self.z_c = np.array([])
        # Upper surface and camber line
        for x in x_chord:
            self.x_c = np.append(self.x_c, x)
            self.z_c = np.append(self.z_c, get_camber(x))
            self.x = np.append(self.x, get_coord_u(x)[0])
            self.z = np.append(self.z, get_coord_u(x)[1])
        # Lower surface
        for x in x_chord[::-1]:
            self.x = np.append(self.x, get_coord_l(x)[0])
            self.z = np.append(self.z, get_coord_l(x)[1])
        return None


class Spar(base.Component):
    """Contains a single spar's data."""
    def __init__(self, parent, name, loc_percent=0.30, material=mt.aluminium):
        """Set spar location as percent of chord length."""
        super().__init__(parent, name)
        parent.spars.append(self)
        self.material = material
        self.cap_area = float()
        # bi.bisect_left: returns index of first value in parent.x > loc
        # This ensures that spar geom intersects with airfoil geom.
        loc = loc_percent * parent.chord
        # Spar upper coordinates
        spar_u = bi.bisect_left(parent.x, loc) - 1
        self.x = np.append(self.x, parent.x[spar_u])
        self.z = np.append(self.z, parent.z[spar_u])
        # Spar lower coordinates
        spar_l = bi.bisect_left(parent.x[::-1], loc)
        self.x = np.append(self.x, parent.x[-spar_l])
        self.z = np.append(self.z, parent.z[-spar_l])
        return None

    def set_cap_area(self, cap_area):
        self.cap_area = cap_area
        return None

    def set_mass(self, mass):
        self.mass = mass
        return None


class Stringer(base.Component):
    """Contains the coordinates of all stringers."""
    def __init__(self, parent, name):
        super().__init__(parent, name)
        parent.stringers = self
        self.x_start = []
        self.x_end = []
        self.z_start = []
        self.z_end = []
        self.diameter = float()
        self.area = float()

    def add_coord(self, airfoil, den_u_1=4, den_u_2=4, den_l_1=4, den_l_2=4):
        """Add equally distributed stringers to four airfoil locations
        (upper nose, lower nose, upper surface, lower surface).

        Parameters:
        airfoil_coord: packed airfoil coordinates
        spar_coord: packed spar coordinates
        den_u_1: upper nose number of stringers
        den_u_2: upper surface number of stringers
        den_l_1: lower nose number of stringers
        den_l_2: lower surface number of stringers

        Returns:
        None
        """

        # Find distance between leading edge and first upper stringer
        interval = airfoil.spars[0].x[0] / (den_u_1 + 1)
        # initialise first self.stringer_x at first interval
        x = interval
        # Add upper stringers from leading edge until first spar.
        for _ in range(0, den_u_1):
            # Index of the first value of airfoil.x > x
            i = bi.bisect_left(airfoil.x, x)
            self.x = np.append(self.x, airfoil.x[i])
            self.z = np.append(self.z, airfoil.z[i])
            x += interval
            # Add upper stringers from first spar until last spar
            # TODO: stringer placement if only one spar is created
        interval = (airfoil.spars[-1].x[0] -
                    airfoil.spars[0].x[0]) / (den_u_2 + 1)
        x = interval + airfoil.spars[0].x[0]
        for _ in range(0, den_u_2):
            i = bi.bisect_left(airfoil.x, x)
            self.x = np.append(self.x, airfoil.x[i])
            self.z = np.append(self.z, airfoil.z[i])
            x += interval

        # Find distance between leading edge and first lower stringer
        interval = airfoil.spars[0].x[1] / (den_l_1 + 1)
        x = interval
        # Add lower stringers from leading edge until first spar.
        for _ in range(0, den_l_1):
            i = bi.bisect_left(airfoil.x[::-1], x)
            self.x = np.append(self.x, airfoil.x[-i])
            self.z = np.append(self.z, airfoil.z[-i])
            x += interval
            # Add lower stringers from first spar until last spar
        interval = (airfoil.spars[-1].x[1] -
                    airfoil.spars[0].x[1]) / (den_l_2 + 1)
        x = interval + airfoil.spars[0].x[1]
        for _ in range(0, den_l_2):
            i = bi.bisect_left(airfoil.x[::-1], x)
            self.x = np.append(self.x, airfoil.x[-i])
            self.z = np.append(self.z, airfoil.z[-i])
            x += interval
        return None

    def add_area(self, area):
        self.area = area
        return None

    def add_mass(self, mass):
        self.mass = len(self.x) * mass + len(self.x) * mass
        return None

    def add_webs(self, thickness):
        """Add webs to stringers."""
        for _ in range(len(self.x) // 2):
            self.x_start.append(self.x[_])
            self.x_end.append(self.x[_ + 1])
            self.z_start.append(self.z[_])
            self.z_end.append(self.z[_ + 1])
            self.thickness = thickness
        return None

    def info_print(self, round=2):
        super().info_print(round)
        print('Stringer Area:\n', np.around(self.area, round))
        return None


def plot_geom(airfoil):
    """This function plots the airfoil's + sub-components' geometry."""
    fig, ax = plt.subplots()

    # Plot chord
    x = [0, airfoil.chord]
    y = [0, 0]
    ax.plot(x, y, linewidth='1')
    # Plot quarter chord
    ax.plot(airfoil.chord / 4,
            0,
            '.',
            color='g',
            markersize=24,
            label='Quarter-chord')
    # Plot mean camber line
    ax.plot(airfoil.x_c,
            airfoil.z_c,
            '-.',
            color='r',
            linewidth='2',
            label='Mean camber line')
    # Plot airfoil surfaces
    ax.plot(airfoil.x, airfoil.z, color='b', linewidth='1')

    try:  # Plot spars
        for spar in airfoil.spars:
            x = (spar.x)
            y = (spar.z)
            ax.plot(x, y, '-', color='y', linewidth='4')
    except AttributeError:
        print('No spars to plot.')
    try:  # Plot stringers
        for i in range(len(airfoil.stringers.x)):
            x = airfoil.stringers.x[i]
            y = airfoil.stringers.z[i]
            ax.plot(x, y, '.', color='y', markersize=12)
    except AttributeError:
        print('No stringers to plot.')

    ax.set(title='NACA ' + str(airfoil.naca_num) + ' airfoil',
           xlabel='X axis',
           ylabel='Z axis')

    plt.grid(axis='both', linestyle=':', linewidth=1)
    plt.gca().set_aspect('equal', adjustable='box')
    plt.gca().legend(bbox_to_anchor=(1, 1),
                     bbox_transform=plt.gcf().transFigure)
    plt.show()
    return fig, ax