"""
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 creator.base as base

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


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

    The coordinates are saved as two lists
    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.
    """

    # TODO: default values in separate module
    def __init__(self, chord, semi_span, material):
        super().__init__()
        # self.x = np.array([])
        # self.z = np.array([])
        # self.chord = chord
        """Create airfoil from its chord and semi-span."""
        self.chord = chord if chord > 20 else 20
        if chord <= 20:
            logging.debug('Chord too small, using minimum value of 20.')
        self.semi_span = semi_span
        self.material = material
        self.naca_num = int()

    def __str__(self):
        return type(self).__name__

    def add_naca(self, naca_num):
        """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, airfoil, loc_percent, material):
        """Set spar location as percent of chord length."""
        super().__init__()
        super().set_material(material)
        self.cap_area = float()
        loc = loc_percent * airfoil.chord
        # bi.bisect_left: returns index of first value in airfoil.x > loc
        # This ensures that spar geom intersects with airfoil geom.
        # Spar upper coordinates
        spar_u = bi.bisect_left(airfoil.x, loc) - 1
        self.x = np.append(self.x, airfoil.x[spar_u])
        self.z = np.append(self.z, airfoil.z[spar_u])
        # Spar lower coordinates
        spar_l = bi.bisect_left(airfoil.x[::-1], loc)
        self.x = np.append(self.x, airfoil.x[-spar_l])
        self.z = np.append(self.z, airfoil.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):
        super().__init__()
        self.x_start = []
        self.x_end = []
        self.z_start = []
        self.z_end = []
        self.diameter = float()
        self.area = float()

    def add_coord(self, airfoil, spars, stringer_u_1, stringer_u_2,
                  stringer_l_1, stringer_l_2):
        """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
        stringer_u_1: upper nose number of stringers
        stringer_u_2: upper surface number of stringers
        stringer_l_1: lower nose number of stringers
        stringer_l_2: lower surface number of stringers

        Returns:
        None
        """

        # Find distance between leading edge and first upper stringer
        interval = spars.x[0][0] / (stringer_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, stringer_u_1):
            # Index of the first value of airfoil.x > x
            i = bi.bisect_left(airfoil.x, x)
            self.x.append(airfoil.x[i])
            self.z.append(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.spar.x[-1][0] -
                    airfoil.spar.x[0][0]) / (stringer_u_2 + 1)
        x = interval + airfoil.spar.x[0][0]
        for _ in range(0, stringer_u_2):
            i = bi.bisect_left(airfoil.x, x)
            self.x.append(airfoil.x[i])
            self.z.append(airfoil.z[i])
            x += interval

        # Find distance between leading edge and first lower stringer
        interval = airfoil.spar.x[0][1] / (stringer_l_1 + 1)
        x = interval
        # Add lower stringers from leading edge until first spar.
        for _ in range(0, stringer_l_1):
            i = bi.bisect_left(airfoil.x[::-1], x)
            self.x.append(airfoil.x[-i])
            self.z.append(airfoil.z[-i])
            x += interval
            # Add lower stringers from first spar until last spar
        interval = (airfoil.spar.x[-1][1] -
                    airfoil.spar.x[0][1]) / (stringer_l_2 + 1)
        x = interval + airfoil.spar.x[0][1]
        for _ in range(0, stringer_l_2):
            i = bi.bisect_left(airfoil.x[::-1], x)
            self.x.append(airfoil.x[-i])
            self.z.append(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):
        super().info_print(round)
        print('Stringer Area:\n', np.around(self.area, round))
        return None


def plot_geom(airfoil, spars, stringers):
    """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')

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

    # Graph formatting
    # plot_bound = np.amax(airfoil.x)
    ax.set(
        title='NACA ' + str(airfoil.naca_num) + ' airfoil',
        xlabel='X axis',
        # xlim=[-0.10 * plot_bound, 1.10 * plot_bound],
        ylabel='Z axis')
    # ylim=[-(1.10 * plot_bound / 2), (1.10 * plot_bound / 2)])

    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


def main():
    return None


if __name__ == '__main__':
    main()