import gzip
import shutil
import math
import numpy
import os
import random
import sys
import struct
import time

import numpy as np
from scipy.sparse import csr_matrix

from urllib.request import urlretrieve

from .dataset_io import (
    xbin_mmap, download_accelerated, download, sanitize,
    knn_result_read, range_result_read, read_sparse_matrix,
    write_sparse_matrix,
)

BASEDIR = "data/"

class Dataset():
    def prepare(self):
        """
        Download and prepare dataset, queries, groundtruth.
        """
        pass
    def get_dataset_fn(self):
        """
        Return filename of dataset file.
        """
        pass
    def get_dataset(self):
        """
        Return memmapped version of the dataset.
        """
        pass
    def get_dataset_iterator(self, bs=512, split=(1, 0)):
        """
        Return iterator over blocks of dataset of size at most 512.
        The split argument takes a pair of integers (n, p) where p = 0..n-1
        The dataset is split in n shards, and the iterator returns only shard #p
        This makes it possible to process the dataset independently from several
        processes / threads.
        """
        pass
    def get_queries(self):
        """
        Return (nq, d) array containing the nq queries.
        """
        pass
    def get_private_queries(self):
        """
        Return (private_nq, d) array containing the private_nq private queries.
        """
        pass
    def get_groundtruth(self, k=None):
        """
        Return (nq, k) array containing groundtruth indices
        for each query."""
        pass

    def search_type(self):
        """
        "knn" or "range" or "knn_filtered"
        """
        pass

    def distance(self):
        """
        "euclidean" or "ip" or "angular"
        """
        pass

    def data_type(self):
        """
        "dense" or "sparse"
        """
        pass

    def default_count(self):
        """ number of neighbors to return """
        return 10

    def short_name(self):
        return f"{self.__class__.__name__}-{self.nb}"

    def __str__(self):
        return (
            f"Dataset {self.__class__.__name__} in dimension {self.d}, with distance {self.distance()}, "
            f"search_type {self.search_type()}, size: Q {self.nq} B {self.nb}")

#############################################################################
# Datasets for the competition
##############################################################################

class DatasetCompetitionFormat(Dataset):
    """
    Dataset in the native competition format, that is able to read the
    files in the https://big-ann-benchmarks.com/ page.
    The constructor should set all fields. The functions below are generic.

    For the 10M versions of the dataset, the database files are downloaded in
    part and stored with a specific suffix. This is to avoid having to maintain
    two versions of the file.
    """

    def prepare(self, skip_data=False, original_size=10**9):
        if not os.path.exists(self.basedir):
            os.makedirs(self.basedir)

        # start with the small ones...
        for fn in [self.qs_fn, self.gt_fn]:
            if fn is None:
                continue
            if fn.startswith("https://"):
                sourceurl = fn
                outfile = os.path.join(self.basedir, fn.split("/")[-1])
            else:
                sourceurl = os.path.join(self.base_url, fn)
                outfile = os.path.join(self.basedir, fn)
            if os.path.exists(outfile):
                print("file %s already exists" % outfile)
                continue
            download(sourceurl, outfile)

        # private qs url
        if self.private_qs_url:
            outfile = os.path.join(self.basedir, self.private_qs_url.split("/")[-1])
            if os.path.exists(outfile):
                print("file %s already exists" % outfile)
            else:
                download(self.private_qs_url, outfile)

        # private gt url
        if self.private_gt_url:
            outfile = os.path.join(self.basedir, self.private_gt_url.split("/")[-1])
            if os.path.exists(outfile):
                print("file %s already exists" % outfile)
            else:
                download(self.private_gt_url, outfile)

        if skip_data:
            return

        fn = self.ds_fn
        sourceurl = os.path.join(self.base_url, fn)
        outfile = os.path.join(self.basedir, fn)
        if os.path.exists(outfile):
            print("file %s already exists" % outfile)
            return
        if self.nb == 10**9:
            download_accelerated(sourceurl, outfile)
        else:
            # download cropped version of file
            file_size = 8 + self.d * self.nb * np.dtype(self.dtype).itemsize
            outfile = outfile + '.crop_nb_%d' % self.nb
            if os.path.exists(outfile):
                print("file %s already exists" % outfile)
                return
            download(sourceurl, outfile, max_size=file_size)
            # then overwrite the header...
            header = np.memmap(outfile, shape=2, dtype='uint32', mode="r+")

            assert header[0] == original_size
            assert header[1] == self.d
            header[0] = self.nb

    def get_dataset_fn(self):
        fn = os.path.join(self.basedir, self.ds_fn)
        if os.path.exists(fn):
            return fn
        if self.nb != 10**9:
            fn += '.crop_nb_%d' % self.nb
            return fn
        else:
            raise RuntimeError("file not found")

    def get_dataset_iterator(self, bs=512, split=(1,0)):
        nsplit, rank = split
        i0, i1 = self.nb * rank // nsplit, self.nb * (rank + 1) // nsplit
        filename = self.get_dataset_fn()
        x = xbin_mmap(filename, dtype=self.dtype, maxn=self.nb)
        assert x.shape == (self.nb, self.d)
        for j0 in range(i0, i1, bs):
            j1 = min(j0 + bs, i1)
            yield sanitize(x[j0:j1])

    def get_data_in_range(self, start, end):
        assert start >= 0
        assert end <= self.nb
        filename = self.get_dataset_fn()
        x = xbin_mmap(filename, dtype=self.dtype, maxn=self.nb)
        return x[start:end]

    def search_type(self):
        return "knn"

    def data_type(self):
        return "dense"

    def get_groundtruth(self, k=None):
        assert self.gt_fn is not None
        fn = self.gt_fn.split("/")[-1]   # in case it's a URL
        assert self.search_type() in ("knn", "knn_filtered")

        I, D = knn_result_read(os.path.join(self.basedir, fn))
        assert I.shape[0] == self.nq
        if k is not None:
            assert k <= 100
            I = I[:, :k]
            D = D[:, :k]
        return I, D

    def get_dataset(self):
        assert self.nb <= 10**7, "dataset too large, use iterator"
        slice = next(self.get_dataset_iterator(bs=self.nb))
        return sanitize(slice)

    def get_queries(self):
        filename = os.path.join(self.basedir, self.qs_fn)
        x = xbin_mmap(filename, dtype=self.dtype)
        assert x.shape == (self.nq, self.d)
        return sanitize(x)

    def get_private_queries(self):
        assert self.private_qs_url is not None
        fn = self.private_qs_url.split("/")[-1]   # in case it's a URL
        filename = os.path.join(self.basedir, fn)
        x = xbin_mmap(filename, dtype=self.dtype)
        assert x.shape == (self.private_nq, self.d)
        return sanitize(x)

    def get_private_groundtruth(self, k=None):
        assert self.private_gt_url is not None
        fn = self.private_gt_url.split("/")[-1]   # in case it's a URL
        assert self.search_type() == "knn"

        I, D = knn_result_read(os.path.join(self.basedir, fn))
        assert I.shape[0] == self.private_nq
        if k is not None:
            assert k <= 100
            I = I[:, :k]
            D = D[:, :k]
        return I, D

subset_url = "https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/"

class SSNPPDataset(DatasetCompetitionFormat):
    def __init__(self, nb_M=1000):
        # assert nb_M in (10, 1000)
        self.nb_M = nb_M
        self.nb = 10**6 * nb_M
        self.d = 256
        self.nq = 100000
        self.dtype = "uint8"
        self.ds_fn = "FB_ssnpp_database.u8bin"
        self.qs_fn = "FB_ssnpp_public_queries.u8bin"
        self.gt_fn = (
            "FB_ssnpp_public_queries_1B_GT.rangeres" if self.nb_M == 1000 else
            subset_url + "GT_100M/ssnpp-100M" if self.nb_M == 100 else
            subset_url + "GT_10M/ssnpp-10M" if self.nb_M == 10 else
            None
        )

        self.base_url = "https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/"
        self.basedir = os.path.join(BASEDIR, "FB_ssnpp")

        self.private_nq = 100000
        self.private_qs_url = "https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/FB_ssnpp_heldout_queries_3307fba121460a56.u8bin"
        self.private_gt_url = "https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/GT_1B_final_2bf4748c7817/FB_ssnpp.bin"

    def search_type(self):
        return "range"

    def default_count(self):
        """ for range search, this returns the squared range search radius """
        return 96237

    def distance(self):
        return "euclidean"

    def get_groundtruth(self, k=None):
        """ override the ground-truth function as this is the only range search dataset """
        assert self.gt_fn is not None
        fn = self.gt_fn.split("/")[-1]   # in case it's a URL
        return range_result_read(os.path.join(self.basedir, fn))

    def get_private_groundtruth(self, k=None):
        """ override the ground-truth function as this is the only range search dataset """
        assert self.private_gt_url is not None
        fn = self.private_gt_url.split("/")[-1]   # in case it's a URL
        return range_result_read(os.path.join(self.basedir, fn))

