import skimage.io as io
import skimage.transform as skt
import numpy as np
from PIL import Image
from src.models.class_patcher import patcher
from src.utils.imgproc import *
from skimage.color import rgb2hsv, hsv2rgb

class patcher(patcher):
    def __init__(self, body='./body/body_felt.png', **options):
        super().__init__('フェルト', body=body, pantie_position=[11, 633], **options)
        self.mask = io.imread('./mask/mask_felt.png')
        self.bra = io.imread('./mask/bra_felt.png') / 255
        self.bra_shade = io.imread('./material/bra_felt_shade.png') / 255
        self.bra_position = [65, 38]

    def convert(self, image):
        pantie = np.array(image)
        patch = np.copy(pantie[-140:-5, 546:, :])
        [pr, pc, d] = patch.shape
        pantie[-120:, 546:] = 0
        pantie[125:125 + pr, :pc, :] = patch
        pantie = pantie[:-100]
        arrx = np.zeros(25)
        arry = np.linspace(0, 1, 25)**2 * -29
        front = pantie[:220, :360]
        front = resize(affine_transform_by_arr(front, arrx, arry), [.75, .75])[:, 5:]
        back = pantie[:, 268:][:, ::-1]
        back = resize(affine_transform_by_arr(back, arrx, arry), [.75, .75])[:, 5:]
        front = np.uint8(front * 255)
        back = np.uint8(back * 255)

        front = np.concatenate([front[:, ::-1], front], axis=1)
        back = np.concatenate([back[:, ::-1], back], axis=1)
        [fr, fc, d] = front.shape
        [br, bc, d] = back.shape
        dy = 26
        dx = 2
        pantie = np.zeros([fr + br + dy, fc + dx, d], dtype=np.uint8)
        pantie[:fr, dx:dx + fc] = front
        pantie[fr + dy:, :bc] = back
        pantie = np.bitwise_and(pantie, self.mask)
        return Image.fromarray(pantie)

    def pick_color(self, arr):
        return np.mean(np.mean(arr, axis=0), axis=0)

    def gen_bra(self, image):
        pantie = np.array(image)
        front = pantie[20:100, 30:80, :3] / 255.0
        front_shade = pantie[130:150, 0:40, :3] / 255.0
        front_color = self.pick_color(front)
        front_shade_color = self.pick_color(front_shade)
        front_shade_color = rgb2hsv(front_shade_color[None, None])
        front_shade_color[0, 0, 1] *= front_shade_color[0, 0, 2] / 0.3
        if front_shade_color[0, 0, 1] > 0.7:
            front_shade_color[0, 0, 1] *= 0.7
        front_shade_color[0, 0, 2] *= front_shade_color[0, 0, 2] / 0.4
        front_shade_color = np.clip(hsv2rgb(front_shade_color)[0, 0], 0, 1)

        # painting colors
        bra = self.bra[:, :, :3] * front_color
        bra_shade = (self.bra_shade[:, :, -1])[:, :, None] * front_shade_color

        # overlaying layers
        bra = alpha_brend(bra_shade, bra, self.bra_shade[:, :, -1])
        bra = np.dstack((bra, self.bra[:, :, -1]))
        return Image.fromarray(np.uint8(np.clip(bra, 0, 1) * 255))

    def patch(self, image, transparent=False):
        pantie = self.convert(image)
        if transparent:
            patched = Image.new("RGBA", self.body_size)
        else:
            patched = self.body.copy()
        patched = self.paste(patched, pantie, self.pantie_position)
        patched = self.paste(patched, self.gen_bra(image), self.bra_position)
        return patched


import data.shaders.skybox_shader as s
import numpy, pygame
from OpenGL.GL import *

class skybox_renderer():
	SIZE = 500.0
	VERTICES = numpy.array([-SIZE, SIZE, -SIZE, -SIZE, -SIZE, -SIZE, SIZE, -SIZE, -SIZE, SIZE, -SIZE, -SIZE, SIZE, SIZE, -SIZE, -SIZE, SIZE, -SIZE, -SIZE, -SIZE, SIZE, -SIZE, -SIZE, -SIZE, -SIZE, SIZE, -SIZE, -SIZE, SIZE, -SIZE, -SIZE, SIZE, SIZE, -SIZE, -SIZE, SIZE, SIZE, -SIZE, -SIZE, SIZE, -SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, -SIZE, SIZE, -SIZE, -SIZE, -SIZE, -SIZE, SIZE, -SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, -SIZE, SIZE, -SIZE, -SIZE, SIZE, -SIZE, SIZE, -SIZE, SIZE, SIZE, -SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, SIZE, -SIZE, SIZE, SIZE, -SIZE, SIZE, -SIZE, -SIZE, -SIZE, -SIZE, -SIZE, -SIZE, SIZE, SIZE, -SIZE, -SIZE, SIZE, -SIZE, -SIZE, -SIZE, -SIZE, SIZE, SIZE, -SIZE, SIZE], dtype = "float32")
	DAY_TEXTURE_FILES = ["right", "left", "top", "bottom", "back", "front"]
	NIGHT_TEXTURE_FILES = ["night_right", "night_left", "night_top", "night_bottom", "night_back", "night_front"]
	def __init__(self, loader, projection_matrix):
		self.cube = loader.load_to_vao(self.VERTICES, 3)
		self.day_texture = loader.load_cube_map(self.DAY_TEXTURE_FILES)
		self.night_texture = loader.load_cube_map(self.NIGHT_TEXTURE_FILES)
		self.shader = s.skybox_shader()
		self.shader.start()
		self.shader.connect_texture_units()
		self.shader.load_projection_matrix(projection_matrix)
		self.shader.stop()
	def render(self, camera, r, g, b, clock):
		self.shader.start()
		self.shader.load_view_matrix(camera, clock)
		self.shader.load_fog_color(r, g, b)
		glBindVertexArray(self.cube.get_vao_id())
		glEnableVertexAttribArray(0)
		self.bind_textures()
		glDrawArrays(GL_TRIANGLES, 0, self.cube.get_vertex_count())
		glDisableVertexAttribArray(0)
		glBindVertexArray(0)
		self.shader.stop()
	def bind_textures(self):
		time = pygame.time.get_ticks()
		time = time % 24000
		blend_factor = 0.0
		if time >= 0 and time < 5000:
			texture_1 = self.night_texture
			texture_2 = self.night_texture
			blend_factor = float(float((time - 0)) / (5000 - 0))
		elif time >= 5000 and time < 8000:
			texture_1 = self.night_texture
			texture_2 = self.day_texture
			blend_factor = float(float((time - 5000)) / (8000 - 5000))
		elif time >= 8000 and time < 21000:
			texture_1 = self.day_texture
			texture_2 = self.day_texture
			blend_factor = float(float((time - 8000)) / (21000 - 8000))
		else:
			texture_1 = self.day_texture
			texture_2 = self.night_texture
			blend_factor = float(float((time - 21000)) / (24000 - 21000))
		glActiveTexture(GL_TEXTURE0)
		glBindTexture(GL_TEXTURE_CUBE_MAP, texture_1)
		glActiveTexture(GL_TEXTURE1)
		glBindTexture(GL_TEXTURE_CUBE_MAP, texture_2)
		self.shader.load_blend_factor(blend_factor)
	def clean_up(self):
		self.shader.clean_up()


from __future__ import division
from qm.dftbplus.dftbplus import DFTBplus
from qm.dftbplus.dftbpar import spin_w, spin_w_lc, max_l
from misc import au_to_A, call_name
import os, shutil, re, textwrap
import numpy as np

