import os
import sys
import time
import logging
import argparse
import itertools
import matplotlib; matplotlib.use('Agg')
import multiprocessing
import pandas as pd
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

from operator import add
from functools import partial
from collections import Counter
from matplotlib import font_manager
from GetConfig import getConfig

config = getConfig()

'''
Description:
This module is used for network analysis.
Network analysis is an efficient method to show relationship of haplotypes.
Y haplotype data is required for analysis.
It is recommended to use Y-STR haplotype under a specific haplogroup.
Network analysis can show the evolutionary relationship between haplotypes.
This command can output network figure, and fdi file.
'''

def net_parser():

    parser = argparse.ArgumentParser('net', description='(c) Y-LineageTracker: Network analysis')
    # function used for network analysis
    parser.add_argument('net',
                        help='Perform network analysis from Y haplotype data.')
    # required, input file, alignment sequence or matrix
    input = parser.add_mutually_exclusive_group(required=True)
    input.add_argument('--seq',
                        type=str,
                        action='store',
                        help='seq: Haplotype data in sequence alignment format')
    input.add_argument('--matrix',
                        type=str,
                        action='store',
                        help='martix: Y-STR haplotype data in matrix format')
    # optional, format of sequence alignment
    parser.add_argument('--seq-format',
                        required=False,
                        type=str,
                        dest='format',
                        action='store',
                        choices=['fasta', 'phylip', 'nexus', 'meg', 'vcf'],
                        help='--seq-format: The format of sequence file. This option is only required for the sequence alignment file. ')
    # optional, population file
    parser.add_argument('-p', '--population',
                        required=False,
                        type=str,
                        action='store',
                        help='population: A file containing sample ID and population information of each individual.')
    # optional, type of tree
    parser.add_argument('--tree-type',
                        required=False,
                        type=str,
                        dest='treetype',
                        action='store',
                        default='mjn',
                        choices=['mjn', 'msn'],
                        help='tree-type: The algorithm used to construct a network tree.')
    # optional, how to code missing site
    parser.add_argument('--gap',
                        required=False,
                        type=str,
                        action='store',
                        default='allele',
                        choices=['allele', 'missing'],
                        help='gap: The gap coding method. You can choose to code gaps as a new type of allele or missing data')
    # optional, the search for new median nodes
    parser.add_argument('--search',
                        required=False,
                        type=int,
                        action='store',
                        default=0,
                        help='search: How the search for new median nodes is performed: the larger this parameter, the wider the search. ')
    # optional, output fdi file
    parser.add_argument('--fdi',
                        required=False,
                        action='store_true',
                        help='fdi: Output fdi file that can be used as the input file in Network software, to show network plot or estimate TMRCA directly.')
    # optional, filter sites by misisng rate value
    parser.add_argument('--filter',
                        required=False,
                        type=float,
                        action='store',
                        default=0.7,
                        help='filter: The cutoff value of the site to be filtered, sites with a missing rate value greater than this value will be filtered')
    # optional, how to show alternative link in figure
    parser.add_argument('--alt',
                        required=False,
                        type=str,
                        action='store',
                        default='hide',
                        choices=['hide', 'show', 'dashed'],
                        help='alt: Whether to show alternative links between nodes in the network figure.')
    # optional, size of median nodes
    parser.add_argument('--mediansize',
                        required=False,
                        type=float,
                        action='store',
                        default=0.0,
                        help='mediansize: The median node size in the median-joining network compared to base size in the network figure. The default is 0.')
    # optional, keep all nodes in same size
    parser.add_argument('--euqal-size',
                        required=False,
                        dest='euqalsize',
                        action='store_true',
                        help='euqal-size: Keep all nodes the same size in the network figure.')
    # optional, the prefix of output
    parser.add_argument('-o', '--output',
                        required=False,
                        type=str,
                        action='store',
                        help='output: The prefix of output files. ')

    args = parser.parse_args()

    return args

# check missing rate is filter argument is used
def check_misisng_rate(missing_rate):

    if missing_rate < 0 or missing_rate> 1:
        print('[Y-LineageTracker] [Error] Cutoff of missing rate should be a float between 0 and 1')
        sys.exit()

# print program information and write to log file
def set_log(args_log, log_file):

    logger = logging.getLogger()
    logger.setLevel(level=logging.INFO)

    handler = logging.FileHandler(log_file, mode='w')
    handler.setLevel(logging.INFO)
    formatter = logging.Formatter('[%(asctime)s] - [%(levelname)s]: %(message)s')
    handler.setFormatter(formatter)

    console = logging.StreamHandler()
    console.setLevel(logging.INFO)

    logger.addHandler(handler)
    logger.addHandler(console)

    log_info = ['[Y-LineageTracker] [Network]',
                '[Y-LineageTracker] Run Date: ' + time.asctime(time.localtime(time.time())),
                '[Y-LineageTracker] Haplotype File: %s' % args_log.seq,
                '[Y-LineageTracker] Haplotype File Format: %s' % args_log.format,
                '[Y-LineageTracker] Gap euqal to %s' % args_log.gap,
                '[Y-LineageTracker] Searching parameter: %d' % args_log.search,
                '[Y-LineageTracker] Missing rate of allele to be filtered: %s' % args_log.filter,
                '[Y-LineageTracker] Alternative link type : %s' % args_log.alt,
                '[Y-LineageTracker] Median node size: %s' % args_log.mediansize]

    if args_log.euqalsize:
        log_info.append('[Y-LineageTracker] Keep all node the same size')
    else:
        log_info.append('[Y-LineageTracker] Plot node according to sample size')

    print('\n')
    for i in log_info:
        logger.info(i)