class BigANNDataset(DatasetCompetitionFormat):
    def __init__(self, nb_M=1000):
        self.nb_M = nb_M
        self.nb = 10**6 * nb_M
        self.d = 128
        self.nq = 10000
        self.dtype = "uint8"
        self.ds_fn = "base.1B.u8bin"
        self.qs_fn = "query.public.10K.u8bin"
        self.gt_fn = (
            "GT.public.1B.ibin" if self.nb_M == 1000 else
            subset_url + "GT_100M/bigann-100M" if self.nb_M == 100 else
            subset_url + "GT_10M/bigann-10M" if self.nb_M == 10 else
            None
        )
        # self.gt_fn = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/bigann/public_query_gt100.bin" if self.nb == 10**9 else None
        self.base_url = "https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/bigann/"
        self.basedir = os.path.join(BASEDIR, "bigann")

        self.private_nq = 10000
        self.private_qs_url = "https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/bigann/query.private.799253207.10K.u8bin"
        self.private_gt_url = "https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/GT_1B_final_2bf4748c7817/bigann-1B.bin"

    def distance(self):
        return "euclidean"

class Deep1BDataset(DatasetCompetitionFormat):
    def __init__(self, nb_M=1000):
        self.nb_M = nb_M
        self.nb = 10**6 * nb_M
        self.d = 96
        self.nq = 10000
        self.dtype = "float32"
        self.ds_fn = "base.1B.fbin"
        self.qs_fn = "query.public.10K.fbin"
        self.gt_fn = (
            "https://storage.yandexcloud.net/yandex-research/ann-datasets/deep_new_groundtruth.public.10K.bin" if self.nb_M == 1000 else
            subset_url + "GT_100M/deep-100M" if self.nb_M == 100 else
            subset_url + "GT_10M/deep-10M" if self.nb_M == 10 else
            None
        )
        self.base_url = "https://storage.yandexcloud.net/yandex-research/ann-datasets/DEEP/"
        self.basedir = os.path.join(BASEDIR, "deep1b")

        self.private_nq = 30000
        self.private_qs_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/deep1b/query.heldout.30K.fbin"
        self.private_gt_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/deep1b/gt100-heldout.30K.fbin"

        self.private_nq_large = 1000000
        self.private_qs_large_url = "https://storage.yandexcloud.net/yr-secret-share/ann-datasets-5ac0659e27/DEEP/query.private.1M.fbin"

    def distance(self):
        return "euclidean"

class Text2Image1B(DatasetCompetitionFormat):
    def __init__(self, nb_M=1000):
        self.nb_M = nb_M
        self.nb = 10**6 * nb_M
        self.d = 200
        self.nq = 100000
        self.dtype = "float32"
        self.ds_fn = "base.1B.fbin"
        self.qs_fn = "query.public.100K.fbin"
        self.gt_fn = (
            "https://storage.yandexcloud.net/yandex-research/ann-datasets/t2i_new_groundtruth.public.100K.bin" if self.nb_M == 1000 else
            subset_url + "GT_100M/text2image-100M" if self.nb_M == 100 else
            subset_url + "GT_10M/text2image-10M" if self.nb_M == 10 else
            None
        )
        self.base_url = "https://storage.yandexcloud.net/yandex-research/ann-datasets/T2I/"
        self.basedir = os.path.join(BASEDIR, "text2image1B")

        self.private_nq = 30000
        self.private_qs_url = "https://storage.yandexcloud.net/yandex-research/ann-datasets/T2I/query.heldout.30K.fbin"
        self.private_gt_url = "https://storage.yandexcloud.net/yandex-research/ann-datasets/T2I/gt100-heldout.30K.fbin"

        self.private_nq_large = 1000000
        self.private_qs_large_url = "https://storage.yandexcloud.net/yr-secret-share/ann-datasets-5ac0659e27/T2I/query.private.1M.fbin"

    def distance(self):
        return "ip"

    def get_query_train(self, maxn=10**6):
        xq_train = np.memmap(
            BASEDIR + "/text2image1B/query.learn.50M.fbin", offset=8,
            dtype='float32', shape=(maxn, 200), mode='r')
        return np.array(xq_train)

class MSTuringANNS(DatasetCompetitionFormat):
    def __init__(self, nb_M=1000):
        self.nb_M = nb_M
        self.nb = 10**6 * nb_M
        self.d = 100
        self.nq = 100000
        self.dtype = "float32"
        self.ds_fn = "base1b.fbin"
        self.qs_fn = "query100K.fbin"
        self.gt_fn = (
            "query_gt100.bin" if self.nb_M == 1000 else
            "msturing-gt-100M" if self.nb_M == 100 else # back up subset_url + "GT_100M/msturing-100M"
            "msturing-gt-10M" if self.nb_M == 10 else # back up subset_url + "GT_100M/msturing-10M"
            "msturing-gt-1M" if self.nb_M == 1 else
            None
        )
        self.base_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/MSFT-TURING-ANNS/"
        self.basedir = os.path.join(BASEDIR, "MSTuringANNS")

        self.private_nq = 10000
        self.private_qs_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/MSFT-TURING-ANNS/testQuery10K.fbin"
        self.private_gt_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/MSFT-TURING-ANNS/gt100-private10K-queries.bin"

        self.private_nq_large = 99605
        self.private_qs_large_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/MSFT-TURING-ANNS/testQuery99605.fbin"
        self.private_gt_large_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/MSFT-TURING-ANNS/gt100-private99605-queries.bin"

    def distance(self):
        return "euclidean"

class MSTuringClustered10M(DatasetCompetitionFormat):
    def __init__(self):
        self.nb = 10**6 * 10
        self.d = 100
        self.nq = 10000
        self.dtype = "float32"
        self.ds_fn = "msturing-10M-clustered.fbin"
        self.qs_fn = "testQuery10K.fbin"
        self.gt_fn = "clu_msturing10M_gt100"

        self.base_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp23/clustered_data/msturing-10M-clustered/"
        self.basedir = os.path.join(BASEDIR, "MSTuring-10M-clustered")

        self.private_gt_url = None
        self.private_qs_url = None

    def distance(self):
        return "euclidean"

    def prepare(self, skip_data=False, original_size=10 ** 9):
        return super().prepare(skip_data, original_size = self.nb)

class MSTuringClustered30M(DatasetCompetitionFormat):
    def __init__(self):
        self.nb = 29998994
        self.d = 100
        self.nq = 10000
        self.dtype = "float32"
        self.ds_fn = "30M-clustered64.fbin"
        self.qs_fn = "testQuery10K.fbin"
        self.gt_fn = "clu_msturing30M_gt100"

        self.base_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp23/clustered_data/msturing-30M-clustered/"
        self.basedir = os.path.join(BASEDIR, "MSTuring-30M-clustered")

        self.private_gt_url = None
        self.private_qs_url = None

    def distance(self):
        return "euclidean"

    def prepare(self, skip_data=False, original_size=10 ** 9):
        return super().prepare(skip_data, original_size = self.nb)

class MSSPACEV1B(DatasetCompetitionFormat):
    def __init__(self, nb_M=1000):
        self.nb_M = nb_M
        self.nb = 10**6 * nb_M
        self.d = 100
        self.nq = 29316
        self.dtype = "int8"
        self.ds_fn = "spacev1b_base.i8bin"
        self.qs_fn = "query.i8bin"
        self.gt_fn = (
            "public_query_gt100.bin" if self.nb_M == 1000 else
            "msspacev-gt-100M" if self.nb_M == 100 else # backup subset_url + "GT_100M/msspacev-100M"
            "msspacev-gt-10M" if self.nb_M == 10 else # backup subset_url + "GT_10M/msspacev-10M"
            "msspacev-gt-1M" if self.nb_M == 1 else
            None
        )
        self.base_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/spacev1b/"
        self.basedir = os.path.join(BASEDIR, "MSSPACEV1B")

        self.private_nq = 30000
        self.private_qs_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/spacev1b/private_query_30k.bin"
        self.private_gt_url = "https://comp21storage.blob.core.windows.net/publiccontainer/comp21/spacev1b/gt100_private_query_30k.bin"

    def distance(self):
        return "euclidean"

class RandomClusteredDS(DatasetCompetitionFormat):
    def __init__(self, basedir="random-clustered"):
        self.nb = 10000
        self.nq = 1000
        self.d = 20
        self.dtype = 'float32'
        self.ds_fn = f"clu-random.fbin"
        self.qs_fn = f"queries_1000_20.fbin"
        self.gt_fn = f"clu_random_gt100"

        self.base_url="https://comp21storage.blob.core.windows.net/publiccontainer/comp23/clustered_data/random-xs-clustered/"

        self.basedir = os.path.join(BASEDIR, f"{basedir}{self.nb}")
        if not os.path.exists(self.basedir):
            os.makedirs(self.basedir)

        self.private_gt_url = None
        self.private_qs_url = None

    def search_type(self):
        return "knn"

    def distance(self):
        return "euclidean"

    def __str__(self):
        return f"RandomClustered({self.nb})"

    def default_count(self):
        return 10

    def prepare(self, skip_data=False, original_size=10 ** 9):
        return super().prepare(skip_data, original_size = self.nb)

