#!/usr/bin/env python
# -*- coding: latin-1; -*-

'''

Copyright 2018 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.

algorithm.py
Defines the coupling algorithms of CUPyDO.
Authors : David THOMAS, Marco Lucio CERQUAGLIA, Romain BOMAN

'''

# ----------------------------------------------------------------------
#  Imports
# ----------------------------------------------------------------------

from math import *
import numpy as np
import scipy as sp
import traceback
import copy

import ccupydo
from utilities import *
from interfaceData import FlexInterfaceData

np.set_printoptions(threshold=np.nan)

# ----------------------------------------------------------------------
#    Algorithm class
# ----------------------------------------------------------------------

class Algorithm:
    """
    Des.
    """

    def __init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, deltaT, totTime, timeIterTreshold=-1, mpiComm=None):
        """
        Des.
        """

        mpiPrint('\n***************************** Initializing FSI algorithm *****************************', mpiComm)

        self.mpiComm = mpiComm
        self.manager = Manager
        self.FluidSolver = FluidSolver
        self.SolidSolver = SolidSolver
        self.interfaceInterpolator = InterfaceInterpolator

        self.globalTimer = Timer()
        self.communicationTimer = Timer()
        self.meshDefTimer = Timer()
        self.fluidSolverTimer = Timer()
        self.solidSolverTimer = Timer()
        self.solidRemeshingTimer = Timer()
        self.fluidRemeshingTimer = Timer()

        self.deltaT = deltaT
        self.totTime = totTime
        self.timeIterTreshold = timeIterTreshold

        self.time = 0.0
        self.timeIter = 0

        if self.mpiComm != None:
            self.myid = self.mpiComm.Get_rank()
            self.mpiSize = self.mpiComm.Get_size()
        else:
            self.myid = 0
            self.mpiSize = 1

        self.solidHasRun = False

    def setFSIInitialConditions(self):
        """
        Des.
        """

        if self.manager.mechanical:
            if self.manager.computationType == 'unsteady':
                if self.myid in self.manager.getSolidSolverProcessors():
                    self.SolidSolver.setInitialDisplacements()
                self.interfaceInterpolator.getDisplacementFromSolidSolver()
                self.interfaceInterpolator.interpolateSolidDisplacementOnFluidMesh()
                self.interfaceInterpolator.setDisplacementToFluidSolver(self.time)
                self.FluidSolver.setInitialMeshDeformation()
            else:
                self.interfaceInterpolator.getDisplacementFromSolidSolver()

        if self.manager.thermal:
            if self.interfaceInterpolator.chtTransferMethod == 'hFFB' or self.interfaceInterpolator.chtTransferMethod == 'TFFB':
                self.FluidSolver.setInitialInterfaceHeatFlux()
            elif self.interfaceInterpolator.chtTransferMethod == 'hFTB' or self.interfaceInterpolator.chtTransferMethod == 'FFTB':
                self.FluidSolver.setInitialInterfaceTemperature()
            self.FluidSolver.boundaryConditionsUpdate()

    def fluidToSolidMechaTransfer(self):
        """
        Des.
        """

        self.communicationTimer.start()
        self.interfaceInterpolator.getLoadsFromFluidSolver()
        self.interfaceInterpolator.interpolateFluidLoadsOnSolidMesh()
        self.interfaceInterpolator.setLoadsToSolidSolver(self.time)
        self.communicationTimer.stop()
        self.communicationTimer.cumul()

    def solidToFluidMechaTransfer(self):
        """
        Des.
        """

        self.communicationTimer.start()
        self.interfaceInterpolator.interpolateSolidDisplacementOnFluidMesh()
        self.interfaceInterpolator.setDisplacementToFluidSolver(self.time)
        self.communicationTimer.stop()
        self.communicationTimer.cumul()

    def solidToFluidThermalTransfer(self):
        """
        Des.
        """

        self.communicationTimer.start()
        if self.interfaceInterpolator.chtTransferMethod == 'TFFB' or self.interfaceInterpolator.chtTransferMethod == 'hFFB':
            self.interfaceInterpolator.interpolateSolidHeatFluxOnFluidMesh()
            self.interfaceInterpolator.setHeatFluxToFluidSolver(self.time)
        elif self.interfaceInterpolator.chtTransferMethod == 'FFTB' or self.interfaceInterpolator.chtTransferMethod == 'hFTB':
            self.interfaceInterpolator.interpolateSolidTemperatureOnFluidMesh()
            self.interfaceInterpolator.setTemperatureToFluidSolver(self.time)
        self.communicationTimer.stop()
        self.communicationTimer.cumul()

    def fluidToSolidThermalTransfer(self):
        """
        Des.
        """

        self.communicationTimer.start()
        if self.interfaceInterpolator.chtTransferMethod == 'TFFB':
            self.interfaceInterpolator.getTemperatureFromFluidSolver()
            self.interfaceInterpolator.interpolateFluidTemperatureOnSolidMesh()
            self.interfaceInterpolator.setTemperatureToSolidSolver(self.time)
        elif self.interfaceInterpolator.chtTransferMethod == 'FFTB':
            self.interfaceInterpolator.getHeatFluxFromFluidSolver()
            self.interfaceInterpolator.interpolateFluidHeatFluxOnSolidMesh()
            self.interfaceInterpolator.setHeatFluxToSolidSolver(self.time)
        elif self.interfaceInterpolator.chtTransferMethod == 'hFFB' or self.interfaceInterpolator.chtTransferMethod == 'hFTB':
            self.interfaceInterpolator.getRobinTemperatureFromFluidSolver()
            self.interfaceInterpolator.interpolateFluidRobinTemperatureOnSolidMesh()
            self.interfaceInterpolator.setRobinHeatFluxToSolidSolver(self.time)
        self.communicationTimer.stop()
        self.communicationTimer.cumul()

class AlgorithmExplicit(Algorithm):
    """
    Des.
    """

    def __init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, deltaT, totTime, timeIterTreshold=-1, mpiComm=None):
        """
        Des.
        """

        Algorithm.__init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, deltaT, totTime, timeIterTreshold, mpiComm)

    def run(self):
        """
        Des.
        """

        # --- Initialize output manager --- #
        self.iniRealTimeData()

        mpiPrint('\n**********************************', self.mpiComm)
        mpiPrint('*         Begin FSI computation            *', self.mpiComm)
        mpiPrint('**********************************\n', self.mpiComm)

        self.globalTimer.start()

        #If no restart
        mpiPrint('Setting FSI initial conditions...', self.mpiComm)
        self.setFSIInitialConditions()
        mpiPrint('\nFSI initial conditions are set', self.mpiComm)

        try:
            if self.manager.computationType == 'unsteady':
                self.__unsteadyRun()
            else:
                raise Exception('Explicit coupling is only valid for unsteady computations!')
        except:
            mpiPrint('\nA DIVINE ERROR OCCURED...EXITING COMPUTATION\n', self.mpiComm)
            traceback.print_exc()
        finally:
            self.globalTimer.stop()
            self.globalTimer.cumul()

            mpiBarrier(self.mpiComm)

            mpiPrint('\n*************************', self.mpiComm)
            mpiPrint('*    End FSI computation    *', self.mpiComm)
            mpiPrint('*************************\n', self.mpiComm)

            self.printExitInfo()

            # --- Exit the solid solver --- #
            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.exit()

            # --- Exit the fluid solver --- #
            self.FluidSolver.exit()

            # --- Exit computation --- #
            mpiBarrier(self.mpiComm)

    def __unsteadyRun(self):
        """
        Des.
        """

        #If no restart
        nbTimeIter = int((self.totTime/self.deltaT)-1)

        mpiPrint('Begin time integration\n', self.mpiComm)

        # --- External temporal loop --- #
        while self.timeIter <= nbTimeIter:

            mpiPrint("\n>>>> Time iteration {} <<<<".format(self.timeIter), self.mpiComm)

            # --- Preprocess the temporal iteration --- #
            self.FluidSolver.preprocessTimeIter(self.timeIter)
            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.preprocessTimeIter(self.timeIter)

            # --- Internal FSI loop --- #
            self.fsiCoupling()
            # --- End of FSI loop --- #

            mpiBarrier(self.mpiComm)

            # --- Update the fluid and solid solver for the next time step --- #
            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.update()
            self.FluidSolver.update(self.deltaT)

            # --- Write fluid and solid solution, update FSI history  ---#
            self.FluidSolver.save(self.timeIter)

            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.save()

            self.writeRealTimeData()

            # --- Perform some remeshing if necessary
            if self.myid in self.manager.getSolidSolverProcessors():
                self.solidRemeshingTimer.start()
                self.SolidSolver.remeshing()
                self.solidRemeshingTimer.stop()
                self.solidRemeshingTimer.cumul()

            self.fluidRemeshingTimer.start()
            self.FluidSolver.remeshing()
            self.fluidRemeshingTimer.stop()
            self.fluidRemeshingTimer.cumul()
            # ---

            self.timeIter += 1
            self.time += self.deltaT
        # --- End of the temporal loop --- #

    def fsiCoupling(self):
        """
        FSI explicit coupling
        """

        if self.timeIter > self.timeIterTreshold:
            mpiPrint('\n*************** Enter explicit FSI coupling ***************', self.mpiComm)

        if self.manager.mechanical:
            # --- Solid to fluid mechanical transfer --- #
            self.interfaceInterpolator.getDisplacementFromSolidSolver()
            self.solidToFluidMechaTransfer()
            # --- Fluid mesh morphing --- #
            mpiPrint('\nPerforming mesh deformation...\n', self.mpiComm)
            self.meshDefTimer.start()
            self.FluidSolver.meshUpdate(self.timeIter)
            self.meshDefTimer.stop()
            self.meshDefTimer.cumul()

        if self.manager.thermal and self.solidHasRun:
            # --- Solid to fluid thermal transfer --- #
            self.solidToFluidThermalTransfer()
        self.FluidSolver.boundaryConditionsUpdate()

        # --- Fluid solver call  --- #
        mpiPrint('\nLaunching fluid solver...', self.mpiComm)
        self.fluidSolverTimer.start()
        self.FluidSolver.run(self.time-self.deltaT, self.time)
        self.fluidSolverTimer.stop()
        self.fluidSolverTimer.cumul()
        mpiBarrier(self.mpiComm)

        if self.timeIter > self.timeIterTreshold:
            if self.manager.mechanical:
                # --- Fluid to solid mechanical transfer --- #
                mpiPrint('\nProcessing interface fluid loads...\n', self.mpiComm)
                self.fluidToSolidMechaTransfer()
            if self.manager.thermal:
                # --- Fluid to solid thermal transfer --- #
                mpiPrint('\nProcessing interface thermal quantities...\n', self.mpiComm)
                self.fluidToSolidThermalTransfer()
            mpiBarrier(self.mpiComm)

            # --- Solid solver call for FSI subiteration --- #
            mpiPrint('\nLaunching solid solver...\n', self.mpiComm)
            if self.myid in self.manager.getSolidSolverProcessors():
                self.solidSolverTimer.start()
                self.SolidSolver.run(self.time-self.deltaT, self.time)
                self.solidSolverTimer.stop()
                self.solidSolverTimer.cumul()
            self.solidHasRun = True

    def iniRealTimeData(self):
        """
        Des
        """

        if self.myid in self.manager.getSolidSolverProcessors():
            self.SolidSolver.initRealTimeData()

    def writeRealTimeData(self):
        """
        Des
        """

        if self.myid == 0:
            self.FluidSolver.saveRealTimeData(self.time, 0)
            if self.timeIter >= self.timeIterTreshold:
                self.SolidSolver.saveRealTimeData(self.time, 0)

    def printExitInfo(self):
        """
        Des
        """

        mpiPrint('[cpu FSI total]: ' + str(self.globalTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI fluid mesh mapping]: ' + str(self.interfaceInterpolator.mappingTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI fluid mesh deformation]: ' + str(self.meshDefTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI communications]: ' + str(self.communicationTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI fluid solver]: ' + str(self.fluidSolverTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI solid solver]: ' + str(self.solidSolverTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI fluid remeshing]: ' + str(self.fluidRemeshingTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI solid remeshing]: ' + str(self.solidRemeshingTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[Time steps FSI]: ' + str(self.timeIter), self.mpiComm)
        mpiPrint('[Successful Run FSI]: ' + str(self.time >= (self.totTime - 2*self.deltaT)), self.mpiComm) # NB: self.totTime - 2*self.deltaT is the extreme case that can be encountered due to rounding effects!
        mpiPrint('[Mean n. of FSI Iterations]: ' + str(1), self.mpiComm)

        if self.myid == 0 :
            self.FluidSolver.printRealTimeData(self.time, 0)
            self.SolidSolver.printRealTimeData(self.time, 0)

        mpiPrint('RES-FSI-FSIhistory: ' + str(self.timeIter) + '\t' + str(self.time) + '\t' + str(1.0) + '\t' + str(1) + '\n', self.mpiComm)

