//! R8-a: the embedded3 3D cut-cell fluid as the fluid side of a
//! partitioned FSI loop. The body is the embedded flag test's (the circle
//! wall to wall united with a capsule of half-thickness 10 mm around the
//! flag's centreline, its apex on the benchmark's tip A), but the
//! centreline is no longer prescribed: the harness sets it per coupled
//! step from the 2D structure (span-uniform; the centreline's element-corner
//! nodes at reference x 0.25 … 0.59, so the capsule's apex sits on A as the
//! flag test's tip inset puts it), as the pair of lines at the
//! step's start and end; the body's φ and surface velocity at any time in
//! between are the linear blend of the two (the solver asks at the step's
//! two ends only).
//!
//! The loads: the operator route (`Mask::cut_wall_force`) with the R8-a
//! load sink installed — every contribution the route sums, with its
//! position, is returned for the harness to distribute onto the flag.

use std::cell::RefCell;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock};

use rtx_cfd::solvers::incompressible::ConvectionScheme;
use rtx_cfd::solvers::incompressible::embedded3::exchange::set_load_sink;
use rtx_cfd::solvers::incompressible::embedded3::step::device::{
    DeviceSnapshot, DeviceStep, loads_device_enabled,
};
use rtx_cfd::solvers::incompressible::embedded3::{
    Body, Boundaries, DeviceSdf, Field, Fluid, Grid, Parameters, Side, Solver, StepResult,
    WallScheme,
};

pub const H: f64 = 0.41;
pub const L: f64 = 2.5;
pub const CX: f64 = 0.2;
pub const CY: f64 = 0.2;
pub const R_CYL: f64 = 0.05;
/// The capsule's half-thickness (the flag's half-thickness).
pub const HALF: f64 = 0.01;
pub const RHO: f64 = 1000.0;
pub const NU: f64 = 1e-3;
/// The flag test's CFL velocity (its dt convention, kept for comparability).
const U_CFL: f64 = 2.25;

/// A centreline at time `t`: points (x, y) and their velocities.
#[derive(Clone, Debug)]
pub struct Line {
    pub t: f64,
    pub pts: Vec<[f64; 2]>,
    pub vel: Vec<[f64; 2]>,
}

/// The step's two lines.
pub struct Lines {
    pub a: Line,
    pub b: Line,
}

impl Lines {
    /// The blend at `t` (the start line before `a.t`, the end line after `b.t`).
    fn at(&self, t: f64) -> Line {
        let eps = 1e-12 * (1.0 + t.abs());
        if (t - self.b.t).abs() <= eps || t >= self.b.t || self.b.t <= self.a.t {
            return self.b.clone();
        }
        if t <= self.a.t + eps {
            return self.a.clone();
        }
        let s = (t - self.a.t) / (self.b.t - self.a.t);
        let mix = |p: &[[f64; 2]], q: &[[f64; 2]]| -> Vec<[f64; 2]> {
            p.iter()
                .zip(q)
                .map(|(p, q)| [p[0] + s * (q[0] - p[0]), p[1] + s * (q[1] - p[1])])
                .collect()
        };
        assert_eq!(
            self.a.pts.len(),
            self.b.pts.len(),
            "blend needs equal point counts"
        );
        Line {
            t,
            pts: mix(&self.a.pts, &self.b.pts),
            vel: mix(&self.a.vel, &self.b.vel),
        }
    }
}

/// Bumped on every `set_lines` (the host closures' per-thread cache key).
static VERSION: AtomicU64 = AtomicU64::new(0);

/// The capsule's signed distance at (x, y) and the centreline velocity at
/// the closest point (the flag test's `flag_2d_recorded`, cached per thread
/// and time).
fn capsule(lines: &RwLock<Lines>, x: f64, y: f64, t: f64) -> (f64, (f64, f64)) {
    thread_local! {
        static CACHE: RefCell<(u64, u64, Vec<[f64; 4]>)> = const { RefCell::new((u64::MAX, u64::MAX, Vec::new())) };
    }
    let ver = VERSION.load(Ordering::Acquire);
    CACHE.with(|cell| {
        let mut c = cell.borrow_mut();
        if c.0 != t.to_bits() || c.1 != ver {
            let line = lines.read().expect("lines").at(t);
            c.2 = line
                .pts
                .iter()
                .zip(&line.vel)
                .map(|(p, v)| [p[0], p[1], v[0], v[1]])
                .collect();
            c.0 = t.to_bits();
            c.1 = ver;
        }
        let pts = &c.2;
        let mut best = f64::INFINITY;
        let mut v_best = (0.0, 0.0);
        for m in 0..pts.len() - 1 {
            let [ax, ay, avx, avy] = pts[m];
            let [bx, by, bvx, bvy] = pts[m + 1];
            let (ex, ey) = (bx - ax, by - ay);
            let l2 = ex * ex + ey * ey;
            if l2 == 0.0 {
                continue;
            }
            let u = (((x - ax) * ex + (y - ay) * ey) / l2).clamp(0.0, 1.0);
            let (px, py) = (ax + u * ex, ay + u * ey);
            let d = ((x - px).powi(2) + (y - py).powi(2)).sqrt();
            if d < best {
                best = d;
                v_best = (avx + u * (bvx - avx), avy + u * (bvy - avy));
            }
        }
        (best - HALF, v_best)
    })
}

