"""
:mod:`mirgecom.eos` provides implementations of gas equations of state.

Equations of State
^^^^^^^^^^^^^^^^^^
This module is designed provide Equation of State objects used to compute and
manage the relationships between and among state and thermodynamic variables.

.. autoclass:: EOSDependentVars
.. autoclass:: GasEOS
.. autoclass:: MixtureEOS
.. autoclass:: IdealSingleGas
.. autoclass:: PyrometheusMixture
"""

__copyright__ = """
Copyright (C) 2021 University of Illinois Board of Trustees
"""

__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

from dataclasses import dataclass
import numpy as np
from pytools import memoize_in
from meshmode.mesh import BTAG_ALL, BTAG_NONE  # noqa
from mirgecom.fluid import ConservedVars, make_conserved
from abc import ABCMeta, abstractmethod

@dataclass
class EOSDependentVars:
    """State-dependent quantities for :class:`GasEOS`.

    Prefer individual methods for model use, use this
    structure for visualization or probing.

    .. attribute:: temperature
    .. attribute:: pressure
    """

    temperature: np.ndarray
    pressure: np.ndarray

class GasEOS(metaclass=ABCMeta):
    r"""Abstract interface to equation of state class.

    Equation of state (EOS) classes are responsible for
    computing relations between fluid or gas state variables.

    Each interface call takes an :class:`~mirgecom.fluid.ConservedVars` object
    array representing the simulation state quantities. Each EOS class
    implementation should document its own state data requirements.

    .. automethod:: pressure
    .. automethod:: temperature
    .. automethod:: sound_speed
    .. automethod:: internal_energy
    .. automethod:: gas_const
    .. automethod:: dependent_vars
    .. automethod:: total_energy
    .. automethod:: kinetic_energy
    .. automethod:: gamma
    .. automethod:: transport_model
    .. automethod:: get_internal_energy
    """

    @abstractmethod
    def pressure(self, cv: ConservedVars):
        """Get the gas pressure."""

    @abstractmethod
    def temperature(self, cv: ConservedVars):
        """Get the gas temperature."""

    @abstractmethod
    def sound_speed(self, cv: ConservedVars):
        """Get the gas sound speed."""

    @abstractmethod
    def gas_const(self, cv: ConservedVars):
        r"""Get the specific gas constant ($R_s$)."""

    @abstractmethod
    def heat_capacity_cp(self, cv: ConservedVars):
        r"""Get the specific heat capacity at constant pressure ($C_p$)."""

    @abstractmethod
    def heat_capacity_cv(self, cv: ConservedVars):
        r"""Get the specific heat capacity at constant volume ($C_v$)."""

    @abstractmethod
    def internal_energy(self, cv: ConservedVars):
        """Get the thermal energy of the gas."""

    @abstractmethod
    def total_energy(self, cv: ConservedVars, pressure: np.ndarray):
        """Get the total (thermal + kinetic) energy for the gas."""

    @abstractmethod
    def kinetic_energy(self, cv: ConservedVars):
        """Get the kinetic energy for the gas."""

    @abstractmethod
    def gamma(self, cv: ConservedVars):
        """Get the ratio of gas specific heats Cp/Cv."""

    @abstractmethod
    def transport_model(self):
        """Get the transport model if it exists."""

    @abstractmethod
    def get_internal_energy(self, temperature, *, mass, species_mass_fractions):
        """Get the fluid internal energy from temperature and mass."""

    def dependent_vars(self, cv: ConservedVars) -> EOSDependentVars:
        """Get an agglomerated array of the dependent variables."""
        return EOSDependentVars(
            pressure=self.pressure(cv),
            temperature=self.temperature(cv),
        )

class MixtureEOS(GasEOS):
    r"""Abstract interface to gas mixture equation of state class.

    This EOS base class extends the GasEOS base class to include the
    necessary interface for dealing with gas mixtures.

    .. automethod:: get_density
    .. automethod:: get_species_molecular_weights
    .. automethod:: get_production_rates
    .. automethod:: species_enthalpies
    .. automethod:: get_species_source_terms
    """

    @abstractmethod
    def get_density(self, pressure, temperature, species_mass_fractions):
        """Get the density from pressure, temperature, and species fractions (Y)."""

    @abstractmethod
    def get_species_molecular_weights(self):
        """Get the species molecular weights."""

    @abstractmethod
    def species_enthalpies(self, cv: ConservedVars):
        """Get the species specific enthalpies."""

    @abstractmethod
    def get_production_rates(self, cv: ConservedVars):
        """Get the production rate for each species."""

    @abstractmethod
    def get_species_source_terms(self, cv: ConservedVars):
        r"""Get the species mass source terms to be used on the RHS for chemistry."""

class IdealSingleGas(GasEOS):
    r"""Ideal gas law single-component gas ($p = \rho{R}{T}$).

    The specific gas constant, R, defaults to the air-like 287.1 J/(kg*K),
    but can be set according to simulation units and materials.

    Each interface call expects that the :class:`~mirgecom.fluid.ConservedVars`
    object representing the simulation conserved quantities contains at least
    the canonical conserved quantities mass ($\rho$), energy ($\rho{E}$), and
    momentum ($\rho\vec{V}$).

    .. automethod:: __init__
    .. automethod:: pressure
    .. automethod:: temperature
    .. automethod:: sound_speed
    .. automethod:: internal_energy
    .. automethod:: gas_const
    .. automethod:: dependent_vars
    .. automethod:: total_energy
    .. automethod:: kinetic_energy
    .. automethod:: gamma
    .. automethod:: transport_model
    .. automethod:: get_internal_energy
    """

    def __init__(self, gamma=1.4, gas_const=287.1, transport_model=None):
        """Initialize Ideal Gas EOS parameters."""
        self._gamma = gamma
        self._gas_const = gas_const
        self._transport_model = transport_model

    def transport_model(self):
        """Get the transport model object for this EOS."""
        return self._transport_model

    def gamma(self, cv: ConservedVars = None):
        """Get specific heat ratio Cp/Cv."""
        return self._gamma

    def heat_capacity_cp(self, cv: ConservedVars = None):
        r"""Get specific heat capacity at constant pressure.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector of
            species masses, ($\rho{Y}_\alpha$).
        """
        return self._gas_const * self._gamma / (self._gamma - 1)

    def heat_capacity_cv(self, cv: ConservedVars = None):
        r"""Get specific heat capacity at constant volume.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector of
            species masses, ($\rho{Y}_\alpha$).
        """
        return self._gas_const / (self._gamma - 1)

    def gas_const(self, cv: ConservedVars = None):
        """Get specific gas constant R."""
        return self._gas_const

    def kinetic_energy(self, cv: ConservedVars):
        r"""Get kinetic (i.e. not internal) energy of gas.

        The kinetic energy is calculated as:

        .. math::

            k = \frac{1}{2\rho}(\rho\vec{V} \cdot \rho\vec{V})

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The kinetic energy of the fluid flow
        """
        mom = cv.momentum
        return (0.5 * np.dot(mom, mom) / cv.mass)

    def internal_energy(self, cv: ConservedVars):
        r"""Get internal thermal energy of gas.

        The internal energy (e) is calculated as:

        .. math::

            e = \rho{E} - \frac{1}{2\rho}(\rho\vec{V} \cdot \rho\vec{V})

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The internal energy of the fluid material
        """
        return (cv.energy - self.kinetic_energy(cv))

    def pressure(self, cv: ConservedVars):
        r"""Get thermodynamic pressure of the gas.

        Gas pressure (p) is calculated from the internal energy (e) as:

        .. math::

            p = (\gamma - 1)e

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The fluid pressure
        """
        return self.internal_energy(cv) * (self._gamma - 1.0)

    def sound_speed(self, cv: ConservedVars):
        r"""Get the speed of sound in the gas.

        The speed of sound (c) is calculated as:

        .. math::

            c = \sqrt{\frac{\gamma{p}}{\rho}}

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The speed of sound in the fluid
        """
        actx = cv.array_context
        return actx.np.sqrt(self._gamma / cv.mass * self.pressure(cv))

    def temperature(self, cv: ConservedVars):
        r"""Get the thermodynamic temperature of the gas.

        The thermodynamic temperature (T) is calculated from
        the internal energy (e) and specific gas constant (R)
        as:

        .. math::

            T = \frac{(\gamma - 1)e}{R\rho}

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The fluid temperature
        """
        return (
            (((self._gamma - 1.0) / self._gas_const)
             * self.internal_energy(cv) / cv.mass)
        )

    def total_energy(self, cv, pressure):
        r"""
        Get gas total energy from mass, pressure, and momentum.

        The total energy density (rhoE) is calculated from
        the mass density (rho) , pressure (p) , and
        momentum (rhoV) as:

        .. math::

            \rho{E} = \frac{p}{(\gamma - 1)} +
            \frac{1}{2}\rho(\vec{v} \cdot \vec{v})

        .. note::

            The total_energy function computes cv.energy from pressure,
            mass, and momentum in this case. In general in the EOS we need
            DV = EOS(CV), and inversions CV = EOS(DV). This is one of those
            inversion interfaces.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$).

        pressure: :class:`~meshmode.dof_array.DOFArray`
            The fluid pressure

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The total energy of the fluid (i.e. internal + kinetic)
        """
        return (pressure / (self._gamma - 1.0)
                + self.kinetic_energy(cv))

    def get_internal_energy(self, temperature, mass, species_mass_fractions=None):
        r"""Get the gas thermal energy from temperature, and fluid density.

        The gas internal energy $e$ is calculated from:

        .. math::

            e = R_s T \frac{\rho}{\left(\gamma - 1\right)}

        Parameters
        ----------
        temperature: :class:`~meshmode.dof_array.DOFArray`
            The fluid temperature
        mass: float or :class:`~meshmode.dof_array.DOFArray`
            The fluid mass density
        species_mass_fractions:
            Unused
        """
        return self._gas_const * mass * temperature / (self._gamma - 1)