class RandomRangeDS(DatasetCompetitionFormat):
    def __init__(self, nb, nq, d):
        self.nb = nb
        self.nq = nq
        self.d = d
        self.dtype = 'float32'
        self.ds_fn = f"data_{self.nb}_{self.d}"
        self.qs_fn = f"queries_{self.nq}_{self.d}"
        self.gt_fn = f"gt_{self.nb}_{self.nq}_{self.d}"
        self.basedir = os.path.join(BASEDIR, f"random{self.nb}")
        if not os.path.exists(self.basedir):
            os.makedirs(self.basedir)

    def prepare(self, skip_data=False):
        import sklearn.datasets
        import sklearn.model_selection
        from sklearn.neighbors import NearestNeighbors

        print(f"Preparing datasets with {self.nb} random points and {self.nq} queries.")

        X, _ = sklearn.datasets.make_blobs(
            n_samples=self.nb + self.nq, n_features=self.d,
            centers=self.nq, random_state=1)

        data, queries = sklearn.model_selection.train_test_split(
            X, test_size=self.nq, random_state=1)

        with open(os.path.join(self.basedir, self.ds_fn), "wb") as f:
            np.array([self.nb, self.d], dtype='uint32').tofile(f)
            data.astype('float32').tofile(f)
        with open(os.path.join(self.basedir, self.qs_fn), "wb") as f:
            np.array([self.nq, self.d], dtype='uint32').tofile(f)
            queries.astype('float32').tofile(f)

        print("Computing groundtruth")

        nbrs = NearestNeighbors(n_neighbors=100, metric="euclidean", algorithm='brute').fit(data)
        D, I = nbrs.kneighbors(queries)

        nres = np.count_nonzero((D < math.sqrt(self.default_count())) == True, axis=1)
        DD = np.zeros(nres.sum())
        II = np.zeros(nres.sum(), dtype='int32')

        s = 0
        for i, l in enumerate(nres):
            DD[s : s + l] = D[i, 0 : l]
            II[s : s + l] = I[i, 0 : l]
            s += l

        with open(os.path.join(self.basedir, self.gt_fn), "wb") as f:
            np.array([self.nq, nres.sum()], dtype='uint32').tofile(f)
            nres.astype('int32').tofile(f)
            II.astype('int32').tofile(f)
            DD.astype('float32').tofile(f)

    def get_groundtruth(self, k=None):
        """ override the ground-truth function as this is the only range search dataset """
        assert self.gt_fn is not None
        fn = self.gt_fn.split("/")[-1]   # in case it's a URL
        return range_result_read(os.path.join(self.basedir, fn))

    def search_type(self):
        return "range"

    def default_count(self):
        return 49

    def distance(self):
        return "euclidean"

    def __str__(self):
        return f"RandomRange({self.nb})"

class YFCC100MDataset(DatasetCompetitionFormat):
    """ the 2023 competition """

    def __init__(self, filtered=True, dummy=False):
        self.filtered = filtered
        nb_M = 10
        self.nb_M = nb_M
        self.nb = 10**6 * nb_M
        self.d = 192
        self.nq = 100000
        self.dtype = "uint8"
        private_key = 12345
        if dummy:
            # for now it's dummy because we don't have the descriptors yet
            self.ds_fn = "dummy2.base.10M.u8bin"
            self.qs_fn = "dummy2.query.public.100K.u8bin"
            self.qs_private_fn = "dummy2.query.private.%d.100K.u8bin" % private_key
            self.ds_metadata_fn = "dummy2.base.metadata.10M.spmat"
            self.qs_metadata_fn = "dummy2.query.metadata.public.100K.spmat"
            self.qs_private_metadata_fn = "dummy2.query.metadata.private.%d.100K.spmat" % private_key
            if filtered:
                # no subset as the database is pretty small.
                self.gt_fn = "dummy2.GT.public.ibin"
            else:
                self.gt_fn = "dummy2.unfiltered.GT.public.ibin"

        else:
            # with Zilliz' CLIP descriptors
            self.ds_fn = "base.10M.u8bin"
            self.qs_fn = "query.public.100K.u8bin"
            self.qs_private_fn = "query.private.%d.100K.u8bin" % private_key
            self.ds_metadata_fn = "base.metadata.10M.spmat"
            self.qs_metadata_fn = "query.metadata.public.100K.spmat"
            self.qs_private_metadata_fn = "query.metadata.private.%d.100K.spmat" % private_key
            if filtered:
                # no subset as the database is pretty small.
                self.gt_fn = "GT.public.ibin"
            else:
                self.gt_fn = "unfiltered.GT.public.ibin"

            # data is uploaded but download script not ready.
        self.base_url = "https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/yfcc100M/"
        self.basedir = os.path.join(BASEDIR, "yfcc100M")

        self.private_nq = 100000
        self.private_qs_url = self.base_url + ""
        self.private_gt_url = self.base_url + ""

        self.metadata_base_url = self.base_url + self.ds_metadata_fn
        self.metadata_queries_url = self.base_url + self.qs_metadata_fn
        self.metadata_private_queries_url = ""

    def prepare(self, skip_data=False):
        super().prepare(skip_data, 10**7)
        for fn in (self.metadata_base_url, self.metadata_queries_url, self.metadata_private_queries_url):
            if fn:
                outfile = os.path.join(self.basedir, fn.split("/")[-1])
                if os.path.exists(outfile):
                    print("file %s already exists" % outfile)
                else:
                    download(fn, outfile)

    def get_dataset_metadata(self):
        return read_sparse_matrix(os.path.join(self.basedir, self.ds_metadata_fn))

    def get_queries_metadata(self):
        return read_sparse_matrix(os.path.join(self.basedir, self.qs_metadata_fn))

    def get_private_queries_metadata(self):
        return read_sparse_matrix(os.path.join(self.basedir, self.qs_private_metadata_fn))

    def distance(self):
        return "euclidean"

    def search_type(self):
        if self.filtered:
            return "knn_filtered"
        else:
            return "knn"

def _strip_gz(filename):
    if not filename.endswith('.gz'):
        raise RuntimeError(f"expected a filename ending with '.gz'. Received: {filename}")
    return filename[:-3]

def _gunzip_if_needed(filename):
    if filename.endswith('.gz'):
        print('unzipping', filename, '...', end=" ", flush=True)

        with gzip.open(filename, 'rb') as f_in, open(_strip_gz(filename), 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)

        os.remove(filename)
        print('done.')