class HaploNetwork(object):
    '''
    This class is used for network analysis.
    Main functions of this class include:
    1. construct median-joining network from haplotypes
    2. plot network by networkx package
    3. write network to fdi file
    4. output network info
    '''

    def __init__(self, gap, path):

        self.logger = logging.getLogger()
        self.skip_letters = config.get('Statistics', 'SkipLetters').split(',')
        self.gap = gap
        self.path = path
        self.net_info = open(self.path+'.info.net', 'w')

    # convert STR numbers to 0-1 code haplotype
    def _convert_allel_to_binary(self, haplotype_data, filter):

        # create a empty dataframe to save haplotypes
        sequence_df = pd.DataFrame(index=haplotype_data.index,
                                   columns=['Haplotype'])
        sequence_df['Haplotype'] = ''

        STR_code_pos = {}
        start_pos = 0
        hap_length = 0
        allel_num = haplotype_data.index.size
        for STR in haplotype_data.columns.tolist():
            str_allels = haplotype_data[STR].astype(str).tolist()
            # filter STR according to filter parameter
            if sum([str_allels.count(i) for i in self.skip_letters])/allel_num >= filter:
                continue
            if len(set(set(str_allels)&set(self.skip_letters))) > 1:
                str_allels = ['.' if i in self.skip_letters else i for i in str_allels]
            repeats_type = sorted(list(set(str_allels)))
            # convert STR repeat numbers to 0-1 code
            repeats_type_num = len(repeats_type)
            max_length = len('{0:b}'.format(repeats_type_num-1))
            binary_set = [] # creat a empty list to save 0-1 code
            for i in range(repeats_type_num):
                binary_num = '{0:b}'.format(i)
                binary_num_length = len(binary_num)
                if binary_num_length < max_length:
                    binary_num = (max_length-binary_num_length)*'0'+binary_num
                binary_set.append(binary_num)
            if self.gap == 'missing':
                replaced_allels = ['-'*max_length if i == '.' else binary_set[repeats_type.index(i)] for i in str_allels]
            else:
                replaced_allels = [binary_set[repeats_type.index(i)] for i in str_allels]
            sequence_df['Haplotype'] = list(map(add, sequence_df['Haplotype'].tolist(), replaced_allels))
            for pos in range(start_pos, start_pos+max_length):
                STR_code_pos[pos] = STR
            start_pos += max_length
            hap_length += 1

        self.net_info.write('#HapLength:%d\n\n' % hap_length)

        return sequence_df, STR_code_pos

    # calculate frequency of all kinds of haplotypes and the proportion in each population
    def _calculate_distance_matrix(self, sequence_df):

        # get haplotypes and name for them
        haplotypes = sorted(list(set(sequence_df['Haplotype'])))
        sp_sequence_df = pd.DataFrame(index=range(len(haplotypes)), columns=['Haplotype'])
        sp_sequence_df['Haplotype'] = haplotypes

        # calculate distance between two haplotypes
        distance_matrix = pd.DataFrame(index=range(len(haplotypes)), columns=range(len(haplotypes)))
        hap_num = 0
        for hap1_name, hap1 in enumerate(haplotypes):
            hap_num += 1
            for hap2_name, hap2 in enumerate(haplotypes[hap_num:]):
                if self.gap == 'missing':
                    distance = sum(l1 != l2 for l1, l2 in zip(hap1, hap2) if l1 not in self.skip_letters and l2 not in self.skip_letters)
                    if distance == 0:
                        distance = 1
                else:
                    distance = sum(l1 != l2 for l1, l2 in zip(hap1, hap2))
                hap2_name = hap2_name + hap_num
                distance_matrix.at[hap2_name, hap1_name] = distance

        self.net_info.write('#HapSamples:\n')

        sample_sp_code = {}
        for i in haplotypes:
            idx = sequence_df['Haplotype'].tolist().index(i)
            sample = sequence_df.index[idx]
            sample_sp_code[sample] = i
            all_samples = ','.join(sequence_df[sequence_df['Haplotype']==i].index.tolist())
            self.net_info.write('%d\t%s\n' %(idx, all_samples))

        self.net_info.write('\n')

        return distance_matrix, sp_sequence_df, sample_sp_code

    # get msn tree
    def _calculate_minimum_spanding_tree(self, n, distance_matrix, return_pruned=True):

        if n < 2:
            self.logger.error('[Y-LineageTracker] Program stopped since there are less than 2 types of haplotypes')
            sys.exit()

        # create a empty edge matrix to save hapotype and their distance
        edge_num = n-1
        edge_matrix = np.empty(shape=(edge_num, 3))
        edge_matrix[:] = np.nan

        # get the distance from small to large between of two haplotypes
        forest = np.linspace(0, n-1, n)

        add_num = 0
        while add_num < edge_num:
            ci, ri = np.unravel_index(np.nanargmin(distance_matrix.T.values), distance_matrix.T.shape)
            row_record = forest[ri]
            col_record = forest[ci]
            if row_record != col_record:
                edge_matrix[add_num] = [ci, ri, distance_matrix.at[ri, ci]] # col hap, row hap, step
                forest[forest==col_record] = row_record
                add_num += 1
            distance_matrix.at[ri, ci] = np.nan

        if return_pruned:
            return edge_matrix, distance_matrix
        else:
            return edge_matrix

    # get median haplotypes
    @staticmethod
    def calculate_median_vector(triplet, sp_sequence_df, max_n_cost, epsilon, n, gap, new_link_data, shared_median_haps, initial_lambda, n_cost, lock):

        median_haps = pd.DataFrame(columns=['Haplotype'])
        temp_link_data = np.empty(shape=(0, 3))
        median_num = n

        for item in triplet:
            hap_in_order = item
            # check whether two haplotypes can be linked
            condition1 = any(hap_in_order[0]==np.array(new_link_data)[:, 0]) and any(hap_in_order[1]==np.array(new_link_data)[:, 1])
            condition2 = any(hap_in_order[0]==np.array(new_link_data)[:, 0]) and any(hap_in_order[2]==np.array(new_link_data)[:, 1])
            condition3 = any(hap_in_order[1]==np.array(new_link_data)[:, 0]) and any(hap_in_order[2]==np.array(new_link_data)[:, 1])
            if (condition1+condition2+condition3) < 2:
                continue
            # get alleles that are variable
            sub_seq = sp_sequence_df.loc[list(hap_in_order)]
            # get the median haplotype
            sub_seq_length = len(sub_seq.at[hap_in_order[0], 'Haplotype'])
            median_seq = ''.join([Counter(sub_seq['Haplotype'].map(lambda x: x[i])).most_common(1)[0][0] for i in range(sub_seq_length)])
            # check whether median sequence exist
            if median_seq in sp_sequence_df['Haplotype'].tolist():
                continue
            if not median_haps.empty:
                if median_seq in shared_median_haps:
                    continue
            # calculate the distance between median sequence and haplotypes
            if gap == 'missing':
                median_distance = [sum (l1 != l2 for l1, l2 in zip(i, median_seq) if l1 not in self.skip_letters and l2 not in self.skip_letters) for i in sub_seq['Haplotype']]
                if 0 in median_distance:
                    median_distance = [1 if dist==0 else dist for dist in median_distance]
            else:
                median_distance = [sum (l1 != l2 for l1, l2 in zip(i, median_seq)) for i in sub_seq['Haplotype']]
            connection_cost = sum(median_distance)
            # update n cost
            lock.acquire()
            n_cost.value += 1
            if n_cost.value > max_n_cost:
                print('[Y-LineageTracker] Program stopped since there are too many connnections')
                sys.exit()
            if connection_cost < initial_lambda.value:
                # update initial lambda
                initial_lambda.value = connection_cost
            if connection_cost <= initial_lambda.value + epsilon:
                hap_to_add1 = hap_in_order
                hap_to_add2 = [median_num]*3
                step_to_add = median_distance
                link_data_to_add = np.array([hap_to_add1, hap_to_add2, step_to_add]).T
                temp_link_data = np.concatenate([temp_link_data, link_data_to_add])
                median_haps.loc[median_num] = median_seq
                median_num += 1
                if median_seq not in shared_median_haps:
                    shared_median_haps.append(median_seq)
            lock.release()

        return median_haps, temp_link_data

    # construct network from distance matrix
    def _calculate_network(self, distance_matrix, sp_sequence_df, tree_type, epsilon):

        # initialization
        n = sp_sequence_df.index.size
        haplotype_length = len(sp_sequence_df.at[0, 'Haplotype'])
        median_num = n
        MST, pruned_distance_matrix = self._calculate_minimum_spanding_tree(n ,distance_matrix)

        if n < 3 or tree_type == 'msn':
            self.logger.warning('There are less than 3 types of haplotypes')
            alt_data = np.empty(shape=(0, 3))
            return MST, alt_data

        MST_step = MST[:, 2]
        lower_limit = min(MST_step)
        upper_limit = max(MST_step) + epsilon
        new_link_data = []
        for i in pruned_distance_matrix.columns:
            for j in pruned_distance_matrix.index:
                step = pruned_distance_matrix.at[j, i]
                if (not np.isnan(step)) and step > lower_limit and step <= upper_limit:
                    new_link_data.append([i, j, step])
        if new_link_data:
            new_link_data = np.concatenate([MST, np.array(new_link_data)])
        else:
            new_link_data = MST

        start_n_cost = 0
        start_initial_lambda = float('inf')
        tripleted = []

        while True:

            # get triplet for calculation
            triplet_new = list(itertools.combinations(range(n), 3))
            triplet = [i for i in triplet_new if i not in tripleted]
            tripleted.extend(triplet)

            process_num = 10
            pool = multiprocessing.Pool(processes=process_num)

            # global variables to be shared
            manager = multiprocessing.Manager()
            n_cost = manager.Value('i', start_n_cost)
            initial_lambda  = manager.Value('f', start_initial_lambda)
            shared_new_link_data = manager.list(new_link_data)
            shared_median_haps = manager.list([])
            LOCK = manager.Lock()

            # split triplet to sections
            triplet_interval = np.linspace(0, len(triplet), process_num, dtype=int)
            #triplet_sections = [triplet[triplet_interval[n]: triplet_interval[n+1]] if n == 0 else triplet[(triplet_interval[n]+1): triplet_interval[n+1]] for n in range(process_num-1)]
            triplet_sections = [triplet[triplet_interval[n]: triplet_interval[n+1]] for n in range(process_num-1)]

            # multiprocess triplet data
            pt = partial(self.calculate_median_vector, # function
                         sp_sequence_df=sp_sequence_df, # variable for read
                         max_n_cost=40000, # variable for read
                         epsilon=epsilon, # variable for read
                         n=n, # variable for read
                         gap=self.gap, # variable for read
                         new_link_data=new_link_data, # variable for writing
                         shared_median_haps=shared_median_haps, # variable for writing
                         initial_lambda=initial_lambda, # variable for writing
                         n_cost=n_cost, # variable for writing
                         lock=LOCK) # process lock
            median_haps_list, temp_link_data_list = zip(*pool.map(pt, triplet_sections))

            pool.close()
            pool.join()

            # intergrate data
            median_haps = pd.concat(median_haps_list)
            temp_link_data = np.concatenate(temp_link_data_list)
            if median_haps.empty:
                self.logger.info('[Y-LineageTracker] Get all median haplotypes')
                break
            else:
                # code for median vectors
                start_num = sp_sequence_df.index.size
                end_num = start_num + median_haps.index.size
                median_haps.index = range(start_num, end_num)
                temp_link_data[:, 2] = [val for val in range(start_num, end_num) for _ in (0, 1, 2)]
                # add to haplotype information
                sp_sequence_df = pd.concat([sp_sequence_df, median_haps])
                new_link_data = np.concatenate([new_link_data, temp_link_data])
                # update for iteration
                start_n_cost = n_cost.value
                start_initial_lambda = initial_lambda.value

        sp_sequence_df = sp_sequence_df.drop_duplicates().reset_index()

        self.logger.info('[Y-LineageTracker] Calculating distance matrix...')
        # calculate and add median haplotype to distance matrix
        hap_num = 0
        for hap1_name, hap1 in zip(sp_sequence_df.index, sp_sequence_df['Haplotype']):
            hap_num += 1
            for hap2_name, hap2 in zip(sp_sequence_df.index[hap_num:], sp_sequence_df['Haplotype'][hap_num:]):
                if self.gap == 'missing':
                    distance = sum(l1 != l2 for l1, l2 in zip(hap1, hap2) if l1 not in self.skip_letters and l2 not in self.skip_letters)
                    if distance == 0:
                        distance = 1
                else:
                    distance = sum(l1 != l2 for l1, l2 in zip(hap1, hap2))
                distance_matrix.at[hap2_name, hap1_name] = distance

        N = sp_sequence_df.index.size
        dnet = distance_matrix.copy().applymap(lambda x: 0 if not np.isnan(x) else x)
        forest = np.linspace(0, N-1, N)
        network_data = np.empty(shape=(0, 3))
        alt_data = np.empty(shape=(0, 3))
        steps = sorted(list(set(filter(lambda x: x==x, distance_matrix.to_numpy().reshape(-1)))))

        edge_num = 0
        add_new_data = False

        self.logger.info('[Y-LineageTracker] Constructing relationship between nodes')
        while len(set(forest)) > 1:

            step_num = steps[edge_num]
            step_df = distance_matrix[distance_matrix==step_num]
            step_count = step_df.notnull().sum().sum()
            for i in range(step_count):
                ci, ri = np.unravel_index(np.nanargmin(step_df.T.values), step_df.T.shape)
                step_df.at[ri, ci] = np.nan
                row_record = forest[ri]
                col_record = forest[ci]
                if row_record != col_record:

                    for col in np.where(forest==col_record)[0]:
                        if ci > col:
                            num_col = dnet.at[ci, col]
                        elif ci < col:
                            num_col = dnet.at[col, ci]
                        else:
                            num_col = 0
                        for row in np.where(forest==row_record)[0]:
                            if ri > row:
                                num_row = dnet.at[ri, row]
                            elif ri < row:
                                num_row = dnet.at[row, ri]
                            else:
                                num_row = 0
                            dnet_num = num_col + num_row + step_num
                            dnet.at[row, col] = dnet_num
                    forest[forest==col_record] = row_record
                    add_new_data = True
                else:
                    if step_num < dnet.at[ri, ci]:
                        add_new_data = True
                        # update dnet
                        cols = np.append(network_data[network_data[:, 1]==ci, 0], network_data[network_data[:, 0]==ci, 1])
                        rows = np.append(network_data[network_data[:, 1]==ri, 0], network_data[network_data[:, 0]==ri, 1])
                        for col in np.append(ci, cols).astype(int):
                            if ci > col:
                                num_col = dnet.at[ci, col]
                            elif ci < col:
                                num_col = dnet.at[col, ci]
                            else:
                                num_col = 0
                            for row in np.append(ri, rows).astype(int):
                                if col == row:
                                    continue
                                else:
                                    if ri > row:
                                        num_row = dnet.at[ri, row]
                                    elif ri < row:
                                        num_row = dnet.at[row, ri]
                                    else:
                                        num_row = 0
                                    dnet_num = num_col + num_row + step_num
                                    if row > col:
                                        if dnet.at[row, col] > dnet_num:
                                            dnet.at[row, col] = dnet_num
                                    elif row < col:
                                        if dnet.at[col, row] > dnet_num:
                                            dnet.at[col, row] = dnet_num
                if add_new_data:
                    network_to_add = np.array([[ci, ri, step_num]])
                    network_data = np.concatenate([network_data, network_to_add])
                    add_new_data = False
            edge_num += 1
            if edge_num > 5000:
                self.logger.error('Program stopped since there are too many connnections')
                sys.exit()

        network_data = network_data.astype(int)
        net_MST = self._calculate_minimum_spanding_tree(N ,distance_matrix, False)
        net_MST = net_MST.astype(int)

        if len(network_data) > len(net_MST):
            alt_data = np.array([i for i in network_data if i[:2].tolist() not in net_MST[:, :2].tolist()])
            network_data = np.array([i for i in network_data if i[:2].tolist() in net_MST[:, :2].tolist()])
        else:
            alt_data = np.array([])

        self.logger.info('[Y-LineageTracker] Calculating median haplotypes finished')

        return network_data, alt_data, sp_sequence_df

    # prune network for visualization
    def _get_network_for_plot(self, network_data, alt_data, hap_num, alt_type):

        # specify network data according to type of alternative link
        if alt_data.size != 0:
            alt_type = alt_type
            pruned_network_data = network_data.copy()
            max_number = np.max(pruned_network_data[:, :2]) + 1
            while True:
                to_be_deleted_row = []
                for i in range(hap_num, max_number):
                    if np.count_nonzero(pruned_network_data[:, :2]==i) == 1:
                        row_index = np.where(pruned_network_data[:, :2]==i)[0]
                        to_be_deleted_row.append(row_index)
                pruned_network_data = np.delete(pruned_network_data, to_be_deleted_row, axis=0)
                if len(to_be_deleted_row) == 0:
                    break
            # get network data for visulization
            if alt_type == 'hide':
                network_data = pruned_network_data
        else:
            alt_type = 'hide'
            pruned_network_data = network_data.copy()

        self.net_info.write('#HapLinks:\n')
        [self.net_info.write('\t'.join(i)+'\n') for i in pruned_network_data.astype(str)]
        self.net_info.close()
        node_num = len(np.unique(network_data[:, :2]))

        return network_data, node_num

    # get mutation number
    def _get_mutations(self, network_line, sp_sequence_df, STR_code_pos):

        seq1 = sp_sequence_df.at[network_line[0], 'Haplotype']
        seq2 = sp_sequence_df.at[network_line[1], 'Haplotype']
        mutation = [STR_code_pos[n] for n, s in enumerate(zip(seq1, seq2)) if s[0] != s[1]]

        return mutation

    # write and output fdi file
    def _get_fdi_file(self, network_data, torso_num, node_num, sequence_df, sp_sequence_df, sample_sp_code, pos, hap_size, hap_freq, STR_code_pos, pops):

        fdi_file = open(self.path+'.fdi', 'w')
        medians = sorted(list(set(set(range(len(sample_sp_code), sp_sequence_df.index.size)) & set(pos.keys()))))

        # head info
        center_x,  center_y = np.mean(np.array(list(pos.values())).T, axis=1)
        fdi_file.write('SHOW_NODENAMES;TRUE\n'\
                       'NODES_PROPORTIONAL;TRUE\n'\
                       'SHOW_ONLY_TORSO;FALSE\n'\
                       'SHOW_MEDIANS;FALSE\n'\
                       'SHOW_MEDIAN_NAMES;TRUE\n'\
                       'SHOW_CHARACTERS;FALSE\n'\
                       'SHOW_LINES;FALSE\n'\
                       'MIN_CIRC_RADIUS;4\n'\
                       'MAX_CIRC_RADIUS;50\n'\
                       'MED_RADIUS;3\n'\
                       '\n'\
                       'SPECIAL_FILE_TYPE;FALSE\n'\
                       'DIAGRAM_CENTER;%d;%d\n'\
                       'IMAGE_CENTER;%d;%d\n'\
                       % (center_x, center_y, center_x, center_y))

        # taxon info
        from FilesIO import set_color_num
        from matplotlib import colors
        rgb_colors = [list(map(lambda x: round(x*255), c)) for c in set_color_num(len(pops))]
        colors = [i[0]+i[1]*256+i[2]*256*256 for i in rgb_colors]
        fdi_file.write('NUMBER_OF_TAXA;%d\n' % (sequence_df.index.size+len(medians))) # including median nodes
        plot_num = 0
        for name in sequence_df.index:
            if name in sample_sp_code.keys():
                size = hap_size[plot_num]
                freq = hap_freq[plot_num]
                pos_x, pos_y = list(map(round, pos[plot_num]))
                active = 'TRUE'
                plot_num += 1
            else:
                size = 1
                pos_x, pos_y = 0, 0
                active = 'FALSE'
                freq = None
            taxon_info = 'TAXON_NAME;%s;'\
                         'TAXON_NAME_HEIGHT;6;'\
                         'TAXON_NAME_COLOR;0;'\
                         'TAXON_FREQUENCY;%d;'\
                         'TAXON_ORIG_FREQUENCY;1;'\
                         'TAXON_GEOGRAPHY;;'\
                         'TAXON_PHENOTYPE;;'\
                         'TAXON_LINEAGE;;'\
                         'TAXON_GROUP1;Group1;;'\
                         'TAXON_GROUP2;Group2;;'\
                         'TAXON_GROUP3;Group3;;'\
                         'TAXON_X;%d;'\
                         'TAXON_Y;%d;'\
                         % (name, size, pos_x, pos_y)
            if freq:
                pie_info = ''
                colored_pie_num = 0
                for pie_num, pie_freq in enumerate(freq):
                    if pie_freq == 0:
                        continue
                    else:
                        colored_pie_num += 1
                        pie_color_code = colors[pie_num]
                        pie_info += 'TAXON_COLOR_PIE%d;%d;'\
                                    'TAXON_PIE_FREQUENCY%d;%d;'\
                                    'TAXON_STYLE_PIE%d;SOLID;'\
                                    % (colored_pie_num, pie_color_code, colored_pie_num, pie_freq, colored_pie_num)
            else:
                pie_info = 'TAXON_COLOR_PIE1;65535;'\
                           'TAXON_PIE_FREQUENCY1;1;'\
                           'TAXON_STYLE_PIE1;SOLID;'\

            line_info = 'TAXON_LINE_WIDTH;1;'\
                        'TAXON_LINE_COLOR;0;'\
                        'TAXON_LINE_STYLE;SOLID;'\
                        'TAXON_ACTIVE;%s\n'\
                        % (active)

            fdi_file.write(taxon_info + pie_info + line_info)

        # median info
        for n, m in enumerate(medians):
            pos_x, pos_y = list(map(round, pos[m]))
            median_info = 'TAXON_NAME;mv%d;'\
                          'TAXON_NAME_HEIGHT;6;'\
                          'TAXON_NAME_COLOR;0;'\
                          'TAXON_FREQUENCY;0;'\
                          'TAXON_ORIG_FREQUENCY;0;'\
                          'TAXON_GEOGRAPHY;;'\
                          'TAXON_PHENOTYPE;;'\
                          'TAXON_LINEAGE;;'\
                          'TAXON_GROUP1;Group1;;'\
                          'TAXON_GROUP2;Group2;;'\
                          'TAXON_GROUP3;Group3;;'\
                          'TAXON_X;%d;'\
                          'TAXON_Y;%d;'\
                          'TAXON_COLOR_PIE1;255;'\
                          'TAXON_PIE_FREQUENCY1;0;'\
                          'TAXON_STYLE_PIE1;SOLID;'\
                          'TAXON_LINE_WIDTH;1;'\
                          'TAXON_LINE_COLOR;0;'\
                          'TAXON_LINE_STYLE;'\
                          'SOLID;TAXON_ACTIVE;TRUE\n'\
                          % (n, pos_x, pos_y)
            fdi_file.write(median_info)

        fdi_file.write('\n')
        fdi_file.write('NUMBER_OF_LINKS;%d\n' % len(network_data))
        for n, link in enumerate(network_data):
            if n < torso_num:
                in_torso = 'TRUE'
            else:
                in_torso = 'FALSE'
            link_names = [list(sample_sp_code.keys())[name_code] if name_code < len(sample_sp_code) else 'mv'+str(medians.index(name_code)) for name_code in link[:2]]
            link_info = 'LINK_TAXON1;%s;'\
                        'LINK_TAXON2;%s;'\
                        'LINK_IN_TORSO;%s;'\
                        'LINK_TORSO_PART;1;'\
                        'NUMBER_OF_MUTATIONS;%d;'\
                        % (link_names[0], link_names[1], in_torso ,link[2])
            mutation_data = self._get_mutations(link, sp_sequence_df, STR_code_pos)
            mutation_info = ''
            for mutation_n, mutation in enumerate(mutation_data):
                mutation_info += 'MUTATION%d_NAME;%s;'\
                                 'MUTATION%d_NAME_HEIGHT;6;'\
                                 'MUTATION%d_NAME_COLOR;255;'\
                                 'LINK_LINE_WIDTH;1;'\
                                 'LINK_LINE_STYLE;SOLID;'\
                                 'LINK_LINE_COLOR;8421504;'\
                                 'LINK_ACTIVE;TRUE'\
                                 % (mutation_n+1, mutation, mutation_n+1, mutation_n+1)
            fdi_file.write(link_info+mutation_info+'\n')

        # other info
        # median nodes not included
        fdi_file.write('\n')
        other_info = 'NUMBER_OF_HELP_POINT_LINKS;0\n'\
                     '\n'\
                     'LONG_EDGES;0\n'\
                     'TOTAL_TAXA;%d\n'\
                     'TOTAL_HAPLOTYPES;%d\n'\
                     % (sequence_df.index.size, len(sample_sp_code))
        fdi_file.write(other_info)

        # equal info
        equal_info = ''
        equal_num = 0
        for sample, hap in sample_sp_code.items():
            equal_samples = sequence_df[sequence_df['Haplotype']==hap].index.tolist()
            for equal in equal_samples:
                if equal != sample:
                    equal_num += 1
                    equal_info += 'EQUIVALENT_TAXA;%s;%s\n' % (sample, equal)
        fdi_file.write('NUMBER_EQUIVALENT_TAXA;%d\n' % equal_num)
        fdi_file.write(equal_info)

        # end info
        end_info = 'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Phenotype;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Geography;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Lineage;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Group1;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Group2;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Group3;\n'
        fdi_file.write(end_info)

        fdi_file.close()

    # summarize parameters for figure
    def _get_figure_parameters(self, sequence_df, population_data, network_data, alt_data, node_num, original_network_size, alt_type, euqal_size, label):

        haplotypes = sorted(list(set(sequence_df['Haplotype'])))
        # calculate frequency of each haplotypes
        if euqal_size:
            hap_size = [1]*len(haplotypes)
        else:
            hap_size = [sequence_df['Haplotype'].tolist().count(hap) for hap in haplotypes]
            hap_size_per = [size / sequence_df.index.size for size in hap_size ]

        if isinstance(population_data, pd.DataFrame):
            pops = sorted(list(set(population_data[label])))
            # calculate haplotype frequency in each population
            hap_pop_data = pd.DataFrame(index=range(len(haplotypes)), columns=pops)
            for pop in pops:
                pop_inds = population_data[population_data[label]==pop]['SampleID'].tolist()
                pop_data = sequence_df.loc[pop_inds]
                hap_pop_data[pop] = [pop_data['Haplotype'].tolist().count(hap) for hap in haplotypes]
            hap_freq = [hap_pop_data.loc[i].tolist() for i in hap_pop_data.index]
            hap_freq_per = [(hap_pop_data.loc[i]/sum(hap_pop_data.loc[i])).round(6).tolist() for i in hap_pop_data.index]
        else:
            hap_freq_per = None
            pops = None

        # set length of each node
        length_data = np.empty(shape=(0, 3))
        num = np.max(network_data[:, :2])+1
        to_be_deleted_row = []
        for row_index, i in enumerate(network_data):
            step = i[2]
            if step > 1:
                start = i[0]
                end = i[1]
                for i in range(step):
                    if i == 0:
                        length = [start, num, 1]
                    elif i == step-1:
                        length = [num, end, 1]
                    else:
                        length = [num, num+1, 1]
                        num += 1
                    length_data = np.concatenate((length_data, np.array([length])))
                num += 1
                to_be_deleted_row.append(row_index)
                original_network_size -= 1
        network_data = np.delete(network_data, to_be_deleted_row, axis=0)
        length_data = length_data.astype(int)

        # set figure size and network class
        import warnings
        warnings.filterwarnings('ignore')
        font_files = font_manager.findSystemFonts(fontpaths=[os.path.split(os.path.realpath(__file__))[0] + '/sans-serif'])
        font_list = font_manager.createFontList(font_files)
        font_manager.fontManager.ttflist.extend(font_list)
        plt.rcParams['font.family']= 'sans-serif'
        plt.rcParams['font.sans-serif'] = ['Arial'] # font
        plt.rcParams['legend.handlelength'] = 1
        plt.rcParams['legend.handleheight'] = 1
        fig_len = len(length_data)/3
        fig, ax = plt.subplots(figsize=(fig_len, fig_len))
        trans = ax.transData.transform
        trans2 = fig.transFigure.inverted().transform
        DG = nx.DiGraph()

        # construct network by adding nodes and edges
        style_list = []
        for i, j in enumerate(network_data):
            if i < original_network_size:
                DG.add_edge(j[0], j[1])
                style_list.append('solid')
            else:
                DG.add_edge(j[0], j[1])
                if alt_type != 'dashed':
                    style_list.append('dashed')
                else:
                    style_list.append('solid')
        # construct network by adding node of length
        for i in length_data:
            DG.add_edge(i[0], i[1])
            style_list.append('solid')

        try:
            import pygraphviz as pg
            pos = nx.nx_agraph.graphviz_layout(DG, prog='neato')
        except ImportError:
            pre_pos = nx.spring_layout(DG)
            pos = {}
            for i in pre_pos.keys():
                pos[i] = trans(pre_pos[i])

        return length_data, DG, pos, hap_size, hap_freq, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len

    # plot a tree and replace nodes to pies
    def plot_network(self, length_data, DG, pos, node_num, median_size_rate, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len, label):

        import warnings
        warnings.filterwarnings('ignore')

        # set figure parameters
        # mdian haplotype size
        base_size = 5*20
        median_size = base_size * median_size_rate
        # all node size
        hap_num = len(hap_size_per)
        max_size = max(hap_size_per)
        size_rate = base_size/max_size
        size = [hap_size_per[i]*size_rate if i < hap_num else median_size for i in DG.nodes][:node_num]
        size = size + [0.2] * length_data.size

        try:
            # color and line type
            if hap_freq_per:
                from matplotlib.patches import Patch
                from FilesIO import set_color_num
                nx.draw(DG, pos=pos, arrows=False, node_color='black', node_size=min(size), style=style_list)
                # set colors
                colors = set_color_num(len(pops))
                legend_elements = []
                for color, pop in zip(colors , pops):
                    legend_elements.append(Patch(facecolor=color, edgecolor='black', label=pop))
                plt.legend(handles=legend_elements, frameon=False, bbox_to_anchor=(-0.04, 0.5), prop={'size': fig_len})
                pie_size = 0.025
                move = pie_size/2.0
                for i in DG:
                    if i < hap_num:
                        xx, yy = trans(pos[i])
                        xa, ya = trans2((xx,yy))
                        pie_ax = plt.axes([xa-move, ya-move, pie_size, pie_size])
                        pie_ax.set_aspect('equal')
                        fracs = hap_freq_per[i]
                        if 1 in fracs:
                            one_colors= [colors[fracs.index(1)]]
                            one_fracs = [1]
                            pie_ax.pie(one_fracs, colors=one_colors, radius=(1/(max_size/hap_size_per[i]))**0.5, wedgeprops={'edgecolor':'black','linewidth': 0.5, 'linestyle': 'solid'})
                        else:
                            pie_ax.pie(fracs, colors=colors, radius=(1/(max_size/hap_size_per[i]))**0.5, wedgeprops={'edgecolor':'black','linewidth': 0.5, 'linestyle': 'solid'})
            else:
                colors = ['yellow' if i < hap_num else 'black' for i in DG.nodes] # color
                nx.draw(DG, pos=pos, arrows=False, node_color=colors, node_size=size, style=style_list, linewidths=1)
                node_ax = plt.gca() # to get the current axis
                node_ax.collections[0].set_edgecolor('black')
                node_ax.collections[0].set_linewidth(0.5)

            # output
            if label:
                plt.savefig(self.path+'.net.fig.%s.pdf' % label, dpi=100, bbox_inches='tight')
            else:
                plt.savefig(self.path+'.net.fig.pdf', dpi=100, bbox_inches='tight')
            plt.close()
        except AttributeError:
            print('[Y-LineageTracker] Cannot generate figure since matplotlib.cbook version bugs')

    # main function of network analysis
    def calculate_and_output(self, haplotype_data, population_data, arguments, label):

        # convert data to 0-1 code and calculate distance matrix
        sequence_df, STR_code_pos = self._convert_allel_to_binary(haplotype_data, arguments.filter)
        distance_matrix, sp_sequence_df, sample_sp_code = self._calculate_distance_matrix(sequence_df)
        # calculate network
        self.logger.info('[Y-LineageTracker] Calculating median haplotypes')
        network_data, alt_data, sp_sequence_df2 = self._calculate_network(distance_matrix, sp_sequence_df, arguments.treetype, arguments.search)
        plot_network_data, node_num = self._get_network_for_plot(network_data, alt_data, sp_sequence_df.index.size, arguments.alt)

        # visulization
        self.logger.info('[Y-LineageTracker] Start to plot network')

        if label != 'Group':
            length_data, DG, pos, hap_size, hap_freq, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len = self._get_figure_parameters(sequence_df, population_data, plot_network_data, alt_data, node_num, network_data.size, arguments.alt, arguments.euqalsize, label)
            if arguments.fdi:
                self._get_fdi_file(plot_network_data, plot_network_data.size, node_num, sequence_df, sp_sequence_df2, sample_sp_code, pos, hap_size, hap_freq, STR_code_pos, pops)
            self.plot_network(length_data, DG, pos, node_num, arguments.mediansize, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len, label)
        else:
            label = ['Population', 'Group']
            for i in label:
                length_data, DG, pos, hap_size, hap_freq, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len = self._get_figure_parameters(sequence_df, population_data, plot_network_data, alt_data, node_num, network_data.size, arguments.alt, arguments.euqalsize, label)
                if arguments.fdi:
                    self._get_fdi_file(plot_network_data, plot_network_data.size, node_num, sequence_df, sp_sequence_df2, sample_sp_code, pos, hap_size, hap_freq, STR_code_pos, pops)
                self.plot_network(length_data, DG, pos, node_num, arguments.mediansize, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len, label)

        self.logger.info('[Y-LineageTracker] Network analysis finished')

