"""
Module for parsing PicoTech PicoLog files, such as those generated by the TC-08
thermocouple reader.

Usage
`````
Available since ``yadg-5.1``.

.. autopydantic_model:: dgbowl_schemas.yadg.dataschema_5_1.filetype.PicoLog_tc08

Schema
``````
.. code-block:: yaml

    datatree.DataTree:
      coords:
        uts:            !!float           # Unix timestamp, optional
      data_vars:
        {{ name }}      (uts)             # Temperature of the named thermocouple

Metadata
````````
Metadata about the TC-08 device is read from the overall metadata file. The following
parameters are currently parsed:

- ``Model``: model number of the network analyseer
- ``Serial``: serial number of the device
- ``Software Version``: software version with which the file was created


Uncertainties
`````````````
Uncertainties in temperatures are set to 2.2°C, which is the usual value for a Type-K
thermocouple.

Notes on file structure
```````````````````````
The ``.picolog`` files are actually ``.tar`` archives, which include the following
files:

- a ``metadata.json`` file, containing information about the device and the
  thermocouples plugged in,
- a ``data-map.json.gz`` file, containing information about chunks of data corresponding
  to each thermocouple,
- a set of zipped binary files with data averaged at different time resolutions

Currently, the data is extracted from the ``1.0.gz`` file only.

.. codeauthor::
    Peter Kraus

"""

import os
import json
import numpy as np
import tarfile
import gzip
import tempfile
import xarray as xr
from datatree import DataTree

from yadg import dgutils


def extract(
    *,
    fn: str,
    **kwargs: dict,
) -> DataTree:
    with tarfile.open(fn, mode="r") as tf:
        with tempfile.TemporaryDirectory() as tempdir:
            tf.extractall(tempdir)

            # Get file metadata
            with open(os.path.join(tempdir, "metadata.json")) as inf:
                metadata = json.load(inf)

            # Get data from 1.0.gz
            with gzip.open(os.path.join(tempdir, "data-map.json.gz"), "rb") as inp:
                dmap = json.loads(inp.read())
            with gzip.open(os.path.join(tempdir, "1.0.gz"), "rb") as inp:
                raw = inp.read()

    uts = dgutils.str_to_uts(timestamp=metadata["startDate"], timezone=None)

    # Here we assume that "items" only has one element
    devdata = metadata["devices"]["items"][0]
    attrs = {
        "Software Version": metadata["appVersion"],
        "Model": devdata["__type"],
        "Serial": devdata["serialInternal"],
    }

    # Get necessary device metadata from the file metadata
    devices = {}
    for k, v in devdata.items():
        if isinstance(v, dict) and v["enabled"]:
            devices[k] = {
                "id": v["channelIdentifier"],
                "type": v["channelType"],
                "name": v["description"],
                "xmul": v["captureConfig"]["interval"] // 1000,
                "npts": v["capturedSamples"],
            }

    # Convert bytes to floats
    points = {}
    for tag, params in dmap:
        _, __, namestr = tag.split(".")
        id, res, time = namestr.split("/")
        archive, start, length = params
        if id not in points:
            points[id] = np.empty(0, dtype=">f4")
        if res == "1" and archive == "1.0.gz":
            new = np.frombuffer(raw, offset=start, dtype=">f4", count=length // 4)
            points[id] = np.concatenate((points[id], new))
        elif res == "1":
            raise RuntimeError(f"Resolution of 1 but archive is {archive!r}.")

    # Push the data into the Dataset
    ds = xr.Dataset(coords={"uts": (["uts"], [])})
    for id, data in points.items():
        for k, meta in devices.items():
            if meta["id"] == id:
                #  The type of the device should be thermocouple
                if meta["type"] == "thermocouple":
                    unit = {"units": "degC"}
                else:
                    raise RuntimeError("Unknown type {meta['type']!r}.")
                yvals = data[~np.isnan(data)]
                ydevs = np.ones(len(yvals)) * 2.2
                xvals = np.arange(len(yvals)) * meta["xmul"] + uts
                newds = xr.Dataset(
                    data_vars={
                        meta["name"]: (["uts"], yvals, unit),
                        f"{meta['name']}_std_err": (["uts"], ydevs, unit),
                    },
                    coords={"uts": (["uts"], xvals)},
                )
                ds = xr.merge((ds, newds))
    for var in ds.variables:
        if f"{var}_std_err" in ds.variables:
            ds[var].attrs["ancillary_variables"] = f"{var}_std_err"
        elif var.endswith("_std_err"):
            end = var.index("_std_err")
            if var[:end] in ds.variables:
                ds[var].attrs["standard_name"] = f"{var[:end]} standard_error"

    ds.attrs = dict(original_metadata=attrs)
    return DataTree(ds)
