//! The Turek–Hron flag as a 3-D solid: a Hex20 plate `[x0, x1] × [y0, y1]
//! × [z0, z1]` (length × thickness × span), root face `x = x0` clamped,
//! for R8's coupled 3-D FSI (omni-cortex roadmap, item R8-b).
//!
//! Nothing here is a new solver: [`NonlinearDynamicAnalysis`] and the
//! total-Lagrangian St. Venant–Kirchhoff path are dimension-generic, so
//! the 3-D flag steps through exactly the machinery the 2-D harness uses
//! (`set_nodal_forces`, `step(state) → state`, Newmark γ/β, the
//! line-search/subdivision rescue). This module adds what the 2-D harness
//! builds by hand:
//!
//! * the structured Hex20 mesh on the `(2nx+1) × (2ny+1) × (2nz+1)`
//!   serendipity lattice (x-major node order, so the sequential DOF
//!   numbering keeps the banded LU's bandwidth at one x-slab);
//! * the root clamp, with the lateral faces either **free** (the real
//!   plate, the physical 3-D problem) or **plane strain** (`u_z = 0` at
//!   every node: a z-independent field is exactly representable by Hex20
//!   and the 3-D energy then equals span × the 2-D plane-strain energy, so
//!   this reproduces the 2-D Quad8 model to rounding — the self-consistency
//!   gate, "run the 3-D problem as the 2-D problem first");
//! * the wetted surface (bottom, top, tip, and the two lateral faces) as
//!   Quad8 faces with outward orientation, and the consistent nodal forces
//!   of a traction field integrated over the *current* (deformed) faces —
//!   the interface load a partitioned coupling feeds to
//!   [`NonlinearDynamicStepper::set_nodal_forces`](super::NonlinearDynamicStepper::set_nodal_forces).
//!
//! New code only: no existing solver path changes.

use super::{AnalysisConfig, ConvergenceCriteria, NonlinearDynamicAnalysis};
use crate::assembly::dof_mapping::DofComponent;
use crate::boundary::dirichlet::{DirichletBC, DirichletType};
use crate::boundary::{BoundaryCondition, BoundaryConditionSet, SpatialFunction};
use crate::error::FeaResult;
use crate::materials::{LinearElastic, MaterialDatabase};
use crate::mesh::{Element, ElementType, MaterialId, Mesh, Node, NodeId};
use nalgebra::{DVector, Vector3};
use std::collections::BTreeMap;

/// Geometry and resolution of the plate. `nx`, `ny`, `nz` are Hex20
/// element counts along length, thickness and span.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Flag3dSpec {
    pub x0: f64,
    pub x1: f64,
    pub y0: f64,
    pub y1: f64,
    pub z0: f64,
    pub z1: f64,
    pub nx: usize,
    pub ny: usize,
    pub nz: usize,
}

impl Flag3dSpec {
    /// The Turek–Hron flag (`[0.25, 0.6] × [0.19, 0.21]`) extruded over
    /// `z ∈ [z0, z0 + span]`.
    pub fn turek_hron(span: f64, z0: f64, nx: usize, ny: usize, nz: usize) -> Self {
        Self {
            x0: 0.25,
            x1: 0.6,
            y0: 0.19,
            y1: 0.21,
            z0,
            z1: z0 + span,
            nx,
            ny,
            nz,
        }
    }

    pub fn span(&self) -> f64 {
        self.z1 - self.z0
    }
}

/// How the lateral faces `z = z0, z1` are held.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LateralFaces {
    /// Traction-free: the physical 3-D plate.
    Free,
    /// `u_z = 0` at every node: exact 2-D plane strain (the 2-D model's
    /// definition), for the self-consistency gate.
    PlaneStrain,
}

/// Which wetted face of the plate a [`SurfaceFace`] lies on.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum FlagSide {
    /// `y = y0`
    Bottom,
    /// `y = y1`
    Top,
    /// `x = x1`
    Tip,
    /// `z = z0`
    SideLow,
    /// `z = z1`
    SideHigh,
}

/// A Quad8 face of the Hex20 mesh: corners counter-clockwise seen from
/// outside, then mid-edge nodes `(0-1, 1-2, 2-3, 3-0)` — the library's
/// Quad8 order, oriented so `∂x/∂ξ × ∂x/∂η` points out of the solid.
#[derive(Debug, Clone, PartialEq)]
pub struct SurfaceFace {
    pub side: FlagSide,
    pub nodes: [NodeId; 8],
}

/// The structured Hex20 plate.
#[derive(Debug, Clone)]
pub struct Flag3d {
    pub spec: Flag3dSpec,
    pub mesh: Mesh,
    lattice: Vec<Option<NodeId>>,
    dims: [usize; 3],
}

