# ==============================================================================
# pyvale: the python validation engine
# License: MIT
# Copyright (C) 2025 The Computer Aided Validation Team
# ==============================================================================

"""
Riley: DIC UQ from Exodus
================================================================================
Here we render the same plate with a hole in tension from the last example but 
we load the data from an exodus file and generate the UVs using Riley's tools.
"""

import copy
from dataclasses import replace
from pathlib import Path

import numpy as np
import riley

import pyvale.dataio as io
from pyvale import render

# %%
# 1. Load Exodus data and build a textured surface mesh
# ------------------------------------------------------------
exodus_sim = riley.load_exodus(
    riley.data.platehole_exodus_path(),
    disp_keys=("disp_x", "disp_y", "disp_z"),
)
block = exodus_sim.elem_blocks["connect1"]
assert exodus_sim.disp is not None
simulation = io.SimData(
    coords=exodus_sim.coords,
    connect={"connect1": block.connect},
    node_vars={
        f"disp_{axis}": values
        for axis, values in zip("xyz", exodus_sim.disp, strict=True)
    },
)

texture = riley.load_texture_mono_u8(riley.data.speckle_texture_path())

surface_mesh = render.meshes3d_from_simdata(
    simulation,
    {"connect1": riley.ConnectConvention(
        block.elem_type, riley.EConnectAxis.ROW, 1, riley.ENodeOrder.EXODUS,
    )},
    displacement_keys=("disp_x", "disp_y", "disp_z"),
)["connect1"]

uv_plane = render.UVPlane(
    normal=np.array((0.0, 0.0, -1.0), dtype=np.float64),
    origin=np.array((0.0, 0.0, 0.0), dtype=np.float64),
)
uvs = render.uv_project_planar_centered(
    surface_mesh.coords,
    (2056, 2464),
    span=0.8,
    plane=uv_plane,
)
surface_mesh.shader = riley.TextureShader(uvs=uvs, texture=texture)

# To save time we only render the first and last frame. If you want to render
# all frames comment this out.
frame_indices = render.first_last_frame_indices(
    surface_mesh.displacements.shape[0]
)
surface_mesh.displacements = render.select_frames(
    surface_mesh.displacements, frame_indices
)

# %%
# 2. Create and position a distorted stereo camera pair
# ------------------------------------------------------------
pixels_num = (2464, 2056)
pixels_size = (3.45e-6, 3.45e-6)
focal_length = 50.0e-3
roi_centre = tuple(render.mesh_center(surface_mesh))

rot_world_0 = (0.0, 0.0, 0.0)
pos_world_0 = tuple(
    render.cam_pos_frame_points(
        surface_mesh.coords,
        pixels_num,
        pixels_size,
        focal_length,
        rot_world_0,
        fov_scale=0.65,
    )
)

rot_world_1 = (0.0, float(np.deg2rad(20.0)), 0.0)
pos_world_1 = tuple(
    render.cam_pos_frame_points(
        surface_mesh.coords,
        pixels_num,
        pixels_size,
        focal_length,
        rot_world_1,
        fov_scale=0.65,
    )
)

distortion_model = int(render.EDistortionModel.BROWN_CONRADY)

camera_0 = riley.Camera(
    pixels_num=pixels_num,
    pixels_size=pixels_size,
    pos_world=pos_world_0,
    rot_world=rot_world_0,
    roi_cent_world=roi_centre,
    focal_length=focal_length,
    sub_sample=2,
    distortion_model=distortion_model,
    distortion_k1=-0.2,
    distortion_k2=0.1,
    distortion_p1=0.0001,
    distortion_p2=-0.0001,
)

camera_1 = copy.deepcopy(camera_0)
camera_1.rot_world = rot_world_1
camera_1.pos_world = pos_world_1

cameras = [camera_0, camera_1]

# %%
# 3. Configure and build the renderer
# ------------------------------------------------------------

config = riley.create_raster_config(
    num_frames=surface_mesh.displacements.shape[0],
    total_threads=8,
    save_strategy=riley.SaveStrategy.disk,
)

config.background_value = 128.0
config.tile_size_max = 128
config.save_scaling = riley.ScaleStrategy.none

output_dir = (
    Path.cwd() / "pyvale-output"
    / "render3d_ex1e_riley_dic_from_exodus"
)

renderer = render.Riley(config, output_dir)

# %%
# 4. Build the scene and render the extracted surface
# ------------------------------------------------------------
result = renderer.render(
    render.Scene3D(meshes=[surface_mesh], cameras=cameras)
)

# %%
# 5. Save the stereo pair in Riley's exchange format
# ------------------------------------------------------------
riley.save_stereo_pair(
    str(output_dir),
    "stereo_data_opengl.csv",
    cameras[0],
    cameras[1],
)

riley.save_stereo_pair(
    str(output_dir),
    "stereo_data_opencv.csv",
    replace(cameras[0], coord_sys=riley.CameraCoordSys.opencv),
    replace(cameras[1], coord_sys=riley.CameraCoordSys.opencv),
)

print(f"Rendered Exodus driven DIC images to {output_dir}")
print(f"{result.images=}")

# %%
# The first frame from both cameras is combined side by side below.
#
# .. image:: ../../../../_static/render3d_ex1e_riley_dic_from_exodus.png
#    :alt: Riley stereo DIC render loaded from Exodus data
#    :width: 900px
#    :align: center