class AlgorithmBGSStaticRelax(Algorithm):
    """
    Des.
    """

    def __init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, Criterion, nbFSIIterMax, deltaT, totTime, timeIterTreshold=-1, omegaBoundList=[1.0,1.0], mpiComm=None):
        """
        Des.
        """

        Algorithm.__init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, deltaT, totTime, timeIterTreshold, mpiComm)

        self.criterion = Criterion

        if type(omegaBoundList) == list:        #A list specified by the user or the default one
            self.omegaBoundMecha = omegaBoundList[0]
            self.omegaBoundThermal = omegaBoundList[1]
        else:                                   #The user just put one value
            self.omegaBoundMecha = omegaBoundList
            self.omegaBoundThermal = 1.0
        self.omegaMinMecha = 1e-12
        self.omegaMinThermal = 1e-12
        self.omegaMecha = self.omegaBoundMecha
        self.omegaThermal = self.omegaBoundThermal

        self.writeInFSIloop = False

        self.FSIIter = 0
        self.errValue = 0.0
        self.FSIConv = False
        self.totNbOfFSIIt = 0
        self.nbFSIIterMax = nbFSIIterMax

        self.predictor = True
        self.predictorOrder = 2
        self.alpha_0 = 1.0
        self.alpha_1 = 0.5

        self.solidInterfaceVelocity = None
        self.solidInterfaceVelocitynM1 = None
        self.solidInterfaceResidual = None
        self.solidHeatFluxResidual = None
        self.solidTemperatureResidual = None

    def initInterfaceData(self):
        """
        Des.
        """

        ns = self.interfaceInterpolator.getNs()
        d = self.interfaceInterpolator.getd()

        # --- Initialize data for prediction (mechanical only) --- #
        if self.predictor and self.manager.mechanical:
            self.solidInterfaceVelocity = FlexInterfaceData(ns+d, 3, self.mpiComm)
            self.solidInterfaceVelocitynM1 = FlexInterfaceData(ns+d, 3, self.mpiComm)

        # --- Initialize coupling residuals --- #
        if self.manager.mechanical:
            self.solidInterfaceResidual = FlexInterfaceData(ns+d, 3, self.mpiComm)
        if self.manager.thermal:
            self.solidHeatFluxResidual = FlexInterfaceData(ns+d, 3, self.mpiComm)
            self.solidTemperatureResidual = FlexInterfaceData(ns+d, 1, self.mpiComm)

    def run(self):
        """
        Des.
        """

        # --- Initialize interface data --- #
        self.initInterfaceData()

        # --- Initialize output manager --- #
        self.iniRealTimeData()

        mpiPrint('\n**********************************', self.mpiComm)
        mpiPrint('*         Begin FSI computation            *', self.mpiComm)
        mpiPrint('**********************************\n', self.mpiComm)

        self.globalTimer.start()

        #If no restart
        mpiPrint('Setting FSI initial conditions...', self.mpiComm)
        self.setFSIInitialConditions()
        mpiPrint('\nFSI initial conditions are set', self.mpiComm)

        try:
            if self.manager.computationType == 'unsteady':
                self.__unsteadyRun()
            else:
                self.time = self.totTime
                self.timeIter = 1
                self.deltaT = self.totTime
                self.writeInFSIloop = True
                self.fsiCoupling()
                self.totNbOfFSIIt = self.FSIIter
        except:
            mpiPrint('\nA DIVINE ERROR OCCURED...EXITING COMPUTATION\n', self.mpiComm)
            traceback.print_exc()
        finally:
            self.globalTimer.stop()
            self.globalTimer.cumul()

            mpiBarrier(self.mpiComm)

            mpiPrint('\n*************************', self.mpiComm)
            mpiPrint('*    End FSI computation    *', self.mpiComm)
            mpiPrint('*************************\n', self.mpiComm)

            self.printExitInfo()

            # --- Exit the solid solver --- #
            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.exit()

            # --- Exit the fluid solver --- #
            self.FluidSolver.exit()

            # --- Exit computation --- #
            mpiBarrier(self.mpiComm)

    def __unsteadyRun(self):
        """
        Des.
        """

        #If no restart
        nbTimeIter = int((self.totTime/self.deltaT)-1)

        mpiPrint('Begin time integration\n', self.mpiComm)

        # --- External temporal loop --- #
        while self.timeIter <= nbTimeIter:

            mpiPrint("\n>>>> Time iteration {} <<<<".format(self.timeIter), self.mpiComm)

            # --- Preprocess the temporal iteration --- #
            self.FluidSolver.preprocessTimeIter(self.timeIter)
            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.preprocessTimeIter(self.timeIter)

            # --- Internal FSI loop --- #
            self.fsiCoupling()
            # --- End of FSI loop --- #

            mpiBarrier(self.mpiComm)

            if self.timeIter > 0:
                self.totNbOfFSIIt += self.FSIIter

            # --- Update the fluid and solid solver for the next time step --- #
            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.update()
            self.FluidSolver.update(self.deltaT)

            # --- Write fluid and solid solution, update FSI history  ---#
            self.FluidSolver.save(self.timeIter)

            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.save()

            self.writeRealTimeData()

            # --- Perform some remeshing if necessary
            if self.myid in self.manager.getSolidSolverProcessors():
                self.solidRemeshingTimer.start()
                self.SolidSolver.remeshing()
                self.solidRemeshingTimer.stop()
                self.solidRemeshingTimer.cumul()

            self.fluidRemeshingTimer.start()
            self.FluidSolver.remeshing()
            self.fluidRemeshingTimer.stop()
            self.fluidRemeshingTimer.cumul()
            # ---

            if self.timeIter >= self.timeIterTreshold and self.predictor:
                # --- Displacement predictor for the next time step and update of the solid solution --- #
                mpiPrint('\nSolid displacement prediction for next time step', self.mpiComm)
                self.solidDisplacementPredictor()

            self.timeIter += 1
            self.time += self.deltaT
        # --- End of the temporal loop --- #

    def iniRealTimeData(self):
        """
        Des
        """

        if self.myid in self.manager.getSolidSolverProcessors():
            self.SolidSolver.initRealTimeData()
        histFile = open('FSIhistory.ascii', "w")
        histFile.write("TimeIter\tTime\tFSIError\tCHTError\tFSINbIter\tomegaMecha\tomegaThermal\n")
        histFile.close()

    def writeRealTimeData(self):
        """
        Des
        """

        if self.myid == 0:
            self.FluidSolver.saveRealTimeData(self.time, self.FSIIter)
            if self.timeIter >= self.timeIterTreshold:
                self.SolidSolver.saveRealTimeData(self.time, self.FSIIter)
            histFile = open('FSIhistory.ascii', "a")
            histFile.write(str(self.timeIter) + '\t' + str(self.time) + '\t' + str(self.errValue) + '\t' + str(self.errValue_CHT) + '\t' + str(self.FSIIter) + '\t' + str(self.omegaMecha) + '\t' + str(self.omegaThermal) + '\n')
            histFile.close()

    def getMeanNbOfFSIIt(self):
        """
        Des
        """
        if self.manager.computationType == 'unsteady':
            if self.timeIter > 1:
                return float(self.totNbOfFSIIt)/(self.timeIter-1)
            else:
                return 0.0
        else:
            return self.FSIIter

    def printExitInfo(self):
        """
        Des
        """

        mpiPrint('[cpu FSI total]: ' + str(self.globalTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI fluid mesh mapping]: ' + str(self.interfaceInterpolator.mappingTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI fluid mesh deformation]: ' + str(self.meshDefTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI communications]: ' + str(self.communicationTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI fluid solver]: ' + str(self.fluidSolverTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI solid solver]: ' + str(self.solidSolverTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI fluid remeshing]: ' + str(self.fluidRemeshingTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[cpu FSI solid remeshing]: ' + str(self.solidRemeshingTimer.cumulTime) + ' s', self.mpiComm)
        mpiPrint('[Time steps FSI]: ' + str(self.timeIter), self.mpiComm)
        mpiPrint('[Successful Run FSI]: ' + str(self.time >= (self.totTime - 2*self.deltaT)), self.mpiComm) # NB: self.totTime - 2*self.deltaT is the extreme case that can be encountered due to rounding effects!
        mpiPrint('[Mean n. of FSI Iterations]: ' + str(self.getMeanNbOfFSIIt()), self.mpiComm)

        if self.myid == 0 :
            self.FluidSolver.printRealTimeData(self.time, self.FSIIter)
            self.SolidSolver.printRealTimeData(self.time, self.FSIIter)

        mpiPrint('RES-FSI-FSIhistory: ' + str(self.timeIter) + '\t' + str(self.time) + '\t' + str(self.errValue) + '\t' + str(self.FSIIter) + '\n', self.mpiComm)

    def fsiCoupling(self):
        """
        Block Gauss Seidel (BGS) method for strong coupling FSI
        """

        if self.timeIter > self.timeIterTreshold:
            nbFSIIter = self.nbFSIIterMax
            mpiPrint('\n*************** Enter Block Gauss Seidel (BGS) method for strong coupling FSI ***************', self.mpiComm)
        else:
             nbFSIIter = 1

        self.FSIIter = 0
        self.FSIConv = False
        self.errValue = 1e12
        self.errValue_CHT = 1e6

        while ((self.FSIIter < nbFSIIter) and (not self.criterion.isVerified(self.errValue, self.errValue_CHT))):
            mpiPrint("\n>>>> FSI iteration {} <<<<\n".format(self.FSIIter), self.mpiComm)

            if self.manager.mechanical:
                # --- Solid to fluid mechanical transfer --- #
                self.solidToFluidMechaTransfer()
                # --- Fluid mesh morphing --- #
                mpiPrint('\nPerforming mesh deformation...\n', self.mpiComm)
                self.meshDefTimer.start()
                self.FluidSolver.meshUpdate(self.timeIter)
                self.meshDefTimer.stop()
                self.meshDefTimer.cumul()
            if self.manager.thermal and self.solidHasRun:
                # --- Solid to fluid thermal transfer --- #
                self.solidToFluidThermalTransfer()
            self.FluidSolver.boundaryConditionsUpdate()

            # --- Fluid solver call for FSI subiteration --- #
            mpiPrint('\nLaunching fluid solver...', self.mpiComm)
            self.fluidSolverTimer.start()
            self.FluidSolver.run(self.time-self.deltaT, self.time)
            self.fluidSolverTimer.stop()
            self.fluidSolverTimer.cumul()
            mpiBarrier(self.mpiComm)

            if self.timeIter > self.timeIterTreshold:
                if self.manager.mechanical:
                    # --- Fluid to solid mechanical transfer --- #
                    mpiPrint('\nProcessing interface fluid loads...\n', self.mpiComm)
                    self.fluidToSolidMechaTransfer()
                if self.manager.thermal:
                    # --- Fluid to solid thermal transfer --- #
                    mpiPrint('\nProcessing interface thermal quantities...\n', self.mpiComm)
                    self.fluidToSolidThermalTransfer()
                mpiBarrier(self.mpiComm)

                # --- Solid solver call for FSI subiteration --- #
                mpiPrint('\nLaunching solid solver...\n', self.mpiComm)
                if self.myid in self.manager.getSolidSolverProcessors():
                    self.solidSolverTimer.start()
                    self.SolidSolver.run(self.time-self.deltaT, self.time)
                    self.solidSolverTimer.stop()
                    self.solidSolverTimer.cumul()
                self.solidHasRun = True

                if self.manager.mechanical:
                    # --- Compute the mechanical residual --- #
                    res = self.computeSolidInterfaceResidual()
                    self.errValue = self.criterion.update(res)
                    mpiPrint('\nFSI error value : {}\n'.format(self.errValue), self.mpiComm)
                else:
                    self.errValue = 0.0
                if self.manager.thermal:
                    # --- Compute the thermal residual --- #
                    res_CHT = self.computeSolidInterfaceResidual_CHT()
                    self.errValue_CHT = self.criterion.updateThermal(res_CHT)
                    mpiPrint('\nCHT error value : {}\n'.format(self.errValue_CHT), self.mpiComm)
                else:
                    self.errValue_CHT = 0.0

                # --- Monitor the coupling convergence --- #
                self.FSIConv = self.criterion.isVerified(self.errValue, self.errValue_CHT)

                if self.manager.mechanical:
                    # --- Relaxe the solid position --- #
                    mpiPrint('\nProcessing interface displacements...\n', self.mpiComm)
                    self.relaxSolidPosition()

                if self.manager.thermal:
                    # --- Relaxe thermal data --- #
                    self.relaxCHT()

            if self.writeInFSIloop == True:
                self.writeRealTimeData()

            self.FSIIter += 1
            if self.manager.computationType != 'unsteady':
                self.time += self.deltaT

            # --- Update the solvers for the next BGS iteration --- #
            if self.myid in self.manager.getSolidSolverProcessors():
                self.SolidSolver.bgsUpdate()
            self.FluidSolver.bgsUpdate()

        if self.timeIter > self.timeIterTreshold:
            mpiPrint('\n*************** BGS is converged ***************', self.mpiComm)

    def computeSolidInterfaceResidual(self):
        """
        Des.
        """

        ns = self.interfaceInterpolator.getNs()
        d = self.interfaceInterpolator.getd()

        # --- Get the predicted (computed) solid interface displacement from the solid solver --- #
        predictedDisplacement = FlexInterfaceData(ns+d, 3, self.mpiComm)

        if self.myid in self.manager.getSolidInterfaceProcessors():
            localSolidInterfaceDisp_X, localSolidInterfaceDisp_Y, localSolidInterfaceDisp_Z = self.SolidSolver.getNodalDisplacements()
            for iVertex in range(self.manager.getNumberOfLocalSolidInterfaceNodes()):
                iGlobalVertex = self.manager.getGlobalIndex('solid', self.myid, iVertex)
                predictedDisplacement[iGlobalVertex] = [localSolidInterfaceDisp_X[iVertex], localSolidInterfaceDisp_Y[iVertex], localSolidInterfaceDisp_Z[iVertex]]

        predictedDisplacement.assemble()

        # --- Calculate the residual (vector and norm) --- #
        mpiPrint("\nCompute FSI residual based on solid interface displacement.", self.mpiComm)
        self.solidInterfaceResidual.set(predictedDisplacement - self.interfaceInterpolator.solidInterfaceDisplacement)

        return self.solidInterfaceResidual

    def computeSolidInterfaceResidual_CHT(self):
        """
        Des.
        """

        ns = self.interfaceInterpolator.getNs()
        d = self.interfaceInterpolator.getd()

        predictedHF = FlexInterfaceData(ns+d, 3, self.mpiComm)
        predictedTemp = FlexInterfaceData(ns+d, 1, self.mpiComm)

        if self.myid in self.manager.getSolidInterfaceProcessors():
            localSolidInterfaceHeatFlux_X, localSolidInterfaceHeatFlux_Y, localSolidInterfaceHeatFlux_Z = self.SolidSolver.getNodalHeatFluxes()
            localSolidInterfaceTemperature = self.SolidSolver.getNodalTemperatures()
            for iVertex in range(self.manager.getNumberOfLocalSolidInterfaceNodes()):
                iGlobalVertex = self.manager.getGlobalIndex('solid', self.myid, iVertex)
                predictedHF[iGlobalVertex] = [localSolidInterfaceHeatFlux_X[iVertex], localSolidInterfaceHeatFlux_Y[iVertex], localSolidInterfaceHeatFlux_Z[iVertex]]
                predictedTemp[iGlobalVertex] = [localSolidInterfaceTemperature[iVertex]]

        predictedHF.assemble()
        predictedTemp.assemble()

        if self.interfaceInterpolator.chtTransferMethod == 'hFFB' or self.interfaceInterpolator.chtTransferMethod == 'TFFB':
            mpiPrint("\nCompute CHT residual based on solid interface heat flux.", self.mpiComm)
            self.solidHeatFluxResidual.set(predictedHF - self.interfaceInterpolator.solidInterfaceHeatFlux)
            return self.solidHeatFluxResidual
        elif self.interfaceInterpolator.chtTransferMethod == 'hFTB' or self.interfaceInterpolator.chtTransferMethod == 'FFTB':
            mpiPrint("\nCompute CHT residual based on solid interface temperature.", self.mpiComm)
            self.solidTemperatureResidual.set(predictedTemp - self.interfaceInterpolator.solidInterfaceTemperature)
            return self.solidTemperatureResidual
        else:
            return None

    def solidDisplacementPredictor(self):
        """
        Des
        """

        # --- Get the velocity (current and previous time step) of the solid interface from the solid solver --- #
        if self.myid in self.manager.getSolidInterfaceProcessors():
            localSolidInterfaceVel_X, localSolidInterfaceVel_Y, localSolidInterfaceVel_Z = self.SolidSolver.getNodalVelocity()
            localSolidInterfaceVelNm1_X, localSolidInterfaceVelNm1_Y, localSolidInterfaceVelNm1_Z = self.SolidSolver.getNodalVelocityNm1()
            for iVertex in range(self.manager.getNumberOfLocalSolidInterfaceNodes()):
                iGlobalVertex = self.manager.getGlobalIndex('solid', self.myid, iVertex)
                self.solidInterfaceVelocity[iGlobalVertex] = [localSolidInterfaceVel_X[iVertex], localSolidInterfaceVel_Y[iVertex], localSolidInterfaceVel_Z[iVertex]]
                self.solidInterfaceVelocitynM1[iGlobalVertex] = [localSolidInterfaceVelNm1_X[iVertex], localSolidInterfaceVelNm1_Y[iVertex], localSolidInterfaceVelNm1_Z[iVertex]]

        self.solidInterfaceVelocity.assemble()
        self.solidInterfaceVelocitynM1.assemble()

        # --- Predict the solid position for the next time step --- #
        if self.predictorOrder == 1:
            mpiPrint("First order predictor.", self.mpiComm)
            self.interfaceInterpolator.solidInterfaceDisplacement += (self.alpha_0*self.deltaT*self.solidInterfaceVelocity)
        else:
            mpiPrint("Second order predictor.", self.mpiComm)
            self.interfaceInterpolator.solidInterfaceDisplacement += (self.alpha_0*self.deltaT*self.solidInterfaceVelocity + self.alpha_1*self.deltaT*(self.solidInterfaceVelocity-self.solidInterfaceVelocitynM1))

    def setOmegaMecha(self):
        """
        Des.
        """

        self.omegaMecha = self.omegaBoundMecha

        mpiPrint('Static under-relaxation summary, mechanical : {}'.format(self.omegaMecha), self.mpiComm)

    def setOmegaThermal(self):
        """
        Des.
        """

        self.omegaThermal = self.omegaBoundThermal

        mpiPrint('Static under-relaxation summary, thermal : {}'.format(self.omegaThermal), self.mpiComm)

    def relaxSolidPosition(self):
        """
        Des.
        """

        # --- Set the relaxation parameter --- #
        self.setOmegaMecha()

        # --- Relax the solid interface position --- #
        self.interfaceInterpolator.solidInterfaceDisplacement += (self.omegaMecha*self.solidInterfaceResidual)

    def relaxCHT(self):
        """
        Des.
        """

        self.setOmegaThermal()

        if self.interfaceInterpolator.chtTransferMethod == 'hFFB' or self.interfaceInterpolator.chtTransferMethod == 'TFFB':
            self.interfaceInterpolator.solidInterfaceHeatFlux += (self.omegaThermal*self.solidHeatFluxResidual)
        elif self.interfaceInterpolator.chtTransferMethod == 'hFTB' or self.interfaceInterpolator.chtTransferMethod == 'FFTB':
            self.interfaceInterpolator.solidInterfaceTemperature += (self.omegaThermal*self.solidTemperatureResidual)

class AlgorithmBGSAitkenRelax(AlgorithmBGSStaticRelax):

    def __init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, Criterion, nbFSIIterMax, deltaT, totTime, timeIterTreshold=-1, omegaBoundList=[1.0, 1.0], mpiComm=None):
        """
        Des.
        """

        AlgorithmBGSStaticRelax.__init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, Criterion, nbFSIIterMax, deltaT, totTime, timeIterTreshold, omegaBoundList, mpiComm)

        self.solidInterfaceResidualkM1 = None
        self.solidHeatFluxResidualkM1 = None
        self.solidTemperatureResidualkM1 = None
        self.aitkenCritMecha = 'max'
        self.aitkenCritThermal = 'max'

    def initInterfaceData(self):
        """
        Des.
        """

        AlgorithmBGSStaticRelax.initInterfaceData(self)
        ns = self.interfaceInterpolator.getNs()
        d = self.interfaceInterpolator.getd()

        self.solidInterfaceResidualkM1 = FlexInterfaceData(ns+d, 3, self.mpiComm)
        self.solidHeatFluxResidualkM1 = FlexInterfaceData(ns+d, 3, self.mpiComm)
        self.solidTemperatureResidualkM1 = FlexInterfaceData(ns+d, 1, self.mpiComm)

    def setOmegaMecha(self):
        """
        Des.
        """

        if self.FSIIter != 0:
            # --- Compute the dynamic Aitken coefficient --- #
            deltaInterfaceResidual = self.solidInterfaceResidual - self.solidInterfaceResidualkM1

            prodScalRes_X, prodScalRes_Y, prodScalRes_Z = deltaInterfaceResidual.dot(self.solidInterfaceResidualkM1)
            prodScalRes = prodScalRes_X + prodScalRes_Y + prodScalRes_Z

            deltaInterfaceResidual_NormX, deltaInterfaceResidual_NormY, deltaInterfaceResidual_NormZ = deltaInterfaceResidual.norm()
            deltaResNormSquare = deltaInterfaceResidual_NormX**2 + deltaInterfaceResidual_NormY**2 + deltaInterfaceResidual_NormZ**2

            if deltaResNormSquare != 0.:
                self.omegaMecha *= -prodScalRes/deltaResNormSquare
            else:
                self.omegaMecha = self.omegaMinMecha

        else:
            # --- Initiate omega with min/max bounding --- #
            if self.aitkenCritMecha == 'max':
                self.omegaMecha = max(self.omegaBoundMecha, self.omegaMecha)
            else:
                self.omega = min(self.omegaBoundMecha, self.omegaMecha)

        self.omegaMecha = min(self.omegaMecha, 1.0)
        self.omegaMecha = max(self.omegaMecha, self.omegaMinMecha)

        mpiPrint('Aitken under-relaxation summary, mechanical : {}'.format(self.omegaMecha), self.mpiComm)

        # --- Update the value of the residual for the next FSI iteration --- #
        self.solidInterfaceResidual.copy(self.solidInterfaceResidualkM1)

    def setOmegaThermal(self):
        """
        Des.
        """

        if self.FSIIter != 0:
            # --- Compute the dynamic Aitken coefficient --- #
            if self.interfaceInterpolator.chtTransferMethod == 'hFFB' or self.interfaceInterpolator.chtTransferMethod == 'TFFB':
                deltaHeatFluxResidual = self.solidHeatFluxResidual - self.solidHeatFluxResidualkM1

                prodScalRes_X, prodScalRes_Y, prodScalRes_Z = deltaHeatFluxResidual.dot(self.solidHeatFluxResidualkM1)
                prodScalRes = prodScalRes_X + prodScalRes_Y + prodScalRes_Z

                deltaHeatFluxResidual_NormX, deltaHeatFluxResidual_NormY, deltaHeatFluxResidual_NormZ = deltaHeatFluxResidual.norm()
                deltaResNormSquare = deltaHeatFluxResidual_NormX**2 + deltaHeatFluxResidual_NormY**2 + deltaHeatFluxResidual_NormZ**2
            elif self.interfaceInterpolator.chtTransferMethod == 'hFTB' or self.interfaceInterpolator.chtTransferMethod == 'FFTB':
                deltaTemperatureResidual = self.solidTemperatureResidual - self.solidTemperatureResidualkM1

                tempDot = deltaTemperatureResidual.dot(self.solidTemperatureResidualkM1)
                prodScalRes = tempDot[0]

                tempNorm = deltaTemperatureResidual.norm()
                deltaTemperatureResidual_Norm = tempNorm[0]
                deltaResNormSquare = deltaTemperatureResidual_Norm**2

            if deltaResNormSquare != 0.:
                self.omegaThermal *= -prodScalRes/deltaResNormSquare
            else:
                self.omegaThermal = self.omegaMinThermal

        else:
            # --- Initiate omega with min/max bounding --- #
            if self.aitkenCritThermal == 'max':
                self.omegaThermal = max(self.omegaBoundThermal, self.omegaThermal)
            else:
                self.omegaThermal = min(self.omegaBoundThermal, self.omegaThermal)

        self.omegaThermal = min(self.omegaThermal, 1.0)
        self.omegaThermal = max(self.omegaThermal, self.omegaMinThermal)

        mpiPrint('Aitken under-relaxation summary, thermal : {}'.format(self.omegaThermal), self.mpiComm)

        # --- Update the value of the residual for the next FSI iteration --- #
        self.solidHeatFluxResidual.copy(self.solidInterfaceResidualkM1)
        self.solidTemperatureResidual.copy(solidTemperatureResidualkM1)