fn cylinder(x: f64, y: f64) -> f64 {
    ((x - CX).powi(2) + (y - CY).powi(2)).sqrt() - R_CYL
}

/// One load contribution: position, component, part (0 pressure, 1 wall
/// shear, 2 diffusive exchange, 3 convective exchange), force on the body
/// (N, over the whole z extent).
pub type Contribution = ([f64; 3], usize, usize, f64);

pub struct E3Fluid {
    pub device: DeviceStep,
    pub field: Field,
    pub grid: Grid,
    pub h: f64,
    pub dt: f64,
    /// The duct's z extent.
    pub width: f64,
    /// The z extent the last `loads` integrated over (per unit span).
    pub load_width: f64,
    pub lines: Arc<RwLock<Lines>>,
    sink: Arc<Mutex<Vec<Contribution>>>,
}

impl E3Fluid {
    /// The fluid on the benchmark channel at rung `ny`: `nz_slab > 0` the
    /// thin slab periodic in z (the flag as a 2D problem), 0 the full 0.41 m
    /// duct with slip side walls; the 2D inflow (parabolic in y, Ū 1) in
    /// both. `speed` bounds the flag's surface speed (the narrow band).
    /// `start`: a saved state (its time, its line, its fields — extruded
    /// onto every plane when the saved nz differs); the rest flow otherwise.
    pub fn build(
        ny: usize,
        nz_slab: usize,
        speed: f64,
        rest: Line,
        start: Option<&super::state::Saved>,
    ) -> Self {
        let rest = match start {
            Some(s) => s.line.clone(),
            None => rest,
        };
        let h = H / ny as f64;
        let nx = (L / h).round() as usize;
        let nz = if nz_slab > 0 {
            nz_slab
        } else {
            (H / h).round() as usize
        };
        // The flag test's step (its CFL velocity; `speed` is the band's bound only).
        let dt = (0.3 * h / U_CFL).min(0.5 * h * h / (6.0 * NU))
            * super::env_f("RTX_E3FSI_DT_SCALE", 1.0);
        let boundaries = if nz_slab > 0 {
            Boundaries {
                x1: Side::PressureOutlet,
                z0: Side::Periodic,
                z1: Side::Periodic,
                ..Boundaries::default()
            }
        } else {
            Boundaries {
                x1: Side::PressureOutlet,
                z0: Side::SlipWall,
                z1: Side::SlipWall,
                ..Boundaries::default()
            }
        };
        let mut solver = Solver::new(
            Fluid {
                density: RHO,
                viscosity: RHO * NU,
                reference_velocity: 1.0,
                reference_length: 2.0 * R_CYL,
            },
            Parameters {
                corrector_steps: super::env_f("RTX_E3FSI_CORRECTORS", 3.0) as usize,
                inner_stop_factor: super::env_f("RTX_E3FSI_INNER", 1e-3),
                tolerance: 1e-8,
                convection_scheme: ConvectionScheme::TvdVanAlbada,
                wall_scheme: WallScheme::CutCell,
                boundaries,
                max_surface_speed: Some(speed),
                ..Parameters::default()
            },
        );
        let inflow = |y: f64| 6.0 * y * (H - y) / (H * H);
        solver.set_boundary_velocity(move |x, y, _z, _t| {
            if x <= 0.0 {
                (inflow(y), 0.0, 0.0)
            } else {
                (0.0, 0.0, 0.0)
            }
        });
        let lines = Arc::new(RwLock::new(Lines {
            a: rest.clone(),
            b: rest,
        }));
        VERSION.fetch_add(1, Ordering::AcqRel);
        let (l1, l2, l3) = (lines.clone(), lines.clone(), lines.clone());
        let body = Body::from_sdf(move |x, y, _z, t| cylinder(x, y).min(capsule(&l1, x, y, t).0))
            .with_surface_velocity(move |x, y, _z, t| {
                let (df, (vx, vy)) = capsule(&l2, x, y, t);
                if df <= cylinder(x, y) {
                    (vx, vy, 0.0)
                } else {
                    (0.0, 0.0, 0.0)
                }
            });
        let width = nz as f64 * h;
        let body = body.with_device_sdf(move |t| {
            let line = l3.read().expect("lines").at(t);
            DeviceSdf {
                cyl: [CX, CY, R_CYL],
                cyl_cut: false,
                flag_cut: false,
                zc: 0.5 * width,
                span: width,
                r_edge: h,
                half: HALF,
                fillet: 0.0,
                poly: line.pts,
                vel: line.vel,
                // R8-c's plate body (merged alongside): the span-uniform harness keeps the polyline.
                plate: None,
            }
        });
        solver.set_moving_body(body);
        let g = Grid::cubic(nx, ny, nz, h);
        let mut field = Field::new(g);
        for k in 0..nz {
            for j in 0..ny {
                let u0 = inflow((j as f64 + 0.5) * h);
                for i in 0..=nx {
                    field.u[g.uface(k, j, i)] = u0;
                }
            }
        }
        if let Some(s) = start {
            s.fill(&mut field);
            solver.set_time(s.t);
        }
        solver.initialize(&mut field);
        let mut device = DeviceStep::new(solver, g);
        device.upload(&field);
        println!(
            "  R8-a fluid: ny {ny}, {nx}×{ny}×{nz} = {} cells ({}), h {h:.4e}, dt {dt:.4e}, speed bound {speed} m/s",
            g.cells(),
            if nz_slab > 0 {
                "slab, periodic z"
            } else {
                "full duct, slip sides"
            }
        );
        Self {
            device,
            field,
            grid: g,
            h,
            dt,
            width,
            load_width: width,
            lines,
            sink: Arc::new(Mutex::new(Vec::new())),
        }
    }

