import h5py
import numpy as np
import xml.etree.ElementTree as ET
from pathlib import Path

# ---------------------------------------------------------------
# Parameters - edit these to match your simulation
# ---------------------------------------------------------------

# Base point - the point about which the torsional moment is computed.
# Typically the centroid of the building footprint at ground level.
base_point = np.array([0.0, 0.0, 0.0])  # (x, y, z) in metres

# Nominal area and volume for normalisation.
# For a rectangular building of width b, depth d, and height h:
#   nominal_area   = b * h
#   nominal_volume = b * d * h
# Example for a building of 40x40x160
nominal_area   = 40*160     # m^2
nominal_volume = 40*40*160  # m^3

# File paths - adjust to the location of your downloaded files
base = Path(".")
body_h5   = base / "bodies.body_Cp body.h5"
body_xdmf = base / "bodies.body_Cp body.xdmf"
output_csv = base / "base_load_coefficients.csv"
# ---------------------------------------------------------------


# ----- Step 1: read mesh from XDMF / HDF5 -----
print("Reading mesh geometry from XDMF ...")
tree = ET.parse(body_xdmf)
root = tree.getroot()
domain = root.find("Domain")
collection = domain.find("Grid")
first_grid = collection.findall("Grid")[0]

# Topology (triangle connectivity)
topo_elem = first_grid.find("Topology")
topo_data = topo_elem.find("DataItem")
topo_ref = topo_data.text.strip()  # e.g. "file.h5:/path"
topo_h5_file, topo_h5_path = topo_ref.split(":")

# Geometry (vertex coordinates)
geom_elem = first_grid.find("Geometry")
geom_data = geom_elem.find("DataItem")
geom_ref = geom_data.text.strip()
geom_h5_file, geom_h5_path = geom_ref.split(":")

# Resolve HDF5 paths relative to the XDMF directory
xdmf_dir = body_xdmf.parent
topo_h5_path_full = xdmf_dir / topo_h5_file
geom_h5_path_full = xdmf_dir / geom_h5_file

with h5py.File(topo_h5_path_full, "r") as f:
    triangles = f[topo_h5_path][:]  # (N_tri, 3) - vertex indices

with h5py.File(geom_h5_path_full, "r") as f:
    vertices = f[geom_h5_path][:]  # (N_vert, 3) - xyz coordinates


# ----- Step 2: compute face normals, areas, and centroids -----
print("Computing face normals, areas, and centroids ...")
v0 = vertices[triangles[:, 0]]
v1 = vertices[triangles[:, 1]]
v2 = vertices[triangles[:, 2]]

edge1 = v1 - v0
edge2 = v2 - v0
cross = np.cross(edge1, edge2)  # (N_tri, 3)

# Area = 0.5 * |cross product|
area = 0.5 * np.linalg.norm(cross, axis=1)  # (N_tri,)

# Unit outward normal
normals = cross / (2.0 * area[:, np.newaxis])  # (N_tri, 3)

# Face centroid = mean of three vertices
centroids = (v0 + v1 + v2) / 3.0  # (N_tri, 3)

# Area-weighted normal vectors (A_i * n_hat_i)
area_normals = area[:, np.newaxis] * normals  # (N_tri, 3)

# Lever arm from base point to each face centroid
lever_arm = centroids - base_point[np.newaxis, :]  # (N_tri, 3)

n_tri = len(triangles)
print(f"  {n_tri} triangles")
print(f"  Nominal area   = {nominal_area} m^2")
print(f"  Nominal volume = {nominal_volume} m^3")


# ----- Step 3: read Cp and integrate force / moment coefficients -----
print("Reading Cp and integrating load coefficients ...")
with h5py.File(body_h5, "r") as f:
    if "cp" not in f:
        raise RuntimeError(
            "No 'cp' group found in the body HDF5 file. "
            "Run compute_cp.py first to add pressure coefficients."
        )
    cp_keys = sorted(f["cp"].keys(), key=lambda k: float(k[1:]))
    n_steps = len(cp_keys)

    # Pre-allocate output arrays
    times = np.empty(n_steps)
    Cfx = np.empty(n_steps)
    Cfy = np.empty(n_steps)
    CMz = np.empty(n_steps)

    for i, key in enumerate(cp_keys):
        cp = f[f"cp/{key}"][:]  # (N_tri,)
        times[i] = float(key[1:])  # strip the leading 't'

        # Force coefficient per face: cf_i = -cp_i * A_i * n_hat_i / A_nom
        # The negative sign accounts for pressure acting inward
        # on the surface (positive cp = compression = inward force).
        # Summing over all faces gives the global force coefficient.
        Cfx[i] = -(cp * area_normals[:, 0]).sum() / nominal_area
        Cfy[i] = -(cp * area_normals[:, 1]).sum() / nominal_area

        # Torsional moment coefficient:
        # CMz = sum( rx * cfy_i - ry * cfx_i ) / V_nom
        fx_per_face = -cp * area_normals[:, 0]
        fy_per_face = -cp * area_normals[:, 1]
        CMz[i] = (
            lever_arm[:, 0] * fy_per_face
            - lever_arm[:, 1] * fx_per_face
        ).sum() / nominal_volume

        if (i + 1) % 100 == 0 or (i + 1) == n_steps:
            print(f"  {i + 1}/{n_steps} timesteps", end="\r")

print(f"\nIntegrated {n_steps} timesteps over {n_tri} faces.")


# ----- Step 4: write results to CSV -----
print(f"Writing results to {output_csv} ...")
header = "time,Cfx,Cfy,CMz"
data = np.column_stack([times, Cfx, Cfy, CMz])
np.savetxt(output_csv, data, delimiter=",", header=header, comments="",
           fmt="%.6e")

print("Done.")
print(f"  Cfx range: [{Cfx.min():.4f}, {Cfx.max():.4f}]")
print(f"  Cfy range: [{Cfy.min():.4f}, {Cfy.max():.4f}]")
print(f"  CMz range: [{CMz.min():.4f}, {CMz.max():.4f}]")