class AlgorithmIQN_ILS(AlgorithmBGSAitkenRelax):
    """
    Des.
    """

    def __init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, Criterion, nbFSIIterMax, deltaT, totTime, timeIterTreshold=-1, omegaBoundList= [1.0, 1.0], nbTimeToKeep=0, computeTangentMatrixBasedOnFirstIt = False, mpiComm=None):
        """
        Des.
        """

        AlgorithmBGSAitkenRelax.__init__(self, Manager, FluidSolver, SolidSolver, InterfaceInterpolator, Criterion, nbFSIIterMax, deltaT, totTime, timeIterTreshold, omegaBoundList, mpiComm)

        # --- Number of previous time steps used in the approximation of the tangent matrix --- #
        self.nbTimeToKeep = nbTimeToKeep

        # --- Option which allows to build the tangent matrix of a given time step using differences with respect to the first FSI iteration (delta_r_k = r_k+1 - r_0) instead of the previous iteration (delta_r_k = r_k+1 - r_k) --- #
        self.computeTangentMatrixBasedOnFirstIt = computeTangentMatrixBasedOnFirstIt

        # --- Option which determines the way the c coefficients are computed either using Degroote's QR decompoistion or simply using np.linalg.lstsq
        self.useQR = True
        self.tollQR = 1.0e-1 # Tolerance employed for the QR decomposition
        self.qrFilter = 'Haelterman' # Type of QR filtering employed. Possible choices are 'Degroote1', 'Degroote2', and 'Haelterman' (see 'qrSolve()' function below)

        self.maxNbOfItReached = False
        self.convergenceReachedInOneIt = False

        # --- Global V and W matrices for IQN-ILS algorithm, including information from previous time steps --- #
        self.V = []
        self.W = []

    def qrSolve(self, V, W, res):

        if self.qrFilter == 'Degroote1': # QR filtering as described by J. Degroote et al. Computers and Structures, 87, 793-801 (2009).
            Q, R = sp.linalg.qr(V, mode='economic')
            s = np.dot(np.transpose(Q), -res)
            toll = self.tollQR*sp.linalg.norm(R, 2)
            c = solve_upper_triangular_mod(R, s, toll)

        elif self.qrFilter == 'Degroote2': # QR filtering as described by J. Degroote et al. CMAME, 199, 2085-2098 (2010).
            Q, R, V, W = QRfiltering(V, W, self.tollQR)
            s = np.dot(np.transpose(Q), -res)
            c = np.linalg.solve(R, s)

        elif self.qrFilter == 'Haelterman': # 'Modified' QR filtering as described by R. Haelterman et al. Computers and Structures, 171, 9-17 (2016).
            Q, R, V, W = QRfiltering_mod(V, W, self.tollQR)
            s = np.dot(np.transpose(Q), -res)
            c = np.linalg.solve(R, s)

        else:
            raise NameError('IQN-ILS Algorithm: the QR filtering technique is unknown!')

        return c, W

    def fsiCoupling(self):
        """
        Interface Quasi Newton - Inverse Least Square (IQN-ILS) method for strong coupling FSI
        """

        if self.timeIter > self.timeIterTreshold:
            nbFSIIter = self.nbFSIIterMax
            mpiPrint('\n*************** Enter Interface Quasi Newton - Inverse Least Square (IQN-ILS) method for strong coupling FSI ***************', self.mpiComm)
        else:
             nbFSIIter = 1

        self.FSIIter = 0
        self.FSIConv = False
        self.errValue = 1.0
        self.errValue_CHT = 1e6 # Just for compatibility. CHT not yet implemented for the IQN-ILS algorithm.

        ns = self.interfaceInterpolator.getNs()

        # --- Initialize all the quantities used in the IQN-ILS method --- #
        res = FlexInterfaceData(ns, 3, self.mpiComm)
        solidInterfaceResidual0 = FlexInterfaceData(ns, 3, self.mpiComm)

        solidInterfaceDisplacement_tilde = FlexInterfaceData(ns, 3, self.mpiComm)
        solidInterfaceDisplacement_tilde1 = FlexInterfaceData(ns, 3, self.mpiComm)

        delta_ds = FlexInterfaceData(ns, 3, self.mpiComm)

        Vk_mat = np.zeros((self.manager.nDim*ns,1))
        Wk_mat = np.zeros((self.manager.nDim*ns,1))

        delta_ds_loc_X = np.zeros(0)
        delta_ds_loc_Y = np.zeros(0)
        delta_ds_loc_Z = np.zeros(0)

        if (self.nbTimeToKeep!=0 and self.timeIter > 1): # If information from previous time steps is re-used then Vk = V, Wk = W
            Vk = copy.deepcopy(self.V)
            Wk = copy.deepcopy(self.W)
        else: # If information from previous time steps is not re-used then Vk and Wk are empty lists of np.array()
            Vk = []
            Wk = []

        nIt = 0

        while ((self.FSIIter < nbFSIIter) and (not self.criterion.isVerified(self.errValue,self.errValue_CHT))):
            mpiPrint("\n>>>> FSI iteration {} <<<<\n".format(self.FSIIter), self.mpiComm)

            # --- Solid to fluid mechanical transfer --- #
            self.solidToFluidMechaTransfer()
            # --- Fluid mesh morphing --- #
            mpiPrint('\nPerforming mesh deformation...\n', self.mpiComm)
            self.meshDefTimer.start()
            self.FluidSolver.meshUpdate(self.timeIter)
            self.meshDefTimer.stop()
            self.meshDefTimer.cumul()

            # --- Fluid solver call for FSI subiteration --- #
            mpiPrint('\nLaunching fluid solver...', self.mpiComm)
            self.fluidSolverTimer.start()
            self.FluidSolver.run(self.time-self.deltaT, self.time)
            self.fluidSolverTimer.stop()
            self.fluidSolverTimer.cumul()
            mpiBarrier(self.mpiComm)

            if self.timeIter > self.timeIterTreshold:
                # --- Fluid to solid mechanical transfer --- #
                mpiPrint('\nProcessing interface fluid loads...\n', self.mpiComm)
                self.fluidToSolidMechaTransfer()
                mpiBarrier(self.mpiComm)

                # --- Solid solver call for FSI subiteration --- #
                mpiPrint('\nLaunching solid solver...\n', self.mpiComm)
                if self.myid in self.manager.getSolidSolverProcessors():
                    self.solidSolverTimer.start()
                    self.SolidSolver.run(self.time-self.deltaT, self.time)
                    self.solidSolverTimer.stop()
                    self.solidSolverTimer.cumul()

                # --- Compute and monitor the FSI residual --- #
                res = self.computeSolidInterfaceResidual()
                self.errValue = self.criterion.update(res)
                mpiPrint('\nFSI error value : {}\n'.format(self.errValue), self.mpiComm)
                self.FSIConv = self.criterion.isVerified(self.errValue)

                # --- Initialize d_tilde for the construction of the Wk matrix -- #
                if self.myid in self.manager.getSolidInterfaceProcessors():
                    localSolidInterfaceDisp_X, localSolidInterfaceDisp_Y, localSolidInterfaceDisp_Z = self.SolidSolver.getNodalDisplacements()
                    for iVertex in range(self.manager.getNumberOfLocalSolidInterfaceNodes()):
                        iGlobalVertex = self.manager.getGlobalIndex('solid', self.myid, iVertex)
                        solidInterfaceDisplacement_tilde[iGlobalVertex] = [localSolidInterfaceDisp_X[iVertex], localSolidInterfaceDisp_Y[iVertex], localSolidInterfaceDisp_Z[iVertex]]

                solidInterfaceDisplacement_tilde.assemble()

                if ((self.FSIIter == 0 and (self.nbTimeToKeep == 0 or (self.nbTimeToKeep != 0 and (self.maxNbOfItReached or self.convergenceReachedInOneIt or self.timeIter == 1)))) or self.timeIter < 1): # If information from previous time steps is re-used then this step is only performed at the first iteration of the first time step, otherwise it is performed at the first iteration of every time step
                    # --- Relax the solid position --- #
                    mpiPrint('\nProcessing interface displacements...\n', self.mpiComm)
                    self.relaxSolidPosition()
                else:
                    # --- Construct Vk and Wk matrices for the computation of the approximated tangent matrix --- #
                    mpiPrint('\nCorrect solid interface displacements using IQN-ILS method...\n', self.mpiComm)

                    # --- Start gathering on root process --- #
                    res_X_Gat, res_Y_Gat, res_Z_Gat = mpiGatherInterfaceData(res, ns, self.mpiComm, 0)
                    solidInterfaceResidual0_X_Gat, solidInterfaceResidual0_Y_Gat, solidInterfaceResidual0_Z_Gat = mpiGatherInterfaceData(solidInterfaceResidual0, ns, self.mpiComm, 0)
                    solidInterfaceDisplacement_tilde_X_Gat, solidInterfaceDisplacement_tilde_Y_Gat, solidInterfaceDisplacement_tilde_Z_Gat = mpiGatherInterfaceData(solidInterfaceDisplacement_tilde, ns, self.mpiComm, 0)
                    solidInterfaceDisplacement_tilde1_X_Gat, solidInterfaceDisplacement_tilde1_Y_Gat, solidInterfaceDisplacement_tilde1_Z_Gat = mpiGatherInterfaceData(solidInterfaceDisplacement_tilde1, ns, self.mpiComm, 0)

                    if self.myid == 0:
                        if self.FSIIter > 0: # Either information from previous time steps is re-used or not, Vk and Wk matrices are enriched only starting from the second iteration of every FSI loop
                            if self.manager.nDim == 3:
                                delta_res = np.concatenate([res_X_Gat - solidInterfaceResidual0_X_Gat, res_Y_Gat - solidInterfaceResidual0_Y_Gat, res_Z_Gat - solidInterfaceResidual0_Z_Gat], axis=0)
                                delta_d = np.concatenate([solidInterfaceDisplacement_tilde_X_Gat - solidInterfaceDisplacement_tilde1_X_Gat, solidInterfaceDisplacement_tilde_Y_Gat - solidInterfaceDisplacement_tilde1_Y_Gat, solidInterfaceDisplacement_tilde_Z_Gat - solidInterfaceDisplacement_tilde1_Z_Gat], axis = 0)
                            else:
                                delta_res = np.concatenate([res_X_Gat - solidInterfaceResidual0_X_Gat, res_Y_Gat - solidInterfaceResidual0_Y_Gat], axis=0)
                                delta_d = np.concatenate([solidInterfaceDisplacement_tilde_X_Gat - solidInterfaceDisplacement_tilde1_X_Gat, solidInterfaceDisplacement_tilde_Y_Gat - solidInterfaceDisplacement_tilde1_Y_Gat], axis = 0)

                            Vk.insert(0, delta_res)
                            Wk.insert(0, delta_d)

                            nIt+=1

                        Vk_mat = np.vstack(Vk).T
                        Wk_mat = np.vstack(Wk).T

                        if (Vk_mat.shape[1] > self.manager.nDim*ns and self.qrFilter == 'Degroote1'): # Remove extra columns if number of iterations (i.e. columns of Vk and Wk) is larger than number of interface degrees of freedom
                            mpiPrint('WARNING: IQN-ILS Algorithm using \'Degroote1\' QR filter. Approximated stiffness matrix number of columns exceeds the number of degrees of freedom at FSI interface. Extra columns (the oldest ones!) are deleted for next iterations to avoid overdetermined problem!', self.mpiComm)
                            Vk_mat = np.delete(Vk_mat, np.s_[(self.manager.nDim*ns-Vk_mat.shape[1]):], 1)
                            Wk_mat = np.delete(Wk_mat, np.s_[(self.manager.nDim*ns-Wk_mat.shape[1]):], 1)

                        dummy_V = Vk_mat.copy()
                        dummy_W = Wk_mat.copy()

                        if self.manager.nDim == 3:
                            dummy_Res = np.concatenate([res_X_Gat, res_Y_Gat, res_Z_Gat], axis=0)
                        else:
                            dummy_Res = np.concatenate([res_X_Gat, res_Y_Gat], axis=0)

                        if self.useQR: # Technique described by Degroote et al.
                            c, dummy_W = self.qrSolve(dummy_V, dummy_W, dummy_Res)
                        else:
                            c = np.linalg.lstsq(dummy_V, -dummy_Res)[0] # Classical QR decomposition: NOT RECOMMENDED!

                        if self.manager.nDim == 3:
                            delta_ds_loc = np.split((np.dot(dummy_W,c).T + np.concatenate([res_X_Gat, res_Y_Gat, res_Z_Gat], axis=0)),3,axis=0)

                            delta_ds_loc_X = delta_ds_loc[0]
                            delta_ds_loc_Y = delta_ds_loc[1]
                            delta_ds_loc_Z = delta_ds_loc[2]
                        else:
                            delta_ds_loc = np.split((np.dot(dummy_W,c).T + np.concatenate([res_X_Gat, res_Y_Gat], axis=0)),2,axis=0)

                            delta_ds_loc_X = delta_ds_loc[0]
                            delta_ds_loc_Y = delta_ds_loc[1]
                            delta_ds_loc_Z = np.zeros(ns)

                        for iVertex in range(delta_ds_loc_X.shape[0]):
                            iGlobalVertex = self.manager.getGlobalIndex('solid', self.myid, iVertex)
                            delta_ds[iGlobalVertex] = [delta_ds_loc_X[iVertex], delta_ds_loc_Y[iVertex], delta_ds_loc_Z[iVertex]]

                    # --- Go back to parallel run --- #
                    mpiBarrier(self.mpiComm)
                    delta_ds.assemble()
                    self.interfaceInterpolator.solidInterfaceDisplacement += delta_ds

                if self.computeTangentMatrixBasedOnFirstIt:
                    if self.FSIIter == 0:
                        res.copy(solidInterfaceResidual0)
                        solidInterfaceDisplacement_tilde.copy(solidInterfaceDisplacement_tilde1)
                else:
                    res.copy(solidInterfaceResidual0)
                    solidInterfaceDisplacement_tilde.copy(solidInterfaceDisplacement_tilde1)

            if self.writeInFSIloop == True:
                self.writeRealTimeData()

            self.FSIIter += 1

        # if comm.myself == rootProcess

        # update of the matrices V and W at the end of the while
        if self.nbTimeToKeep != 0 and self.timeIter >= 1:

            # --- Trick to avoid breaking down of the simulation in the rare cases when, in the initial time steps, FSI convergence is reached without iterating (e.g. starting from a steady condition and using very small time steps), leading to empty V and W matrices ---
            if not (self.FSIIter == 1 and self.FSIConv and len(self.V)==0):

                self.convergenceReachedInOneIt = False

                # --- Managing situations where FSI convergence is not reached ---
                if (self.FSIIter >= nbFSIIter and not self.FSIConv):
                    mpiPrint('WARNING: IQN-ILS using information from {} previous time steps reached max number of iterations. Next time step is run without using any information from previous time steps!'.format(self.nbTimeToKeep), self.mpiComm)

                    self.maxNbOfItReached = True
                    self.V = []
                    self.W = []
                else:
                    self.maxNbOfItReached = False

                    mpiPrint('\nUpdating V and W matrices...\n', self.mpiComm)

                    self.V.insert(0, Vk_mat[:,0:nIt].T)
                    self.W.insert(0, Wk_mat[:,0:nIt].T)

                    if (self.timeIter > self.nbTimeToKeep and len(self.V) > self.nbTimeToKeep):
                        del self.V[-1]
                        del self.W[-1]
                # ---
            else:
                mpiPrint('\nWARNING: IQN-ILS algorithm convergence reached in one iteration at the beginning of the simulation. V and W matrices cannot be built. BGS will be employed for the next time step!\n', self.mpiComm)
                self.convergenceReachedInOneIt = True
            # ---

        # --- Update the FSI history file --- #
        if self.timeIter > self.timeIterTreshold:
            mpiPrint('\n*************** IQN-ILS is converged ***************', self.mpiComm)