const GAUSS3: [(f64, f64); 3] = [
    (-0.774_596_669_241_483_4, 5.0 / 9.0),
    (0.0, 8.0 / 9.0),
    (0.774_596_669_241_483_4, 5.0 / 9.0),
];

/// Quad8 serendipity shape functions and their `(ξ, η)` derivatives, the
/// library's node order.
fn quad8(xi: f64, eta: f64) -> ([f64; 8], [[f64; 2]; 8]) {
    let corners = [(-1.0, -1.0), (1.0, -1.0), (1.0, 1.0), (-1.0, 1.0)];
    let mut n = [0.0; 8];
    let mut d = [[0.0; 2]; 8];
    for (a, &(xa, ya)) in corners.iter().enumerate() {
        let (p, q) = (1.0 + xa * xi, 1.0 + ya * eta);
        let r = xa * xi + ya * eta - 1.0;
        n[a] = 0.25 * p * q * r;
        d[a][0] = 0.25 * xa * (q * r + p * q);
        d[a][1] = 0.25 * ya * (p * r + p * q);
    }
    // Mid-edge nodes: (0, -1), (1, 0), (0, 1), (-1, 0).
    n[4] = 0.5 * (1.0 - xi * xi) * (1.0 - eta);
    d[4] = [-xi * (1.0 - eta), -0.5 * (1.0 - xi * xi)];
    n[5] = 0.5 * (1.0 + xi) * (1.0 - eta * eta);
    d[5] = [0.5 * (1.0 - eta * eta), -(1.0 + xi) * eta];
    n[6] = 0.5 * (1.0 - xi * xi) * (1.0 + eta);
    d[6] = [-xi * (1.0 + eta), 0.5 * (1.0 - xi * xi)];
    n[7] = 0.5 * (1.0 - xi) * (1.0 - eta * eta);
    d[7] = [-0.5 * (1.0 - eta * eta), -(1.0 - xi) * eta];
    (n, d)
}

impl Flag3d {
    /// Build the mesh (material 0 on every element).
    pub fn build(spec: Flag3dSpec) -> FeaResult<Self> {
        assert!(
            spec.nx > 0 && spec.ny > 0 && spec.nz > 0,
            "element counts must be positive"
        );
        let dims = [2 * spec.nx + 1, 2 * spec.ny + 1, 2 * spec.nz + 1];
        let mut mesh = Mesh::new(3)?;
        let mut lattice = vec![None; dims[0] * dims[1] * dims[2]];
        for i in 0..dims[0] {
            for j in 0..dims[1] {
                for k in 0..dims[2] {
                    if (i % 2) + (j % 2) + (k % 2) > 1 {
                        continue; // not a serendipity node
                    }
                    let x = spec.x0 + (spec.x1 - spec.x0) * i as f64 / (dims[0] - 1) as f64;
                    let y = spec.y0 + (spec.y1 - spec.y0) * j as f64 / (dims[1] - 1) as f64;
                    let z = spec.z0 + (spec.z1 - spec.z0) * k as f64 / (dims[2] - 1) as f64;
                    lattice[(i * dims[1] + j) * dims[2] + k] =
                        Some(mesh.add_node(Node::new_3d(x, y, z)));
                }
            }
        }
        let mut flag = Self {
            spec,
            mesh,
            lattice,
            dims,
        };
        for ex in 0..spec.nx {
            for ey in 0..spec.ny {
                for ez in 0..spec.nz {
                    let (a, b, c) = (2 * ex, 2 * ey, 2 * ez);
                    let at = |i, j, k| flag.lattice_node(i, j, k).expect("serendipity node");
                    let nodes = vec![
                        at(a, b, c),
                        at(a + 2, b, c),
                        at(a + 2, b + 2, c),
                        at(a, b + 2, c),
                        at(a, b, c + 2),
                        at(a + 2, b, c + 2),
                        at(a + 2, b + 2, c + 2),
                        at(a, b + 2, c + 2),
                        at(a + 1, b, c),
                        at(a + 2, b + 1, c),
                        at(a + 1, b + 2, c),
                        at(a, b + 1, c),
                        at(a + 1, b, c + 2),
                        at(a + 2, b + 1, c + 2),
                        at(a + 1, b + 2, c + 2),
                        at(a, b + 1, c + 2),
                        at(a, b, c + 1),
                        at(a + 2, b, c + 1),
                        at(a + 2, b + 2, c + 1),
                        at(a, b + 2, c + 1),
                    ];
                    flag.mesh.add_element(Element::new(
                        ElementType::Hex20,
                        nodes,
                        MaterialId(0),
                    )?)?;
                }
            }
        }
        Ok(flag)
    }

