#!/usr/bin/python3

"""
Last update: March 2021
Author: Luigi D'Ascenzo, PhD - The Scripps Research Institute, La Jolla (CA)
Contact info: dascenzoluigi@gmail.com
GitHub project repository: https://github.com/ldascenzo/pytheas

***DESCRIPTION***
Final step of the Pytheas in silico digest library generation. m/z values for precursor ions and all the derived
CID fragment ions are added to each target and decoy sequences. Isotopic labeling can be added via an additional
nucleotides dictionary file. Charge tables for precursor ions and fragment ions (if different from the default) can
also be provided.
More information on the elemental masses, CID fragmentation of RNA and the available options/parameters can be found
in the Digest section of Pytheas manual.

***OUTPUT***
1) Digest_[name_fasta].txt -> in silico digest library file ready to be used for matching, with all
                 the info on each precursor ion and its derived fragmentation ions
"""

import os
import sys

import numpy as np
import pandas as pd
from shutil import move

import consolidate_tools as ct
import ntpath
import platform

##########
# ELEMENT MASS DICTIONARY
##########
"""
Masses from:
(1) http://physics.nist.gov/cgi-bin/Compositions/stand_alone.pl?ele=&ascii=html&isotype=some
(2) http://www.sisweb.com/referenc/source/exactmaa.htm

Note that for each element with no specified isotope the mass of the most abundant isotope is used.
In the case of O, C, H and N the most abundant isotope is >=99% abundant.
S32 has 95% abundance
Se80 has only 49% abundance, so you could consider adding multiple isotopes if working with Se
Se80 has only 49% abundance, so you could consider adding multiple isotopes if working with Se
"""
ele_mass = {"C": 12.0000000, "H": 1.007825032, "N": 14.0030740044, "O": 15.9949146196, "S": 31.9720711744,
            "P": 30.9737619984, "Se": 79.9165218, "H2": 2.0141017781, "C13": 13.0033548351, "N15": 15.0001088989,
            "O18": 17.9991596128}

