# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Lance Authors
import csv
import os
import shutil
from datetime import datetime

import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.parquet as pq
from lance.file import LanceFileReader, LanceFileWriter

"""
Data generation (in-memory)
import duckdb
import pyarrow.parquet as pq

con = duckdb.connect(database=":memory:")
con.execute("INSTALL tpch; LOAD tpch")
con.execute("CALL dbgen(sf=1)")

tables = ["lineitem", "nation", "region", "part", "supplier", "customer", "partsupp",
    "orders"]
for t in tables:
    res = con.query("SELECT * FROM " + t)
    pq.write_table(res.to_arrow_table(), tpch_dir_path + t + ".parquet")
"""


def measure_pyarrow_read_time(path, num_trials, verbose=False):
    """
    Measures the time required to read a parquet file using pyarrow,
    averaged over multiple trials
    """
    parquet_read_time = 0
    for trial in range(1, num_trials + 1):
        start = datetime.now()
        pa.parquet.read_table(path)
        end = datetime.now()
        parquet_read_time += (end - start).total_seconds()

        if verbose:
            print(
                f"Parquet Read Time for trial {trial}: {(end - start).total_seconds()}s"
            )

    avg_read_time = parquet_read_time / num_trials

    file_size_bytes = os.path.getsize(path)
    file_size_mb = file_size_bytes / (1024 * 1024)

    print(
        f"Parquet Read Time: {avg_read_time:.2f}s for file {path} of size: "
        f"{file_size_mb:.2f} MiB"
    )
    return avg_read_time


def measure_parquet_write_time(table, path, num_trials, verbose=False):
    """
    Measures the time required to write a parquet file using pyarrow
    """
    if os.path.exists(path):
        if os.path.isdir(path):
            shutil.rmtree(path)
        else:
            os.remove(path)

    parquet_write_time = 0
    for trial in range(1, num_trials + 1):
        start = datetime.now()
        pq.write_table(table, path, compression="snappy")
        end = datetime.now()
        parquet_write_time += (end - start).total_seconds()

        if verbose:
            print(
                f"Parquet Write Time for trial {trial}: \
                    {(end - start).total_seconds()}s"
            )

    avg_write_time = parquet_write_time / num_trials

    file_size_bytes = os.path.getsize(path)
    file_size_mb = file_size_bytes / (1024 * 1024)

    print(
        f"Parquet Write Time: {avg_write_time:.2f}s for file {path} of size: "
        f"{file_size_mb:.2f} MiB"
    )
    return avg_write_time


def measure_lance_read_time(
    path, num_trials, batch_size, verbose=False, lance_version="2.0"
):
    """
    Measures the time required to read a lance file using lance,
    averaged over multiple trials
    """
    lance_read_time = 0
    for trial in range(1, num_trials + 1):
        start = datetime.now()
        reader = LanceFileReader(path)
        for batch in reader.read_all(batch_size=batch_size).to_batches():
            pass
        end = datetime.now()

        lance_read_time += (end - start).total_seconds()

        if verbose:
            print(
                f"Lance Read Time for trial {trial}: {(end - start).total_seconds()}s"
            )

    avg_read_time = lance_read_time / num_trials

    file_size_bytes = os.path.getsize(path)
    file_size_mb = file_size_bytes / (1024 * 1024)

    print(
        f"Lance {lance_version} Read Time: {avg_read_time:.2f}s for file {path} "
        f"of size: {file_size_mb:.2f} MiB"
    )
    return avg_read_time


def measure_lance_write_time(
    dataset, path, num_trials=10, verbose=False, lance_version="2.0"
):
    """
    Takes a lance dataset object as input, and writes a lance
    file to disk in a directory
    """

    if os.path.exists(path):
        if os.path.isdir(path):
            shutil.rmtree(path)
        else:
            os.remove(path)

    lance_write_time = 0
    for trial in range(1, num_trials + 1):
        start = datetime.now()
        with LanceFileWriter(path, dataset.schema, version=lance_version) as writer:
            for batch in dataset.to_batches():
                writer.write_batch(batch)
        end = datetime.now()
        lance_write_time += (end - start).total_seconds()

        if verbose:
            print(
                f"Lance Write Time for trial {trial}: {(end - start).total_seconds()}s"
            )

    avg_write_time = lance_write_time / num_trials
    if os.path.exists(path):
        file_size_bytes = os.path.getsize(path)
        file_size_mb = file_size_bytes / (1024 * 1024)

        print(
            f"Lance {lance_version} Write Time: {avg_write_time:.2f}s for file {path} "
            f"of size: {file_size_mb:.2f} MiB"
        )
    else:
        raise (f"Lance {lance_version} write file {path} failed")

    return avg_write_time


