import random
import subprocess
from time import perf_counter as clock
from pathlib import Path

import numpy as np

# Constants

STEP = 1000 * 100  # the size of the buffer to fill the table, in rows
SCALE = 0.1  # standard deviation of the noise compared with actual
# values
NI_NTIMES = 1  # The number of queries for doing a mean (non-idx cols)
# COLDCACHE = 10   # The number of reads where the cache is considered 'cold'
# WARMCACHE = 50   # The number of reads until the cache is considered 'warmed'
# READ_TIMES = WARMCACHE+50    # The number of complete calls to DB.query_db()
# COLDCACHE = 50   # The number of reads where the cache is considered 'cold'
# WARMCACHE = 50   # The number of reads until the cache is considered 'warmed'
# READ_TIMES = WARMCACHE+50    # The number of complete calls to DB.query_db()
MROW = 1000 * 1000

# Test values
COLDCACHE = 5  # The number of reads where the cache is considered 'cold'
WARMCACHE = 5  # The number of reads until the cache is considered 'warmed'
READ_TIMES = 10  # The number of complete calls to DB.query_db()

# global variables
rdm_cod = ["lin", "rnd"]
prec = 6  # precision for printing floats purposes


def get_nrows(nrows_str):
    powers = {"k": 3, "m": 6, "g": 9}
    try:
        return int(float(nrows_str[:-1]) * 10 ** powers[nrows_str[-1]])
    except KeyError:
        raise ValueError(
            "value of nrows must end with either 'k', 'm' or 'g' suffixes."
        )