    /// The node at lattice index `(i, j, k)` (`0..=2n` per direction), if
    /// that lattice point carries a serendipity node.
    pub fn lattice_node(&self, i: usize, j: usize, k: usize) -> Option<NodeId> {
        if i >= self.dims[0] || j >= self.dims[1] || k >= self.dims[2] {
            return None;
        }
        self.lattice[(i * self.dims[1] + j) * self.dims[2] + k]
    }

    /// Lattice sizes `(2nx+1, 2ny+1, 2nz+1)`.
    pub fn lattice_dims(&self) -> [usize; 3] {
        self.dims
    }

    /// The node nearest to `p` (reference coordinates).
    pub fn nearest_node(&self, p: Vector3<f64>) -> NodeId {
        self.mesh
            .nodes
            .iter()
            .min_by(|a, b| {
                let da = (a.1.position() - p).norm();
                let db = (b.1.position() - p).norm();
                da.partial_cmp(&db).unwrap()
            })
            .map(|(&id, _)| id)
            .expect("non-empty mesh")
    }

    /// The Turek–Hron point A `(x1, (y0+y1)/2)` on the mid-span line
    /// (the nearest lattice node: exact for even `nz`, else the nearer of
    /// the two mid-span candidates).
    pub fn point_a(&self) -> NodeId {
        let s = &self.spec;
        self.nearest_node(Vector3::new(s.x1, 0.5 * (s.y0 + s.y1), 0.5 * (s.z0 + s.z1)))
    }

    /// Nodes on the clamped root face `x = x0`.
    pub fn root_nodes(&self) -> Vec<NodeId> {
        let mut out = Vec::new();
        for j in 0..self.dims[1] {
            for k in 0..self.dims[2] {
                if let Some(id) = self.lattice_node(0, j, k) {
                    out.push(id);
                }
            }
        }
        out
    }

    /// The root clamp (`u = 0` on `x = x0`) plus the lateral condition.
    pub fn clamp_root(&self, lateral: LateralFaces) -> BoundaryConditionSet {
        let zero = || DirichletType::Spatial(SpatialFunction(Box::new(|_| 0.0)));
        let dirichlet = |nodes: Vec<NodeId>, component| {
            BoundaryCondition::Dirichlet(DirichletBC {
                nodes,
                components: vec![component],
                condition_type: zero(),
                time_range: None,
                ramping_factor: 1.0,
                gradual_enforcement: false,
            })
        };
        let root = self.root_nodes();
        let mut set = BoundaryConditionSet::new();
        for component in [
            DofComponent::DisplacementX,
            DofComponent::DisplacementY,
            DofComponent::DisplacementZ,
        ] {
            set.add_condition(dirichlet(root.clone(), component));
        }
        if lateral == LateralFaces::PlaneStrain {
            let all: Vec<NodeId> = self.mesh.nodes.keys().copied().collect();
            set.add_condition(dirichlet(all, DofComponent::DisplacementZ));
        }
        set
    }

    /// One linear-elastic material (the TL path reads its Lamé pair and
    /// density).
    pub fn materials(e: f64, nu: f64, rho: f64) -> MaterialDatabase {
        let mut db = MaterialDatabase::new();
        db.add_material(
            MaterialId(0),
            LinearElastic::new(e, nu).with_density(rho),
            None,
        );
        db
    }

    /// The coupled march's structure: total-Lagrangian SVK, Newmark with
    /// the given `γ` and `β = (γ + ½)²/4`, 60 Newton iterations — the 2-D
    /// harness's settings. `num_steps` only matters for
    /// [`NonlinearDynamicAnalysis::run`]; a coupling uses
    /// [`NonlinearDynamicAnalysis::stepper`].
    pub fn dynamic_analysis(
        &self,
        e: f64,
        nu: f64,
        rho: f64,
        lateral: LateralFaces,
        dt: f64,
        num_steps: usize,
        gamma: f64,
    ) -> NonlinearDynamicAnalysis {
        let beta = (gamma + 0.5).powi(2) / 4.0;
        NonlinearDynamicAnalysis::new(
            self.mesh.clone(),
            Self::materials(e, nu, rho),
            self.clamp_root(lateral),
            dt,
            num_steps,
            AnalysisConfig::default(),
        )
        .with_total_lagrangian()
        .with_convergence_criteria(ConvergenceCriteria {
            max_iterations: 60,
            ..ConvergenceCriteria::default()
        })
        .with_newmark_parameters(gamma, beta)
    }