    pub fn time(&self) -> f64 {
        self.device.solver.time()
    }

    /// The step's start and end lines (their times are the lines' own).
    pub fn set_lines(&self, a: Line, b: Line) {
        *self.lines.write().expect("lines") = Lines { a, b };
        VERSION.fetch_add(1, Ordering::AcqRel);
    }

    pub fn step(&mut self) -> StepResult {
        self.device.advance(self.dt)
    }

    pub fn snapshot(&mut self) -> DeviceSnapshot {
        self.device.snapshot()
    }

    pub fn restore(&mut self, snap: &DeviceSnapshot) {
        self.device.restore(snap);
    }

    /// The operator-route force on the whole body per unit span and the
    /// route's contributions (N over the whole z extent).
    pub fn loads(&mut self) -> ([f64; 3], Vec<Contribution>) {
        self.device.download(&mut self.field);
        let t = self.time();
        self.sink.lock().expect("sink").clear();
        let s = self.sink.clone();
        let previous = set_load_sink(Some(Box::new(move |pos, c, part, v| {
            s.lock().expect("sink").push((pos, c, part, v));
        })));
        assert!(previous.is_none(), "a load sink was already installed");
        let mask = self.device.solver.mask().expect("mask");
        let body = self.device.solver.body().expect("body");
        // `RTX_E3FSI_LOAD_PLANES=n`: the route on the n mid planes only (per
        // unit span of those planes) — a cost knob for a spanwise-uniform
        // flow; the whole span by default.
        let nz = self.grid.nz;
        let n = (super::env_f("RTX_E3FSI_LOAD_PLANES", 0.0) as usize).min(nz);
        // R8-f (`RTX_E3_LOADS_DEVICE=1`): the same route on the device, its
        // summands replayed into the sink in the host's order (exact zeros
        // dropped); the host route when the device cannot serve it.
        let planes = (n > 0 && n < nz).then(|| ((nz - n) / 2, (nz - n) / 2 + n));
        let on_device = if loads_device_enabled() {
            self.device.cut_wall_force_device(RHO * NU, planes)
        } else {
            None
        };
        let f = if let Some(f) = on_device {
            if let Some((k0, k1)) = planes {
                self.load_width = (k1 - k0) as f64 * self.h;
                f
            } else {
                self.load_width = self.width;
                [f[0] / self.width, f[1] / self.width, f[2] / self.width]
            }
        } else if n > 0 && n < nz {
            let k0 = (nz - n) / 2;
            self.load_width = n as f64 * self.h;
            mask.cut_wall_force_per_span(body, &self.field, RHO * NU, t, (k0, k0 + n))
                .expect("cut wall force per span")
        } else {
            self.load_width = self.width;
            let f = mask
                .cut_wall_force(body, &self.field, RHO * NU, t)
                .expect("cut wall force");
            [f[0] / self.width, f[1] / self.width, f[2] / self.width]
        };
        set_load_sink(None);
        let contributions = std::mem::take(&mut *self.sink.lock().expect("sink"));
        (f, contributions)
    }
}