def benchmark_tpch_lineitem(
    dataset_dir, num_trials=10, verbose=False, lance_version="2.0"
):
    name = benchmark_tpch_lineitem.__name__
    dataset_path = dataset_dir + "lineitem.parquet"
    output = csv.writer(open(name + ".csv", "w"))

    dataset = ds.dataset(dataset_path)
    num_rows = dataset.count_rows()
    parquet_size = os.path.getsize(dataset_path)
    parquet_size_mb = parquet_size / (1024 * 1024)

    pyarrow_read_time = measure_pyarrow_read_time(dataset_path, num_trials=10)

    dataset = ds.dataset(dataset_path)
    lance_path = "/tmp/tpch.lanceV" + lance_version
    measure_lance_write_time(
        dataset, path=lance_path, num_trials=1, lance_version=lance_version
    )
    lance_size = os.path.getsize(lance_path)
    lance_size_mb = lance_size / (1024 * 1024)

    lance_read_time = measure_lance_read_time(
        lance_path, num_trials=10, batch_size=1024 * 8, lance_version=lance_version
    )

    output.writerows(
        [
            [
                "name",
                "dataset",
                "num_trials",
                "num_rows",
                "parquet_size",
                "lance_size",
                "pyarrow_read_time",
                "lance_read_time",
            ],
            [
                name,
                "tpch_lineitem",
                num_trials,
                num_rows,
                f"{parquet_size_mb:.2f} MiB",
                f"{lance_size_mb:.2f} MiB",
                f"{pyarrow_read_time:.2f} seconds",
                f"{lance_read_time:.2f} seconds",
            ],
        ]
    )


def benchmark_tpch_encodings(
    dataset_dir,
    dataset_name,
    encoding_type,
    num_trials=10,
    verbose=False,
    lance_version="2.0",
):
    """
    Loads numeric columns from TPCH tables and benchmarks the read times for
    parquet and lance files
    """

    benchmark_name = benchmark_tpch_encodings.__name__
    output = csv.writer(
        open(benchmark_name + "_" + dataset_name + "_" + encoding_type + ".csv", "w")
    )
    output.writerow(
        [
            "name",
            "dataset",
            "lance_encoding_type",
            "num_trials",
            "num_rows",
            "parquet_size",
            "lance_size",
            "parquet_write_time",
            "lance_write_time",
            "parquet_read_time",
            "lance_read_time",
        ]
    )

    tables = [
        "nation",
        "region",
        "part",
        "supplier",
        "customer",
        "partsupp",
        "orders",
        "lineitem",
    ]

    for table_name in tables:
        print("Table: {}".format(table_name))
        PATH = dataset_dir + table_name + ".parquet"

        orig_dataset = ds.dataset(PATH)
        orig_schema = orig_dataset.schema

        if encoding_type == "plain_numeric":
            target_types = [pa.int32(), pa.int64(), pa.float32(), pa.float64()]
        elif encoding_type == "plain_non_numeric":
            target_types = [pa.string()]
        elif encoding_type == "plain_timestamp":
            target_types = [pa.date32()]
        elif encoding_type == "integer":
            target_types = [
                pa.int8(),
                pa.int16(),
                pa.int32(),
                pa.int64(),
                pa.uint8(),
                pa.uint16(),
                pa.uint32(),
                pa.uint64(),
            ]

        target_columns = [
            field.name for field in orig_schema if field.type in target_types
        ]

        if len(target_columns) == 0:
            continue
        table = orig_dataset.to_table(columns=target_columns)
        parquet_path = "/tmp/parquet_ds.parquet"

        parquet_write_time = measure_parquet_write_time(
            table, parquet_path, num_trials=num_trials
        )

        dataset = ds.dataset(parquet_path)
        num_rows = dataset.count_rows()
        parquet_size = os.path.getsize(parquet_path)
        parquet_size_mb = parquet_size / (1024 * 1024)

        parquet_read_time = measure_pyarrow_read_time(
            parquet_path, num_trials=num_trials
        )

        lance_path = "/tmp/tpch.lanceV" + lance_version
        lance_write_time = measure_lance_write_time(
            table, lance_path, num_trials=num_trials, lance_version=lance_version
        )
        lance_size = os.path.getsize(lance_path)
        lance_size_mb = lance_size / (1024 * 1024)

        lance_read_time = measure_lance_read_time(
            lance_path,
            num_trials=num_trials,
            batch_size=1024 * 8,
            lance_version=lance_version,
        )

        output.writerow(
            [
                benchmark_name,
                dataset_name + "_" + table_name,
                encoding_type,
                num_trials,
                num_rows,
                f"{parquet_size_mb:.2f} MiB",
                f"{lance_size_mb:.2f} MiB",
                f"{parquet_write_time:.2f} seconds",
                f"{lance_write_time:.2f} seconds",
                f"{parquet_read_time:.2f} seconds",
                f"{lance_read_time:.2f} seconds",
            ]
        )

        os.remove(parquet_path)

        if os.path.isdir(lance_path):
            shutil.rmtree(lance_path)
        else:
            os.remove(lance_path)