# --- Solid test algorithm ---
class FsiSolidTestAlgorithm:
    def __init__(self, _solid):
        self.solid = _solid

    def run(self):
        # --------------------------
        # fake FSI solver
        # --------------------------

        t1 = 0.0  # initial time
        dt = 0.5  # time step size
        nt = 10

        # we want nt time steps
        for j in range(nt):

            # each time step is arbitrarily calculated twice (for testing purpose)
            for i in range(2):

                t2=t1+dt  # time to be calculated

                self.solid.fakeFluidSolver(t2)  # creates some dummy loads for time=t2

                # run solid solver
                print '='*80
                print "running from %f to %f: try #%d" % (t1,t2,i+1)
                print '='*80
                self.solid.run(t1,t2)

                # gets the deformed interface
                dx, dy, dz = self.solid.getNodalDisplacements()
                print dx
                print dy
                print dz

            self.solid.update()
            self.solid.save()

            t1=t2 # fsi loop has converged - time t2 is accepted

        # end.


# coding=utf-8

"""

Tensorflow implementation of Doc2VecC algorithm wrapper class

:author: Fido Wang (wanghailin317@gmail.com)
:refer: https://openreview.net/pdf?id=B1Igu2ogg

"""

from __future__ import print_function
import tensorflow as tf
import numpy as np
from option import Option
import math
import time
import collections
from itertools import compress
import random

