#    Copyright 2023 ONERA - contact luis.bernardos@onera.fr
#
#    This file is part of MOLA.
#
#    MOLA is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Lesser General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    MOLA is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public License
#    along with MOLA.  If not, see <http://www.gnu.org/licenses/>.

from typing import Union
import numpy as np
from mola.logging import mola_logger, MolaAssertionError
from mola import naming_conventions as names
from mola.cfd.preprocess.mesh import tools as mesh_tools
from mola.workflow import WorkflowManager

class WorkflowTurbomachineryManager(WorkflowManager):

    def add_isospeed_line(self, throttles, ShaftRotationSpeed=None, ParallelMode=False, initialize_from_previous=True):
        '''
        Add N operating points to compute on a same iso-speed line, where N is the size of **throttles**.
        The outflow boundary conditions is automatically detected with its `Type` starting by `Outflow`.

        .. warning::
            Cuurently, this method is designed to work with only one outflow boundary.

        Parameters
        ----------
        throttles : list or numpy.ndarray
            Values to modify to determine operating points to compute. 
            Depending on the outflow boundary condition, it could be the value of 
            `Pressure`, `MassFlow` or `valve_relax` for instance.
        ShaftRotationSpeed : float, optional
            Value of rotation speed for the iso-speed line, in rad/s.
            If not given, it is taken from Workflow attribute `ApplicationContext`.
        ParallelMode : bool, optional
            If True, all operating points are computed in parallel.
            If False, operating points on a same iso-speed line are computed sequentially.
            By default False.
        initialize_from_previous : bool, optional
            If True, each operating point on an iso-speed line is initialized by the previous point.
            The first operating point is initialized according the Workflow attribute `Initialization`.
            By default True.

        Example
        -------

        .. code-block:: python

            from mola.workflow.rotating_component import turbomachinery 

            workflow = turbomachinery.Workflow(...)
            manager = turbomachinery.WorkflowManager(workflow, '/the/path/to/my/parametric/study/')

            # Lines specific to this manager
            manager.add_isospeed_line([1e5, 1.1e5, 1.2e5])
            manager.add_isospeed_line([1e5, 1.1e5, 1.2e5], ShaftRotationSpeed=5000.) # new isospeed line for another rotation speed
            
            manager.prepare()  # at the end of this command, all the directories are deployed and job files are ready to be submitted
            manager.submit()  # submit all the jobs generated by the previous line
        '''
        if ShaftRotationSpeed is None:
            ShaftRotationSpeed = self.base_workflow.ApplicationContext.get('ShaftRotationSpeed')

        job_name = self.base_workflow.RunManagement.get('JobName', 'isospeed')

        if ParallelMode and initialize_from_previous:
            mola_logger.user_warning('Because ParallelMode=True, initialize_from_previous is set to False.')
            initialize_from_previous = False

        outflow_family, throttle_key = self._get_family_and_throttle_key()

        if not ParallelMode:
            self.new_job(f'isospeed_{30/np.pi*ShaftRotationSpeed:.0f}rpm')

        for throttle in throttles:
            if ParallelMode:
                self.new_job(f'{throttle_key}_{throttle:.2f}')
            self.add_variations(
                [
                    ('RunManagement|JobName', f'{job_name}_{throttle:.2f}'),
                    ('RunManagement|RunDirectory', f'{throttle_key}_{throttle:.2f}'),
                    ('ApplicationContext|ShaftRotationSpeed', ShaftRotationSpeed),
                    (f'BoundaryConditions|Family={outflow_family}|{throttle_key}', throttle),
                ], 
                initialize_from_previous = initialize_from_previous
                )
        
    def _get_family_and_throttle_key(self):
        outflow_bc = mesh_tools.get_bc_from_bc_type(self.base_workflow, 'Outflow*')
        outflow_family = outflow_bc["Family"]
        throttle_key = self._get_throttle_key(outflow_bc)
        return outflow_family, throttle_key
    
    @staticmethod
    def _get_throttle_key(outflow_bc):

        if outflow_bc['Type'] == 'OutflowPressure':
            return 'Pressure'

        elif outflow_bc['Type'] == 'OutflowMassFlow':
            return 'MassFlow'
        
        elif outflow_bc['Type'] == 'OutflowRadialEquilibrium':

            for key in ['MassFlow', 'PressureAtHub', 'PressureAtShroud', 'PressureAtSpecifiedHeight']:
                if key in outflow_bc:
                    return key
            
            # otherwise, a ValveLaw must be defined
            if outflow_bc['ValveLaw']['Type'] == 'Linear':
                return 'PressureRef'
            elif outflow_bc['ValveLaw']['Type'] == 'Quadratic':
                return 'ValveCoefficient'
            
        else:
            raise MolaAssertionError(f"Outflow BC type {outflow_bc['Type']} not supported by WorkflowTurbomachineryManager")
    
    def gather_performance(self, 
                           stage:Union[tuple, str], 
                           filename:Union[str, None]=None,
                           update_from_remote_machine=True) -> dict:
        upstream_plane, downstream_plane = self._get_planes_names_for_perfo(stage)

        suffix = names.CONTAINER_OUTPUT_FIELDS_AT_VERTEX.split('#')[1]

        queries = [
            f'CGNSTree/Averages0D/{downstream_plane}/{names.CONTAINER_OUTPUT_FIELDS_AT_VERTEX}/Massflow',
            f'CGNSTree/Averages0D/{downstream_plane}/Comparison#{upstream_plane}#{suffix}/StagnationPressureRatio',
            f'CGNSTree/Averages0D/{downstream_plane}/Comparison#{upstream_plane}#{suffix}/IsentropicEfficiency',
        ]
        perfo_data = self.gather_signals(queries, filename=filename, keep_last_point=True, update_from_remote_machine=update_from_remote_machine)

        perfo_data = self._rearange_performance(perfo_data)

        return perfo_data
    
    def _get_planes_names_for_perfo(self, stage:Union[tuple, str]) -> tuple:
        # Get upstream and downstream rows from stage
        if isinstance(stage, (tuple, list)):
            assert len(stage) == 2
            row1, row2 = stage
        elif isinstance(stage, str):
            row1 = row2 = stage
        else:
            raise MolaAssertionError(f'stage must be either a str or a tuple, and the given type is {type(stage)}')
        
        # Check that rows are compatible with workflow
        for row in (row1, row2):
            if row not in self.base_workflow.ApplicationContext['Rows']:
                raise MolaAssertionError(f'The given row "{row}" is not in ApplicationContext["Rows"]')

        upstream_plane = None
        downstream_plane = None
        for extraction in self.base_workflow.Extractions:
            if extraction['Type'] == 'IsoSurface':
                try:
                    if (extraction['OtherOptions']['ReferenceRow'] == row1
                        and extraction['OtherOptions']['tag'] == 'InletPlane'):
                        upstream_plane = extraction['Name']
                except: 
                    pass
                try:
                    if (extraction['OtherOptions']['ReferenceRow'] == row2
                        and extraction['OtherOptions']['tag'] == 'OutletPlane'):
                        downstream_plane = extraction['Name']
                except: 
                    pass

        if not upstream_plane:
            raise MolaAssertionError(f'Cannot find upstream plane on {row1}')
        if not downstream_plane:
            raise MolaAssertionError(f'Cannot find downstream plane on {row2}')

        return upstream_plane, downstream_plane

    def _rearange_performance(self, perfo):
        ordered_perfo = dict()
        for dir_iso, perfo_on_iso in self.rearange_signals(perfo).items():
            rpm = float(dir_iso.replace('isospeed_', '').replace('rpm', ''))

            ordered_perfo[rpm] = perfo_on_iso
            cases = ordered_perfo[rpm].pop('case') 
            throttle_list = [float(case.split('_')[-1]) for case in cases]
            ordered_perfo[rpm]['throttle'] = np.array(throttle_list)
        
        return ordered_perfo

    def plot_isospeed_lines(self, perfo, filename='isoSpeedLines.png'):
        import matplotlib.pyplot as plt

        linestyles = [dict(linestyle=ls, marker=mk) for mk in ['o', 's', 'd', 'h']
                                                for ls in ['-', ':', '--', '-.']]
        fig, ax1 = plt.subplots()

        # Total pressure ratio
        color = 'teal'
        ax1.set_xlabel('MassFlow (kg/s)')
        ax1.set_ylabel('Total pressure ratio (-)', color=color)
        i = 0
        for rpm, perfo_iso in perfo.items():
            ax1.plot(perfo_iso['MassFlow'], perfo_iso['PressureStagnationRatio'],
                    color=color, 
                    label=f'{rpm} rpm', 
                    **linestyles[i])
            i += 1
        ax1.tick_params(axis='y', labelcolor=color)

        # Isentropic efficiency
        color = 'firebrick'
        ax2 = ax1.twinx()
        ax2.set_ylabel('Isentropic efficiency (-)', color=color)
        i = 0
        for rpm, perfo_iso in perfo.items():
            ax2.plot(perfo_iso['MassFlow'], perfo_iso['EfficiencyIsentropic'],
                    color=color, 
                    label=f'{rpm} rpm', 
                    **linestyles[i])
            # To display legend in black
            ax2.plot([], [], color='k', label=f'{rpm} rpm', **linestyles[i])
            i += 1
        ax2.tick_params(axis='y', labelcolor=color)
        ax2.set_ylim(top=1)

        if len(perfo) > 1:
            ax2.legend(loc='center left', bbox_to_anchor=(1.1, 0.5))

        fig.tight_layout()
        plt.savefig(filename, dpi=300)
        plt.show()