"""Тесты для загрузки локальных данных по дивидендам."""
import pandas as pd
import pytest

from poptimizer.data.adapters.loaders import dividends
from poptimizer.data.ports import col, outer

@pytest.mark.asyncio
async def test_loader_raise_on_wrong_name():
    """Не верное название таблицы."""
    table_name = outer.TableName(outer.CPI, "AKRN")

    with pytest.raises(outer.DataError, match="Некорректное имя таблицы для обновления"):
        loader = dividends.DividendsLoader()
        await loader.get(table_name)

@pytest.mark.asyncio
async def test_loader(mocker):
    """Форматирование загруженного DataFrame."""
    fake_client = mocker.patch.object(dividends.resources, "get_mongo_client")
    fake_db = fake_client.return_value.__getitem__.return_value  # noqa: WPS609
    fake_collection = fake_db.__getitem__.return_value  # noqa: WPS609
    fake_collection.aggregate.return_value = mocker.AsyncMock()
    mocker.patch.object(dividends.pd, "DataFrame", return_value=pd.DataFrame(columns=[1, 2]))

    loader = dividends.DividendsLoader()
    df = await loader.get(outer.TableName(outer.DIVIDENDS, "AKRN"))

    assert df.index.name == col.DATE
    assert df.columns == ["AKRN"]


import torch
import torch.nn as nn
import torch.nn.functional as F

class DualPolyEncoder(nn.Module):
    def __init__(self, bert):
        super().__init__()
        self.bert = bert
        self.codes_number = 16
        self.hidden_size = 768
        self.poly_codes_for_context = nn.Embedding(self.codes_number,
                                                   self.hidden_size)
        self.poly_codes_for_candidate = nn.Embedding(self.codes_number,
                                                     self.hidden_size)

        self.init_weights()

    def init_weights(self):
        initrange = 0.1
        self.poly_codes_for_context.weight.data.uniform_(-initrange, initrange)
        self.poly_codes_for_candidate.weight.data.uniform_(
            -initrange, initrange)

    def dot_attention(self, q, k, v):
        attention_weights = torch.bmm(q, k.transpose(1, 2))
        scaled_attention_weights = F.softmax(attention_weights, dim=0)
        output = torch.bmm(scaled_attention_weights, v)
        return output

    def mean_pooling(self, last_hidden, attention_mask):
        pass

    def forward(self,
                context_input_ids,
                context_attention_mask,
                candidate_input_ids,
                candidate_attention_mask,
                labels=None):
        batch_size, _ = context_input_ids.shape
        device = context_input_ids.device

        context_output = self.bert(context_input_ids, context_attention_mask)
        context_hidden_state = context_output[0]
        context_cls = context_output[1][:, None, :]
        candidate_output = self.bert(candidate_input_ids, candidate_attention_mask)
        candidate_hidden_state = candidate_output[0]
        candidate_cls = candidate_output[1][:, None, :]
        poly_code_ids = torch.arange(self.codes_number, dtype=torch.int).unsqueeze(0).expand(batch_size, self.codes_number).to(device)

        poly_codes_for_context = self.poly_codes_for_context(poly_code_ids)
        # poly_codes_for_context.shape -> [16, 16, 768]
        # context_hidden_state.shape -> [16, 500, 768]
        # candidate_cls.shape -> [16, 1, 768]

        ctxt_poly_embed = self.dot_attention(poly_codes_for_context, context_hidden_state, context_hidden_state)
        ctxt_embed = self.dot_attention(candidate_cls, ctxt_poly_embed, ctxt_poly_embed)
        negative_score = torch.bmm(candidate_cls, ctxt_embed.transpose(1, 2)).flatten()

        poly_codes_for_candidate = self.poly_codes_for_candidate(poly_code_ids)
        cand_poly_embed = self.dot_attention(poly_codes_for_candidate, candidate_hidden_state, candidate_hidden_state)
        cand_embed = self.dot_attention(context_cls, cand_poly_embed, cand_poly_embed)
        positive_score = torch.bmm(context_cls, cand_embed.transpose(1, 2)).flatten()

        final_score = positive_score - negative_score

        if labels is not None:
            criterion = nn.BCEWithLogitsLoss()
            loss = criterion(final_score, labels)
            return loss
        else:
            return final_score


#==== Imports ====#
import unittest
import numpy as np
import pandas as pd
from scipy.io import arff
from sklearn.model_selection import train_test_split as ts
from sklearn import metrics
from sklearn.tree import DecisionTreeClassifier
#=================#

#==== Functions ====#
def decision_tree_model(data,args=None):
    dt = DecisionTreeClassifier()
    dt.fit(data[0],data[1])
    return dt

def dummy_data(filename = 'CM1.arff.txt' ):
    #File
    data = arff.loadarff(filename)
    loaddata = pd.DataFrame(data[0])

    #Software metrics and decisions
    SM = np.array(loaddata[['LOC_BLANK','BRANCH_COUNT','CALL_PAIRS','LOC_CODE_AND_COMMENT']]) #Software metrics
    L = np.array(loaddata['Defective']) #label
    SM2 = np.array(loaddata[['LOC_BLANK','BRANCH_COUNT','CALL_PAIRS']])
    SM_train, SM_test, L_train, L_test = ts(SM,L,test_size = 0.1)
    SM2_train, SM2_test, L2_train, L2_test = ts(SM2,L,test_size = 0.1)
    L_train=L_train.astype('str')
    L_test=L_test.astype('str')
    return [SM_train,L_train], SM2_test, SM_test,L_test
#===================#

class dt_test(unittest.TestCase):

    def test_type(self):
        dummy = dummy_data()
        dt = decision_tree_model(dummy[0])
        dt2 = DecisionTreeClassifier()
        self.assertTrue(type(dt) == type(dt2))

    def test_empty(self):
        try:
            dt = decision_tree_model([])
        except IndexError:
            fail = True
        self.assertTrue(fail)

    def test_fit(self):
        dummy = dummy_data()
        lr = decision_tree_model(dummy[0])
        try:
            lr.predict(dummy[1])
        except ValueError:
            fail = True
        self.assertTrue(fail)

    def test_normal(self):
        dummy = dummy_data()
        dt = decision_tree_model(dummy[0])
        prediction = dt.predict(dummy[2])
        res = metrics.accuracy_score(dummy[3], prediction)
        print(res)
        self.assertTrue(res >= 0.75)

