# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License

import re
import numpy as np
import warnings
import collections
import os
from monty.io import zopen
import fnmatch
from collections import defaultdict
from pymatgen.electronic_structure.core import Spin, Orbital
from pymatgen.io.vasp.outputs import Vasprun
from pymatgen.electronic_structure.dos import Dos, LobsterCompleteDos
from pymatgen.core.structure import Structure
from pymatgen.electronic_structure.bandstructure import Kpoint
from pymatgen.io.vasp.inputs import Kpoints
from collections import OrderedDict
from pymatgen.electronic_structure.bandstructure import LobsterBandStructureSymmLine

"""
Module for reading Lobster output files. For more information
on LOBSTER see www.cohp.de.
"""

__author__ = "Marco Esters, Janine George"
__copyright__ = "Copyright 2017, The Materials Project"
__version__ = "0.2"
__maintainer__ = "Marco Esters, Janine George"
__email__ = "esters@uoregon.edu, janine.george@uclouvain.be"
__date__ = "Dec 13, 2017"

class Cohpcar:
    """
    Class to read COHPCAR/COOPCAR files generated by LOBSTER.

    Args:
        are_coops: Determines if the file is a list of COHPs or COOPs.
          Default is False for COHPs.

        filename: Name of the COHPCAR file. If it is None, the default
          file name will be chosen, depending on the value of are_coops.

    .. attribute: cohp_data

         Dict that contains the COHP data of the form:
           {bond: {"COHP": {Spin.up: cohps, Spin.down:cohps},
                   "ICOHP": {Spin.up: icohps, Spin.down: icohps},
                   "length": bond length,
                   "sites": sites corresponding to the bond}
         Also contains an entry for the average, which does not have
         a "length" key.

    .. attribute: efermi

         The Fermi energy in eV.

    .. attribute: energies

         Sequence of energies in eV. Note that LOBSTER shifts the energies
         so that the Fermi energy is at zero.

    .. attribute: is_spin_polarized

         Boolean to indicate if the calculation is spin polarized.

    .. attribute: orb_res_cohp

        orb_cohp[label] = {bond_data["orb_label"]: {"COHP": {Spin.up: cohps, Spin.down:cohps},
                                                     "ICOHP": {Spin.up: icohps, Spin.down: icohps},
                                                     "orbitals": orbitals,
                                                     "length": bond lengths,
                                                     "sites": sites corresponding to the bond}}

    """

    def __init__(self, are_coops=False, filename=None):
        self.are_coops = are_coops
        if filename is None:
            filename = "COOPCAR.lobster" if are_coops \
                else "COHPCAR.lobster"

        with zopen(filename, "rt") as f:
            contents = f.read().split("\n")

        # The parameters line is the second line in a COHPCAR file. It
        # contains all parameters that are needed to map the file.
        parameters = contents[1].split()
        # Subtract 1 to skip the average
        num_bonds = int(parameters[0]) - 1
        self.efermi = float(parameters[-1])
        if int(parameters[1]) == 2:
            spins = [Spin.up, Spin.down]
            self.is_spin_polarized = True
        else:
            spins = [Spin.up]
            self.is_spin_polarized = False

        # The COHP data start in row num_bonds + 3
        data = np.array([np.array(row.split(), dtype=float)
                         for row in contents[num_bonds + 3:]]).transpose()
        self.energies = data[0]

        cohp_data = {"average": {"COHP": {spin: data[1 + 2 * s * (num_bonds + 1)]
                                          for s, spin in enumerate(spins)},
                                 "ICOHP": {spin: data[2 + 2 * s * (num_bonds + 1)]
                                           for s, spin in enumerate(spins)}}}
        orb_cohp = {}

        # present for Lobster versions older than Lobster 2.2.0
        veryold = False
        # the labeling had to be changed: there are more than one COHP for each atom combination
        # this is done to make the labeling consistent with ICOHPLIST.lobster
        bondnumber = 0
        for bond in range(num_bonds):
            bond_data = self._get_bond_data(contents[3 + bond])

            label = str(bondnumber)

            orbs = bond_data["orbitals"]
            cohp = {spin: data[2 * (bond + s * (num_bonds + 1)) + 3]
                    for s, spin in enumerate(spins)}

            icohp = {spin: data[2 * (bond + s * (num_bonds + 1)) + 4]
                     for s, spin in enumerate(spins)}
            if orbs is None:
                bondnumber = bondnumber + 1
                label = str(bondnumber)
                cohp_data[label] = {"COHP": cohp, "ICOHP": icohp,
                                    "length": bond_data["length"],
                                    "sites": bond_data["sites"]}

            elif label in orb_cohp:
                orb_cohp[label].update({bond_data["orb_label"]:
                                            {"COHP": cohp,
                                             "ICOHP": icohp,
                                             "orbitals": orbs,
                                             "length": bond_data["length"],
                                             "sites": bond_data["sites"]}})
            else:
                # present for Lobster versions older than Lobster 2.2.0
                if bondnumber == 0:
                    veryold = True
                if veryold:
                    bondnumber += 1
                    label = str(bondnumber)

                orb_cohp[label] = {bond_data["orb_label"]: {"COHP": cohp,
                                                            "ICOHP": icohp,
                                                            "orbitals": orbs,
                                                            "length": bond_data["length"],
                                                            "sites": bond_data["sites"]}}

        # present for lobster older than 2.2.0
        if veryold:
            for bond in orb_cohp:
                cohp_data[bond] = {"COHP": None, "ICOHP": None,
                                   "length": bond_data["length"],
                                   "sites": bond_data["sites"]}

        self.orb_res_cohp = orb_cohp if orb_cohp else None
        self.cohp_data = cohp_data

    @staticmethod
    def _get_bond_data(line):
        """
        Subroutine to extract bond label, site indices, and length from
        a LOBSTER header line. The site indices are zero-based, so they
        can be easily used with a Structure object.

        Example header line: No.4:Fe1->Fe9(2.4524893531900283)
        Example header line for orbtial-resolved COHP:
            No.1:Fe1[3p_x]->Fe2[3d_x^2-y^2](2.456180552772262)

        Args:
            line: line in the COHPCAR header describing the bond.

        Returns:
            Dict with the bond label, the bond length, a tuple of the site
            indices, a tuple containing the orbitals (if orbital-resolved),
            and a label for the orbitals (if orbital-resolved).
        """

        orb_labs = ["s", "p_y", "p_z", "p_x", "d_xy", "d_yz", "d_z^2",
                    "d_xz", "d_x^2-y^2", "f_y(3x^2-y^2)", "f_xyz",
                    "f_yz^2", "f_z^3", "f_xz^2", "f_z(x^2-y^2)", "f_x(x^2-3y^2)"]

        line = line.rsplit("(", 1)
        # bondnumber = line[0].replace("->", ":").replace(".", ":").split(':')[1]
        length = float(line[-1][:-1])

        sites = line[0].replace("->", ":").split(":")[1:3]
        site_indices = tuple(int(re.split(r"\D+", site)[1]) - 1
                             for site in sites)

        # species = tuple(re.split(r"\d+", site)[0] for site in sites)
        if "[" in sites[0]:
            orbs = [re.findall(r"\[(.*)\]", site)[0] for site in sites]
            orbitals = [tuple((int(orb[0]), Orbital(orb_labs.index(orb[1:]))))
                        for orb in orbs]
            orb_label = "%d%s-%d%s" % (orbitals[0][0], orbitals[0][1].name,
                                       orbitals[1][0], orbitals[1][1].name)

        else:
            orbitals = None
            orb_label = None

        # a label based on the species alone is not feasible, there can be more than one bond for each atom combination
        # label = "%s" % (bondnumber)

        bond_data = {"length": length, "sites": site_indices,
                     "orbitals": orbitals, "orb_label": orb_label}
        return bond_data