def main():

    start = time.perf_counter()
    arguments = net_parser()

    from FilesIO import check_seq_matrix_input, check_hap_input, set_haplotype_index, check_population, check_overlap, check_pop_header, time_count
    # check input files
    input, format = check_seq_matrix_input(arguments.seq, arguments.format, arguments.matrix)

    # check haplotype data
    haplotype_data = check_hap_input(input, 'haplotype', format)

    # check filter argument
    check_misisng_rate(arguments.filter)

    from FilesIO import get_out_path
    # get path of output files
    path = get_out_path(arguments.output, 'Net')
    log_file = path + '.NetLog.log'

    # set log file
    set_log(arguments, log_file)

    # check if there is a population file
    if arguments.population:
        population_data = check_population(arguments.population)
        population_data = check_pop_header(population_data)
        haplotype_data, population_data = check_overlap(haplotype_data, population_data)
        label = 'Population'
        if 'Group' in population_data.columns:
            label = 'Group'
    else:
        population_data = None
        label = ''

    # set index of haplotype data
    haplotype_data = set_haplotype_index(haplotype_data, format)

    # calculate network and visualize result
    net = HaploNetwork(arguments.gap, path)
    net.calculate_and_output(haplotype_data, population_data, arguments, label)

    time_count(start)

if __name__ == '__main__':
    main()