MAX_SENTENCE_SAMPLE = 100

def generate_batch_doc2VecC_tail(doc_ids, word_ids, doc_len, batch_size, window_size, sample_size):
    """
    batch generator for PV-DM (Distributed Memory Model of Paragraph Vectors)
    :param doc_ids: list of document indices
    :param word_ids: list of word indices
    :param doc_len: record accumulated length of each doc
    :param batch_size: number of words in each mini-batch
    :param window_size: number of words before the target word
    :return: list of tuple of (batch, labels, batch_doc_sample, num_sampled)
    """
    data_index = 0
    assert batch_size % window_size == 0
    span = window_size + 1
    buffer = collections.deque(maxlen=span)
    buffer_doc = collections.deque(maxlen=span)
    batches = np.ndarray(shape=(batch_size, window_size + 1), dtype=np.int32)
    labels = np.ndarray(shape=(batch_size, 1), dtype=np.int32)
    batch_doc = np.ndarray(shape=(batch_size, sample_size), dtype=np.int32)
    mask = [1] * span
    mask[-1] = 0
    i = 0

    while data_index < len(word_ids):
        if len(set(buffer_doc)) == 1 and len(buffer_doc) == span:
            doc_id = buffer_doc[-1]
            batches[i, :] = list(compress(buffer, mask)) + [doc_id]
            labels[i, 0] = buffer[-1]
            batch_doc[i, :] = random.sample(word_ids[doc_len[doc_id]:doc_len[doc_id + 1]],
                                       sample_size)
            i += 1
        buffer.append(word_ids[data_index])
        buffer_doc.append(doc_ids[data_index])
        data_index = (data_index + 1) % len(word_ids)
        if i == batch_size:
            yield batches, labels, batch_doc