class DB:

    def __init__(self, nrows, rng, userandom):
        self.step = STEP
        self.scale = SCALE
        self.rng = rng
        self.userandom = userandom
        self.filename = Path("-".join([rdm_cod[userandom], nrows]))
        self.nrows = get_nrows(nrows)

    def get_db_size(self):
        sout = subprocess.run(
            "sync;du -s %s" % self.filename, shell=True, stdout=subprocess.PIPE
        ).stdout
        return int(sout.split(maxsplit=1)[0])

    def print_mtime(self, t1, explain):
        mtime = clock() - t1
        print(f"{explain}: {mtime:.6f}")
        print(f"Krows/s: {self.nrows / 1000 / mtime:.6f}")

    def print_qtime(self, colname, ltimes):
        qtime1 = ltimes[0]  # First measured time
        qtime2 = ltimes[-1]  # Last measured time
        print(f"Query time for {colname}: {qtime1:.6f}")
        print(f"Mrows/s: {self.nrows / MROW / qtime1:.6f}")
        print(f"Query time for {colname} (cached): {qtime2:.6f}")
        print(f"Mrows/s (cached): {self.nrows / MROW / qtime2:.6f}")

    def norm_times(self, ltimes):
        "Get the mean and stddev of ltimes, avoiding the extreme values."
        lmean = ltimes.mean()
        lstd = ltimes.std()
        ntimes = ltimes[ltimes < lmean + lstd]
        nmean = ntimes.mean()
        nstd = ntimes.std()
        return nmean, nstd

    def print_qtime_idx(self, colname, ltimes, repeated, verbose):
        if repeated:
            r = "[REP] "
        else:
            r = "[NOREP] "
        ltimes = np.array(ltimes)
        _ = len(ltimes)
        qtime1 = ltimes[0]  # First measured time
        ctimes = ltimes[1:COLDCACHE]
        cmean, cstd = self.norm_times(ctimes)
        wtimes = ltimes[WARMCACHE:]
        wmean, wstd = self.norm_times(wtimes)
        if verbose:
            print("Times for cold cache:\n", ctimes)
            # print "Times for warm cache:\n", wtimes
            hist1, hist2 = np.histogram(wtimes)
            print(f"Histogram for warm cache: {hist1}\n{hist2}")
        print(f"{r}1st query time for {colname}: {qtime1:.{prec}f}")
        print(
            f"{r}Query time for {colname} (cold cache): "
            f"{cmean:.{prec}f} +- {cstd:.{prec}f}"
        )
        print(
            f"{r}Query time for {colname} (warm cache): "
            f"{wmean:.{prec}f} +- {wstd:.{prec}f}"
        )

    def print_db_sizes(self, init, filled, indexed):
        table_size = (filled - init) / 1024
        indexes_size = (indexed - filled) / 1024
        print(f"Table size (MB): {table_size:.3f}")
        print(f"Indexes size (MB): {indexes_size:.3f}")
        print(f"Full size (MB): {table_size + indexes_size:.3f}")

    def fill_arrays(self, start, stop):
        arr_f8 = np.arange(start, stop, dtype="float64")
        arr_i4 = np.arange(start, stop, dtype="int32")
        if self.userandom:
            arr_f8 += np.random.normal(0, stop * self.scale, size=stop - start)
            arr_i4 = np.array(arr_f8, dtype="int32")
        return arr_i4, arr_f8

    def create_db(self, dtype, kind, optlevel, verbose):
        self.con = self.open_db(remove=1)
        self.create_table(self.con)
        init_size = self.get_db_size()
        t1 = clock()
        self.fill_table(self.con)
        table_size = self.get_db_size()
        self.print_mtime(t1, "Insert time")
        self.index_db(dtype, kind, optlevel, verbose)
        indexes_size = self.get_db_size()
        self.print_db_sizes(init_size, table_size, indexes_size)
        self.close_db(self.con)

    def index_db(self, dtype, kind, optlevel, verbose):
        if dtype == "int":
            idx_cols = ["col2"]
        elif dtype == "float":
            idx_cols = ["col4"]
        else:
            idx_cols = ["col2", "col4"]
        for colname in idx_cols:
            t1 = clock()
            self.index_col(self.con, colname, kind, optlevel, verbose)
            self.print_mtime(t1, "Index time (%s)" % colname)

    def query_db(
        self,
        niter,
        dtype,
        onlyidxquery,
        onlynonidxquery,
        avoidfscache,
        verbose,
        inkernel,
    ):
        self.con = self.open_db()
        if dtype == "int":
            reg_cols = ["col1"]
            idx_cols = ["col2"]
        elif dtype == "float":
            reg_cols = ["col3"]
            idx_cols = ["col4"]
        else:
            reg_cols = ["col1", "col3"]
            idx_cols = ["col2", "col4"]
        if avoidfscache:
            rseed = int(np.random.randint(self.nrows))
        else:
            rseed = 19
        # Query for non-indexed columns
        np.random.seed(rseed)
        base = np.random.randint(self.nrows)
        if not onlyidxquery:
            for colname in reg_cols:
                ltimes = []
                random.seed(rseed)
                for i in range(NI_NTIMES):
                    t1 = clock()
                    results = self.do_query(self.con, colname, base, inkernel)
                    ltimes.append(clock() - t1)
                if verbose:
                    print("Results len:", results)
                self.print_qtime(colname, ltimes)
            # Always reopen the file after *every* query loop.
            # Necessary to make the benchmark to run correctly.
            self.close_db(self.con)
            self.con = self.open_db()
        # Query for indexed columns
        if not onlynonidxquery:
            for colname in idx_cols:
                ltimes = []
                np.random.seed(rseed)
                rndbase = np.random.randint(self.nrows, size=niter)
                # First, non-repeated queries
                for i in range(niter):
                    base = rndbase[i]
                    t1 = clock()
                    results = self.do_query(self.con, colname, base, inkernel)
                    # results, tprof = self.do_query(
                    #    self.con, colname, base, inkernel)
                    ltimes.append(clock() - t1)
                if verbose:
                    print("Results len:", results)
                self.print_qtime_idx(colname, ltimes, False, verbose)
                # Always reopen the file after *every* query loop.
                # Necessary to make the benchmark to run correctly.
                self.close_db(self.con)
                self.con = self.open_db()
                ltimes = []
                # Second, repeated queries
                #                 for i in range(niter):
                #                     t1=time()
                #                     results = self.do_query(
                #                         self.con, colname, base, inkernel)
                # results, tprof = self.do_query(self.con, colname, base, inkernel)
                #                     ltimes.append(time()-t1)
                #                 if verbose:
                #                     print "Results len:", results
                #                 self.print_qtime_idx(colname, ltimes, True, verbose)
                # Print internal PyTables index tprof statistics
                # tprof = np.array(tprof)
                # tmean, tstd = self.norm_times(tprof)
                # print "tprof-->", round(tmean, prec), "+-", round(tstd, prec)
                # print "tprof hist-->", \
                #    np.histogram(tprof)
                # print "tprof raw-->", tprof
                # Always reopen the file after *every* query loop.
                # Necessary to make the benchmark to run correctly.
                self.close_db(self.con)
                self.con = self.open_db()
        # Finally, close the file.
        self.close_db(self.con)

    def close_db(self, con):
        con.close()


