"""GX-TXT tutorial: fit a 2D point curve, build a profile, extrude it,
triangulate the extrusion, and publish the 3D mesh as a canonical geometry.mesh Dataset.

Designed for the GX-TXT Python Script Workspace / analysis.python_script workflow.

SPDX-License-Identifier: GPL-3.0-or-later
"""

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",
    "NUMERIC_API_v1@1:math.numeric+complex_linear_algebra",
    "DATASET_PUBLICATION_API_v2@2",
    "DATASET_FILE_PUBLICATION_API_v1@1",
    "SESSION_PERSISTENCE_API_v1@1",
])

ctx = gxtxt_context()
source_id = ctx["inputs"]["primary"]["dataset_id"]
if not source_id:
    raise RuntimeError("Profile tutorial: select one primary Dataset.")

print(f"GX-TXT Python profile tutorial; source Dataset={source_id}")

schema = gxtxt_api_call("DATASET_VARIABLE_SCHEMA_API_v3", "list", [source_id])

def variable_by_canonical(items, canonical):
    for item in items:
        if item.get("canonical") == canonical:
            return item
    raise RuntimeError(
        f"Profile tutorial: Dataset {source_id} has no canonical variable {canonical!r}."
    )

x_desc = variable_by_canonical(schema, "x_mm")
h_desc = variable_by_canonical(schema, "half_height_mm")
x_id = x_desc["variable_id"]
h_id = h_desc["variable_id"]

description = gxtxt_api_call("CANONICAL_DATASET_API_v2", "describe", [source_id])
row_count = int(description["row_count"])
block = gxtxt_api_call(
    "CANONICAL_DATASET_API_v2",
    "read_range",
    [source_id, [x_id, h_id], 0, row_count],
)

x = np.asarray(block[x_id]["values"], dtype=float).reshape(-1)
h = np.asarray(block[h_id]["values"], dtype=float).reshape(-1)
good = np.isfinite(x) & np.isfinite(h)
x = x[good]
h = h[good]

if len(x) < 6:
    raise RuntimeError("Profile tutorial: at least six finite points are required.")

order = np.argsort(x)
x = x[order]
h = h[order]

if np.any(np.diff(x) <= 0):
    raise RuntimeError("Profile tutorial: X coordinates must be strictly increasing.")
if np.any(h < 0):
    raise RuntimeError("Profile tutorial: half-height must be non-negative.")

x_min = float(x[0])
x_max = float(x[-1])
length = x_max - x_min
if not (length > 0):
    raise RuntimeError("Profile tutorial: invalid X span.")

u = (x - x_min) / length

def profile_basis(values):
    values = np.asarray(values, dtype=float).reshape(-1)
    q = values * (1.0 - values)
    return np.column_stack([
        q,
        q * values,
        q * values**2,
        q * values**3,
    ])

design = profile_basis(u)
rhs = h.reshape(-1, 1)

fit = gxtxt_api_call(
    "NUMERIC_API_v1",
    "linear_matrix",
    [
        "solve",
        {"real": design},
        {"real": rhs},
        {"rank_policy": "error"},
    ],
)

coefficients = np.asarray(fit["value_real"], dtype=float).reshape(-1)
if len(coefficients) != 4:
    raise RuntimeError("Profile tutorial: unexpected coefficient count.")

print("Profile coefficients:", coefficients)
print("Fit residual norm:", fit["residual_norm"])
print("Design condition number:", fit["condition_number"])

profile_samples = 81
x_profile = np.linspace(x_min, x_max, profile_samples)
u_profile = (x_profile - x_min) / length
evaluation_matrix = profile_basis(u_profile)

evaluation = gxtxt_api_call(
    "NUMERIC_API_v1",
    "linear_matrix",
    [
        "multiply",
        {"real": evaluation_matrix},
        {"real": coefficients.reshape(-1, 1)},
    ],
)

half_height = np.asarray(evaluation["value_real"], dtype=float).reshape(-1)
half_height = np.maximum(half_height, 0.0)
half_height[0] = 0.0
half_height[-1] = 0.0

profile_dataset_id = "tutorial.python.fitted_profile"
definition = {
    "dataset_id": profile_dataset_id,
    "contract": "table.generic",
    "storage": "auto",
    "replace_owned": True,
    "producer_version": "1.0.0",
    "source_dataset_ids": source_id,
    "source_variable_ids": [x_id, h_id],
    "notes": "Dense fitted half-profile used to generate the tutorial extrusion mesh.",
}

writer = gxtxt_api_call("DATASET_PUBLICATION_API_v2", "begin", [definition])
try:
    gxtxt_api_call(
        "DATASET_PUBLICATION_API_v2",
        "declare_dimension",
        [writer, "row", len(x_profile), {}],
    )

    x_schema = {
        "dtype": "float64",
        "rank": 1,
        "dimension_ids": "row",
        "coordinate_role": "coordinate",
        "semantic_role": "independent",
        "unit": x_desc.get("unit", ""),
    }
    h_schema = {
        "dtype": "float64",
        "rank": 1,
        "dimension_ids": "row",
        "coordinate_role": "data",
        "semantic_role": "derived",
        "unit": h_desc.get("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_half_height", h_schema],
    )
    gxtxt_api_call(
        "DATASET_PUBLICATION_API_v2", "write_variable",
        [writer, "v_x", x_profile, len(x_profile)],
    )
    gxtxt_api_call(
        "DATASET_PUBLICATION_API_v2", "write_variable",
        [writer, "v_half_height", half_height, len(half_height)],
    )

    published_profile = 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])

print("Published fitted profile Dataset:", published_profile["dataset_id"])