class Icohplist:
    """
    Class to read ICOHPLIST/ICOOPLIST files generated by LOBSTER.

    Args:
        are_coops: Determines if the file is a list of ICOHPs or ICOOPs.
          Defaults to False for ICOHPs.

        filename: Name of the ICOHPLIST file. If it is None, the default
          file name will be chosen, depending on the value of are_coops.

    .. attribute: are_coops
         Boolean to indicate if the populations are COOPs or COHPs.

    .. attribute: is_spin_polarized
         Boolean to indicate if the calculation is spin polarized.

    .. attribute: Icohplist
        Dict containing the listfile data of the form:
           {bond: "length": bond length,
                  "number_of_bonds": number of bonds
                  "icohp": {Spin.up: ICOHP(Ef) spin up, Spin.down: ...}}

    .. attribute: IcohpCollection
        IcohpCollection Object

    """

    def __init__(self, are_coops=False, filename=None):

        self.are_coops = are_coops
        if filename is None:
            filename = "ICOOPLIST.lobster" if are_coops \
                else "ICOHPLIST.lobster"

        # LOBSTER list files have an extra trailing blank line
        # and we don't need the header.
        with zopen(filename) as f:
            data = f.read().split("\n")[1:-1]
        if len(data) == 0:
            raise IOError("ICOHPLIST file contains no data.")

        # Which Lobster version?
        if len(data[0].split()) == 8:
            version = '3.1.1'
        elif len(data[0].split()) == 6:
            version = '2.2.1'
            warnings.warn('Please consider using the new Lobster version. See www.cohp.de.')
        else:
            raise ValueError

        # If the calculation is spin polarized, the line in the middle
        # of the file will be another header line.
        if "distance" in data[len(data) // 2]:
            num_bonds = len(data) // 2
            if num_bonds == 0:
                raise IOError("ICOHPLIST file contains no data.")
            self.is_spin_polarized = True
        else:
            num_bonds = len(data)
            self.is_spin_polarized = False

        list_labels = []
        list_atom1 = []
        list_atom2 = []
        list_length = []
        list_translation = []
        list_num = []
        list_icohp = []
        for bond in range(num_bonds):
            line = data[bond].split()
            icohp = {}
            if version == '2.2.1':
                label = "%s" % (line[0])
                atom1 = str(line[1])
                atom2 = str(line[2])
                length = float(line[3])
                icohp[Spin.up] = float(line[4])
                num = int(line[5])
                translation = [0, 0, 0]
                if self.is_spin_polarized:
                    icohp[Spin.down] = float(data[bond + num_bonds + 1].split()[4])

            elif version == '3.1.1':
                label = "%s" % (line[0])
                atom1 = str(line[1])
                atom2 = str(line[2])
                length = float(line[3])
                translation = [int(line[4]), int(line[5]), int(line[6])]
                icohp[Spin.up] = float(line[7])
                num = int(1)

                if self.is_spin_polarized:
                    icohp[Spin.down] = float(data[bond + num_bonds + 1].split()[7])

            list_labels.append(label)
            list_atom1.append(atom1)
            list_atom2.append(atom2)
            list_length.append(length)
            list_translation.append(translation)
            list_num.append(num)
            list_icohp.append(icohp)

        # to avoid circular dependencies
        from pymatgen.electronic_structure.cohp import IcohpCollection
        self._icohpcollection = IcohpCollection(are_coops=are_coops, list_labels=list_labels, list_atom1=list_atom1,
                                                list_atom2=list_atom2, list_length=list_length,
                                                list_translation=list_translation, list_num=list_num,
                                                list_icohp=list_icohp, is_spin_polarized=self.is_spin_polarized)

    @property
    def icohplist(self):
        """
        Returns: icohplist compatible with older version of this class
        """
        icohplist_new = {}
        for key, value in self._icohpcollection._icohplist.items():
            icohplist_new[key] = {"length": value._length, "number_of_bonds": value._num,
                                  "icohp": value._icohp, "translation": value._translation}
        return icohplist_new

    @property
    def icohpcollection(self):
        """
        Returns: IcohpCollection object
        """
        return self._icohpcollection

class Doscar:
    """
    Class to deal with Lobster's projected DOS and local projected DOS.
    The beforehand quantum-chemical calculation was performed with VASP

    Args:
        doscar: DOSCAR filename, typically "DOSCAR.lobster"
        vasprun: vasprun filename, typically "vasprun.xml"
        dftprogram: so far only "vasp" is implemented

    .. attribute:: completedos

        LobsterCompleteDos Object

    .. attribute:: pdos
        List of Dict including numpy arrays with pdos. Access as pdos[atomindex]['orbitalstring']['Spin.up/Spin.down']

    .. attribute:: tdos
        Dos Object of the total density of states

    .. attribute:: energies
        numpy array of the energies at which the DOS was calculated (in eV, relative to Efermi)

    .. attribute:: tdensities
        tdensities[Spin.up]: numpy array of the total density of states for the Spin.up contribution at each of the energies
        tdensities[Spin.down]: numpy array of the total density of states for the Spin.down contribution at each of the energies

        if is_spin_polarized=False:
        tdensities[Spin.up]: numpy array of the total density of states

    .. attribute:: is_spin_polarized
        Boolean. Tells if the system is spin polarized

    """

    def __init__(self, doscar="DOSCAR.lobster", vasprun="vasprun.xml", dftprogram="Vasp"):

        self._doscar = doscar
        if dftprogram == "Vasp":
            self._vasprun = vasprun
            self._VASPRUN = Vasprun(filename=self._vasprun, ionic_step_skip=None,
                                    ionic_step_offset=0, parse_dos=False,
                                    parse_eigen=False, parse_projected_eigen=False,
                                    parse_potcar_file=False, occu_tol=1e-8,
                                    exception_on_bad_xml=True)
            self._final_structure = self._VASPRUN.final_structure
            self._is_spin_polarized = self._VASPRUN.is_spin

        self._parse_doscar()

    def _parse_doscar(self):
        doscar = self._doscar

        tdensities = {}
        f = open(doscar)
        natoms = int(f.readline().split()[0])
        efermi = float([f.readline() for nn in range(4)][3].split()[17])
        dos = []
        orbitals = []
        for atom in range(natoms + 1):
            line = f.readline()
            ndos = int(line.split()[2])
            orbitals.append(line.split(';')[-1].split())
            line = f.readline().split()
            cdos = np.zeros((ndos, len(line)))
            cdos[0] = np.array(line)
            for nd in range(1, ndos):
                line = f.readline().split()
                cdos[nd] = np.array(line)
            dos.append(cdos)
        f.close()

        doshere = np.array(dos[0])
        energies = doshere[:, 0]
        if not self._is_spin_polarized:
            tdensities[Spin.up] = doshere[:, 1]
            pdoss = []
            spin = Spin.up
            for atom in range(natoms):
                pdos = defaultdict(dict)
                data = dos[atom + 1]
                _, ncol = data.shape
                orbnumber = 0
                for j in range(1, ncol):
                    orb = orbitals[atom + 1][orbnumber]
                    pdos[orb][spin] = data[:, j]
                    orbnumber = orbnumber + 1
                pdoss.append(pdos)
        else:
            tdensities[Spin.up] = doshere[:, 1]
            tdensities[Spin.down] = doshere[:, 2]
            pdoss = []
            for atom in range(natoms):
                pdos = defaultdict(dict)
                data = dos[atom + 1]
                _, ncol = data.shape
                orbnumber = 0
                for j in range(1, ncol):
                    if j % 2 == 0:
                        spin = Spin.down
                    else:
                        spin = Spin.up
                    orb = orbitals[atom + 1][orbnumber]
                    pdos[orb][spin] = data[:, j]
                    if j % 2 == 0:
                        orbnumber = orbnumber + 1
                pdoss.append(pdos)
        self._efermi = efermi
        self._pdos = pdoss
        self._tdos = Dos(efermi, energies, tdensities)
        self._energies = energies
        self._tdensities = tdensities
        final_struct = self._final_structure

        pdossneu = {final_struct[i]: pdos for i, pdos in enumerate(self._pdos)}

        self._completedos = LobsterCompleteDos(final_struct, self._tdos, pdossneu)

    @property
    def completedos(self):
        return self._completedos

    @property
    def pdos(self):
        return self._pdos

    @property
    def tdos(self):
        return self._tdos

    @property
    def energies(self):
        return self._energies

    @property
    def tdensities(self):
        return self._tdensities

    @property
    def is_spin_polarized(self):
        return self._is_spin_polarized

class Charge:
    """Class to read CHARGE files generated by LOBSTER
        Args:
            filename: filename for the CHARGE file, typically "CHARGE.lobster"

        .. attribute: atomlist
            List of atoms in CHARGE.lobster
        .. attribute: types
            List of types of atoms in CHARGE.lobster
        .. attribute: Mulliken
            List of Mulliken charges of atoms in CHARGE.lobster
        .. attribute: Loewdin
            List of Loewdin charges of atoms in CHARGE.Loewdin
        .. attribute: num_atoms
            Number of atoms in CHARGE.lobster

    """

    def __init__(self, filename="CHARGE.lobster"):
        with zopen(filename) as f:
            data = f.read().split("\n")[3:-3]
        if len(data) == 0:
            raise IOError("CHARGES file contains no data.")

        self.num_atoms = len(data)
        self.atomlist = []
        self.types = []
        self.Mulliken = []
        self.Loewdin = []
        for atom in range(0, self.num_atoms):
            line = data[atom].split()
            self.atomlist.append(line[1] + line[0])
            self.types.append(line[1])
            self.Mulliken.append(float(line[2]))
            self.Loewdin.append(float(line[3]))

    def get_structure_with_charges(self, structure_filename):
        """
        get a Structure with Mulliken and Loewdin charges as site properties
        Args:
            structure_filename: filename of POSCAR
        Returns:
            Structure Object with Mulliken and Loewdin charges as site properties
        """

        struct = Structure.from_file(structure_filename)
        Mulliken = self.Mulliken
        Loewdin = self.Loewdin
        site_properties = {"Mulliken Charges": Mulliken, "Loewdin Charges": Loewdin}
        new_struct = struct.copy(site_properties=site_properties)
        return new_struct

class Lobsterout:
    """
    Class to read in the lobsterout and evaluate the spilling, save the basis, save warnings, save infos
    Args:
      filename: filename of lobsterout

      .. attribute: basis_functions
        list of basis functions that were used in lobster run as strings

      .. attribute: basis_type
         list of basis type that were used in lobster run as strings

      .. attribute: chargespilling
         list of charge spilling (first entry: result for spin 1, second entry: result for spin 2 or not present)

      .. attribute: dftprogram
         string representing the dft program used for the calculation of the wave function

      .. attribute: elements
         list of strings of elements that were present in lobster calculation

      .. attribute: has_CHARGE
         Boolean, indicates that CHARGE.lobster is present

      .. attribute: has_COHPCAR
        Boolean, indicates that COHPCAR.lobster and ICOHPLIST.lobster are present

      .. attribute: has_COOPCAR
        Boolean, indicates that COOPCAR.lobster and ICOOPLIST.lobster are present

      .. attribute: has_DOSCAR
        Boolean, indicates that DOSCAR.lobster is present

      .. attribute: has_Projection
        Boolean, indcates that projectionData.lobster is present

      .. attribute: has_bandoverlaps
        Boolean, indcates that bandOverlaps.lobster is present

      .. attribute: has_density_of_energies
        Boolean, indicates that DensityOfEnergy.lobster is present

      .. attribute: has_fatbands
        Boolean, indicates that fatband calculation was performed

      .. attribute: has_grosspopulation
        Boolean, indicates that GROSSPOP.lobster is present

      .. attribute: info_lines
        string with additional infos on the run

      .. attribute: info_orthonormalization
        string with infos on orthonormalization

      .. attribute: is_restart_from_projection
        Boolean that indicates that calculation was restartet from existing projection file

      .. attribute: lobster_version
        string that indicates Lobster version

      .. attribute: number_of_spins
        Integer indicating the number of spins

      .. attribute: number_of_threads
       integer that indicates how many threads were used

      .. attribute: timing
        dict with infos on timing

      .. attribute: totalspilling
        list of values indicating the total spilling for spin channel 1 (and spin channel 2)

      .. attribute: warninglines
        string with all warnings

    """

    def __init__(self, filename="lobsterout"):

        warnings.warn("Make sure the lobsterout is read in correctly. This is a brand new class.")
        # read in file
        with zopen(filename) as f:
            data = f.read().split("\n")  # [3:-3]
        if len(data) == 0:
            raise IOError("lobsterout does not contain any data")

        LobsterDict = {}
        # check if Lobster starts from a projection
        self.is_restart_from_projection = self._starts_from_projection(data=data)

        self.lobster_version = self._get_lobster_version(data=data)

        self.number_of_threads = int(self._get_threads(data=data))
        self.dftprogram = self._get_dft_program(data=data)

        self.number_of_spins = self._get_number_of_spins(data=data)
        chargespilling, totalspilling = self._get_spillings(data=data, number_of_spins=self.number_of_spins)
        self.chargespilling = chargespilling
        self.totalspilling = totalspilling

        elements, basistype, basisfunctions = self._get_elements_basistype_basisfunctions(data=data)
        self.elements = elements
        self.basis_type = basistype
        self.basis_functions = basisfunctions

        wall_time, user_time, sys_time = self._get_timing(data=data)
        timing = {}
        timing['walltime'] = wall_time
        timing['usertime'] = user_time
        timing['sys_time'] = sys_time
        self.timing = timing

        warninglines = self._get_all_warning_lines(data=data)
        self.warninglines = warninglines

        orthowarning = self._get_warning_orthonormalization(data=data)
        self.info_orthonormalization = orthowarning
        # print(orthowarning)

        infos = self._get_all_info_lines(data=data)
        self.info_lines = infos

        self.has_DOSCAR = self._has_DOSCAR(data=data)
        self.has_COHPCAR = self._has_COOPCAR(data=data)
        self.has_COOPCAR = self._has_COHPCAR(data=data)
        self.has_CHARGE = self._has_CHARGE(data=data)
        self.has_Projection = self._has_projection(data=data)
        self.has_bandoverlaps = self._has_bandoverlaps(data=data)
        self.has_fatbands = self._has_fatband(data=data)
        self.has_grosspopulation = self._has_grosspopulation(data=data)
        self.has_density_of_energies = self._has_density_of_energies(data=data)

    def get_doc(self):
        """

        Returns: LobsterDict with all the information stored in lobsterout
        """

        LobsterDict = {}
        # check if Lobster starts from a projection
        LobsterDict['restart_from_projection'] = self.is_restart_from_projection
        LobsterDict['lobster_version'] = self.lobster_version
        LobsterDict['threads'] = self.number_of_threads
        LobsterDict['Dftprogram'] = self.dftprogram

        number_of_spins = self.number_of_spins
        LobsterDict['chargespilling'] = self.chargespilling
        LobsterDict['totalspilling'] = self.totalspilling

        LobsterDict['elements'] = self.elements
        LobsterDict['basistype'] = self.basis_type
        LobsterDict['basisfunctions'] = self.basis_functions

        LobsterDict['timing'] = self.timing

        LobsterDict['warnings'] = self.warninglines

        LobsterDict['orthonormalization'] = self.info_orthonormalization

        LobsterDict['infos'] = self.info_lines

        LobsterDict['hasDOSCAR'] = self.has_DOSCAR
        LobsterDict['hasCOHPCAR'] = self.has_COHPCAR
        LobsterDict['hasCOOPCAR'] = self.has_COOPCAR
        LobsterDict['hasCHARGE'] = self.has_CHARGE
        LobsterDict['hasProjection'] = self.has_Projection
        LobsterDict['hasbandoverlaps'] = self.has_bandoverlaps
        LobsterDict['hasfatband'] = self.has_fatbands
        LobsterDict['hasGrossPopuliation'] = self.has_grosspopulation
        LobsterDict['hasDensityOfEnergies'] = self.has_density_of_energies

        return LobsterDict

    def _get_lobster_version(self, data):
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 1:
                if splitrow[0] == "LOBSTER":
                    return splitrow[1]

    def _has_bandoverlaps(self, data):
        if 'WARNING: I dumped the band overlap matrices to the file bandOverlaps.lobster.' in data:
            return True
        else:
            return False

    def _starts_from_projection(self, data):
        if 'loading projection from projectionData.lobster...' in data:
            return True
        else:
            return False

    def _has_DOSCAR(self, data):
        if 'writing DOSCAR.lobster...' in data and not 'SKIPPING writing DOSCAR.lobster...' in data:
            return True
        else:
            return False

    def _has_COOPCAR(self, data):
        if 'writing COOPCAR.lobster and ICOOPLIST.lobster...' in data and not 'SKIPPING writing COOPCAR.lobster and ICOOPLIST.lobster...' in data:
            return True
        else:
            return False

    def _has_COHPCAR(self, data):
        if 'writing COHPCAR.lobster and ICOHPLIST.lobster...' in data and not 'SKIPPING writing COHPCAR.lobster and ICOHPLIST.lobster...' in data:
            return True
        else:
            return False

    def _has_CHARGE(self, data):
        # weitere optionen testen -> auch hier kann uebersprungen werden
        if not 'SKIPPING writing CHARGE.lobster...' in data:
            return True
        else:
            return False

    def _has_grosspopulation(self, data):
        if 'writing CHARGE.lobster and GROSSPOP.lobster...' in data:
            return True
        else:
            return False

    def _has_projection(self, data):
        if 'saving projection to projectionData.lobster...' in data:
            return True
        else:
            return False

    def _has_fatband(self, data):
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 1:
                if splitrow[1] == 'FatBand':
                    return True
        return False

    def _has_density_of_energies(self, data):
        if "writing DensityOfEnergy.lobster..." in data:
            return True
        else:
            return False

    def _get_dft_program(self, data):
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 4:
                if splitrow[3] == "program...":
                    return splitrow[4]

    def _get_number_of_spins(self, data):
        if "spillings for spin channel 2" in data:
            # print('two spin channels')
            return 2
        else:
            return 1

    def _get_threads(self, data):
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 11:
                if (splitrow[11]) == "threads":
                    return splitrow[10]

    def _get_spillings(self, data, number_of_spins):
        charge_spilling = []
        total_spilling = []

        for row in data:
            splitrow = row.split()

            if len(splitrow) > 2:
                if splitrow[2] == 'spilling:':
                    # print(splitrow)
                    if splitrow[1] == 'charge':
                        charge_spilling.append(np.float(splitrow[3].replace('%', '')) / 100.0)
                    if splitrow[1] == 'total':
                        total_spilling.append(np.float(splitrow[3].replace('%', '')) / 100.0)

            if len(charge_spilling) == number_of_spins and len(total_spilling) == number_of_spins:
                break;

        return charge_spilling, total_spilling

    def _get_elements_basistype_basisfunctions(self, data):
        begin = False
        end = False
        elements = []
        basistype = []
        basisfunctions = []
        for row in data:

            if begin and not end:
                # print(row)
                splitrow = row.split()
                if splitrow[0] not in ['INFO:', 'WARNING:', 'setting', 'calculating', 'post-processing', 'saving',
                                       'spillings', 'writing']:

                    elements.append(splitrow[0])
                    basistype.append(splitrow[1].replace('(', '').replace(')', ''))
                    # last sign is a ''
                    basisfunctions.append(splitrow[2:])
                else:
                    end = True
            if "setting up local basis functions..." in row:
                begin = True
                # print(row)
        return elements, basistype, basisfunctions

    def _get_timing(self, data):
        # will give back wall, user and sys time
        begin = False
        # end=False
        # time=[]

        for row in data:
            splitrow = row.split()
            if 'finished' in splitrow:
                begin = True
            if begin:
                if 'wall' in splitrow:
                    wall_time = (splitrow[2:10])
                if 'user' in splitrow:
                    user_time = (splitrow[0:8])
                if 'sys' in splitrow:
                    sys_time = (splitrow[0:8])

        wall_time_dict = {"h": wall_time[0], "min": wall_time[2], "s": wall_time[4], "ms": wall_time[6]}
        user_time_dict = {"h": user_time[0], "min": user_time[2], "s": user_time[4], "ms": user_time[6]}
        sys_time_dict = {"h": sys_time[0], "min": sys_time[2], "s": sys_time[4], "ms": sys_time[6]}

        return wall_time_dict, user_time_dict, sys_time_dict

    def _get_warning_orthonormalization(self, data):
        orthowarning = []
        for row in data:
            splitrow = row.split()
            if 'orthonormalized' in splitrow:
                orthowarning.append(" ".join(splitrow[1:]))
        return orthowarning

    def _get_all_warning_lines(self, data):
        warnings = []
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 0:
                if splitrow[0] == 'WARNING:':
                    warnings.append(" ".join(splitrow[1:]))

        return warnings
        # print(warnings)

    def _get_all_info_lines(self, data):
        infos = []
        for row in data:
            splitrow = row.split()
            if len(splitrow) > 0:
                if splitrow[0] == 'INFO:':
                    infos.append(" ".join(splitrow[1:]))

        return infos

class Fatband:
    """
    reads in FATBAND_x_y.lobster files
    Args:
      filenames (list or string): can be a list of file names or a path to a folder folder from which all "FATBAND_*" files will be read
      vasprun: corresponding vasprun file
      Kpointsfile: KPOINTS file for bandstructure calculation, typically "KPOINTS"

      .. attribute: efermi
      efermi that was read in from vasprun.xml

      .. attribute: eigenvals
      {Spin.up:[][],Spin.down:[][]}, the first index of the array
            [][] refers to the band and the second to the index of the
            kpoint. The kpoints are ordered according to the order of the
            kpoints array. If the band structure is not spin polarized, we
            only store one data set under Spin.up.

      .. attribute: is_spinpolarized
      Boolean that tells you whether this was a spin-polarized calculation

      .. attribute: kpoints_array
        list of kpoint as numpy arrays, in frac_coords of the
            given lattice by default

      .. attribute: label_dict
     (dict) of {} this link a kpoint (in frac coords or
            cartesian coordinates depending on the coords).

      .. attribute: lattice
      lattice object of reciprocal lattice as read in from vasprun.xml

      .. attribute: nbands
      number of bands used in the calculation

      .. attribute: p_eigenvals
      dict of orbital projections as {spin: array of dict}.
      The indices of the array are [band_index, kpoint_index].
      The dict is then built the following way:
      {"string of element": "string of orbital as read in from FATBAND file"}
      If the band structure is not spin polarized, we only store one data set under Spin.up.

      .. attribute: structure
      structure read in from vasprun.xml

    """

    def __init__(self, filenames=".", vasprun='vasprun.xml', Kpointsfile='KPOINTS'):

        warnings.warn('Make sure all relevant FATBAND files were generated and read in!')
        warnings.warn('Use Lobster 3.2.0 or newer for fatband calculations!')

        VASPRUN = Vasprun(filename=vasprun, ionic_step_skip=None,
                          ionic_step_offset=0, parse_dos=True,
                          parse_eigen=False, parse_projected_eigen=False,
                          parse_potcar_file=False, occu_tol=1e-8,
                          exception_on_bad_xml=True)
        self.structure = VASPRUN.final_structure
        self.lattice = self.structure.lattice.reciprocal_lattice
        self.efermi = VASPRUN.efermi
        kpoints_object = Kpoints.from_file(Kpointsfile)

        atomtype = []
        atomnames = []
        orbital_names = []

        if type(filenames) is not type([]) or filenames is None:
            filenames_new = []
            if filenames is None:
                filenames = '.'
            for file in os.listdir(filenames):
                if fnmatch.fnmatch(file, 'FATBAND_*.lobster'):
                    filenames_new.append(os.path.join(filenames, file))
            filenames = filenames_new
        if (len(filenames)) == 0:
            raise ValueError("No FATBAND files in folder or given")
        for ifilename, filename in enumerate(filenames):
            with zopen(filename, "rt") as f:
                contents = f.read().split("\n")

            # TODO: could be replaced for future versions of Lobster, get atomname from filename
            atomnames.append(os.path.split(filename)[1].split('_')[1].capitalize())
            parameters = contents[0].split()
            atomtype.append(re.split(r"[0-9]+", parameters[3])[0].capitalize())
            orbital_names.append(parameters[4])

        # get atomtype orbital dict
        atom_orbital_dict = {}
        for iatom, atom in enumerate(atomnames):
            if atom not in atom_orbital_dict:
                atom_orbital_dict[atom] = []
            atom_orbital_dict[atom].append(orbital_names[iatom])
        # test if there are the same orbitals twice or if two different formats were used or if all necessary orbitals are there
        for key, items in atom_orbital_dict.items():
            if len(set(items)) != len(items):
                raise (ValueError("The are two FATBAND files for the same atom and orbital. The program will stop."))
            split = []
            for item in items:
                split.append(item.split("_")[0])
            for orb, number in collections.Counter(split).items():
                if number != 1 and number != 3 and number != 5 and number != 7:
                    raise (ValueError(
                        "Make sure all relevant orbitals were generated and that no duplicates (2p and 2p_x) are present"))

        kpoints_array = []
        for ifilename, filename in enumerate(filenames):
            with zopen(filename, "rt") as f:
                contents = f.read().split("\n")

            if ifilename == 0:
                self.nbands = int(parameters[6])
                self.number_kpts = kpoints_object.num_kpts - int(contents[1].split()[2]) + 1

            if len(contents[1:]) == self.nbands + 2:
                self.is_spinpolarized = False
            elif len(contents[1:]) == self.nbands * 2 + 2:
                self.is_spinpolarized = True
            else:
                linenumbers = []
                for iline, line in enumerate(contents[1:self.nbands * 2 + 4]):
                    # print(line)
                    # if line in ['\n', '\r\n']:
                    #    linenumbers.append(iline)
                    if line.split()[0] == '#':
                        linenumbers.append(iline)

                if ifilename == 0:
                    if len(linenumbers) == 2:
                        self.is_spinpolarized = True
                    else:
                        self.is_spinpolarized = False

            if ifilename == 0:
                eigenvals = {}
                eigenvals[Spin.up] = [[collections.defaultdict(float)
                                       for i in range(self.number_kpts)]
                                      for j in range(self.nbands)]
                if self.is_spinpolarized:
                    eigenvals[Spin.down] = [[collections.defaultdict(float)
                                             for i in range(self.number_kpts)]
                                            for j in range(self.nbands)]

                p_eigenvals = {}
                p_eigenvals[Spin.up] = [
                    [{str(e): {str(orb): collections.defaultdict(float) for orb in atom_orbital_dict[e]}
                      for e in atomnames}
                     for i in range(self.number_kpts)]
                    for j in range(self.nbands)]

                if self.is_spinpolarized:
                    p_eigenvals[Spin.down] = [
                        [{str(e): {str(orb): collections.defaultdict(float) for orb in atom_orbital_dict[e]}
                          for e in atomnames}
                         for i in range(self.number_kpts)]
                        for j in range(self.nbands)]

            ikpoint = -1
            for iline, line in enumerate(contents[1:-1]):
                if line.split()[0] == '#':
                    Kpointnumber = int(line.split()[2])
                    KPOINT = np.array([float(line.split()[4]), float(line.split()[5]), float(line.split()[6])])
                    if ifilename == 0:
                        kpoints_array.append(KPOINT)

                    linenumber = 0
                    iband = 0
                    ikpoint += 1
                if linenumber == self.nbands:
                    iband = 0
                if line.split()[0] != '#':

                    if linenumber < self.nbands:
                        if ifilename == 0:
                            eigenvals[Spin.up][iband][ikpoint] = float(line.split()[1]) + self.efermi

                        p_eigenvals[Spin.up][iband][ikpoint][atomnames[ifilename]][orbital_names[ifilename]] = float(
                            line.split()[2])
                    if linenumber >= self.nbands and self.is_spinpolarized:
                        # print(line.split())
                        # print(iband)
                        if ifilename == 0:
                            eigenvals[Spin.down][iband][ikpoint] = float(line.split()[1]) + self.efermi
                        p_eigenvals[Spin.down][iband][ikpoint][atomnames[ifilename]][
                            orbital_names[ifilename]] = float(line.split()[2])

                    linenumber += 1
                    iband += 1

        self.kpoints_array = kpoints_array
        self.eigenvals = eigenvals
        self.p_eigenvals = p_eigenvals

        label_dict = {}
        for ilabel, label in enumerate(kpoints_object.labels[-self.number_kpts:], start=0):

            if label is not None:
                label_dict[label] = kpoints_array[ilabel]

        self.label_dict = label_dict

    def get_bandstructure(self):
        """
        returns a LobsterBandStructureSymmLine object which can be plotted with a normal BSPlotter
        """

        return LobsterBandStructureSymmLine(kpoints=self.kpoints_array, eigenvals=self.eigenvals, lattice=self.lattice,
                                            efermi=self.efermi, labels_dict=self.label_dict,
                                            structure=self.structure,
                                            projections=self.p_eigenvals)


#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Mon Jun  4 13:04:11 2018

@author: andrew
"""

import glob
import os
import psf
import numpy as np
import sys
from astropy.io import fits
import initialize

def hotpants(location):
    '''
    subtract using hotpants
    '''
    images = glob.glob(location + "/data/*_A_.fits")
    template = glob.glob(location + "/templates/*.fits")
    residuals = glob.glob(location + "/residuals/*residual_.fits")
    psf_data = glob.glob(location + '/psf/*')
    images_names = [(i.split('/')[-1])[:-5] for i in images]
    res_names = [(r.split('/')[-1])[:-14] for r in residuals]
    resids = [res for res in images_names if res not in res_names]
    ims = []
    for rs in resids:
        ims.append(location+'/data/'+rs+'.fits')
    if ims != []:
        if len(psf_data) == 2*(len(images)+1):
            if len(template) == 1:
                for i in ims:
                    hotpants_subtract(i, default=True)
            else:
                print("-> Error with number of templates")
                sys.exit()
        else:
            print("-> Error: Need PSFs before running subtraction\n-> Run psf.py first")
            print("-> If any images have been manually removed from the data directory, delete all contents of the psf directory and run OasisPy again\n")
            sys.exit()
    else:
        print("-> Images have already been subtracted")

def hotpants_subtract(image, default=True, opt_ind=0, spatial_deg=0, params=None):
    location = image.split('/')[:-2]
    location = '/'.join(location)
    name = (image.split('/')[-1])[:-5]
    template = glob.glob("%s/templates/*.fits" % location)
    template = template[0]
    image_header = fits.getheader(image)
    sig_image = psf.fwhm(image)/2.355
    template_header = fits.getheader(template)
    sig_template = psf.fwhm_template(template)/2.355
    tu = template_header['SATURATE']
    tg = template_header['GAIN']
    tmi = "%s/hotpants_mask.fits" % location
    if os.path.exists(tmi) == False:
        try: weight_check = fits.getval(template, 'WEIGHT')
        except: weight_check = 'N'
        temp_mask = fits.getdata(template, 1)
        if weight_check == 'Y':
            temp_mask = (temp_mask - 1) * -1
        hotpants_hdu = fits.PrimaryHDU(temp_mask)
        hotpants_hdu.writeto("%s/hotpants_mask.fits" % (location), overwrite=True)
    iu = image_header['SATURATE']
    ig = image_header['GAIN']
    try: weight_check = fits.getval(image, 'WEIGHT')
    except: weight_check = 'N'
    if weight_check == 'Y':
        im_mask = fits.getdata(image, 1)
        im_mask = (im_mask-1)*-1
        im_mask_hdu = fits.PrimaryHDU(im_mask)
        im_mask_hdu.writeto("%s/configs/temp_htpts_mask.fits" % (location), overwrite=True)
        imi = "%s/configs/temp_htpts_mask.fits" % (location)
    else:
        imi = "%s[1]" % (image)
    cwd = os.getcwd()
    os.chdir(os.path.dirname(initialize.__file__))
    if default == True:
        omi = "%s/residuals/BPM.fits" % location
        outim = "%s/residuals/%sresidual_.fits" % (location, name)
        tl = -1000
        il = -1000
        if sig_template < sig_image:
            sigma_match = np.sqrt((sig_image)**2-(sig_template)**2)
            s1 = .5*sigma_match
            s2 = sigma_match
            s3 = 2*sigma_match
            os.system("./hotpants -inim %s -tmplim %s -outim %s -tl %d -il %d -tu %d -tg %d -tmi %s -iu %d -ig %d -imi %s -omi %s -ng 3 6 %.5f 4 %.5f 2 %.5f -bgo 0" % (image, template, outim, tl, il, tu, tg, tmi, iu, ig, imi, omi, s1, s2, s3))
        elif sig_template >= sig_image:
            sigma_match = np.sqrt((sig_template)**2-(sig_image)**2)
            s1 = .5*sigma_match
            s2 = sigma_match
            s3 = 2*sigma_match
            os.system("./hotpants -inim %s -tmplim %s -outim %s -tl %d -il %d -tu %d -tg %d -tmi %s -iu %d -ig %d -imi %s -omi %s -ng 3 6 %.5f 4 %.5f 2 %.5f -bgo 0" % (template, image, outim, il, tl, iu, ig, imi, tu, tg, tmi, omi, s1, s2, s3))
            invert(outim)
        res_data = fits.getdata(outim)
        res_header = fits.getheader(outim)
        BPM = fits.getdata(omi)
        BPM_median = np.median(BPM)
        BPM_std = np.std(BPM)
        BPM[BPM > (BPM_median + BPM_std)] = 0
        BPM[BPM <= (BPM_median + BPM_std)] = 1
        resHDU = fits.PrimaryHDU(res_data, header=res_header)
        resMask = fits.ImageHDU(BPM)
        resList = fits.HDUList([resHDU, resMask])
        resList.writeto(outim, overwrite=True)
        os.remove(omi)
        os.remove("%s/configs/temp_htpts_mask.fits" % location)
        hdu = fits.open("%s/residuals/%sresidual_.fits" % (location, name), mode='update')
        hdr = hdu[0].header
        hdr.set('OPTIMIZE', 'N')
        hdu.close()
    elif default == False:
        omi = "%s/temp/BPM.fits" % location
        outim = "%s/temp/conv.fits" % location
        if 0 <= opt_ind <= 1:
            s1, s2, s3 = 0.7, 1, 2.5
        elif opt_ind == 2:
            s1, s2, s3 = 0.7, 1.5, 3
        elif opt_ind >= 3:
            s1, s2, s3 = 0.7, 2, 4
        if sig_template < sig_image:
            os.system("./hotpants -inim %s -tmplim %s -outim %s -tl %d -il %d -tu %d -tg %d -tmi %s -iu %d -ig %d -imi %s -omi %s -ng 3 6 %.5f 4 %.5f 2 %.5f -bgo %d -r %d -nsx %d -nsy %d -rss %d -ko %d" % (image, template, outim, params['tl'][0], params['il'][0], tu, tg, tmi, iu, ig, imi, omi, s1, s2, s3, params['bgo'][0], (params['r'])[opt_ind], (params['nsx'])[spatial_deg], (params['nsy'])[spatial_deg], (params['rss'])[opt_ind], (params['ko'])[spatial_deg]))
        else:
            os.system("./hotpants -inim %s -tmplim %s -outim %s -tl %d -il %d -tu %d -tg %d -tmi %s -iu %d -ig %d -imi %s -omi %s -ng 3 6 %.5f 4 %.5f 2 %.5f -bgo %d -r %d -nsx %d -nsy %d -rss %d -ko %d" % (template, image, outim, params['il'][0], params['tl'][0], iu, ig, imi, tu, tg, tmi, omi, s1, s2, s3, params['bgo'][0], (params['r'])[opt_ind], (params['nsx'])[spatial_deg], (params['nsy'])[spatial_deg], (params['rss'])[opt_ind], (params['ko'])[spatial_deg]))
            invert(outim)
    else:
        print("\n-> Error: Default must be boolean value\n-> Exiting...")
        os.chdir(cwd)
        sys.exit()
    os.chdir(cwd)

def invert(image):
    data = fits.getdata(image)
    try: mask = fits.getdata(image, 1)
    except: mask = None
    hdr = fits.getheader(image)
    if mask == None:
        hdu = fits.PrimaryHDU(data*-1, header=hdr)
    else:
        hduData = fits.PrimaryHDU(data*-1, header=hdr)
        hduMask = fits.ImageHDU(mask)
        hdu = fits.HDUList([hduData, hduMask])
    hdu.writeto(image, overwrite=True)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Aug 17 04:13:10 2021

@author: erik.edmonds
"""
import os
from tqdm import tqdm
import numpy as np
import pandas as pd
import cv2
import time
import re

from deepface.detectors import FaceDetector
from deepface import DeepFace
from deepface.extendedmodels import Age
from deepface.commons import functions, realtime, distance as dst
#------------------------

def operation(db_path=''):
    face_detector = FaceDetector.build_model('opencv')
    print("Detector backend is ", 'opencv')

       	#------------------------

    input_shape = (224, 224); input_shape_x = input_shape[0]; input_shape_y = input_shape[1]

    text_color = (255,255,255)

    employees = []
       	#check passed db folder exists
    if os.path.isdir(db_path) == True:
        for r, d, f in os.walk(db_path): # r=root, d=directories, f = files
            for file in f:
                if ('.jpg' in file):
                    exact_path = r + "/" + file
                    employees.append(exact_path)

    if len(employees) == 0:
        print("WARNING: There is no image in this path ( ", db_path,") . Face recognition will not be performed.")

    	#------------------------

    if len(employees) > 0:

        model = DeepFace.build_model('VGG-Face')
        print('VGG-Face'," is built")

    		#------------------------

        input_shape = functions.find_input_shape(model)
        input_shape_x = input_shape[0]; input_shape_y = input_shape[1]

        threshold = dst.findThreshold('VGG-Face', 'VGG-Face')

    	#------------------------
    	#facial attribute analysis models
    tic = time.time()

    emotion_model = DeepFace.build_model('Emotion')
    print("Emotion model loaded")

    age_model = DeepFace.build_model('Age')
    print("Age model loaded")

    gender_model = DeepFace.build_model('Gender')
    print("Gender model loaded")
    race_model = DeepFace.build_model('Race')
    print("Race model loaded")
    toc = time.time()

    print("Facial attibute analysis models loaded in ",toc-tic," seconds")

    	#------------------------

    	#find embeddings for employee list

    tic = time.time()

    	#-----------------------

    pbar = tqdm(range(0, len(employees)), desc='Finding embeddings')

    	#TODO: why don't you store those embeddings in a pickle file similar to find function?

    embeddings = []
    	#for employee in employees:
    for index in pbar:
        employee = employees[index]
        pbar.set_description("Finding embedding for %s" % (employee.split("/")[-1]))
        embedding = []

        img = functions.preprocess_face(img = employee, target_size = (input_shape_y, input_shape_x), enforce_detection = False, detector_backend= 'opencv')
        img_representation = model.predict(img)[0,:]

        embedding.append(employee)
        embedding.append(img_representation)
        embeddings.append(embedding)

    df = pd.DataFrame(embeddings, columns = ['employee', 'embedding'])
    df['VGG-Face'] = 'VGG-Face'

    toc = time.time()

    print("Embeddings found for given data set in ", toc-tic," seconds")

    	#-----------------------

    pivot_img_size = 112 #face recognition result image

    	#-----------------------

    freeze = False
    face_detected = False
    face_included_frames = 0
    freezed_frame = 0
    tic = time.time()

    cap = cv2.VideoCapture(0) #webcam

    while(True):
        ret, img = cap.read()

        if img is None:
            break

        raw_img = img.copy()
        resolution = img.shape; resolution_x = img.shape[1]; resolution_y = img.shape[0]

        if freeze == False:
            faces = FaceDetector.detect_faces(face_detector, 'opencv', img, align = False)

            if len(faces) == 0:
                face_included_frames = 0
        else:
            faces = []

        detected_faces = []
        face_index = 0
        for face, (x, y, w, h) in faces:
            if w > 130:
                face_detected = True
                if face_index == 0:
                    face_included_frames = face_included_frames + 1 #increase frame for a single face

                cv2.rectangle(img, (x,y), (x+w,y+h), (67,67,67), 1) #draw rectangle to main image

                cv2.putText(img, str(5 - face_included_frames), (int(x+w/4),int(y+h/1.5)), cv2.FONT_HERSHEY_SIMPLEX, 4, (255, 255, 255), 2)

                detected_face = img[int(y):int(y+h), int(x):int(x+w)] #crop detected face

    				#-------------------------------------

                detected_faces.append((x,y,w,h))
                face_index = face_index + 1

    				#-------------------------------------

        if face_detected == True and face_included_frames == 5 and freeze == False:
            freeze = True
            base_img = raw_img.copy()
            detected_faces_final = detected_faces.copy()
            tic = time.time()

        if freeze == True:

            toc = time.time()
            if (toc - tic) < 5:

                if freezed_frame == 0:
                    freeze_img = base_img.copy()
    					#freeze_img = np.zeros(resolution, np.uint8) #here, np.uint8 handles showing white area issue

                    for detected_face in detected_faces_final:
                        x = detected_face[0]; y = detected_face[1]
                        w = detected_face[2]; h = detected_face[3]

                        cv2.rectangle(freeze_img, (x,y), (x+w,y+h), (67,67,67), 1) #draw rectangle to main image

    						#-------------------------------

                        custom_face = base_img[y:y+h, x:x+w]


"""Файл с основным классом бота, отвечающим за работу бота."""
import time

import numpy as np
import pyscreenshot as sc
from pynput.mouse import Button as MButton, Controller as MController

class Bot:
    def __init__(self, screen, chest, color):
        self.SCREEN = screen  # Место экрана которое будет обрабатываться
        self.CHEST = chest  # Место где находится сундук
        self.COLOR_FOR_FIND = np.array(color)  # цвет который будет искаться на месте обработки
        self.MOUSE_CONTROLLER = MController()

    def bot_execute(self):
        counter = 0
        im = sc.grab(bbox=self.SCREEN)
        array = np.array(im)
        for i in range(array.shape[0]):
            for j in range(array.shape[1]):
                # если искомый пиксель есть в обрабатываемой области экрана, увеличиваем счетчик
                if np.array_equal(array[i][j], self.COLOR_FOR_FIND):
                    counter += 1
        # если в обрабатываемой области достаточно искомых пикселей, перемещаем курсор в нужное места экрана
        # и производим клин
        if counter > 1000:
            mouse_position = self.MOUSE_CONTROLLER.position
            self.MOUSE_CONTROLLER.position = self.CHEST
            self.MOUSE_CONTROLLER.click(MButton.left, 1)
            time.sleep(0.1)
            self.MOUSE_CONTROLLER.position = mouse_position

    def debug(self):
        """
        Метод вызываемый для отладки.

        Производит отображения области экрана которая будет обрабатываться, место куда будет установлен курсор
        и количество пикселей того или иного цвета.
        """
        time.sleep(1)
        im = sc.grab(bbox=self.SCREEN)
        im.show()

        array = np.array(im)
        unique = dict()
        for i in range(array.shape[0]):
            for j in range(array.shape[1]):
                # ToDo when running on a monotonous surface (processing an area of the screen of the same color),
                #  it throws an error
                """Traceback (most recent call last):
                File "/usr/lib64/python3.8/runpy.py", line 193, in _run_module_as_main
                return _run_code(code, main_globals, None,
                File "/usr/lib64/python3.8/runpy.py", line 86, in _run_code
                exec(code, run_globals)
                File "/home/andrey/projects/tes/env/lib/python3.8/site-packages/autochests_bot/__main__.py", line 11, in <module>
                start_bot()
                File "/home/andrey/projects/tes/env/lib/python3.8/site-packages/autochests_bot/start_bot.py", line 15, in start_bot
                bot.debug()
                File "/home/andrey/projects/tes/env/lib/python3.8/site-packages/autochests_bot/bot_controll/chests_bot.py", line 50, in debug
                key = f"{array[i][j][0]},{array[i][j][1]},{array[i][j][2]}"
                IndexError: invalid index to scalar variable"""
                # Todo array print
                """
                [[0 1 1 ... 5 5 5]
                [1 1 1 ... 5 5 5]
                [4 1 1 ... 5 5 5]
                ...
                [5 5 5 ... 5 5 5]
                [5 5 5 ... 5 5 5]
                [5 5 5 ... 5 5 5]]
                """
                # ToDo write error handling
                key = f"{array[i][j][0]},{array[i][j][1]},{array[i][j][2]}"
                if unique.get(key, None):
                    unique[key] += 1
                else:
                    unique[key] = 1

        sorted_unique = {k: v for k, v in sorted(unique.items(), key=lambda item: item[1])}
        print(sorted_unique)

        position = self.MOUSE_CONTROLLER.position
        self.MOUSE_CONTROLLER.position = self.CHEST
        time.sleep(5)
        self.MOUSE_CONTROLLER.position = position


#!/usr/bin/env python
"""Quantize function.
"""

import ctypes
import datetime
import logging
from os.path import abspath, dirname
import torch
import numpy as np
from numba import cuda
from quant_impl import fake_quantize

_USE_GFPQ_QUANT_LIB = (torch.cuda.device_count() <= 1)

class GFPQParamSt(ctypes.Structure):
    r"""GFPQ param, corresponds with struct GFPQ_PARAM_ST in gfpq.hpp"""
    _fields_ = [("mode", ctypes.c_int), ("param", ctypes.c_byte * 16)]

class _types:
    r"""Some alias types."""
    handle = ctypes.c_void_p
    stream = ctypes.c_void_p

class QuantAndDeQuantGPU():
    r"""quantize and dequantize data with GFPG library.
    """
    def __init__(self,
                 libquant_path=dirname(abspath(__file__)) +
                 "/gpu/lib/libgfpq_gpu.so",
                 libcublas_path="libcublas.so",
                 bit_width=8,
                 param_mode=0):
        global _USE_GFPQ_QUANT_LIB
        self._bit_width = bit_width
        if _USE_GFPQ_QUANT_LIB:
            self._libquant = ctypes.cdll.LoadLibrary(libquant_path)
            self._libcublas = ctypes.cdll.LoadLibrary(libcublas_path)
            self._libcublas.cublasCreate_v2.restype = int
            self._libcublas.cublasCreate_v2.argtypes = [ctypes.c_void_p]
            self._cublas_handle = _types.handle()
            self._libcublas.cublasCreate_v2(ctypes.byref(self._cublas_handle))
            self._param = GFPQParamSt()
            self._stream = cuda.stream()
            self._param.mode = param_mode

    def __call__(self, tensor, mode=0):
        r""" Converts float weights to quantized weights.

        Args:
            - tensor: input data
            - mode: GFPQ mode for param
                GFPQ_MODE_INIT(0): There is no valid parameter in param[].
                    Generate the parameter and filled in param[].
                GFPQ_MODE_UPDATE(1): There is parameter in param[]. Generate
                    new parameter, update param[] when the new parameter is
                    better.
                GFPQ_MODE_APPLY_ONLY(2): There is parameter in param[]. Don't
                    generate parameter. Just use the param[].
        """

        global _USE_GFPQ_QUANT_LIB
        if _USE_GFPQ_QUANT_LIB:
            try:
                if isinstance(tensor, tuple):
                    for tensor_item in tensor:
                        data_cuda_array = cuda.as_cuda_array(
                            tensor_item.data.detach())
                        data_p = data_cuda_array.device_ctypes_pointer
                        self._param.mode = mode
                        ret = self._libquant.HI_GFPQ_QuantAndDeQuant_GPU_PY(
                            data_p, data_cuda_array.size, self._bit_width,
                            ctypes.byref(self._param), self._stream.handle,
                            self._cublas_handle)
                else:
                    data_cuda_array = cuda.as_cuda_array(tensor.data.detach())
                    data_p = data_cuda_array.device_ctypes_pointer
                    self._param.mode = mode
                    ret = self._libquant.HI_GFPQ_QuantAndDeQuant_GPU_PY(
                        data_p, data_cuda_array.size, self._bit_width,
                        ctypes.byref(self._param), self._stream.handle,
                        self._cublas_handle)
            except:
                pass
            finally:
                if ret != 0:
                    _USE_GFPQ_QUANT_LIB = False
                    logger = logging.getLogger(__name__)
                    logger.setLevel(logging.WARNING)
                    logger.warning(
                        """Failed to quantize data with default HiSVP GFPQ library,
                        Use implemented quantization algorithm instead.""")
                    if isinstance(tensor, tuple):
                        for tensor_item in tensor:
                            tensor_item.data = fake_quantize(
                                tensor_item.data.detach().clone(), self._bit_width)
                    else:
                        tensor.data = fake_quantize(tensor.data.detach().clone(),
                                                    self._bit_width)
        else:
            if isinstance(tensor, tuple):
                for tensor_item in tensor:
                    tensor_item.data = fake_quantize(tensor_item.data.detach().clone(),
                                                     self._bit_width)
            else:
                tensor.data = fake_quantize(tensor.data.detach().clone(),
                                            self._bit_width)
        return tensor

_QUANT_HANDLE = QuantAndDeQuantGPU()

def _fuse_conv_bn_weights(conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b):
    """ fuse convolution and batch norm's weight.

    Args:
        conv_w (torch.nn.Parameter): convolution weight.
        conv_b (torch.nn.Parameter): convolution bias.
        bn_rm (torch.nn.Parameter): batch norm running mean.
        bn_rv (torch.nn.Parameter): batch norm running variance.
        bn_eps (torch.nn.Parameter): batch norm epsilon.
        bn_w (torch.nn.Parameter): batch norm weight.
        bn_b (torch.nn.Parameter): batch norm weight.

    Returns:
        conv_w(torch.nn.Parameter): fused convolution weight.
        conv_b(torch.nn.Parameter): fused convllution bias.
    """

    if conv_b is None:
        conv_b = bn_rm.new_zeros(bn_rm.shape)
    bn_var_rsqrt = torch.rsqrt(bn_rv + bn_eps)

    conv_w = conv_w * \
        (bn_w * bn_var_rsqrt).reshape([-1] + [1] * (len(conv_w.shape) - 1))
    conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b

    return torch.nn.Parameter(conv_w), torch.nn.Parameter(conv_b)

def _fuse_conv_bn(conv, bn):
    conv.weight, conv.bias = \
        _fuse_conv_bn_weights(conv.weight, conv.bias,
                             bn.running_mean, bn.running_var, bn.eps, bn.weight, bn.bias)
    return conv

def _fuse_modules(model):
    r"""Fuses a list of modules into a single module

    Fuses only the following sequence of modules:
    conv, bn
    All other sequences are left unchanged.
    For these sequences, fuse modules on weight level, keep model structure unchanged.

    Arguments:
        model: Model containing the modules to be fused

    Returns:
        model with fused modules.

    """
    children = list(model.named_children())
    conv_module = None
    conv_name = None

    for name, child in children:
        if isinstance(child, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d,
                              torch.nn.BatchNorm3d)):
            if isinstance(conv_module, (torch.nn.Conv2d, torch.nn.Conv3d)):
                conv_module = _fuse_conv_bn(conv_module, child)
                model._modules[conv_name] = conv_module
                child.eval()
                child.running_mean = child.running_mean.new_full(
                    child.running_mean.shape, 0)
                child.running_var = child.running_var.new_full(
                    child.running_var.shape, 1)
                if child.weight is not None:
                    child.weight.data = child.weight.data.new_full(
                        child.weight.shape, 1)
                if child.bias is not None:
                    child.bias.data = child.bias.data.new_full(
                        child.bias.shape, 0)
                child.track_running_stats = False
                child.momentum = 0
                child.eps = 0
            conv_module = None
        elif isinstance(child, (torch.nn.Conv2d, torch.nn.Conv3d)):
            conv_module = child
            conv_name = name
        else:
            _fuse_modules(child)
    return model