class SSR(DFTBplus):
    """ Class for SSR method of DFTB+

        :param object molecule: Molecule object
        :param boolean l_scc: Include self-consistent charge (SCC) scheme
        :param double scc_tol: Stopping criteria for the SCC iterations
        :param integer scc_max_iter: Maximum number of SCC iterations
        :param boolean l_onsite: Include onsite correction to SCC term
        :param boolean l_range_sep: Include long-range corrected functional
        :param string lc_method: Algorithms for LC-DFTB
        :param boolean l_ssr22: Use SSR(2,2) calculation?
        :param boolean l_ssr44: Use SSR(4,4) calculation?
        :param string guess: Initial guess method for SCC scheme
        :param string guess_file: Initial guess file for eigenvetors
        :param boolean l_state_interactions: Include state-interaction terms to SA-REKS
        :param double shift: Level shifting value in SCC iterations
        :param double,list tuning: Scaling factor for atomic spin constants
        :param string cpreks_grad_alg: Algorithms used in CP-REKS equations
        :param double cpreks_grad_tol: Tolerance used in the conjugate-gradient based algorithms
        :param boolean l_save_memory: Save memory in cache used in CP-REKS equations
        :param string embedding: Charge-charge embedding options
        :param boolean l_periodic: Use periodicity in the calculations
        :param double,list cell_length: The lattice vectors of periodic unit cell
        :param string sk_path: Path for Slater-Koster files
        :param string install_path: Path for DFTB+ install directory
        :param integer nthreads: Number of threads in the calculations
        :param string version: Version of DFTB+
    """
    def __init__(self, molecule, l_scc=True, scc_tol=1E-6, scc_max_iter=1000, l_onsite=False, \
        l_range_sep=False, lc_method="MatrixBased", l_ssr22=False, l_ssr44=False, guess="h0", \
        guess_file="./eigenvec.bin", l_state_interactions=False, shift=0.3, tuning=None, \
        cpreks_grad_alg="pcg", cpreks_grad_tol=1E-8, l_save_memory=False, embedding=None, \
        l_periodic=False, cell_length=[0., 0., 0., 0., 0., 0., 0., 0., 0.], sk_path="./", \
        install_path="./", nthreads=1, version="20.1"):
        # Initialize DFTB+ common variables
        super(SSR, self).__init__(molecule, sk_path, install_path, nthreads, version)

        # Initialize DFTB+ SSR variables
        self.l_scc = l_scc
        self.scc_tol = scc_tol
        self.scc_max_iter = scc_max_iter

        self.l_onsite = l_onsite
        if (self.l_onsite):
            error_message = "Onsite-correction not implemented!"
            error_vars = f"l_onsite = {self.l_onsite}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.l_range_sep = l_range_sep
        self.lc_method = lc_method.lower()

        self.l_ssr22 = l_ssr22
        # TODO : logical variable (l_ssr22, l_ssr44) must be changed for generality
        if (not self.l_ssr22):
            error_message = "Use (2,2) active space for SSR!"
            error_vars = f"l_ssr22 = {self.l_ssr22}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")
        self.l_ssr44 = l_ssr44
        if (self.l_ssr44):
            error_message = "(4,4) active space for SSR not implemented!"
            error_vars = f"l_ssr44 = {self.l_ssr44}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        # Set initial guess for eigenvectors
        self.guess = guess.lower()
        self.guess_file = guess_file
        if not (self.guess in ["h0", "read"]):
            error_message = "Invalid initial guess for DFTB/SSR!"
            error_vars = f"guess = {self.guess}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.l_state_interactions = l_state_interactions
        self.shift = shift

        # Set scaling factor for atomic spin constants
        self.tuning = tuning
        if (self.tuning != None):
            if (len(self.tuning) != len(self.atom_type)):
                error_message = "Number of elements for scaling factor must be equal to number of atom types!"
                error_vars = f"len(tuning) = {len(self.tuning)}, len(atom_type) = {len(self.atom_type)}"
                raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.cpreks_grad_alg = cpreks_grad_alg.lower()
        self.cpreks_grad_tol = cpreks_grad_tol
        self.l_save_memory = l_save_memory

        self.embedding = embedding
        if (self.embedding != None):
            self.embedding = self.embedding.lower()

        if not (self.embedding in [None, "mechanical", "electrostatic"]):
            error_message = "Invalid charge embedding for QM/MM calculation!"
            error_vars = f"embedding = {self.embedding}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.l_periodic = l_periodic
        self.a_axis = np.copy(cell_length[0:3])
        self.b_axis = np.copy(cell_length[3:6])
        self.c_axis = np.copy(cell_length[6:9])

        # Set 'l_nacme' and 're_calc' with respect to the computational method
        # DFTB/SSR can produce NACs, so we do not need to get NACME from CIoverlap
        # DFTB/SSR can compute the gradient of several states simultaneously.
        molecule.l_nacme = False
        self.re_calc = False

    def get_data(self, molecule, base_dir, bo_list, dt, istep, calc_force_only):
        """ Extract energy, gradient and nonadiabatic couplings from SSR method

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer,list bo_list: List of BO states for BO calculation
            :param double dt: Time interval
            :param integer istep: Current MD step
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        self.copy_files(istep)
        super().get_data(base_dir, calc_force_only)
        self.write_xyz(molecule)
        self.get_input(molecule, istep, bo_list)
        self.run_QM(base_dir, istep, bo_list)
        self.extract_QM(molecule, bo_list)
        self.move_dir(base_dir)

    def copy_files(self, istep):
        """ Copy necessary scratch files in previous step

            :param integer istep: Current MD step
        """
        # Copy required files to read initial guess
        if (self.guess == "read" and istep >= 0):
            # After T = 0.0 s
            shutil.copy(os.path.join(self.scr_qm_dir, "eigenvec.bin"), \
                os.path.join(self.scr_qm_dir, "../eigenvec.bin.pre"))

    def get_input(self, molecule, istep, bo_list):
        """ Generate DFTB+ input files: geometry.gen, dftb_in.hsd

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
        """
        # Make 'geometry.gen' file
        os.system("xyz2gen geometry.xyz")
        if (self.l_periodic):
            # Substitute C to S in first line
            file_be = open('geometry.gen', 'r')
            file_af = open('tmp.gen', 'w')
            first_row = True
            for row in file_be:
                if (first_row):
                    row = f'{molecule.nat_qm} S\n'
                    first_row = False
                file_af.write(row)
            # Add gamma-point and cell lattice information
            geom_periodic = textwrap.dedent(f"""\
            {0.0:15.8f} {0.0:15.8f} {0.0:15.8f}
            {self.a_axis[0]:15.8f} {self.a_axis[1]:15.8f} {self.a_axis[2]:15.8f}
            {self.b_axis[0]:15.8f} {self.b_axis[1]:15.8f} {self.b_axis[2]:15.8f}
            {self.c_axis[0]:15.8f} {self.c_axis[1]:15.8f} {self.c_axis[2]:15.8f}
            """)
            file_af.write(geom_periodic)
            file_be.close()
            file_af.close()
            os.rename('tmp.gen', 'geometry.gen')

        # Make 'point_charges.xyz' file used in electrostatic charge embedding of QM/MM
        if (self.embedding == "electrostatic"):
            # Make 'point_charges.xyz' file
            input_geom_pc = ""
            for iat in range(molecule.nat_qm, molecule.nat):
                input_geom_pc += "".join([f"{i:15.8f}" for i in molecule.pos[iat] * au_to_A])
                input_geom_pc += f"  {molecule.mm_charge[iat - molecule.nat_qm]:8.4f}\n"

            # Write 'point_charges.xyz' file
            file_name = "point_charges.xyz"
            with open(file_name, "w") as f:
                f.write(input_geom_pc)

        # Make 'dftb_in.hsd' file
        input_dftb = ""

        # Geometry Block
        input_geom = textwrap.dedent(f"""\
        Geometry = GenFormat{{
          <<< 'geometry.gen'
        }}
        """)
        input_dftb += input_geom

        # Hamiltonian Block
        input_ham_init = textwrap.dedent(f"""\
        Hamiltonian = DFTB{{
        """)
        input_dftb += input_ham_init

        # SCC-DFTB option
        if (self.l_scc):
            input_ham_scc = textwrap.indent(textwrap.dedent(f"""\
              SCC = Yes
              SCCTolerance = {self.scc_tol}
              MaxSCCIterations = {self.scc_max_iter}
            """), "  ")
            input_dftb += input_ham_scc

            # Read atomic spin constants used in DFTB/SSR
            if (self.l_range_sep):
                spin_constant = ("\n" + " " * 14).join([f"  {itype} = {{ {spin_w_lc[f'{itype}']} }}" for itype in self.atom_type])
            else:
                spin_constant = ("\n" + " " * 14).join([f"  {itype} = {{ {spin_w[f'{itype}']} }}" for itype in self.atom_type])
            input_ham_spin = textwrap.indent(textwrap.dedent(f"""\
              SpinConstants = {{
                ShellResolvedSpin = Yes
              {spin_constant}
              }}
            """), "  ")
            input_dftb += input_ham_spin

            # Long-range corrected DFTB (LC-DFTB) option
            if (self.l_range_sep):
                input_ham_lc = textwrap.indent(textwrap.dedent(f"""\
                  RangeSeparated = LC{{
                    Screening = {self.lc_method}{{}}
                  }}
                """), "  ")
                input_dftb += input_ham_lc

            # Add point charges to Hamiltonian used in electrostatic charge embedding of QM/MM
            if (self.embedding == "electrostatic"):
                input_ham_pc = textwrap.indent(textwrap.dedent(f"""\
                  ElectricField = PointCharges{{
                    CoordsAndCharges [Angstrom] = DirectRead{{
                      Records = {molecule.nat_mm}
                      File = "point_charges.xyz"
                    }}
                  }}
                """), "  ")
                input_dftb += input_ham_pc

        if (self.l_periodic):
            input_ham_periodic = textwrap.indent(textwrap.dedent(f"""\
              KPointsAndWeights = {{
                0.0 0.0 0.0 1.0
              }}
            """), "  ")
            input_dftb += input_ham_periodic

        angular_momentum = ("\n" + " " * 10).join([f"  {itype} = '{max_l[f'{itype}']}'" for itype in self.atom_type])
        input_ham_basic = textwrap.dedent(f"""\
          Charge = {molecule.charge}
          MaxAngularMomentum = {{
          {angular_momentum}
          }}
          SlaterKosterFiles = Type2FileNames{{
            Prefix = '{self.sk_path}'
            Separator = '-'
            Suffix = '.skf'
            LowerCaseTypeName = No
          }}
        }}
        """)
        input_dftb += input_ham_basic

        # Analysis Block
        input_analysis = textwrap.dedent(f"""\
        Analysis = {{
          CalculateForces = Yes
          WriteBandOut = Yes
          WriteEigenvectors = Yes
          MullikenAnalysis = Yes
        }}
        """)
        input_dftb += input_analysis

        # Options Block
        input_options = textwrap.dedent(f"""\
        Options = {{
          WriteDetailedXml = Yes
          WriteDetailedOut = Yes
          TimingVerbosity = -1
        }}
        """)
        input_dftb += input_options

        # REKS Block

        # Energy functional options
        if (self.l_ssr22):
            # Set active space and energy functional used in SA-REKS
            space = "SSR22"
            if (molecule.nst == 1):
                energy_functional = "{ 'PPS' }"
                all_states = "No"
            elif (molecule.nst == 2):
                energy_functional = "{ 'PPS' 'OSS' }"
                all_states = "No"
            elif (molecule.nst == 3):
                energy_functional = "{ 'PPS' 'OSS' }"
                all_states = "Yes"
            else:
                error_message = "SSR(2,2) can calculate up to 3 electronic states!"
                error_vars = f"Molecule.nstates = {molecule.nst}"
                raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

            # Include state-interaction terms to SA-REKS; SI-SA-REKS
            if (self.l_state_interactions):
                do_ssr = "Yes"
            else:
                do_ssr = "No"

            # Read 'eigenvec.bin' from previous step
            if (self.guess == "read"):
                if (istep == -1):
                    if (os.path.isfile(self.guess_file)):
                        # Copy guess file to currect directory
                        shutil.copy(self.guess_file, os.path.join(self.scr_qm_dir, "eigenvec.bin"))
                        restart = "Yes"
                    else:
                        restart = "No"
                elif (istep >= 0):
                    # Move previous file to currect directory
                    os.rename("../eigenvec.bin.pre", "./eigenvec.bin")
                    restart = "Yes"
            elif (self.guess == "h0"):
                restart = "No"

            # Scale the atomic spin constants
            if (self.tuning != None):
                spin_tuning = ""
                spin_tuning += "{"
                for scale_W in self.tuning:
                    spin_tuning += f" {scale_W} "
                spin_tuning += "}"
            else:
                spin_tuning = "{}"

        # CP-REKS algorithm options
        if (self.cpreks_grad_alg == "pcg"):
            cpreks_alg = "ConjugateGradient"
            preconditioner = "Yes"
        elif (self.cpreks_grad_alg == "cg"):
            cpreks_alg = "ConjugateGradient"
            preconditioner = "No"
        elif (self.cpreks_grad_alg == "direct"):
            cpreks_alg = "Direct"
            preconditioner = "No"
        else:
            error_message = "Invalid algorithms for CP-REKS problem!"
            error_vars = f"cpreks_grad_alg = {self.cpreks_grad_alg}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        # Save memory in cache to reduce computational cost in CP-REKS
        if (self.l_save_memory):
            memory = "Yes"
        else:
            memory = "No"

        # NAC calculation options
        if (molecule.nst == 1 or not self.l_state_interactions):
            # Single-state REKS or SA-REKS state
            self.nac = "No"
        else:
            # SSR state
            if (self.calc_coupling):
                # SHXF, SH, Eh need NAC calculations
                self.nac = "Yes"
            else:
                # BOMD do not need NAC calculations
                self.nac = "No"

        # Relaxed density options; It is determined automatically
        if (molecule.l_qmmm and self.embedding == "electrostatic"):
            relaxed_density = "Yes"
        else:
            relaxed_density = "No"

        input_reks = textwrap.dedent(f"""\
        REKS = {space}{{
          Energy = {{
            Functional = {energy_functional}
            StateInteractions = {do_ssr}
            IncludeAllStates = {all_states}
          }}
          TargetState = {bo_list[0] + 1}
          ReadEigenvectors = {restart}
          FONmaxIter = 50
          shift = {self.shift}
          SpinTuning = {spin_tuning}
          Gradient = {cpreks_alg}{{
            CGmaxIter = 100
            Tolerance = {self.cpreks_grad_tol}
            Preconditioner = {preconditioner}
            SaveMemory = {memory}
          }}
          RelaxedDensity = {relaxed_density}
          NonAdiabaticCoupling = {self.nac}
          VerbosityLevel = 1
        }}
        """)
        input_dftb += input_reks

        # ParserOptions Block
        if (self.version == "19.1"):
            error_message = "SSR not implemented in this version!"
            error_vars = f"version = {self.version}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")
        elif (self.version == "20.1"):
            parser_version = 8

        input_parseroptions = textwrap.dedent(f"""\
        ParserOptions = {{
          ParserVersion = {parser_version}
        }}
        """)
        input_dftb += input_parseroptions

        # Write 'dftb_in.hsd' file
        file_name = "dftb_in.hsd"
        with open(file_name, "w") as f:
            f.write(input_dftb)

    def run_QM(self, base_dir, istep, bo_list):
        """ Run SSR calculation and save the output files to qm_log directory

            :param string base_dir: Base directory
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
        """
        # Set run command
        qm_command = os.path.join(self.qm_path, "dftb+")
        # OpenMP setting
        os.environ["OMP_NUM_THREADS"] = f"{self.nthreads}"
        command = f"{qm_command} > log"
        # Run DFTB+ method for molecular dynamics
        os.system(command)

        # Copy the output file to 'qm_log' directory
        tmp_dir = os.path.join(base_dir, "qm_log")
        if (os.path.exists(tmp_dir)):
            log_step = f"log.{istep + 1}.{bo_list[0]}"
            shutil.copy("log", os.path.join(tmp_dir, log_step))

    def extract_QM(self, molecule, bo_list):
        """ Read the output files to get BO information

            :param object molecule: Molecule object
            :param integer,list bo_list: List of BO states for BO calculation
        """
        # Read 'log' file
        file_name = "log"
        with open(file_name, "r") as f:
            log_out = f.read()

        # Read 'detailed.out' file
        file_name = "detailed.out"
        with open(file_name, "r") as f:
            detailed_out = f.read()

        # Energy
        if (molecule.nst == 1):
            # Single-state REKS
            tmp_e = 'Spin' + '\n\s+\w+\s+([-]\S+)(?:\s+\S+){3}' * molecule.nst
            energy = re.findall(tmp_e, log_out)
            energy = np.array(energy, dtype=np.float64)
        else:
            if (self.l_state_interactions):
                # SSR state
                energy = re.findall('SSR state\s+\S+\s+([-]\S+)', log_out)
                energy = np.array(energy, dtype=np.float64)
            else:
                # SA-REKS state
                tmp_e = 'Spin' + '\n\s+\w+\s+([-]\S+)(?:\s+\S+){3}' * molecule.nst
                energy = re.findall(tmp_e, log_out)
                energy = np.array(energy[0], dtype=np.float64)

        for ist in range(molecule.nst):
            molecule.states[ist].energy = energy[ist]

        # Force
        if (self.nac == "Yes"):
            # SHXF, SH, Eh : SSR state
            if (molecule.l_qmmm):
                for ist in range(molecule.nst):
                    tmp_g = f' {ist + 1} st state \(SSR\)' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                    grad = re.findall(tmp_g, log_out)
                    grad = np.array(grad[0], dtype=np.float64)
                    grad = grad.reshape(molecule.nat_qm, 3, order='C')
                    molecule.states[ist].force[0:molecule.nat_qm] = - grad
                    if (self.embedding == "electrostatic"):
                        tmp_f = 'Forces on external charges' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_mm
                        force = re.findall(tmp_f, detailed_out)
                        force = np.array(force[0], dtype=np.float64)
                        force = force.reshape(molecule.nat_mm, 3, order='C')
                        molecule.states[ist].force[molecule.nat_qm:molecule.nat] = force
            else:
                for ist in range(molecule.nst):
                    tmp_g = f' {ist + 1} st state \(SSR\)' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                    grad = re.findall(tmp_g, log_out)
                    grad = np.array(grad[0], dtype=np.float64)
                    grad = grad.reshape(molecule.nat_qm, 3, order='C')
                    molecule.states[ist].force = - np.copy(grad)
        else:
            # BOMD : SSR state, SA-REKS state or single-state REKS
            if (molecule.l_qmmm):
                tmp_g = f' {bo_list[0] + 1} state \(\w+[-]*\w+\)' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                grad = re.findall(tmp_g, log_out)
                grad = np.array(grad[0], dtype=np.float64)
                grad = grad.reshape(molecule.nat_qm, 3, order='C')
                molecule.states[bo_list[0]].force[0:molecule.nat_qm] = - grad
                if (self.embedding == "electrostatic"):
                    tmp_f = 'Forces on external charges' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_mm
                    force = re.findall(tmp_f, detailed_out)
                    force = np.array(force[0], dtype=np.float64)
                    force = force.reshape(molecule.nat_mm, 3, order='C')
                    molecule.states[bo_list[0]].force[molecule.nat_qm:molecule.nat] = force
            else:
                tmp_g = f' {bo_list[0] + 1} state \(\w+[-]*\w+\)' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                grad = re.findall(tmp_g, log_out)
                grad = np.array(grad[0], dtype=np.float64)
                grad = grad.reshape(molecule.nat_qm, 3, order='C')
                molecule.states[bo_list[0]].force = - np.copy(grad)

        # NAC
        if (self.nac == "Yes"):
            kst = 0
            for ist in range(molecule.nst):
                for jst in range(ist + 1, molecule.nst):
                    tmp_c = 'non-adiabatic coupling' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                    nac = re.findall(tmp_c, log_out)
                    nac = np.array(nac[kst], dtype=np.float64)
                    nac = nac.reshape(molecule.nat_qm, 3, order='C')
                    molecule.nac[ist, jst] = nac
                    molecule.nac[jst, ist] = - nac
                    kst += 1


# coding: UTF-8
import sys
import time
import itertools
import json
import numpy as np
from numpy.fft import fft
from collections import deque

import threading as th
import flask as fl
from flask import Flask, render_template, request, Response, redirect, url_for
from server import MuseServer
from search_yt_by_word import Search
from detect_movement import pravila

app = Flask(__name__)
calculatedBeat = 0
mainLock = th.RLock()
search = Search()
currentSong = ""
yt_id = ""
itunes_link = ""

acc_data_num = 0
dq = deque()
headLock = th.RLock()

stopserve = False
stopserveLock = th.Lock()

def serve():
    global calculatedBeat, currentSong, yt_id, itunes_link, stopserve

    print("Starting server!!!")

    t = 5
    # time.sleep(t)
    while 1:
        stopserveLock.acquire()
        if stopserve:
            stopserve = False
            stopserveLock.release()
            print("Returning server1")
            return
        stopserveLock.release()

        acc_data_num = 0
        for i in range(10):
            stopserveLock.acquire()
            if stopserve:
                stopserve = False
                stopserveLock.release()
                print("Returning server2")
                return
            stopserveLock.release()

            server.acc_lock.acquire()
            tmp = list(zip(*server.acc_list))
            raw_tmp = tmp[:]
            server.acc_lock.release()

            n2 = len(tmp)
            tmp = tmp[acc_data_num:]
            acc_data_num = n2
            podatki = pravila(raw_tmp, duration=t/10)
            headLock.acquire()
            if "ndesno" in podatki:
                dq.append("ndesno")
            headLock.release()

            time.sleep(t / 10.0)

        server.acc_lock.acquire()
        tmp = list(zip(*server.acc_list))
        server.acc_list = []
        server.acc_lock.release()

        if tmp and tmp[0]:
            size = len(tmp[0])
            fs = [i / t for i in range(size // 2)]
            tmp = [np.array(x) - sum(x) / len(x) for x in tmp]
            acc_fft = [fft(x)[: size // 2] for x in tmp]
            imax = [sorted(enumerate(x), key=lambda k: abs(k[1]))[-1] for x in acc_fft]
            print(fs[imax[0][0]], fs[imax[0][0]] * 60)

            mainLock.acquire()
            calculatedBeat = fs[imax[0][0]] * 60
            local_beat = calculatedBeat
            mainLock.release()

            search_result = search.search_all(local_beat)

            mainLock.acquire()
            currentSong, yt_id, itunes_link = search_result
            mainLock.release()

ser = th.Thread(target=serve)

@app.route('/restart')
def restart():
    print("restarting serve")
    global calculatedBeat, yt_id, ser, stopserve, server

    mainLock.acquire()
    calculatedBeat = 0
    yt_id = ""
    search.clear()
    mainLock.release()

    server.acc_lock.acquire()
    server.acc_list = []
    server.acc_lock.release()

    stopserveLock.acquire()
    stopserve = True
    stopserveLock.release()
    while True:
        stopserveLock.acquire()
        local = stopserve
        stopserveLock.release()
        if not local:
            ser = th.Thread(target=serve)
            ser.start()
            break

    return "100"

@app.route('/getnext')
def getnext():
    print("bpm")
    bpm = int(request.args.get("bpm", 80))
    print("bpm2")
    try:
        search_result = search.search_all(bpm)
    except Exception as e:
        print(e)
    print(search_result)

    song, yt, itunes = search_result

    return "{\"bpm\": %d, \"song\": \"%s\", \"yt_id\": \"%s\", " \
           "\"itunes_link\": \"%s\"}\n\n" % (bpm, song, yt, itunes)

@app.route('/')
def index():
    global calculatedBeat, yt_id

    mainLock.acquire()
    calculatedBeat = 0
    yt_id = ""
    search.clear()
    mainLock.release()

    return render_template('index.html')

@app.route('/status')
def status():
    def events():
        while 1:
            server.senstate_lock.acquire()
            event = server.senstate
            server.senstate_lock.release()
            if event:
                yield "event: status\ndata: %s\n\n" % (json.dumps(event))

            time.sleep(0.2)  # an artificial delay
    return Response(events(), content_type='text/event-stream')

@app.route('/headswipe')
def headswipe():
    def events():
        while 1:
            headLock.acquire()
            event = len(dq) and dq.pop()
            dq.clear()
            headLock.release()
            if event != 0:
                print("head")
                yield "event: head\ndata: %s\n\n" % (event)

            time.sleep(0.1)  # an artificial delay
    return Response(events(), content_type='text/event-stream')

@app.route('/blink')
def blink():
    def events():
        while 1:
            server.blink_lock.acquire()
            event = len(server.blink_list) and server.blink_list.pop()
            server.blink_lock.release()
            if event != 0:
                yield "event: blink\ndata: %s\n\n" % (event)

            time.sleep(0.05)  # an artificial delay
    return Response(events(), content_type='text/event-stream')

@app.route('/jaw')
def jaw():
    def events():
        while 1:
            server.jaw_lock.acquire()
            event = len(server.jaw_list) and server.jaw_list.pop()
            server.jaw_lock.release()
            if event != 0:
                yield "event: jawclench\ndata: %s\n\n" % (event)

            time.sleep(0.05)  # an artificial delay
    return Response(events(), content_type='text/event-stream')

@app.route('/bpm')
def bpm():
    def events():
        while 1:
            mainLock.acquire()
            tmpCalculatedBeat = calculatedBeat
            mainLock.release()
            if tmpCalculatedBeat != 0 and yt_id != "" and currentSong != "":
                yield "event: calculated\ndata: {\"bpm\": %d, \"song\": \"%s\", \"yt_id\": \"%s\", \"itunes_link\": \"%s\"}\n\n" % (tmpCalculatedBeat, currentSong, yt_id, itunes_link)

            server.acc_lock.acquire()
            tmpMove = 0
            if server.acc_list and server.acc_list[-1]:
                tmpMove = server.acc_list[-1][0]
            server.acc_lock.release()
            if tmpMove != 0:
                yield "event: move\ndata: %d\n\n" % (tmpMove)

            #yield "event: test\ndata: %d\n\n" % (calculatedBeat)
            time.sleep(0.02)  # an artificial delay
    return Response(events(), content_type='text/event-stream')

if __name__ == '__main__':
    server = MuseServer()
    server.start()

    ser.start()
    print("Starting local server")
    # serve()
    app.run("0.0.0.0", port=8080, threaded=True)


# importing above defined libraries to
# implement the functionalities
#from pillow import ImageGrab, ImageOps
from pynput.mouse import Button,Controller
import pyscreenshot as ImageGrab,ImageOps
import pyautogui
import time
import numpy as np

class cordinates():

	# coordinates of replay button to start the game
	replaybutton =(337, 249)
	# this coordinates represent the top-right coordinates
	# that will be used to define the front box
	dinasaur = (151, 283 )

def restartGame():

	# using pyautogui library, we are clicking on the
	# replay button without any user interaction
	pyautogui.click(cordinates.replaybutton)

	# we will keep our Bot always down that
	# will prevent him to get hit by bird
	pyautogui.keyDown('down')

def press_space():

	# releasing the Down Key
	pyautogui.keyUp('down')

	# pressing Space to overcome Bush
	pyautogui.keyDown('space')

	# so that Space Key will be recognized easily
	time.sleep(0.05)

	# printing the "Jump" statement on the
	# terminal to see the current output
	print("jump")
	time.sleep(0.10)

	# releasing the Space Key
	pyautogui.keyUp('space')

	# again pressing the Down Key to keep my Bot always down
	pyautogui.keyDown('down')

def imageGrab():
	# defining the coordinates of box in front of dinosaur
	box = (cordinates.dinasaur[0]+70, cordinates.dinasaur[1],
		cordinates.dinasaur[0]+160, cordinates.dinasaur[1]+2)

	# grabbing all the pixels values in form of RGB tupples
	image = ImageGrab.grab(box)

	# converting RGB to Grayscale to
	# make processing easy and result faster
	grayImage = ImageOps.grayscale(image)

	# using numpy to get sum of all grayscale pixels
	a = np.array(grayImage.getcolors())

	# returning the sum
	print(a.sum())
	return a.sum()

#mouse=Controller()
# function to restart the game
restartGame()
while True:
	# 435 is the sum of white pixels values of box.
	# You may get different value is you are taking bigger
	# or smaller box than the box taken in this article.
	# if value returned by "imageGrab" function is not equal to 435,
	# it means either bird or bush is coming towards dinosaur
	#print(pyautogui.position())
	#print(mouse.position)
	if(imageGrab()!= 435):
		press_space()
		# time to recognize the operation performed by above function
		time.sleep(0.1)


#!/usr/bin/env python
# coding: utf-8

# # 2. Object Oriented Programming

# In this lecture, we'll go over the basics of writing object oriented code. Because this is a large shift in programming syntax (and thinking), we'll be spending several lectures building up intuition for it.

# ## Part I: Methods and Attributes

# *Writing* classes can be intimidating at first, and doing so involves some general boilerplate that can be confusing at first.

# *Using* classes, on the other hand, is generally fairly intuitive. One reason for this:
#
# ***Everything in Python is an Object***. This means we are fairly used to using objects, even if we don't know it.

# Let's begin with a simple example. I'm going to initialize a basic `numpy` array.

# In[2]:

import numpy as np

arr = np.ones((5,6))*15.5
arr

# This simple array, `arr`, is an *object*. Objects in python have both *attributes* and *methods*. Both are accessed via what we call *dot notation*. Using dot notation, we can access an attribute of an object as follows:

# In[3]:

arr.shape

# `shape` is an attribute of all `numpy` arrays. By executing `arr.shape`, we retrieved this attribute. For documented classes, available attributes are generally listed in the documentation. The are also generally available for tab completion.
#
# Attributes are called as such because they are not *callable*. We can query them, but they are more like "properties" of the object, which are often used to make our code more robust (e.g., instead of hard-coding numbers into a code, we can flexibly adjust based on the sizes and shapes of the arrays within).

# Alternatively, `methods` of an object are callable functions that the object knows how to run, generally upon itself (with or without additional information). Returning to our `arr` array, we can call a method that reshapes the array:

# In[9]:

arr.reshape(6,5)

# We use the same dot notation as for attribute access, but we then give parenthesis and arguments as necessary, as we would when calling a function. We can see here that the output of this method is a reshaping of the array. Most Python objects (especially the basic datatypes, such as `lists`, `arrays`, `dictionaries`, and even `floats`, `ints`, and `strings`) have many useful attributes and methods attached to them, and knowing them can save you time in your coding.

# In examples previous, we've seen plenty of quick method and attribute use --- for example, taking a string and running `.startswith()` on it, or `.strip()`.
#
# But how do we give ourselves the ability to write `classes` which give us this sort of attribute/method access flexibility?

# ## Part II: Defining Classes

# Classes are defined via the `class` call, and all contain, at minimum, a method known as `__init__`.

# In[10]:

class MyObject():
    def __init__(self):
        pass

# Methods of a class are defined just like normal functions, but are all indented within the block of the class. They also all contain `self` as the first argument of the method. Why is this?

# When we created an array above, we did something known as *instantiating an object*. Somewhere in the `numpy` library is a master template for what all `numpy arrays` look like, and we put in some arguments to determine what our specific instance of the class would look like. (In our context, this is the value of the array, 15.5, and its shape).

# After this, we were able to to call methods of the class via its name (that we set it to in our code) plus the dot notation. The mapping of the *internal* methods of the class to the *external* method call is handled by the use of the `self` argument.

# As long as we remember to set `self` as the first argument of every method (followed by any other arguments you would typically add), our classes should work properly. There's one other critical element of using self:

# In[11]:

class MyObject():
    def __init__(self,parameter):
        self.parameter = parameter

# We've now given our class it's first initialized attribute, `parameter`. This parameter is fed in when we *instantiate* the object, not as any additional call:

# In[12]:

obj1 = MyObject(parameter=5)

# In[13]:

obj1.parameter

# By setting the *class attribute* `self.parameter = parameter`, we make it accessible both anywhere within the scope of our class (even in other methods) as well as outside, as a dot-accessible attribute. Any variable otherwise created in any class method that doesn't have `self.` in front of it won't be accessible outside the class as an attribute. Notice we never make reference to `__init__()` when interacting with the object from the outside.

# Let's move to a slightly more physical example:

# In[16]:

class Body():
    def __init__(self,name,mass,radius):
        self.name = name
        self.mass = mass
        self.radius = radius

# In[17]:

import astropy.units as u
sun = Body('Sun',1*u.Msun,1*u.Rsun)

# In[18]:

sun.mass

# In[19]:

sun.radius

# Our new `Body` class takes the inputs to the initialization function and decides that they should immediately be class attributes (accessible anywhere in the class or outside, as demonstrated above.)
#
# We need not set every input to `__init__` to *be* a class attribute. For example, if a class reads in a string pointer to some `fits` file, we may load the file and then not care about having access to that input string anymore.

# ## Creating Methods

# The init method just handles stuff we need to put into to the class to get it up and running. Let's define our first "real" method.

# In[26]:

class Body():
    def __init__(self,name,mass,radius):
        self.name = name
        self.mass = mass
        self.radius = radius
    def calc_density(self):
        vol = (4/3.) * np.pi * self.radius**3
        return (self.mass / vol).to((u.g/u.cm**3))

# In[27]:

sun = Body('Sun',1*u.Msun,1*u.Rsun)
sun.calc_density()

# You may be thinking that we could, technically, define a class *attribute* called density since it depends only on the previously defined parameters. It's true that we could've done this, but, in this case, changing the mass attribute wouldn't propogate to the density. By forcing the calculation each time, we ensure that changes to the mass or radius then are reflected in a *query* of the density.

# ### Exercise
#
# Copy the class below, add an input attribute $x$, representing position. Assume it is a vector containing the x, y, z position of the object.
# Then, add a method to the class `calc_grav`, which takes as input another instance of a `Body` class, and uses its own position and mass, and the other body's position and mass, to calculate the force $F_{\rm grav}$ between them via Newton's law. Create two bodies and test your methods!

# In[ ]:

class Body():
    def __init__(self,mass,radius):
        self.mass = mass
        self.radius = radius
    def calc_density(self):
        vol = (4/3.) * np.pi * self.radius**3
        return self.mass / vol

# # PSF Photometry

# We're all familiar (from 255) with aperture photometry. Aperture photometry sets a radius within which to sum the flux from some astronomical source.
#
# One question to ask when performing aperture photometry is what radius to choose. Too large, and you introduce noise from pixels that are sky. Too small, and you miss some flux from the actual star.
#
# We also know that the true "point source" of a star is spread out by the point spread function of the telescope/optics/detector. If we can *estimate* the shape (width) of the PSF in an image of sources, then we can construct a model which weights each pixel based on it's estimated fractional contribution to the stellar flux. Thus, pixels near the center of the stellar distribution are summed in full, while those near the edge are not included (or only just barely).

# To construct a pipeline that uses this technique, we'll need to
#
# - Read in an astronomical image with stars (and other point sources)
# - Perform the releant cleanings and calibrations to the image
# - Algorithmically locate the peaks (point sources) in the image, and carefully calculate their centers
# - Estimate the PSF and construct the normalized model
# - Multiply the weight model by each source and calculate the sum using a weighting formula (using cutouts and our known positions)
#

# *Jamboard Discussion*: Based on the list above, discuss some reasons why moving to `classes` might be beneficial in trying to create such a pipeline.
#
# bit.ly/330jamboard

# ### Formulas Needed

# **Centroid**
#
# $$
# x_{\rm com} = \frac{\sum x_i \hat{f}_i}{\sum \hat{f}_i}
# $$
#
# $$
# y_{\rm com} = \frac{\sum y_i \hat{f}_i}{\sum \hat{f}_i}
# $$
#
# where the $x_i$ are the pixel numbers and $\hat{f}_i$ are the fluxes in those pixels. This formula amounts to a light weighted mean --- pixels with larger fluxes will push the mean toward their pixel value. (This is also just a center of mass formula).

# **2nd Moment**
#
# The centroid is the first moment of the light distribution. The second moment gives us an estimate of the width of a distribution. (To use stats terms, a mean is a first moment measurement, the standard deviation is a second moment measurement).
#
# We will use the following formulas in the lab to determine the x and y widths of distributions:
#
# $$
# \sigma_x = \left[\frac{\sum x_i^2 \hat{f}_i}{\sum \hat{f}_i} - x_{\rm com}^2\right]^{1/2}
# $$
#
# $$
# \sigma_y = \left[\frac{\sum y_i^2 \hat{f}_i}{\sum \hat{f}_i} - y_{\rm com}^2\right]^{1/2}
# $$

# **PSF Photometry**
#
# Finally, the actual step of PSF photometry is performed via
#
# $$
# f_{\rm PSF} = \frac{\sum \hat{f_i} p_i}{\sum p_i^2},
# $$
#
# where $\hat{f}_i$ are our pixel fluxes from before, and $p_i$ is our model for the PSF. One can model the 2D PSF with several functional forms. The Moffat profile is one of the closest matches to real stars, but it requires more than one shape parameter. Since our second moment only gives us one, we'll use a 2D Gaussian in the lab.


from flask import Flask, Response
import cv2
import os
import argparse
import numpy as np

from maindash import recent_locations, get_config
from stream_analysis.motion_detection.models import *  # set ONNX_EXPORT in models.py
from stream_analysis.motion_detection.utils.datasets import *
from stream_analysis.motion_detection.utils.utils import *

class Detector(object):

    def __init__(self,weights_arg,cfg_arg,names_arg,device_arg,source_arg,save_img=False):

        self.weights = weights_arg
        self.cfg = cfg_arg
        self.source = source_arg
        self.names = names_arg
        self.device = device_arg
        self.webcam = self.source == '0' or self.source.startswith('rtsp') or self.source.startswith('http') or self.source.endswith('.txt')
        print('in')

        self.imgsz = 512
        self.out = 'output'
        self.half = False
        self.view_img = False
        self.save_txt = False
        self.augment = False
        self.agnostic_nms = False
        self.classes = None
        self.conf_thres = 0.3
        self.fourcc = 'mp4v'
        self.iou_thres = 0.6

        # Initialize
        self.device = torch_utils.select_device(device=self.device)
        if os.path.exists(self.out):
            shutil.rmtree(self.out)  # delete output folder
        os.makedirs(self.out)  # make new output folder
        # Initialize model
        self.model = Darknet(self.cfg, self.imgsz)

        # Load weights
        if self.weights.endswith('.pt'):  # pytorch format
            self.model.load_state_dict(torch.load(self.weights, map_location=self.device)['model'])
        else:  # darknet format
            load_darknet_weights(self.model, self.weights)

        # Second-stage classifier
        self.classify = False
        if self.classify:
            self.modelc = torch_utils.load_classifier(name='resnet101', n=2)  # initialize
            self.modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=self.device)['model'])  # load weights
            modelc.to(self.device).eval()

        # Eval mode
        self.model.to(self.device).eval()

        # Fuse Conv2d + BatchNorm2d layers
        # model.fuse()

        # Half precision
        self.half = self.half and self.device.type != 'cpu'  # half precision only supported on CUDA
        if self.half:
            self.model.half()

        # Set Dataloader
        self.vid_path, self.vid_writer = None, None
        if self.webcam:
            self.view_img = True
            torch.backends.cudnn.benchmark = True  # set True to speed up constant image size inference
            self.dataset = LoadStreams(self.source, img_size=self.imgsz)
        else:
            self.save_img = True
            self.dataset = LoadImages(self.source, img_size=self.imgsz)

        # Get names and colors
        self.names = load_classes(self.names)
        self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(self.names))]

        # Run inference
        t0 = time.time()
        self.img = torch.zeros((1, 3, self.imgsz, self.imgsz), device=self.device)  # init img
        _ = self.model(self.img.half() if self.half else self.img.float()) if self.device.type != 'cpu' else None  # run once
        print(self.dataset)

    def get_frame(self,path, img, im0s, vid_cap):
        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = torch_utils.time_synchronized()
        pred = self.model(img, augment=self.augment)[0]
        t2 = torch_utils.time_synchronized()

        # to float
        if self.half:
            pred = pred.float()

        # Apply NMS
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres,
                                   multi_label=False, classes=self.classes, agnostic=self.agnostic_nms)

        # Apply Classifier
        if self.classify:
            pred = apply_classifier(pred, self.modelc, img, im0s)

        # Process detections
        for i, det in enumerate(pred):  # detections for image i
            label_det = []
            if self.webcam:  # batch_size >= 1
                p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
            else:
                p, s, im0 = path, '', im0s

            save_path = str(Path(self.out) / Path(p).name)
            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  #  normalization gain whwh
            if det is not None and len(det):
                # Rescale boxes from imgsz to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].detach().unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += '%g %ss, ' % (n, self.names[int(c)])  # add to string

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    if self.save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        with open(save_path[:save_path.rfind('.')] + '.txt', 'a') as file:
                            file.write(('%g ' * 5 + '\n') % (cls, *xywh))  # label format

                    if self.view_img:  # Add bbox to image
                        label = '%s %.2f' % (self.names[int(cls)], conf)
                        label_det.append(label)
                        plot_one_box(xyxy, im0, label=label, color=self.colors[int(cls)])

            label_det.reverse()
            # Stream results
            if self.view_img:
                return (p, im0, det, label_det)

def gen(camera):

    config = get_config()
    transformation_matrix = config["transformation_matrix"]

    for path, img, im0s, vid_cap in camera.dataset:
        (p, im0, det, labels) = camera.get_frame(path, img, im0s, vid_cap)

        if det is not None:

            for i in range(len(labels)):

                cl = 'green'
                if labels[i].split(' ')[0] == 'black_car':
                    cl = 'black'
                else: cl = 'red'

                recent_locations.pop(0)
                new_point = np.array([det[i][0], det[i][1], 1])
                new_point = np.matmul(transformation_matrix,new_point)

                recent_locations.append((int(new_point[0]), int(new_point[1]), cl))

        dim = tuple(config["camera_aspect_ratio"])
        im0 = cv2.resize(im0, dim, interpolation = cv2.INTER_AREA)
        _ , jpeg = cv2.imencode('.jpg', im0)
        frame = jpeg.tobytes()

        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')


from copy import deepcopy
import matplotlib.pylab as pylab
import seaborn as sns
import numpy as np
import argparse
import os
from tqdm import tqdm
import torch
import pickle
from glob import glob
#from inhibitor import *
import skimage.transform
import elephant_fix
from pycocotools.coco import COCO
from numpy import array, argsort, zeros

from vision_utils.boxutils import *
from vision_utils.generic_utils import dict_to_param_str, str2bool, find
from maskrcnn_benchmark.structures.bounding_box import BoxList
from maskrcnn_benchmark.structures.boxlist_ops import cat_boxlist
from maskrcnn_benchmark.structures.boxlist_ops import boxlist_iou

from postprocessing import *
from mask_ops import *

# this makes our figures bigger
d = .5
pylab.rcParams['figure.figsize'] = 20 * d, 12 * d
pylab.rc('image', interpolation='bilinear')

from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.data.datasets.coco import COCODataset
from predictor import COCODemo
from maskrcnn_benchmark.structures.boxlist_ops import boxlist_nms
from matplotlib.pyplot import imshow
from maskrcnn_benchmark.data.datasets.evaluation.coco import coco_eval
from maskrcnn_benchmark.data.datasets.coco import COCODataset

config_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda:0"])

def re_detect2(img, preds, min_score=0, max_score=1, inflation_factor=1, is_abs=False,
               min_keep_iou=.5):
    scores = preds.extra_fields['scores']
    labels = preds.extra_fields['labels'].numpy()
    masking_color = 128
    masked_images = []
    new_dets_map = {}
    for ibox, score in tqdm(enumerate(scores), total=len(preds)):
        bb = preds[[ibox]]
        if score > max_score or score < min_score:
            continue
        cur_label = bb.extra_fields['labels'][0]
        cur_mask = img * 0
        mask_from_a_box(cur_mask, bb.bbox[0], inflation_factor, is_abs=is_abs)
        masked_img = img * cur_mask + (1 - cur_mask) * masking_color
        preds_masked = coco_demo.compute_prediction(masked_img[:, :, ::-1])
        remove_masks(preds_masked)
        if len(preds_masked) > 0:
            ious = boxlist_iou(bb, preds_masked)
            ious_above_thresh = find((ious >= min_keep_iou).numpy().ravel())
            # throw away low-iou matches
            preds_masked = preds_masked[find(
                (ious >= min_keep_iou).numpy().ravel())]
            if len(preds_masked) > 0:
                new_dets_map[ibox] = dict(source_ind=ibox,
                                          source=bb, result=preds_masked,
                                          result_ious=ious[0,
                                                           ious_above_thresh],
                                          min_score=min_score,
                                          max_score=max_score, inflation_factor=inflation_factor,
                                          is_abs=is_abs, min_keep_iou=min_keep_iou)
    return new_dets_map

def merge_extra_detections2(preds, new_dets_map, filter_by_label=False, ignore_orig_score_above=1,
                            ignore_orig_score_below=0):
    scores = preds.extra_fields['scores']
    new_preds = deepcopy(preds)
    n_updated = 0
    n0 = 0
    n1 = 0
    n2 = 0
    for ibox, score in enumerate(scores):
        if ibox not in new_dets_map:
            continue
        n0 += 1
        b = new_dets_map[ibox]
        res = b['result']
        src = b['source']
        res_labels = res.extra_fields['labels']
        src_labels = src.extra_fields['labels']
        if filter_by_label:
            is_same_label = find((res_labels == src_labels).numpy().ravel())
            if len(is_same_label) == 0:
                continue
            res = res[is_same_label]
        if len(res) == 0:
            continue
        n2 += 1
        cur_score = src.extra_fields['scores']

        if cur_score > ignore_orig_score_above or cur_score < ignore_orig_score_below:
            continue
        m = res.extra_fields['scores'].max()
        if m > cur_score:
            new_preds.extra_fields['scores'][ibox] = m
            n_updated += 1
    return new_preds, n_updated

def predict_and_show(v):
    image = U.load_image(v)
    img_orig = image.copy()
    dd = 1
    r = 1
    fix, ax = plt.subplots(2, 1, figsize=(12 * dd, r * 20 * dd))
    predictions = coco_demo.compute_prediction(image[:, :, ::-1])
    coco_demo.confidence_threshold = .2
    top_predictions = coco_demo.select_top_predictions(predictions)
    img_w_boxes = coco_demo.overlay_boxes(image, top_predictions, box_width=2)
    img_w_boxes_and_names = coco_demo.overlay_class_names(
        image, top_predictions)
    ax.flat[1].imshow(img_orig)
    coco.showAnns(U.get_anns(v))
    ax.flat[0].imshow(img_w_boxes_and_names)
    [r.axis('off') for r in ax.flat]
    return img_orig, top_predictions

def remove_masks(p):
    if 'mask' in p.extra_fields:
        del p.extra_fields['mask']
    return p

def remove_masks_file(fname):
    p = pickle.load(open(fname, 'rb'))

    if type(p['preds']) is list and len(p['preds']) == 0:
        p['preds'] = BoxList(torch.Tensor(0, 4), (300, 300))
        pickle.dump(p, open(fname, 'wb'))
    if 'mask' in p['preds'].extra_fields:
        del p['preds'].extra_fields['mask']
        pickle.dump(p, open(fname, 'wb'))
    return p

def detect_masked(img, coco_demo, preds, args, masking_color=128):
    masking_thresh = args.masking_thresh
    if args.all_at_once:
        if args.masking_type == 'bbox':
            mask = mask_from_boxes(img, preds, args)
            # mask = get_blocked_image(img,preds.bbox,inflation=masking_args.inflation_factor,is_abs=False,
            # block_color=masking_color)
        else:
            # mask = mask_from_masks(img,preds) # inflation factor not used here.
            if len(preds) == 0:
                mask = img * 0
            else:
                mask = mask_from_masks(img, preds, args)
        if args.negative_masking:
            mask = 1 - mask
        masked_img = img * mask + (1 - mask) * masking_color
        preds_masked = coco_demo.compute_prediction(masked_img[:, :, ::-1])
        return preds_masked
    else:
        all_preds = []
        ff = preds.extra_fields
        # if 'mask' in preds.extra_fields:
        # del preds.extra_fields['mask']
        n_above = sum(ff['scores'].numpy() >= args.masking_thresh)
        for bb, score, mask in tqdm(zip(preds.bbox, ff['scores'], ff['mask']), desc='masking one by one...', total=n_above):
            if score < args.masking_thresh:
                continue
            if args.masking_type == 'bbox':
                cur_mask = img * 0
                mask_from_a_box(cur_mask, bb, args.inflation_factor)
            else:
                cur_mask = np.repeat(mask.numpy().transpose(1, 2, 0), 3, 2)
            if args.negative_masking:
                cur_mask = 1 - cur_mask
            masked_img = img * cur_mask + (1 - cur_mask) * masking_color

            preds_masked = coco_demo.compute_prediction(masked_img[:, :, ::-1])

            all_preds.append(preds_masked)
            if 'mask' in preds_masked.extra_fields:
                del preds_masked.extra_fields['mask']
        if len(all_preds) > 0:
            all_preds = cat_boxlist(all_preds)
            all_preds = boxlist_nms(
                all_preds,
                .5,
                max_proposals=100,
                score_field="scores")
        else:
            all_preds = BoxList(torch.Tensor(0, 4), img.shape[:2])
        return all_preds

def initial_detection(coco_dataset, coco_demo, args, max_images=10**9, savedir=None):
    D = os.path.join(savedir, 'orig')
    os.makedirs(D, exist_ok=True)
    pickle.dump(args, open(os.path.join(D, 'args.pkl'), 'wb'))
    N = min(len(coco_dataset), max_images)

    # for i, (img, target, idx) in tqdm(enumerate(coco_dataset), total=N):
    for i in tqdm(range(N)):
        if i == max_images:
            break
        idx = i

        k = coco_dataset.id_to_img_map[idx]
        out_path = os.path.join(D, '{}.pkl'.format(k))
        if os.path.isfile(out_path):
            continue
        img, target, idx_ = coco_dataset[i]
        assert(idx_ == idx)
        image = np.array(img)
        with torch.no_grad():
            img = image.copy()
            preds = coco_demo.compute_prediction(img[:, :, ::-1])
            # if 'mask' in preds.extra_fields:
            # del preds.extra_fields['mask']
        pickle.dump(dict(idx=idx, k=k, preds=preds), open(out_path, 'wb'))

def masked_detection(coco_dataset, coco_demo, args):

    savedir = args.save_dir
    D = os.path.join(savedir, 'orig')

    masking_args = dict(masking_type=args.masking_type, negative_masking=args.negative_masking,
                        inflation_factor=args.inflation_factor, all_at_once=args.all_at_once, masking_thresh=args.masking_thresh)

    out_D = os.path.join(savedir, dict_to_param_str(masking_args))
    os.makedirs(out_D, exist_ok=True)
    pickle.dump(args, open(os.path.join(out_D, 'args.pkl'), 'wb'))

    orig_res_paths = glob(D + '/*.pkl')
    print('!!!!!!!!!!!!!!')
    print(out_D)

    N = min(len(coco_dataset), args.max_images)

    # for i,(img,target,idx) in tqdm(enumerate(coco_dataset),total=N):
    for i in tqdm(range(N)):
        if i == args.max_images:
            break
        idx = i
        k = coco_dataset.id_to_img_map[idx]
        out_path = os.path.join(out_D, '{}.pkl'.format(k))
        if os.path.isfile(out_path):
            continue
        p = os.path.join(D, '{}.pkl'.format(k))
        #cur_res = remove_masks(p)
        cur_res = pickle.load(open(p, 'rb'))
        idx, k_, pred = cur_res['idx'], cur_res['k'], cur_res['preds']
        img, target, idx = coco_dataset[i]
        image = np.array(img)
        preds = detect_masked(image, coco_demo, pred, args, masking_color=128)
        if len(preds) > 0:
            if 'mask' in preds.extra_fields:
                del preds.extra_fields['mask']
        pickle.dump(dict(idx=idx, k=k, preds=preds), open(out_path, 'wb'))

def masked_detection_2(coco_dataset, coco_demo, args):
    savedir = args.save_dir
    D = os.path.join(savedir, 'orig')

    masking_args = dict(masking_type=args.masking_type, negative_masking=args.negative_masking,
                        inflation_factor=args.inflation_factor, all_at_once=args.all_at_once,
                        masking_thresh=args.masking_thresh, inflation_is_abs=args.inflation_is_abs)

    out_D = os.path.join(savedir, dict_to_param_str(masking_args))
    os.makedirs(out_D, exist_ok=True)
    pickle.dump(args, open(os.path.join(out_D, 'args.pkl'), 'wb'))

    print(out_D)

    N = min(len(coco_dataset), args.max_images)
    # for i,(img,target,idx) in tqdm(enumerate(coco_dataset),total=N):
    for i in tqdm(range(N)):
        if i == args.max_images:
            break
        idx = i
        k = coco_dataset.id_to_img_map[idx]
        out_path = os.path.join(out_D, '{}.pkl'.format(k))
        if os.path.isfile(out_path):
            continue
        p = os.path.join(D, '{}.pkl'.format(k))
        cur_res = pickle.load(open(p, 'rb'))

        idx, k_, pred = cur_res['idx'], cur_res['k'], cur_res['preds']
        pred = remove_masks(pred)
        img, target, idx = coco_dataset[i]
        img = np.array(img)
        new_dets_map = re_detect2(img, pred, min_score=0, max_score=1, inflation_factor=args.inflation_factor,
                                  is_abs=args.inflation_is_abs, min_keep_iou=.5)

        pickle.dump(new_dets_map, open(out_path, 'wb'))

def areas_in_mask(preds, M):
    M_int = skimage.transform.integral_image(M[:, :, 0])
    bb = preds.bbox.numpy().astype(int)
    starts = [(b, a) for a, b in bb[:, :2]]
    ends = [(b, a) for a, b in bb[:, 2:]]
    return skimage.transform.integrate(M_int, starts, ends), bb_areas(bb)

def get_label_names(preds):
    return [coco_demo.CATEGORIES[q] for q in top_predictions.extra_fields['labels']]

sns.set_context('paper', font_scale=1.5)

def re_detect(img, preds, min_score=0, max_score=1, inflation_factor=1, is_abs=False):
    scores = preds.extra_fields['scores']
    masking_color = 128
    all_preds = []
    masked_images = []
    target_boxes = []
    # for bb,score in tqdm_notebook(zip(preds.bbox,scores),total=len(preds)):
    for bb, score in zip(preds.bbox, scores):
        if score >= max_score or score <= min_score:
            continue
        cur_mask = img * 0
        mask_from_a_box(cur_mask, bb, inflation_factor, is_abs=is_abs)
        #cur_mask = 1-cur_mask
        masked_img = img * cur_mask + (1 - cur_mask) * masking_color
        preds_masked = coco_demo.compute_prediction(masked_img[:, :, ::-1])
        remove_masks(preds_masked)
        if len(preds_masked) > 0:
            all_preds.append(preds_masked)
        masked_images.append(masked_img)
        target_boxes.append(bb)
    return all_preds, masked_images, target_boxes

config_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda:0"])
cfg.MODEL.ROI_HEADS.DETECTIONS_PER_IMG = 1000
coco_demo = COCODemo(
    cfg,
    min_image_size=800,
    confidence_threshold=0.5)

def make_parser():
    parser = argparse.ArgumentParser(description='Elephant In the Room Fixes')
    parser.add_argument('-dataType', default='val2017')
    parser.add_argument('-dataDir', default='/home/amir/data/coco')
    parser.add_argument('-max_images', type=int, default=10,
                        help='maximal number of images')
    parser.add_argument('-masking_type', default='bbox', choices=['bbox', 'mask'],
                        help='mask using bounding box (or "mask" to use the predicted mask)')
    parser.add_argument('-masking_thresh', type=float, default=0.5,
                        help='minimal score for detection to be used as a mask')
    parser.add_argument("-inflation_is_abs", type=str2bool, default=False,
                        help="use absolute instead of relative inflation for bounding boxes")
    parser.add_argument('-negative_masking', action='store_true',
                        help='mask only the detection (otherwise mask everything but the detection)')

    parser.add_argument('-all_at_once', action='store_false',
                        help='mask all detections then run (otherwise one by one))')

    parser.add_argument('-inflation_factor', type=float, default=1.0,
                        help='how much to change size (multiplicative) of bbox/mask before masking')
    parser.add_argument('-save_dir', default='/home/amir/elephant_fix_results')
    return parser

def main():
    parser = make_parser()
    args = parser.parse_args()
    # loading the mscoco validation set .

    dataType = args.dataType
    dataDir = args.dataDir
    annFile = '{}/annotations/instances_{}.json'.format(dataDir, dataType)
    coco = COCO(annFile)
    cats = coco.loadCats(coco.getCatIds())
    #nms=[cat['name'] for cat in cats]
    imgDir = '{}/images/{}'.format(dataDir, dataType)
    #print('COCO categories: \n{}\n'.format(' '.join(nms)))
    #nms = set([cat['supercategory'] for cat in cats])
    #print('COCO supercategories: \n{}'.format(' '.join(nms)))

    coco_demo = COCODemo(
        cfg,
        min_image_size=800,
        confidence_threshold=0.5)

    coco_dataset = COCODataset(
        annFile, dataDir + '/images/' + args.dataType, True)
    initial_detection(coco_dataset, coco_demo, args,
                      max_images=args.max_images, savedir=args.save_dir)
    print(args)
    #masked_detection( coco_dataset, coco_demo, args)
    masked_detection_2(coco_dataset, coco_demo, args)

    print('all_at_once:', args.all_at_once)

if __name__ == "__main__":
    main()


#!/usr/bin/env python

import os
import sys
import math
import struct
import numpy as np

from npose_util import *

from pyrosetta import *
from pyrosetta.rosetta import *

_xx = (0, 0)
_xy = (0, 1)
_xz = (0, 2)
_yx = (1, 0)
_yy = (1, 1)
_yz = (1, 2)
_zx = (2, 0)
_zy = (2, 1)
_zz = (2, 2)
_x = (0, 3)
_y = (1, 3)
_z = (2, 3)

def xform_to_matrix(xform):
    mat = numeric.xyzMatrix_double_t()
    mat.xx = xform[_xx]
    mat.xy = xform[_xy]
    mat.xz = xform[_xz]
    mat.yx = xform[_yx]
    mat.yy = xform[_yy]
    mat.yz = xform[_yz]
    mat.zx = xform[_zx]
    mat.zy = xform[_zy]
    mat.zz = xform[_zz]

    return mat

def matrix_to_xform(mat):
    xform = np.identity(4)
    xform[_xx] = mat.xx
    xform[_xy] = mat.xy
    xform[_xz] = mat.xz
    xform[_yx] = mat.yx
    xform[_yy] = mat.yy
    xform[_yz] = mat.yz
    xform[_zx] = mat.zx
    xform[_zy] = mat.zy
    xform[_zz] = mat.zz

    return xform

def xform_to_vector(xform):
    vec = numeric.xyzVector_double_t()
    vec.x = xform[_x]
    vec.y = xform[_y]
    vec.z = xform[_z]

    return vec

def vector_to_xform(vec):
    xform = np.identity(4)
    xform[_x] = vec.x
    xform[_y] = vec.y
    xform[_z] = vec.z

    return xform

def to_vector(xyz):
    vec = numeric.xyzVector_double_t(xyz[0], xyz[1], xyz[2])
    return vec

def to_float_vector(xyz):
    vec = numeric.xyzVector_float_t(xyz[0], xyz[1], xyz[2])
    return vec

def from_vector(vec):
    xyz = np.array([0, 0, 0]).astype(float)
    xyz[0] = vec.x
    xyz[1] = vec.y
    xyz[2] = vec.z
    return xyz

# Stub of the last 3 heavy atoms
def get_business_stub(res):
    nheavy = res.nheavyatoms()
    stub = get_stub_from_n_ca_c(from_vector(res.xyz(nheavy)), from_vector(res.xyz(nheavy-1)), from_vector(res.xyz(nheavy-2)))
    return stub

def get_stub_from_res(res):
    return get_stub_from_n_ca_c(from_vector(res.xyz("N")), from_vector(res.xyz("CA")), from_vector(res.xyz("C")))

def get_atoms_and_radii(pose, subset=None, atom_ids_too=False):

    atoms = []
    radii = []
    ids = []

    for seqpos in range(1, pose.size()+1):
        if ( not subset is None ):
            if ( not subset[seqpos] ):
                continue
        res = pose.residue(seqpos)
        for iatom in range(1, res.natoms()+1):
            atoms.append(from_vector(res.xyz(iatom)))
            radii.append(res.atom_type(iatom).lj_radius())
            ids.append((seqpos, iatom))

    if ( atom_ids_too ):
        return np.array(atoms), np.array(radii), np.array(ids)

    return np.array(atoms), np.array(radii)

def get_pose_cas(pose, subset=None):
    cas = []

    for seqpos in range(1, pose.size()+1):
        if ( not subset is None ):
            if ( not subset[seqpos] ):
                continue
        res = pose.residue(seqpos)

        cas.append(from_vector(res.xyz("CA")))

    return np.array(cas)

def get_sasa_surface(pose, probe_size=2.2, resl=0.5):

    atomic_depth = core.scoring.atomic_depth.AtomicDepth(pose, probe_size, False, resl, False)

    vertices = atomic_depth.get_surface_vertices()
    verts = np.zeros((len(vertices), 3))
    vert_normals = np.zeros((len(vertices), 3))

    for i in range(len(vertices)):
        vert = vertices[i]
        verts[i, 0] = vert.x
        verts[i, 1] = vert.y
        verts[i, 2] = vert.z
        vert_normals[i,0] = vert.pn.x
        vert_normals[i,1] = vert.pn.y
        vert_normals[i,2] = vert.pn.z

    faces = atomic_depth.get_surface_faces()
    face_centers = np.zeros((len(faces), 3))
    face_normals = np.zeros((len(faces), 3))
    face_areas = np.zeros(len(faces))

    for i in range(len(faces)):
        face = faces[i]
        face_centers[i] = (verts[face.a] + verts[face.b] + verts[face.c]) / 3
        face_normals[i,0] = face.pn.x
        face_normals[i,1] = face.pn.y
        face_normals[i,2] = face.pn.z
        face_areas[i] = face.area

    return verts, vert_normals, face_centers, face_normals, face_areas

def npose_from_pose(pose, return_npose_to_pose=False, allow_non_protein=False, allow_missing_atoms=False, dtype=np.float32):
    npose_to_pose = []
    for i in range(1, pose.size()+1):
        if ( allow_non_protein or pose.residue(i).is_protein() ):
            npose_to_pose.append(i)

    npose = np.zeros((len(npose_to_pose)*R, 4), dtype)
    npose.fill(np.nan)
    npose[:,3] = 1

    cb_fills = []

    for inpose, seqpos in enumerate(npose_to_pose):
        offset = inpose * R
        res = pose.residue(seqpos)
        for iat, atom_name in enumerate(ATOM_NAMES):
            if ( not res.has( atom_name ) ):

                if ( atom_name == "CB" ):
                    cb_fills.append(inpose)
                    continue
                else:
                    if ( not allow_missing_atoms ):
                        raise ValueError('Pose missing atom %s'%atom_name)

            npose[offset+iat,:3] = from_vector(res.xyz(atom_name))

    if ( len(cb_fills) > 0 ):
        cb_fills = np.array(cb_fills)
        ncac = npose.reshape(-1, R, 4)[cb_fills][:,[N, CA, C], :3]
        tpose = get_stubs_from_n_ca_c(ncac[:,0], ncac[:,1], ncac[:,2])
        cbs = build_CB(tpose)
        offsets = cb_fills*R+CB
        npose[offsets] = cbs

    if ( return_npose_to_pose ):
        return npose, npose_to_pose
    return npose

def pose_from_npose(npose, scratch_pose=None):

    size = nsize(npose)

    if ( scratch_pose is None ):
        scratch_pose = pose_from_sequence("A"*size, "centroid")
    else:
        assert( scratch_pose.size() == size )

    idvec = utility.vector1_core_id_AtomID()
    for i in range(size):
        for j in range(7):
            idvec.append(core.id.AtomID(j+1, i+1))

    vec = utility.vector1_numeric_xyzVector_double_t()

    Hs = build_H(npose)
    by_res = npose.reshape(-1, R, 4)

    for i in range(size):

        vec.append(to_vector(by_res[i,N,:3]))
        vec.append(to_vector(by_res[i,CA,:3]))
        vec.append(to_vector(by_res[i,C,:3]))
        vec.append(to_vector(by_res[i,O,:3]))
        vec.append(to_vector(by_res[i,CB,:3]))
        vec.append(to_vector(by_res[i,CB,:3]))
        vec.append(to_vector(Hs[i]))

    scratch_pose.batch_set_xyz(idvec, vec)

    return scratch_pose


import tensorflow as tf
from parameters import FLAGS
import numpy as np

class network(object):

	############################################################################################################################
	def __init__(self, embeddings):

		with tf.device('/device:GPU:0'):

			self.prediction = []

			# create GRU cells
			with tf.variable_scope("tweet"):
				self.cell_fw = tf.nn.rnn_cell.GRUCell(num_units=FLAGS.rnn_cell_size, activation=tf.sigmoid)
				self.cell_bw = tf.nn.rnn_cell.GRUCell(num_units=FLAGS.rnn_cell_size, activation=tf.sigmoid)

			# RNN placeholders
			self.reg_param = tf.placeholder(tf.float32, shape=[])

			num_of_total_filters = len(FLAGS.filter_sizes.split(",")) * FLAGS.num_filters
			total_tweets = FLAGS.batch_size * FLAGS.tweet_per_user

			# weigths
			self.weights = {'fc1': tf.Variable(tf.random_normal([2 * FLAGS.rnn_cell_size, FLAGS.num_classes]), name="fc1-weights"),
					'fc1-cnn': tf.Variable(tf.random_normal([num_of_total_filters, FLAGS.num_classes]),name="fc1-weights"),
					'att1-w': tf.Variable(tf.random_normal([2 * FLAGS.rnn_cell_size, 2 * FLAGS.rnn_cell_size]), name="att1-weights"),
					'att1-v': tf.Variable(tf.random_normal([2 * FLAGS.rnn_cell_size]), name="att1-vector"),
					'att2-w': tf.Variable(tf.random_normal([2 * FLAGS.rnn_cell_size, 2 * FLAGS.rnn_cell_size]), name="att2-weights"),
					'att2-v': tf.Variable(tf.random_normal([2 * FLAGS.rnn_cell_size]), name="att2-vector"),
					'att2-cnn-w': tf.Variable(tf.random_normal([num_of_total_filters, num_of_total_filters]), name="att2-weights"),
					'att2-cnn-v': tf.Variable(tf.random_normal([num_of_total_filters]), name="att2-vector"),
					}
			# biases
			self.bias = {'fc1': tf.Variable(tf.random_normal([FLAGS.num_classes]), name="fc1-bias-noreg"),
				     'fc1-cnn': tf.Variable(tf.random_normal([FLAGS.num_classes]), name="fc1-bias-noreg"),
				     'att1-w': tf.Variable(tf.random_normal([2 * FLAGS.rnn_cell_size]), name="att1-bias-noreg"),
				     'att2-w': tf.Variable(tf.random_normal([2 * FLAGS.rnn_cell_size]), name="att2-bias-noreg"),
				     'att1-cnn-w': tf.Variable(tf.random_normal([num_of_total_filters]), name="att1-bias-noreg"),
				     'att2-cnn-w': tf.Variable(tf.random_normal([num_of_total_filters]), name="att2-bias-noreg")
						 }

			# initialize the computation graph for the neural network
			# self.rnn()
			#self.rnn_with_attention()
			self.cnn(embeddings.shape[0])
			self.architecture()
			self.backward_pass()

    ############################################################################################################################
	def architecture(self):

		with tf.device('/device:GPU:0'):
			#user level attention
			self.att_context_vector_char = tf.tanh(tf.tensordot(self.cnn_output, self.weights["att2-cnn-w"], axes=1) + self.bias["att2-cnn-w"])
			self.attentions_char = tf.nn.softmax(tf.tensordot(self.att_context_vector_char, self.weights["att2-cnn-v"], axes=1))
			self.attention_output_char = tf.reduce_sum(self.cnn_output * tf.expand_dims(self.attentions_char, -1), 1)

			# FC layer for reducing the dimension to 2(# of classes)
			self.logits = tf.tensordot(self.attention_output_char, self.weights["fc1-cnn"], axes=1) + self.bias["fc1-cnn"]

			# predictions
			self.prediction = tf.nn.softmax(self.logits)

			# calculate accuracy
			self.correct_pred = tf.equal(tf.argmax(self.prediction, 1), tf.argmax(self.input_y, 1))
			self.accuracy = tf.reduce_mean(tf.cast(self.correct_pred, tf.float32))

			return self.prediction

    ############################################################################################################################
	def backward_pass(self):

		with tf.device('/device:GPU:0'):
			# calculate loss
			self.loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.logits, labels=self.input_y))

			# add L2 regularization
			self.l2 = self.reg_param * sum(
				tf.nn.l2_loss(tf_var)
				for tf_var in tf.trainable_variables()
				if not ("noreg" in tf_var.name or "bias" in tf_var.name)
			)
			self.loss += self.l2

			# optimizer
			self.optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate)
			self.train = self.optimizer.minimize(self.loss)

			return self.accuracy, self.loss, self.train

	############################################################################################################################
	def rnn(self):
		# embedding layer
		self.rnn_input = tf.nn.embedding_lookup(self.tf_embeddings, self.X)

		# rnn layer
		(self.outputs, self.output_states) = tf.nn.bidirectional_dynamic_rnn(self.cell_fw, self.cell_bw, self.rnn_input, self.sequence_length, dtype=tf.float32,scope="tweet")

		# concatenate the backward and forward cells
		self.rnn_output_raw = tf.concat([self.output_states[0], self.output_states[1]], 1)

		#reshape the output for the next layers
		self.rnn_output = tf.reshape(self.rnn_output_raw, [FLAGS.batch_size, FLAGS.tweet_per_user, 2*FLAGS.rnn_cell_size])

		return self.rnn_output

    ############################################################################################################################
	def rnn_with_attention(self):
		# embedding layer
		self.rnn_input = tf.nn.embedding_lookup(self.tf_embeddings, self.X)

		# rnn layer
		(self.outputs, self.output_states) = tf.nn.bidirectional_dynamic_rnn(self.cell_fw, self.cell_bw, self.rnn_input, self.sequence_length, dtype=tf.float32,scope="tweet")

		# concatenate the backward and forward cells
		self.concat_outputs = tf.concat(self.outputs, 2)

		# attention layer
		self.att_context_vector = tf.tanh(tf.tensordot(self.concat_outputs, self.weights["att1-w"], axes=1) + self.bias["att1-w"])
		self.attentions = tf.nn.softmax(tf.tensordot(self.att_context_vector, self.weights["att1-v"], axes=1))
		self.attention_output_raw = tf.reduce_sum(self.concat_outputs * tf.expand_dims(self.attentions, -1), 1)

		#reshape the output for the next layers
		self.attention_output = tf.reshape(self.attention_output_raw, [FLAGS.batch_size, FLAGS.tweet_per_user, 2*FLAGS.rnn_cell_size])

		return self.attention_output

	############################################################################################################################
	def captioning(self):
		pass

	############################################################################################################################
	def cnn(self, vocab_size):

		with tf.device('/device:GPU:0'):

			# CNN placeholders
			self.input_x = tf.placeholder(tf.int32, [FLAGS.batch_size*FLAGS.tweet_per_user, FLAGS.sequence_length], name="input_x")
			self.input_y = tf.placeholder(tf.float32, [FLAGS.batch_size, FLAGS.num_classes], name="input_y")

			filter_sizes = [int(size) for size in FLAGS.filter_sizes.split(",")]

			# Embedding layer
			with tf.name_scope("embedding"):
				W = tf.Variable(tf.random_uniform([vocab_size, FLAGS.char_embedding_size], -1.0, 1.0), name="W")
				self.embedded_chars = tf.nn.embedding_lookup(W, self.input_x)
				self.embedded_chars_expanded = tf.expand_dims(self.embedded_chars, -1)

			# Create a convolution + maxpool layer for each filter size
			pooled_outputs = []
			for i, filter_size in enumerate(filter_sizes):
				with tf.name_scope("conv-maxpool-%s" % filter_size):
					# Convolution Layer
					filter_shape = [filter_size, FLAGS.char_embedding_size, 1, FLAGS.num_filters]
					W = tf.Variable(tf.truncated_normal(filter_shape, stddev=0.1), name="W")
					b = tf.Variable(tf.constant(0.1, shape=[FLAGS.num_filters]), name="b-noreg")
					conv = tf.nn.conv2d(
					self.embedded_chars_expanded,
					W,
					strides=[1, 1, 1, 1],
					padding="VALID",
					name="conv")
					# Apply nonlinearity
					h = tf.nn.relu(tf.nn.bias_add(conv, b), name="relu")
					# Maxpooling over the outputs
					pooled = tf.nn.max_pool(
						h,
						ksize=[1, FLAGS.sequence_length - filter_size + 1, 1, 1],
						strides=[1, 1, 1, 1],
						padding='VALID',
						name="pool")
					pooled_outputs.append(pooled)

			# Combine all the pooled features
			num_filters_total = FLAGS.num_filters * len(filter_sizes)
			self.h_pool = tf.concat(pooled_outputs, 3)
			self.h_flat_pool = tf.reshape(self.h_pool, [-1, num_filters_total])

			self.cnn_output = tf.reshape(self.h_flat_pool, [FLAGS.batch_size, FLAGS.tweet_per_user, num_filters_total])

			return self.cnn_output


"""
NeurodataLab LLC 25.12.2019
Created by Andrey Belyaev
"""
import cv2
import gi
import numpy as np
import asyncio
import itertools
import signal
import sys
import traceback
import aiohttp
import requests
from argparse import ArgumentParser
from functools import partial
from typing import Any, Dict

gi.require_version('Gst', '1.0')
gi.require_version('GstRtspServer', '1.0')
from gi.repository import Gst, GstRtspServer, GObject

# Factory class
class SensorFactory(GstRtspServer.RTSPMediaFactory):
    pass

# Server class
class GstServer(GstRtspServer.RTSPServer):
    pass

class SensorFactory(GstRtspServer.RTSPMediaFactory):
    def __init__(self, **properties):
        super(SensorFactory, self).__init__(**properties)  # Init super class
        self.cap = cv2.VideoCapture(2)  # Initialize webcam. You may have to change 0 to your webcam number
        self.frame_number = 0  # Current frame number
        self.fps = 30  # output streaming fps
        self.duration = 1 / self.fps * Gst.SECOND  # duration of a frame in nanoseconds
        self.launch_string = 'appsrc name=source is-live=true block=true format=GST_FORMAT_TIME ' \
                             'caps=video/x-raw,format=BGR,width=1280,height=720,framerate={}/1 ' \
                             '! videoconvert ! video/x-raw,format=I420 ' \
                             '! x264enc speed-preset=ultrafast tune=zerolatency ' \
                             '! rtph264pay config-interval=1 name=pay0 pt=96'.format(self.fps)

        self.christmas_image = cv2.imread('christmas.png', cv2.IMREAD_COLOR)
        self.christmas_image = cv2.resize(self.christmas_image, None, fx=0.1, fy=0.1)
        self.scale_factor = 0.5#TODO: this is the default, provide the ability to change it
        self.bodypix_url = "http://127.0.0.1:9000"#TODO: and here
        background = cv2.imread('background.jpg');#TODO: same here
        self.hologram = True #TODO: and more
        self.width = 1280
        self.height = 720

        #TODO: this came out of loadImages() it should go back there
        self.images: Dict[str, Any] = {}
        if background is not None:
                background = cv2.resize(background, (self.width, self.height))
                background = itertools.repeat(background)

        self.images["background"] = background
        self.session = requests.Session()

    def on_need_data(self, src, lenght):
        if self.cap.isOpened():  # Check webcam is opened
            ret, frame = self.cap.read()  # Read next frame
            frame = self.draw_on_frame(frame)  # Draw something on frame frame
            if ret:  # If read success
                data = frame.tostring()  # Reformat frame to string
                buf = Gst.Buffer.new_allocate(None, len(data), None)  # Allocate memory
                buf.fill(0, data)  # Put new data in memory
                buf.duration = self.duration  # Set data duration
                timestamp = self.frame_number * self.duration  # Current frame timestamp
                buf.pts = buf.dts = int(timestamp)
                buf.offset = timestamp  # Set frame timestamp
                self.frame_number += 1  # Increase current frame number
                retval = src.emit('push-buffer', buf)  # Push allocated memory to source container
                if retval != Gst.FlowReturn.OK:  # Check pushing process
                    print(retval)  # Print error message

    def _get_mask(self, frame, session):
        frame = cv2.resize(frame, (0, 0), fx=self.scale_factor,
                           fy=self.scale_factor)
        _, data = cv2.imencode(".png", frame)
        with session.post(
            url=self.bodypix_url, data=data.tostring(),
            headers={"Content-Type": "application/octet-stream"}
        ) as r:
            mask = np.frombuffer(r.content, dtype=np.uint8)
            mask = mask.reshape((frame.shape[0], frame.shape[1]))
            mask = cv2.resize(
                mask, (0, 0), fx=1 / self.scale_factor,
                fy=1 / self.scale_factor, interpolation=cv2.INTER_NEAREST
            )
            mask = cv2.dilate(mask, np.ones((10, 10), np.uint8), iterations=1)
            mask = cv2.blur(mask.astype(float), (30, 30))
            return mask

    def shift_image(self, img, dx, dy):
        img = np.roll(img, dy, axis=0)
        img = np.roll(img, dx, axis=1)
        if dy > 0:
            img[:dy, :] = 0
        elif dy < 0:
            img[dy:, :] = 0
        if dx > 0:
            img[:, :dx] = 0
        elif dx < 0:
            img[:, dx:] = 0
        return img

    def hologram_effect(self, img):
        # add a blue tint
        holo = cv2.applyColorMap(img, cv2.COLORMAP_WINTER)
        # add a halftone effect
        bandLength, bandGap = 2, 3
        for y in range(holo.shape[0]):
            if y % (bandLength+bandGap) < bandLength:
                holo[y,:,:] = holo[y,:,:] * np.random.uniform(0.1, 0.3)
        # add some ghosting
        holo_blur = cv2.addWeighted(holo, 0.2, self.shift_image(holo.copy(), 5, 5), 0.8, 0)
        holo_blur = cv2.addWeighted(holo_blur, 0.4, self.shift_image(holo.copy(), -5, -5), 0.6, 0)
        # combine with the original color, oversaturated
        out = cv2.addWeighted(img, 0.5, holo_blur, 0.6, 0)
        return out

    def draw_on_frame(self, frame):
        #idxs = np.where(np.logical_and(self.christmas_image > 0, self.christmas_image < 255))
        #for i in range(3):
        #    for j in range(2):
        #        cur_idxs = (idxs[0] + i * self.christmas_image.shape[0],
        #                    idxs[1] + j * (frame.shape[1] - self.christmas_image.shape[1]),
        #                    idxs[2])
        #        frame[cur_idxs] = self.christmas_image[idxs]

        # fetch the mask with retries (the app needs to warmup and we're lazy)
        # e v e n t u a l l y c o n s i s t e n t
        mask = None
        while mask is None:
            try:
                mask = self._get_mask(frame, self.session)
            except Exception as e:
                print(f"Mask request failed, retrying: {e}")
                traceback.print_exc()

        if self.hologram:
            frame = self.hologram_effect(frame)

        # composite the mask and background
        background = next(self.images["background"])
        for c in range(frame.shape[2]):
            frame[:, :, c] = frame[:, :, c] * mask + background[:, :, c] * (1 - mask)

        return frame

    def do_create_element(self, url):
        return Gst.parse_launch(self.launch_string)  # Launch gst plugin

    def do_configure(self, rtsp_media):
        self.frame_number = 0  # Set current frame number to zero
        appsrc = rtsp_media.get_element().get_child_by_name('source')  # get source from gstreamer
        appsrc.connect('need-data', self.on_need_data)  # set data provider

class GstServer(GstRtspServer.RTSPServer):
    def __init__(self, **properties):
        super(GstServer, self).__init__(**properties)  # Init super class
        self.factory = SensorFactory()  # Create factory
        self.set_service("3002")  # Set service port
        self.factory.set_shared(True)  # Set shared to true
        self.get_mount_points().add_factory("/test", self.factory)  # Add routing to access factory
        self.attach(None)

if __name__ == '__main__':
    loop = GObject.MainLoop()  # Create infinite loop for gstreamer server
    GObject.threads_init()  # Initialize server threads for asynchronous requests
    Gst.init(None)  # Initialize GStreamer

    server = GstServer()  # Initialize server
    loop.run()  # Start infinite loop


from keras.models import load_model
import cv2
import numpy as np
from random import choice
#import Speech_Recognition
#from qibullet import SimulationManager
#from qibullet import PepperVirtual
#from qibullet import Camera

#model hasnt been trained with data for no user hand
REV_CLASS_MAP = {
		0: "rock",
		1: "paper",
		2: "scissors"
	}

#gives each value of r/p/s a value of 0-3
def mapper(val):
    return REV_CLASS_MAP[val]

#calculates winner of game, returns who won
def calculate_winner(move1, move2):
    if move1 == move2:
        return "Tie"

    if move1 == "rock":
        if move2 == "scissors":
            return "User"
        if move2 == "paper":
            return "Computer"

    if move1 == "paper":
        if move2 == "rock":
            return "User"
        if move2 == "scissors":
            return "Computer"

    if move1 == "scissors":
        if move2 == "paper":
            return "User"
        if move2 == "rock":
            return "Computer"

def rock_paper_scissors():

	#voiceInp=Speech_Recognition.voiceRecog()
	#voiceInp.voice()
	#loads trained model
	model = load_model("rock-paper-scissors-model.h5")

	#loads videocapture device
	cap = cv2.VideoCapture(0)

	prev_move = None

	#beings rps game with user
	while True:
		ret, frame = cap.read()
		if not ret:
			continue

		# area which user will play within
		cv2.rectangle(frame, (100, 100), (500, 500), (66, 182, 245), 2)
		# area which user will play within
		cv2.rectangle(frame, (800, 100), (1200, 500), (255, 0, 0), 2)

		# extract the region of image within the user play area
		roi = frame[100:500, 100:500]
		img = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)#converts from brg to rgb so that model is able to compare between frames of video feed and computers choice
		img = cv2.resize(img, (150, 150))#changes the frame from 400x400 to 150x150 - matches model image sizes

		# predict the move made
		pred = model.predict(np.array([img]))#uses the tensorflow model to guess which out of r/p/s user chose
		move_code = np.argmax(pred[0])#assigns prediction from model which is saved as either 0,1,2 depending on REV_CLASS_MAP
		user_move_name = mapper(move_code)#changes 0,1,2 to its corresponding r/p/s value

		# Choose winner of game
		if prev_move != user_move_name:#only makes computer randomly choose a new r/p/s after the computer recognises that a new sign is shown
			computer_move_name = choice(['rock', 'paper', 'scissors'])#selects r/p/s randomly from array and assigns it as the users choice
			winner = calculate_winner(user_move_name, computer_move_name)

		prev_move = user_move_name

		# display the information
		cv2.putText(frame, "Your Move: " + user_move_name,
					(50, 50), 3, 1.2, (66, 182, 245), 2, cv2.LINE_AA)
		cv2.putText(frame, "Computer's Move: " + computer_move_name,
					(750, 50), 3, 1.2, (255, 0, 0), 2, cv2.LINE_AA)
		cv2.putText(frame, "Winner: " + winner,
					(400, 600), 3, 2, (0, 0, 255), 4, cv2.LINE_AA)

		if computer_move_name != "none":
			icon = cv2.imread(
				"images/{}.png".format(computer_move_name))
			icon = cv2.resize(icon, (400, 400))
			frame[100:500, 800:1200] = icon

		cv2.imshow("Rock Paper Scissors", frame)

		k = cv2.waitKey(10)
		if k == ord('q'):#quits program if user presses q
			break
		#if voiceInp.voiceWord == "exit":
		#print("Thank you for playing rock paper scissors")
		#playing=False

	cap.release()
	cv2.destroyAllWindows()