class Masses:
    def __init__(self, ion_mode, light_alphabet, heavy_alphabet, MS1_charge_tabl, MS2_charge_tabl, MS1_mzlow,
                 MS1_mzhigh, MS2_mzlow, MS2_mzhigh, CID_series, mz_consolidation, MS1_consolidation_ppm,
                 MS2_consolidation_ppm, fasta_name):
        self.ion_mode, self.nts_alphabet_light, self.nts_alphabet_heavy = ion_mode, light_alphabet, heavy_alphabet
        self.MS1_charges, self.MS2_charges = MS1_charge_tabl, MS2_charge_tabl
        self.MS1_mzlow, self.MS1_mzhigh, self.MS2_mzlow, self.MS2_mzhigh = MS1_mzlow, MS1_mzhigh, MS2_mzlow, MS2_mzhigh
        self.CID_series, self.mz_consolidation = CID_series, mz_consolidation
        self.MS1_consolidation_ppm = round(MS1_consolidation_ppm, 1)
        self.MS2_consolidation_ppm = round(MS2_consolidation_ppm, 1)
        self.fasta_name = fasta_name

    def __controls(self):
        """
        Make sure the right options with right values are assigned for the correct execution of the program
        Also check if the nucleotide dictionaries for light and heavy atoms are consistent with what used in the
        preceding steps
        """
        # Check the consistency of nucleotides alphabets between this scripts and the preceding ones for light atoms
        df_light = self.read_excel_input(self.nts_alphabet_light)

        d = DictDiffer(read_csv(), dict(zip(df_light.ID, df_light.ID_ext)))

        if len(d.changed()) != 0:
            print(
                "WARNING! One or more nts IDs from the input nts_light have different ID_ext compared "
                "to the alphabet used in the previous scripts. List of IDs with differences: {}"
                "  Please check carefully your input dictionary and rerun the script\n".format(d.changed()))

        if len(d.added()) != 0:
            print(
                "WARNING! One or more nts IDs from the input nts_light are missing compared "
                "to the alphabet used in the previous scripts. List of missing IDs: {}"
                "  Please check carefully your input dictionary and rerun the script\n".format(d.added()))

        if len(d.removed()) != 0:
            print(
                "WARNING! One or more nts IDs from the input nts_light are not present "
                "in the alphabet used in the previous scripts. List of absent IDs: {}"
                "  Please check carefully your input dictionary and rerun the script\n".format(d.removed()))

        if self.nts_alphabet_heavy:

            # Check the consistency of nucleotides alphabets between this scripts and the preceding ones for heavy
            # atoms
            df_heavy = self.read_excel_input(self.nts_alphabet_heavy)

            d = DictDiffer(read_csv(), dict(zip(df_heavy.ID, df_heavy.ID_ext)))

            if len(d.changed()) != 0:
                print(
                    "WARNING! One or more nts IDs from the input nts_heavy have different ID_ext compared "
                    "to the alphabet used in the previous scripts. List of IDs with differences: {}"
                    "  Please check carefully your input dictionary and rerun the script\n".format(d.changed()))

            if len(d.added()) != 0:
                print(
                    "WARNING! One or more nts IDs from the input nts_heavy are missing compared "
                    "to the alphabet used in the previous scripts. List of missing IDs: {}"
                    "  Please check carefully your input dictionary and rerun the script\n".format(d.added()))

            if len(d.removed()) != 0:
                print(
                    "WARNING! One or more nts IDs from the input nts_heavy are not present "
                    "in the alphabet used in the previous scripts. List of absent IDs: {}"
                    "  Please check carefully your input dictionary and rerun the script\n".format(d.removed()))

    def read_excel_input(self, input_alphabet):
        """
        Generate a dataframe with all the info on the nucleotides from the input file nts_light
        """
        # Checking that the nts_light file given in argument exists
        if not os.path.exists(input_alphabet):
            print("ERROR! File {} does not exist. Execution terminated without generating any output".format(
                input_alphabet))
            sys.exit(1)

        # Create a dataframe with info from Excel spreadsheet
        df = pd.read_excel(input_alphabet, header=12)

        # Drop rows with NaN values
        df = df[pd.notnull(df['ID'])]

        # Transform all ID values in string (so numbers can be used as one letter code for bases)
        df = df.astype({"ID": str})

        return df

    def lines_with_charge(self, input_lines):
        """
        Write new lines for each precursor ion charged status
        """
        new_lines_charge, MS1_charges = [], charge_table(self.MS1_charges)
        MS1_maxlength, MS1_minlength = max(int(s) for s in MS1_charges.keys()), min(int(s) for s in MS1_charges.keys())

        for line in input_lines:

            # The window of fragments length is defined by the intervals given in the charge table input files
            length = len(line.split()[4])
            if MS1_minlength <= length <= MS1_maxlength:

                for ele in MS1_charges[str(len(line.split()[4]))]:
                    new_lines_charge.append(add_charge(line, self.ion_mode + str(ele)))

        return new_lines_charge

    def lines_final(self, input_lines, mzlow, mzhigh):
        """
        Prepare the final output lines
        """
        nts_light = nts_mass(self.read_excel_input(self.nts_alphabet_light))[0]
        output_lines = []
        if self.nts_alphabet_heavy:
            nts_heavy = nts_mass(self.read_excel_input(self.nts_alphabet_heavy))[0]
        else:
            nts_heavy = nts_light

        for line in input_lines:
            h_mass, l_mass = 0, 0

            # Add the basic masses of the nucleotides (neutral state) present in the fragment
            for a in line.split()[5]:
                l_mass, h_mass = l_mass + nts_light[a], h_mass + nts_heavy[a]

            # Add masses for the 5'end of the fragments
            if line.split()[8] == "OH":
                l_mass, h_mass = l_mass + ele_mass["O"] + ele_mass["H"], h_mass + ele_mass["O"] + ele_mass["H"]

            elif line.split()[8] == "P":
                l_mass, h_mass = l_mass + ele_mass["P"] + ele_mass["O"] * 4 + ele_mass["H"] * 2, h_mass + ele_mass[
                    "P"] + ele_mass["O"] * 4 + ele_mass["H"] * 2

            else:
                l_mass, h_mass = l_mass + ele_mass["P"] + ele_mass["O"] * 3 + ele_mass["H"], h_mass + ele_mass[
                    "P"] + ele_mass["O"] * 3 + ele_mass["H"]

            # Add masses for the 3' end of the fragments
            # If fragment has 3' OH subtract from its mass -PO3 (all fragments are considered neutral by default)
            if line.split()[7] == "OH":
                l_mass, h_mass = l_mass - ele_mass["P"] - ele_mass["O"] * 3, h_mass - ele_mass["P"] - ele_mass["O"] * 3

            # If fragment has 3' P adds a -H to the mass (all fragments are considered neutral by default)
            elif line.split()[7] == "P":
                l_mass, h_mass = l_mass + ele_mass["H"], h_mass + ele_mass["H"]

            # If fragment has 3'>P subtract from its mass -OH (all fragments are considered neutral by default)
            elif line.split()[7] == "cP":
                l_mass, h_mass = l_mass - ele_mass["O"] - ele_mass["H"], h_mass - ele_mass["O"] - ele_mass["H"]

            # Add the right amount of hydrogen atoms depending on the fragment charge (starting charge is 0 by default)
            if self.ion_mode == "+":
                l_mass, h_mass = l_mass + ele_mass["H"] * (int(line.split()[3][1])), h_mass + ele_mass["H"] * (
                    int(line.split()[3][1]))

            # Subtract the right amount of hydrogen atoms depending on the fragment charge
            # (starting charge is 0 by default)
            if self.ion_mode == "-":
                l_mass, h_mass = l_mass - ele_mass["H"] * (int(line.split()[3][1])), h_mass - ele_mass["H"] * (
                    int(line.split()[3][1]))

            # Calculate the m/z for each fragment
            l_mass, h_mass = l_mass / int(line.split()[3][1:]), h_mass / int(line.split()[3][1:])

            # Add the line to the list only if its light mass is > mzlow and < mzhigh
            if mzlow < l_mass < mzhigh:
                if self.nts_alphabet_heavy:
                    output_lines.append(add_masses(line, l_mass, h_mass))

                else:
                    output_lines.append(add_masses(line, l_mass, '-'))

        return output_lines

    def header_info_MS1(self, input_file):
        """
        Write the header section for the MS1 output file
        """
        header_lines = []

        if self.ion_mode == "+":
            header_lines.append("#ION_MODE positive\n")

        elif self.ion_mode == "-":
            header_lines.append("#ION_MODE negative\n")

        header_lines.append("#MZLOW_MS1 " + str(self.MS1_mzlow) + "\n")
        header_lines.append("#MZHIGH_MS1 " + str(self.MS1_mzhigh) + "\n")
        header_lines.append("#NTS_LIGHT " + str(self.nts_alphabet_light) + "\n")

        if self.nts_alphabet_heavy:
            header_lines.append("#NTS_HEAVY " + str(self.nts_alphabet_heavy) + "\n")

        for line in input_file:

            if line[0] == "#":
                header_lines.append(line)

        header_lines.append(
            "m/z_light m/z_heavy molecule residue_start residue_end charge miss sequence sequence_mod 3'end 5'end "
            "num_copy molecule_location\n")

        return header_lines

    def fragment_MS2_masses(self, line, MS2_charge_dic, nts_dic_heavy):
        """
        Calculate the masses of CID derived fragment ions
        """
        nts_dic_light = nts_mass(self.read_excel_input(self.nts_alphabet_light))[0]
        output_masses, charge_prec, mass_prec = [], int(line.split()[5][1:]), np.float64(line.split()[0])
        chem3prime, chem5prime = line.split()[10], line.split()[9]

        # Correction to the charges for charged losses
        charged_loss = charge_prec - 1

        # Mass correction to apply for positive/negative mode in case of charged losses
        if self.ion_mode == '-':
            freebase_correction = 0
            # In case of a charged loss in negative mode, the mass needs no adjustments
            charge_correction = 0
        else:
            # Correct the m/z of free bases in positive mode adding two protons since they are charged +1
            freebase_correction = 2 * ele_mass["H"]
            # Subtract two protons in positive mode when a positive charge is lost due to charged loss
            charge_correction = - 2 * ele_mass['H']

        # Add the ion M-H2O as first entry of the MS2 ions list
        MH2O_mz = mass_prec + (self.dic_M_x_series['M-H2O']) / charge_prec

        # Keep the ion if its m/z is in the given interval for MS2 ions
        if self.MS2_mzhigh >= MH2O_mz >= self.MS2_mzlow:
            output_masses.append('M-H2O({}):{}'.format(line.split()[5], str(round_masses(MH2O_mz))))

        if chem3prime == 'P' or chem5prime == 'P':
            # Add the neutral phosphate loss ions M-PO3H and M-PO4H3 to the MS2 list if 3' end is P or OH
            MPO3H_mz = mass_prec + (self.dic_M_x_series['M-PO3H']) / charge_prec

            # Keep the ion if its m/z is in the given interval for MS2 ions
            if self.MS2_mzhigh >= MPO3H_mz >= self.MS2_mzlow:
                output_masses.append('M-P({}):{}'.format(line.split()[5], str(round_masses(MPO3H_mz))))

            MPO4H3_mz = mass_prec + (self.dic_M_x_series['M-PO4H3']) / charge_prec

            # Keep the ion if its m/z is in the given interval for MS2 ions
            if self.MS2_mzhigh >= MPO4H3_mz >= self.MS2_mzlow:
                output_masses.append('M-H2O-P({}):{}'.format(line.split()[5], str(round_masses(MPO4H3_mz))))

            # Add the charged phosphate loss ions M-PO3- and M-PO4H2- to the MS2 list
            # if 3' end is P and if the precursor has a charge greater than 1
            if charge_prec > 1:
                MPO3_mz = (mass_prec * charge_prec + self.dic_M_x_series['M-PO3'] + charge_correction) / charged_loss

                # Keep the ion if its m/z is in the given interval for MS2 ions
                if self.MS2_mzhigh >= MPO3_mz >= self.MS2_mzlow:
                    output_masses.append(
                        'M-P({}):{}'.format(line.split()[5][0] + str(charged_loss), str(round_masses(MPO3_mz))))

                MPO4H2_mz = (mass_prec * charge_prec + self.dic_M_x_series['M-PO4H2'] +
                             charge_correction) / charged_loss

                # Keep the ion if its m/z is in the given interval for MS2 ions
                if self.MS2_mzhigh >= MPO4H2_mz >= self.MS2_mzlow:
                    output_masses.append(
                        'M-H2O-P({}):{}'.format(line.split()[5][0] + str(charged_loss), str(round_masses(MPO4H2_mz))))

        elif chem3prime == 'cP':
            # Add the charged phosphate loss ions M-PO3- to the MS2 list
            # if 3' end is cP and if the precursor has a charge greater than 1
            if charge_prec > 1:
                MPO3_mz = (mass_prec * charge_prec + self.dic_M_x_series['M-PO3'] + charge_correction) / charged_loss
                # Keep the ion if its m/z is in the given interval for MS2 ions
                if self.MS2_mzhigh >= MPO3_mz >= self.MS2_mzlow:
                    output_masses.append(
                        'M-P({}):{}'.format(line.split()[5][0] + str(charged_loss), str(round_masses(MPO3_mz))))

        # Add the free bases B- ions and M-B-/M-BH/M-P-B to the MS2 list
        unique_seq = set(list(line.split()[7]))
        last_nt = list(line.split()[7])[-1]

        for b in unique_seq:

            if line.split()[1] == "light":
                output_string = '{}({}1):{}'.format(b, self.ion_mode, str(round_masses(self.nts_dic_onlyB_light[b] +
                                                                                freebase_correction)))
                if output_string not in output_masses:
                    output_masses.append(output_string)

                # The charged base loss ion M-B- is generated only for ions with charge > 1
                if charge_prec > 1:
                    MB_mz = (np.float64(line.split()[0]) * charge_prec + charge_correction - self.nts_dic_onlyB_light[
                        b]) / charged_loss

                    # Keep the ion if its m/z is in the given interval for MS2 ions
                    if self.MS2_mzhigh >= MB_mz >= self.MS2_mzlow:
                        output_string = 'M-{}({}):{}'.format(b, line.split()[5][0] + str(charged_loss),
                                                             str(round_masses(MB_mz)))
                        if output_string not in output_masses:
                            output_masses.append(output_string)

                    if chem3prime == 'P' or chem5prime == 'P':
                        # Add the charged simultaneous loss of both P and B
                        MBP_mz = (np.float64(line.split()[0]) * charge_prec + charge_correction -
                                  self.nts_dic_onlyB_light[b] + self.dic_M_x_series['M-PO3H']) / charged_loss

                        # Keep the ion if its m/z is in the given interval for MS2 ions
                        if self.MS2_mzhigh >= MBP_mz >= self.MS2_mzlow:
                            output_string = 'M-P-{}({}):{}'.format(b, line.split()[5][0] + str(charged_loss),
                                                                   str(round_masses(MBP_mz)))
                            if output_string not in output_masses:
                                output_masses.append(output_string)

                    elif chem3prime == 'cP':
                        # Add the charged simultaneous loss of both P and the 3' base B
                        MBP_mz = (np.float64(line.split()[0]) * charge_prec + charge_correction -
                                  self.nts_dic_onlyB_light[last_nt] + self.dic_M_x_series['M-PO3H']) / charged_loss

                        # Keep the ion if its m/z is in the given interval for MS2 ions
                        if self.MS2_mzhigh >= MBP_mz >= self.MS2_mzlow:
                            output_string = 'M-P-{}({}):{}'.format(last_nt, line.split()[5][0] + str(charged_loss),
                                                                   str(round_masses(MBP_mz)))
                            if output_string not in output_masses:
                                output_masses.append(output_string)

                        # Add the special case of base double loss (with one being the 3'end base) for cP 3' chemistry
                        if b in line.split()[7][:-1]:
                            MBB_mz = (np.float64(line.split()[0]) * charge_prec + charge_correction -
                                      self.nts_dic_onlyB_light[b] - self.nts_dic_onlyB_light[last_nt]) / charged_loss

                            # Keep the ion if its m/z is in the given interval for MS2 ions
                            if self.MS2_mzhigh >= MBB_mz >= self.MS2_mzlow:
                                output_string = 'M-{}-{}({}):{}'.format(b, last_nt, line.split()[5][0] +
                                                                        str(charged_loss), str(round_masses(MBB_mz)))
                                if output_string not in output_masses:
                                    output_masses.append(output_string)

                # The neutral base loss ion M-B is generated for all the precursor ions
                MBH_mz = np.float64(line.split()[0]) - (self.nts_dic_onlyB_light[b] + ele_mass["H"]) / charge_prec

                # Keep the ion if its m/z is in the given interval for MS2 ions
                if self.MS2_mzhigh >= MBH_mz >= self.MS2_mzlow:
                    output_masses.append('M-{}({}):{}'.format(b, line.split()[5], str(round_masses(MBH_mz))))

                if chem3prime == 'P' or chem5prime == 'P':
                    # Add the neutral simultaneous loss of both P and B
                    MBHP_mz = np.float64(line.split()[0]) - (
                            self.nts_dic_onlyB_light[b] + ele_mass["H"] - self.dic_M_x_series['M-PO3H']) / charge_prec

                    # Keep the ion if its m/z is in the given interval for MS2 ions
                    if self.MS2_mzhigh >= MBHP_mz >= self.MS2_mzlow:
                        output_string = 'M-P-{}({}):{}'.format(b, line.split()[5], str(round_masses(MBHP_mz)))
                        if output_string not in output_masses:
                            output_masses.append(output_string)

                elif chem3prime == 'cP':
                    # Add the special case of base double loss (with one being the 3'end base) for cP 3' chemistry
                    if b in line.split()[7][:-1]:
                        MBHBH_mz = np.float64(line.split()[0]) - (self.nts_dic_onlyB_light[b] + ele_mass["H"] +
                                                                  self.nts_dic_onlyB_light[last_nt]) / charge_prec

                        # Keep the ion if its m/z is in the given interval for MS2 ions
                        if self.MS2_mzhigh >= MBHBH_mz >= self.MS2_mzlow:
                            output_string = 'M-{}-{}({}):{}'.format(b, last_nt, line.split()[5],
                                                                    str(round_masses(MBHBH_mz)))
                            if output_string not in output_masses:
                                output_masses.append(output_string)

            elif line.split()[1] == "heavy":
                output_masses.append('{}({}1):{}'.format(b, self.ion_mode, str(round_masses(self.nts_dic_onlyB_heavy[b])
                                                                               + freebase_correction)))

                # The charged base loss ion M-B- is generated only for ions with charge > 1
                if charge_prec > 1:
                    MBh_mz = (np.float64(line.split()[0]) * charge_prec + charge_correction - self.nts_dic_onlyB_heavy[
                        b]) / charged_loss

                    # Keep the ion if its m/z is in the given interval for MS2 ions
                    if self.MS2_mzhigh >= MBh_mz >= self.MS2_mzlow:
                        output_string = 'M-{}({}):{}'.format(b, line.split()[5][0] + str(charged_loss),
                                                             str(round_masses(MBh_mz)))
                        if output_string not in output_masses:
                            output_masses.append(output_string)

                    if chem3prime == 'P' or chem5prime == 'P':
                        # Add the charged simultaneous loss of both P and B
                        MBPh_mz = (np.float64(line.split()[0]) * charge_prec + charge_correction -
                                   self.nts_dic_onlyB_heavy[b] + self.dic_M_x_series['M-PO3H']) / charged_loss

                        # Keep the ion if its m/z is in the given interval for MS2 ions
                        if self.MS2_mzhigh >= MBPh_mz >= self.MS2_mzlow:
                            output_string = 'M-P-{}({}):{}'.format(b, line.split()[5][0] + str(charged_loss),
                                                                   str(round_masses(MBPh_mz)))
                            if output_string not in output_masses:
                                output_masses.append(output_string)

                    elif chem3prime == 'cP':
                        # Add the charged simultaneous loss of both P and the 3' base B
                        MBPh_mz = (np.float64(line.split()[0]) * charge_prec + charge_correction -
                                   self.nts_dic_onlyB_heavy[last_nt] + self.dic_M_x_series['M-PO3H']) / charged_loss

                        # Keep the ion if its m/z is in the given interval for MS2 ions
                        if self.MS2_mzhigh >= MBPh_mz >= self.MS2_mzlow:
                            output_string = 'M-P-{}({}):{}'.format(last_nt, line.split()[5][0] + str(charged_loss),
                                                                   str(round_masses(MBPh_mz)))
                            if output_string not in output_masses:
                                output_masses.append(output_string)

                        # Add the special case of base double loss (with one being the 3'end base) for cP 3' chemistry
                        if b in line.split()[7][:-1]:
                            MBBh_mz = (np.float64(line.split()[0]) * charge_prec + charge_correction -
                                       self.nts_dic_onlyB_heavy[b] - self.nts_dic_onlyB_heavy[last_nt]) / charged_loss

                            # Keep the ion if its m/z is in the given interval for MS2 ions
                            if self.MS2_mzhigh >= MBBh_mz >= self.MS2_mzlow:
                                output_string = 'M-{}-{}({}):{}'.format(b, last_nt, line.split()[5][0] +
                                                                        str(charged_loss), str(round_masses(MBBh_mz)))
                                if output_string not in output_masses:
                                    output_masses.append(output_string)

                MBHh_mz = np.float64(line.split()[0]) - (
                        self.nts_dic_onlyB_heavy[b] + ele_mass["H"]) / charge_prec

                # Keep the ion if its m/z is in the given interval for MS2 ions
                if self.MS2_mzhigh >= MBHh_mz >= self.MS2_mzlow:
                    output_string = 'M-{}({}):{}'.format(b, line.split()[5], str(round_masses(MBHh_mz)))
                    if output_string not in output_masses:
                        output_masses.append(output_string)

                if chem3prime == 'P' or chem5prime == 'P':
                    # Add the neutral simultaneous loss of both P and B
                    MBHPh_mz = np.float64(line.split()[0]) - (self.nts_dic_onlyB_heavy[b] + ele_mass["H"] -
                                                              self.dic_M_x_series['M-PO3H']) / charge_prec

                    # Keep the ion if its m/z is in the given interval for MS2 ions
                    if self.MS2_mzhigh >= MBHPh_mz >= self.MS2_mzlow:
                        output_string = 'M-P-{}({}):{}'.format(b, line.split()[5], str(round_masses(MBHPh_mz)))
                        if output_string not in output_masses:
                            output_masses.append(output_string)

                elif chem3prime == 'cP':
                    # Add the special case of base double loss (with one being the 3'end base) for cP 3' chemistry
                    if b in line.split()[7][:-1]:
                        MBHBHh_mz = np.float64(line.split()[0]) - (self.nts_dic_onlyB_heavy[b] + ele_mass["H"] +
                                                                   self.nts_dic_onlyB_heavy[last_nt]) / charge_prec

                        # Keep the ion if its m/z is in the given interval for MS2 ions
                        if self.MS2_mzhigh >= MBHBHh_mz >= self.MS2_mzlow:
                            output_string = 'M-{}-{}({}):{}'.format(b, last_nt, line.split()[5],
                                                                    str(round_masses(MBHBHh_mz)))
                            if output_string not in output_masses:
                                output_masses.append(output_string)

        # Loop through all the CID series to output the sequence-defining fragment ion masses
        for ion in list(set(self.CID_series)):
            sequence = list(line.split()[7])

            # If a ion from a,a-B, b,c or d series has to be calculated, the fragment sequence is read from
            # first to last nucleotide
            if ion == "a" or ion == "a-B" or ion == "b" or ion == "c" or ion == "d":
                sequence = sequence

            # If a ion different from a,a-B, b,c or d series has to be calculated, the fragment sequence is read from
            # last to first nucleotide
            else:
                sequence = list(reversed(sequence))

            for i in range(1, len(sequence)):

                # Keep only MS2 fragments with length within the specified input (charges_MS2)
                if str(i) in MS2_charge_dic.keys():

                    mass_fragment, mass_fragment_3OH = 0, 0

                    # Calculate the mass of the fragment up to the last nucleotide (standard masses for neutral nts)
                    for nt in sequence[:i - 1]:

                        # Add the masses for light nts
                        if line.split()[1] == "light":
                            mass_fragment += nts_dic_light[nt]

                        # Add the masses for isotopically labeled heavy nts
                        else:
                            mass_fragment += nts_dic_heavy[nt]

                    # Add/remove the proper amount of mass depending on the CID ion type for light nts
                    if line.split()[1] == "light":
                        # Dealing with the special case of 3'OH series
                        if ion == "y-P" or ion == "z-P":
                            mass_fragment += nts_dic_light[sequence[i - 1]] + self.dic_CID_series_masses[ion[0]]

                        else:
                            mass_fragment += nts_dic_light[sequence[i - 1]] + self.dic_CID_series_masses[ion]

                        # If also the base is lost, remove the mass corresponding to the base
                        # (stored in nts_dic_onlyB_light)
                        if ion == "a-B":
                            mass_fragment -= self.nts_dic_onlyB_light[sequence[i - 1]]

                    # Add/remove the proper amount of mass depending on the CID ion type for heavy nts
                    else:

                        # Dealing with the special case of 3'OH series
                        if ion == "y-P" or ion == "z-P":
                            mass_fragment += nts_dic_heavy[sequence[i - 1]] + self.dic_CID_series_masses[ion[0]]

                        else:
                            mass_fragment += nts_dic_heavy[sequence[i - 1]] + self.dic_CID_series_masses[ion]

                        # If also the base is lost, remove the mass corresponding to the base
                        # (stored in nts_dic_onlyB_heavy)
                        if ion == "a-B":
                            mass_fragment -= self.nts_dic_onlyB_heavy[sequence[i - 1]]

                    if ion == "a" or ion == "a-B" or ion == "b" or ion == "c" or ion == "d":

                        # Edit the mass of the ion based on the 5' end
                        # Add the 5'-OH end mass to the fragment
                        if chem5prime == "OH":
                            mass_fragment += ele_mass["O"] + ele_mass["H"]

                        # Add the 5'-P end mass to the fragment
                        if chem5prime == "P":
                            mass_fragment += ele_mass["O"] * 4 + ele_mass["H"] * 2 + ele_mass["P"]

                    else:

                        # Add the 3'-OH end mass to the fragment
                        if chem3prime == "OH":
                            mass_fragment = mass_fragment - ele_mass["P"] - ele_mass["O"] * 3

                        # Add the 3'-P end mass to the fragment
                        elif chem3prime == "P":
                            mass_fragment += ele_mass["H"]

                            # Add 3'OH-y and 3'-OH-z ions in case the 3' end is P
                            if ion == "y-P" or ion == "z-P":
                                mass_fragment_3OH = mass_fragment + self.dic_CID_series_masses[ion]

                        # Add the 3'>P end mass to the fragment
                        elif chem3prime == "cP":
                            mass_fragment = mass_fragment - ele_mass["O"] - ele_mass["H"]

                    # Calculate m/z for MS2 ions with charges specified in the charge table input
                    for char in MS2_charge_dic[str(i)]:

                        # Sanity check to exclude MS2 ions with charges larger than the precursor
                        if int(char) <= charge_prec:

                            # Add the right amount of hydrogen atoms depending on the fragment charge
                            # (starting charge is 0 by default)
                            if self.ion_mode == "+":
                                mass_charged = mass_fragment + ele_mass["H"] * int(char)

                                if mass_fragment_3OH:
                                    mass_charged_3OH = mass_fragment_3OH + ele_mass["H"] * int(char)

                            # Remove the right amount of hydrogen atoms depending on the fragment charge
                            # (starting charge is 0 by default)
                            elif self.ion_mode == "-":
                                mass_charged = mass_fragment - ele_mass["H"] * int(char)

                                if mass_fragment_3OH:
                                    mass_charged_3OH = mass_fragment_3OH - ele_mass["H"] * int(char)

                            mz_fragment = mass_charged / int(char)

                            if mass_fragment_3OH:
                                mz_fragment_3OH = mass_charged_3OH / int(char)

                            else:
                                mz_fragment_3OH = 0

                            # Keep the ion if its m/z is in the given interval for MS2 ions
                            if self.MS2_mzhigh >= mz_fragment >= self.MS2_mzlow and \
                                    ion != "y-P" and ion != "z-P":

                                # Output the a-B series with the number after the a
                                if ion != 'a-B':
                                    output_masses.append(ion + str(i) + "(" + self.ion_mode + str(char) + "):" + str(
                                        round_masses(mz_fragment)))

                                else:
                                    output_masses.append(
                                        ion[0] + str(i) + "-B(" + self.ion_mode + str(char) + "):" + str(
                                            round_masses(mz_fragment)))

                            if self.MS2_mzhigh >= mz_fragment_3OH >= self.MS2_mzlow:

                                # Output the a-B series with the number after the a
                                if ion != 'a-B' and ion != 'y-P' and ion != 'z-P':
                                    output_masses.append(ion + str(i) + "(" + self.ion_mode + str(char) + "):" + str(
                                        round_masses(mz_fragment_3OH)))

                                else:
                                    output_masses.append(
                                        ion[0] + str(i) + ion[1:] + "(" + self.ion_mode + str(char) + "):" + str(
                                            round_masses(mz_fragment_3OH)))

        return output_masses

    def final_lines_MS2(self, lines_MS2):
        """
        Add the m/z of fragments ions to the final lines for the output
        """
        output_lines = []

        # Define two variables to keep track of the number of unique sequences of targets and decoys
        global tot_targets
        tot_targets = 0
        global tot_decoys
        tot_decoys = 0

        if self.nts_alphabet_heavy:
            nts_dic_heavy = nts_mass(self.read_excel_input(self.nts_alphabet_heavy))[0]
        else:
            nts_dic_heavy = {}

        for line in lines_MS2:
            output_lines.append(
                " ".join(line.split() + self.fragment_MS2_masses(line, self.MS2_charge_table, nts_dic_heavy) +
                         list("\n")))

            # Add the targets and decoys to the respective counters
            split = line.split()
            if 'decoy' in split[2]:
                tot_decoys += 1

            else:
                tot_targets += 1

        return output_lines

    def header_info_MS2(self, input_file):
        """
        Create the header for the output file
        """
        header_lines = []

        for line in input_file:

            if line[0] == "#" and "MIN_LENGTH" not in line:
                header_lines.append(line)

        if self.ion_mode == "+":
            header_lines.append("#ION_MODE positive\n")

        elif self.ion_mode == "-":
            header_lines.append("#ION_MODE negative\n")

        header_lines.append("#CID_SERIES " + " ".join(sorted(list(set(self.CID_series)))) + "\n")
        header_lines.append("#MZLOW_MS2 " + str(self.MS2_mzlow) + "\n")
        header_lines.append("#MZHIGH_MS2 " + str(self.MS2_mzhigh) + "\n")
        header_lines.append("#MZLOW_MS1 " + str(self.MS1_mzlow) + "\n")
        header_lines.append("#MZHIGH_MS1 " + str(self.MS1_mzhigh) + "\n")
        header_lines.append("#ELEMENTAL_COMPOSITION_LIGHT " + str(self.nts_alphabet_light) + "\n")

        if self.nts_alphabet_heavy:
            header_lines.append("#ELEMENTAL_COMPOSITION_HEAVY " + str(self.nts_alphabet_heavy) + "\n")

        header_lines.append("#SEQX_CONSOLIDATION " + str(self.mz_consolidation) + "\n")

        if self.mz_consolidation:
            header_lines.append("#MS1_SEQX_PPM " + str(self.MS1_consolidation_ppm) + "\n")
            header_lines.append("#MS2_SEQX_PPM " + str(self.MS2_consolidation_ppm) + "\n")

        return header_lines

    def final_output(self):
        """
        Create the MS1 output file
        """
        # Control if options are correctly given and if there are differences
        # between nts dictionaries specified here and in previous scripts
        self.__controls()
        out_files = []

        if os.path.exists("./output.3.MS1"):
            charge_table(self.MS1_charges)
            # Write all the lines inside the standard output file "Digest_MS1"
            open("./Digest_MS1.txt", 'w').writelines(
                self.header_info_MS1(open("./output.3.MS1", 'r')) +
                self.lines_final(self.lines_with_charge(lines_with_chemistry(inlines_MS1(open("./output.3.MS1", 'r')))),
                                 self.MS1_mzlow, self.MS1_mzhigh))
            out_files.append("Digest_MS1.txt")

        ################################# MS2 SPECIFIC WORKFLOW #############################################
        """
        This part is executed only under the condition that the file output.3.MS2 is present in the directory
        """

        if os.path.exists("./output.3.MS2"):
            # Create two new dictionaries for a-B fragments (nucleobases)
            self.nts_dic_onlyB_light = nts_mass(self.read_excel_input(self.nts_alphabet_light))[1]

            if self.nts_alphabet_heavy:
                self.nts_dic_onlyB_heavy = nts_mass(self.read_excel_input(self.nts_alphabet_heavy))[1]

            lines_MS2 = []

            for line in self.lines_final(self.lines_with_charge(lines_with_chemistry(inlines_MS2())), self.MS1_mzlow,
                                         self.MS1_mzhigh):

                split = line.split()
                # Only fragments longer than the specified cutoff for MS2 ions are kept
                if len(split[7]) >= min_length_MS2:

                    # Only fragments with m/z in specified MS1 interval are considered
                    if self.MS1_mzhigh >= np.float64(split[0]) >= self.MS1_mzlow:
                        lines_MS2.append("{} {} {} {} {} {}".format(split[0], 'light', " ".join(split[2:9]), split[10],
                                                                    split[9], " ".join(split[11:])))

                    if self.nts_alphabet_heavy:
                        if self.MS1_mzhigh >= np.float64(split[1]) >= self.MS1_mzlow:
                            lines_MS2.append("{} {} {} {} {} {}".format(split[1], 'heavy', " ".join(split[2:9]),
                                                                        split[10], split[9], " ".join(split[11:])))

            # Determine a dictionary with the amount of masses to subtract (-) or add (+) to nts for the corresponding
            # type of CID ion. The difference is compared to a nt in neutral state within a chain, no 3' or 5' ends
            self.dic_CID_series_masses = {"a": -(ele_mass["P"] + ele_mass["O"] * 4 + ele_mass["H"] * 2),
                                          "a-B": -(ele_mass["P"] + ele_mass["O"] * 4 + ele_mass["H"] * 3),
                                          "b": -(ele_mass["P"] + ele_mass["O"] * 3),
                                          "c": -(ele_mass["O"] + ele_mass["H"]),
                                          "d": ele_mass["H"], "z": -ele_mass["H"],
                                          "y": ele_mass["O"] + ele_mass["H"], "x": ele_mass["P"] + ele_mass["O"] * 3,
                                          "w": ele_mass["P"] + ele_mass["O"] * 4 + ele_mass["H"] * 2,
                                          "z-P": -(ele_mass["H"] + ele_mass["P"] + ele_mass["O"] * 3),
                                          "y-P": -(ele_mass["H"] + ele_mass["P"] + ele_mass["O"] * 3)}

            # Determine a dictionary with the amount of masses to subtract (-) or add (+) to precursor ions (M) to
            # obtain the "M-xx" series
            self.dic_M_x_series = {"M-H2O": -(ele_mass["H"] * 2 + ele_mass["O"]),
                                   "M-PO3H": -(ele_mass["O"] * 3 + ele_mass["P"] + ele_mass["H"]),
                                   "M-PO4H3": -(ele_mass["P"] + ele_mass["O"] * 4 + ele_mass["H"] * 3),
                                   "M-PO3": -(ele_mass["O"] * 3 + ele_mass["P"]),
                                   "M-PO4H2": -(ele_mass["P"] + ele_mass["O"] * 4 + ele_mass["H"] * 2)}

            self.MS2_charge_table = charge_table(self.MS2_charges)

            # Write the output lines in the output file
            body_output = self.final_lines_MS2(lines_MS2)

            # Add the info on unique sequences and decoys in the header
            final_header = (self.header_info_MS2(open(os.getcwd() + "/output.3.MS2", 'r')) +
                            ["#TARGETS {}\n#DECOYS {}\n".format(tot_targets, tot_decoys)] +
                            ["m/z isotope molecule_ID residue_start residue_end charge miss sequence "
                            "sequence_mod 5'end 3'end num_copy molecule_location CID_series_fragment(charge):m/z\n"])

            # Separate the input files if multiple fasta files are selected, based on the running OS
            if platform.system() == 'Windows':
                fasta_input = self.fasta_name.split(';')

            else:
                fasta_input = self.fasta_name.split(':')

            # Determine the output name from the input fasta file
            output_name = filename_from_path(fasta_input[0])[:-6]

            open(os.getcwd() + "/Digest_" + output_name + ".txt", 'w').writelines(final_header + body_output)
            out_files.append("Digest_" + output_name + ".txt")

        else:
            print("WARNING! MS2 file output.3.MS2 from 3_consolidate.py is missing, MS2 digest will not be calculated")

        # m/z based consolidation, where nucleotides that in a particular alphabet are too close
        # in masses ppm to be indistinguishable in the matching, are reported as 'X' within their sequence
        if self.mz_consolidation:
            digest_lines = ct.mz_consolidate(self.nts_alphabet_light, 'Digest_' + output_name + '.txt', 'light',
                                             self.MS1_consolidation_ppm,
                                             self.MS2_consolidation_ppm, read_csv())
            if digest_lines:
                open('test_consolidate.txt', 'w').writelines(digest_lines)
                os.remove('./Digest_' + output_name + '.txt')
                move('./test_consolidate.txt', './Digest_' + output_name + '.txt')

            if self.nts_alphabet_heavy:
                digest_lines = ct.mz_consolidate(self.nts_alphabet_heavy, 'Digest_' + output_name + '.txt', 'heavy',
                                                 self.MS1_consolidation_ppm,
                                                 self.MS2_consolidation_ppm, read_csv())

                open('test_consolidate.txt', 'w').writelines(digest_lines)
                os.remove('./Digest_' + output_name + '.txt')
                move('./test_consolidate.txt', './Digest_' + output_name + '.txt')
            print("SeqX consolidation COMPLETED!\n")

        else:
            # Check if any combination of two nucleotides in the alphabet(s)
            # provided has masses within 0.5 Dalton, warning the user and asking to run consolidation
            ct.check_Da_nucleotides(self.nts_alphabet_light)
            if self.nts_alphabet_heavy:
                ct.check_Da_nucleotides(self.nts_alphabet_heavy)

