r"""
Cyclic-symmetry modal on a rotor sector
=======================================

A rotor with :math:`N` identical sectors spanning 360° has a spectrum
that factors by *harmonic index* :math:`k \in \{0, 1, \dots, N/2\}`.
Analysing one base sector plus the constraint

.. math::

    u_\text{high} = e^{i k \alpha}\, R(\alpha)\, u_\text{low},
    \qquad \alpha = 2\pi / N,

recovers the full-rotor spectrum at a fraction of the DOF count.

This example builds a simple 4-sector rotor from a pyvista annular
sector, wraps it in a :class:`~femorph_solver.CyclicModel`, and runs
:meth:`~femorph_solver.CyclicModel.solve_modal` across all harmonic
indices.  ``CyclicModel`` auto-detects the low/high cyclic faces from
the rotation axis and angle, so the caller never has to assemble a
manual DOF pair list.
"""

from __future__ import annotations

import numpy as np
import pyvista as pv

import femorph_solver as fs
from femorph_solver import ELEMENTS

# %%
# Build a 90° annular sector
# --------------------------
N_SECTORS = 4
ALPHA = 2.0 * np.pi / N_SECTORS

N_T, N_R, N_Z = 12, 6, 3
theta = np.linspace(0.0, ALPHA, N_T + 1)
r = np.linspace(0.8, 1.0, N_R + 1)
z = np.linspace(0.0, 0.1, N_Z + 1)

# (r, θ, z) order gives a CCW-from-above hex winding in the (r, θ)
# plane — swapping r and θ would flip the winding inside-out.
R, T, Z = np.meshgrid(r, theta, z, indexing="ij")
X = R * np.cos(T)
Y = R * np.sin(T)

grid = pv.StructuredGrid(X, Y, Z).cast_to_unstructured_grid()

# %%
# Wrap as a CyclicModel
# ---------------------
# ``CyclicModel.from_grid`` builds the underlying base-sector
# :class:`~femorph_solver.Model`, detects the cyclic face pairs from
# the rotation axis (default ``"z"``) + ``n_sectors``, and validates
# that every low-face point has a rotated counterpart on the high
# face within ``pair_tol``.
cm = fs.CyclicModel.from_grid(grid, n_sectors=N_SECTORS)

# Per-sector physics — element kernel + isotropic material — is set
# on the inner Model exactly as for a non-cyclic analysis.
cm.assign(ELEMENTS.HEX8, material={"EX": 2.1e11, "PRXY": 0.30, "DENS": 7850.0})

# %%
# Solve every harmonic index
# --------------------------
# ``modal_solve`` defaults to all harmonics ``k = 0 … N/2``, returning
# one :class:`~femorph_solver.solvers.cyclic.CyclicModalResult` per
# harmonic.  Each result carries ``frequency`` and ``mode_shapes``.
results = cm.solve_modal(n_modes=4)

print(f"Cyclic modal on {N_SECTORS}-sector rotor:")
for res in results:
    fs_str = ", ".join(f"{f:.2f}" for f in res.frequency[:4])
    print(f"  k = {res.harmonic_index}: f [Hz] = [{fs_str}]")

# %%
# Visualise mode 0 of k = 1
# -------------------------
r1 = next(res for res in results if res.harmonic_index == 1)
phi = r1.mode_shapes[:, 0].reshape(-1, 3)
grid.point_data["mode_k1"] = np.abs(phi)

plotter = pv.Plotter(off_screen=True)
plotter.add_mesh(
    grid.warp_by_vector("mode_k1", factor=0.05 / (np.max(np.abs(phi)) + 1e-30)),
    scalars=np.linalg.norm(phi, axis=1),
    cmap="viridis",
    scalar_bar_args={"title": f"|mode| (k=1, f={r1.frequency[0]:.1f} Hz)"},
)
plotter.add_axes()
plotter.show()