class PyrometheusMixture(MixtureEOS):
    r"""Ideal gas mixture ($p = \rho{R}_\mathtt{mix}{T}$).

    This is the :mod:`pyrometheus`-based EOS. Please refer to the :any:`documentation
    of Pyrometheus <pyrometheus>` for underlying implementation details.

    Each interface call expects that the :class:`mirgecom.fluid.ConservedVars` object
    representing the simulation conserved quantities contains at least the
    canonical conserved quantities mass ($\rho$), energy ($\rho{E}$), and
    momentum ($\rho\vec{V}$), and the vector of species masses, ($\rho{Y}_\alpha$).

    .. important::
        When using this EOS, users should take care to match solution initialization
        with the appropriate units that are used in the user-provided `Cantera`
        mechanism input files.

    .. automethod:: __init__
    .. automethod:: pressure
    .. automethod:: temperature
    .. automethod:: sound_speed
    .. automethod:: internal_energy
    .. automethod:: gas_const
    .. automethod:: dependent_vars
    .. automethod:: total_energy
    .. automethod:: kinetic_energy
    .. automethod:: gamma
    .. automethod:: transport_model
    .. automethod:: get_internal_energy
    .. automethod:: get_density
    .. automethod:: get_species_molecular_weights
    .. automethod:: get_production_rates
    .. automethod:: species_enthalpies
    .. automethod:: get_species_source_terms
    """

    def __init__(self, pyrometheus_mech, temperature_guess=300.0,
                 transport_model=None):
        """Initialize Pyrometheus-based EOS with mechanism class.

        Parameters
        ----------
        pyrometheus_mech: :class:`pyrometheus.Thermochemistry`
            The :mod:`pyrometheus`  mechanism :class:`~pyrometheus.Thermochemistry`
            object that is generated by the user with a call to
            *pyrometheus.get_thermochem_class*. To create the mechanism
            object, users need to provide a mechanism input file. Several built-in
            mechanisms are provided in `mirgecom/mechanisms/` and can be used through
            the :meth:`mirgecom.mechanisms.get_mechanism_cti`.

        tguess: float
            This provides a constant starting temperature for the Newton iterations
            used to find the mixture temperature. It defaults to 300.0 Kelvin. This
            parameter is important for the performance and proper function of the
            code. Users should set a tguess that is close to the average temperature
            of the simulated domain.  Ideally, we would use the last computed
            temperature for the guess, but doing so requires restart infrastructure
            that is TBD.
        """
        self._pyrometheus_mech = pyrometheus_mech
        self._tguess = temperature_guess
        self._transport_model = transport_model

    def transport_model(self):
        """Get the transport model object for this EOS."""
        return self._transport_model

    def heat_capacity_cp(self, cv: ConservedVars):
        r"""Get mixture-averaged specific heat capacity at constant pressure.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector of
            species masses, ($\rho{Y}_\alpha$).
        """
        temp = self.temperature(cv)
        y = cv.species_mass_fractions
        return self._pyrometheus_mech.get_mixture_specific_heat_cp_mass(temp, y)

    def heat_capacity_cv(self, cv: ConservedVars):
        r"""Get mixture-averaged specific heat capacity at constant volume.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector of
            species masses, ($\rho{Y}_\alpha$).
        """
        temp = self.temperature(cv)
        y = cv.species_mass_fractions
        return (
            self._pyrometheus_mech.get_mixture_specific_heat_cp_mass(temp, y)
            / self.gamma(cv)
        )

    def gamma(self, cv: ConservedVars):
        r"""Get mixture-averaged specific heat ratio for mixture $\frac{C_p}{C_p - R_s}$.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector of
            species masses, ($\rho{Y}_\alpha$).
        """
        temperature = self.temperature(cv)
        y = cv.species_mass_fractions
        cp = self._pyrometheus_mech.get_mixture_specific_heat_cp_mass(temperature, y)
        rspec = self.gas_const(cv)
        return cp / (cp - rspec)

    def gas_const(self, cv: ConservedVars):
        r"""Get specific gas constant $R_s$.

        The mixture specific gas constant is calculated
        as $R_s = \frac{R}{\sum{\frac{{Y}_\alpha}{{M}_\alpha}}}$ by the
        :mod:`pyrometheus` mechanism provided by the user. ${M}_\alpha$ are the
        species molar masses.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).
        """
        y = cv.species_mass_fractions
        return self._pyrometheus_mech.get_specific_gas_constant(y)

    def kinetic_energy(self, cv: ConservedVars):
        r"""Get kinetic (i.e. not internal) energy of gas.

        The kinetic energy is calculated as:

        .. math::

            k = \frac{1}{2\rho}(\rho\vec{V} \cdot \rho\vec{V})

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).
        """
        mom = cv.momentum
        return (0.5 * np.dot(mom, mom) / cv.mass)

    def internal_energy(self, cv: ConservedVars):
        r"""Get internal thermal energy of gas.

        The internal energy ($e$) is calculated as:

        .. math::

            e = \rho{E} - \frac{1}{2\rho}(\rho\vec{V} \cdot \rho\vec{V})

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            Internal energy of the fluid
        """
        return (cv.energy - self.kinetic_energy(cv))

    def get_density(self, pressure, temperature, species_mass_fractions):
        r"""Get the density from pressure, temperature, and species fractions (Y).

        The gas density $\rho$ is calculated from pressure, temperature and $R$ as:

        .. math::

            \rho = \frac{p}{R_s T}

        Parameters
        ----------
        pressure: :class:`~meshmode.dof_array.DOFArray`
            The fluid pressure
        temperature: :class:`~meshmode.dof_array.DOFArray`
            The fluid temperature
        species_mass_fractions: numpy.ndarray
            Object array of species mass fractions

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The total fluid mass density
        """
        return self._pyrometheus_mech.get_density(pressure, temperature,
                                                  species_mass_fractions)

    def get_internal_energy(self, temperature, species_mass_fractions,
                            mass=None):
        r"""Get the gas thermal energy from temperature, and species fractions (Y).

        The gas internal energy $e$ is calculated from:

        .. math::

            e = R_s T \sum{Y_\alpha h_\alpha}

        Parameters
        ----------
        temperature: :class:`~meshmode.dof_array.DOFArray`
            The fluid temperature
        species_mass_fractions: numpy.ndarray
            Object array of species mass fractions
        mass:
            Unused
        """
        return self._pyrometheus_mech.get_mixture_internal_energy_mass(
            temperature, species_mass_fractions)

    def get_species_molecular_weights(self):
        """Get the species molecular weights."""
        return self._pyrometheus_mech.wts

    def species_enthalpies(self, cv: ConservedVars):
        """Get the species specific enthalpies."""
        return self._pyrometheus_mech.get_species_enthalpies_rt(self.temperature(cv))

    def get_production_rates(self, cv: ConservedVars):
        r"""Get the production rate for each species.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).

        Returns
        -------
        numpy.ndarray
            The chemical production rates for each species
        """
        temperature = self.temperature(cv)
        y = cv.species_mass_fractions
        return self._pyrometheus_mech.get_net_production_rates(
            cv.mass, temperature, y)

    def pressure(self, cv: ConservedVars):
        r"""Get thermodynamic pressure of the gas.

        Gas pressure ($p$) is calculated from the internal energy ($e$) as:

        .. math::

            p = (\gamma_{\mathtt{mix}} - 1)e

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The pressure of the fluid.
        """
        @memoize_in(cv, (PyrometheusMixture.pressure,
                         type(self._pyrometheus_mech)))
        def get_pressure():
            temperature = self.temperature(cv)
            y = cv.species_mass_fractions
            return self._pyrometheus_mech.get_pressure(cv.mass, temperature, y)
        return get_pressure()

    def sound_speed(self, cv: ConservedVars):
        r"""Get the speed of sound in the gas.

        The speed of sound ($c$) is calculated as:

        .. math::

            c = \sqrt{\frac{\gamma_{\mathtt{mix}}{p}}{\rho}}

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The speed of sound in the fluid.
        """
        @memoize_in(cv, (PyrometheusMixture.sound_speed,
                         type(self._pyrometheus_mech)))
        def get_sos():
            actx = cv.array_context
            return actx.np.sqrt((self.gamma(cv) * self.pressure(cv)) / cv.mass)
        return get_sos()

    def temperature(self, cv: ConservedVars):
        r"""Get the thermodynamic temperature of the gas.

        The thermodynamic temperature ($T$) is calculated from
        the internal energy ($e$) and specific gas constant ($R_s$)
        as:

        .. math::

            T = \frac{(\gamma_{\mathtt{mix}} - 1)e}{R_s \rho}

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The temperature of the fluid.
        """
        @memoize_in(cv, (PyrometheusMixture.temperature,
                         type(self._pyrometheus_mech)))
        def get_temp():
            y = cv.species_mass_fractions
            e = self.internal_energy(cv) / cv.mass
            return self._pyrometheus_mech.get_temperature(e, self._tguess,
                                                          y, True)
        return get_temp()

    def total_energy(self, cv, pressure):
        r"""
        Get gas total energy from mass, pressure, and momentum.

        The total energy density ($\rho E$) is calculated from
        the mass density ($\rho$) , pressure ($p$) , and
        momentum ($\rho\vec{V}$) as:

        .. math::

            \rho E = \frac{p}{(\gamma_{\mathtt{mix}} - 1)} +
            \frac{1}{2}\rho(\vec{v} \cdot \vec{v})

        .. note::

            The total_energy function computes cv.energy from pressure,
            mass, and momentum in this case. In general in the EOS we need
            DV = EOS(CV), and inversions CV = EOS(DV). This is one of those
            inversion interfaces.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).
        pressure: :class:`~meshmode.dof_array.DOFArray`
            The fluid pressure

        Returns
        -------
        :class:`~meshmode.dof_array.DOFArray`
            The total energy fo the fluid (i.e. internal + kinetic)
        """
        return (pressure / (self.gamma(cv) - 1.0)
                + self.kinetic_energy(cv))

    def get_species_source_terms(self, cv: ConservedVars):
        r"""Get the species mass source terms to be used on the RHS for chemistry.

        Parameters
        ----------
        cv: :class:`~mirgecom.fluid.ConservedVars`
            :class:`~mirgecom.fluid.ConservedVars` containing at least the mass
            ($\rho$), energy ($\rho{E}$), momentum ($\rho\vec{V}$), and the vector
            of species masses, ($\rho{Y}_\alpha$).

        Returns
        -------
        :class:`~mirgecom.fluid.ConservedVars`
            Chemistry source terms
        """
        omega = self.get_production_rates(cv)
        w = self.get_species_molecular_weights()
        dim = cv.dim
        species_sources = w * omega
        rho_source = 0 * cv.mass
        mom_source = 0 * cv.momentum
        energy_source = 0 * cv.energy

        return make_conserved(dim, rho_source, energy_source, mom_source,
                              species_sources)