class DictDiffer(object):
    """
    Calculate the difference between two dictionaries as:
    (1) items added
    (2) items removed
    (3) keys same in both but changed values
    (4) keys same in both and unchanged values
    """

    def __init__(self, current_dict, past_dict):
        self.current_dict, self.past_dict = current_dict, past_dict
        self.set_current, self.set_past = set(current_dict.keys()), set(past_dict.keys())
        self.intersect = self.set_current.intersection(self.set_past)

    def added(self):
        return self.set_current - self.intersect

    def removed(self):
        return self.set_past - self.intersect

    def changed(self):
        return set(o for o in self.intersect if self.past_dict[o] != self.current_dict[o])

    def unchanged(self):
        return set(o for o in self.intersect if self.past_dict[o] == self.current_dict[o])

def ppm_range(value, difference):
    """
    Calculate incertitude on MS1/MS2 masses equivalent to given ppms
    """
    return difference * 1000000 / value

def round_masses(mass):
    """
    Round all the masses for the output to 8 digits after comma
    """
    if mass == '-':
        return '-'
    else:
        return round(mass, 8)

def read_csv(input_csv='nts_light.csv'):
    """
    Produce a dictionary nts_alphabet
    mod_alphabet contains all ID : ID_ext couples, thus the one letter and extended codes for each nucleobase
    """
    if not os.path.exists(input_csv):
        print(
            "ERROR! File {} with info on nucleotides from script 2_modify.py is missing. "
            "Execution terminated without output".format(input_csv))
        sys.exit(1)

    else:
        # Read the csv file with the nucleotides dictionary
        df = pd.read_csv(input_csv, usecols=['ID', 'ID_ext'])

        # Drop rows with NaN values
        df = df[pd.notnull(df['ID'])]

        return dict(zip(df.ID, df.ID_ext))