T = int(input())
arr = list(map(int, input().split()))
count = 0

for i in range(T) :
  n = 2
  if (arr[i] > 1) :
    count += 1
  num_by = int(arr[i] ** 0.5)
  while (n <= num_by) :
    if (arr[i] % n == 0) :
      count -= 1
      break
    n += 1
print(count)

Solved/코딩 테스트 준비 - 기초/수학/소수찾기.py


'''
Use of this source code is governed by a MIT-style license that can be found in the LICENSE file.
Created on Jul 6, 2017
@author: Niels Lubbes
'''

from linear_series.class_poly_ring import PolyRing
from linear_series.class_linear_series import LinearSeries

class TestClassLinearSeries:

    def test__subs( self ):

        ring = PolyRing( 'x,y,z', True )
        ring.ext_num_field( 't^2 + 1' )
        ring.ext_num_field( 't^3 + a0' )

        a0, a1 = ring.root_gens()
        x, y, z = ring.gens()

        pol_lst = [x ** 2 + a0 * y * z, y + a1 * z + x ]
        ls = LinearSeries( pol_lst, ring )
        assert str( ls.pol_lst ) == '[x^2 + a0*y*z, x + y + a1*z]'

        ls.subs( {x: x + 1} )
        assert str( ls.pol_lst ) == '[x^2 + a0*y*z + 2*x + 1, x + y + a1*z + 1]'

    def test__chart( self ):

        ring = PolyRing( 'x,y,z', True )
        ring.ext_num_field( 't^2 + 1' )
        ring.ext_num_field( 't^3 + a0' )

        ls = LinearSeries( ['x^2+a0*y*z', 'x+y+a1*z '], ring )
        xls = ls.copy().chart( ['x'] )

        assert str( xls.pol_lst ) == '[a0*y*z + 1, y + a1*z + 1]'

    def test__diff( self ):

        ring = PolyRing( 'x,y,z', True )
        ring.ext_num_field( 't^2 + 1' )
        ring.ext_num_field( 't^3 + a0' )

        ls = LinearSeries( ['a0*x^2+y^3*z^3', 'x+z+z^3*y^4'], ring )
        xls = ls.copy().chart( ['x'] )

        assert str( xls.diff( 3, 2 ).pol_lst ) == '[36*z, 144*y*z]'

    def test__translate_to_origin( self ):

        ring = PolyRing( 'x,y,z', True )
        ring.ext_num_field( 't^2 + 1' )
        ring.ext_num_field( 't^3 + a0' )

        ls = LinearSeries( ['y^2', 'y*x'], ring )
        xls = ls.copy().chart( ['x'] )

        a0, a1 = ring.coerce( 'a0,a1' )

        xls.translate_to_origin( ( a0 + 1, a1 + a0 ) )
        assert str( xls ) == '{ 2, <<y^2 + (2*a0 + 2)*y + 2*a0, y + a0 + 1>>, QQ( <a0|t^2 + 1>, <a1|t^2 + a0*t - 1> )[y, z] }'

    def test__blow_up_origin( self ):

        ring = PolyRing( 'x,y,z', True )
        ring.ext_num_field( 't^2 + 1' )
        ring.ext_num_field( 't^3 + a0' )

        ls = LinearSeries( ['x^2', 'x*z + y^2'], ring )
        zls = ls.copy().chart( ['z'] )

        zls, m = zls.blow_up_origin( 's' )
        assert str( zls ) == '{ 2, <<x, x*y^2 + 1>>, QQ( <a0|t^2 + 1>, <a1|t^2 + a0*t - 1> )[x, y] }'
        assert m == 1

    def test__quo( self ):

        ring = PolyRing( 'x,y,z', True )
        ring.ext_num_field( 't^2 + 1' )
        ring.ext_num_field( 't^3 + a0' )

        ls = LinearSeries( ['x^2', 'x + x^2'], ring )
        assert str( ls.quo( 'x' ) ) == '{ 2, <<x, x + 1>>, QQ( <a0|t^2 + 1>, <a1|t^2 + a0*t - 1> )[x, y, z] }'