def freeze_bn(m, freeze_bn_affine=True):
    """Freeze batch normalization.
        reference: https://arxiv.org/abs/1806.08342

    Args:
        - m (nn.module): torch module
        - freeze_bn_affine (bool, optional): Freeze affine scale and
        translation factor or not. Defaults: True.
    """

    if isinstance(
            m,
        (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d)):

        m.eval()
        if freeze_bn_affine:
            m.weight.requires_grad = False
            m.bias.requires_grad = False

def merge_freeze_bn(model):
    """merge batch norm's weight into convolution, then freeze it.

    Args:
        model (nn.module): model.

    Returns:
        [nn.module]: model.
    """
    model = _fuse_modules(model)
    model.apply(freeze_bn)
    return model

def unquant_weight(m):
    """ unquantize weight before update weight, avoid training turbulence.

    Args:
        - m (nn.module): torch module.
    """
    try:
        if hasattr(m, "weight_origin") and m.weight is not None:
            m.weight.data.copy_(m.weight_origin.data)
    except AttributeError:
        pass
    except TypeError:
        pass

def quant_dequant_weight(m):
    """ quant weight manually.

    Args:
        - m (nn.module): torch module.
    """
    global _QUANT_HANDLE
    global _USE_GFPQ_QUANT_LIB
    quant_handle = _QUANT_HANDLE
    if not _USE_GFPQ_QUANT_LIB:
        quant_handle = QuantAndDeQuantGPU()
    try:
        if hasattr(m, "weight_origin") and m.weight is not None:
            m.weight_origin.data.copy_(m.weight.data)
            m.weight.data = quant_handle(m.weight.data.detach().clone())
    except AttributeError:
        pass
    except TypeError:
        pass