def nts_mass(df):
    """
    Create a dictionary with the masses of all nts based on their atomic composition
    """
    mass_nts, mass_base, mass_backbone = {}, {}, {}

    for index, row in df.iterrows():
        # Calculate mass of the base only
        mass_b = (np.float64(row['C']) * ele_mass["C"] + np.float64(row['O']) * ele_mass["O"] + np.float64(
            row['H']) *
                  ele_mass["H"] + np.float64(row['N']) * ele_mass["N"] +
                  np.float64(row['P']) * ele_mass["P"] + np.float64(row['S']) * ele_mass["S"] + np.float64(
                    row['Se']) *
                  ele_mass["Se"] + np.float64(row['C13']) * ele_mass["C13"]
                  + np.float64(row['O18']) * ele_mass["O18"] + np.float64(row['N15']) * ele_mass[
                      "N15"] + np.float64(
                    row['H2']) * ele_mass["H2"])

        mass_base[row['ID']] = mass_b

        # Calculate mass of the backbone only
        mass_back = (np.float64(row['C.1']) * ele_mass["C"] + np.float64(row['O.1']) * ele_mass["O"] + np.float64(
            row['H.1']) * ele_mass["H"] + np.float64(row['N.1']) * ele_mass["N"] +
                     np.float64(row['P.1']) * ele_mass["P"] + np.float64(row['S.1']) * ele_mass["S"] + np.float64(
                    row['Se.1']) * ele_mass["Se"] + np.float64(row['C13.1']) * ele_mass["C13"]
                     + np.float64(row['O18.1']) * ele_mass["O18"] + np.float64(row['N15.1']) * ele_mass[
                         "N15"] + np.float64(row['H2.1']) * ele_mass["H2"])

        mass_backbone[row['ID']] = mass_back

        # Calculate mass of the whole nucleotide
        mass_nts[row['ID']] = mass_b + mass_back

    return mass_nts, mass_base, mass_backbone