class Doc2VecC(object):
    """
    Doc2VecC embedding class
    """

    def __init__(self, options):
        assert (isinstance(options, Option))
        self._options = options
        self._session = None
        self.saver = None
        self._cost = None
        self._optimizer = None
        self._word_embeddings = None
        self._para_embeddings = None
        self.vocab = None
        self.vocab_size = 0
        self.document_size = 0
        self.__inputs, self.__labels, self.__lr = None, None, None
        self.__cost = None
        self.__optimizer = None
        self.__summary = None
        self.__normalized_word_embeddings = None

    def setVocab(self, vocab):
        self.vocab = vocab
        self.vocab_size = len(vocab)
        return self

    def setDocSize(self, doc_size):
        assert (isinstance(doc_size, int))
        self.document_size = doc_size
        return self

    def useSubSampling(self, switch=True, threshold=1e-5):
        self.use_sub_sampling = switch
        self.sub_sampling_threshold = 1e-5
        return self

    def _get_batches(self, doc_ids, word_ids):
        opts = self._options
        return generate_batch_doc2VecC_tail(doc_ids, word_ids, doc_ids, opts.batch_size, opts.window_size, opts.sentence_sample)

    def _get_inputs(self):
        """
        Create TF Placeholders for input, targets, and learning rate.
         :return: Tuple of Placeholders (input, targets, learning rate)
        """
        opts = self._options
        inputs_ = tf.placeholder(tf.int32, [None, opts.window_size], name='input')
        doc_inputs_ = tf.placeholder(tf.int32, [None, None], name='doc_input')
        labels_ = tf.placeholder(tf.int32, [None, 1], name='label')
        lr_ = tf.placeholder(tf.float32, name='learning_rate')
        return inputs_, doc_inputs_, labels_, lr_

    def _get_embedding_layer(self, input_data, doc_input_data):
        """
        Create embedding for <input_data> and <doc_input_data>.
        :param input_data: TF placeholder for text input.
        :return: Embedded input tensor.
        """
        opts = self._options
        word_embedding = tf.Variable(tf.random_uniform((self.vocab_size, opts.embed_dim), -1.0, 1.0))
        embed = []

        temp = tf.zeros([opts.batch_size, opts.embed_dim])
        embed_d = []
        for n in range(opts.sentence_sample):
            temp = tf.add(temp, tf.nn.embedding_lookup(word_embedding, doc_input_data[:, n]))
        embed_d.append(temp)

        if opts.concat == 'True':
            combined_embed_vector_length = opts.embed_dim * opts.window_size + opts.embed_dim
            for j in range(opts.window_size):
                embed_w = tf.nn.embedding_lookup(word_embedding, input_data[:, j])
                embed.append(embed_w)
            embed.append(embed_d)
        else:
            combined_embed_vector_length = opts.embed_dim
            embed_w = tf.zeros([opts.batch_size, opts.embed_dim])
            for j in range(opts.window_size):
                embed_w += tf.nn.embedding_lookup(word_embedding, input_data[:, j])
            embed_w += embed_d
            embed.append(embed_w)

        return tf.concat(embed, 1), word_embedding, combined_embed_vector_length

    def build_graph(self):
        """
        Create Graph and Initialize tf Session for training
        """
        train_graph = tf.Graph()
        opts = self._options
        with train_graph.as_default():
            self.__inputs, self.__doc_inputs, self.__labels,  self.__lr = self._get_inputs()
            embed, word_embeddings, combined_embed_vector_length = self._get_embedding_layer(
                self.__inputs, self.__doc_inputs)

            norm_w = tf.sqrt(tf.reduce_sum(tf.square(word_embeddings), 1, keep_dims=True))
            self.__normalized_word_embeddings = word_embeddings / norm_w

            weights = tf.Variable(
                tf.truncated_normal((self.vocab_size, combined_embed_vector_length),
                                    stddev=1.0 / math.sqrt(combined_embed_vector_length))
            )
            biases = tf.Variable(tf.zeros(self.vocab_size))

            if opts.loss == 'softmax':
                loss = tf.nn.sampled_softmax_loss(weights=weights,
                                                  biases=biases,
                                                  labels=self.__labels,
                                                  inputs=embed,
                                                  num_sampled=opts.negative_sample_size,
                                                  num_classes=opts.vocab_size)
                tf.summary.scalar("Softmax loss", loss)
            else:
                loss = tf.nn.nce_loss(weights=weights,
                                      biases=biases,
                                      labels=self.__labels,
                                      inputs=embed,
                                      num_sampled=opts.negative_sample_size,
                                      num_classes=opts.vocab_size)
                tf.summary.scalar("NCE loss", loss)

            self.__cost = tf.reduce_mean(loss)

            if opts.train_method == 'Adam':
                self.__optimizer = tf.train.AdamOptimizer(self.__lr).minimize(self.__cost)
            else:
                self.__optimizer = tf.train.GradientDescentOptimizer(self.__lr).minimize(self.__cost)

            self.__summary = tf.summary.merge_all()

            self._session = tf.Session(graph=train_graph)
            self.saver = tf.train.Saver()
        return self

    def fit(self, docs):
        opts = self._options
        iteration = 1
        loss = 0
        doc_ids = [[i] * len(j) for i, j in enumerate(docs)]
        doc_ids = [item for sublist in doc_ids for item in sublist]
        doc_lens = [0] + [len(i) for i in docs]
        for i in range(1, len(doc_lens)):
            doc_lens[i] += doc_lens[i-1]
        word_ids = [item for sublist in docs for item in sublist]

        with self._session as session:
            session.run(tf.global_variables_initializer())
            for e in range(1, opts.epochs_to_train + 11):
                batches = self._get_batches(doc_ids, word_ids, doc_lens)
                start = time.time()
                lr = opts.learning_rate if e <= opts.epochs_to_train else opts.learning_rate * (
                    e - opts.epochs_to_train / 10)
                for x, y, m, l in batches:
                    opts.doc_batch_len = l
                    feed = {self.__inputs: x,
                            self.__labels: y,
                            self.__doc_inputs: m,
                            self.__lr: lr}
                    train_loss, _ = session.run([self.__cost, self.__optimizer], feed_dict=feed)

                    loss += train_loss
                    if iteration % opts.statistics_interval == 0:
                        end = time.time()
                        print("Epoch {}/{}".format(e, opts.epochs_to_train + 11),
                              "Iteration: {}".format(iteration),
                              "Avg. Training loss: {:.4f}".format(loss * 1.0 / opts.statistics_interval),
                              "{:.4f} sec/batch".format((end - start) * 1.0 / opts.statistics_interval))
                        loss = 0
                        start = time.time()
                    if iteration % opts.checkpoint_interval == 0:
                        self.saver.save(self._session,
                                        "doc2vecc",
                                        global_step=iteration)
                    iteration += 1
            self._word_embeddings = self.__normalized_word_embeddings.eval()
            self.saver(self._session, "final_doc2vecc")

    def transform_w(self, word_index):
        return self._word_embeddings[word_index, :]

    def transform_doc(self, word_indexs):
        doc_embeddings = [self._word_embeddings[i, :] for i in word_indexs]
        return doc_embeddings


import numpy as np
import scipy
import pandas as pd
import anndata
import scanpy as sc
import scnym.utils as utils
import scnym.model as model
import os
import os.path as osp
import _pickle as pickle
import random
from scnym import main
from anndata import AnnData
from metalearner_parser import make_parser
UNLABELED_TOKEN = 'Unlabeled'