class SparseDataset(DatasetCompetitionFormat):
    """ the 2023 competition
        Sparse vectors for sparse max inner product search
        Data is based on MSMARCO passage retrieval data (text passages and queries),
        embedded via the SPLADE model.

        The class overrides several methods since the sparse format is different than other datasets.
    """

    def __init__(self, version="small"):

        versions = {"small": (100000, "base_small.csr.gz", "base_small.dev.gt"),
                    "1M": (1000000, "base_1M.csr.gz", "base_1M.dev.gt"),
                    "full": (8841823, "base_full.csr.gz", "base_full.dev.gt")}

        assert version in versions, f'version="{version}" is invalid. Please choose one of {list(versions.keys())}.'

        self.nb = versions[version][0]
        self.nq = 6980
        self.private_nq = 0  # TBD

        self.ds_fn = versions[version][1]
        self.qs_fn = "queries.dev.csr.gz"

        self.qs_private_fn = ""  # TBD

        self.base_url = "https://storage.googleapis.com/ann-challenge-sparse-vectors/csr/"
        self.basedir = os.path.join(BASEDIR, "sparse")

        self.gt_fn = versions[version][2]
        self.private_gt = ""  # TBD

        self.d = np.nan # this is only for compatibility with printing the name of the class

    def prepare(self, skip_data=False):
        # downloads the datasets and unzips (if necessary).

        if not os.path.exists(self.basedir):
            os.makedirs(self.basedir)

        # start with the small ones...
        for fn in [self.qs_fn, self.gt_fn]:
            if fn is None:
                continue

            sourceurl = os.path.join(self.base_url, fn)
            outfile = os.path.join(self.basedir, fn)
            if outfile.endswith('.gz'):
                # check if the unzipped file already exists
                if os.path.exists(_strip_gz(outfile)):
                    print("unzipped version of file %s already exists" % outfile)
                    continue

            if os.path.exists(outfile):
                print("file %s already exists" % outfile)
                _gunzip_if_needed(outfile)
                continue
            download(sourceurl, outfile)
            _gunzip_if_needed(outfile)
        # # private qs url: todo

        if skip_data:
            return

        fn = self.ds_fn
        sourceurl = os.path.join(self.base_url, fn)
        outfile = os.path.join(self.basedir, fn)
        if outfile.endswith('.gz'):
            # check if the unzipped file already exists
            unzipped_outfile = _strip_gz(outfile)
            if os.path.exists(unzipped_outfile):
                print("unzipped version of file %s already exists" % outfile)
                return
        if os.path.exists(outfile):
            print("file %s already exists" % outfile)
            _gunzip_if_needed(outfile)
            return
        download(sourceurl, outfile)
        _gunzip_if_needed(outfile)

    def get_dataset_fn(self):
        fn = _strip_gz(os.path.join(self.basedir, self.ds_fn))
        if os.path.exists(fn):
            return fn
        raise RuntimeError("file not found")

    def get_dataset_iterator(self, bs=512, split=(1, 0)):
        assert split == (1,0), 'No sharding supported yet.'  # todo

        filename = self.get_dataset_fn()

        x = read_sparse_matrix(filename, do_mmap=True)
        assert x.shape[0] == self.nb

        for j0 in range(0, self.nb, bs):
            j1 = min(j0 + bs, self.nb)
            yield x[j0:j1, :]

        # i0, i1 = self.nb * rank // nsplit, self.nb * (rank + 1) // nsplit
        # x = xbin_mmap(filename, dtype=self.dtype, maxn=self.nb)
        # assert x.shape == (self.nb, self.d)

        # for j0 in range(i0, i1, bs):
        #     j1 = min(j0 + bs, i1)
        #     yield sanitize(x[j0:j1])

    def get_groundtruth(self, k=None):
        assert self.gt_fn is not None
        assert self.search_type() == "knn"

        I, D = knn_result_read(os.path.join(self.basedir, self.gt_fn))
        assert I.shape[0] == self.nq
        if k is not None:
            assert k <= 10
            I = I[:, :k]
            D = D[:, :k]
        return I, D

    def get_dataset(self):
        assert self.nb <= 10 ** 6, "dataset too large, use iterator"
        return next(self.get_dataset_iterator(bs=self.nb))

    def get_queries(self):
        filename = os.path.join(self.basedir, self.qs_fn)
        print(filename)
        x = read_sparse_matrix(_strip_gz(filename), do_mmap=False)  # read the queries file. It is a small file, so no need to mmap
        assert x.shape[0] == self.nq
        return x

    def get_private_queries(self):
        raise RuntimeError("not implemented yet")

    def get_private_groundtruth(self, k=None):
        raise RuntimeError("not implemented yet")

    def distance(self):
        return "ip"

    def search_type(self):
        return "knn"

    def data_type(self):
        return "sparse"

class RandomDS(DatasetCompetitionFormat):
    def __init__(self, nb, nq, d, basedir="random"):
        self.nb = nb
        self.nq = nq
        self.d = d
        self.dtype = 'float32'
        self.ds_fn = f"data_{self.nb}_{self.d}"
        self.qs_fn = f"queries_{self.nq}_{self.d}"
        self.gt_fn = f"gt_{self.nb}_{self.nq}_{self.d}"
        self.basedir = os.path.join(BASEDIR, f"{basedir}{self.nb}")
        if not os.path.exists(self.basedir):
            os.makedirs(self.basedir)

    def prepare(self, skip_data=False):
        import sklearn.datasets
        import sklearn.model_selection
        from sklearn.neighbors import NearestNeighbors

        print(f"Preparing datasets with {self.nb} random points and {self.nq} queries.")

        X, _ = sklearn.datasets.make_blobs(
            n_samples=self.nb + self.nq, n_features=self.d,
            centers=self.nq, random_state=1)

        data, queries = sklearn.model_selection.train_test_split(
            X, test_size=self.nq, random_state=1)

        with open(os.path.join(self.basedir, self.ds_fn), "wb") as f:
            np.array([self.nb, self.d], dtype='uint32').tofile(f)
            data.astype('float32').tofile(f)
        with open(os.path.join(self.basedir, self.qs_fn), "wb") as f:
            np.array([self.nq, self.d], dtype='uint32').tofile(f)
            queries.astype('float32').tofile(f)

        print("Computing groundtruth")

        nbrs = NearestNeighbors(n_neighbors=100, metric="euclidean", algorithm='brute').fit(data)
        D, I = nbrs.kneighbors(queries)
        with open(os.path.join(self.basedir, self.gt_fn), "wb") as f:
            np.array([self.nq, 100], dtype='uint32').tofile(f)
            I.astype('uint32').tofile(f)
            D.astype('float32').tofile(f)

    def search_type(self):
        return "knn"

    def distance(self):
        return "euclidean"

    def __str__(self):
        return f"Random({self.nb})"

    def default_count(self):
        return 10