if __name__ == '__main__':

    # TestClassLinearSeries().test__subs()
    # TestClassLinearSeries().test__chart()
    # TestClassLinearSeries().test__diff()
    # TestClassLinearSeries().test__translate_to_origin()
    # TestClassLinearSeries().test__blow_up_origin()
    # TestClassLinearSeries().test__quo()
    pass


import sqlite3
db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
sql=db.cursor()
db.row_factory=sqlite3.Row
# вставлять везде этот обработчик ошибок
# db = None
# try:
# except sqlite3.Error as e:
#     if db: db.rollback()
#     print (e)
# finally:
#     db.commit()
#     if db: db.close()

#функция проверки зарегистрирован ли ползователь
def check_registration(username):
    db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
    sql=db.cursor()
    db = None
    try:
        sql.execute(f"SELECT name FROM worker WHERE username == '{username}'")
        from_worker = sql.fetchone()
        sql.execute(f"SELECT name FROM client WHERE username == '{username}'")
        from_client = sql.fetchone()
        if from_worker != None or from_client != None:
            return True
        else:
            return False
    except sqlite3.Error as e:
        if db: db.rollback()
        print (e)
    finally:
        if db: db.close()

#функция проверки роли пользователя
def check_role(username):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT 'worker' as source, username FROM worker WHERE username='{username}' UNION ALL SELECT 'client' as source, username FROM client WHERE username='{username}'")
        rows = sql.fetchone()
        # print(rows)
        return rows[0]
    except:
        return 'Пользователь не найден'

#запись работника в бд
def reg_worker(date, username):
    db = None

    worker_information = date[username]
    username = worker_information['username'].replace(' ', '_')
    first_name = worker_information['first_name'].replace(' ', '_')
    resume = worker_information['resume'].replace(' ', '_')
    phone = worker_information['phone'].replace(' ', '_')
    age = worker_information['age']
    chat_id = worker_information['chat_id']
    if resume == '/next': #обработка пропуска
        resume = 'Не указано'
    try:
        db = sqlite3.connect('DataBases/workers.db')
        sql = db.cursor()

        sql.execute(f"INSERT INTO worker (username , name , specialization, phone_number, age, rate, chat_id) VALUES('{username}', '{first_name}','{resume}','{phone}', '{age}', '{0}', '{chat_id}')")
        sql.execute(f"UPDATE main_info SET workers_count = workers_count + 1")

    except sqlite3.Error as e:
        if db: db.rollback()
        print (e)

    finally:
        db.commit()
        if db: db.close()

def reg_client(client_info,username):               #функция записи заказчика в базу данных, принимает словарь и имя пользователя
    username=client_info[username]['username']
    chat_id=client_info[username]['chat_id']
    phone_number=client_info[username]['phone_number']
    name=client_info[username]['name']
    try:
        sql.execute(f"INSERT INTO client(name,phone_number,chat_id,username) VALUES('{name}','{phone_number}','{chat_id}','{username}')")
        sql.execute(f"UPDATE main_info SET client_count = client_count + 1")
    except sqlite3.Error as e:
        if db: db.rollback()
        print (e)
    finally:
        db.commit()
        db.close()

#изменение имени у работника
def worker_change_name_bd(username, new_username):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"UPDATE worker SET name = '{new_username}' WHERE username = '{username}'")
        return f'Успешно, ваше имя измененно на {new_username}'
    except sqlite3.Error as e:
        if db: db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()
#изменение описания у работника
def worker_change_description_bd(username, new_description):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"UPDATE worker SET specialization = '{new_description}' WHERE username = '{username}'")
        return f'Успешно, ваше описание измененно'
    except sqlite3.Error as e:
        if db: db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()
#изменение номера телефона у работника
def worker_change_phone_bd(username, new_phone_namber):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"UPDATE worker SET phone_number = '{new_phone_namber}' WHERE username = '{username}'")
        return f'Успешно, ваш номер успешно изменен на {new_phone_namber}'
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()
#изменение возраста у работника
def worker_change_age_bd(username, new_age):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"UPDATE worker SET age = '{new_age}' WHERE username = '{username}'")
        return f'Успешно, ваш возраст успешно изменен на{new_age}'
    except sqlite3.Error as e:
        if db: db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

#изменение профиля у заказчика
#изменение имени у заказчика
def client_change_name_bd(username, new_username):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"UPDATE client SET name = '{new_username}' WHERE username = '{username}'")
        print ('name(1)')
        return f'Успешно, ваше имя измененно на {new_username}'
    except sqlite3.Error as e:
        if db: db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()
#изменение номера телефона у заказчика
def client_change_phone_bd(username, new_phone_namber):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"UPDATE client SET phone_number = '{new_phone_namber}' WHERE username = '{username}'")
        db.commit()
        return f'Успешно, ваш номер успешно изменен на {new_phone_namber}'
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

def all_vacancy_find_work_db():
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT id, client, title, description, adress, workers_count, age, price FROM 'order' WHERE status = '1'") #достаю все записи из двух столбцов
        all_vacancy = sql.fetchall()
        return all_vacancy
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

# print(all_vacancy_find_work_db())

#Добавление записи в таблицу черновая работа
def add_work_black(work_dict,username):               #функция записи заказчика в базу данных, принимает словарь и имя пользователя
    # username=client_info[username]['username']
    # chat_id=client_info[username]['chat_id']
    # phone_number=client_info[username]['phone_number']
    # name=client_info[username]['name']
    username=username
    type_work=title=work_dict[username]['type']
    title=work_dict[username]['title']
    description=work_dict[username]['description']
    adress=work_dict[username]['adress']
    count_workers=work_dict[username]['workers_count']
    age=work_dict[username]['age']
    price=work_dict[username]['price']

    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT id FROM client WHERE username=='{username}'")
        client=sql.fetchone()[0]
        print(client)

        sql.execute(f"INSERT INTO 'order' (client,title,description,adress,workers_count,age,price,type,status) VALUES('{client}','{title}','{description}','{adress}',{count_workers},'{age}','{price}','{type_work}',0)")
        sql.execute(f"UPDATE main_info SET orders_count = orders_count + 1")
        sql.execute(f"UPDATE main_info SET orders_sum = orders_sum + {price}")
        last_id = sql.execute('SELECT last_insert_rowid()').fetchone()[0]
        return last_id
    except sqlite3.Error as e:
        if db: db.rollback()
        print (e)
    finally:
        db.commit()
        db.close()

# add_work_black({'Hzorhz': {'type': 'Почасовая', 'title': 'название', 'description': 'Описание', 'adress': 'адрес', 'workers_count': 6, 'age': 12, 'price': 1789}},'Hzorhz')

def profile_worker_db(username):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT name, specialization, phone_number, age FROM worker WHERE username = '{username}'")
        date_worker = sql.fetchone()
        print(date_worker)
        return date_worker
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()
profile_worker_db('Alexei0212022')

def profile_client_db(username):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT name, phone_number FROM client WHERE username = '{username}'")
        date_client = sql.fetchone()
        return date_client
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

# print (profile_client_db('Alexei0212022'))

#достает chatid заказчика для пересылки ему отклика работника
def find_chatid_client_for_worker_db(id_order):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT client FROM 'order' WHERE id = '{id_order}'")
        client_id = sql.fetchone()[0]
        sql.execute(f'SELECT chat_id FROM client WHERE id = {client_id}')
        chatid_client = sql.fetchone()[0]
        return chatid_client
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

        db.close()

#Функция, которая возвращает чат айди админа

def select_admin_chat_id():
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT chat_id FROM admin")
        admin_chat_id = sql.fetchone()[0]

        return admin_chat_id
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

        db.close()
#Функция одобрения работы и перевода ее в активное состояние. Также возвращает чат айди клиента, чтобы отправить ему сообщение о том, что заявка опубликована
def accept_work_db(order_id):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"UPDATE 'order' SET status = 1 WHERE id = {order_id}")
        client_id=sql.execute(f"SELECT client FROM 'order' WHERE id = {order_id}").fetchone()[0]

        chat_id=sql.execute(f"SELECT chat_id FROM 'client' WHERE id = {client_id}").fetchone()[0]

        return chat_id
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

#Функция удаления работы из бд, если админ отклонил ее
def reject_work_db(order_id):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()

        client_id=sql.execute(f"SELECT client FROM 'order' WHERE id = {order_id}").fetchone()[0]

        chat_id=sql.execute(f"SELECT chat_id FROM 'client' WHERE id = {client_id}").fetchone()[0]

        sql.execute(f"DELETE FROM 'order' WHERE id ={order_id}")
        return chat_id
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

def select_worker_chat_id(username):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT chat_id FROM worker WHERE username = '{username}'")
        worker_chatid = sql.fetchone()[0]
        return worker_chatid
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

def full_info_order(order_id):
    try:
        db=sqlite3.connect('DataBases/workers.db',check_same_thread=False)
        sql=db.cursor()
        sql.execute(f"SELECT * FROM 'order' WHERE id = {order_id}")
        order_info = sql.fetchone()
        return order_info
    except sqlite3.Error as e:
        db.rollback()
        print (e)
        return 'Упс, что то пошло не так'
    finally:
        db.commit()
        db.close()

print(full_info_order(1))


"""
# 110
Easy
Balanced Binary Tree

https://leetcode.com/problems/balanced-binary-tree/?envType=daily-question&envId=2023-09-01

Given a binary tree, determine if it is height-balanced.

** height-balanced **
Every node's left sub-tree and right sub-tree differ not more than one

Example 1:
Input: root = [3,9,20,null,null,15,7]
Output: true

Example 2:
Input: root = [1,2,2,3,3,null,null,4,4]
Output: false

Example 3:
Input: root = []
Output: true

Constraints:
The number of nodes in the tree is in the range [0, 5000].
-10^4 <= Node.val <= 10^4
"""

from typing import Optional
from data_structure.tree import TreeNode, list_to_binary_tree

def is_balanced(root: Optional[TreeNode]) -> bool:
    def depth(root: TreeNode):
        if not root:
            return 0

        left_depth = depth(root.left)
        if left_depth < 0:
            return -1

        right_depth = depth(root.right)
        if right_depth < 0 or abs(left_depth - right_depth) > 1:
            return -1

        return max(left_depth, right_depth) + 1

    return depth(root) >= 0