def _quantizing_activation(module, input, output):
    if isinstance(
            module,
        (torch.nn.ReLU, torch.nn.ELU, torch.nn.LeakyReLU, torch.nn.PReLU)):
        global _QUANT_HANDLE
        global _USE_GFPQ_QUANT_LIB
        quant_handle = _QUANT_HANDLE
        if not _USE_GFPQ_QUANT_LIB:
            quant_handle = QuantAndDeQuantGPU()
        # print("quantizing activation.")
        # print(output[0][0][0])
        output_type = output.dtype
        module.activation_max_value = torch.max(torch.max(torch.abs(output)), module.activation_max_value.to(output_type))
        # print(module.activation_max_value)
        tensor_t = torch.cat((output, torch.ones(output[0].shape).cuda().unsqueeze(0) * module.activation_max_value))
        output.data = quant_handle(tensor_t.float())[:-1]
        output = output.to(output_type)
        # print(output[0][0][0])

def _quantizing_data(module, input):
    global _QUANT_HANDLE
    global _USE_GFPQ_QUANT_LIB
    quant_handle = _QUANT_HANDLE
    if not _USE_GFPQ_QUANT_LIB:
        quant_handle = QuantAndDeQuantGPU()
    # print("quantizing data.")
    # print(input[0][0][0])
    # print("quantizing data.")
    # print(input[0][0][0])
    # input_type = input.dtype
    if isinstance(input, tuple):
        for item in input:
            item_type = item.dtype
            item = quant_handle(item.float())
            item.to(item_type)
    else:
        input = quant_handle(input.float())
    # input = input.to(input_type)
    # print(input[0][0][0])

