#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright 2020 University of Liège
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


## Polar
# Adrien Crovato
#
# Compute aerodynamic load coefficients as a function of angle of attack

import math
import dart.utils as dU
import tbox.utils as tU
import fwk
from fwk.coloring import ccolors
from ..core import init_dart

class Polar:
    """Utility to compute the polar of a lifting surface

    Attributes
    ----------
    tms : fwk.Timers object
        dictionary of timers
    dim : int
        problem dimensions
    format : str
        output format type
    slice : array of float
        arrays of y-coordinates to perform slices (for 3D only)
    tag : int
        index of physical tag to slice (see gmsh)
    mach : float
        freestream Mach number
    alphas : array of float
        range of angles of attack to compute
    Cls : array of float
        lift coefficients for the different AoAs
    Cds : array of float
        drag coefficients for the different AoAs
    Cms : array of float
        moment coefficients for the different AoAs
    """
    def __init__(self, p):
        """Setup and configure

        Parameters
        ----------
        p : dict
            dictionary of parameters to configure DART
        """
        # basic init
        self.tms = fwk.Timers()
        self.tms["total"].start()
        _dart = init_dart(p, scenario='aerodynamic', task='analysis')
        self.dim, self.msh, self.wrt, self.pbl, self.bnd, self.sol = (_dart.get(key) for key in ['dim', 'msh', 'wrt', 'pbl', 'bnd', 'sol'])
        # save some parameters for later use
        self.format = p.get('Format', 'vtk')
        self.slice = p.get('Slice')
        self.tag = p.get('TagId')
        self.mach = self.pbl.M_inf
        self.alphas = tU.myrange(p['AoA_begin'], p['AoA_end'], p['AoA_step'])

    def run(self):
        """Compute the polar
        """
        # initialize the loop
        self.Cls = []
        self.Cds = []
        self.Cms = []
        print(ccolors.ANSI_BLUE + 'Sweeping AoA from', self.alphas[0], 'to', self.alphas[-1], ccolors.ANSI_RESET)
        for i in range(len(self.alphas)):
            # define current angle of attack
            alpha = self.alphas[i]*math.pi/180
            acs = '_{:04d}'.format(i)
            # update problem and reset ICs to improve robustness in transonic cases
            self.pbl.update(alpha)
            self.sol.reset()
            # run the solver and save the results
            print(ccolors.ANSI_BLUE + 'Running the solver for', self.alphas[i], 'degrees', ccolors.ANSI_RESET)
            self.tms["solver"].start()
            self.sol.run()
            self.tms["solver"].stop()
            self.sol.save(self.wrt, acs)
            # extract Cp
            if self.dim == 2:
                Cp = dU.extract(self.bnd.groups[0].tag.elems, self.sol.Cp)
                tU.write(Cp, f'Cp_{self.msh.name}_airfoil{acs}.dat', '%1.4e', ',', 'alpha = '+str(alpha*180/math.pi)+' deg\nx, y, cp', '')
            elif self.dim == 3 and self.format == 'vtk' and self.slice:
                dU.write_slices(self.msh.name, self.slice, self.tag, acs)
            # extract force coefficients
            self.Cls.append(self.sol.Cl)
            self.Cds.append(self.sol.Cd)
            self.Cms.append(self.sol.Cm)

    def disp(self):
        """Display the results and draw the polar
        """
        # display results
        print(ccolors.ANSI_BLUE + 'Airfoil polar' + ccolors.ANSI_RESET)
        print("       M    alpha       Cl       Cd       Cm")
        i = 0
        while i < len(self.alphas):
            print("{0:8.2f} {1:8.1f} {2:8.4f} {3:8.4f} {4:8.4f}".format(self.mach, self.alphas[i], self.Cls[i], self.Cds[i], self.Cms[i]))
            i = i+1
        print('\n')
        # display timers
        self.tms["total"].stop()
        print(ccolors.ANSI_BLUE + 'CPU statistics' + ccolors.ANSI_RESET)
        print(self.tms)
        # plot results
        if self.alphas[0] != self.alphas[-1]:
            tU.plot(self.alphas, self.Cls, {'xlabel': 'alpha', 'ylabel': 'Cl'})
            tU.plot(self.alphas, self.Cds, {'xlabel': 'alpha', 'ylabel': 'Cd'})
            tU.plot(self.alphas, self.Cms, {'xlabel': 'alpha', 'ylabel': 'Cm'})