def scnym_train(
    adata: AnnData,
    config: dict,
) -> None:
    '''Train an scNym model.

    Parameters
    ----------
    adata : AnnData
        [Cells, Genes] experiment containing annotated
        cells to train on.
    config : dict
        configuration options.

    Returns
    -------
    None.
    Saves model outputs to `config["out_path"]` and adds model results
    to `adata.uns["scnym_train_results"]`.

    Notes
    -----
    This method should only be directly called by advanced users.
    Most users should use `scnym_api`.

    See Also
    --------
    scnym_api
    '''
    # determine if unlabeled examples are present
    n_unlabeled = np.sum(
        adata.obs[config['groupby']] == UNLABELED_TOKEN
    )
    if n_unlabeled == 0:
        print('No unlabeled data was found.')
        print(f'Did you forget to set some examples as `"{UNLABELED_TOKEN}"`?')
        print('Proceeding with purely supervised training.')
        print()

        unlabeled_counts = None
        unlabeled_genes  = None

        X = utils.get_adata_asarray(adata)
        y = pd.Categorical(
            np.array(adata.obs[config['groupby']]),
            categories=np.unique(adata.obs[config['groupby']]),
        ).codes
        class_names = np.unique(adata.obs[config['groupby']])
        # set all samples for training
        train_adata = adata
    else:
        print(f'{n_unlabeled} unlabeled observations found.')
        print('Using unlabeled data as a target set for semi-supervised, adversarial training.')
        print()

        target_bidx = adata.obs[config['groupby']] == UNLABELED_TOKEN

        train_adata = adata[~target_bidx, :]
        target_adata = adata[target_bidx, :]

        print('training examples: ', train_adata.shape)
        print('target   examples: ', target_adata.shape)

        X = utils.get_adata_asarray(train_adata)
        y = pd.Categorical(
            np.array(train_adata.obs[config['groupby']]),
            categories=np.unique(train_adata.obs[config['groupby']]),
        ).codes
        unlabeled_counts = utils.get_adata_asarray(target_adata)
        class_names = np.unique(train_adata.obs[config['groupby']])

    print('X: ', X.shape)
    print('y: ', y.shape)

    if 'scNym_split' not in adata.obs_keys():
        # perform a 90/10 train test split
        traintest_idx = np.random.choice(
            X.shape[0],
            size=int(np.floor(0.9*X.shape[0])),
            replace=False
        )
        val_idx = np.setdiff1d(np.arange(X.shape[0]), traintest_idx)
    else:
        train_idx = np.where(
            train_adata.obs['scNym_split'] == 'train'
        )[0]
        test_idx = np.where(
            train_adata.obs['scNym_split'] == 'test',
        )[0]
        val_idx = np.where(
            train_adata.obs['scNym_split'] == 'val'
        )[0]

        if len(train_idx) < 100 or len(test_idx) < 10 or len(val_idx) < 10:
            msg = 'Few samples in user provided data split.\n'
            msg += f'{len(train_idx)} training samples.\n'
            msg += f'{len(test_idx)} testing samples.\n'
            msg += f'{len(val_idx)} validation samples.\n'
            msg += 'Halting.'
            raise RuntimeError(msg)
        # `fit_model()` takes a tuple of `traintest_idx`
        # as a training index and testing index pair.
        traintest_idx = (
            train_idx,
            test_idx,
        )

    # check if domain labels were manually specified
    if config.get('domain_groupby', None) is not None:
        domain_groupby = config['domain_groupby']
        # check that the column actually exists
        if domain_groupby not in adata.obs.columns:
            msg = f'no column {domain_groupby} exists in `adata.obs`.\n'
            msg += 'if domain labels are specified, a matching column must exist.'
            raise ValueError(msg)
        # get the label indices as unique integers using pd.Categorical
        # to code each unique label with an int
        domains = np.array(
            pd.Categorical(
                adata.obs[domain_groupby],
                categories=np.unique(adata.obs[domain_groupby]),
            ).codes,
            dtype=np.int32,
        )
        # split domain labels into source and target sets for `fit_model`
        input_domain = domains[~target_bidx]
        unlabeled_domain = domains[target_bidx]
        print('Using user provided domain labels.')
        n_source_doms = len(np.unique(input_domain))
        n_target_doms = len(np.unique(unlabeled_domain))
        print(
            f'Found {n_source_doms} source domains and {n_target_doms} target domains.'
        )
    else:
        # no domains manually supplied, providing `None` to `fit_model`
        # will treat source data as one domain and target data as another
        input_domain = None
        unlabeled_domain = None

    # check if pre-trained weights should be used to initialize the model
    if config['trained_model'] is None:
        pretrained = None
    elif 'pretrained_' in config['trained_model']:
        msg = 'pretrained model fetching is not supported for training.'
        raise NotImplementedError(msg)
    else:
        # setup a prediction model
        pretrained = osp.join(
            config['trained_model'],
            '00_best_model_weights.pkl',
        )
        if not osp.exists(pretrained):
            msg = f'{pretrained} file not found.'
            raise FileNotFoundError(msg)

    acc, loss = main.fit_model(
        X=X,
        y=y,
        traintest_idx=traintest_idx,
        val_idx=val_idx,
        batch_size=config['batch_size'],
        n_epochs=config['n_epochs'],
        lr=config['lr'],
        optimizer_name=config['optimizer_name'],
        weight_decay=config['weight_decay'],
        ModelClass=model.CellTypeCLF,
        out_path=config['out_path'],
        mixup_alpha=config['mixup_alpha'],
        unlabeled_counts=unlabeled_counts,
        input_domain=input_domain,
        unlabeled_domain=unlabeled_domain,
        unsup_max_weight=config['unsup_max_weight'],
        unsup_mean_teacher=config['unsup_mean_teacher'],
        ssl_method=config['ssl_method'],
        ssl_kwargs=config['ssl_kwargs'],
        pretrained=pretrained,
        patience=config.get('patience', None),
        save_freq=config.get('save_freq', None),
        tensorboard=config.get('tensorboard', False),
        weighted_classes=config['weighted_classes'],
        balanced_classes=config['balanced_classes'],
        verbose=config['verbose'],
        in_metalearner=config['in_metalearner'],
        **config['model_kwargs'],
    )

    # add the final model results to `adata`
    results = {
        'model_path': osp.realpath(osp.join(config['out_path'], '00_best_model_weights.pkl')),
        'final_acc': acc,
        'final_loss': loss,
        'n_genes': adata.shape[1],
        'n_cell_types': len(np.unique(y)),
        'class_names': class_names,
        'gene_names': adata.var_names.tolist(),
        'model_kwargs': config['model_kwargs'],
        'traintest_idx': traintest_idx,
        'val_idx': val_idx,
    }
    assert osp.exists(results['model_path'])

    adata.uns['scNym_train_results'] = results

    # save the final model results to disk
    train_results_path = osp.join(
        config['out_path'],
        'scnym_train_results.pkl',
    )

    with open(train_results_path, 'wb') as f:
        pickle.dump(
            results,
            f
        )
    return

if __name__ == "__main__":
    parser = make_parser()
    configs = parser.parse_args()

    # path = "/home/svcapp/tbrain_x/azimuth/pbmc_multimodal_raw.h5ad"
    # unlabeled_path = "/home/svcapp/tbrain_x/SKT_data_corrected/IO_response/geninus_raw_final.h5ad"
    # path = "filtered_azimuth_and_geninus.h5ad"

#     configs = {
#     'out_path': "ssl_balanced/",
#     'groupby': "celltype.l2",
#     'n_epochs': 1000,
#     'patience': 40,
#     'lr': 0.1,
#     'optimizer_name': 'adadelta',
#     'weight_decay': 1e-4,
#     'batch_size': 512,
#     'mixup_alpha': 0.3,
#     'unsup_max_weight': 0.5,
#     'unsup_mean_teacher': False,
#     'weighted_classes': False,
#     'balanced_classes': True,
#     'verbose': 0,
#     'ssl_method': 'mixmatch',
#     'ssl_kwargs': {
#         'augment_pseudolabels': False,
#         'augment': 'log1p_drop',
#         'unsup_criterion': 'mse',
#         'n_augmentations': 1,
#         'T': 0.5,
#         'ramp_epochs': 100,
#         'burn_in_epochs': 0,
#         'dan_criterion': True,
#         'dan_ramp_epochs': 20,
#         'dan_max_weight': 0.1,
#         'min_epochs': 20,
#         'pseudolabel_min_confidence':0.5
#     },
#     'model_kwargs' : {
#         'n_hidden': 256,
#         'n_layers': 2,
#         'init_dropout': 0.0,
#         'residual': False,
#     },
#     'tensorboard': True,
# }

    # unlabeled_adata = anndata.read_h5ad(unlabeled_path)
    # unlabeled_adata.X = scipy.sparse.csr_matrix(unlabeled_adata.X, dtype=np.float32)
    # sc.pp.normalize_per_cell(unlabeled_adata, counts_per_cell_after=2e3)

    adata = anndata.read_h5ad(configs.h5ad_path)
    # adata = anndata.read_h5ad(path)
    # sc.pp.normalize_per_cell(adata, counts_per_cell_after=2e3)
    # sc.pp.filter_genes(adata, min_cells=20000)
    # sc.pp.filter_cells(adata, min_genes=100)

    # num_filtered = 161753

    # sample_idx = random.sample(range(num_filtered, adata.shape[0]), (adata.shape[0]-num_filtered)//4)
    # adata = anndata.concat( [adata[:num_filtered], adata[sample_idx]])
    # adata.obs["celltype.l2"][sample_idx] = "Unlabeled"

    # adata.obs["domain"] = "azimuth"
    # gene_names = adata.var_names.tolist()
    # gene_idx = np.zeros(len(gene_names), dtype=np.int32)
    # for i, gene in enumerate(gene_names):
    #     gene_idx[i] = np.where(unlabeled_adata.var.index == gene_names[i])[0]
    # unlabeled_adata = unlabeled_adata[:, gene_idx]
    # unlabeled_adata.obs["celltype.l2"] = "Unlabeled"
    # unlabeled_adata.obs["domain"] = "geninus"

    # sc.pp.filter_cells(unlabeled_adata, min_genes=100)

    # num_unlabel = unlabeled_adata.shape[0]
    # end_idx = num_unlabel // 5

    # for i in range(4):
    #     print(i)
    #     adata = anndata.concat([adata, unlabeled_adata[:end_idx]])
    #     print(f"{i}th concat finished")
    #     unlabeled_adata = unlabeled_adata[end_idx:]
    # print("try concat")
    # adata = anndata.concat([adata, unlabeled_adata])
    # print("concat end")
    # del unlabeled_adata
    # print("del end")

    # configs['trained_model'] = None
    # configs['domain_groupby'] = None
    # configs['trained_model'] = "sl_only/"
    # configs['domain_groupby'] = "domain"
    os.makedirs(configs.out_path, exist_ok=True)

    print("train start")
    config_dict = vars(configs)
    config_dict['ssl_kwargs'] = {}
    config_dict['model_kwargs'] = {}
    keys = list(config_dict.keys())
    for key in keys:
        if 'ssl_kwargs/' in key:
            config_dict['ssl_kwargs'][key.split('/')[1]] = config_dict.pop(key)
        if 'model_kwargs/' in key:
            config_dict['model_kwargs'][key.split('/')[1]] = config_dict.pop(key)
    class_balance_solution = config_dict.pop('class_balance_solution')
    config_dict['balanced_classes'] = class_balance_solution == "balanced"
    config_dict['weighted_classes'] = class_balance_solution == "weighted"

    scnym_train(adata, config_dict)