#====================( Main )====================
if __name__=='__main__':
    unittest.main()


import os
from datetime import datetime
from urllib.parse import urljoin

import pyarrow.hdfs
from django.conf import settings
from django.core.files import File
from django.core.files.storage import Storage
from django.core.signals import setting_changed
from django.utils import timezone
from django.utils._os import safe_join
from django.utils.deconstruct import deconstructible
from django.utils.encoding import filepath_to_uri
from django.utils.functional import cached_property
from django.utils.timezone import now
from unidecode import unidecode

def clean_char(char, replace=''):
    """
    Remove ['.',',','@'] and accents to char
    :param char: char
    :return: clean char
    """
    bad_tokens = [",", ";", ".",
                  "!", "'", ".",
                  "-", '"', "@",
                  r"\n", r"\r", '?',
                  '_']
    for delete in bad_tokens:
        char = char.replace(delete, replace)
    return char

@deconstructible
class HadoopStorage(Storage):
    hadoop_host: str = getattr(settings, 'HADOOP_HOST', 'localhost')
    hadoop_port: int = getattr(settings, 'HADOOP_PORT', 8020)
    hadoop_user: str = getattr(settings, 'HADOOP_USER', 'hadoop')
    hadoop_home: str = getattr(settings, 'HADOOP_HOME', '/usr/lib/hadoop/')
    replications: int = 3

    def __init__(self, location=None, base_url=None, file_permissions_mode=None,
                 directory_permissions_mode=None):
        self._location = location
        self._base_url = base_url
        self._file_permissions_mode = file_permissions_mode
        self._directory_permissions_mode = directory_permissions_mode

        # e['ARROW_LIBHDFS_DIR'] = '/opt/hadoop/hadoop/lib/native/'
        os.environ['HADOOP_HOME'] = self.hadoop_home
        # e['JAVA_HOME'] = "/usr/lib/jvm/java-11-openjdk-amd64"
        os.environ['CLASSPATH'] = f'{self.hadoop_home}/bin/hdfs classpath --glob'

        self.hdfs = pyarrow.hdfs.connect(
            self.hadoop_host, self.hadoop_port, user=self.hadoop_user,
            extra_conf={'dfs.replication': str(self.replications)})
        setting_changed.connect(self._clear_cached_properties)

    def _clear_cached_properties(self, setting, **kwargs):
        """Reset setting based property values."""
        if setting == 'MEDIA_ROOT':
            self.__dict__.pop('base_location', None)
            self.__dict__.pop('location', None)
        elif setting == 'MEDIA_URL':
            self.__dict__.pop('base_url', None)
        elif setting == 'FILE_UPLOAD_PERMISSIONS':
            self.__dict__.pop('file_permissions_mode', None)
        elif setting == 'FILE_UPLOAD_DIRECTORY_PERMISSIONS':
            self.__dict__.pop('directory_permissions_mode', None)

    def _value_or_setting(self, value, setting):
        return setting if value is None else value

    @property
    def path_prefix(self):
        """
        Set a prefix for all directories, usefull for pull apart directories by user session
        """
        return ''

    @cached_property
    def base_location(self):
        return self._value_or_setting(self._location, settings.MEDIA_ROOT)

    @cached_property
    def location(self):
        return os.path.abspath(self.base_location)

    @cached_property
    def base_url(self):
        if self._base_url is not None and not self._base_url.endswith('/'):
            self._base_url += '/'
        return self._value_or_setting(self._base_url, settings.MEDIA_URL)

    @cached_property
    def file_permissions_mode(self):
        return self._value_or_setting(self._file_permissions_mode, settings.FILE_UPLOAD_PERMISSIONS)

    @cached_property
    def directory_permissions_mode(self):
        return self._value_or_setting(self._directory_permissions_mode, settings.FILE_UPLOAD_DIRECTORY_PERMISSIONS)

    def _open(self, name, mode='rb'):
        return File(self.hdfs.open(self.path(name), mode))

    def _save(self, name, content):
        full_path = self.path(name)
        hdfs = self.hdfs
        # Create any intermediate directories that do not exist.
        directory = os.path.dirname(full_path)
        if not hdfs.exists(directory):
            hdfs.mkdir(directory)

        # There's a potential race condition between get_available_name and
        # saving the file; it's possible that two threads might return the
        # same name, at which point all sorts of fun happens. So we need to
        # try to create the file, but if it already exists we have to go back
        # to get_available_name() and try again.
        hdfs.upload(full_path, content)

        # Store filenames with forward slashes, even on Windows.
        return name.replace('\\', '/')

    def delete(self, name):
        assert name, "The name argument is not allowed to be empty."
        name = self.path(name)
        # If the file or directory exists, delete it from the filesystem.
        if self.hdfs.exists(name):
            self.hdfs.delete(name)

    def get_available_name(self, name, max_length=100):
        # set a timestamp because hadoop override for same names
        dir_name, file_name = os.path.split(name)

        file_root, file_ext = os.path.splitext(file_name)
        file_root = unidecode(clean_char(file_root, '_'))

        # Validate if the name with the additional datetime ('%d%m%Y_%H%M%S%f') and the id_user_isis
        #  has more than 100 characters

        if max_length and (len(name) + 26) >= max_length:
            # Truncate file_root
            aux = len(name) - len(file_root)
            truncation = 74 - aux
            file_root = file_root[:truncation]

        file_name = "%s_%s%s" % (now().strftime('%d%m%Y_%H%M%S%f'), file_root, file_ext)
        return os.path.join(self.path_prefix, dir_name, file_name)

    def exists(self, name):
        return self.hdfs.exists(self.path(name))

    def listdir(self, path):
        path = self.path(path)

        directories, files = [], []
        for entry in self.hdfs.ls(path):
            if self.hdfs.info(entry)['kind'] == 'directory':
                directories.append(entry)
            elif self.hdfs.info(entry)['kind'] == 'file':
                files.append(entry)
        return directories, files

    def path(self, name):
        return safe_join(self.location, name)

    def size(self, name):
        return self.hdfs.info(self.path(name))['size']

    def url(self, name):
        if self.base_url is None:
            raise ValueError("This file is not accessible via a URL.")
        url = filepath_to_uri(name)
        if url is not None:
            url = url.lstrip('/')
        return urljoin(self.base_url, url)

    def _datetime_from_timestamp(self, ts):
        """
        If timezone support is enabled, make an aware datetime object in UTC;
        otherwise make a naive one in the local timezone.
        """
        if settings.USE_TZ:
            # Safe to use .replace() because UTC doesn't have DST
            return datetime.utcfromtimestamp(ts).replace(tzinfo=timezone.utc)
        else:
            return datetime.fromtimestamp(ts)

    def get_accessed_time(self, name):
        return self._datetime_from_timestamp(self.hdfs.info(self.path(name))['last_accessed'])

    def get_created_time(self, name):
        # not exists created time on hadoop
        return self._datetime_from_timestamp(self.hdfs.info(self.path(name))['last_modified'])

    def get_modified_time(self, name):
        return self._datetime_from_timestamp(self.hdfs.info(self.path(name))['last_modified'])