upper = np.column_stack([x_profile, half_height])
lower = np.column_stack([
    x_profile[-2:0:-1],
    -half_height[-2:0:-1],
])
ring_xy = np.vstack([upper, lower])
ring_count = len(ring_xy)

extrusion_depth = 12.0
z_front = -0.5 * extrusion_depth
z_back = +0.5 * extrusion_depth

front_ring = np.column_stack([
    ring_xy[:, 0],
    ring_xy[:, 1],
    np.full(ring_count, z_front),
])
back_ring = np.column_stack([
    ring_xy[:, 0],
    ring_xy[:, 1],
    np.full(ring_count, z_back),
])

center_xy = np.mean(ring_xy, axis=0)
vertices = np.vstack([
    front_ring,
    back_ring,
    [center_xy[0], center_xy[1], z_front],
    [center_xy[0], center_xy[1], z_back],
])

front_center = 2 * ring_count
back_center = front_center + 1

triangles = []
for i in range(ring_count):
    j = (i + 1) % ring_count
    triangles.append((i, j, ring_count + j))
    triangles.append((i, ring_count + j, ring_count + i))

for i in range(ring_count):
    j = (i + 1) % ring_count
    triangles.append((front_center, j, i))
    triangles.append((back_center, ring_count + i, ring_count + j))

triangles = np.asarray(triangles, dtype=np.int64)

if np.min(triangles) < 0 or np.max(triangles) >= len(vertices):
    raise RuntimeError("Profile tutorial: generated triangle index out of range.")

print(
    f"Mesh: vertices={len(vertices)} triangles={len(triangles)} "
    f"depth={extrusion_depth:g}"
)

mesh_file = Path(ctx["output_dir"]) / "extruded_profile_mesh.vtp"

def write_ascii_vtp(path, points, faces):
    connectivity = " ".join(
        str(int(index)) for face in faces for index in face
    )
    offsets = " ".join(
        str(3 * (i + 1)) for i in range(len(faces))
    )
    coordinates = " ".join(
        f"{float(value):.12g}" for point in points for value in point
    )

    document = f"""<?xml version="1.0"?>
<VTKFile type="PolyData" version="0.1" byte_order="LittleEndian">
  <PolyData>
    <Piece NumberOfPoints="{len(points)}"
           NumberOfVerts="0"
           NumberOfLines="0"
           NumberOfStrips="0"
           NumberOfPolys="{len(faces)}">
      <Points>
        <DataArray type="Float64" NumberOfComponents="3" format="ascii">
          {coordinates}
        </DataArray>
      </Points>
      <Polys>
        <DataArray type="Int32" Name="connectivity" format="ascii">
          {connectivity}
        </DataArray>
        <DataArray type="Int32" Name="offsets" format="ascii">
          {offsets}
        </DataArray>
      </Polys>
    </Piece>
  </PolyData>
</VTKFile>
"""
    path.write_text(document, encoding="utf-8")

write_ascii_vtp(mesh_file, vertices, triangles)

mesh_dataset_id = "tutorial.python.extruded_profile_mesh"
mesh_result = gxtxt_api_call(
    "DATASET_FILE_PUBLICATION_API_v1",
    "publish",
    [{
        "dataset_id": mesh_dataset_id,
        "contract": "geometry.mesh",
        "format": "vtp",
        "source_file": str(mesh_file),
        "storage_kind": "binary",
        "media_type": "application/vnd.vtk",
        "producer_version": "1.0.0",
        "source_dataset_ids": [source_id, profile_dataset_id],
        "source": "Python Script Workspace profile extrusion tutorial",
        "provenance": (
            "Triangle mesh created by fitting the source half-profile through "
            "NUMERIC_API_v1, mirroring it into a closed section, and extruding "
            "the section along Z."
        ),
        "notes": (
            f"vertices={len(vertices)}; triangles={len(triangles)}; "
            f"extrusion_depth={extrusion_depth:g}"
        ),
        "structured_metadata": {
            "geometry.kind": "extruded_profile_triangle_mesh",
            "geometry.vertex_count": str(len(vertices)),
            "geometry.triangle_count": str(len(triangles)),
            "geometry.source_profile_dataset": profile_dataset_id,
            "geometry.extrusion_depth": f"{extrusion_depth:.12g}",
            "geometry.coordinate_unit": x_desc.get("unit", ""),
        },
        "replace_owned": True,
    }],
)

print("Published 3D mesh Dataset:", mesh_result["dataset_id"])
print("Mesh SHA-256:", mesh_result["sha256"])

gxtxt_api_call("SESSION_PERSISTENCE_API_v1", "commit", [])

report = Path(ctx["output_dir"]) / "profile_mesh_report.txt"
report.write_text(
    "\n".join([
        "GX-TXT Python profile-to-mesh tutorial",
        "======================================",
        "",
        f"Source Dataset: {source_id}",
        f"Fitted profile Dataset: {profile_dataset_id}",
        f"Mesh Dataset: {mesh_dataset_id}",
        f"Input points: {len(x)}",
        f"Dense profile samples: {len(x_profile)}",
        f"Ring vertices: {ring_count}",
        f"Mesh vertices: {len(vertices)}",
        f"Mesh triangles: {len(triangles)}",
        f"Extrusion depth: {extrusion_depth:g} {x_desc.get('unit', '')}",
        f"Fit residual norm: {fit['residual_norm']}",
        f"Condition number: {fit['condition_number']}",
        "",
    ]),
    encoding="utf-8",
)

print("PASS PROFILE_MESH_TUTORIAL", source_id, profile_dataset_id, mesh_dataset_id)