from itertools import groupby
import numbers
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from models.layers.meta_sequential import MetaSequential
from models.layers.meta_patch import MetaPatchConv2d, make_meta_patch_conv2d_block

class HyperGen(nn.Module):
    """ Hypernetwork generator comprised of a backbone network, weight mapper, and a decoder.

    Args:
        backbone (nn.Module factory): Backbone network
        weight_mapper (nn.Module factory): Weight mapper network.
        in_nc (int): input number of channels.
        num_classes (int): output number of classes.
        kernel_sizes (int): the kernel size of the decoder layers.
        level_layers (int): number of layers in each level of the decoder.
        level_channels (list of int, optional): If specified, sets the output channels of each level in the decoder.
        expand_ratio (int): inverted residual block's expansion ratio in the decoder.
        groups (int, optional): Number of blocked connections from input channels to output channels.
        weight_groups (int, optional): per level signal to weights in the decoder.
        inference_hflip (bool): If true, enables horizontal flip of input tensor.
        inference_gather (str): Inference gather type: ``mean'' or ``max''.
        with_out_fc (bool): If True, add a final fully connected layer to the decoder.
        decoder_groups (int, optional): per level groups in the decoder.
        decoder_dropout (float): If specified, enables dropout with the given probability.
        coords_res (list of tuple of int, optional): list of inference resolutions for caching positional embedding.
    """
    def __init__(self, backbone, weight_mapper, in_nc=3, num_classes=3, kernel_sizes=3, level_layers=1, expand_ratio=1,
                 groups=1, inference_hflip=False, inference_gather='mean', with_out_fc=False, decoder_dropout=None):
        super(HyperGen, self).__init__()
        self.inference_hflip = inference_hflip
        self.inference_gather = inference_gather

        self.backbone = backbone()
        feat_channels = [in_nc] + self.backbone.feat_channels[:-1]
        self.decoder = MultiScaleDecoder(feat_channels, 3, num_classes, kernel_sizes, level_layers,
                                         with_out_fc=with_out_fc, out_kernel_size=1, expand_ratio=expand_ratio,
                                         dropout=decoder_dropout)
        self.weight_mapper = weight_mapper(self.backbone.feat_channels[-1], self.decoder.param_groups)

    @property
    def hyper_params(self):
        return self.decoder.hyper_params

    def process_single_tensor(self, x, hflip=False):
        x = torch.flip(x, [-1]) if hflip else x
        features = self.backbone(x)
        weights = self.weight_mapper(features[-1])
        x = [x] + features[:-1]
        x = self.decoder(x, weights)
        x = torch.flip(x, [-1]) if hflip else x

        return x

    def gather_results(self, x, y=None):
        assert x is not None
        if y is None:
            return x
        if self.inference_gather == 'mean':
            return (x + y) * 0.5
        else:
            return torch.max(x, y)

    def forward(self, x):
        assert isinstance(x, (list, tuple, torch.Tensor)), f'x must be of type list, tuple, or tensor'
        if isinstance(x, torch.Tensor):
            return self.process_single_tensor(x)

        # Note: the first pyramid will determine the output resolution
        out_res = x[0].shape[2:]
        out = None
        for p in x:
            if self.inference_hflip:
                p = torch.max(self.process_single_tensor(p), self.process_single_tensor(p, hflip=True))
            else:
                p = self.process_single_tensor(p)

            # Resize current image to output resolution if necessary
            if p.shape[2:] != out_res:
                p = F.interpolate(p, out_res, mode='bilinear', align_corners=False)

            out = self.gather_results(p, out)

        return out