import os,sys
import vs_pfm.docking.pybel_rmsd as pybel_rmsd
import vs_pfm.docking.arbalign_rmsd as arbalign_rmsd
import pybel
import vs_pfm.general_functions as gf
import vs_pfm.path_const as pc
import subprocess
import tempfile
import vs_pfm.docking.add_charge_knowledge_based as charge
import rdkit.Chem as Chem
import pandas as pd
import pdb2pqr

def LigProtonation(input_name, input_format,output_format,output_name=None, method='pybel',ph=7):
    '''
    protonate ligands. This function can be used in both DockPrep and Conformer generation.
    '''

    if method == 'pybel' or method == None:
        print('Calculating protonation state for the ligand using pybel')
        if input_format == 'smi':
            mol = pybel.readstring('smi',input_name)
        elif not input_format.endswith('_block'):
            mol = next(pybel.readfile(input_format, input_name))
        else:
            real_format = input_format.split('_')[0]
            mol = pybel.readstring(real_format,input_name)

        mol.OBMol.AddHydrogens( False, True, ph) # (polaronly=False, correctforph=True, PH=5)
        if output_format != 'smi' and output_format != 'rdkit_mol':
            if not output_name:
                return mol.write(output_format)
            else:
                mol.write(output_format, output_name, overwrite=True)

        elif output_format == 'smi':
            return mol.write(output_format).split('\t')[0]
        else:
            sdf_block = mol.write('sdf')
            rdkit_mol = Chem.MolFromMolBlock(sdf_block,removeHs=False)
            if not rdkit_mol:
                mol2_block = mol.write('mol2')
                rdkit_mol = Chem.MolFromMol2Block(mol2_block,removeHs=False)
                if not rdkit_mol:
                    smi_h = mol.write('smi').split('\t')[0]
                    rdkit_mol = Chem.MolFromSmiles(smi_h)
                    rdkit_mol = Chem.AddHs(rdkit_mol,addCoords=True)
            return rdkit_mol

     # add h using knowledge based method (specifically developed for gc3 and gc4 compounds)
    elif method =='knowledge':

        if not input_format.endswith('_block'):
            if input_format == 'sdf':
                lig_mol = Chem.MolFromMolFile(input_name)

            elif input_format == 'smi':
                lig_mol = Chem.MolFromSmiles(input_name)
            else:
                lig_sdf= gf.format_convert_to_temp(input_name,input_format,'sdf')
                lig_mol = Chem.MolFromMolFile(lig_sdf)

        else:
            real_format = input_format.split('_')[0]
            if real_format == 'sdf':
                lig_mol = Chem.MolFromMolBlock(input_name)
            elif real_format == 'pdb':
                lig_mol = Chem.MolFromPDBBlock(input_name)
            else:
                raise ValueError(input_format + ' is not a supported format block.')

        lig_mol_charged = charge.add_charge(lig_mol)
        lig_mol_charged_h = Chem.AddHs(lig_mol_charged,addCoords=True)
        if output_format == 'pdb':
            Chem.MolToPDBFile(lig_mol_charged_h,output_name)
        elif output_format == 'sdf':
            Chem.MolToMolFile(lig_mol_charged_h,output_name)
        elif output_format == 'smi':
            return Chem.MolToSmiles(lig_mol_charged_h)
        elif output_format == 'rdkit_mol':
            return lig_mol_charged_h

        else:
            raise ValueError('ligand output format not supported')
        # only remove lig_sdf if its a temp file
        if not (input_format == 'sdf' or input_format == 'smi' or input_format == 'pdb' or input_format.endswith('_block')):
            os.system('rm ' + lig_sdf)
    else:
        raise ValueError(method + ' not supported for ligand protonation state calculation')