    /// The Quad8 faces of the given sides (the root face is never
    /// wetted). Order: side, then element index.
    pub fn surface_faces(&self, sides: &[FlagSide]) -> Vec<SurfaceFace> {
        let [di, dj, dk] = self.dims;
        let (ilast, jlast, klast) = (di - 1, dj - 1, dk - 1);
        let mut faces = Vec::new();
        // base lattice point + the two in-face axes (ξ, η) as unit steps.
        let mut push = |side, base: [usize; 3], u: [usize; 3], v: [usize; 3]| {
            let p = |cu: usize, cv: usize| {
                let q = [
                    base[0] + cu * u[0] + cv * v[0],
                    base[1] + cu * u[1] + cv * v[1],
                    base[2] + cu * u[2] + cv * v[2],
                ];
                self.lattice_node(q[0], q[1], q[2]).expect("face node")
            };
            faces.push(SurfaceFace {
                side,
                nodes: [
                    p(0, 0),
                    p(2, 0),
                    p(2, 2),
                    p(0, 2),
                    p(1, 0),
                    p(2, 1),
                    p(1, 2),
                    p(0, 1),
                ],
            });
        };
        const X: [usize; 3] = [1, 0, 0];
        const Y: [usize; 3] = [0, 1, 0];
        const Z: [usize; 3] = [0, 0, 1];
        for &side in sides {
            match side {
                FlagSide::Bottom => {
                    for a in (0..ilast).step_by(2) {
                        for c in (0..klast).step_by(2) {
                            push(side, [a, 0, c], X, Z); // x × z = −y
                        }
                    }
                }
                FlagSide::Top => {
                    for a in (0..ilast).step_by(2) {
                        for c in (0..klast).step_by(2) {
                            push(side, [a, jlast, c], Z, X); // z × x = +y
                        }
                    }
                }
                FlagSide::Tip => {
                    for b in (0..jlast).step_by(2) {
                        for c in (0..klast).step_by(2) {
                            push(side, [ilast, b, c], Y, Z); // y × z = +x
                        }
                    }
                }
                FlagSide::SideLow => {
                    for a in (0..ilast).step_by(2) {
                        for b in (0..jlast).step_by(2) {
                            push(side, [a, b, 0], Y, X); // y × x = −z
                        }
                    }
                }
                FlagSide::SideHigh => {
                    for a in (0..ilast).step_by(2) {
                        for b in (0..jlast).step_by(2) {
                            push(side, [a, b, klast], X, Y); // x × y = +z
                        }
                    }
                }
            }
        }
        faces
    }

    /// Every wetted face: bottom, top, tip, and both lateral faces.
    pub fn wetted_faces(&self) -> Vec<SurfaceFace> {
        self.surface_faces(&[
            FlagSide::Bottom,
            FlagSide::Top,
            FlagSide::Tip,
            FlagSide::SideLow,
            FlagSide::SideHigh,
        ])
    }

    /// Consistent nodal forces `f_a = ∫ N_a t(x, n) da` of a traction
    /// field over the given faces in the configuration `X + u`
    /// (`displacement` in the stepper's global DOF numbering, read through
    /// `node_dofs`; `None` = the reference configuration). `traction`
    /// receives the current point and the current outward unit normal —
    /// a pressure `p` is `|x, n| -p(x) * n`. 3 × 3 Gauss per face.
    /// Returns one entry per touched node, sorted by node id.
    pub fn face_nodal_forces(
        &self,
        faces: &[SurfaceFace],
        displacement: Option<(&DVector<f64>, &dyn Fn(NodeId) -> Vec<usize>)>,
        traction: &dyn Fn(Vector3<f64>, Vector3<f64>) -> Vector3<f64>,
    ) -> Vec<(NodeId, Vector3<f64>)> {
        let position = |id: NodeId| -> Vector3<f64> {
            let x = self.mesh.get_node(id).expect("face node").position();
            match displacement {
                Some((u, dofs)) => {
                    let d = dofs(id);
                    x + Vector3::new(u[d[0]], u[d[1]], u[d[2]])
                }
                None => x,
            }
        };
        let mut out: BTreeMap<NodeId, Vector3<f64>> = BTreeMap::new();
        for face in faces {
            let xs: Vec<Vector3<f64>> = face.nodes.iter().map(|&id| position(id)).collect();
            for &(xi, wx) in &GAUSS3 {
                for &(eta, wy) in &GAUSS3 {
                    let (n, d) = quad8(xi, eta);
                    let mut x = Vector3::zeros();
                    let mut t1 = Vector3::zeros();
                    let mut t2 = Vector3::zeros();
                    for a in 0..8 {
                        x += xs[a] * n[a];
                        t1 += xs[a] * d[a][0];
                        t2 += xs[a] * d[a][1];
                    }
                    let cross = t1.cross(&t2);
                    let jac = cross.norm();
                    let t = traction(x, cross / jac);
                    let w = wx * wy * jac;
                    for a in 0..8 {
                        *out.entry(face.nodes[a]).or_insert_with(Vector3::zeros) += t * (n[a] * w);
                    }
                }
            }
        }
        out.into_iter().collect()
    }
}