def inlines_MS1(input_file):
    """
    Read the input file for MS1 (output.3.MS1)
    """
    input_lines = []

    for line in input_file:

        if line[0] != "#":

            if line.split()[2].isdigit():
                input_lines.append(line.rstrip())

    return input_lines

def reorganize_lines(line):
    """
    Order the precursor ion field
    """
    split = line.split()
    new_line = "{} {} {} {} {} {} {} {} {} {}".format(split[-3], split[-2], split[-1], split[2], split[0], split[1],
                                                      split[3], split[4], split[5], split[6])
    return new_line

def lines_with_chemistry(input_lines):
    """
    Prepare lines for the output lines including 3' and 5' chemistry
    """
    new_lines = []

    for line in input_lines:
        new_lines.append(reorganize_lines(line))

    return new_lines

def add_charge(line, charge):
    """
    Add information on the charge to the output lines
    """
    new_line = "{} {} {}".format(" ".join(line.split()[:3]), charge, " ".join(line.split()[3:]))

    return new_line

def charge_table(input_file):
    """
    Create a dictionary with the charge table for MS1/MS2 ions
    """
    charges_dic = {}

    with open(input_file, 'r') as infile:
        for line in infile:

            if line[0].isdigit():
                charges_dic[line.split()[0]] = []

                for ele in line.split()[1].split(','):

                    if ele.isdigit():
                        charges_dic[line.split()[0]].append(ele)

    if 'MS1' in input_file or 'ms1' in input_file:
        # Define the window for charges on all MS1 fragments, derived from the input charge table file
        global MS1_minz, MS1_z
        MS1_minz, MS1_z = min(int(s) for s in min(charges_dic.items(), key=lambda x: x[1])[1]), max(
            int(s) for s in max(charges_dic.items(), key=lambda x: x[1])[1])

    elif 'MS2' in input_file or 'ms2' in input_file:
        # Define the window for charges on all MS2 fragments, derived from the input charge table file
        global MS2_minz, MS2_z
        MS2_minz, MS2_z = min(int(s) for s in min(charges_dic.items(), key=lambda x: x[1])[1]), max(
            int(s) for s in max(charges_dic.items(), key=lambda x: x[1])[1])

    return charges_dic

def add_masses(line, mass_light, mass_heavy):
    """
    Add m/z information in the output lines
    """
    new_line = "{} {} {}\n".format(round_masses(mass_light), round_masses(mass_heavy), line)

    return new_line