class DockPrep(object):
    '''
    prepare receptor and ligand for vina
    1. add h
    2. convert to pdbqt
    3. get config file for receptor structure
    4. generate conformation of ligand

    '''

    def __init__(self,input_file=None,input_type=None):
        '''
        :param input_file: structure input file (input_file should has suffix to indicates its format)
        :param input_type: ligand or receptor
        '''
        self.input_format = input_file.split('.')[-1]

        self.input_file = input_file # regardless of structure input format, the file info is saved as self.input_file
        self.input_type = input_type

    def add_h(self, output_prefix, output_format='pdb', method=None,ph=7):
        '''
        add h to receptor/ligand
        :param output_prefix: output file prefix
        :param output_format: output file format
        :param method: method for calculating protonation state.
                    Default:
                    protein--pdb2pqr no other options
                    ligand--pybel
                            other options: knowledge (functional group based protonation method only developed for cats system )
        :return:
        '''

        if self.input_type == 'receptor':
            # add h to file and store as pqr file
            if method == None or method == 'pdb2pqr':
                pdb2pqr = gf.get_tool('pdb2pqr_cli')
                if not self.input_format == 'pdb':
                    temp_pdb_file = gf.format_convert_to_temp(self.input_file,self.input_format,'pdb')
                    pdb_file = temp_pdb_file
                else:
                    pdb_file = self.input_file
                pqr_file = output_prefix + '.pqr'
                print('Calculating charge for the receptor using pdb2pqr')

                subprocess.Popen([pdb2pqr,'--ff=amber','--ph-calc-method=propka31','--with-ph=' + str(ph),pdb_file, pqr_file]\
                    ).communicate()
                try:
                    os.system('rm -f ' + temp_pdb_file)
                except:
                    pass

                # delete last two column in pqr file since pythonsh receptor4.py cannot recognize them
                lines = open(output_prefix + '.pqr','r').readlines()
                new_lines = [l[0:54] + '\n' for l in lines]
                output_file = open(output_prefix + '.pdb','w')
                output_file.writelines(new_lines)
                output_file.close()

                if not output_format == 'pdb':
                    gf.convert_format(output_prefix+'.pdb','pdb',output_prefix + '.' + output_format, output_format)
                    os.system('rm -f ' + output_prefix + '.pdb')

            else:
                raise ValueError(method + ' is not supported to calculate protein protonation state!')

        elif self.input_type == 'ligand':

            output_file = output_prefix + '.' + output_format
            LigProtonation(self.input_file,self.input_format,output_format,output_file,method=method,ph=ph)

        else:
            raise ValueError('wrong input type for add_h')

        self.pdb_h = output_prefix + '.' + output_format

    def prepare_pdbqt(self, input,output, flex=True):
        '''
        convert receptor/ligand to pdbqt file, should have all h added
        :param output: output file name
        :param flex: if ligand is considered as flexible (only used for ligand pdbqt generation)
        :return:
        '''
        prepare_receptor = gf.get_tool('prepare_receptor4.py')
        prepare_ligand = gf.get_tool('prepare_ligand4.py')

        input_format = input.split('.')[-1]
        if self.input_type == 'receptor':
            print('Generating pdbqt for the receptor')
            cmd = prepare_receptor + ' -r ' + input + ' -o ' + output #+ ' -U lps'
            os.system(cmd)
        elif self.input_type == 'ligand':
            print('Generating pdbqt for the ligand')
            pdbqt_dir = os.path.dirname (os.path.realpath (output))
            if input_format != 'pdb':
                cmd = 'obabel -i' + input_format + ' ' + input + ' -opdb > ' + pdbqt_dir + '/temp.pdb'
                os.system(cmd)
                input = pdbqt_dir + '/temp.pdb'

            if flex == True:
                cmd = prepare_ligand + ' -l ' + input + ' -o ' + output
            else:
                cmd = prepare_ligand + ' -l ' + input + ' -o ' + output  + ' -Z'
            os.system(cmd)
            if os.path.exists(pdbqt_dir + '/temp.pdb'):
                cmd = 'rm ' + pdbqt_dir + '/temp.pdb'
                os.system(cmd)

class DockRun(object):
    '''
    1. prepare receptor and ligand for vina
    2. get config file for vina
    3. run docking use smina
    4. analyze vina results
    '''

    def __init__(self, receptor_pdbqt, ligand_pdbqt, output_prefix, config=None):

        self.receptor_pdbqt = receptor_pdbqt
        self.ligand_pdbqt = ligand_pdbqt
        if config:
            self.config = config
            self.remove_config = False
        else:
            self.config = tempfile.NamedTemporaryFile(delete=False).name
            self.get_config()
            self.remove_config = True
        self.output_prefix = output_prefix

    def get_config(self):
            '''
            prepare vina config file using receptor atom positions
            :param config_file: config file name
            :return:
            '''
            print('Define searching space using entire receptor')
            lines = [l for l in open(self.receptor_pdbqt,'r').readlines() if l.startswith('ATOM') or l.startswith('HETATM')]

            x = [float(l[30:38]) for l in lines]
            y = [float(l[38:46]) for l in lines]
            z = [float(l[46:54]) for l in lines]
            center_x = (max(x) + min(x))/2
            center_y = (max(y) + min(y))/2
            center_z = (max(z) + min(z))/2
            size_x = abs(max(x)-center_x) + 10
            size_y = abs(max(y)-center_y) + 10
            size_z = abs(max(z)-center_z) + 10
            config = open(self.config, 'w')
            config.write('center_x = %.2f \n' % center_x)
            config.write('center_y = %.2f \n' % center_y)
            config.write('center_z = %.2f \n' % center_z)
            config.write('size_x = %.2f \n' % size_x)
            config.write('size_y = %.2f \n' % size_y)
            config.write('size_z = %.2f \n' % size_z)
            config.close()

    def run_smina(self, option=[]):
        '''
        run smina for self.receptor_pdbqt and self.ligand_pdbqt
        :param option: a list of arguments for smina or vina
        :return:
        '''
        print('Run Smina docking')
        smina = gf.get_tool('smina')
        out_pdb = self.output_prefix + '.pdb'
        smina_env = os.environ.copy()
        smina_env['OMP_NUM_THREADS'] = '1'
        option_list = [smina, '-r', self.receptor_pdbqt, '-l', self.ligand_pdbqt, '-o',out_pdb,'--seed','10', '--config', self.config] + option # OMP_NUM_THREADS is to limit the number of threads smina can use to avoid killing nodes on NYU hpc
        subprocess.Popen(option_list,env=smina_env).communicate()

        if self.remove_config:
            os.system('rm -f ' + self.config)