def _quantizing_weight(module, input):
    global _QUANT_HANDLE
    global _USE_GFPQ_QUANT_LIB
    quant_handle = _QUANT_HANDLE
    if not _USE_GFPQ_QUANT_LIB:
        quant_handle = QuantAndDeQuantGPU()
    # print("quantizing weight.")
    # print(module.weight[0][0][0])
    module.weight_origin.data.copy_(module.weight.data)
    module.weight.data = quant_handle(module.weight.data.detach().clone())
    # print(module.weight[0][0][0])

def register_quantization_hook(model,
                               quant_weight=True,
                               quant_activation=True,
                               quant_data=False):
    """register quantization hook for model.

    Args:
        model (:class:`Module`): Module.

    Returns:
        Module: self
    """

    #  weight quantizing.
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)

    for _, module in model._modules.items():
        if len(list(module.children())) > 0:
            register_quantization_hook(module, quant_weight, quant_activation)
        else:
            if quant_weight and hasattr(
                    module,
                    "weight") and module.weight is not None and not isinstance(
                        module, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d,
                                 torch.nn.BatchNorm3d)):
                module.register_buffer('weight_origin', module.weight.detach().clone())
                if quant_data:
                    module.register_forward_pre_hook(_quantizing_data)
                    logger.info("Quantizing input data of %s", str(module))
                module.register_forward_pre_hook(_quantizing_weight)
                logger.info("Quantizing weight of %s", str(module))

            if quant_activation and isinstance(
                    module, (torch.nn.ReLU, torch.nn.ELU, torch.nn.LeakyReLU, torch.nn.PReLU)):
                module.register_buffer("activation_max_value", torch.tensor(0, dtype=torch.float).cuda())
                module.register_forward_hook(_quantizing_activation)
                logger.info("Quantizing activation of %s", str(module))

    return model