def inlines_MS2():
    """
    Read the line from the input file for MS2
    """
    input_file = open("./output.3.MS2", 'r')
    input_lines = []

    for line in input_file:
        # Extracting the minimum length of fragments considered for MS2 fragmentation
        if "MIN_LENGTH" in line:
            global min_length_MS2
            min_length_MS2 = int(line.split()[1])

        if line[0] != "#" and line.split():

            if line.split()[2].isdigit():
                input_lines.append(line[:-1])

    return input_lines

def filename_from_path(path):
    head, tail = ntpath.split(path)
    return tail or ntpath.basename(head)


import os
import cv2
import dlib
import numpy as np
import argparse
import inception_resnet_v1
import tensorflow as tf
from imutils.face_utils import FaceAligner
from imutils.face_utils import rect_to_bb

def get_args():
    parser = argparse.ArgumentParser(description="This script detects faces from web cam input, "
                                                 "and estimates age and gender for the detected faces.",
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--weight_file", type=str, default=None,
                        help="path to weight file (e.g. weights.18-4.06.hdf5)")
    parser.add_argument("--depth", type=int, default=16,
                        help="depth of network")
    parser.add_argument("--width", type=int, default=8,
                        help="width of network")
    args = parser.parse_args()
    return args

def draw_label(image, point, label, font=cv2.FONT_HERSHEY_SIMPLEX,
               font_scale=1, thickness=2):
    size = cv2.getTextSize(label, font, font_scale, thickness)[0]
    x, y = point
    cv2.rectangle(image, (x, y - size[1]), (x + size[0], y), (255, 0, 0), cv2.FILLED)
    cv2.putText(image, label, point, font, font_scale, (255, 255, 255), thickness)

def main(sess,age,gender,train_mode,images_pl):
    args = get_args()
    depth = args.depth
    k = args.width

    # for face detection
    detector = dlib.get_frontal_face_detector()
    predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
    fa = FaceAligner(predictor, desiredFaceWidth=160)

    # load model and weights
    img_size = 160

    # capture video
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    while True:
        # get video frame
        ret, img = cap.read()

        if not ret:
            print("error: failed to capture image")
            return -1

        input_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img_h, img_w, _ = np.shape(input_img)

        # detect faces using dlib detector
        detected = detector(input_img, 1)
        faces = np.empty((len(detected), img_size, img_size, 3))

        for i, d in enumerate(detected):
            x1, y1, x2, y2, w, h = d.left(), d.top(), d.right() + 1, d.bottom() + 1, d.width(), d.height()
            xw1 = max(int(x1 - 0.4 * w), 0)
            yw1 = max(int(y1 - 0.4 * h), 0)
            xw2 = min(int(x2 + 0.4 * w), img_w - 1)
            yw2 = min(int(y2 + 0.4 * h), img_h - 1)
            cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2)
            # cv2.rectangle(img, (xw1, yw1), (xw2, yw2), (255, 0, 0), 2)
            faces[i, :, :, :] = fa.align(input_img, gray, detected[i])
            # faces[i,:,:,:] = cv2.resize(img[yw1:yw2 + 1, xw1:xw2 + 1, :], (img_size, img_size))
        #
        if len(detected) > 0:
            # predict ages and genders of the detected faces
            ages,genders = sess.run([age, gender], feed_dict={images_pl: faces, train_mode: False})

        # draw results
        for i, d in enumerate(detected):
            label = "{}, {}".format(int(ages[i]), "F" if genders[i] == 0 else "M")
            draw_label(img, (d.left(), d.top()), label)

        cv2.imshow("result", img)
        key = cv2.waitKey(1)

        if key == 27:
            break

def load_network(model_path):
    sess = tf.Session()
    images_pl = tf.placeholder(tf.float32, shape=[None, 160, 160, 3], name='input_image')
    images_norm = tf.map_fn(lambda frame: tf.image.per_image_standardization(frame), images_pl)
    train_mode = tf.placeholder(tf.bool)
    age_logits, gender_logits, _ = inception_resnet_v1.inference(images_norm, keep_probability=0.8,
                                                                 phase_train=train_mode,
                                                                 weight_decay=1e-5)
    gender = tf.argmax(tf.nn.softmax(gender_logits), 1)
    age_ = tf.cast(tf.constant([i for i in range(0, 101)]), tf.float32)
    age = tf.reduce_sum(tf.multiply(tf.nn.softmax(age_logits), age_), axis=1)
    init_op = tf.group(tf.global_variables_initializer(),
                       tf.local_variables_initializer())
    sess.run(init_op)
    saver = tf.train.Saver()
    ckpt = tf.train.get_checkpoint_state(model_path)
    if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)
        print("restore model!")
    else:
        pass
    return sess,age,gender,train_mode,images_pl

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", "--M", default="./models", type=str, help="Model Path")
    args = parser.parse_args()
    sess, age, gender, train_mode,images_pl = load_network(args.model_path)
    print("age is=", age)
    main(sess,age,gender,train_mode,images_pl)


import warnings
import mmcv
import numpy as np
import torch
from torch.nn.modules.utils import _pair

from mmdet.core.anchor.builder import ANCHOR_GENERATORS
from mmdet.core.anchor import AnchorGenerator

@ANCHOR_GENERATORS.register_module(force=True)
class SSDAnchorGenerator(AnchorGenerator):
    """Anchor generator for SSD

    Args:
        strides (list[int]  | list[tuple[int, int]]): Strides of anchors
            in multiple feature levels.
        ratios (list[float]): The list of ratios between the height and width
            of anchors in a single level.
        basesize_ratio_range (tuple(float)): Ratio range of anchors.
        input_size (int): Size of feature map, 300 for SSD300,
            512 for SSD512.
        scale_major (bool): Whether to multiply scales first when generating
            base anchors. If true, the anchors in the same row will have the
            same scales. It is always set to be False in SSD.
    """

    def __init__(self,
                 strides,
                 ratios,
                 basesize_ratio_range,
                 input_size=300,
                 scale_major=True):
        assert len(strides) == len(ratios)
        assert mmcv.is_tuple_of(basesize_ratio_range, float)

        self.strides = [_pair(stride) for stride in strides]
        self.input_size = max(input_size) if isinstance(input_size, (list,tuple)) else input_size
        self.centers = [(stride[0] / 2., stride[1] / 2.)
                        for stride in self.strides]
        self.basesize_ratio_range = basesize_ratio_range

        # calculate anchor ratios and sizes
        min_ratio, max_ratio = basesize_ratio_range
        min_ratio = int(min_ratio * 100)
        max_ratio = int(max_ratio * 100)
        step = int(np.floor(max_ratio - min_ratio) / (self.num_levels - 2))
        min_sizes = []
        max_sizes = []
        for ratio in range(int(min_ratio), int(max_ratio) + 1, step):
            min_sizes.append(int(self.input_size * ratio / 100))
            max_sizes.append(int(self.input_size * (ratio + step) / 100))
        if self.input_size == 300:
            if basesize_ratio_range[0] == 0.15:  # SSD300 COCO
                min_sizes.insert(0, int(self.input_size * 7 / 100))
                max_sizes.insert(0, int(self.input_size * 15 / 100))
            elif basesize_ratio_range[0] == 0.2:  # SSD300 VOC
                min_sizes.insert(0, int(self.input_size * 10 / 100))
                max_sizes.insert(0, int(self.input_size * 20 / 100))
            else:
                min_sizes.insert(0, int(self.input_size * basesize_ratio_range[0] * 0.4))
                max_sizes.insert(0, int(self.input_size * basesize_ratio_range[0]))
                warnings.warn(
                    'according to original SSD, basesize_ratio_range[0] should be either 0.15'
                    'or 0.2 when input_size is 300, got '
                    f'{basesize_ratio_range[0]}.')
        elif self.input_size == 512:
            if basesize_ratio_range[0] == 0.1:  # SSD512 COCO
                min_sizes.insert(0, int(self.input_size * 4 / 100))
                max_sizes.insert(0, int(self.input_size * 10 / 100))
            elif basesize_ratio_range[0] == 0.15:  # SSD512 VOC
                min_sizes.insert(0, int(self.input_size * 7 / 100))
                max_sizes.insert(0, int(self.input_size * 15 / 100))
            else:
                min_sizes.insert(0, int(self.input_size * basesize_ratio_range[0] * 0.4))
                max_sizes.insert(0, int(self.input_size * basesize_ratio_range[0]))
                warnings.warn('according to original SSD, basesize_ratio_range[0] should be either 0.1'
                                 'or 0.15 when input_size is 512, got'
                                 f' {basesize_ratio_range[0]}.')
        else:
            if basesize_ratio_range[0] == 0.1:  # SSD512 COCO
                min_sizes.insert(0, int(self.input_size * 4 / 100))
                max_sizes.insert(0, int(self.input_size * 10 / 100))
            else:
                min_sizes.insert(0, int(self.input_size * basesize_ratio_range[0] * 0.4))
                max_sizes.insert(0, int(self.input_size * basesize_ratio_range[0]))

        anchor_ratios = []
        anchor_scales = []
        for k in range(len(self.strides)):
            scales = [1., np.sqrt(max_sizes[k] / min_sizes[k])]
            anchor_ratio = [1.]
            for r in ratios[k]:
                anchor_ratio += [1 / r, r]  # 4 or 6 ratio
            anchor_ratios.append(torch.Tensor(anchor_ratio))
            anchor_scales.append(torch.Tensor(scales))

        self.base_sizes = min_sizes
        self.scales = anchor_scales
        self.ratios = anchor_ratios
        self.scale_major = scale_major
        self.center_offset = 0
        self.base_anchors = self.gen_base_anchors()
        # added for proto export
        self.min_sizes = min_sizes
        self.max_sizes = max_sizes

    def gen_base_anchors(self):
        """Generate base anchors.

        Returns:
            list(torch.Tensor): Base anchors of a feature grid in multiple \
                feature levels.
        """
        multi_level_base_anchors = []
        for i, base_size in enumerate(self.base_sizes):
            base_anchors = self.gen_single_level_base_anchors(
                base_size,
                scales=self.scales[i],
                ratios=self.ratios[i],
                center=self.centers[i])
            indices = list(range(len(self.ratios[i])))
            indices.insert(1, len(indices))
            base_anchors = torch.index_select(base_anchors, 0,
                                              torch.LongTensor(indices))
            multi_level_base_anchors.append(base_anchors)
        return multi_level_base_anchors

    def __repr__(self):
        """str: a string that describes the module"""
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
        repr_str += f'{indent_str}strides={self.strides},\n'
        repr_str += f'{indent_str}scales={self.scales},\n'
        repr_str += f'{indent_str}scale_major={self.scale_major},\n'
        repr_str += f'{indent_str}input_size={self.input_size},\n'
        repr_str += f'{indent_str}scales={self.scales},\n'
        repr_str += f'{indent_str}ratios={self.ratios},\n'
        repr_str += f'{indent_str}num_levels={self.num_levels},\n'
        repr_str += f'{indent_str}base_sizes={self.base_sizes},\n'
        repr_str += f'{indent_str}basesize_ratio_range='
        repr_str += f'{self.basesize_ratio_range})'
        return repr_str