class DockAnlys(object):

    def __init__(self, dock_file):
        '''
        get docking score and rmsds from smina output and rmsd calculators
        dock_file: absolute path to smina docking output file
        '''
        num_pose = self.check_num_pose(dock_file)
        self.score = [] * num_pose
        self.rmsd = [] * num_pose
        self.dock_file = dock_file
        self.dock_format = dock_file.split('.')[-1]
        self.dock_folder = '/'.join(dock_file.split('/')[0:-1])
        if not (self.dock_format == 'pdb' or self.dock_format == 'pdbqt'):
            raise ValueError('dock pose format not supported')
        self.prefix = dock_file.split('/')[-1].split('.')[0]

    @staticmethod
    def check_num_pose(dock_file):
        num_pose = len([l for l in open(dock_file,'r').readlines() if l.startswith('MODEL ')])
        return num_pose

    @staticmethod
    def separate_pdb(dock_file):
        lines = open(dock_file,'r').readlines()
        all_models = []
        ptr = 0

        while ptr < len(lines):
            line = lines[ptr]
            if line.startswith('MODEL '):
                model_lines = []
                model_lines.append(line)
            elif line.startswith('ENDMDL'):
                all_models.append(model_lines)

            else:
                model_lines.append(line)
            ptr += 1
        return all_models

    def get_rmsd(self,ref_file,option='no match',ref_match_sdf=None):
        '''
        Calculate rmsd between reference pdb and dock poses (in pdb format )
        param: option:
                if option == 'match': use arbalign to macth order and then use pybel to calculate rmsd
                elif option == 'no match': use pybel to calculate rmsd
        param ref_match_sdf:
                if option == 'match', ref_match_sdf is the ref sdf with matched order
        '''

        ref_format = ref_file.split('.')[-1]
        ref_temp = False
        if option == 'no match':
            if not ref_format == 'sdf':
                ref = gf.format_convert_to_temp(ref_file,ref_format,'sdf')
                ref_temp = True # label if ref file is temp file for future removal of temp files
            else:
                ref = ref_file
            dock_sdf= gf.format_convert_to_temp(self.dock_file, self.dock_format, 'sdf')
            rmsds = pybel_rmsd.cal_RMSD_sym(dock_sdf,ref)
        else:
            if not ref_format == 'pdb':
                ref = gf.format_convert_to_temp(ref_file,ref_format, 'pdb')
                ref_temp = True
            else:
                ref = ref_file
            all_models = self.separate_pdb(self.dock_file)
            rmsds = []
            for ind, model in enumerate(all_models):
                with open(tempfile.NamedTemporaryFile(delete=False).name,'w') as temp:
                    temp.writelines(model)
                    temp.flush()
                    temp.close()
                    if ind == 0:
                        obj_align = arbalign_rmsd.Arbalign(temp.name,ref,self.dock_format,'pdb',output_dir=self.dock_folder + '/temp')
                        obj_align.arbalign(del_file=False)
                        matchObj = arbalign_rmsd.MatchOrder('c-temp.xyz', 'c-temp-aligned_to-c-ref-noHydrogens.xyz',
                                  file_dir=self.dock_folder +'/temp')
                        matchObj.match(ref_match_sdf)

                    temp_sdf = gf.format_convert_to_temp(temp.name,'pdb','sdf')

                    rmsd = pybel_rmsd.cal_RMSD_sym(temp_sdf,ref_match_sdf)[0]
                    rmsds.append(rmsd)
                    os.system('rm -f ' + temp.name)
                    os.system('rm -f ' + temp_sdf)

            os.system('rm -fr ' + self.dock_folder + '/temp')

        if ref_temp == True:
            os.system('rm -f '+ ref)

        try:
            os.system('rm -f '+ dock_sdf)
        except:
            pass
        self.rmsd = rmsds

    def get_score(self):
        if self.dock_format == 'pdbqt': # docking by vina
            scores = [float(l.split()[3]) for l in open(self.dock_file,'r').readlines() if l.startswith('REMARK VINA RESULT')]
        elif self.dock_format == 'pdb': ## docking by smina
            scores = [float(l.split()[2]) for l in open(self.dock_file,'r').readlines() if l.startswith('REMARK minimizedAffinity')]
        self.score = scores

    def write_results(self,result_csv):
        df = pd.DataFrame()
        df['rmsd'] = self.rmsd
        df['score'] = self.score
        df.to_csv(result_csv,index_label='pose')


#! /usr/bin/env python

"""This module contains functions related to photometry that may be generally
useful to analysis of multiple CARs
"""
from astropy.stats import sigma_clipped_stats
from astropy.visualization import SqrtStretch
from astropy.visualization.mpl_normalize import ImageNormalize
import matplotlib.pyplot as plt
import numpy as np
from photutils import CircularAnnulus, CircularAperture, DAOStarFinder, aperture_photometry

def do_photometry(image, source_table, aperture_radius=3, subtract_background=False, annulus_radii=(6, 8)):
    """Perform aperture photometry on the input data using the source
    locations specified in the source_table

    Parameters
    ----------
    image : numpy.ndarray
        2D image

    source_table : astropy.table.Table
        Table of source locations (i.e. output from photutils's
        DAOStarFinder)

    subtract_background : bool
        If True, use photutils to estimate and subtract background

    Returns
    -------
    source_table : astropy.table.Table
        Updated source_table including photometry results
    """
    # Create list of tuples giving source locations
    positions = [(tab['xcentroid'], tab['ycentroid']) for tab in source_table]

    apertures = CircularAperture(positions, r=aperture_radius)

    # If background is to be subtracted, calculate the sigma-clipped median
    # value of the pixels in the annuli around the sources
    if subtract_background:
        annulus_apertures = CircularAnnulus(positions, r_in=annulus_radii[0], r_out=annulus_radii[1])
        annulus_masks = annulus_apertures.to_mask(method='center')

        bkg_median = []
        for mask in annulus_masks:
            annulus_data = mask.multiply(image)
            annulus_data_1d = annulus_data[mask.data > 0]
            _, median_sigclip, _ = sigma_clipped_stats(annulus_data_1d)
            bkg_median.append(median_sigclip)
        bkg_median = np.array(bkg_median)

    # Do the photometry and add to the input table
    phot_table = aperture_photometry(image, apertures)
    source_table['aperture_sum'] = phot_table['aperture_sum']

    if subtract_background:
        # Add columns with that background subtraction info
        source_table['annulus_median'] = bkg_median
        source_table['aper_bkg'] = bkg_median * apertures.area
        source_table['aper_sum_bkgsub'] = source_table['aperture_sum'] - source_table['aper_bkg']
    return source_table