# -*- coding: utf-8 -*-
import numpy as np
import pickle
from .log_util import LogInfo

class WordEmbeddingUtil(object):

    def __init__(self, emb_dir, dim_emb):
        self.word_dict_fp = '%s/word_emb.indices' % emb_dir
        self.word_emb_mat_fp = '%s/word_emb.glove_%d.npy' % (emb_dir, dim_emb)
        self.mid_dict_fp = '%s/mid_emb.indices' % emb_dir
        self.mid_emb_mat_fp = '%s/mid_emb.glove_%d.npy' % (emb_dir, dim_emb)
        self.dep_name_fp = '%s/dep_names.txt' % emb_dir

        self.dim_emb = dim_emb
        self.word_idx_dict = None
        self.word_emb_matrix = None
        self.n_words = None
        self.mid_idx_dict = None
        self.mid_emb_matrix = None
        self.n_mids = None
        self.dep_name_dict = {}

        self.load_word_indices()
        self.load_mid_indices()
        self.load_dep_names()

    def load_word_indices(self):
        if self.word_idx_dict is None:
            LogInfo.logs('Loading <word, idx> pairs from [%s] ... ', self.word_dict_fp)
            with open(self.word_dict_fp, 'rb') as br:
                self.word_idx_dict = pickle.load(br)
            LogInfo.logs('%d <word, idx> loaded.', len(self.word_idx_dict))
            self.n_words = len(self.word_idx_dict)

    def load_mid_indices(self):
        if self.mid_idx_dict is None:
            LogInfo.logs('Loading <mid, idx> pairs from [%s] ... ', self.mid_dict_fp)
            with open(self.mid_dict_fp, 'rb') as br:
                self.mid_idx_dict = pickle.load(br)
            LogInfo.logs('%d <mid, idx> loaded.', len(self.mid_idx_dict))
            self.n_mids = len(self.mid_idx_dict)

    def load_word_embeddings(self):
        if self.word_emb_matrix is None:
            LogInfo.logs('Loading word embeddings for [%s] ...', self.word_emb_mat_fp)
            self.word_emb_matrix = np.load(self.word_emb_mat_fp)
            LogInfo.logs('%s word embedding loaded.', self.word_emb_matrix.shape)
            assert self.word_emb_matrix.shape == (self.n_words, self.dim_emb)

    def load_mid_embeddings(self):
        if self.mid_emb_matrix is None:
            LogInfo.logs('Loading mid embeddings for [%s] ...', self.mid_emb_mat_fp)
            self.mid_emb_matrix = np.load(self.mid_emb_mat_fp)
            LogInfo.logs('%s mid embedding loaded.', self.mid_emb_matrix.shape)
            assert self.mid_emb_matrix.shape == (self.n_mids, self.dim_emb)

    def load_dep_names(self):
        with open(self.dep_name_fp, 'r') as br:
            for line in br.readlines():
                dep, name = line.strip().split('\t')
                self.dep_name_dict[dep] = name
        LogInfo.logs('%d dependency name loaded.', len(self.dep_name_dict))

    def get_phrase_emb(self, phrase):
        """
        Calculate the embedding of a new phrase, by averaging the embeddings of all observed words
        """
        self.load_word_embeddings()
        if phrase == '':
            return None
        spt = phrase.split(' ')
        idx_list = []
        for wd in spt:
            if wd in self.word_idx_dict:
                idx_list.append(self.word_idx_dict[wd])
        if len(idx_list) == 0:
            return None
        emb = np.mean(self.word_emb_matrix[idx_list], axis=0)  # (n_words, dim_emb) ==> (dim_emb, )
        emb = emb / np.linalg.norm(emb)  # normalization
        return emb

    def produce_active_word_embedding(self, active_word_dict, dep_simulate=False):
        active_size = len(active_word_dict)
        word_emb_matrix = np.random.uniform(low=-0.1, high=0.1,
                                            size=(active_size, self.dim_emb)).astype('float32')
        self.load_word_indices()
        self.load_word_embeddings()
        for tok, active_idx in active_word_dict.items():
            if tok in self.word_idx_dict:
                local_idx = self.word_idx_dict[tok]
                word_emb_matrix[active_idx] = self.word_emb_matrix[local_idx]
            elif dep_simulate:       # try to simulate the embedding of dependency label
                dep_tok = tok.replace('!', '')
                if dep_tok in self.dep_name_dict:       # this is a known dependency label
                    dep_name_spt = self.dep_name_dict[dep_tok].split(' ')
                    filt_dep_name_spt = list(filter(lambda x: x in self.word_idx_dict, dep_name_spt))
                    if len(filt_dep_name_spt) == 0:
                        continue
                    # Now simulate dependency embedding by its name
                    simu_emb = np.zeros((self.dim_emb,), dtype='float32')
                    for term in filt_dep_name_spt:
                        term_idx = self.word_idx_dict[term]
                        simu_emb += self.word_emb_matrix[term_idx]
                    simu_emb /= len(filt_dep_name_spt)
                    if not tok.startswith('!'):
                        word_emb_matrix[active_idx] = simu_emb
                    else:       # reversed label
                        word_emb_matrix[active_idx] = -1. * simu_emb

        LogInfo.logs('dependency simulate = %s.', dep_simulate)
        LogInfo.logs('%s active word embedding created.', word_emb_matrix.shape)
        return word_emb_matrix

    def produce_active_mid_embedding(self, active_mid_dict):
        active_size = len(active_mid_dict)
        mid_emb_matrix = np.random.uniform(low=-0.1, high=0.1,
                                           size=(active_size, self.dim_emb)).astype('float32')
        self.load_mid_indices()
        self.load_mid_embeddings()
        for tok, active_idx in active_mid_dict.items():
            if tok in self.mid_idx_dict:
                local_idx = self.mid_idx_dict[tok]
                mid_emb_matrix[active_idx] = self.mid_emb_matrix[local_idx]
        LogInfo.logs('%s active mid embedding created.', mid_emb_matrix.shape)
        return mid_emb_matrix


import logging
import os

import cv2
import numpy as np
import wget

import tensorflow as tf

#logging
logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
log = logging.getLogger("fer-capture-log")

#global values
expression_dict = {0: "Angry", 1: "Disgust", 2: "Fear", 3: "Happy", 4: "Sad", 5: "Surprise", 6: "Neutral"}
face_casc = os.path.join(os.path.dirname(__file__), "haarcascade_frontalface_default.xml")

default_model_path = wget.download("https://storage.googleapis.com/id-public-read/model.h5")

def path_to_img(image_path):
    """
    This function takes a filepath-like object and returns an OpenCV image array.
    """
    return cv2.imread(image_path)

def detect_faces(frame):
    """
    This function detects faces in the given frame.

    Parameters:
        frame (array): OpenCV image array.

    Returns:
        gray (array): Grasyscale version of input frame
        faces (list); List of detected faces as list of tuples.
    """
    face_cascade = cv2.CascadeClassifier(face_casc)
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    log.info("Beginning facial detection.")
    faces = face_cascade.detectMultiScale(
        frame,
        scaleFactor=1.2,
        minNeighbors=5,
        minSize=(48, 48),
        flags = cv2.CASCADE_SCALE_IMAGE
    )
    log.info("Detected {} faces.".format(len(faces)))
    return gray, faces

def face_check(img, model, show=False):
    """
    This function performs an FER routine on the given input using the loaded model.

    Parameters:
        img   (array): OpenCV image array.
        model (model): Keras model.

    Returns:
        data (dict): Dictionary containing detected faces along with the predicted expressions.
    """
    #begin analysis
    frame = img
    gray, faces = detect_faces(frame)
    data = {"faces" : []}
    log.info("Beginning expression recognition.")

    for i, (x, y, w, h) in enumerate(faces):
        roi_gray = gray[y:y + h, x:x + w]
        cropped_img = np.expand_dims(np.expand_dims(cv2.resize(roi_gray, (48, 48)), -1), 0)
        cv2.normalize(cropped_img, cropped_img, alpha=0, beta=1, norm_type=cv2.NORM_L2, dtype=cv2.CV_32F)
        prediction = model.predict(cropped_img)[0]
        predictions = {
            "angry": float(prediction[0]),
            "disgusted": float(prediction[1]),
            "fearful": float(prediction[2]),
            "happy": float(prediction[3]),
            "sad": float(prediction[4]),
            "surprised": float(prediction[5]),
            "neutral": float(prediction[6])
        }
        data["faces"].append({
            "id" : i,
            "xywh": (x, y, w, h),
            "predictions": predictions
        })
        if show:
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 4)
            cv2.putText(
                frame,
                "id: {}  |  {}: {}".format(
                    i,
                    max(predictions, key=predictions.get),
                    round(
                        predictions[
                            max(predictions, key=predictions.get)
                        ],
                        3
                    )
                ),
                (x, y-10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (36,255,12),
                2
            )
    log.info("Detection completed successfully!")
    if show:
        cv2.imshow('Facial Detection', frame)
        cv2.waitKey(0)
    return data

def check_image(image_path, model_path=default_model_path, show=False):
    """
    This function handles the FER routine for individual images.

    Parameters:
        model_path (filepath): /path/to/model.h5
        image_path (filepath): /path/to/image.jpeg
        show (bool): Whether or not to display the media being analyzed.

    Returns:
        data (dict): Dictionary containing detected faces along with the predicted expressions.
    """
    img = path_to_img(image_path)
    model = tf.keras.models.load_model(model_path)
    return face_check(img, model, show)

def check_stream(input=0, model_path=default_model_path, show=False):
    """
    This function handles the FER routine for videos/streams.

    Parameters:
        model_path (filepath): /path/to/model.h5
        input (filepath or int):
            when int: stream from webcam device.
            when filepath: stream from videofile.
        show (bool): Whether or not to display the media being analyzed.
    Returns:
        data (list/dict): List of dictionaries containing detected faces along with the predicted expressions from each frame.
    """
    cap = cv2.VideoCapture(input)
    model = tf.keras.models.load_model(model_path)
    data = []
    while True:
        _, img = cap.read()
        if type(img) is not "NoneType":
            try:
                data.append(face_check(img, model, show))
            except Exception as e:
                log.error(e)
                break
        else:
            break
    cap.release()
    return data