if __name__ == "__main__":
    root = list_to_binary_tree([3, 9, 20, None, None, 15, 7])
    assert is_balanced(root)

    root = list_to_binary_tree([1, 2, 2, 3, 3, None, None, 4, 4])
    assert not is_balanced(root)

    assert is_balanced(None)

    root = list_to_binary_tree([1, 2, 2, 3, None, None, 3, 4, None, None, 4])
    assert not is_balanced(root)


from sys import argv

from constants import transitions, constantsHL
from utils import to_decimal, to_xlsx, to_pretty

def setup_matrix(seq: str) -> list:
    """
    Set up simple 2d matrix with 2 rows and len(seq) columns
    """
    return [[0 for _ in range(len(seq))] for _ in range(2)]

def calculate(const_hl: dict, transitions_: dict, matrix: list, seq: str) -> tuple:
    """
    Calculates matrix, HL sequence and result probability
    """
    best = []
    result_value = None
    for idx, elem in enumerate(seq):
        if idx == 0:
            h_value = transitions_['start'] + const_hl['H'][elem]
            l_value = transitions_['start'] + const_hl['L'][elem]
        else:
            h_value = const_hl['H'][elem] + max(
                matrix[0][idx - 1] + transitions_['HH'],
                matrix[1][idx - 1] + transitions_['LH']
            )
            l_value = const_hl['L'][elem] + max(
                matrix[0][idx - 1] + transitions_['HL'],
                matrix[1][idx - 1] + transitions_['LL']
            )
        matrix[0][idx] = h_value
        matrix[1][idx] = l_value
        if to_decimal(h_value) > to_decimal(l_value):
            best.append('H')
        elif to_decimal(h_value) == to_decimal(l_value):
            best.append(best[-1])
        else:
            best.append('L')

        result_value = max(h_value, l_value)

    return matrix, best, 2 ** result_value

def main():
    try:
        sequence = argv[1]
    except IndexError:
        sequence = input('Input sequence: ')

    matrix = setup_matrix(sequence)
    calculated, best, result = calculate(constantsHL, transitions, matrix, sequence)
    pretty = to_pretty(calculated)
    file = to_xlsx(pretty, sequence, best, result)
    print(f'Successfully calculated! Look for the result in {file}')

if __name__ == '__main__':
    main()


import pytest
from runtool.datatypes import (
    Algorithm,
    Algorithms,
    Dataset,
    Datasets,
    Experiment,
    Experiments,
)
from runtool.recurse_config import Versions

ALGORITHM = Algorithm(
    {
        "image": "012345678901.dkr.ecr.eu-west-1.amazonaws.com/gluonts/cpu:latest",
        "instance": "ml.m5.xlarge",
        "hyperparameters": {
            "prediction_length": 7,
            "freq": "D",
        },
    }
)

DATASET = Dataset(
    {
        "path": {
            "train": "s3://gluonts-run-tool/gluon_ts_datasets/constant/train/data.json",
            "test": "s3://gluonts-run-tool/gluon_ts_datasets/constant/test/data.json",
        }
    }
)

EXPERIMENT = Experiment.from_nodes(ALGORITHM, DATASET)
ALGORITHMS = lambda num: Algorithms([ALGORITHM] * num)
DATASETS = lambda num: Datasets([DATASET] * num)
EXPERIMENTS = lambda num: Experiments([EXPERIMENT] * num)
VERSIONS = lambda item, num: Versions([item] * num)

def perform_typeerror_test(data, to_multiply):
    for item in to_multiply:
        with pytest.raises(TypeError):
            data * item
        with pytest.raises(TypeError):
            item * data

def test_versions_algorithm_mul_versions_dataset():
    assert VERSIONS(ALGORITHM, 1) * VERSIONS(DATASET, 1) == EXPERIMENTS(1)
    assert VERSIONS(DATASET, 1) * VERSIONS(ALGORITHM, 1) == EXPERIMENTS(1)

def test_versions_algorithm_mul_versions_datasets():
    assert VERSIONS(ALGORITHM, 1) * VERSIONS(DATASETS(4), 1) == EXPERIMENTS(4)
    assert VERSIONS(DATASETS(4), 1) * VERSIONS(ALGORITHM, 1) == EXPERIMENTS(4)

def test_versions_algorithms_mul_versions_dataset():
    assert VERSIONS(ALGORITHMS(4), 1) * VERSIONS(DATASET, 1) == EXPERIMENTS(4)
    assert VERSIONS(DATASET, 1) * VERSIONS(ALGORITHMS(4), 1) == EXPERIMENTS(4)

def test_versions_algorithms_mul_versions_datasets():
    assert VERSIONS(ALGORITHMS(4), 1) * VERSIONS(
        DATASETS(4), 1
    ) == EXPERIMENTS(16)
    assert VERSIONS(DATASETS(4), 1) * VERSIONS(
        ALGORITHMS(4), 1
    ) == EXPERIMENTS(16)

def test_versions_int_mul_versions_datasets():
    perform_typeerror_test(VERSIONS(DATASET, 1), VERSIONS(1, 3))

def test_versions_str_mul_versions_datasets():
    perform_typeerror_test(VERSIONS(DATASET, 1), VERSIONS("hello", 3))


"""Datasets based on accesing indexable list-like objects"""
import tables

class MultiIndexableDataset:
    def __init__(self,arrays,transforms=None):
        self.arrays = arrays
        self.transforms = transforms

    def __getitem__(self, item):
        if self.transforms is None:
            return tuple([array[item] for array in self.arrays])

        return tuple(array[item] if self.transforms[item] is None
                      else self.transforms[item](array[item])
                      for array in self.arrays)

class MultiHDF5TablesDataset(MultiIndexableDataset):
    def __init__(self, files, array_nodes, load_in_memory=None, transforms=None):
        assert len(files) == len(array_nodes)

        if load_in_memory is None:
            self.mem = (False,)*len(files)
        else:
            self.mem = load_in_memory
        assert len(self.mem) == len(files)

        self.files = []
        arrays = []
        for i,file in enumerate(files):
            newfile = tables.open_file(file,mode="r")
            self.files.append(newfile)
            newarray = newfile.get_node(array_nodes[i])
            assert isinstance(newarray,tables.array.Array), \
                "Node {} in file {} is not an array".format(array_nodes[i],file)
            arrays.append(newarray)

        assert len(set(len(a) for a in arrays)) == 1, "All arrays must be of the same length"

        super(MultiHDF5TablesDataset, self).__init__(arrays, transforms=transforms)

    def __len__(self):
        return len(self.arrays[0])

    def close_files(self):
        for file in self.files:
            file.close()


import tvm
import numpy as np
from tvm import tg
from tvm.tensor_graph.core import compute, GraphOp, GraphTensor, ForwardGraph, \
  make_fwd_graph, make_tir_graph
from tvm.tensor_graph.testing.models import resnet
from tvm.tensor_graph.nn import CELoss, SGD

def zero_pad2d(inputs, padding=0):
  """
  Zero padding for 2d tensor

  Args:
  -----------------------------
  inputs : GraphNode
      shape [batch, channel, height, width]
  padding: (optional:0) int or tuple
      expected: (h_pad_up, h_pad_down, w_pad_up, w_pad_down)
  -----------------------------

  Returns:
  -----------------------------
  GraphOp
      shape [batch, channel, padded_height, padded_width]
  -----------------------------
  """
  padding = (padding, padding, padding, padding) if isinstance(padding, (int, tvm.tir.IntImm)) else padding
  assert isinstance(padding, tuple), "type(padding)={}".format(type(padding))
  if len(padding) == 2:
    padding = (padding[0], padding[0], padding[1], padding[1])
  assert (len(padding) == 4)
  batch_size, in_channel, height, width = inputs.shape
  padded_shape = (batch_size, in_channel, height + padding[0] + padding[1], width + padding[2] + padding[3])
  def _inner_zero_pad2d(batch_size, in_channel, h, w, inputs, requires_grad=True):
    # Warning, we use "float32" as type of 0
    padding_zero = tvm.tir.expr.const(0, "float32")
    return compute(padded_shape,
      lambda b, c, h, w: tvm.te.if_then_else(
            tvm.te.all(h >= padding[0], h < height + padding[0], w >= padding[2], w < width + padding[2]),
            inputs[b, c, h - padding[0], w - padding[2]],
            padding_zero
            ),
      name="zero_pad2d",
      tag="TG_AUTOGEN",
      requires_grad=requires_grad)
  return GraphOp(padded_shape, [] , [inputs], _inner_zero_pad2d, name="zero_pad2d")