def find_sources(data, threshold=30, fwhm=3.0, show_sources=True, save_sources=False, plot_name='sources.png'):
    """Use photutils' DAOFind to locate sources in the input image.

    Parameters
    ----------
    data : numpy.ndarray
        2D image

    threshold : float
        Number of sigma above the background used to determine the
        presence of a source

    show_sources : bool
        If True, create an image with overlain marks showing sources

    save_sources : bool
        If True, save the image of the source locations at ```plot_name```

    plot_name : str
        Name of file to save the image with marked sources

    Returns
    -------
    sources : astropy.table.Table
        Table of source positions
    """
    mean, median, std = sigma_clipped_stats(data, sigma=3.0)
    daofind = DAOStarFinder(fwhm=fwhm, threshold=threshold*std)
    sources = daofind(data - median)
    if sources is not None:
        print('{} sources found.'.format(len(sources)))

        if show_sources or save_sources:
            positions = np.transpose((sources['xcentroid'], sources['ycentroid']))
            apertures = CircularAperture(positions, r=4.)
            norm = ImageNormalize(stretch=SqrtStretch())
            plt.imshow(data, cmap='Greys', origin='lower', norm=norm)
            apertures.plot(color='blue', lw=1.5, alpha=0.5)
            if show_sources:
                plt.show()
            if save_sources:
                plt.savefig(plot_name)
                print('Plot saved to {}'.format(plot_name))
            plt.close()
    else:
        print('No sources found.')

    return sources

def fwhm(wavelength):
    """Calculate the approximate FWHM, in arcsec, for a given wavelength.
    This equation was calculated by hand using approximate values for the
    FWHM found on JDox:
    https://jwst-docs.stsci.edu/near-infrared-camera/nircam-predicted-performance/nircam-point-spread-functions#NIRCamPointSpreadFunctions-PSFFWHM

    Parameters
    ----------
    wavelength : float
        Wavelength, in microns

    Returns
    -------
    fwhm : float
        FWHM of the PSF, in units of arcsec
    """
    return 0.03 * wavelength + 0.004

def get_fwhm(filter_name):
    """Calculate the FWHM in units of pixels given a filter name

    Parameters
    ----------
    filter_name : str
        e.g. F200W

    Returns
    -------
    FWHM of the PSF, in units of pixels
    """
    filter_name = filter_name.upper()
    try:
        wave = float(filter_name[1:4]) / 100.
    except ValueError:
        raise ValueError("ERROR: Cannot get a wavelength from the filter name: {}".format(filter_name))

    fwhm_arcsec = fwhm(wave)
    if wave < 2.4:
        pix_scale = 0.031  # arcsec/pixel
    else:
        pix_scale = 0.062  # arcsec/pixel
    return fwhm_arcsec / pix_scale


# 20210712 build

# TODO
# 剪切板助手/切换二维码和设置历史剪切板
# 退出最小化到状态栏/全局Esc和鼠标移动到左上角
# 图标/二维码切换

# FIXME
# 关闭和重启时位置可能漂移

# TODO FUTURE
# 分段和翻页
# 定时切换回到图标/鼠标放置除外
# 可配置参数文件
# 支持文件

import os
from threading import Thread

import wx
import pynput
import qrcode

__ver__ = 'Ver 0.01'

CONFIG = 'config.ini'

def qrmake(qrstr):
    qr = qrcode.QRCode(box_size=2, border=2)
    qr.add_data(qrstr)
    qr.best_fit()
    qr.makeImpl(False, 3) # 3 or 6
    img = qr.make_image()
    return img

class ClipboardHistroy:
    def __init__(self, parent):
        self.data = []

    def Latest(self):
        do = wx.TextDataObject()
        if wx.TheClipboard.Open():
            ret = wx.TheClipboard.GetData(do)
            wx.TheClipboard.Close()
            if ret:
                text = do.GetText()
                if text not in self.data[:1]: # maybe data is empty
                    if text in self.data[1:]:
                        self.data.remove(text)
                    self.data.insert(0, text)
                    return text

class Monitor:
    def __init__(self, func):
        self.func = func
        t1 = Thread(target=self.MouseListener)
        t2 = Thread(target=self.KeyboardListener)
        t1.start()
        t2.start()

    def refresh(self, *args):
        self.func()

    def ignore(self, *args):
        return

    def MouseListener(self):
        with pynput.mouse.Listener(on_move=self.ignore, on_click=self.refresh, on_scroll=self.ignore) as self.ml:
            self.ml.join()

    def KeyboardListener(self):
        with pynput.keyboard.Listener(on_press=self.ignore, on_release=self.refresh) as self.kl:
            self.kl.join()

    def Close(self):
        pynput.keyboard.Listener.stop(self.kl)
        pynput.mouse.Listener.stop(self.ml)

class MyPanel(wx.Panel):
    def __init__(self, parent):
        wx.Panel.__init__(self, parent)
        self.parent = parent
        self.bmp = wx.StaticBitmap(self)
        self.history = ClipboardHistroy(self)
        self.monitor = Monitor(self.SetQrCode)
        self.SetQrCode()

    def SetQrCode(self):
        text = self.history.Latest()
        if text:
            text = text[:1000] # TODO 增加翻页
            qr = qrmake(text)
            img_pil = qr.convert('RGB')
            img_wx = wx.Image(img_pil.size, img_pil.tobytes())

            w0, h0 = self.bmp.GetSize()
            self.bmp.SetBitmap(wx.Bitmap(img_wx))

            w, h = self.bmp.GetSize()
            x, y = self.parent.GetPosition()
            self.parent.SetSize((w+2, h+2)) # TODO 自动合适大小
            self.parent.SetPosition((int(x+(w0-w)/2), int(y+(h0-h)/2))) # TODO 自动合适大小

            tt = wx.ToolTip('Length: %d\n%s'%(len(text), text))
            self.bmp.SetToolTip(tt)