"""
Data generation (sift)
    nvecs = 1000000
    ndims = 128
    with open("../sift/sift_base.fvecs", mode="rb") as fobj:
        buf = fobj.read()
        data = np.array(struct.unpack("<128000000f",
                        buf[4 : 4 + 4 * nvecs * ndims])).reshape((nvecs, ndims))
        dd = dict(zip(range(nvecs), data))

    table = vec_to_table(dd)
    parquet_path = "/tmp/sift1m.parquet"

"""


def benchmark_sift_vector_encodings(
    dataset_dir,
    dataset_name,
    encoding_type,
    num_trials=10,
    verbose=False,
    lance_version="2.0",
):
    benchmark_name = benchmark_sift_vector_encodings.__name__
    output = csv.writer(open(benchmark_name + "_" + encoding_type + ".csv", "w"))
    output.writerow(
        [
            "name",
            "dataset",
            "lance_encoding_type",
            "num_trials",
            "num_rows",
            "parquet_size",
            "lance_size",
            "parquet_write_time",
            "lance_write_time",
            "parquet_read_time",
            "lance_read_time",
        ]
    )

    PATH = dataset_dir + "sift1m.parquet"
    orig_dataset = ds.dataset(PATH)

    table = orig_dataset.to_table()
    parquet_path = "/tmp/parquet_ds.parquet"

    parquet_write_time = measure_parquet_write_time(
        table, parquet_path, num_trials=num_trials
    )

    dataset = ds.dataset(parquet_path)
    num_rows = dataset.count_rows()
    parquet_size = os.path.getsize(parquet_path) / (1024**3)

    parquet_read_time = measure_pyarrow_read_time(parquet_path, num_trials=num_trials)

    lance_path = "/tmp/tpch.lanceV" + lance_version
    lance_write_time = measure_lance_write_time(
        table, lance_path, num_trials=num_trials, lance_version=lance_version
    )
    lance_size = os.path.getsize(lance_path) / (1024**3)

    lance_read_time = measure_lance_read_time(
        lance_path,
        num_trials=num_trials,
        batch_size=1024 * 8,
        lance_version=lance_version,
    )

    output.writerow(
        [
            benchmark_name,
            dataset_name,
            encoding_type,
            num_trials,
            num_rows,
            parquet_size,
            parquet_write_time,
            parquet_read_time,
            lance_size,
            lance_write_time,
            lance_read_time,
        ]
    )


# Should contain all table files
tpch_dir_path = "/tmp/TPCH_SF1/"
dataset_name = "tpch_sf1"

if __name__ == "__main__":
    # Version "2.1" test is temporarily disabled due to currently issues when
    # activate it
    for lance_version in ["2.0"]:
        benchmark_tpch_lineitem(tpch_dir_path, lance_version)
        benchmark_tpch_encodings(
            tpch_dir_path,
            dataset_name=dataset_name,
            encoding_type="plain_numeric",
            lance_version=lance_version,
        )
        benchmark_tpch_encodings(
            tpch_dir_path,
            dataset_name=dataset_name,
            encoding_type="plain_non_numeric",
            lance_version=lance_version,
        )
        benchmark_tpch_encodings(
            tpch_dir_path,
            dataset_name=dataset_name,
            encoding_type="plain_timestamp",
            lance_version=lance_version,
        )
        benchmark_tpch_encodings(
            tpch_dir_path,
            dataset_name=dataset_name,
            encoding_type="integer",
            lance_version=lance_version,
        )

        """
        benchmark_sift_vector_encodings(
            "/home/ubuntu/test/",
            dataset_name="sift1m",
            encoding_type="plain_fixed_size_list",
            num_trials=5,
            lance_version=lance_version,
        )
        """
