"""Real Session research workflow using only public GX-TXT APIs."""
from concurrent.futures import ThreadPoolExecutor
import csv
from pathlib import Path

import numpy as np
from gxtxt import gxtxt_api_call, gxtxt_api_require, gxtxt_context

gxtxt_api_require([
    "CANONICAL_DATASET_API_v2@2",
    "DATASET_VARIABLE_SCHEMA_API_v3@3",
    "DATASET_PUBLICATION_API_v2@2",
    "PLOT_DATASET_API_v3@3:plot.dataset+plotspec",
    "PROGRESS_API_v1@1",
    "PARALLEL_EXECUTION_API_v1@1",
    "DATASET_FILE_PUBLICATION_API_v1@1",
])
ctx = gxtxt_context()
source_id = ctx["inputs"]["primary"]["dataset_id"]
assert source_id
schema = gxtxt_api_call("DATASET_VARIABLE_SCHEMA_API_v3", "list", [source_id])
x_id, y_id = schema[0]["variable_id"], schema[1]["variable_id"]
source = gxtxt_api_call("CANONICAL_DATASET_API_v2", "describe", [source_id])
block = gxtxt_api_call("CANONICAL_DATASET_API_v2", "read_range", [source_id, [x_id, y_id], 0, int(source["row_count"])])
x = np.asarray(block[x_id]["values"], dtype=float).reshape(-1)
y = np.asarray(block[y_id]["values"], dtype=float).reshape(-1)
assert len(x) == len(y) and len(x) > 0

policy = gxtxt_api_call("PARALLEL_EXECUTION_API_v1", "resources", [])
workers = min(2, int(policy["default_workers"]), int(policy["max_workers"])) if policy["enabled"] else 1
def compute(chunk):
    return chunk * 2.0
if workers > 1:
    chunks = np.array_split(y, workers)
    with ThreadPoolExecutor(max_workers=workers) as pool:
        derived = np.concatenate(list(pool.map(compute, chunks)))
else:
    derived = compute(y)
assert np.allclose(derived, y * 2.0)

dataset_id = "python.numeric.derived"
definition = {"dataset_id": dataset_id, "contract": "table.generic", "storage": "hdf5",
              "replace_owned": True, "source_dataset_ids": source_id, "source_variable_ids": y_id}
writer = gxtxt_api_call("DATASET_PUBLICATION_API_v2", "begin", [definition])
try:
    gxtxt_api_call("DATASET_PUBLICATION_API_v2", "declare_dimension", [writer, "row", len(x), {}])
    x_schema = {"dtype": "float64", "rank": 1, "dimension_ids": "row", "coordinate_role": "coordinate", "unit": schema[0]["unit"]}
    y_schema = {"dtype": "float64", "rank": 1, "dimension_ids": "row", "coordinate_role": "data", "unit": schema[1]["unit"]}
    gxtxt_api_call("DATASET_PUBLICATION_API_v2", "declare_variable", [writer, "v_x", x_schema])
    gxtxt_api_call("DATASET_PUBLICATION_API_v2", "declare_variable", [writer, "v_derived", y_schema])
    gxtxt_api_call("DATASET_PUBLICATION_API_v2", "write_variable", [writer, "v_x", x, len(x)])
    gxtxt_api_call("DATASET_PUBLICATION_API_v2", "write_variable", [writer, "v_derived", derived, len(derived)])
    published = gxtxt_api_call("DATASET_PUBLICATION_API_v2", "finalize", [writer])
    writer = None
finally:
    if writer is not None:
        gxtxt_api_call("DATASET_PUBLICATION_API_v2", "abort", [writer])
assert published["dataset_id"] == dataset_id

spec = gxtxt_api_call("PLOT_DATASET_API_v3", "new_spec", ["plot.dataset.xy", "Python derived data"])
spec["panels"][0]["series"] = [{"id": "derived", "dataset_id": dataset_id,
                                  "x_variable_id": "v_x", "y_variable_id": "v_derived",
                                  "axis": "left", "kind": "line"}]
reference = gxtxt_api_call("PLOT_DATASET_API_v3", "upsert_plot", ["python_result", spec])
gxtxt_api_call("PLOT_DATASET_API_v3", "render", [reference])

artifact = Path(ctx["output_dir"]) / "derived_measurements.csv"
with artifact.open("w", encoding="utf-8", newline="") as stream:
    csv_writer = csv.writer(stream)
    csv_writer.writerow(["source_x", "derived_y"])
    csv_writer.writerows(zip(x, derived))
file_result = gxtxt_api_call("DATASET_FILE_PUBLICATION_API_v1", "publish", [{
    "dataset_id": "python.numeric.csv", "contract": "table.csv", "format": "csv",
    "source_file": str(artifact), "storage_kind": "text", "media_type": "text/csv",
    "source_dataset_ids": source_id, "replace_owned": True,
    "provenance": "Python research script calculated y * 2 from the bound Dataset",
}])
assert file_result["dataset_id"] == "python.numeric.csv"
gxtxt_api_call("PROGRESS_API_v1", "publish", [{
    "state": "completed", "mode": "determinate", "current": 1, "total": 1,
    "module_id": ctx["module_id"], "run_id": ctx["run_id"],
    "message": "Published Python Dataset, graph and CSV artifact",
}])
gxtxt_api_call("SESSION_PERSISTENCE_API_v1", "commit", [])
print("PASS PYTHON_DATASET_WORKFLOW", source_id, dataset_id, reference["variant_id"], "workers", workers)