class MyFrame(wx.Frame):
    def __init__(self):
        wx.Frame.__init__(self, None, -1, 'QR Assistant - %s'%__ver__,style = wx.STAY_ON_TOP)

        self.delta = (0,0)

        self.panel = MyPanel(self)
        self.LoadLocation()
        self.Binding(self.panel.bmp)
        self.Show()

    def OnExit(self, evt):
        with open(CONFIG, 'w') as f:
            f.write(str(self.GetPosition()))
        self.panel.monitor.Close()
        self.Close()

    def LoadLocation(self):
        if os.path.exists(CONFIG):
            with open(CONFIG) as f:
                xy = eval(f.read())
                self.SetPosition(xy)
        else:
            self.Center()

    def OnLeftDown(self, evt):
        self.CaptureMouse()
        x, y = self.ClientToScreen(evt.GetPosition())
        originx, originy = self.GetPosition()
        dx = x - originx
        dy = y - originy
        self.delta = ((dx, dy))

    def OnLeftUp(self, evt):
        if self.HasCapture():
            self.ReleaseMouse()

    def OnMouseMove(self, evt):
        if evt.Dragging() and evt.LeftIsDown():
            x, y = self.ClientToScreen(evt.GetPosition())
            fp = (x - self.delta[0], y - self.delta[1])
            self.Move(fp)

    def Binding(self, widget):
        widget.Bind(wx.EVT_RIGHT_UP,  self.OnExit)
        widget.Bind(wx.EVT_LEFT_DOWN, self.OnLeftDown)
        self.Bind(wx.EVT_LEFT_UP,     self.OnLeftUp)
        self.Bind(wx.EVT_MOTION,      self.OnMouseMove)

app = wx.App()
frm = MyFrame()
app.MainLoop()


import csv
import matplotlib.pyplot as plt
from datetime import datetime

filename = 'data/sitka_weather_2018_simple.csv'

with open(filename) as f:
    reader = csv.reader(f)
    header = next(reader)