if __name__ == "__main__":
    import sys
    import getopt

    try:
        import psyco

        psyco_imported = 1
    except Exception:
        psyco_imported = 0

    usage = (
        """usage: %s [-T] [-P] [-v] [-f] [-k] [-p] [-m] [-c] [-q] [-i] [-I] [-S] [-x] [-z complevel] [-l complib] [-R range] [-N niter] [-n nrows] [-d datadir] [-O level] [-t kind] [-s] col -Q [suplim]
            -T use Pytables
            -P use Postgres
            -v verbose
            -f do a profile of the run (only query functionality)
            -k do a profile for kcachegrind use (out file is 'indexed_search.kcg')
            -p use "psyco" if available
            -m use random values to fill the table
            -q do a query (both indexed and non-indexed versions)
            -i do a query (just indexed one)
            -I do a query (just in-kernel one)
            -S do a query (just standard one)
            -x choose a different seed for random numbers (i.e. avoid FS cache)
            -c create the database
            -z compress with zlib (no compression by default)
            -l use complib for compression (zlib used by default)
            -R select a range in a field in the form "start,stop" (def "0,10")
            -N number of iterations for reading
            -n sets the number of rows (in krows) in each table
            -d directory to save data (default: data.nobackup)
            -O set the optimization level for PyTables indexes
            -t select the index type: "medium" (default) or "full", "light", "ultralight"
            -s select a type column for operations ('int' or 'float'. def all)
            -Q do a repeteated query up to 10**value
            \n"""
        % sys.argv[0]
    )

    try:
        opts, pargs = getopt.getopt(
            sys.argv[1:], "TPvfkpmcqiISxz:l:R:N:n:d:O:t:s:Q:"
        )
    except Exception:
        sys.stderr.write(usage)
        sys.exit(1)

    # default options
    usepytables = 0
    usepostgres = 0
    verbose = 0
    doprofile = 0
    dokprofile = 0
    usepsyco = 0
    userandom = 0
    docreate = 0
    optlevel = 0
    kind = "medium"
    docompress = 0
    complib = "zlib"
    doquery = False
    onlyidxquery = False
    onlynonidxquery = False
    inkernel = True
    avoidfscache = 0
    # rng = [-10, 10]
    rng = [-1000, -1000]
    repeatquery = 0
    repeatvalue = 0
    krows = "1k"
    niter = READ_TIMES
    dtype = "all"
    datadir = "data.nobackup"

    # Get the options
    for option in opts:
        if option[0] == "-T":
            usepytables = 1
        elif option[0] == "-P":
            usepostgres = 1
        elif option[0] == "-v":
            verbose = 1
        elif option[0] == "-f":
            doprofile = 1
        elif option[0] == "-k":
            dokprofile = 1
        elif option[0] == "-p":
            usepsyco = 1
        elif option[0] == "-m":
            userandom = 1
        elif option[0] == "-c":
            docreate = 1
        elif option[0] == "-q":
            doquery = True
        elif option[0] == "-i":
            doquery = True
            onlyidxquery = True
        elif option[0] == "-I":
            doquery = True
            onlynonidxquery = True
        elif option[0] == "-S":
            doquery = True
            onlynonidxquery = True
            inkernel = False
        elif option[0] == "-x":
            avoidfscache = 1
        elif option[0] == "-z":
            docompress = int(option[1])
        elif option[0] == "-l":
            complib = option[1]
        elif option[0] == "-R":
            rng = [int(i) for i in option[1].split(",")]
        elif option[0] == "-N":
            niter = int(option[1])
        elif option[0] == "-n":
            krows = option[1]
        elif option[0] == "-d":
            datadir = option[1]
        elif option[0] == "-O":
            optlevel = int(option[1])
        elif option[0] == "-t":
            if option[1] in ("full", "medium", "light", "ultralight"):
                kind = option[1]
            else:
                print(
                    "kind should be either 'full', 'medium', 'light' or "
                    "'ultralight'"
                )
                sys.exit(1)
        elif option[0] == "-s":
            if option[1] in ("int", "float"):
                dtype = option[1]
            else:
                print("column should be either 'int' or 'float'")
                sys.exit(1)
        elif option[0] == "-Q":
            repeatquery = 1
            repeatvalue = int(option[1])

    # If not database backend selected, abort
    if not usepytables and not usepostgres:
        print("Please select a backend:")
        print("PyTables: -T")
        print("Postgres: -P")
        sys.exit(1)

    # Create the class for the database
    if usepytables:
        from pytables_backend import PyTablesDB

        db = PyTablesDB(
            krows, rng, userandom, datadir, docompress, complib, kind, optlevel
        )
    elif usepostgres:
        from postgres_backend import PostgresDB

        db = PostgresDB(krows, rng, userandom)

    if not avoidfscache:
        # in order to always generate the same random sequence
        np.random.seed(20)

    if verbose:
        if userandom:
            print("using random values")
        if onlyidxquery:
            print("doing indexed queries only")

    if psyco_imported and usepsyco:
        psyco.bind(db.create_db)
        psyco.bind(db.query_db)

    if docreate:
        if verbose:
            print("writing %s rows" % krows)
        db.create_db(dtype, kind, optlevel, verbose)

    if doquery:
        print("Calling query_db() %s times" % niter)
        if doprofile:
            import pstats
            import cProfile

            cProfile.run(
                "db.query_db(niter, dtype, onlyidxquery, onlynonidxquery, "
                "avoidfscache, verbose, inkernel)",
                "indexed_search.prof",
            )
            stats = pstats.Stats("indexed_search.prof")
            stats.strip_dirs()
            stats.sort_stats("time", "calls")
            if verbose:
                stats.print_stats()
            else:
                stats.print_stats(20)
        elif dokprofile:
            import cProfile

            import lsprofcalltree

            prof = cProfile.Profile()
            prof.run(
                "db.query_db(niter, dtype, onlyidxquery, onlynonidxquery, "
                "avoidfscache, verbose, inkernel)"
            )
            kcg = lsprofcalltree.KCacheGrind(prof)
            with Path("indexed_search.kcg").open("w") as ofile:
                kcg.output(ofile)
        elif doprofile:
            import hotshot
            import hotshot.stats

            prof = hotshot.Profile("indexed_search.prof")
            benchtime, stones = prof.run(
                "db.query_db(niter, dtype, onlyidxquery, onlynonidxquery, "
                "avoidfscache, verbose, inkernel)"
            )
            prof.close()
            stats = hotshot.stats.load("indexed_search.prof")
            stats.strip_dirs()
            stats.sort_stats("time", "calls")
            stats.print_stats(20)
        else:
            db.query_db(
                niter,
                dtype,
                onlyidxquery,
                onlynonidxquery,
                avoidfscache,
                verbose,
                inkernel,
            )

    if repeatquery:
        # Start by a range which is almost None
        db.rng = [1, 1]
        if verbose:
            print("range:", db.rng)
        db.query_db(
            niter,
            dtype,
            onlyidxquery,
            onlynonidxquery,
            avoidfscache,
            verbose,
            inkernel,
        )
        for i in range(repeatvalue):
            for j in (1, 2, 5):
                rng = j * 10**i
                db.rng = [-rng / 2, rng / 2]
                if verbose:
                    print("range:", db.rng)
                #                 if usepostgres:
                #                     os.system(
                #                         "echo 1 > /proc/sys/vm/drop_caches;"
                #                         " /etc/init.d/postgresql restart")
                #                 else:
                #                     os.system("echo 1 > /proc/sys/vm/drop_caches")
                db.query_db(
                    niter,
                    dtype,
                    onlyidxquery,
                    onlynonidxquery,
                    avoidfscache,
                    verbose,
                    inkernel,
                )