class MultiScaleDecoder(nn.Module):
    """ Dynamic multi-scale decoder.

    Args:
        feat_channels (list of int): per level input feature channels.
        signal_channels (list of int): per level input signal channels.
        num_classes (int): output number of classes.
        kernel_sizes (int): the kernel size of the layers.
        level_layers (int): number of layers in each level.
        level_channels (list of int, optional): If specified, sets the output channels of each level.
        norm_layer (nn.Module): Type of feature normalization layer
        act_layer (nn.Module): Type of activation layer
        out_kernel_size (int): kernel size of the final output layer.
        expand_ratio (int): inverted residual block's expansion ratio.
        groups (int, optional): number of blocked connections from input channels to output channels.
        weight_groups (int, optional): per level signal to weights.
        with_out_fc (bool): If True, add a final fully connected layer.
        dropout (float): If specified, enables dropout with the given probability.
        coords_res (list of tuple of int, optional): list of inference resolutions for caching positional embedding.
    """
    def __init__(self, feat_channels, in_nc=3, num_classes=3, kernel_sizes=3, level_layers=1, norm_layer=nn.BatchNorm2d,
                 act_layer=nn.ReLU6(inplace=True), out_kernel_size=1, expand_ratio=1, with_out_fc=False, dropout=None):
        super(MultiScaleDecoder, self).__init__()
        if isinstance(kernel_sizes, numbers.Number):
            kernel_sizes = (kernel_sizes,) * len(feat_channels)
        if isinstance(level_layers, numbers.Number):
            level_layers = (level_layers,) * len(feat_channels)
        assert len(kernel_sizes) == len(feat_channels), \
            f'kernel_sizes ({len(kernel_sizes)}) must be of size {len(feat_channels)}'
        assert len(level_layers) == len(feat_channels), \
            f'level_layers ({len(level_layers)}) must be of size {len(feat_channels)}'
        self.level_layers = level_layers
        self.levels = len(level_layers)
        self.layer_params = []
        feat_channels = feat_channels[::-1]  # Reverse the order of the feature channels

        # For each level
        prev_channels = 0
        for level in range(self.levels):
            curr_ngf = feat_channels[level]
            prev_channels += curr_ngf  # Accommodate the previous number of channels
            curr_layers = []
            kernel_size = kernel_sizes[level]

            # For each layer in the current level
            for layer in range(self.level_layers[level]):
                if (not with_out_fc) and (level == (self.levels - 1) and (layer == (self.level_layers[level] - 1))):
                    curr_ngf = num_classes
                if kernel_size > 1:
                    curr_layers.append(HyperPatchInvertedResidual(
                        prev_channels + 2, curr_ngf, kernel_size, expand_ratio=expand_ratio, norm_layer=norm_layer,
                        act_layer=act_layer))
                else:
                    curr_layers.append(make_meta_patch_conv2d_block(prev_channels + 2, curr_ngf, kernel_size))
                prev_channels = curr_ngf

            # Add level layers to module
            self.add_module(f'level_{level}', MetaSequential(*curr_layers))

        # Add the last layer
        if with_out_fc:
            out_fc_layers = [nn.Dropout2d(dropout, True)] if dropout is not None else []
            out_fc_layers.append(
                MetaPatchConv2d(prev_channels, num_classes, out_kernel_size, padding=out_kernel_size // 2))
            self.out_fc = MetaSequential(*out_fc_layers)
        else:
            self.out_fc = None

        # Calculate number of hyper parameters, weight ranges, and total number of hyper parameters per level
        self.hyper_params = 0
        self._ranges = [0]
        self.param_groups = []
        for level in range(self.levels):
            level_layers = getattr(self, f'level_{level}')
            self.hyper_params += level_layers.hyper_params
            self._ranges.append(self.hyper_params)
            self.param_groups.append(level_layers.hyper_params)
        if with_out_fc:
            self.hyper_params += self.out_fc.hyper_params
            self.param_groups.append(self.out_fc.hyper_params)
        self._ranges.append(self.hyper_params)

    def forward(self, x, w):
        assert isinstance(w, (list, tuple))
        assert len(x) <= self.levels

        # For each level
        p = None
        for level in range(len(x)):
            level_w = w[level]
            level_layers = getattr(self, f'level_{level}')

            # Initial layer input
            if p is None:
                p = x[-level - 1]
            else:
                # p = F.interpolate(p, scale_factor=2, mode='bilinear', align_corners=False)  # Upsample x2
                if p.shape[2:] != x[-level - 1].shape[2:]:
                    p = F.interpolate(p, x[-level - 1].shape[2:], mode='bilinear', align_corners=False)  # Upsample
                p = torch.cat((x[-level - 1], p), dim=1)

            # Add image coordinates
            p = torch.cat([get_image_coordinates(p.shape[0], *p.shape[-2:], p.device), p], dim=1)

            # Computer the output for the current level
            p = level_layers(p, level_w)

        # Last layer
        if self.out_fc is not None:
            p = self.out_fc(p, w[-1])

        return p

class HyperPatchInvertedResidual(nn.Module):
    def __init__(self, in_nc, out_nc, kernel_size=3, stride=1, expand_ratio=1, norm_layer=nn.BatchNorm2d,
                 act_layer=nn.ReLU6(inplace=True), padding_mode='reflect'):
        super(HyperPatchInvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = int(round(in_nc * expand_ratio))
        self.use_res_connect = self.stride == 1 and in_nc == out_nc

        layers = []
        if expand_ratio != 1:
            # pw
            layers.append(make_meta_patch_conv2d_block(in_nc, hidden_dim, 1, norm_layer=norm_layer,
                                                       act_layer=act_layer))
        layers.extend([
            # dw
            make_meta_patch_conv2d_block(hidden_dim, hidden_dim, kernel_size, stride=stride, groups=hidden_dim,
                                         norm_layer=norm_layer, act_layer=act_layer, padding_mode=padding_mode),
            # pw-linear
            make_meta_patch_conv2d_block(hidden_dim, out_nc, 1, stride=stride, norm_layer=norm_layer, act_layer=None)
        ])
        self.conv = MetaSequential(*layers)

    @property
    def hyper_params(self):
        return self.conv.hyper_params

    def forward(self, x, w):
        if self.use_res_connect:
            return x + self.conv(x, w)
        else:
            return self.conv(x, w)

def get_image_coordinates(b, h, w, device):
    x = torch.linspace(-1, 1, steps=w, device=device)
    y = torch.linspace(-1, 1, steps=h, device=device)
    grid = torch.stack(torch.meshgrid(y, x)[::-1], dim=0).repeat(b, 1, 1, 1)
    # grid = torch.stack(torch.meshgrid(x, y)[::-1], dim=0).repeat(b, 1, 1, 1)

    return grid

class WeightMapper(nn.Module):
    """ Weight mapper module (called context head in the paper).

    Args:
        in_channels (int): input number of channels.
        out_channels (int): output number of channels.
        levels (int): number of levels operating on different strides.
        bias (bool): if True, enables bias in all convolution operations.
        min_unit (int): used when dividing the signal channels into parts, must be a multiple of this number.
        weight_groups (int): per level signal to weights groups.
    """
    def __init__(self, in_channels, out_channels, levels=2, bias=False, min_unit=8, down_groups=1, flat_groups=1,
                 weight_groups=1, avg_pool=False):
        super(WeightMapper, self).__init__()
        assert levels > 0, 'levels must be greater than zero'
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.levels = levels
        self.bias = bias
        self.avg_pool = avg_pool
        self.down_groups = down_groups
        self.flat_groups = flat_groups
        self.weight_groups = weight_groups

        min_unit = max(min_unit, weight_groups)

        for level in range(self.levels - 1):
            down = nn.Sequential(
                nn.Conv2d(in_channels, in_channels, kernel_size=2, stride=2, bias=bias, groups=down_groups),
                nn.BatchNorm2d(in_channels),
                nn.ReLU(inplace=True))
            self.add_module(f'down_{level}', down)

            up = nn.UpsamplingNearest2d(scale_factor=2)
            self.add_module(f'up_{level}', up)

            flat = [nn.Conv2d(in_channels * 2, in_channels, kernel_size=1, bias=bias, groups=flat_groups),
                    nn.BatchNorm2d(in_channels)]
            if level > 0:
                flat.append(nn.ReLU(inplace=True))
            flat = nn.Sequential(*flat)
            self.add_module(f'flat_{level}', flat)

        out_channels = [next_multiply(c, weight_groups) for c in out_channels]
        self.out_conv = Conv2dMulti(in_channels, out_channels, 1, bias=bias, min_unit=min_unit, groups=weight_groups)

    def forward(self, x):
        if self.levels <= 1:
            return self.out_conv(x)

        # Down stream
        feat = [x]
        for level in range(self.levels - 1):
            down = getattr(self, f'down_{level}')
            feat.append(down(feat[-1]))

        # Average the last feature map
        if self.avg_pool:
            orig_shape = feat[-1].shape
            if orig_shape[-2:] != (1, 1):
                feat[-1] = F.adaptive_avg_pool2d(feat[-1], 1)
                feat[-1] = F.interpolate(feat[-1], orig_shape[-2:], mode='nearest')

        # Up stream
        for level in range(self.levels - 2, -1, -1):
            up = getattr(self, f'up_{level}')
            flat = getattr(self, f'flat_{level}')

            x = up(feat.pop(-1))
            feat[-1] = torch.cat((feat[-1], x), dim=1)
            feat[-1] = flat(feat[-1])

        # Output weights
        w = self.out_conv(feat[-1])
        if self.weight_groups > 1:
            w = [wi[:, :oc] for wi, oc in zip(w, self.out_channels)]

        return w

    def extra_repr(self):
        return f'in_channels={self.in_channels}, out_channels={self.out_channels}, bias={self.bias}'

def next_multiply(x, base):
    return type(x)(np.ceil(x / base) * base)

class Conv2dMulti(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1,
                 bias=True, padding_mode='zeros', min_unit=8):
        super(Conv2dMulti, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.bias = bias
        self._ranges = [0]

        in_nc_parts = divide_feature_legacy(in_channels, out_channels, min_unit)
        for i, out_nc in enumerate(out_channels):
            in_nc = in_nc_parts[i]
            self._ranges.append(self._ranges[-1] + in_nc)

            conv = nn.Conv2d(in_nc, out_nc, kernel_size, stride, padding, dilation, groups, bias, padding_mode)
            self.add_module(f'conv_{i}', conv)

    def forward(self, x):
        out = []
        for i in range(len(self.out_channels)):
            linear = getattr(self, f'conv_{i}')
            out.append(linear(x[:, self._ranges[i]:self._ranges[i + 1]]))

        return out

    def extra_repr(self):
        return f'in_channels={self.in_channels}, out_channels={self.out_channels}, bias={self.bias}'

# Contains bugs (but used in some modules)
def divide_feature_legacy(in_feature, out_features, min_unit=8):
    """ Divides in_feature relative to each of the provided out_features.
    The division of the input feature will be in multiplies of "min_unit".
    The algorithm makes sure that equal output features will get the same portion of the input feature.
    The smallest out feature will receive all the round down overflow (usually the final fc)
    Args:
        in_feature: the input feature to divide
        out_features: the relative sizes of the output features
        min_unit: each division of the input feature will be divisible by this number.
        in_feature must be divisible by this number as well
    Returns:
        np.array: array of integers of the divided input feature in the size of out_features.
    """
    assert in_feature % min_unit == 0, f'in_feature ({in_feature}) must be divisible by min_unit ({min_unit})'
    units = in_feature // min_unit
    indices = np.argsort(out_features)
    out_features_sorted = np.array(out_features)[indices]
    out_feat_groups = [(k, indices[list(g)]) for k, g in groupby(range(len(indices)), lambda i: out_features_sorted[i])]
    out_feat_groups.sort(key=lambda x: x[0] * len(x[1]), reverse=True)
    units_feat_ratio = float(units) / sum(out_features)

    # For each feature group
    remaining_units = units
    out_group_units = []
    for i, out_feat_group in enumerate(out_feat_groups):  # out_feat_group: (out_feature, indices array)
        if i < (len(out_feat_groups) - 1):
            curr_out_feat_size = out_feat_group[0] * len(out_feat_group[1])
            curr_units = max(curr_out_feat_size * units_feat_ratio, 1)
            curr_units = curr_units // len(out_feat_group[1]) * len(out_feat_group[1])  # Make divisible by num elements
            out_group_units.append(curr_units)
            remaining_units -= curr_units
        else:
            out_group_units.append(remaining_units)

    # Final feature division
    divided_in_features = np.zeros(len(out_features), dtype=int)
    for i, out_feat_group in enumerate(out_feat_groups):
        for j in range(len(out_feat_group[1])):
            divided_in_features[out_feat_group[1][j]] = out_group_units[i] // len(out_feat_group[1]) * min_unit

    return divided_in_features

def hyperseg_efficientnet(model_name, pretrained=False, levels=3, down_groups=1, flat_groups=1, weight_groups=1,
                          avg_pool=True, weights_path=None, **kwargs):
    from models.backbones.efficientnet import efficientnet
    from functools import partial

    weight_mapper = partial(WeightMapper, levels=levels, down_groups=down_groups, flat_groups=flat_groups,
                            weight_groups=weight_groups, avg_pool=avg_pool)
    backbone = partial(efficientnet, model_name, pretrained=pretrained, head=None, return_features=True)
    model = HyperGen(backbone, weight_mapper, **kwargs)

    if weights_path is not None:
        checkpoint = torch.load(weights_path)
        state_dict = checkpoint['state_dict']
        model.load_state_dict(state_dict, strict=True)

    return model

def main(model='models.hyperseg_v0_1.hyperseg_efficientnet', res=(512,), pyramids=None,
         train=False):
    from utils.obj_factory import obj_factory
    from utils.utils import set_device
    from utils.img_utils import create_pyramid

    assert len(res) <= 2, f'res must be either a single number or a pair of numbers: "{res}"'
    res = res * 2 if len(res) == 1 else res

    device, gpus = set_device()
    model = obj_factory(model).to(device).train(train)
    x = torch.rand(2, 3, *res).to(device)
    x = create_pyramid(x, pyramids) if pyramids is not None else x
    pred = model(x)
    print(pred.shape)

if __name__ == "__main__":
    # Parse program arguments
    import argparse

    parser = argparse.ArgumentParser('hyperseg test')
    parser.add_argument('-m', '--model', default='models.hyperseg_v0_1.hyperseg_efficientnet',
                        help='model object')
    parser.add_argument('-r', '--res', default=(512,), type=int, nargs='+',
                        metavar='N', help='image resolution')
    parser.add_argument('-p', '--pyramids', type=int, metavar='N',
                        help='number of image pyramids')
    parser.add_argument('-t', '--train', action='store_true',
                        help='If True, sets the model to training mode')
    main(**vars(parser.parse_args()))


import os
import numpy as np
import subprocess
from sklearn.metrics import f1_score, accuracy_score
from .utils_model import *

def predict():
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    data_dir = "media/audios/"
    mode = "unlabeled"  # real, fake, or unlabeled
    pretrained_model_name = 'FinalModel.h5'
    print(f"Loading inference data from {os.path.join(data_dir,mode)}")
    print(f"Loading pretrained model {pretrained_model_name}")

    # preprocess the files
    processed_data = preprocess_from_ray_parallel_inference(
        data_dir, mode, use_parallel=False)
    processed_data = sorted(processed_data, key=lambda x: len(x[0]))

    # Load trained model
    discriminator = Discriminator_Model(
        load_pretrained=True, saved_model_name=pretrained_model_name, real_test_mode=False)

    # Do inference
    if mode == 'unlabeled':

        print("The probability of the clip being real is: {:.2%}".format(
            discriminator.predict_labels(processed_data, raw_prob=True, batch_size=20)[0][0]))

        return round(discriminator.predict_labels(processed_data, raw_prob=True, batch_size=20)[0][0], 3)


#!/usr/bin/env python
from __future__ import print_function

import os
from typing import Any, Dict

import cv2
import matplotlib.pyplot as plt
import multiprocessing
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from docopt import docopt
from PIL import Image
from torch.autograd import Variable

import deeplab.resnet as resnet
import deeplab.datasets as datasets
from deeplab.deeplab_resnet import MS_Deeplab, Res_Deeplab

DOCSTR = """Evaluate ResNet-DeepLab trained on scenes (VOC2012), a total of 21
labels including background.

Usage:
    test.py [options]

Options:
    -h, --help                  Print this message
    --snapPath=<str>            Path to snapshot to evaluate.
    --valPath=<str>             Path to file list of validation images.
    --imgPath=<str>             Path to image directory.
    --gtPath=<str>              Page to ground truth directory.
    --imgExt=<str>              File extension for images.
    --gtExt=<str>               File extension for ground truth
    --numLabels=<int>           Number of different labels in training data,
                                including background
    --gpu=<int>                 GPU number [default: 0]
    --visualize                 Generate visualizations of model outputs.

"""

def fast_hist(ground_truth: np.ndarray, output: np.ndarray, num_labels: int) -> np.ndarray:
    valid_idx = (ground_truth >= 0) & (ground_truth < num_labels)
    gt_contrib = num_labels * ground_truth[valid_idx].astype(int)
    confusion_mat = np.bincount(
        gt_contrib + output[valid_idx], minlength=num_labels**2)
    return confusion_mat.reshape(num_labels, num_labels)

def get_model(num_labels: int, snapshot_path: str, gpu: int) -> MS_Deeplab:
    """Retrieve model with appropriate parameters"""
    return resnet.getDeepLabV2(num_labels, snapshot_path).eval().cuda(gpu)

def make_dataloader(args: Dict[str, str]) -> Any:
    """Construct appropriate dataloader"""
    dataset = datasets.SegmentationDataset(
        datasets.Mode.VAL, args["--valPath"], args["--imgPath"],
        args["--gtPath"], args["--imgExt"], args["--gtExt"])
    return DataLoader(dataset, num_workers=multiprocessing.cpu_count() - 1)

def validate(model: MS_Deeplab, num_labels: int, gpu: int, args: Dict[str, str]) -> None:
    hist = np.zeros((num_labels, num_labels))
    for i, (img, gt) in enumerate(make_dataloader(args)):
        print('processing {}'.format(i))

        # Rescale ground truth to [0, 255] for histogram calculation
        gt = np.array(gt.squeeze() * 255, dtype="uint8")

        with torch.no_grad():
            # TODO: Rescale to 255 only for pretrained models trained w [0-255]
            output = model(Variable(img * 255).cuda(gpu))

        # Resize to match ground truth size and take highest probability label
        output = F.interpolate(
            output[3], (513, 513), mode='bilinear', align_corners=True).cpu()
        output = output.data[0].numpy()[:, :gt.shape[0], :gt.shape[1]]
        output = np.argmax(output, axis=0).astype(np.uint8)

        if args['--visualize']:
            plt.subplot(2, 1, 1)
            plt.imshow(gt)
            plt.subplot(2, 1, 2)
            plt.imshow(output)
            plt.show()

        hist += fast_hist(gt.flatten(), output.flatten(), num_labels)

    miou = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
    print('Mean iou = {}'.format(np.sum(miou) / len(miou)))

def main():
    args = docopt(DOCSTR)
    print(args)

    # Create model
    num_labels = int(args['--numLabels'])
    gpu = int(args['--gpu'])
    model = get_model(num_labels, args['--snapPath'], gpu)

    # Select which validation to use
    validate(model, num_labels, gpu, args)

if __name__ == "__main__":
    main()


import os, time
import argparse,logging
import numpy as np
import mxnet as mx
import cv2, dlib, urllib
from lightened_moon import lightened_moon_feature
from shutil import copyfile

logger = logging.getLogger()
logger.setLevel(logging.INFO)
import pdb

def detectFaces(img, detector):
    faces = detector(img, 1)
    return faces

def detectFacesInCascade(img, face_cascade):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    opencv_faces = face_cascade.detectMultiScale(gray, 1.3, 5, 0, (20,20))

    faces = []
    for (x,y,w,h) in opencv_faces:
    #print x, y, x+w, y+h
        faces.append(dlib.rectangle(int(x), int(y), int(x+w), int(y+h)))

    return faces

def attribute(img, detector, face_cascade, devs, symbol, arg_params, aux_params, argssize):
    h,w,c=img.shape
    scale = max(h/240.0, w/320.0)
    neww = int(w/scale)
    newh = int(h/scale)
    newimg = cv2.resize(img,(neww, newh))

    # read img and drat face rect
    faces = detectFaces(newimg, detector)
    if len(faces) == 0:
        faces = detectFacesInCascade(newimg, face_cascade)

    gray = np.zeros(img.shape[0:2])

    for i in range(len(faces)):
        faces[i] = dlib.rectangle(int(faces[i].left() *scale), int(faces[i].top()*scale), int(faces[i].right()*scale), int(faces[i].bottom()*scale))

    font = cv2.FONT_HERSHEY_SIMPLEX
    if len(faces) > 0:
            max_face = max(faces, key=lambda rect: rect.width() * rect.height())

            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            pad = [0.25, 0.25, 0.25, 0.25]
            left = int(max(0, max_face.left() - max_face.width()*float(pad[0])))
            top = int(max(0, max_face.top() - max_face.height()*float(pad[1])))
            right = int(min(gray.shape[1], max_face.right() + max_face.width()*float(pad[2])))
            bottom = int(min(gray.shape[0], max_face.bottom()+max_face.height()*float(pad[3])))

            gray = gray[top:bottom, left:right]
            #print gray.shape, left, right, top, right,  argssize
            gray = cv2.resize(gray, (argssize, argssize))/255.0
            imgpred= np.expand_dims(np.expand_dims(gray, axis=0), axis=0)
            # get pred
            arg_params['data'] = mx.nd.array(imgpred, devs)
            exector = symbol.bind(devs, arg_params ,args_grad=None, grad_req="null", aux_states=aux_params)
            exector.forward(is_train=False)
            exector.outputs[0].wait_to_read()
            output = exector.outputs[0].asnumpy()
            text = ["5_o_Clock_Shadow","Arched_Eyebrows","Attractive","Bags_Under_Eyes","Bald", "Bangs","Big_Lips","Big_Nose",
            "Black_Hair","Blond_Hair","Blurry","Brown_Hair","Bushy_Eyebrows","Chubby","Double_Chin","Eyeglasses","Goatee",
            "Gray_Hair", "Heavy_Makeup","High_Cheekbones","Male","Mouth_Slightly_Open","Mustache","Narrow_Eyes","No_Beard",
            "Oval_Face","Pale_Skin","Pointy_Nose","Receding_Hairline","Rosy_Cheeks","Sideburns","Smiling","Straight_Hair",
            "Wavy_Hair","Wearing_Earrings","Wearing_Hat","Wearing_Lipstick","Wearing_Necklace","Wearing_Necktie","Young"]
            pred = np.ones(40)
            #print("attribution is:")
            cnt = 0
            for i in range(40):
                #print text[i].rjust(20)+" : \t",
                if output[0][i] < 0:
                    pred[i] = -1
                    #print "No"
                else:
                    pred[i] = 1
                    #print "Yes"
                    cv2.putText(img, text[i], (20, 50 + cnt * 30), font, 1, (255,255,0), 2, cv2.CV_AA)
                    cnt = cnt + 1
    else:
        cv2.putText(img, "No detected faces", (20,  50), font, 1, (255,255,0), 2, cv2.CV_AA)

    for f in faces:
        if f == max_face:
            cv2.rectangle(img, (f.left(), f.top()), (f.right(), f.bottom()), (0,0,255), 2)
        else:
            cv2.rectangle(img, (f.left(), f.top()), (f.right(), f.bottom()), (255,0,0), 2)

    return img

def main(args):
    symbol = lightened_moon_feature(num_classes=40, use_fuse=True)
    devs = None
    if args.gpus is not None:
        print("use gpu...")
        devs = mx.gpu()
    else:
        print("use cpu...")
        devs = mx.cpu()
    _, arg_params, aux_params = mx.model.load_checkpoint(args.model_load_prefix, args.model_load_epoch)
    detector = dlib.get_frontal_face_detector()
    face_cascade = cv2.CascadeClassifier(args.opencv)

    stream = urllib.urlopen(args.url)
    bytes = ''
    temp_count = 16384

    while True:

        #Get a new frame
        imbytes = None
        temp_bytes = stream.read(temp_count)
        if temp_bytes:
            bytes+=temp_bytes
        else:
            stream = urllib.urlopen(args.url)
            bytes = ''
            temp_count = 16384

        a = bytes.find('\xff\xd8')
        b = bytes.find('\xff\xd9')
        if a != -1 and b != -1:

            jpg = bytes[a:b + 2]
            temp_count = b + 2 - a
            bytes = bytes[b + 2:]
            imbytes = jpg

        if imbytes is None:
            continue

        frame = cv2.imdecode(np.fromstring(imbytes, dtype=np.uint8), cv2.IMREAD_COLOR)
        result = attribute(frame, detector, face_cascade, devs, symbol, arg_params, aux_params, args.size)
        cv2.imshow("demo", result)
        cv2.waitKey(10)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="predict the face attribution of one input image")
    parser.add_argument('--gpus', type=str, help='the gpus will be used, e.g "0,1,2,3"')
    parser.add_argument('--img', type=str, default='./test.jpg', help='the input img path')
    parser.add_argument('--url', type=str, default='', help='the input stream url')
    parser.add_argument('--size', type=int, default=128,
                        help='the image size of lfw aligned image, only support squre size')
    parser.add_argument('--opencv', type=str, default='../model/opencv/cascade.xml',
                        help='the opencv model path')
    parser.add_argument('--pad', type=float, nargs='+',
                                 help="pad (left,top,right,bottom) for face detection region")
    parser.add_argument('--model-load-prefix', type=str, default='../model/lightened_moon/lightened_moon_fuse',
                        help='the prefix of the model to load')
    parser.add_argument('--model-load-epoch', type=int, default=82,
                        help='load the model on an epoch using the model-load-prefix')
    args = parser.parse_args()
    logging.info(args)
    main(args)


#!/usr/bin/env python3
import sys, os
import glob
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from collections import Counter
from Bio import SeqIO
from Bio.SeqFeature import SeqFeature, FeatureLocation

# General info about the genome references and the location of the cassettes in the genomes
seq_record5 = SeqIO.read('./data/C5_refSeq.gb', "genbank")
seq_record20 = SeqIO.read('./data/C20_refSeq.gb', "genbank")
seq_record26 = SeqIO.read('./data/C26_refSeq.gb', "genbank")

genome_C5 = str(seq_record5.seq)
genome_C20 = str(seq_record20.seq)
genome_C26 = str(seq_record26.seq)

cass_par = [565510, 572124] # Present in C5 and C20
cass_3b = [780850, 787445]  # Only in C20
cass_26 = [574473, 581077]

# Annotation information to work faster with dataframes
ncbi = pd.read_csv('./data/mpn_annotation.csv', sep='\t', header=None) # Genome annotation
ncbi.columns = ['gene', 'start', 'end', 'strand']
ncbi.set_index('gene', inplace=True)
gold = pd.read_csv('./data/goldsets.csv', sep='\t')  # Gold set

posN = []
posE = []
for gene, cat in zip(gold['gene'], gold['class']):
    if cat=='E':
        posE += range(ncbi.loc[gene][0], ncbi.loc[gene][1]+1)
    else:
        posN += range(ncbi.loc[gene][0], ncbi.loc[gene][1]+1)

# Define the different set of positions in the genome were we will perform the analysis
posN = set(posN)
posE = set(posE)
cas1 = set(range(cass_par[0], cass_par[1]+1))
cas2 = set(range(cass_3b[0], cass_3b[1]+1))
cas3 = set(range(cass_26[0], cass_26[1]+1))
cas12 = cas1.union(cas2)
geno = set(range(1, 816395)).difference(cas1.union(cas2))

### SEVERAL FUNCTIONS TO PARSE THE DATAFRAMES
def reduce_mem_usage(props):
    start_mem_usg = props.memory_usage().sum() / 1024**2
    print("Memory usage of properties dataframe is :",start_mem_usg," MB")
    NAlist = [] # Keeps track of columns that have missing values filled in.
    for col in props.columns:
        if props[col].dtype != object:  # Exclude strings

            # Print current column type
            print("******************************")
            print("Column: ",col)
            print("dtype before: ",props[col].dtype)

            # make variables for Int, max and min
            IsInt = False
            mx = props[col].max()
            mn = props[col].min()

            # Integer does not support NA, therefore, NA needs to be filled
            if not np.isfinite(props[col]).all():
                NAlist.append(col)
                props[col].fillna(mn-1,inplace=True)

            # test if column can be converted to an integer
            asint = props[col].fillna(0).astype(np.int64)
            result = (props[col] - asint)
            result = result.sum()
            if result > -0.01 and result < 0.01:
                IsInt = True

            # Make Integer/unsigned Integer datatypes
            if IsInt:
                if mn >= 0:
                    if mx < 255:
                        props[col] = props[col].astype(np.uint8)
                    elif mx < 65535:
                        props[col] = props[col].astype(np.uint16)
                    elif mx < 4294967295:
                        props[col] = props[col].astype(np.uint32)
                    else:
                        props[col] = props[col].astype(np.uint64)
                else:
                    if mn > np.iinfo(np.int8).min and mx < np.iinfo(np.int8).max:
                        props[col] = props[col].astype(np.int8)
                    elif mn > np.iinfo(np.int16).min and mx < np.iinfo(np.int16).max:
                        props[col] = props[col].astype(np.int16)
                    elif mn > np.iinfo(np.int32).min and mx < np.iinfo(np.int32).max:
                        props[col] = props[col].astype(np.int32)
                    elif mn > np.iinfo(np.int64).min and mx < np.iinfo(np.int64).max:
                        props[col] = props[col].astype(np.int64)

            # Make float datatypes 32 bit
            else:
                props[col] = props[col].astype(np.float32)

            # Print new column type
            print("dtype after: ",props[col].dtype)
            print("******************************")

    # Print final result
    print("___MEMORY USAGE AFTER COMPLETION:___")
    mem_usg = props.memory_usage().sum() / 1024**2
    print("Memory usage is: ",mem_usg," MB")
    print("This is ",100*mem_usg/start_mem_usg,"% of the initial size")
    return props, NAlist

def counthash(fil):
    c=0
    with open(fil, 'r') as fi:
        for line in fi:
            if line.startswith('#'):
                c+=1
            else:
                return c-1

def prok(effect_col):
    """ Remove splice and intron signal events, not meaningful for prokaryotes"""
    return ','.join([ii for ii in effect_col.split(',') if 'splice' not in ii and 'intron' not in ii])

def correspondance(genome, positions):
    """
    Return variants mapped within <sten> which is a list in format [start, end]
    <genome> defines the record to use as reference
    """
    d = {}
    c = 0
    if genome==5:
        seq_record = seq_record5
    elif genome==20:
        seq_record = seq_record20
    else:
        seq_record = seq_record26
    for i in seq_record.features:
        if i.location.start in positions and i.type not in ['Polymorphism', 'misc_feature', 'Site', 'repeat_region', 'gene']:
            d[c] = i
            c+=1
    return d

def map_cassettes():
    pos_cas1 = {i:'intergenic_cassette' for i in cas1}
    for v in correspondance(5, positions=cas1).values():
        if v.type!='CDS':
            name = 'regulator '+v.qualifiers['label'][0]
        else:
            name = v.qualifiers['label'][0]
        for i in range(int(v.location.start), int(v.location.end)+1):
            if pos_cas1[i][0]!='i':
                pos_cas1[i]+= ' | '+name
            else:
                pos_cas1[i] = name
    pos_cas2 = {i:'intergenic_cas1' if i<700000 else 'intergenic_cas2' for i in cas12}
    for v in correspondance(20, positions=cas12).values():
        if v.type!='CDS':
            name = 'regulator '+v.qualifiers['label'][0]
        else:
            name = v.qualifiers['label'][0]
        assigned = False
        for i in range(int(v.location.start), int(v.location.end)+1):
            if not assigned:
                if i<700000:
                    name = 'cas1 '+name
                else:
                    name = 'cas2 '+name
                assigned = True
            if pos_cas2[i][0]!='i':
                pos_cas2[i]+= ' | '+name
            else:
                pos_cas2[i] = name
    pos_cas3 = {i:'intergenic_cassette' for i in cas3}
    for v in correspondance(26, positions=cas3).values():
        if v.type!='CDS':
            name = 'regulator '+v.qualifiers['label'][0]
        else:
            name = v.qualifiers['label'][0]
        for i in range(int(v.location.start), int(v.location.end)+1):
            if pos_cas3[i][0]!='i':
                pos_cas3[i]+= ' | '+name
            else:
                pos_cas3[i] = name
    return {5:pos_cas1, 20:pos_cas2, 26:pos_cas3}

maps = map_cassettes()
gold_hash = {k:v for k, v in zip(list(gold['gene']), list(gold['class']))}

def process_effect(effect, position, genome=5):
    eff = str(effect).split(',')
    rs = {}
    for e in eff:
        string = e.split('|')
        if len(string)>5:
            typ, impact, gene = string[1], string[2], string[4]
            # Only keep gene variant
            if 'up' in string[1] or 'down' in string[1]:
                ann = 'intergenic'
                mut = string[-7]
            else:
                if gene in gold_hash:
                    if gold_hash[gene]=='E':
                        ann = 'essential'
                    else:
                        ann = 'non-essential'
                else:
                    ann = 'gene'
                mut = string[-6]
            if position in maps[genome]:
                ann = 'cassette'
                gene = maps[genome][position]
            if string[0] not in rs:
                rs[string[0]] = [typ, impact, gene, mut, ann]
    return rs

def simplify_effect(df, genome=5):
    """ Simplifies the field effect to be human readable """
    c = 1
    rs = {}

    for i, j in df.iterrows():
        try:
            effd = process_effect(str(j.EFF), j.POS, genome)
        except:
            print(i, j)
        for alt, altn in zip(str(j.ALT).split(','), str(j.ALTN).split(',')):
            try:
                rs[c] = [j.SAMPLE, j.P, int(j.POS), j.QUAL, j.TOT, j.REFN, int(altn), int(altn)/j.REFN*100, str(j.REF), alt]+effd[alt]
                c+=1
            except:
                pass
    new_df = pd.DataFrame.from_dict(rs, orient='index')
    new_df.columns = ['SAMPLE', 'PASS', 'POS', 'QUAL', 'TOT', 'REFN', 'ALTN', 'FRAC', 'REF', 'ALT', 'EFF', 'IMPACT', 'AFF', 'MUT', 'ANN_TYPE']
    return new_df

def orderdf(df, ordered_classes, col='SAMPLE'):
    """ To order by sample"""
    df_list = []
    for i in ordered_classes:
        df_list.append(df[df[col]==i])
    ordered_df = pd.concat(df_list)
    return ordered_df.reset_index(drop=True)

def parse_variations(gre, order, genome=5, fname='./results/data1.pickle'):
    """ Function to parse the variations removing the header information """
    if os.path.isfile(fname):
        print('Loading data for genome {}'.format(genome))
        return pd.read_pickle(fname)
    else:
        print('Processing data for genome {}'.format(genome))
        snpcalls = []
        for fil in glob.glob(gre):
            if 'filter' in gre:
                ide = fil.split('/')[2].split('.')[0]
            else:
                ide = fil.split('/')[1]

            df = pd.read_csv(fil, header=counthash(fil), sep='\t')
            df['TOT'] = [int(i.split(':')[1]) for i in df[ide]]
            df['REFN'] = [int(i.split(':')[3]) for i in df[ide]]

            df['ALTN'] = (df['TOT']-df['REFN'])
            df['ALTS'] = [i.split(':')[5] for i in df[ide]]
            df['FRAC'] = (df['ALTN'])/df['TOT']*100

            df['SAMPLE'] = [ide for i in df[ide]]
            df['P'] = [int(i.split('_')[0].replace('p', '').replace('3IPTG', '18')) for i in df['SAMPLE']]
            df['EFF'] = [prok(info.split('ANN=')[1]) for info in df.INFO]
            df = df[['POS', 'REF', 'ALT', 'QUAL', 'TOT', 'REFN', 'ALTN', 'ALTS', 'FRAC', 'SAMPLE', 'P', 'EFF']]
            snpcalls.append(df)
        snpcalls = pd.concat(snpcalls).sort_values(['P', 'SAMPLE'])
        orderedf = orderdf(snpcalls, order)
        print('Cleaning data for genome {}'.format(genome))
        cleanedf, nalist = reduce_mem_usage(simplify_effect(orderedf, genome=genome))
        print('Saving data for genome {}'.format(genome))
        cleanedf.to_pickle(fname)
        return cleanedf

#### FOR RUNNING DN/DS ANALYSES
# Adaptation of https://github.com/adelq/dnds/blob/master/dnds.py to mycoplasma pneumoniae and vcf files

### This program is written for dN/dS ratio calculation based on VCF and GenBank files.
from Bio import SeqIO
from Bio.Seq import Seq
from math import log

codons = {"TTT": "F", "TTC": "F", "TTA": "L", "TTG": "L", "TCT": "S", "TCC": "S", "TCA": "S", "TCG": "S",
          "TAT": "Y", "TAC": "Y", "TAA": "*", "TAG": "*", "TGT": "C", "TGC": "C", "TGA": "W", "TGG": "W",
          "CTT": "L", "CTC": "L", "CTA": "L", "CTG": "L", "CCT": "P", "CCC": "P", "CCA": "P", "CCG": "P",
          "CAT": "H", "CAC": "H", "CAA": "Q", "CAG": "Q", "CGT": "R", "CGC": "R", "CGA": "R", "CGG": "R",
          "ATT": "I", "ATC": "I", "ATA": "I", "ATG": "M", "ACT": "T", "ACC": "T", "ACA": "T", "ACG": "T",
          "AAT": "N", "AAC": "N", "AAA": "K", "AAG": "K", "AGT": "S", "AGC": "S", "AGA": "R", "AGG": "R",
          "GTT": "V", "GTC": "V", "GTA": "V", "GTG": "V", "GCT": "A", "GCC": "A", "GCA": "A", "GCG": "A",
          "GAT": "D", "GAC": "D", "GAA": "E", "GAG": "E", "GGT": "G", "GGC": "G", "GGA": "G", "GGG": "G"}
BASES = {'A', 'G', 'T', 'C'}
seq_record5 = SeqIO.read('./data/C5_refSeq.gb', "genbank")
seq_record20 = SeqIO.read('./data/C20_refSeq.gb', "genbank")

# copy from https://github.com/adelq/dnds/commit/783e8197541aeb456e19c8ae9effa54cc16b02c0
# copy from https://github.com/adelq/dnds/blob/master/dnds.py
def split_seq(seq, n=3):
    # Returns sequence split into chunks of n characters, default is codons
    return [seq[i:i+n] for i in range(0, len(seq), n)]

def average_list(l1, l2):
    return [float(i1 + i2)/2 for i1, i2 in zip(l1, l2)]

def dna_to_protein(codon):
    # Returns single letter amino acid code for given codon
    return codons[codon]

def translate(seq):
    # Translate a DNA sequence into the 1-letter amino acid sequence
    return "".join([dna_to_protein(codon) for codon in split_seq(seq)])

def is_synonymous(codon1, codon2):
    # Returns boolean whether given codons are synonymous
    return dna_to_protein(codon1) == dna_to_protein(codon2)

def dnds_codon(codon):
    # Returns list of synonymous counts
    syn_list = []
    for i in range(len(codon)):
        base = codon[i]
        other_bases = BASES - {base}
        syn = 0
        for new_base in other_bases:
            new_codon = codon[:i] + new_base + codon[i + 1:]
            syn += int(is_synonymous(codon, new_codon))
        syn_list.append(float(syn)/3)
    return syn_list

def dnds_codon_pair(codon1, codon2):
    # Get the dN/dS for the given codon pair
    return average_list(dnds_codon(codon1), dnds_codon(codon2))

def syn_sum(seq1, seq2):
    # Get the sum of synonymous sites from two DNA sequences
    syn = 0
    codon_list1 = split_seq(seq1)
    codon_list2 = split_seq(seq2)
    for i in range(len(codon_list1)):
        codon1 = codon_list1[i]
        codon2 = codon_list2[i]
        if len(codon2)==3:
            dnds_list = dnds_codon_pair(codon1, codon2)
            syn += sum(dnds_list)
        else:
            return float(syn)
    return float(syn)

def hamming(s1, s2):
    # Return the hamming distance between 2 DNA sequences
    return sum(ch1 != ch2 for ch1, ch2 in zip(s1, s2)) + abs(len(s1) - len(s2))

def codon_subs(codon1, codon2):
    # Returns number of synonymous substitutions in provided codon pair Methodology for multiple substitutions from Dr. Swanson, UWashington https://faculty.washington.edu/wjs18/dnds.ppt
    diff = hamming(codon1, codon2)
    if diff < 1:
        return 0
    elif diff == 1:
        return int(translate(codon1) == translate(codon2))

    syn = 0
    for i in range(len(codon1)):
        base1 = codon1[i]
        base2 = codon2[i]
        if base1 != base2:
            new_codon = codon1[:i] + base2 + codon1[i + 1:]
            syn += int(is_synonymous(codon1, new_codon))
            syn += int(is_synonymous(codon2, new_codon))
    return float(syn) / diff

def substitutions(seq1, seq2):
    # Returns number of synonymous and nonsynonymous substitutions
    dna_changes = hamming(seq1, seq2)
    codon_list1 = split_seq(seq1)
    codon_list2 = split_seq(seq2)
    syn = 0
    for i in range(len(codon_list1)):
        codon1 = codon_list1[i]
        codon2 = codon_list2[i]
        if len(codon2)==3:
            syn += codon_subs(codon1, codon2)
        else:
            return [syn, dna_changes-syn]
    return [syn, dna_changes-syn]

def clean_sequence(seq):
    # Clean up provided sequence by removing whitespace.
    return seq.replace(' ', '')

def dnds(seq1, seq2):
    """Main function to calculate dN/dS between two DNA sequences per Nei &
    Gojobori 1986. This includes the per site conversion adapted from Jukes &
    Cantor 1967.
    """
    # Strip any whitespace from both strings
    seq1 = clean_sequence(seq1)
    seq2 = clean_sequence(seq2)
    # Check that both sequences have the same length
    # assert len(seq1) == len(seq2)
    # Check that sequences are codons
    # assert len(seq1) % 3 == 0
    syn_sites = syn_sum(seq1, seq2)
    non_sites = len(seq1)-syn_sites
    syn_subs, non_subs = substitutions(seq1, seq2)
    pn = float(non_subs)/non_sites
    ps = float(syn_subs)/syn_sites
    dn = -0.75 * log(1 - (4 * pn / 3))
    ds = -0.75 * log(1 - (4 * ps / 3))
    return round(float(pn), 3), round(float(ps), 3), round(float(dn), 3), round(float(ds), 3)

def snp_in_gene(file, record, pos):
    # Confirm whether an SNP in a gene
    output = None
    for records in SeqIO.parse(file, "genbank"):
        if record == records.id:
            for feature in records.features:
                if feature.type == 'CDS' and (feature.location.start <= pos and feature.location.end >= pos):
                    output = feature
                    break
    return output

def get_new_sequence(record, gene, pos):
    for feature in record.features:
        if ('locus_tag' in feature.qualifiers and gene== feature.qualifiers['locus_tag'][0]) or ('label' in feature.qualifiers and gene==feature.qualifiers['label'][0]):
            st = feature.location.start
            en = feature.location.end
            strand = feature.location.strand
            raw_seq = str(record.seq[st:en])
            new_seq = str(record.seq[st:pos[0]-1]) + pos[2] + str(record.seq[pos[0]+len(pos[1])-1:en])
            if feature.location.strand < 0:
                raw_seq = str(Seq(raw_seq).reverse_complement())
                new_seq = str(Seq(new_seq).reverse_complement())
                if gene=='LacI4':
                    raw_seq = raw_seq[:-4]
                    new_seq = new_seq[:-4]
            return raw_seq, new_seq

def evolution_of_sample(df, genome=5):
    if genome==5:
        record_seq = seq_record5
    else:
        record_seq = seq_record20

    rs = {}
    for sample in set(df['SAMPLE']):
        pn, ps, pns, dn, ds, dns = 0, 0, 0, 0, 0, 0
        subdf = df[df['SAMPLE']==sample]
        for pos, ref, alt, gene in zip(list(subdf['POS']),list(subdf['REF']), list(subdf['ALT']), list(subdf['AFF'])):
            try:
                raw_seq, new_seq = get_new_sequence(record_seq, gene, pos=[pos, ref, alt])
                newpn, newps, newdn, newds = dnds(raw_seq, new_seq)
                pn += newpn
                ps += newps
                dn += newdn
                ds += newds
            except:
                pass
        if pn == 0 and ps == 0:
            pns, dns = '/', '/'
        elif ps == 0:
            pns, dns = '+', '+'
        elif pn == 0:
            pns, dns = '-', '-'
        else:
            pns, dns = pn/ps, dn/ds
        rs[sample] = [pn, ps, pns, dn, ds, dns]
    return rs


# coding: utf-8

import socket
import struct
import time
import numpy as np

NUMBER_BYTES_FROM_UR = 1116  # For the UR3.10 version
INTEGER = 1
DOUBLE = 2
ARRAY_OF_6_DOUBLES = 3

class RobotUR(object):
    def __init__(self, tcp_host_ip, tcp_port, rtc_host_ip, rtc_port,joint_acc=1.4,joint_vel=1.05):

        # Default home joint configuration
        # self.home_joint_config = [-np.pi, -np.pi/2, np.pi/2, -np.pi/2, -np.pi/2, 0]
        self.home_joint_config = [-(180.0 / 360.0) * 2 * np.pi, -(84.2 / 360.0) * 2 * np.pi,
                                  (112.8 / 360.0) * 2 * np.pi, -(119.7 / 360.0) * 2 * np.pi,
                                  -(90.0 / 360.0) * 2 * np.pi, 0.0]

        # Default joint speed configuration
        self.joint_acc = joint_acc # Safe: 1.4
        self.joint_vel = joint_vel  # Safe: 1.05

        # Joint tolerance for blocking calls
        self.joint_tolerance = 0.01

        # Network parameters to connect to robot client
        self.tcp_host_ip = tcp_host_ip
        self.tcp_port = tcp_port

        # Network parameters to connect as real-time client to parse state data
        self.rtc_host_ip = rtc_host_ip
        self.rtc_port = rtc_port

        # The state dictionnary to collect all the info from UR robot
        self.dicoVariables = {'actualCartesianCoordinatesOfTool': (444, ARRAY_OF_6_DOUBLES),
                              'actualJointPositions': (252, ARRAY_OF_6_DOUBLES)}

        # Move robot to home pose
        #self.go_home()

    def go_home(self):
        self.move_joints(self.home_joint_config)

    def move_to(self, tool_position, tool_orientation):
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.connect((self.tcp_host_ip, self.tcp_port))
        # Build the movel command
        tcp_command = "movel(p[%f,%f,%f,%f,%f,%f],a=%f,v=%f,t=0,r=0)\n" % (
            tool_position[0], tool_position[1], tool_position[2], tool_orientation[0], tool_orientation[1],
            tool_orientation[2], self.tool_acc, self.tool_vel)
        self.tcp_socket.send(str.encode(tcp_command))
        # Block until robot reaches target tool position
        dicoState = self.get_state()
        actual_tool_pose = dicoState['actualCartesianCoordinatesOfTool']
        while not all([np.abs(actual_tool_pose[j] - tool_position[j]) < self.tool_pose_tolerance[j] for j in range(3)]):
            dicoState = self.get_state()
            actual_tool_pose = dicoState['actualCartesianCoordinatesOfTool']
        self.tcp_socket.close()

    def move_joints(self, joint_configuration):
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.connect((self.tcp_host_ip, self.tcp_port))
        # Build the movej command
        tcp_command = "movej([%f" % joint_configuration[0]
        for joint_idx in range(1, 6):
            tcp_command = tcp_command + (",%f" % joint_configuration[joint_idx])
        tcp_command = tcp_command + "],a=%f,v=%f)\n" % (self.joint_acc, self.joint_vel)
        print(tcp_command)
        self.tcp_socket.send(str.encode(tcp_command))
        # # Block until robot reaches home state
        dicoState = self.get_state()
        actual_joint_positions = dicoState['actualJointPositions']
        while not all([np.abs(actual_joint_positions[j] - joint_configuration[j]) < self.joint_tolerance for j in range(6)]):
            dicoState = self.get_state()
            actual_joint_positions = dicoState['actualJointPositions']
        self.tcp_socket.close()

    # Avec les 6 angles exprimés en degré
    def move_joints_degree(self, joint_configuration):
        self.move_joints([q*np.pi/180 for q in joint_configuration])

    def open_gripper(self, async=False):
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.connect((self.tcp_host_ip, self.tcp_port))
        tcp_command = "set_digital_out(8,False)\n"
        #tcp_command = "set_digital_out(8,False)\n"
        self.tcp_socket.send(str.encode(tcp_command))
        self.tcp_socket.close()
        if not async:
            time.sleep(1.5)

    def close_gripper(self, async=False):
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.connect((self.tcp_host_ip, self.tcp_port))
        tcp_command = "set_digital_out(8,True)\n"
        #tcp_command = "set_digital_out(8,True)\n"
        self.tcp_socket.send(str.encode(tcp_command))
        self.tcp_socket.close()
        if not async:
            time.sleep(1.5)

    def isGripperFullyClosed(self, async=False):
        if async:
            gripper_fully_closed = True
        else:
            time.sleep(1.5)
            gripper_fully_closed = self.check_grasp()
        return gripper_fully_closed

    def check_grasp(self):
        state_data = self.get_state()
        tool_analog_input2 = self.parse_tcp_state_data( 'tool_data')
        return tool_analog_input2 > 0.26

    # Get all the information from robot UR through RealTime port 30003
    # See Excel file for the list of available data
    def get_state(self):
        # Connect as real-time client to parse state data
        rtc_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        rtc_socket.connect((self.rtc_host_ip, self.
                            rtc_port))
        stateData = rtc_socket.recv(NUMBER_BYTES_FROM_UR)
        dico = {}
        for key, (offset, dataType) in self.dicoVariables.items():
            if dataType == INTEGER:
                dico[key] = self.getInteger(stateData, offset)
            if dataType == DOUBLE:
                dico[key] = self.getDouble(stateData, offset)
            if dataType == ARRAY_OF_6_DOUBLES:
                dico[key] = self.getListOf6Double(stateData, offset)
        rtc_socket.close()
        return dico

    def getDouble(self, data, pos):
        return struct.unpack('!d', data[pos:pos + 8])[0]  # d car double

    def getInteger(self, data, pos):
        return struct.unpack('!i', data[pos:pos + 4])[0]  # i car integer

    def getListOf6Double(self, data, pos):
        return struct.unpack('!dddddd', data[pos:pos + 48])  # d car double et 48 car 6 doubles * 8 octets

    def readMessage(self,s):
        msgOK = False
        while msgOK != True:
            header = s.recv(4)
            # Read package header
            data_bytes = bytearray()
            data_bytes.extend(header)
            data_length = struct.unpack("!i", data_bytes[0:4])[0]
            # Read message
            message = s.recv(data_length - 4)  # 4 is for the previously read header
            data_bytes.extend(message)
            if message[0] == 16:
                msgOK = True
        return message, data_length

    def parse_tcp_state_data(self,subpackage):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((self.tcp_host_ip, self.tcp_port))
        data_bytes,data_length = self.readMessage(s)
        s.close()
        robot_message_type = data_bytes[0]
        assert (robot_message_type == 16)
        byte_idx = 1
        # Parse sub-packages
        subpackage_types = {'joint_data': 1, 'cartesian_info': 4, 'force_mode_data': 7, 'tool_data': 2}
        while byte_idx < data_length:
            # package_length = int.from_bytes(data_bytes[byte_idx:(byte_idx+4)], byteorder='big', signed=False)
            package_length = struct.unpack("!i", data_bytes[byte_idx:(byte_idx + 4)])[0]
            byte_idx += 4
            package_idx = data_bytes[byte_idx]
            if package_idx == subpackage_types[subpackage]:
                byte_idx += 1
                break
            byte_idx += package_length - 4

        def parse_joint_data(data_bytes, byte_idx):
            actual_joint_positions = [0, 0, 0, 0, 0, 0]
            target_joint_positions = [0, 0, 0, 0, 0, 0]
            for joint_idx in range(6):
                actual_joint_positions[joint_idx] = struct.unpack('!d', data_bytes[(byte_idx + 0):(byte_idx + 8)])[0]
                target_joint_positions[joint_idx] = struct.unpack('!d', data_bytes[(byte_idx + 8):(byte_idx + 16)])[0]
                byte_idx += 41
            return actual_joint_positions

        def parse_cartesian_info(data_bytes, byte_idx):
            actual_tool_pose = [0, 0, 0, 0, 0, 0]
            for pose_value_idx in range(6):
                actual_tool_pose[pose_value_idx] = struct.unpack('!d', data_bytes[(byte_idx + 0):(byte_idx + 8)])[0]
                byte_idx += 8
            return actual_tool_pose

        def parse_tool_data(data_bytes, byte_idx):
            byte_idx += 2
            tool_analog_input2 = struct.unpack('!d', data_bytes[(byte_idx + 0):(byte_idx + 8)])[0]
            return tool_analog_input2

        parse_functions = {'joint_data': parse_joint_data, 'cartesian_info': parse_cartesian_info,
                           'tool_data': parse_tool_data}
        return parse_functions[subpackage](data_bytes, byte_idx)

#
# Programme principal
#
if __name__ == '__main__':
    IP_UR3 = '10.31.56.102'  # IP and port to robot arm as TCP client (UR5)
    IP_UR5 = '10.31.56.104'
    tcp_port = 30002
    rtc_port = 30003
    print("robot at home")
    monRobot = RobotUR(IP_UR3, tcp_port, IP_UR3, rtc_port,0.5,0.5)
    print(monRobot.parse_tcp_state_data('tool_data'))
    print(monRobot.check_grasp())
    print('on ouvre la pince')
    print(monRobot.check_grasp())
    monRobot.open_gripper()
    time.sleep(5)
    print(monRobot.check_grasp())
    print('on ferme la pince')
    monRobot.close_gripper()
    print(monRobot.check_grasp())
    print('fin du premier mouvement')
    monRobot.move_joints([-3.141593,-1.469567,1.968731,-2.089159,-1.570796,0.000000])
    print('fin du deuxième mouvement')
    monRobot.move_joints([-3.141593,-1.570796,1.570796,0.000000,1.570796,3.141593])
    print("déplacement en degré")
    monRobot.move_joints_degree([-270, -90, 0, -90, 270, 180])
    print('fin du programme de test')