class RandomFilterDS(RandomDS):
    def __init__(self, nb, nq, d):
        super().__init__(nb, nq, d, "random-filter")
        self.ds_metadata_fn = f"data_metadata_{self.nb}_{self.d}"
        self.qs_metadata_fn = f"queries_metadata_{self.nb}_{self.d}"

    def prepare(self, skip_data=False):
        import sklearn.datasets
        import sklearn.model_selection
        from sklearn.neighbors import NearestNeighbors

        print(f"Preparing datasets with {self.nb} random points, {self.nq} queries, and two filters.")

        X, _ = sklearn.datasets.make_blobs(
            n_samples=self.nb + self.nq, n_features=self.d,
            centers=self.nq, random_state=1)

        data, queries = sklearn.model_selection.train_test_split(
            X, test_size=self.nq, random_state=1)

        filter1 = [1, 2]
        filter2 = [3, 4]

        assert self.nb % 2 == 0

        # simple filters, first half of the data matches second
        # half of the queries, and vice versa

        data_filters = [filter1] * (self.nb // 2) + [filter2] * (self.nb // 2)
        query_filters = [filter2] * (self.nq // 2) + [filter1] * (self.nq // 2)

        assert len(data_filters) == data.shape[0]

        with open(os.path.join(self.basedir, self.ds_fn), "wb") as f:
            np.array([self.nb, self.d], dtype='uint32').tofile(f)
            data.astype('float32').tofile(f)
        with open(os.path.join(self.basedir, self.qs_fn), "wb") as f:
            np.array([self.nq, self.d], dtype='uint32').tofile(f)
            queries.astype('float32').tofile(f)

        data_indices = np.array(data_filters).flatten()
        data_indptr = [2 * i for i in range(self.nb)] + [2 * self.nb]
        data_data = [1] * self.nb * 2
        data_metadata_sparse = csr_matrix((data_data, data_indices, data_indptr))

        query_indices = np.array(query_filters).flatten()
        query_indptr = [2 * i for i in range(self.nq)] + [2 * self.nq]
        query_data = [1] * self.nq * 2
        query_metadata_sparse = csr_matrix((query_data, query_indices, query_indptr))

        write_sparse_matrix(data_metadata_sparse,
                            os.path.join(self.basedir, self.ds_metadata_fn))
        write_sparse_matrix(query_metadata_sparse,
                            os.path.join(self.basedir, self.qs_metadata_fn))

        print("Computing groundtruth")

        n_neighbors = 100

        nbrs = NearestNeighbors(n_neighbors=n_neighbors, metric="euclidean", algorithm='brute').fit(data[:self.nb // 2])
        DD, II = nbrs.kneighbors(queries[self.nq // 2:])

        nbrs = NearestNeighbors(n_neighbors=n_neighbors, metric="euclidean", algorithm='brute').fit(data[self.nb // 2: ])
        D, I = nbrs.kneighbors(queries[:self.nq // 2])

        D = np.concatenate((D, DD))
        I = np.concatenate((I + self.nb // 2, II))

        with open(os.path.join(self.basedir, self.gt_fn), "wb") as f:
            np.array([self.nq, n_neighbors], dtype='uint32').tofile(f)
            I.astype('uint32').tofile(f)
            D.astype('float32').tofile(f)

    def get_dataset_metadata(self):
        return read_sparse_matrix(os.path.join(self.basedir, self.ds_metadata_fn))

    def get_queries_metadata(self):
        return read_sparse_matrix(os.path.join(self.basedir, self.qs_metadata_fn))

    def search_type(self):
        return "knn_filtered"

    def __str__(self):
        return f"RandomFilter({self.nb})"

DATASETS = {
    'bigann-1B': lambda : BigANNDataset(1000),
    'bigann-100M': lambda : BigANNDataset(100),
    'bigann-10M': lambda : BigANNDataset(10),

    'deep-1B': lambda : Deep1BDataset(),
    'deep-100M': lambda : Deep1BDataset(100),
    'deep-10M': lambda : Deep1BDataset(10),

    'ssnpp-1B': lambda : SSNPPDataset(1000),
    'ssnpp-10M': lambda : SSNPPDataset(10),
    'ssnpp-100M': lambda : SSNPPDataset(100),
    'ssnpp-1M': lambda : SSNPPDataset(1),

    'text2image-1B': lambda : Text2Image1B(),
    'text2image-1M': lambda : Text2Image1B(1),
    'text2image-10M': lambda : Text2Image1B(10),
    'text2image-100M': lambda : Text2Image1B(100),

    'msturing-1B': lambda : MSTuringANNS(1000),
    'msturing-100M': lambda : MSTuringANNS(100),
    'msturing-10M': lambda : MSTuringANNS(10),
    'msturing-1M': lambda : MSTuringANNS(1),

    'msturing-10M-clustered': lambda: MSTuringClustered10M(),
    'msturing-30M-clustered': lambda: MSTuringClustered30M(),

    'msspacev-1B': lambda : MSSPACEV1B(1000),
    'msspacev-100M': lambda : MSSPACEV1B(100),
    'msspacev-10M': lambda : MSSPACEV1B(10),
    'msspacev-1M': lambda : MSSPACEV1B(1),

    'yfcc-10M': lambda: YFCC100MDataset(),
    'yfcc-10M-unfiltered': lambda: YFCC100MDataset(filtered=False),
    'yfcc-10M-dummy': lambda: YFCC100MDataset(dummy=True),
    'yfcc-10M-dummy-unfiltered': lambda: YFCC100MDataset(filtered=False, dummy=True),

    'sparse-small': lambda: SparseDataset("small"),
    'sparse-1M': lambda: SparseDataset("1M"),
    'sparse-full': lambda: SparseDataset("full"),

    'random-xs': lambda : RandomDS(10000, 1000, 20),
    'random-s': lambda : RandomDS(100000, 1000, 50),

    'random-xs-clustered': lambda: RandomClusteredDS(),

    'random-range-xs': lambda : RandomRangeDS(10000, 1000, 20),
    'random-range-s': lambda : RandomRangeDS(100000, 1000, 50),

    'random-filter-s': lambda : RandomFilterDS(100000, 1000, 50),

}

{'sklearn.datasets': 'sklearn.datasets', 'sklearn.model_selection': 'sklearn.model_selection', 'NearestNeighbors': 'sklearn.neighbors.NearestNeighbors'}


'''
919 · Meeting Rooms II
https://www.lintcode.com/problem/920
'''
import heapq
import random
from typing import List, Tuple

from test_tool import assert_value

class Solution:
    def minMeetingRooms(self, intervals: List[Tuple[int]]) -> int:
        cnt = []
        for s, e in intervals:
            heapq.heappush(cnt, (s, 'start'))
            heapq.heappush(cnt, (e, 'end'))
        res, curr = 0, 0
        while cnt:
            t, i = heapq.heappop(cnt)
            if i == 'start':
                curr += 1
            else:
                curr -= 1
            res = max(curr, res)
        return res

assert_value(2, Solution().minMeetingRooms, intervals=[(0, 30), (5, 10), (15, 20)])
assert_value(1, Solution().minMeetingRooms, intervals=[(2, 7)])


from typing import List

class Solution:
    # 气死了，中位数的方法怎么都搞不好 merge 那，先写个排序的，之后把中位数的补上
    def wiggleSort(self, nums: List[int]) -> None:
        tmp = sorted(nums)
        mid_index = int((len(nums) - 1) / 2)
        i, j, k = mid_index, len(nums) - 1, 0
        while i >= 0 and j > mid_index:
            nums[k] = tmp[i]
            nums[k+1] = tmp[j]
            k += 2
            i -= 1
            j -= 1
        if i == 0:
            nums[k] = tmp[i]

    # """
    # 大小不等的不满足条件，经过交换会产生效果，但是对于大小相等的元素不适用。
    #
    # 要求时间 O(n)，空间 O(1) 是真的好难想 ...
    #
    # 思路：利用快排的思想得到 nums 的中位数，然后把小于中位数的数放在 (i & 1) == 0 的位置，大于中位数的数放在 (i & 1) == 1 的位置。
    #
    # 在 merge 时的注意要点：要从尾开始取元素！
    # 如果从头取元素：4556 -> 4556，不会发生变化
    # 从尾取元素：4456 -> 5645，不会导致第二个元素和第三个元素相等
    # """
    # def wiggleSort(self, nums: List[int]) -> None:
    #     tmp = [0] * len(nums)
    #     for i in range(len(nums)):
    #         tmp[i] = nums[i]
    #     mid_index = self.mid_num(tmp)
    #     i, j, k = mid_index, len(nums)-1, 0
    #     while i >= 0 and j > mid_index:
    #         nums[k] = tmp[i]
    #         nums[k+1] = tmp[j]
    #         k += 2
    #         i -= 1
    #         j -= 1
    #     if i == 0:
    #         nums[k] = tmp[i]
    #
    # def mid_num(self, nums):
    #     N = len(nums)
    #     mid_index = int((N - 1) / 2)
    #     index = self.partition(nums, 0, N-1, mid_index)
    #     while index != mid_index:
    #         if index < mid_index:
    #             index = self.partition(nums, index+1, N-1, mid_index)
    #         else:
    #             index = self.partition(nums, 0, index-1, mid_index)
    #     return mid_index
    #
    # def partition(self, nums, lo, hi, pos):
    #     nums[lo], nums[pos] = nums[pos], nums[lo]
    #     i = lo
    #     for j in range(lo+1, hi+1):
    #         if nums[j] <= nums[lo]:
    #             i += 1
    #             nums[i], nums[j] = nums[j], nums[i]
    #     nums[lo], nums[i] = nums[i], nums[lo]
    #     return i

if __name__ == '__main__':
    s = Solution()
    # nums = [1, 2, 2, 1, 2, 1, 1, 1, 1, 2, 2, 2]
    # nums = [1, 2, 2, 1, 2, 1, 1, 1, 1, 2, 2]
    # nums = [1, 5, 1, 1, 6, 4]
    # WA points
    # nums = [1, 3, 2, 2, 3, 1]
    # nums = [4, 5, 5, 6]
    nums = [5, 3, 1, 2, 6, 7, 8, 5, 5]

    # print(s.partition(nums, 0, len(nums)-1, 0))
    # s.mid_num(nums)
    # print(nums)

    s.wiggleSort(nums)
    print(nums)


import MetadataInfo as mdi
import pandas as pd
import datetime as dt

#define a converter for use with PANDAS' read_csv file
converter={}

# Transform currencies to be more easily usable
chars_to_remove = "$,%"

remove_chars_from = mdi.has_dollarsign + mdi.has_percentsign
remove_chars = lambda x : float(x.translate(str.maketrans('', '', chars_to_remove))) if x else None
for field in remove_chars_from:
    converter[field]=remove_chars

fill_rating = lambda x : x if x else 'NO VALUE ENTERED'
converter['X8']= fill_rating
converter['X9']= fill_rating

fill_residence_status = lambda x : x if x else 'NO VALUE ENTERED'
converter['X12'] = fill_residence_status

fill_employer = fill_residence_status
converter['X10'] = fill_employer

# convert a datetime to number of days, making it scale
def to_days(x):
    reference_date=dt.datetime.strptime('Feb-1-2016','%b-%d-%Y').date()
    dateparts = x.split('-')
    if x:
        if x[0].isdigit():
            diff = reference_date - dt.datetime.strptime(x,"%d-%b").date()
        else:
            diff = reference_date - dt.datetime.strptime(x,"%b-%y").date()
        return float(diff.days / ( 365) )

converter['X23'] = to_days

#removing the years makes it more clearly ordinal
def remove_years(x):
    '''
    Changes the length of employment to obvious ordinal ranks
    :param x: input string for employment length
    :return: number representing ordinal rank
    '''
    if '<' in x:
        return 0
    elif '+' in x:
        return 10
    elif x=='n/a' or x=='':
        return -1
    else:
        return int(x.replace(' years','').replace(' year',''))

converter['X11']=remove_years

# Convert 36 or 60 months into a number, for more plain ordinal ranking
remove_months = lambda x : x.replace(' months','').replace(' ','')
converter[mdi.has_phrase_months] = remove_months

# define the data types we wish the import to assign
field_dtypes={}
for field in mdi.scale_fields:
    field_dtypes[field]=pd.np.float64

for field in mdi.nominal_fields:
    field_dtypes[field] = pd.np.str

for field in mdi.ordinal_fields:
    field_dtypes[field] = pd.np.float64

for field in mdi.string_fields:
    field_dtypes[field] = pd.np.str

for field in mdi.date_fields:
    field_dtypes[field] = pd.np.timedelta64
    converter[field] = to_days

def custom_import_data(rows=None):
    '''
    Import the training data from the CSV file
    :param rows: number of rows to import. If None, import them all
    :return: a Pandas DataFrame containing the data
    '''
    return pd.read_csv('Data for Cleaning & Modeling.csv',nrows=rows, header=0, low_memory=False, converters=converter, dtype=field_dtypes);

def custom_import_test_data(rows=None):
    '''
    Import the test data from the CSV file
    :param rows: number of rows to import. If None, import them all
    :return: a Pandas DataFrame containing the data
    '''
    return pd.read_csv('Holdout for Testing.csv',header=0,nrows=rows, low_memory=False, converters=converter, dtype=field_dtypes);


from cgi import test
from http import client
from urllib import response
import pytest
from library.editorial.models import *
from rest_framework.test import APIClient
from rest_framework.test import RequestsClient

client = APIClient()

@pytest.mark.django_db
def test_editorial_name():
    editorial = Editorial.objects.create(
        name = 'Posada',
        address = 'Tizapan San Angel, Alvaro obregon',
        country = 'CDMX'
        )
    print('editorial name>>>', editorial.name)
    assert editorial.name == 'Posada'
    editorial.delete()

@pytest.mark.django_db
def test_editorial_address():
    editorial = Editorial.objects.create(
        name = 'Posada',
        address = 'Tizapan San Angel, Alvaro obregon',
        country = 'CDMX'
        )
    print('editorial adress>>>', editorial.name)
    assert editorial.address == 'Tizapan San Angel, Alvaro obregon'
    editorial.delete()

@pytest.mark.django_db
def test_editorial_country():
    editorial = Editorial.objects.create(
        name = 'Posada',
        address = 'Tizapan San Angel, Alvaro obregon',
        country = 'CDMX'
        )
    print('editorial country>>>', editorial.name)
    assert editorial.country == 'CDMX'
    editorial.delete()

@pytest.mark.django_db
def test_editorial_post_fail():
    response = client.post("editorial/company/", dict(
        nme = 'Posada',
        dress = 'Tizapan San Angel, Alvaro obregon',
        cuntr = 'CDMX'
    ))
    assert response.status_code == 404

@pytest.mark.django_db
def test_editorial_():
    payload = dict(
        name = 'Posada',
        address = 'Tizapan San Angel, Alvaro obregon',
        country = 'CDMX'
    )
    print(payload)
    response = client.post("/editorial/", payload)

    data = response.data
    assert data["name"] == payload["name"]
    assert data["address"] == payload["address"]
    assert data["country"] == payload["country"]
    assert response.status_code == 201

@pytest.mark.django_db
def test_company_phone():
    company = Company.objects.create(
        phone = '66212121',
        name = '',
        email = '',
        address = '',
        country = '',
        employees_qnt =  '',
        company_role = '',
        )
    print('company country>>>', company.phone)
    assert company.phone == '66212121'
    company.delete()

@pytest.mark.django_db
def test_company_name():
    company = Company.objects.create(
        phone = '',
        name = 'Editorial corp',
        email = '',
        address = '',
        country = '',
        employees_qnt =  '',
        company_role = '',
        )
    print('company country>>>', company.phone)
    assert company.name == 'Editorial corp'
    company.delete()

@pytest.mark.django_db
def test_company_post():
    payload = dict(
        phone = '66212121',
        name = 'Company Corp',
        address = 'nevada 123123',
        country = 'Nevada',
        employees_qnt =  '22',
        company_role = "NEWS",
    )
    print(payload)
    response = client.post("/editorial/company/", payload)
    data = response.data
    assert data["name"] == payload["name"]
    assert data["address"] == payload["address"]
    assert data["phone"] == payload["phone"]
    assert data["address"] == payload["address"]
    assert data["employees_qnt"] == payload["employees_qnt"]
    assert data["company_role"] == payload["company_role"]
    assert response.status_code == 201

@pytest.mark.django_db
def test_company_post_fail():
    response = client.post("editorial/company/", dict(
        phon = '66212121'
    ))
    assert response.status_code == 404


import base_net
import json

ADDRESS = "localhost"
PORT = 50000
LISTEN = 10

class ServerUi(base_net.Server):
    def __init__(self, *args):
        super().__init__(*args)

    def on_receive(self, client, data):
        print(data)
        _json = json.loads(data)
        self.send_all(_json["message"])

if __name__ == "__main__":
    server = ServerUi(ADDRESS, PORT, LISTEN)
    server.start()


#!/usr/bin/env python3

import logging
import time
from argparse import ArgumentParser
from logging.handlers import TimedRotatingFileHandler
from pathlib import Path

from qc_tool.common import CONFIG
from qc_tool.common import create_job_dir
from qc_tool.worker.dispatch import dispatch

LOG_FORMAT = "{asctime} {name}:{levelname} {pathname}:{lineno} {message}"
LOG_TIME_FORMAT = "%Y-%m-%d %H:%M:%S"
LOG_FILENAME = "job.log"

log = logging.getLogger(__name__)

def init_logging(job_dir, log_filename=LOG_FILENAME):
    formatter = logging.Formatter(fmt=LOG_FORMAT, style="{", datefmt=LOG_TIME_FORMAT)
    # All log timestamps should be in UTC.
    formatter.converter = time.gmtime
    # Every job has its own log located in its job dir.
    handler = logging.FileHandler(job_dir.joinpath(log_filename))
    handler.setFormatter(formatter)

def main():
    # Parse command line arguments.
    parser = ArgumentParser()
    parser.add_argument("--job-uuid",
                        help="The identifier of the job being run.",
                        dest="job_uuid",
                        action="store",
                        nargs=1,
                        required=True)
    parser.add_argument("--product",
                        help="Identifier of the product.",
                        dest="product_ident",
                        action="store",
                        nargs=1,
                        required=True)
    parser.add_argument("--skip-steps",
                        help="Comma separated ordinal numbers of job steps to be skipped.",
                        dest="skip_steps",
                        default=None,
                        action="store",
                        nargs=1,
                        required=False)
    parser.add_argument("username",
                        help="The name of the user managing the delivery.",
                        action="store",
                        nargs=1)
    parser.add_argument("filename",
                        help="Path to delivery file relative to INCOMING_DIR.",
                        action="store",
                        nargs=1)
    pargs = parser.parse_args()
    job_uuid = pargs.job_uuid[0]
    username = pargs.username[0]
    filename = pargs.filename[0]
    filepath = CONFIG["incoming_dir"].joinpath(username, filename)
    if pargs.skip_steps is None:
        skip_steps = tuple()
    else:
        skip_steps = tuple(int(i) for i in pargs.skip_steps[0].split(","))

    # Create job dir.
    job_dir = create_job_dir(job_uuid)

    # Set up logging.
    init_logging(job_dir)
    log.info("Logging of the job {:s} has been started.".format(job_uuid))

    # Run the checks.
    dispatch(job_uuid, username, filepath, pargs.product_ident[0], skip_steps)

if __name__ == "__main__":
    main()


import pandas as pd
import matplotlib.pyplot as plt
import os
import scipy.stats

os.chdir(os.path.dirname(__file__))

OUTPUT_DIR = "histograms"
INPUT_FILENAME = "OUT_ALL_CLEAN.csv"

DISTINCT_THRESHOLD = 10
# create output dir
if not os.path.isdir(OUTPUT_DIR):
    os.mkdir(OUTPUT_DIR)

data = pd.read_csv(INPUT_FILENAME, sep=',', low_memory=False)
sequence_types = list(data["Initial Label"].unique())

print(sequence_types)

trfs = data[data["Initial Label"] == "trf"]

data_columns = data.columns.values[2:]

COLOR_PALETTE = ["#e60049", "#0bb4ff", "#50e991", "#e6d800", "#9b19f5", "#ffa300", "#dc0ab4", "#b3d4ff", "#00bfa0"]

kruskal_data = []
MW_data = []
# Kruskal wallis
ZERO_THRESHOLD = 1e-5
for column in data_columns:
    kwdata = []
    for seq_idx, seq_type in enumerate(sequence_types, start=1):
        active_data =  data[data["Initial Label"] == seq_type][column].dropna()
        kwdata.append(active_data)
    # https://stackoverflow.com/questions/51632900/pandas-apply-kruskal-wallis-to-numeric-columns
    kruskal = scipy.stats.kruskal(*kwdata, nan_policy="omit")
    kruskal_data.append([column, *kruskal])

    for seq_idx, seq_type_1 in enumerate(sequence_types[:-1]):
        for seq_type_2 in sequence_types[seq_idx+1:]:
            seq1_data = data[data["Initial Label"] == seq_type_1][column].dropna()
            seq2_data = data[data["Initial Label"] == seq_type_2][column].dropna()
            try:
                MW_res = scipy.stats.mannwhitneyu(seq1_data, seq2_data, use_continuity=False, method="asymptotic")
                if MW_res.pvalue < 0.05:
                    pvalue = MW_res.pvalue
                    if pvalue < ZERO_THRESHOLD:
                        pvalue = 0
                    MW_data.append([column, seq_type_1, seq_type_2, MW_res.statistic, pvalue])
            except TypeError:
                pass

kruskal_df = pd.DataFrame(kruskal_data, columns=("Column", "H statistic", "pvalue"))
kruskal_df[kruskal_df["pvalue"] < ZERO_THRESHOLD]["pvalue"] = 0.0
kruskal_df.to_csv("out_kruskal.csv", sep=';', index=False)

MW_df = pd.DataFrame(MW_data, columns=("Column", "Initial Label 1", "Initial Label 2", "Mann-Whitney U statistic", "pvalue"))
MW_df.to_csv("out_mannwhitney.csv", sep=';', index=False)

for column in data_columns:
    print(f"\tColumn: {column}")

    try:
        plt.figure(figsize=(18.5, 10.5))
        ax_share = None

        n_bins = 20

        sep_data = []

        for seq_idx, seq_type in enumerate(sequence_types, start=1):
            active_data =  data[data["Initial Label"] == seq_type]
            print(f"Running: {seq_type}")

            clr = COLOR_PALETTE[seq_idx-1]

            if ax_share:
                plt.subplot(len(sequence_types), 1, seq_idx, sharex=ax_share, sharey=ax_share)
            else:
                ax_share = plt.subplot(len(sequence_types), 1, seq_idx)

            if len(active_data[column].unique()) <= DISTINCT_THRESHOLD:
                val_counts = dict(active_data[column].value_counts())
                bar_labels = sorted(list(map(str, val_counts.keys())))
                plt.bar(bar_labels, [x / len(active_data[column]) for x in val_counts.values()], color=clr, alpha=0.7, label=seq_type)

            else:

                # n_bins = 0
                # try:
                #     n_bins = math.ceil((active_data[column].max() - active_data[column].min()) / 0.5)
                # except TypeError:
                #     print("\t\tNo numeric data!")
                #     continue

                # except OverflowError:
                #     print("\t\tInf values!")
                #     continue

                # finally:
                #     plt.close()

                h, _, _ = plt.hist(active_data[column], n_bins, facecolor=clr, alpha=0.7, density=True, label=seq_type)
                sep_data.append(h)

            plt.grid(True)
            plt.xlabel(column)
            plt.ylabel('Density')
            # fig.yscale('log')
            #plt.title(f"Column: {column}")
            # plt.show()
            plt.plot([], [], ' ', label=f"mean: {active_data[column].mean()}")
            plt.plot([], [], ' ', label=f"stdev: {active_data[column].std()}")
            plt.legend()

        # if len(sep_data) == len(sequence_types):
        #     kruskal = scipy.stats.kruskal(*sep_data)
        #     kruskal_data.append([column, *kruskal])

        plt.suptitle(f"Column: {column}")
        seq_type = seq_type.replace(" / ", "-")
        column_file_name = str(column).replace("/", "slash").replace("|", "line")

        plot_output_path = os.path.join(OUTPUT_DIR, f"{column_file_name}.png")
        plt.savefig(plot_output_path)
        #plt.show()
        plt.close()

    except (TypeError, ValueError):
        print(f"Could not plot Column: {column}")

# kruskal_df = pd.DataFrame(kruskal_data, columns=("Initial Label", "H statistic", "pvalue"))
# kruskal_df.to_csv("out_kruskal.csv", sep=';', index=False)


import spotipy
from selenium.webdriver.common.keys import Keys
import os
import time

from . browser_control import Browser
from . song import Song

class SpotifyPlayer(Browser):

    def __init__(self, device_name, headless=True):
        super().__init__(headless=headless, landing_page='about:preferences')
        self.device_name = device_name
        self.clickOnElement(None, css='#playDRMContent')
        self.clickOnElement(None, css='#autoplaySettingsButton')
        self.cached_songs = None

        self.sendKeyToBrowser(Keys.UP, post_sleep=0.3)
        for _ in range(4):
            self.sendKeyToBrowser(Keys.TAB, post_sleep=0.3)
        self.sendKeyToBrowser(Keys.ENTER)

    def loadKey(self, key):
        self.key = key
        self.sp = spotipy.Spotify(auth=key)
        print("!!! LOADED SP ---")

    def loadSpotifyDevice(self, key):
        self.loadKey(key)
        file_url = f'file://{os.getcwd()}/music_backend/backend_lib/spotify_player.html?device-name={self.device_name.replace(" ", "%20")}&token={self.key}'
        print(f'FILE_URL: {file_url}')
        self.goToURL(file_url)
        pi_device = None
        for _ in range(10):
            time.sleep(1)
            devices = self.sp.devices()['devices']
            for device in devices:
                if device['name'] == self.device_name:
                    pi_device = device['id']
            if pi_device != None:
                print('Switched to new device for spotify!')
                break

        self.sp.transfer_playback(pi_device, force_play=False)

    def startPlayingSong(self, global_state, song):
        print(f"playing from spotify: |{song.url}|")
        global_state.songStarted()
        sp = spotipy.Spotify(auth=global_state.getSpotifyKey())
        sp.start_playback(uris=[song.url])
        playing = True
        while True:
            changedThisTurn = False
            if not global_state.isPlaying() and playing:
                sp.pause_playback()
                playing = False
                changedThisTurn = True
            elif global_state.isPlaying() and not playing:
                sp.start_playback()
                playing = True
                changedThisTurn = True
            if global_state.shouldSkip():
                sp.pause_playback()
                break
            if not changedThisTurn:
                playing_track = sp.current_user_playing_track()
                if playing and not playing_track['is_playing']:
                    sp.pause_playback()
                    break
            time.sleep(1)
        global_state.songFinished()

    def getSongObjectFromSpotifyUri(self, global_state, song_uri):
        sp = spotipy.Spotify(auth=global_state.getSpotifyKey())
        song = sp.track(song_uri)
        return Song(
                    platform='spotify',
                    url=song['uri'],
                    title=song['name'],
                    artist=song['artists'][0]['name'],
                    artwork_url=song['album']['images'][0]['url'],
                )

    def getLikes(self, global_state):
        if self.cached_songs != None:
            return self.cached_songs
        sp = spotipy.Spotify(auth=global_state.getSpotifyKey())
        liked_tracks = sp.current_user_saved_tracks(limit=30)['items']
        songs = []
        for song in liked_tracks:
            song = song['track']
            songs.append(
                Song(
                    platform='spotify',
                    url=song['uri'],
                    title=song['name'],
                    artist=song['artists'][0]['name'],
                    artwork_url=song['album']['images'][0]['url'],
                )
        )
        self.cached_songs = {'songs': [s.dictRep() for s in songs]}
        return self.cached_songs


# -*- coding: utf-8 -*-
"""
Created on Tue Dec 22 21:03:17 2020

@author: hhh
"""

#%%
'''#%%创建将代码分为不同代码块，ctrl+enter执行当前代码块，
shift+enter执行并将光标移到下一代码块'''
#In[1]
import pandas as pd
import numpy as np
from sklearn.datasets import load_iris
import matplotlib.pyplot as plt
from math import exp
from time import *
from random import sample

#%%
#In[2]
'''读取数据'''
iris = load_iris() #加载iris数据集
df=pd.DataFrame(iris.data,columns=iris.feature_names) #将数据集填入列表中
df["label"]=iris.target

'''每一类别的不同个数'''
df.label.value_counts() #对于panda中的series对象，value_counts()方法返回每一列的不同频数

df1=df[:100]
#修改lebel使其变成-1或1
#败招：不能这样索引，否则改变不了dataframe的值
# for i in range(new_df.shape[0]):
#     new_df.loc[i][-1]=2*new_df.loc[i][-1]-1

#正招：需要将两层索引【i】【-1】改为一层，这样才能改变dataframe的值
df1["label"]=2*df1["label"]-1

#%%
#In[3]
# '''画未分类的二维散点图'''
# plt.scatter(df[:50]['sepal length (cm)'], df[:50]['sepal width (cm)'], label='0')
# plt.scatter(df[50:100]['sepal length (cm)'], df[50:100]['sepal width (cm)'], label='1')
# plt.xlabel('sepal length')
# plt.ylabel('sepal width')
# plt.legend() #图位于右上角“plots”

#%%
#In[3]
class my_svm:
    def __init__(self,train_data,lambda_0=1.0,max_iter=100,kernel_type="linear"):
        #data为一个列表，最后一列为标签
        self.data = train_data
        (self.m,self.n)=self.data.shape
        self.lambda_0 = lambda_0
        self.kernel = kernel_type
        self.max_iter = max_iter
        self.index=[i for i in range(self.m)] #需要抽样获取下标

        #模型参数
        self.X=[self.x(i) for i in range(self.m)]
        self.Y=[self.y(i) for i in range(self.m)]
        self.alpha = np.ones([self.m,1])#*self.lambda_0/2
        self.b=0
        self.E=self.E_()
        self.eps=1.0e-7
        self.length=50

    def x(self,i:int):
        '''将数据转化为向量自变量'''
        return np.transpose(np.array([self.data.loc[i][:-1]]))

    def y(self,i:int):
        '''将标签转化为因变量'''
        return self.data.loc[i][-1]

    def K(self,x:np.array,y:np.array):
        if self.kernel == "linear":
            return np.dot(np.transpose(x),y)
        elif self.kernel == "poly":
            return (np.dot(np.transpose(x),y)+1)**2
        elif self.kernel == "gauss":
            return exp(-np.dot(np.transpose(x-y),x-y)/2)

    def eta(self,i:int,j:int):
        '''计算Kii+Kjj-2Kij'''
        return self.K(self.X[i],self.X[i])+self.K(self.X[j],self.X[j])-2*self.K(self.X[i],self.X[j])

    def g(self,x:np.array):
        '''给定x，算w^(T)*phi(x)+b'''
        g_ = self.b
        for i in range(self.m):
            g_ += self.alpha[i]*self.Y[i]*self.K(self.X[i],x)
        return g_

    def _E(self,i:int):
        '''计算Ei'''
        return self.g(self.X[i])-self.Y[i]

    def E_(self):
        '''将所有的Ei保存到一个列表里'''
        list_ = [self._E(i) for i in range(self.m)];
        return list_

    def violation(self,i:int):
        # 败招1，会导致每次循环的判断时间大大延长
        # '''判断违反KKT的程度'''
        # if alpha[i]==0:
        #     return max(0,1-self.y(i)*self.g(alpha,self.x(i),b))
        # elif alpha[i]==self.lambda_0:
        #     return max(0,self.y(i)*self.g(alpha,self.x(i),b)-1)
        # elif (0<alpha[i])and(alpha[i]<self.lambda_0):
        #     return abs(self.y(i)*self.g(alpha,self.x(i),b)-1)

        if self.alpha[i] == 0:
            return (1 > self.Y[i]*self.g(self.X[i]))
        elif self.alpha[i] == self.lambda_0:
            return (self.Y[i]*self.g(self.X[i]) > 1)
        else:
            return (self.Y[i]*self.g(self.X[i]) != 1)

    def cycle(self):
        '''外层循环，选择第一个变量'''
        #等式下标集
        equal_list=[i for i in range(self.m) if 0<self.alpha[i]<self.lambda_0]
        #不等式下标集
        unequal_list=[i for i in range(self.m) if not(0<self.alpha[i]<self.lambda_0)]
        #合并
        _list=equal_list+unequal_list
        #从等式下标集开始考虑是否满足KKT条件
        for i in range(self.m):
            if self.violation(_list[i]):
                break
        i1 = _list[i]

        '''内层循环，选择第二个变量'''
        if self.E[i1]>=0:
            i2=min(range(self.m), key=lambda x: self.E[x])
        else:
            i2=max(range(self.m), key=lambda x: self.E[x])
        return i1,i2

    def cut(self,L:float,H:float,alpha_0:float):
        '''对alpha_0剪枝,L为下限，H为上限'''
        if alpha_0<L:
            alpha_0=L
        elif alpha_0>H:
            alpha_0=H
        return alpha_0

    def train(self):
        #误差序列，用于决定是否停止训练
        alpha_delta_list=[]

        for t in range(self.max_iter):
            # i1, i2 = self.cycle()
            begin_time = time()
            alpha_i2_delta=0
            k=0
            while alpha_i2_delta==0:
                samples = sample(self.index,2)
                i1, i2 = samples[0], samples[1]

                #得到L与H，并判断是否需要重新选择i1与i2
                if self.eta(i1,i2)<=0:
                    continue
                delta = (self.E[i1]-self.E[i2])/self.eta(i1,i2)
                if self.Y[i1] == self.Y[i2]:
                    L = max(self.alpha[i1]+self.alpha[i2]-self.lambda_0,0)
                    H = min(self.lambda_0,self.alpha[i1]+self.alpha[i2])
                else:
                    L = max(self.alpha[i2]-self.alpha[i1],0)
                    H = min(self.lambda_0,self.lambda_0+self.alpha[i2]-self.alpha[i1])

                #更新alpha，
                #注意将数组中的元素赋给一个值时，该值会随该元素变化而变化
                #若将数组中元素运算处理过后再赋给一个值时，该值则不随该元素变化
                alpha_i1_old = self.alpha[i1]/1.0
                alpha_i2_old = self.alpha[i2]/1.0
                alpha_i2_unc = alpha_i2_old+self.Y[i2]*delta
                alpha_i2_new = self.cut(L,H,alpha_i2_unc)
                alpha_i2_delta = alpha_i2_new-alpha_i2_old
            end_time = time()
            time1_ = end_time-begin_time
            alpha_delta_list.append(abs(alpha_i2_delta))
            #若连续几次的更新量均较小，则认为训练结束，可跳出循环
            if t>=(self.length+1):
                del(alpha_delta_list[0])
                if max(alpha_delta_list)<=self.eps:
                    break

            self.alpha[i2] = alpha_i2_new
            # print(self.alpha[i1].shape)
            # print(alpha_i2_old.shape)
            # print(alpha_i2_new.shape)
            self.alpha[i1] += self.Y[i1]*self.Y[i2]*(alpha_i2_old-self.alpha[i2])
            alpha_i1_delta = self.alpha[i1]-alpha_i1_old
            alpha_i2_delta = self.alpha[i2]-alpha_i2_old

            #更新b,首先计算b1与b2，再根据情况更新
            b1 = self.b-self.E[i1]
            b1 -= self.Y[i1]*self.K(self.X[i1],self.X[i1])*(self.alpha[i1]-alpha_i1_old)
            b1 -= self.Y[i2]*self.K(self.X[i1],self.X[i2])*(self.alpha[i2]-alpha_i2_old)
            b2 = self.b-self.E[i2]
            b2 -= self.Y[i1]*self.K(self.X[i1],self.X[i2])*(self.alpha[i1]-alpha_i1_old)
            b2 -= self.Y[i2]*self.K(self.X[i2],self.X[i2])*(self.alpha[i2]-alpha_i2_old)
            b_old = self.b
            if 0<self.alpha[i1]<self.lambda_0:
                self.b = b1
            elif 0<self.alpha[i2]<self.lambda_0:
                self.b = b2
            else:
                self.b = (b1+b2)/2

            # self.E[i1]=self._E(i1)
            # self.E[i2]=self._E(i2)
            b_delta = self.b-b_old

            # 更新Ei
            begin_time2=time()
            for i in range(self.m):
                self.E[i] += b_delta
                self.E[i] += self.Y[i1]*alpha_i1_delta*self.K(self.X[i],self.X[i1])
                self.E[i] += self.Y[i2]*alpha_i2_delta*self.K(self.X[i],self.X[i2])
            end_time2=time()
            time2_=end_time2-begin_time2
            print("b from",t,"th iteration:",self.b)
            print("i1 from",t,"th iteration:",i1)
            print("i2 from",t,"th iteration:",i2)
            print("delta from",t,"th iteration:",delta)
            print("alpha_i2_delta from",t,"th iteration:",alpha_i2_delta)
            print("alpha_i2_old:",alpha_i2_old)
            print("alpha_i2_new:",self.alpha[i2])
            print("time of choosing variables from",t,"th iteration:",time1_)
            # print("time of updating error from",t,"th iteration:",time2_)
            print("\n")
        print(t)
        return "what a shot!"

    def predict(self,x:np.array):
        if self.g(x)>=0:
            return 1
        else:
            return -1

    def score(self):
        score=0
        for i in range(self.m):
            if self.predict(self.X[i])==self.Y[i]:
                score+=1
        return score/self.m

#%%
begin_time=time()
model=my_svm(df1,max_iter=1000)
model.train()
end_time=time()
print("score=",model.score())
print("training time:",end_time-begin_time)


import telebot
from config_cb import keys,TOKEN
from extensions_cb import ConversionException, CryptoConverter

bot = telebot.TeleBot(TOKEN)

@bot.message_handler(commands=['start', 'help'])
def help(message: telebot.types.Message):
    text = "Привет!\nЯ бот, который конвертирует валюты в режиме реального времени!\
    \nДля начала работы задай Боту следующие параметры в одну строку через пробел: \n<имя валюты> \
    \n<в какую валюту перевести> \
    <количество валюты> \
    \nНАПРИМЕР: доллар евро 1 \
    \nСписок доступных валют можно увидеть по команде:\n /values"
    bot.reply_to(message,text)

@bot.message_handler(commands=['values'])
def values(message: telebot.types.Message):
    text = "Я могу конвертировать такие валюты:"
    for key in keys.keys():
        text = '\n'.join((text,key ))
    bot.reply_to(message,text)

@bot.message_handler(content_types=['text', ])
def convert(message: telebot.types.Message):
    try:
        values = message.text.split(' ')

        if len(values) != 3:
            raise ConversionException('Ошибка в параметрах, нужно три!\nпопробуй снова!')

        quote, base, amount = map(lambda x: x.lower(), values) #Название валют можно вводить с любым регистром
        total_base = CryptoConverter.convert(quote, base, amount)
        counter = float(values[2])
        total_summ = total_base * counter #подсчитываем количество валюты
    except ConversionException as e:
        bot.reply_to(message, f'Ошибка пользователя.\n{e}')
    except Exception as e:
        bot.reply_to(message, f'Не удалось обработать команду\n{e}')
    else:
        text = f'Стоимость {amount} {quote} в {base} равен {total_summ}'
        bot.send_message(message.chat.id, text)

bot.polling()