import os
import logging
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

logger = logging.getLogger()


def read_from_raw_file(abs_filepath: str) -> np.array:
    """
    Extract the data from the little-endian 32bit integer .raw file.
    Args:
        abs_filepath:
            The absolute filepath of the .raw file.
    Returns:
        A numpy array containing 32bit integers corresponding to ADC readings.
    """
    return np.fromfile(abs_filepath, np.dtype("<i4"))


def convert_to_mv(data: np.array, mv_span: int) -> np.array:
    """
    Converts an array of 32bit integers into their corresponding mV values.
    Args:
        data:
            Numpy array corresponding to the 32bit integer signal that will
            be converted to mV.
        mv_span:
            Millivolt span that will be used for the conversion.
    Returns:
        None.
    """
    return (data / np.power(2.0, 32)) * mv_span


def write_mv_data_to_parquet(
    shard: np.array,
    timestamps: np.array,
    timezone: int,
    target_data_directory: str,
    plant_id: str,
) -> None:
    """
    Writes a Numpy array of mV data into a partitioned parquet directory.

    This is the function that is in charge of actually writing data to
    parquet. The input `shard` array should be given in floating point
    millivolt format. The corresponding timestamp and partition columns are
    added to the datastructure before writing, according to the dataset
    schema (defined below).
    The parquet data schema is the following:
        'mV': float64
        'timestamp': timestamp('ms')
        'timezone': int16
        'year': int16
        'month': int16
        'day': int16
        'plant_id': string
    The partitioned parquet dataset directory structure generated by this
    function is the following:

    target_data_directory:
        |_ year=2020
        |   |_ month=3
        |   |   |_ day=10
        |   |   |   |_ plant_id=<PLANT_ID>
        |   |   |   |   |_ <SHARD_UUID>.parquet
        |   |   |   |   |_ <SHARD_UUID>.parquet
        |   |   |   |   |_ ...
        |   |   |_ day=11
        |   |   |_ ...
        |   |_ month=4
        |   |_ ...
        |_ year=2021
        |_ ...
    Args:
        shard:
            The numpy array containing 32bit integers corresponding
            to the signal to be written.
        timestamps:
            A numpy array of np.datetime64('ms') objects corresponding
            to the timestamps in the shard input parameter. These
            timestamps should be in UTC
        timezone:
            The timezone value to accompany the timestamps input parameter
        target_data_directory:
            The absolute path to the target directory to write the
            partitioned parquet dataset to.
        plant_id:
            The plant id to assign to the recording being processed.
            As a convention, the plant id follows the format:
                <GROWER/TRIAL>-<DEVICE_MAC_ADDRESS>-<CHANNEL_NUMBER>
    Returns:
        None.
    """

    logger.info(
        "Processing plant_id=%s mV shard %s--%s",
        plant_id,
        timestamps[0],
        timestamps[-1],
    )
    shard = shard.copy()

    daterange = pd.Series(timestamps).dt.round("ms")

    plant_ids = np.full(shard.shape, plant_id)
    timezones = np.full(shard.shape, timezone)

    data = [
        shard,
        daterange,
        timezones,
        daterange.dt.year,
        daterange.dt.month,
        daterange.dt.day,
        plant_ids,
    ]

    logger.debug(
        "Writing plant_id=%s mV shard %s--%s to Dataset",
        plant_id,
        timestamps[0],
        timestamps[-1],
    )
    psr_schema = pa.schema(
        [
            pa.field("mV", pa.float64()),
            pa.field("timestamp", pa.timestamp("ms")),
            pa.field("timezone", pa.int16()),
            pa.field("year", pa.int16()),
            pa.field("month", pa.int16()),
            pa.field("day", pa.int16()),
            pa.field("plant_id", pa.string()),
        ]
    )
    table = pa.Table.from_arrays(data, schema=psr_schema)

    pq.write_to_dataset(
        table,
        root_path=target_data_directory,
        partition_cols=["year", "month", "day", "plant_id"],
        compression="GZIP",
    )

    logger.debug(
        "DONE plant_id=%s shard %s--%s", plant_id, timestamps[0], timestamps[-1]
    )