def test():
    r""" Test GFPG library QuantAndDeQuantGPU.
    """
    quant_handle = QuantAndDeQuantGPU()
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    tensor = torch.Tensor(np.array([-9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9])).cuda()
    logging.info("Origin Data: ")
    logging.info(tensor)

    start_time = datetime.datetime.now()
    quant_tensor = quant_handle(tensor)
    end_time = datetime.datetime.now()

    logging.info("Quant Data: ")
    logging.info(quant_tensor)

    data_expected = np.array([
        -8.7240619659, 0.0000000000, 1.0000000000, 2.0000000000, 2.9536523819,
        4.0000000000, 4.9674310684, 5.9073047638, 7.0250086784, 8.0000000000,
        8.7240619659
    ])

    logging.info("Data expected:  ")
    logging.info(" ".join([str(v) for v in data_expected]))

    data_diff = quant_tensor.data.detach().cpu().numpy() - data_expected
    flag = "success."
    for num in data_diff:
        if abs(num) > 0.000000001:
            flag = "failed."

    run_time = end_time - start_time
    logging.info("QuantAndDeQuantGPU time: %s", str(run_time))
    logging.info("QuantAndDeQuantGPU %s", flag)


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

import argparse
import argh
import os.path
import collections
import random
import re
import shutil
import socket
import sys
import tempfile
import time
import cloud_logging
from tqdm import tqdm
import gzip
import numpy as np
import tensorflow as tf
from tensorflow import gfile

import go
import dual_net
from gtp_wrapper import make_gtp_instance, MCTSPlayer
import preprocessing
import selfplay_mcts
from utils import logged_timer as timer
import evaluation
import sgf_wrapper
import utils

import qmeas
import goparams
import glbl

# How many positions we should aggregate per 'chunk'.
EXAMPLES_PER_RECORD = goparams.EXAMPLES_PER_RECORD

# How many positions to draw from for our training window.
# AGZ used the most recent 500k games, which, assuming 250 moves/game = 125M
# WINDOW_SIZE = 125000000
#WINDOW_SIZE = 500000
WINDOW_SIZE = goparams.WINDOW_SIZE

def _ensure_dir_exists(directory):
    if directory.startswith('gs://'):
        return
    os.makedirs(directory, exist_ok=True)

def gtp(load_file: "The path to the network model files"=None,
        readouts: 'How many simulations to run per move'=10000,
        #readouts: 'How many simulations to run per move'=2000,
        cgos_mode: 'Whether to use CGOS time constraints'=False,
        verbose=1):
    engine = make_gtp_instance(load_file,
                               readouts_per_move=readouts,
                               verbosity=verbose,
                               cgos_mode=cgos_mode)
    sys.stderr.write("GTP engine ready\n")
    sys.stderr.flush()
    while not engine.disconnect:
        inpt = input()
        # handle either single lines at a time
        # or multiple commands separated by '\n'
        try:
            cmd_list = inpt.split("\n")
        except:
            cmd_list = [inpt]
        for cmd in cmd_list:
            engine_reply = engine.send(cmd)
            sys.stdout.write(engine_reply)
            sys.stdout.flush()

def bootstrap(
        working_dir: 'tf.estimator working directory. If not set, defaults to a random tmp dir'=None,
        model_save_path: 'Where to export the first bootstrapped generation'=None):
    qmeas.start_time('bootstrap')
    glbl.prof.set_operation('bootstrap')
    if working_dir is None:
        with tempfile.TemporaryDirectory() as working_dir:
            _ensure_dir_exists(working_dir)
            _ensure_dir_exists(os.path.dirname(model_save_path))
            dual_net.bootstrap(working_dir)
            dual_net.export_model(working_dir, model_save_path)
    else:
        _ensure_dir_exists(working_dir)
        _ensure_dir_exists(os.path.dirname(model_save_path))
        dual_net.bootstrap(working_dir)
        dual_net.export_model(working_dir, model_save_path)
    glbl.prof.end_operation('bootstrap')
    qmeas.stop_time('bootstrap')