"""Defines the graph for sampling from Coconet."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import time
# internal imports
import numpy as np
import tensorflow as tf
from magenta.models.coconet import lib_graph
from magenta.models.coconet import lib_hparams

FLAGS = tf.app.flags.FLAGS

class CoconetSampleGraph(object):
  """Graph for Gibbs sampling from Coconet."""

  def __init__(self, chkpt_path, placeholders=None):
    """Initializes inputs for the Coconet sampling graph.

    Does not build or restore the graph. That happens lazily if you call run(),
    or explicitly using instantiate_sess_and_restore_checkpoint.

    Args:
      chkpt_path: Checkpoint directory for loading the model.
          Uses the latest checkpoint.
      placeholders: Optional placeholders.
    """
    self.chkpt_path = chkpt_path
    self.hparams = lib_hparams.load_hparams(chkpt_path)
    if placeholders is None:
      self.placeholders = self.get_placeholders()
    else:
      self.placeholders = placeholders
    self.samples = None
    self.sess = None

  def get_placeholders(self):
    hparams = self.hparams
    return dict(
        pianorolls=tf.placeholder(
            tf.bool,
            [None, None, hparams.num_pitches, hparams.num_instruments],
            "pianorolls"),
        # The default value is only used for checking if completion masker
        # should be evoked.  It can't be used directly as the batch size
        # and length of pianorolls are unknown during static time.
        outer_masks=tf.placeholder_with_default(
            np.zeros(
                (1, 1, hparams.num_pitches, hparams.num_instruments),
                dtype=np.float32),
            [None, None, hparams.num_pitches, hparams.num_instruments],
            "outer_masks"),
        sample_steps=tf.placeholder_with_default(0, (), "sample_steps"),
        total_gibbs_steps=tf.placeholder_with_default(
            0, (), "total_gibbs_steps"),
        current_step=tf.placeholder_with_default(0, (), "current_step"),
        temperature=tf.placeholder_with_default(0.99, (), "temperature"))

  @property
  def inputs(self):
    return self.placeholders

  def make_outer_masks(self, outer_masks, input_pianorolls):
    """Returns outer masks, if all zeros created by completion masking."""
    outer_masks = tf.to_float(outer_masks)
    # If outer_masks come in as all zeros, it means there's no masking,
    # which also means nothing will be generated. In this case, use
    # completion mask to make new outer masks.
    outer_masks = tf.cond(
        tf.reduce_all(tf.equal(outer_masks, 0)),
        lambda: make_completion_masks(input_pianorolls),
        lambda: outer_masks)
    return outer_masks

  def build_sample_graph(self, input_pianorolls=None, outer_masks=None,
                         total_gibbs_steps=None):
    """Builds the tf.while_loop based sampling graph.

    Args:
      input_pianorolls: Optional input pianorolls override. If None, uses the
          pianorolls placeholder.
      outer_masks: Optional input outer_masks override. If None, uses the
          outer_masks placeholder.
      total_gibbs_steps: Optional input total_gibbs_steps override. If None,
          uses the total_gibbs_steps placeholder.
    Returns:
      The output op of the graph.
    """
    if input_pianorolls is None:
      input_pianorolls = self.inputs["pianorolls"]
    if outer_masks is None:
      outer_masks = self.inputs["outer_masks"]

    tt = tf.shape(input_pianorolls)[1]
    sample_steps = tf.to_float(self.inputs["sample_steps"])
    if total_gibbs_steps is None:
      total_gibbs_steps = self.inputs["total_gibbs_steps"]
    temperature = self.inputs["temperature"]

    input_pianorolls = tf.to_float(input_pianorolls)
    outer_masks = self.make_outer_masks(outer_masks, input_pianorolls)

    # Calculate total_gibbs_steps as steps * num_instruments if not given.
    total_gibbs_steps = tf.cond(
        tf.equal(total_gibbs_steps, 0),
        lambda: tf.to_float(tt * self.hparams.num_instruments),
        lambda: tf.to_float(total_gibbs_steps))

    # sample_steps is set to total_gibbs_steps if not given.
    sample_steps = tf.cond(
        tf.equal(sample_steps, 0),
        lambda: total_gibbs_steps,
        lambda: tf.to_float(sample_steps))

    def infer_step(pianorolls, step_count):
      """Called by tf.while_loop, takes a Gibbs step."""
      mask_prob = compute_mask_prob_from_yao_schedule(step_count,
                                                      total_gibbs_steps)
      # 1 indicates mask out, 0 is not mask.
      masks = make_bernoulli_masks(tf.shape(pianorolls), mask_prob,
                                   outer_masks)

      logits = self.predict(pianorolls, masks)
      samples = sample_with_temperature(logits, temperature=temperature)

      outputs = pianorolls * (1 - masks) + samples * masks

      check_completion_op = tf.assert_equal(
          tf.where(tf.equal(tf.reduce_max(masks, axis=2), 1.),
                   tf.reduce_max(outputs, axis=2),
                   tf.reduce_max(pianorolls, axis=2)),
          1.)
      with tf.control_dependencies([check_completion_op]):
        outputs = tf.identity(outputs)

      step_count += 1
      return outputs, step_count

    current_step = tf.to_float(self.inputs["current_step"])

    # Initializes pianorolls by evaluating the model once to fill in all gaps.
    logits = self.predict(tf.to_float(input_pianorolls), outer_masks)
    samples = sample_with_temperature(logits, temperature=temperature)
    tf.get_variable_scope().reuse_variables()

    self.samples, current_step = tf.while_loop(
        lambda samples, current_step: current_step < sample_steps,
        infer_step, [samples, current_step],
        shape_invariants=[
            tf.TensorShape([None, None, None, None]),
            tf.TensorShape(None),
        ],
        back_prop=False,
        parallel_iterations=1,
        name="coco_while")
    self.samples.set_shape(input_pianorolls.shape)
    return self.samples

  def predict(self, pianorolls, masks):
    """Evalutes the model once and returns predictions."""
    direct_inputs = dict(
        pianorolls=pianorolls, masks=masks,
        lengths=tf.to_float([tf.shape(pianorolls)[1]]))

    model = lib_graph.build_graph(
        is_training=False,
        hparams=self.hparams,
        direct_inputs=direct_inputs,
        use_placeholders=False)
    self.logits = model.logits
    return self.logits

  def instantiate_sess_and_restore_checkpoint(self):
    """Instantiates session and restores from self.chkpt_path."""
    if self.samples is None:
      self.build_sample_graph()
    sess = tf.Session()
    saver = tf.train.Saver()
    chkpt_fpath = tf.train.latest_checkpoint(self.chkpt_path)
    tf.logging.info("loading checkpoint %s", chkpt_fpath)
    saver.restore(sess, chkpt_fpath)
    tf.get_variable_scope().reuse_variables()
    self.sess = sess
    return self.sess

  def run(self,
          pianorolls,
          masks=None,
          sample_steps=0,
          current_step=0,
          total_gibbs_steps=0,
          temperature=0.99):
    """Given input pianorolls, runs Gibbs sampling to fill in the rest.

    When total_gibbs_steps is 0, total_gibbs_steps is set to
    time * instruments.  If faster sampling is desired on the expanse of sample
    quality, total_gibbs_steps can be explicitly set to a lower number,
    possibly to the value of sample_steps if do not plan on stopping sample
    early to obtain intermediate results.

    This function can be used to return intermediate results by setting the
    sample_steps to when results should be returned and leaving
    total_gibbs_steps to be 0.

    To continue sampling from intermediate results, set current_step to the
    number of steps taken, and feed in the intermediate pianorolls.  Again
    leaving total_gibbs_steps as 0.

    Builds the graph and restores checkpoint if necessary.

    Args:
      pianorolls: a 4D numpy array of shape (batch, time, pitch, instrument)
      masks: a 4D numpy array of the same shape as pianorolls, with 1s
          indicating mask out.  If is None, then the masks will be where have 1s
          where there are no notes, indicating to the model they should be
          filled in.
      sample_steps: an integer indicating the number of steps to sample in this
          call.  If set to 0, then it defaults to total_gibbs_steps.
      current_step: an integer indicating how many steps might have already
          sampled before.
      total_gibbs_steps: an integer indicating the total number of steps that
          a complete sampling procedure would take.
      temperature: a float indicating the temperature for sampling from softmax.

    Returns:
      A dictionary, consisting of "pianorolls" which is a 4D numpy array of
      the sampled results and "time_taken" which is the time taken in sampling.
    """
    if self.sess is None:
      # Build graph and restore checkpoint.
      self.instantiate_sess_and_restore_checkpoint()

    if masks is None:
      masks = np.zeros_like(pianorolls)

    start_time = time.time()
    new_piece = self.sess.run(
        self.samples,
        feed_dict={
            self.placeholders["pianorolls"]: pianorolls,
            self.placeholders["outer_masks"]: masks,
            self.placeholders["sample_steps"]: sample_steps,
            self.placeholders["total_gibbs_steps"]: total_gibbs_steps,
            self.placeholders["current_step"]: current_step,
            self.placeholders["temperature"]: temperature
        })

    label = "independent blocked gibbs"
    time_taken = (time.time() - start_time) / 60.0
    tf.logging.info("exit  %s (%.3fmin)" % (label, time_taken))
    return dict(pianorolls=new_piece, time_taken=time_taken)

def make_completion_masks(pianorolls, outer_masks=1.):
  pianorolls = tf.to_float(pianorolls)
  masks = tf.reduce_all(tf.equal(pianorolls, 0), axis=2, keep_dims=True)
  inner_masks = tf.to_float(masks) + 0 * pianorolls
  return inner_masks * outer_masks

def make_bernoulli_masks(shape, pm, outer_masks=1.):
  bb = shape[0]
  tt = shape[1]
  pp = shape[2]
  ii = shape[3]
  probs = tf.random_uniform([bb, tt, ii])
  masks = tf.tile(tf.to_float(tf.less(probs, pm))[:, :, None, :], [1, 1, pp, 1])
  return masks * outer_masks

def sample_with_temperature(logits, temperature):
  """Either argmax after softmax or random sample along the pitch axis.

  Args:
    logits: a Tensor of shape (batch, time, pitch, instrument).
    temperature: a float  0.0=argmax 1.0=random

  Returns:
    a Tensor of the same shape, with one_hots on the pitch dimension.
  """
  logits = tf.transpose(logits, [0, 1, 3, 2])
  pitch_range = tf.shape(logits)[-1]

  def sample_from_logits(logits):
    with tf.control_dependencies([tf.assert_greater(temperature, 0.0)]):
      logits = tf.identity(logits)
    reshaped_logits = (
        tf.reshape(logits, [-1, tf.shape(logits)[-1]]) / temperature)
    choices = tf.multinomial(reshaped_logits, 1)
    choices = tf.reshape(choices,
                         tf.shape(logits)[:logits.get_shape().ndims - 1])
    return choices

  choices = tf.cond(tf.equal(temperature, 0.0),
                    lambda: tf.argmax(tf.nn.softmax(logits), -1),
                    lambda: sample_from_logits(logits))
  samples_onehot = tf.one_hot(choices, pitch_range)
  return tf.transpose(samples_onehot, [0, 1, 3, 2])

def compute_mask_prob_from_yao_schedule(i, n, pmin=0.1, pmax=0.9, alpha=0.7):
  wat = (pmax - pmin) * i/ n
  return tf.maximum(pmin, pmax - wat / alpha)

def main(unused_argv):
  checkpoint_path = FLAGS.checkpoint
  sampler = CoconetSampleGraph(checkpoint_path)

  batch_size = 1
  decode_length = 4
  target_shape = [batch_size, decode_length, 46, 4]
  pianorolls = np.zeros(target_shape, dtype=np.float32)
  generated_piece = sampler.run(pianorolls, sample_steps=16, temperature=0.99)
  tf.logging.info("num of notes in piece %d", np.sum(generated_piece))

  tf.logging.info("Done.")

if __name__ == "__main__":
  tf.app.run()


# Copyright (C) 2018-2021 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import logging as log

import numpy as np

from mo.front.common.partial_infer.utils import dynamic_dimension, shape_array, dynamic_dimension_value
from mo.front.extractor import bool_to_str
from mo.graph.graph import Node, Graph
from mo.middle.passes.convert_data_type import np_data_type_to_destination_type
from mo.ops.op import Op
from mo.utils.error import Error

class NonMaxSuppression(Op):
    op = 'NonMaxSuppression'

    def __init__(self, graph: Graph, attrs: dict):
        mandatory_props = {
            'type': self.op,
            'op': self.op,
            'version': 'opset5',
            'infer': self.infer,
            'output_type': np.int64,
            'box_encoding': 'corner',
            'in_ports_count': 5,
            'sort_result_descending': 1,
            'force_precision_in_ports': {
                2: 'int64'},
            'type_infer': self.type_infer,
        }
        super().__init__(graph, mandatory_props, attrs)
        version = self.get_opset()
        if version in ['opset1', 'opset3', 'opset4']:
            self.attrs['out_ports_count'] = 1
        elif version == 'opset5':
            self.attrs['out_ports_count'] = 3
        else:
            raise Error('Unsupported operation opset version "{}"'.format(version))

    def backend_attrs(self):
        version = self.get_opset()
        if version in ['opset3', 'opset4', 'opset5']:
            return [('sort_result_descending', lambda node: bool_to_str(node, 'sort_result_descending')),
                    'box_encoding',
                    ('output_type', lambda node: np_data_type_to_destination_type(node.output_type))]
        elif version == 'opset1':
            return [('sort_result_descending', lambda node: bool_to_str(node, 'sort_result_descending')),
                    'box_encoding']
        else:
            raise Error('Unsupported operation opset version "{}"'.format(version))

    @staticmethod
    def infer(node: Node):
        num_of_inputs = len(node.in_ports())
        opset = node.get_opset()
        max_num_of_inputs = 6 if opset == 'opset5' else 5
        input_msg_fmt = 'NonMaxSuppression node {} from {} must have from 2 to {} inputs'
        node_name = node.soft_get('name', node.id)
        inputs_msg = input_msg_fmt.format(node_name, opset, max_num_of_inputs)
        assert 2 <= num_of_inputs <= max_num_of_inputs, inputs_msg

        boxes_shape = node.in_port(0).data.get_shape()
        assert boxes_shape is not None, 'The shape of tensor with boxes is not defined'
        scores_shape = node.in_port(1).data.get_shape()
        assert scores_shape is not None, 'The shape of tensor with scores is not defined'
        assert len(boxes_shape) == 3, 'Length of tensors with boxes must be equal to 3'
        assert len(scores_shape) == 3, 'Length of tensors with scores must be equal to 3'

        # According to the specification of the operation NonMaxSuppression,
        # the input 'max_output_boxes_per_class' (port 2) is optional, with default value 0.
        if num_of_inputs >= 3:
            max_output_boxes_per_class = node.in_port(2).data.get_value()
        else:
            max_output_boxes_per_class = 0

        if not max_output_boxes_per_class:
            log.info('Set default "max_output_boxes_per_class" for node {} to number of boxes'.format(node.name))
            max_output_boxes_per_class = boxes_shape[1]

        # convert the np.array value to a scalar to avoid issue with ragged numpy array generation in the shape
        # calculation formulas below
        if isinstance(max_output_boxes_per_class, np.ndarray):
            max_output_boxes_per_class = max_output_boxes_per_class.item()

        num_classes = scores_shape[1]
        num_input_boxes = boxes_shape[1]
        assert scores_shape[2] is dynamic_dimension or scores_shape[2] == num_input_boxes or scores_shape[2] is None \
               or num_input_boxes is None, 'Number of boxes mismatch for operation {}'.format(node_name)

        if node.get_opset() in ['opset4', 'opset5']:
            max_number_of_boxes = min(num_input_boxes, max_output_boxes_per_class) * boxes_shape[0] * num_classes
        else:
            max_number_of_boxes = min(num_input_boxes, boxes_shape[0] * max_output_boxes_per_class * num_classes)
        node.out_port(0).data.set_shape(shape_array([max_number_of_boxes, 3]))

        if opset == 'opset5':
            node.out_port(0).data.set_shape(shape_array([dynamic_dimension_value, 3]))
            num_of_outputs = len([port for port in node.out_ports().values() if not port.disconnected()])
            if num_of_outputs >= 2 and node.has_port('out', 1):
                node.out_port(1).data.set_shape(shape_array([dynamic_dimension_value, 3]))
            if num_of_outputs >= 3 and node.has_port('out', 2):
                node.out_port(2).data.set_shape(shape_array(1))

    @staticmethod
    def type_infer(node):
        opset = node.get_opset()
        if opset == 'opset5':
            node.out_port(0).set_data_type(node.output_type)
            if node.has_port('out', 1):
                node.out_port(1).set_data_type(np.float32)
            if node.has_port('out', 2):
                node.out_port(2).set_data_type(np.int64)
        elif opset in ['opset3', 'opset4']:
            node.out_port(0).set_data_type(node.output_type)
        else:
            node.out_port(0).set_data_type(np.int64)