def conv2d_nchw(inputs, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
  """Convolution 2d NCHW layout

  Args:
  -----------------------------
  inputs  : GraphNode
      shape [batch, channel, height, width]
  weight  : GraphNode
      shape [out_channel, channel // groups, kernel_height, kernel_width]
  bias    : (optional:None) GraphNode
      shape [out_channel]
  stride  : (optional:1) int or tuple

  padding : (optional:0) int or tuple

  dilation: (optional:1) int

  groups  : (optional:1) int
  -----------------------------

  Returns:
  -----------------------------
  GraphOp
      shape [batch, out_channel, output_height, output_width]
  -----------------------------
  """
  batch_size, in_channel, in_h, in_w = inputs.shape
  out_channel, channel_per_group, k_h, k_w = weight.shape
  assert channel_per_group * groups == in_channel
  out_channel_per_group = out_channel // groups
  assert out_channel_per_group * groups == out_channel

  stride = (stride, stride) if isinstance(stride, (int, tvm.tir.IntImm)) else stride
  padding = (padding, padding) if isinstance(padding, (int, tvm.tir.IntImm)) else padding
  dilation = (dilation, dilation) if isinstance(dilation, (int, tvm.tir.IntImm)) else dilation
  assert isinstance(stride, tuple) and len(stride) == 2
  assert isinstance(padding, tuple) and len(padding) == 2
  assert isinstance(dilation, tuple) and len(dilation) == 2

  out_h = (in_h + 2 * padding[0] - dilation[0] * (k_h - 1) - 1) // stride[0] + 1
  out_w = (in_w + 2 * padding[1] - dilation[1] * (k_w - 1) - 1) // stride[1] + 1

  padded = zero_pad2d(inputs, padding=padding)
  conv_out_shape = (batch_size, out_channel, out_h, out_w)
  def _inner_conv2d_nchw(batch_size, out_channel, out_h, out_w, channel_per_group,
                            k_w, k_h, padded, weight, requires_grad=True):
    rc = tvm.te.reduce_axis((0, channel_per_group), name="rc")
    rw = tvm.te.reduce_axis((0, k_w), name="rw")
    rh = tvm.te.reduce_axis((0, k_h), name="rh")
    return compute(conv_out_shape,
        lambda b, c, h, w: tvm.te.sum(
          (padded[b, c // out_channel_per_group * channel_per_group + rc,
                h * stride[0] + rh * dilation[0], w * stride[1] + rw * dilation[1]]
          * weight[c, rc, rh, rw]),
          axis=[rc, rw, rh]),
          name="conv2d_nchw",
          tag="TG_AUTOGEN",
          requires_grad=requires_grad
        )
  conv_out = GraphOp(
    conv_out_shape,
    [channel_per_group, k_w, k_h],
    [padded, weight],
    _inner_conv2d_nchw,
    name="conv2d_nchw")
  def _inner_bias(batch_size, out_channel, out_h, out_w, conv_out, bias, requires_grad=True):
    return compute(
      (batch_size, out_channel, out_h, out_w),
      lambda b, c, h, w: conv_out[b, c, h, w] + bias[c],
      name="conv2d_nchw_bias",
      tag="TG_AUTOGEN",
      requires_grad=requires_grad
      )
  if bias is not None:
    return GraphOp(conv_out_shape, [], [conv_out, bias], _inner_bias, name="conv2d_bias")
  return conv_out

def test1():
  print("test 1 ##############################")
  A = GraphTensor([1, 1024, 14, 14])
  W = GraphTensor([2048, 1024, 3, 3])
  B = conv2d_nchw(A, W, bias=None, stride=1, padding=1)

  B_tensor, params = B({})
  tvm_tensor = B_tensor.tvm_tensor
  print(tvm_tensor.op.tag)
  print(tvm_tensor.op.body)
  print(tvm_tensor.op.input_tensors[0].op.tag)

  print("Success!")

def test2():
  print("test 2 ##############################")
  batch = 64
  img_shape = [batch, 3, 224, 224]
  num_classes = 1000
  label_shape = [batch, num_classes]
  dtype = "float32"
  model = resnet.resnet50(num_classes=1000)
  img_tensor = GraphTensor(img_shape, dtype=dtype, name="image")
  label_tensor = GraphTensor(label_shape, dtype=dtype, name="label")

  # this is data
  img_np = np.random.uniform(-1, 1, img_shape).astype(dtype)
  label_np = np.random.uniform(-1, 1, [batch, num_classes]).astype(dtype)

  ce_loss = CELoss(label_tensor)
  sgd = SGD(0.002)

  fwd_graph = make_fwd_graph(model, [img_tensor])
  tir_graph = make_tir_graph(fwd_graph, inference=True)

  op_map = {}
  for op in tir_graph.operation_list:
    if op.tag not in op_map:
      op_map[op.tag] = []
    op_map[op.tag].append(op)

  print("Totally", len(op_map), "different tags.")
  for k, v in op_map.items():
    print("Tag:", k, "has", len(v), "operators")

  print("Success")

def test3():
  print("test 3 ##############################")
  batch = 64
  img_shape = [batch, 3, 224, 224]
  num_classes = 1000
  label_shape = [batch, num_classes]
  dtype = "float32"
  model = resnet.resnet50(num_classes=1000)
  img_tensor = GraphTensor(img_shape, dtype=dtype, name="image")
  label_tensor = GraphTensor(label_shape, dtype=dtype, name="label")

  # this is data
  img_np = np.random.uniform(-1, 1, img_shape).astype(dtype)
  label_np = np.random.uniform(-1, 1, [batch, num_classes]).astype(dtype)

  ce_loss = CELoss(label_tensor)
  sgd = SGD(0.002)

  fwd_graph = make_fwd_graph(model, [img_tensor])
  tir_graph = make_tir_graph(fwd_graph, loss=ce_loss, optimizer=sgd, inference=False)

  op_map = {}
  for op in tir_graph.operation_list:
    if op.tag not in op_map:
      op_map[op.tag] = []
    op_map[op.tag].append(op)

  with open("trace.log", "w") as fout:
    print("Totally", len(op_map), "different tags.", file=fout)
    for k, v in op_map.items():
      print("Tag:", k, "has", len(v), "operators", file=fout)

  print("Success")

if __name__ == "__main__":
  test1()
  test2()
  test3()


import re
from operator import attrgetter
from Day7.classes import Node

def string_to_dict(string, pattern):
    regex = re.sub(r'{(.+?)}', r'(?P<_\1>.+)', pattern)
    values = list(re.search(regex, string).groups())
    keys = re.findall(r'{(.+?)}', pattern)
    _dict = dict(zip(keys, values))
    return _dict

def tests():
    assert ({"before": "X", "after": "M"} ==
            string_to_dict("Step X must be finished before step M can begin.",
                           "Step {before} must be finished before step {after} can begin."))
    assert (1 == get_index_from_letter("A"))
    assert (26 == get_index_from_letter("Z"))
    assert (61 == get_step_length("A"))
    assert (86 == get_step_length("Z"))
    assert ([Node("0"), Node("1"), Node("3")] == choose_x_available([Node("1"), Node("3"), Node("0")], 5, {}))
    assert ([Node("0"), Node("3")] == choose_x_available([Node("1"), Node("3"), Node("0")], 5, {1: Node("1")}))

def extract_before_and_after_from_string(line):
    word_map = string_to_dict(line, "Step {before} must be finished before step {after} can begin.")
    return word_map["before"], word_map["after"]

def main():
    with open("input.txt") as task_input:
        content_array = task_input.read().split("\n")
        nodes = {}
        for line in content_array:
            before_value, after_value = extract_before_and_after_from_string(line)
            before_node = nodes[before_value] if before_value in nodes else create_node_put_in_map(before_value, nodes)
            after_node = nodes[after_value] if after_value in nodes else create_node_put_in_map(after_value, nodes)
            common_parents = set(before_node.parents).intersection(after_node.parents)
            for common_parent in common_parents:
                after_node.parents.remove(common_parent)
                common_parent.children.remove(after_node)
            before_node.add_child(after_node)
            after_node.add_parent(before_node)
        print_steps(nodes)
        get_parallel_steps(nodes)

def create_node_put_in_map(value, nodes):
    node = Node(value)
    nodes[value] = node
    return node

def is_available(child, prev):
    return all([parent.value in prev for parent in child.parents])

def print_steps(nodes):
    available = [nodes[node] for node in nodes if len(nodes[node].parents) == 0]
    steps = []
    while (len(available)) > 0:
        next_node = choose_next_from_available(available)
        steps.append(next_node.value)
        available.extend(list(filter(lambda child: is_available(child, steps), next_node.children)))
        available.remove(next_node)
    answer = "".join(steps)
    print("Answer to part 1:", answer)
    assert ("ACHOQRXSEKUGMYIWDZLNBFTJVP" == answer)

def choose_x_available(available, how_many, current_tasks):
    sorted_available = list(filter(lambda task: task not in current_tasks.values(), sorted(available, key=lambda node: node.value)))
    return sorted_available[:how_many]

def check_any_finished_tasks(current_tasks, workers_times, time, steps, available):
    if not time > 0:
        return
    for worker_data in current_tasks:
        if not current_tasks[worker_data]:
            continue
        if workers_times[worker_data] <= time:
            finished_task = current_tasks[worker_data]
            steps.append(finished_task.value)
            new_available = list(filter(lambda child: is_available(child, steps), finished_task.children))
            available.extend(new_available)
            available.remove(finished_task)
            current_tasks[worker_data] = None

def get_next_available_step(available, current_tasks):
    next_steps = choose_x_available(available, 1, current_tasks)
    return next_steps[0] if len(next_steps) > 0 else None

def get_parallel_steps(nodes, workers=5):
    available = [nodes[node] for node in nodes if len(nodes[node].parents) == 0]
    steps = []
    time = 0
    workers_times = {worker: 0 for worker in range(1, workers + 1)}
    current_tasks = {worker: None for worker in range(1, workers + 1)}
    while len(available) > 0 or any(current_tasks.values()):
        check_any_finished_tasks(current_tasks, workers_times, time, steps, available)
        next_step = get_next_available_step(available, current_tasks)
        if not next_step:
            time += 1
            continue
        assigned = False
        while not assigned:
            check_any_finished_tasks(current_tasks, workers_times, time, steps, available)
            for worker in workers_times:
                his_time = workers_times[worker]
                if his_time <= time:
                    workers_times[worker] = time + get_step_length(next_step.value)
                    assigned = True
                    current_tasks[worker] = next_step
                    break
            if not assigned:
                time += 1
        # by this time task should be assigned
    answer = max(workers_times.values())
    print("Answer to part 2:", answer)
    assert(985 == answer)

def get_index_from_letter(letter):
    return ord(letter) - 64

def get_step_length(value):
    return 60 + get_index_from_letter(value)

def choose_next_from_available(nodes):
    return min(nodes, key=attrgetter('value'))

if __name__ == "__main__":
    tests()
    main()


from pathlib import Path
import json
from collections import OrderedDict
import html

from .enum import CellType

def complete_path_split(path, relative_to=None):
    components = []

    path = Path(path)
    if relative_to:
        path = path.relative_to(relative_to)

    while path.name:
        components.append(path.name)

        path = path.parent

    return components

def get_files(suffixes, src):
    suffixes = set(s.lower() for s in suffixes)

    for fp in Path(src).glob('**/*.*'):
        if fp.is_file() and fp.suffix.lower() in suffixes \
                and not any(comp[0] in {'.', '_'} for comp in complete_path_split(fp)):
            yield fp

def read_jupyter(fp):
    with Path(fp).open() as f:
        for cell in json.load(f, object_pairs_hook=OrderedDict).get('cells', []):
            if cell['cell_type'] == 'code':
                for output in cell.get('outputs', []):
                    if output['output_type'] == 'display_data':
                        data = output['data']
                        types = data.keys()
                        if CellType.HTML in types:
                            yield ''.join(data[CellType.HTML])
                        elif CellType.PLAIN in types:
                            yield '

'.format(html.escape(''.join(data[CellType.PLAIN])))
                        else:
                            raise TypeError(repr(types))

            elif cell['cell_type'] == CellType.MARKDOWN:
                yield ''.join(cell['source'])


import os
import pickle

import google_auth_oauthlib.flow
import googleapiclient.discovery
from googleapiclient.discovery import Resource
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials

from .apis import Playlist, PlaylistItem, Video, Thumbnail, Comment, CommentThread, \
I18n, VideoCategory, VideoAbuseReportReason, Channel, ChannelSection, Search, Subscription

class Client:
    def __init__(self) -> None:
        self.client: Resource = ...

    def _init_resources(self):
        """Adds resources to the class."""
        self.playlist = Playlist(self.client)
        self.playlist_item = PlaylistItem(self.client)
        self.video = Video(self.client)
        self.thumbnail = Thumbnail(self.client)
        self.comment = Comment(self.client)
        self.comment_thread = CommentThread(self.client)
        self.i18n = I18n(self.client)
        self.video_category = VideoCategory(self.client)
        self.video_abuse_report_reason = VideoAbuseReportReason(self.client)
        self.channel = Channel(self.client)
        self.channel_section = ChannelSection(self.client)
        self.search = Search(self.client)
        self.subscription = Subscription(self.client)

    def _save_creds(self, creds: Credentials, token_store: str):
        """Saves credentials to a token cache."""
        with open(token_store, 'wb') as f:
            print('Saving Credentials for Future Use...')
            pickle.dump(creds, f)

    def _fetch_new_creds(self, client_secrets_file: str, scopes: list[str]):
        """Fetches new credentials."""

        flow = google_auth_oauthlib.flow.InstalledAppFlow.from_client_secrets_file(
        client_secrets_file, scopes)
        credentials = flow.run_console()

        return credentials

    @classmethod
    def from_client_secrets(cls, client_secrets_file: str, scopes: list[str],
                            token_store: str = "token.pickle"):
        """
        Returns a youtube client.

        Parameters:
            `client_secrets_file`: the path to your client secrets file.
            `scopes`: the scopes of this client.
            `token_store`: the file the program will use to store credentials for later use.

        Returns:
            An instance of the Youtube client, with OAuth2 access.
        """

        inst = cls()

        credentials: Credentials = None

        # * gets stored creds
        if os.path.exists(token_store):
            print("Loading credentials from file...")
            with open(token_store, 'rb') as token:
                credentials = pickle.load(token)

        # * Refreshes the access token/fetches new creds
        # * Saves it in token.pickle at the end
        if not credentials or not credentials.valid:
            try:
                print('Refreshing acess token...')
                request = Request()
                credentials.refresh(request)
                print("Access token refreshed!")
            except:
                print("Fetching new credentials...")
                credentials = inst._fetch_new_creds(client_secrets_file, scopes)
                print("New credentails fetched!")
            finally:
                inst._save_creds(credentials, token_store)

        youtube = googleapiclient.discovery.build(
            "youtube", "v3", credentials=credentials)

        inst.client = youtube
        inst._init_resources()
        return inst


#Просте  симетричне випадкове блукання стартує  з 0.
# Зроблено 100 кроків.
# с) Оцінити теоретичну ймовірність того, що останнє попадання в 0 було до 10, 15, 20, 50-го кроку (включно). Для цього:  змоделювати 1000 траєкторій, порахувати емпіричні імовірності і взяти середнє.
# Порівняти з теоретичними ймовірностями.  Які висновкі можна зробити?  А якщо брати 1000 кроків, і разувати, відповідно, ймовірність того, що останнє попадання в 0 було до 100, 150, 200, 500-го кроку (включно)?

# Теоретична ймовірність того, що останнє попадання в 0 станеться на або до певного кроку, дорівнює:

# Для кроків до 10: Ця ймовірність дорівнює 1/10, оскільки блукання має рівну ймовірність рухатися вліво або вправо на кожному кроці, і нарешті виходити на 0 на або до 10-го кроку має 1/10 шанс.

# Для кроків до 15: Тут ця ймовірність дорівнює 1/15 за аналогічною логікою.

# Для кроків до 20: Та ж логіка дає ймовірність 1/20.

# Для кроків до 50: Та ж логіка дає ймовірність 1/50.

#  щоб побачити, як це відображається в емпіричних ймовірностях, потрібно буде змоделювати 1000 траєкторій та порахувати, скільки з них задовольняють вказаним умовам (останнє попадання в 0 на або до певного кроку).

# 1000 траєкторій з 100 кроками кожна та рахує емпіричні ймовірності для різних умов і порівнює їх з теоретичними.
# результат отримується як різниці між емпіричними і теоретичними результатами. чим більше траєкторій ви моделюєте, тим більш точні емпіричні результати ви отримаєте.

import random

# Кількість траєкторій
n_trajectories = 1000
# Кількість кроків в кожній траєкторії
n_steps = 100

# Лічильники для кожної умови
counts = {'10': 0, '15': 0, '20': 0, '50': 0}

for _ in range(n_trajectories):
    current_position = 0
    last_zero_position = None

    for step in range(n_steps):
        current_position += random.choice([-1, 1])

        if current_position == 0:
            last_zero_position = step

    for condition in counts.keys():
        if last_zero_position is not None and last_zero_position <= int(condition):
            counts[condition] += 1

# Обчислення емпіричних ймовірностей
empirical_probabilities = {condition: count / n_trajectories for condition, count in counts.items()}

# Виведення результатів
for condition, empirical_probability in empirical_probabilities.items():
    print(f"Емпірична ймовірність останнього попадання в 0 до {condition}-го кроку: {empirical_probability}")

# Теоретичні ймовірності
theoretical_probabilities = {'10': 1/10, '15': 1/15, '20': 1/20, '50': 1/50}

# Порівняння з теоретичними результатами
for condition in empirical_probabilities.keys():
    print(f"Теоретична ймовірність останнього попадання в 0 до {condition}-го кроку: {theoretical_probabilities[condition]}")

# відповідь
# Емпірична ймовірність останнього попадання в 0 до 10-го кроку: 0.145
# Емпірична ймовірність останнього попадання в 0 до 15-го кроку: 0.211
# Емпірична ймовірність останнього попадання в 0 до 20-го кроку: 0.243
# Емпірична ймовірність останнього попадання в 0 до 50-го кроку: 0.45
# Теоретична ймовірність останнього попадання в 0 до 10-го кроку: 0.1
# Теоретична ймовірність останнього попадання в 0 до 15-го кроку: 0.06666666666666667
# Теоретична ймовірність останнього попадання в 0 до 20-го кроку: 0.05
# Теоретична ймовірність останнього попадання в 0 до 50-го кроку: 0.02


from IPython import embed
import numpy as np

from specklepy.exceptions import SpecklepyTypeError
from specklepy.utils import combine

class Sequence(object):

    def __init__(self, science_files, sky_files, sky_time_stamps, science_time_stamps, source=None, file_path=None,
                 object=None, setup=None):

        if isinstance(science_files, (list, np.ndarray)):
            self.science_files = science_files
        else:
            raise SpecklepyTypeError('Sequence', 'science_files', type(science_files), 'list')

        if isinstance(sky_files, (list, np.ndarray)):
            self.sky_files = sky_files
        else:
            raise SpecklepyTypeError('Sequence', 'sky_files', type(sky_files), 'list')

        if isinstance(science_time_stamps, (list, np.ndarray)):
            self.science_time_stamps = science_time_stamps
        else:
            raise SpecklepyTypeError('Sequence', 'science_time_stamps', type(science_time_stamps), 'list')

        if isinstance(sky_time_stamps, (list, np.ndarray)):
            self.sky_time_stamps = sky_time_stamps
        else:
            raise SpecklepyTypeError('Sequence', 'sky_time_stamps', type(sky_time_stamps), 'list')

        if source is None:
            if len(sky_files) > 0:
                self.source = 'sky'
            else:
                self.source = 'science'
        elif isinstance(source, str):
            self.source = source
        else:
            raise SpecklepyTypeError('Sequence', 'source', type(source), 'str')

        if file_path is None:
            self.file_path = './'
        elif isinstance(file_path, str):
            self.file_path = file_path
        else:
            raise SpecklepyTypeError('Sequence', 'file_path', type(file_path), 'str')

        if object is None:
            self.object = ''
        else:
            self.object = object

        if setup is None:
            self.setup = ''
        else:
            self.setup = setup

    def __str__(self):
        s = str(type(self)) + "\n"
        s += 'Science files:\n'
        for file in self.science_files:
            s += "> {}\n".format(file)
        s += 'Sky files:\n'
        for file in self.sky_files:
            s += "> {}\n".format(file)
        return s

    @property
    def files(self):
        return np.concatenate((self.sky_files, self.science_files))

    @property
    def time_stamps(self):
        return np.concatenate((self.sky_time_stamps, self.science_time_stamps))

    @property
    def n_sky(self):
        return len(self.sky_files)

    @property
    def n_science(self):
        return len(self.science_files)

    def time_stamps_to_offset_seconds(self):

        # Check whether time stamps free from offset (equivalent to float type)
        if isinstance(self.time_stamps[0], float):
            return 0

        # Find first time stamp
        time_stamps = self.time_stamps
        time_stamps.sort()
        start = time_stamps[0]

        # Remove time offset
        self.sky_time_stamps = combine.time_difference(start, self.sky_time_stamps)
        self.science_time_stamps = combine.time_difference(start, self.science_time_stamps)

    def compute_weights(self, time_scale=300):
        """Compute the weights by the individual sky measurements are weighted for a given science frame.

        Args:
            time_scale (float, optional):
                Time scale for sky background variations. This is used as a reference time scale to compare the offsets
                against.

        Returns:
            weights (np.array):
                The array of weights of shape (n_science, n_sky).
        """

        # Assert that times are offset seconds
        self.time_stamps_to_offset_seconds()

        # Compute weights from time differences
        weights = np.zeros((self.n_science, self.n_sky))
        for ii in range(self.n_science):
            delta_t = self.science_time_stamps[ii] - self.sky_time_stamps
            weights[ii] = np.exp(-1/2 * np.square(np.divide(delta_t, time_scale)))

        # # Transform time offsets into weights
        # weights = np.square(weights)

        # Normalize the weights
        for ii in range(self.n_science):
            weights[ii] /= np.sum(weights[ii])

        return weights

    # def make_master_sky(self):
    #
    #     if len(self.sky_files) == 1:
    #         data = fits.getdata(os.path.join(self.file_path, self.sky_files[0]))
    #         if data.ndim == 3:
    #             std = np.std(data, axis=0)
    #             data = np.mean(data, axis=0)
    #         else:
    #             std = np.zeros(data.shape)
    #         self.master_sky = data
    #         self.master_sky_std = std
    #
    #     else:
    #         for index, file in enumerate(self.sky_files):
    #             data = fits.getdata(os.path.join(self.file_path, file))
    #             if data.ndim == 3:
    #                 std = np.std(data, axis=0)
    #                 data = np.mean(data, axis=0)
    #             else:
    #                 std = np.zeros(data.shape)
    #
    #             if index == 0:
    #                 shape = (len(self.sky_files)) + data.shape
    #                 master_sky = np.zeros(shape)
    #                 master_sky_uncertainty = np.zeros(shape)
    #
    #             master_sky[index] = data
    #             master_sky_uncertainty[index] = std
    #
    #         # Collapse cubes
    #         self.master_sky = np.divide(np.sum(np.multiply(master_sky, master_sky_uncertainty), axis=0),
    #                                     np.sum(master_sky_uncertainty, axis=0))
    #         self.master_sky_std = np.sqrt(np.sum(np.square(master_sky_uncertainty), axis=0))
    #
    # def subtract_master_sky(self, save_to=None, filename_prefix=None):
    #
    #     if filename_prefix is None:
    #         filename_prefix = ''
    #
    #     if not hasattr(self, 'master_sky'):
    #         self.make_master_sky()
    #
    #     for index, file in enumerate(self.science_files):
    #         logger.info('Subtracting sky from file {:2}/{:2}) {}'.format(index+1, len(self.science_files), file))
    #         with fits.open(os.path.join(self.file_path, file)) as hdulist:
    #             hdulist[0].data = np.subtract(hdulist[0].data, self.master_sky)
    #             hdulist[0].header.set('SKYSUB', self.time_stamp())
    #             logger.info(f"Saving sky subtracted data to file {os.path.join(save_to, filename_prefix + file)}")
    #             hdulist[0].writeto(os.path.join(save_to, filename_prefix + file))

    # def time_stamp(self):
    #     """Return a time stamp str of format 'YYYYMMDD_HHMMSS'."""
    #     return datetime.now().strftime('%Y%m%d_%H%M%S')