def train(
        working_dir: 'tf.estimator working directory.',
        chunk_dir: 'Directory where gathered training chunks are.',
        model_save_path: 'Where to export the completed generation.',
        generation_num: 'Which generation you are training.'=0):
    qmeas.start_time('train')
    tf_records = sorted(gfile.Glob(os.path.join(chunk_dir, '*.tfrecord.zz')))
    tf_records = tf_records[-1 * (WINDOW_SIZE // EXAMPLES_PER_RECORD):]

    print("Training from:", tf_records[0], "to", tf_records[-1])

    with timer("Training"):
        dual_net.train(working_dir, tf_records, generation_num)
        dual_net.export_model(working_dir, model_save_path)
    qmeas.stop_time('train')

def validate(
        working_dir: 'tf.estimator working directory',
        *tf_record_dirs: 'Directories where holdout data are',
        checkpoint_name: 'Which checkpoint to evaluate (None=latest)'=None,
        validate_name: 'Name for validation set (i.e., selfplay or human)'=None):
    qmeas.start_time('validate')
    tf_records = []
    with timer("Building lists of holdout files"):
        for record_dir in tf_record_dirs:
            tf_records.extend(gfile.Glob(os.path.join(record_dir, '*.zz')))

    first_record = os.path.basename(tf_records[0])
    last_record = os.path.basename(tf_records[-1])
    with timer("Validating from {} to {}".format(first_record, last_record)):
        dual_net.validate(
            working_dir, tf_records, checkpoint_name=checkpoint_name,
            name=validate_name)
    qmeas.stop_time('validate')

def evaluate(
        black_model: 'The path to the model to play black',
        white_model: 'The path to the model to play white',
        output_dir: 'Where to write the evaluation results'='sgf/evaluate',
        readouts: 'How many readouts to make per move.'=200,
        games: 'the number of games to play'=20,
        verbose: 'How verbose the players should be (see selfplay)' = 1):
    qmeas.start_time('evaluate')
    _ensure_dir_exists(output_dir)

    with timer("Loading weights"):
        black_net = dual_net.DualNetwork(black_model)
        white_net = dual_net.DualNetwork(white_model)

    winners = []
    with timer("%d games" % games):
        winners = evaluation.play_match(
            black_net, white_net, games, readouts, output_dir, verbose)
    qmeas.stop_time('evaluate')
    white_count = 0
    for win in winners:
      if 'W' in win or 'w' in win:
        white_count += 1
    return white_count * 1.0 / games

    # qmeas.report_profiler()

def evaluate_evenly(
        black_model: 'The path to the model to play black',
        white_model: 'The path to the model to play white',
        output_dir: 'Where to write the evaluation results'='sgf/evaluate',
        readouts: 'How many readouts to make per move.'=200,
        games: 'the number of games to play'=20,
        verbose: 'How verbose the players should be (see selfplay)' = 1):
  ''' Returns the white win rate; playes 'games' number of games on both sides. '''
  try:
    result = (evaluate(black_model, white_model, output_dir, readouts, games, verbose) + (1 - evaluate(white_model, black_model, output_dir, readouts, games, verbose)))/ 2.0
  except TypeError:
    # It is remotely possible that in weird twist of fate results in a type
    # error... Possibly due to weird corner cases in the evaluation...
    # Our fall back will be to try agian.
    result = (evaluate(black_model, white_model, output_dir, readouts, games, verbose) + (1 - evaluate(white_model, black_model, output_dir, readouts, games, verbose)))/ 2.0
    # should this really happen twice, the world really doesn't
    # want this to be successful... and we will raise the error.
    # If this is being run by the main loop harness, then the
    # effect of raising here will be to keep the newest model and go back to
    # selfplay.
  return result

def selfplay(
        load_file: "The path to the network model files",
        output_dir: "Where to write the games"="data/selfplay",
        holdout_dir: "Where to write the games"="data/holdout",
        output_sgf: "Where to write the sgfs"="sgf/",
        readouts: 'How many simulations to run per move'=100,
        verbose: '>=2 will print debug info, >=3 will print boards' = 1,
        resign_threshold: 'absolute value of threshold to resign at' = 0.95,
        holdout_pct: 'how many games to hold out for validation' = 0.05):
    qmeas.start_time('selfplay')
    clean_sgf = os.path.join(output_sgf, 'clean')
    full_sgf = os.path.join(output_sgf, 'full')
    _ensure_dir_exists(clean_sgf)
    _ensure_dir_exists(full_sgf)
    _ensure_dir_exists(output_dir)
    _ensure_dir_exists(holdout_dir)

    with timer("Loading weights from %s ... " % load_file):
        network = dual_net.DualNetwork(load_file)

    with timer("Playing game"):
        player = selfplay_mcts.play(
            network, readouts, resign_threshold, verbose)

    output_name = '{}-{}'.format(int(time.time() * 1000 * 1000), socket.gethostname())
    game_data = player.extract_data()
    with gfile.GFile(os.path.join(clean_sgf, '{}.sgf'.format(output_name)), 'w') as f:
        f.write(player.to_sgf(use_comments=False))
    with gfile.GFile(os.path.join(full_sgf, '{}.sgf'.format(output_name)), 'w') as f:
        f.write(player.to_sgf())

    tf_examples = preprocessing.make_dataset_from_selfplay(game_data)

    # Hold out 5% of games for evaluation.
    if random.random() < holdout_pct:
        fname = os.path.join(holdout_dir, "{}.tfrecord.zz".format(output_name))
    else:
        fname = os.path.join(output_dir, "{}.tfrecord.zz".format(output_name))

    preprocessing.write_tf_examples(fname, tf_examples)
    qmeas.stop_time('selfplay')

def selfplay_cache_model(
        network: "The path to the network model files",
        output_dir: "Where to write the games"="data/selfplay",
        holdout_dir: "Where to write the games"="data/holdout",
        output_sgf: "Where to write the sgfs"="sgf/",
        readouts: 'How many simulations to run per move'=100,
        verbose: '>=2 will print debug info, >=3 will print boards' = 1,
        resign_threshold: 'absolute value of threshold to resign at' = 0.95,
        holdout_pct: 'how many games to hold out for validation' = 0.05):
    qmeas.start_time('selfplay')
    clean_sgf = os.path.join(output_sgf, 'clean')
    full_sgf = os.path.join(output_sgf, 'full')
    _ensure_dir_exists(clean_sgf)
    _ensure_dir_exists(full_sgf)
    _ensure_dir_exists(output_dir)
    _ensure_dir_exists(holdout_dir)

    with timer("Playing game"):
        player = selfplay_mcts.play(
            network, readouts, resign_threshold, verbose)

    output_name = '{}-{}'.format(int(time.time() * 1000 * 1000), socket.gethostname())
    game_data = player.extract_data()
    with gfile.GFile(os.path.join(clean_sgf, '{}.sgf'.format(output_name)), 'w') as f:
        f.write(player.to_sgf(use_comments=False))
    with gfile.GFile(os.path.join(full_sgf, '{}.sgf'.format(output_name)), 'w') as f:
        f.write(player.to_sgf())

    tf_examples = preprocessing.make_dataset_from_selfplay(game_data)

    # Hold out 5% of games for evaluation.
    if random.random() < holdout_pct:
        fname = os.path.join(holdout_dir, "{}.tfrecord.zz".format(output_name))
    else:
        fname = os.path.join(output_dir, "{}.tfrecord.zz".format(output_name))

    preprocessing.write_tf_examples(fname, tf_examples)
    qmeas.stop_time('selfplay')

def gather(
        input_directory: 'where to look for games'='data/selfplay/',
        output_directory: 'where to put collected games'='data/training_chunks/',
        examples_per_record: 'how many tf.examples to gather in each chunk'=EXAMPLES_PER_RECORD):
    qmeas.start_time('gather')
    _ensure_dir_exists(output_directory)
    models = [model_dir.strip('/')
              for model_dir in sorted(gfile.ListDirectory(input_directory))[-50:]]
    with timer("Finding existing tfrecords..."):
        model_gamedata = {
            model: gfile.Glob(
                os.path.join(input_directory, model, '*.tfrecord.zz'))
            for model in models
        }
    print("Found %d models" % len(models))
    for model_name, record_files in sorted(model_gamedata.items()):
        print("    %s: %s files" % (model_name, len(record_files)))

    meta_file = os.path.join(output_directory, 'meta.txt')
    try:
        with gfile.GFile(meta_file, 'r') as f:
            already_processed = set(f.read().split())
    except tf.errors.NotFoundError:
        already_processed = set()

    num_already_processed = len(already_processed)

    for model_name, record_files in sorted(model_gamedata.items()):
        if set(record_files) <= already_processed:
            continue
        print("Gathering files for %s:" % model_name)
        for i, example_batch in enumerate(
                tqdm(preprocessing.shuffle_tf_examples(examples_per_record, record_files))):
            output_record = os.path.join(output_directory,
                                         '{}-{}.tfrecord.zz'.format(model_name, str(i)))
            preprocessing.write_tf_examples(
                output_record, example_batch, serialize=False)
        already_processed.update(record_files)

    print("Processed %s new files" %
          (len(already_processed) - num_already_processed))
    with gfile.GFile(meta_file, 'w') as f:
        f.write('\n'.join(sorted(already_processed)))
    qmeas.stop_time('gather')

parser = argparse.ArgumentParser()
argh.add_commands(parser, [gtp, bootstrap, train,
                           selfplay, gather, evaluate, validate])

if __name__ == '__main__':
    cloud_logging.configure()
    argh.dispatch(parser)


import os
import platform
if (platform.node()=='csc-G7-7590'):
    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
else:
    os.environ["CUDA_VISIBLE_DEVICES"] = '3'

import torchvision.models as models
from torch.optim.lr_scheduler import MultiStepLR
import torch
from cscPy.Nets.AligningHandnet import encoderRGB,decPoseNet
from cscPy.dataloader.MVdataloader import MVDataloader
from cscPy.mano.dataloader.synthesizer import ManoSynthesizer
from cscPy.Loss.utils import pose3d_loss,LossHelper,cloud_dis,cloud_dis2
import numpy as np
from cscPy.mano.network.manoArmLayer import MANO_SMPL
from cscPy.mano.network.biomechanicalLoss import BiomechanicalLayer
from cscPy.globalCamera.camera import CameraIntrinsics,perspective_projection
from cscPy.globalCamera.util import fetch_all_sequences,load_rgb_maps,load_depth_maps,get_cameras_from_dir,visualize_better_qulity_depth_map

import cv2
from cscPy.mano.network.utils import *

encoderRGB=encoderRGB().cuda()
decoderPose=decPoseNet().cuda()

# checkpoint = torch.load('./pretrain/rgb2pose.pt')
# decoderPose.load_state_dict(checkpoint['decoderCloud'])
# encoderRGB.load_state_dict(checkpoint['encoderRGB'])
onlysyn=True

train_dataset1 = ManoSynthesizer()
train_dataset2 = MVDataloader()
train_dataset = torch.utils.data.ConcatDataset( [train_dataset2])
def _init_fn(worker_id):np.random.seed(worker_id)
def _init_fn2(worker_id):np.random.seed(worker_id**2+2)
train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=16,num_workers=4, shuffle=True,worker_init_fn=_init_fn)
print('train_loader',len(train_loader))
manoPath='/home/csc/MANO-hand-model-toolkit/mano/models/MANO_RIGHT.pkl'
if not os.path.exists(manoPath):
    manoPath = '/home/shicheng/MANO-hand-model-toolkit/mano/models/MANO_RIGHT.pkl'
mano_right = MANO_SMPL(manoPath, ncomps=45, oriorder=True,device='cuda',userotJoints=True)

mylist=[]
mylist.append({'params':encoderRGB.parameters()})
mylist.append({'params':decoderPose.parameters()})
optimizer = torch.optim.Adam(mylist, lr=1e-3)
scheduler = MultiStepLR(optimizer, milestones=[30,60], gamma=0.1)

def getLatentLoss(z_mean, z_stddev, goalStd=1.0, eps=1e-9):
    latent_loss = 0.5 * torch.sum(z_mean**2 + z_stddev**2 - torch.log(z_stddev**2)  - goalStd, 1)
    return latent_loss
biolayer = BiomechanicalLayer(fingerPlaneLoss=True,fingerFlexLoss=True, fingerAbductionLoss=True)
losshelp=LossHelper()

for epoch in range(80):

    for idx, inp in enumerate(train_loader):
        img,cloud,pose_gt,scale,root,mask=inp['img'].cuda(),inp['cloud'].cuda(),inp['pose3d'].cuda(),inp['scale'].cuda(),\
                                          inp['root'].cuda(),inp['mask'].cuda().reshape(-1)
        # img, cloud, pose_gt, scale, root, mask = inp['img'], inp['cloud'], inp['pose3d'], inp[
        #     'scale'],inp['root'], inp['mask'].reshape(-1)

        # cur=img[0].permute(1,2,0).cpu().numpy()
        # cur=(cur*255).astype(np.uint8)
        # cv2.imshow('cur',cur)
        # cv2.waitKey(1)
        encoderRGB.train()
        decoderPose.train()
        N=img.shape[0]
        z_rgb, mn_rgb, sd_rgb = encoderRGB(img,training=False)  # encode rgb
        # print('img',img)
        # z_rgb = encoderRGB(img,training=False)  # encode rgb
        # print("z_rgb",z_rgb)
        pose_rgb = decoderPose(z_rgb,).reshape(N,21,3)
        latent_loss_rgb = getLatentLoss(mn_rgb, sd_rgb)
        latent_loss_sum = torch.mean(latent_loss_rgb)
        pose_loss_real,eucLoss_real=pose3d_loss(pose_rgb[mask!=1],pose_gt[mask!=1],scale[mask!=1])

        joints=pose_rgb*scale+root
        wrist_trans,local_trans,outjoints=mano_right.matchTemplate2JointsGreedy(joints)
        _, Greedymatchloss = pose3d_loss(joints, outjoints, scale)

        vertexPre, joint_pre = mano_right.get_mano_vertices(wrist_trans.reshape([N, 1, 3, 3]),
                                                            local_trans.view(N, 15, 3, 3),
                                                            torch.zeros([N, 10], dtype=torch.float32).view(N, 10),
                                                            torch.ones([N, 1], dtype=torch.float32).view(N, 1),
                                                            torch.zeros([N,3], dtype=torch.float32).view(N,3),
                                                            pose_type='rot_matrix', mmcp_center=False,
                                                            external_transition=None)

        mmcp = joint_pre[:, 4:5, :].clone()
        joint_pre = (joint_pre - mmcp) / scale
        vertexPre = (vertexPre - mmcp) / scale

        realBioLoss,realBioEudloss=biolayer(joint_pre[mask!=1],scale.reshape(N)[mask!=1])
        realpose_loss_bone, realeucLoss_bone = pose3d_loss(pose_rgb[mask!=1], joint_pre[mask!=1], scale[mask!=1])
        realcdloss, realcdlossEud = cloud_dis2(vertexPre[mask != 1], cloud[mask != 1], scale[mask != 1])

        loss = 0.0001*latent_loss_sum + pose_loss_real+\
               realpose_loss_bone+\
               realBioLoss+\
               realcdloss
        #loss = 0.0001*latent_loss_sum + cloudRec_loss_sum
        dicloss = {"loss": float(loss),
                   "Greedymatch dis": float(Greedymatchloss) * 1000,
                   "real epe": float(eucLoss_real) * 1000, "real loss": float(pose_loss_real),
                   "real cd dis": float(realcdlossEud) * 1000, "real cd loss": float(realcdloss),
                   "real bio dis": float(realBioEudloss) * 1000, "real bio loss": float(realBioLoss),
                   "real bone dis": float(realeucLoss_bone) * 1000, "real bone loss": float(realpose_loss_bone),
                   }
        losshelp.add(dicloss)

        if(idx%5==0):
            print('epoch:{} iteration:{}'.format(epoch,idx))
            losshelp.showcurrent()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # print('scheduler.step()')
    # scheduler.step()
    losshelp.show()
    losshelp.clear()

    print('save model')
    torch.save({
        'epoch': epoch,
        'encoderRGB': encoderRGB.state_dict(),
        'decoderPose': decoderPose.state_dict(),
        'optimizer': optimizer.state_dict()}, './models/' + platform.node() + 'rgb2poseReal.pt')

    #print(epoch, 'epoch mean epeloss', np.mean(epe)*1000)
    # aveloss,epe=[],[]
    # #aveCD,aveEMD=[],[]
    # with torch.no_grad():
    #     for idx,(image, depth, cloud, heatmap, pose_gt, viewRotation, scale) in enumerate(test_loader):
    #         # image, depth, cloud, heatmap, pose_gt, viewRotation, scale = image.cuda(), depth.cuda(), cloud.cuda(), heatmap.cuda(), \
    #         #                                                              pose_gt.cuda(), viewRotation.cuda(), scale.cuda()
    #         pose_gt, scale, image =pose_gt.cuda(), scale.cuda(), image.cuda()
    #         encoderRGB.eval()
    #         decoderPose.eval()
    #         z_rgb, mn_rgb, sd_rgb = encoderRGB(image,training=False)  # encode rgb
    #         pose_rgb = decoderPose(z_rgb, )
    #         pose_loss_rgb, eucLoss_rgb = utils.pose3d_loss(pose_rgb, pose_gt, scale)
    #         epe.append(float(eucLoss_rgb))
    #
    #     print('len(epe)',len(epe[:len(epe)//2]))
    #     print(epoch, 'epoch test mean rgb epeloss', np.mean(epe[:len(epe)//2])*1000)


import os
import io
import lmdb
import tqdm
import glob
import numpy as np
import librosa
import torch
import json
import random
import pandas as pd
from torch.utils.data import Dataset
from typing import Tuple, Optional
from src.data.audio_parser import AudioParser
import soundfile as sf
from src.data.utils import load_audio
import msgpack
import msgpack_numpy as msgnp

# def load_bf_files_from_lmdb(bgfiles_lmdb):
#     env = lmdb.open(bgfiles_lmdb, subdir=os.path.isdir(bgfiles_lmdb),
#                     max_readers=1, readonly=True, lock=False,
#                     readahead=False, meminit=False)
#     with env.begin(write=False) as txn:
#         # self.length = txn.stat()['entries'] - 1
#         length = pa.deserialize(txn.get(b'__len__'))
#         keys = pa.deserialize(txn.get(b'__keys__'))
#     features = []
#     with env.begin(write=False) as txn:
#         for ix in tqdm.tqdm(range(length)):
#             byteflow = txn.get(keys[ix])
#             unpacked = pa.deserialize(byteflow)
#             features.append(unpacked)
#     env.close()
#     return features

class LMDBSpectrogramDataset(Dataset):
    def __init__(self, lmdb_path, labels_map,
                 audio_config, mode="multilabel",
                 augment=False, delimiter=",",
                 mixer=None, transform=None):
        super(LMDBSpectrogramDataset, self).__init__()
        assert os.path.isfile(labels_map)
        assert os.path.splitext(labels_map)[-1] == ".json"
        assert audio_config is not None
        with open(labels_map, 'r') as fd:
            self.labels_map = json.load(fd)
        self.labels_delim = delimiter
        self.parse_audio_config(audio_config)
        if self.background_noise_path is not None:
            if os.path.exists(self.background_noise_path):
                self.bg_files = glob.glob(os.path.join(self.background_noise_path, "*.wav"))
        else:
            self.bg_files = None
        self.mode = mode
        self.spec_parser = AudioParser(n_fft=self.n_fft, win_length=self.win_len,
                                       hop_length=self.hop_len, feature=self.feature_type)
        self.mixer = mixer
        self.transform = transform
        if self.bg_files is not None:
            print("prepping bg_features")
            self.bg_features = []
            for f in tqdm.tqdm(self.bg_files):
                preprocessed_audio = self.__get_audio__(f)
                real, comp = self.__get_feature__(preprocessed_audio)
                self.bg_features.append(real)
        else:
            self.bg_features = None
        self.env = lmdb.open(lmdb_path, subdir=os.path.isdir(lmdb_path),
                             max_readers=1, readonly=True, lock=False,
                             readahead=False, meminit=False)

        with self.env.begin(write=False) as txn:
            # self.length = txn.stat()['entries'] - 1
            self.length = msgpack.unpackb(txn.get(b'__len__'))
            self.keys = msgpack.unpackb(txn.get(b'__keys__'))

    def parse_audio_config(self, audio_config):
        self.sr = audio_config.get("sample_rate", "22050")
        self.n_fft = audio_config.get("n_fft", 511)
        win_len = audio_config.get("win_len", None)
        if not win_len:
            self.win_len = self.n_fft
        else:
            self.win_len = win_len
        hop_len = audio_config.get("hop_len", None)
        if not hop_len:
            self.hop_len = self.n_fft // 2
        else:
            self.hop_len = hop_len
        self.normalize = audio_config.get("normalize", True)
        self.min_duration = audio_config.get("min_duration", None)
        self.background_noise_path = audio_config.get("bg_files", None)
        self.feature_type = audio_config.get("feature", "spectrogram")

    def __get_audio__(self, f):
        audio = load_audio(f, self.sr, self.min_duration)
        return audio

    def __get_feature__(self, audio) -> Tuple[torch.Tensor, torch.Tensor]:
        real, comp = self.spec_parser(audio)
        return real, comp

    def get_bg_feature(self, index: int) -> torch.Tensor:
        if self.bg_features is None:
            return None
        real = self.bg_features[index]
        if self.transform is not None:
            real = self.transform(real)
        return real

    def __get_item_helper__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        with self.env.begin(write=False) as txn:
            byteflow = txn.get(self.keys[index])
        unpacked = msgpack.unpackb(byteflow, object_hook=msgnp.decode)
        flac_audio = unpacked[0]
        lbls = unpacked[1]
        with io.BytesIO(flac_audio) as buf:
            preprocessed_audio, sr = sf.read(buf)
        label_tensor = self.__parse_labels__(lbls)
        real, comp = self.__get_feature__(preprocessed_audio)
        if self.transform is not None:
            real = self.transform(real)
        return real, comp, label_tensor

    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
        real, comp, label_tensor = self.__get_item_helper__(index)
        if self.mixer is not None:
            real, final_label = self.mixer(self, real, label_tensor)
            if self.mode != "multiclass":

                return real, final_label
        return real, label_tensor

    def __parse_labels__(self, lbls: str) -> torch.Tensor:
        if self.mode == "multilabel":
            label_tensor = torch.zeros(len(self.labels_map)).float()
            for lbl in lbls.split(self.labels_delim):
                label_tensor[self.labels_map[lbl]] = 1

            return label_tensor
        elif self.mode == "multiclass":
            return self.labels_map[lbls]

    def __len__(self):
        return self.length

    def get_bg_len(self):
        return len(self.bg_features)