# for index,coloumn in enumerate(header):
#     print(index,coloumn)

    time,highs,lows = [] , []  , []
    for row in reader:
        current_date = datetime.strptime(row[2], '%Y-%m-%d')
        pop = int(row[5])
        lop = int(row[6])
        time.append(current_date)
        highC = ((pop-32)*5//9)
        lowC = ((lop-32)*5//9)
        highs.append(highC)
        lows.append(lowC)
print(highs)

# data visualistaion for hight temperatures
plt.style.use('seaborn')
fig, ax = plt.subplots()
ax.plot(time,highs,c= 'red', alpha =0.5)
ax.plot(time,lows,c = 'blue', alpha =0.5)
ax.fill_between(time, highs,lows , facecolor = 'blue', alpha= 0.1)
#setting
ax.set_title("Daily High and Low temperature " , fontsize = 24)
ax.set_xlabel('',fontsize = 16)
fig.autofmt_xdate()
ax.set_ylabel('HIGH And Low TEMPERATURES (C)', fontsize = 16)
ax.tick_params(axis ='both', which ='major' , labelsize = 16)

plt.show()


import unittest

from pysqler import *

class SearchTestCase(unittest.TestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def compare_sql(self, expected, actually):
        v1 = expected.replace("\n", "")
        v1 = v1.replace(" ", "")

        v1 = v1.lower()

        v2 = actually.replace("\n", "")
        v2 = v2.replace(" ", "")
        v2 = v2.lower()

        self.assertEqual(v1, v2, "sqler.Select() build  sql error")

    def test_sql_select(self):
        query = Select()
        query.select("city", "education", "AVG(age) as avg_age")
        query.from1("people")
        query.where("age", ">", 10)
        query.and_where("job", "like", "%it%")
        query.and_where("birthday", ">", "1988-09-12 12:12:12")
        query.and_where("address", "!=", None)
        query.groupby("city", "education")
        query.orderby("avg_age", "DESC")
        query.limit(8, 10)

        expected = """
        SELECT city,education,AVG(age) as avg_age
        FROM people
        WHERE age > 10 AND job like "％it％"
        AND birthday > "1988-09-12 12:12:12"
        AND address IS NOT null
        GROUP BY city,education
        ORDER BY avg_age DESC
        LIMIT 8,10
        """

        query_str = str(query)
        print(query_str)
        self.compare_sql(expected, query_str)

    def test_sql_insert1(self):
        query = Insert("people")
        query.put("name", "barry")

        query.put("age", 10, value_on_duplicated=20)

        express = Expression()
        express.field("salary")
        express.operator("+")
        express.value(200)
        express.operator("*")
        express.value(3.5)

        query.put("salary", 50000, value_on_duplicated=express)
        query.put("address", "shanghai", value_on_duplicated="china")
        query.put("education", "bachelor")
        query.put("job", None)
        query.put("birthday", "2000-01-01")
        query_str = str(query)
        print(query_str)

        expected = """
        INSERT INTO people ( name,age,salary,address,education,job,birthday )
        VALUES( "barry",10,50000,"shanghai","bachelor",null,"2000-01-01" )
        ON DUPLICATE KEY UPDATE age = 20,salary = salary + 200 * 3.5,
        address = "china"
        """
        self.compare_sql(expected, query_str)

    def test_sql_insert2(self):
        query = Insert("people")
        query.add_columns("name", "age", "salary", "address", "education",
                          "job", "birthday")

        query.add_row("barry", 19, 3100, "shanghai", "bachelor", None,
                      "2010-01-01")

        query.add_row("jack", 24, 3600, "shanghai", "bachelor", "engineer",
                      "2010-01-09")

        query.add_row("bob", 27, 8600, None, "bachelor", "engineer",
                      "1990-01-09")

        query.add_row("edwin", 30, 10600, "beijing", "bachelor", "engineer",
                      "1987-01-09")

        query_str = str(query)
        print(query_str)

        expected = """
        INSERT INTO people ( name,age,salary,address,education,job,birthday )
        VALUES( "barry",19,3100,"shanghai","bachelor",null,"2010-01-01" ),
        ( "jack",24,3600,"shanghai","bachelor","engineer","2010-01-09" ),
        ( "bob",27,8600,null,"bachelor","engineer","1990-01-09" ),
        ( "edwin",30,10600,"beijing","bachelor","engineer","1987-01-09" )

        """
        self.compare_sql(expected, query_str)

    def test_sql_update1(self):
        query = Update("people")
        query.put("name", "barry")
        query.put("age", 10)

        query.where("age", ">", 15)
        query.or_where("age", "<", 5)
        query_str = str(query)
        print(query_str)

        expected = """
        UPDATE people SET name = "barry",age = 10
        WHERE age > 15 OR age < 5
        """
        self.compare_sql(expected, query_str)

    def test_sql_delete1(self):
        query = Delete("people")

        query.where("age", ">", 15)
        query.or_where("name", "in", [9527, "barry", "jack"])
        query_str = str(query)
        print(query_str)

        expected = """
        DELETE FROM people  WHERE age > 15 OR name in (9527,"barry","jack")
        """
        self.compare_sql(expected, query_str)

if __name__ == '__main__':
    unittest.main()


from src.utils.loader import *
from scipy.sparse import *
from src.utils.evaluator import *
from src.MF.NMF import *

def main():
    ds = Dataset()
    ev = Evaluator()
    ev.cross_validation(5, ds.train_final.copy())
    for i in range(0, 5):
        urm, tg_tracks, tg_playlist = ev.get_fold(ds)
        nmf = NMF(urm, features=150, learning_steps=10000)
        nmf.fit(1e-6, 1e-6)
        recs = nmf.predict(list(tg_playlist), list(tg_tracks), ds)
        ev.evaluate_fold(recs)
    map_at_five = ev.get_mean_map()
    print("MAP@5 [Features = 100, Iterations: 10000] ", map_at_five)

    # export csv
    urm = ds.build_train_matrix()
    nmf_exporter = NMF(urm, features=150, learning_steps=10000)
    tg_playlist = list(ds.target_playlists.keys())
    tg_tracks = list(ds.target_tracks.keys())
    # Train the model with the best shrinkage found in cross-validation
    nmf_exporter.fit(1e-6, 1e-6)
    recs = nmf_exporter.predict(tg_playlist, tg_tracks, ds)
    with open('submission_nmf.csv', mode='w', newline='') as out:
        fieldnames = ['playlist_id', 'track_ids']
        writer = csv.DictWriter(out, fieldnames=fieldnames, delimiter=',')
        writer.writeheader()
        for k in tg_playlist:
            track_ids = ''
            for r in recs[k]:
                track_ids = track_ids + r + ' '
            writer.writerow({'playlist_id': k,
                             'track_ids': track_ids[:-1]})

if __name__ == '__main__':
    main()


#!/usr/bin/env python
import rospy
from std_msgs.msg import String
from ar_track_alvar_msgs.msg import AlvarMarkers
from geometry_msgs.msg import Point, Twist
import numpy as np
from tf.transformations import euler_from_quaternion, quaternion_from_euler
import math
import time
j = 0
dat = 0
flag= 0
flag2 = 0
def callback(msg):
    global j, dat
    dat = msg
    j = 0
    ids = []
    while 1:
        try:
            x = int(msg.markers[j].id)
            ids.append(x)
            j += 1
        except IndexError:
            break

## NODE INTITIALIZATION
rospy.init_node('ar_gate', disable_signals=True)

rospy.Subscriber("/ar_pose_marker", AlvarMarkers, callback)
pub = rospy.Publisher('cmd_vel', Twist, queue_size = 1)

speed = Twist()
r = rospy.Rate(10)

while not rospy.is_shutdown():
    print("No. of tags: ",j )
    if j == 2:
        flag = 1
        try:
            global id1, id2
            ## Tag IDS
            id1 = dat.markers[0].id
            id2 = dat.markers[1].id

            ## POSITION OF TAGS
            pos_z1 = dat.markers[0].pose.pose.position.z*100
            pos_x1 = dat.markers[0].pose.pose.position.x*100

            pos_x2 = dat.markers[1].pose.pose.position.x*100
            pos_z2 = dat.markers[1].pose.pose.position.z*100

            ## DISTANCE BETWEEN TAG AND BOT
            dist1 = pow(pow(pos_x1, 2) + pow(pos_z1, 2), 0.5)
            dist2 = pow(pow(pos_x2, 2) + pow(pos_z2, 2), 0.5)

            ## CONDITIONS --

            if dist2>dist1:
                x = dist2-dist1
                if x > 5:
                    speed.linear.x = -0.4
                    speed.angular.z = -pos_x2/500
                    pub.publish(speed)
                else:
                    speed.linear.x = -0.4
                    speed.angular.z = 0
                    pub.publish(speed)

            if dist1> dist2:
                y = dist1 - dist2
                if y > 5:
                    speed.linear.x = -0.4
                    speed.angular.z = -pos_x1/500
                    pub.publish(speed)
                else:
                    speed.linear.x = -0.4
                    speed.angular.z = 0
                    pub.publish(speed)

            print(speed)
        except NameError as e:
            print(e)
            pass

    if j==1:
        or_x = dat.markers[0].pose.pose.orientation.x
        or_y = dat.markers[0].pose.pose.orientation.y
        or_z = dat.markers[0].pose.pose.orientation.z
        or_w = dat.markers[0].pose.pose.orientation.w

        quaternion = (or_x, or_y, or_z, or_w)
        euler = euler_from_quaternion(quaternion)
        pitch = math.degrees(euler[1])+180
        print(pitch)
        if pitch<180:
            speed.angular.z = pitch/360
            speed.linear.x = -0.4
            pub.publish(speed)

        elif pitch>180:
            speed.angular.z = -pitch/360
            speed.linear.x = -0.4
            pub.publish(speed)

    elif j==0 and flag == 1:
        if flag == 1:
            speed.linear.x = -0.4
            speed.angular.z = 0
            pub.publish(speed)
            print(speed)
            time.sleep(4)
            flag = 2
        if flag == 2:
            speed.linear.x = 0
            speed.angular.z = 0
            pub.publish(speed)
            print(speed)

    r.sleep()


def calculateHandlen(hand):
    """
    Returns the length (number of letters) in the current hand.

    hand: dictionary (string int)
    returns: integer
    """
    c=0
    char=0
    for char in hand.values():
        c=c+char
    return c