use crate::riesz::RieszTransform;

/// The monogenic signal: f + i·R₁f + j·R₂f + k·R₃f
///
/// At each point in the grid, this quaternion-valued field
/// decomposes into amplitude, phase, and orientation:
///   - Amplitude: how much geometry is here (energy)
///   - Phase: what kind (edge ≈ π/2, ridge ≈ 0, blob ≈ π)
///   - Orientation: which way the feature points (2 angles)
pub struct MonogenicField {
    pub samples: Vec<MonogenicSample>,
    pub n: usize,
}

#[derive(Debug, Clone, Copy)]
pub struct MonogenicSample {
    /// Original scalar field value.
    pub f: f64,
    /// Riesz transform components.
    pub r1: f64,
    pub r2: f64,
    pub r3: f64,
    /// Local amplitude: sqrt(f² + r1² + r2² + r3²)
    pub amplitude: f64,
    /// Local phase: atan2(|r|, f)
    /// 0 = line/ridge, π/2 = edge/step, π = blob
    pub phase: f64,
    /// Orientation: unit vector (r1, r2, r3) / |(r1, r2, r3)|
    pub orientation: [f64; 3],
}

impl MonogenicField {
    /// Compute the monogenic signal from a scalar field.
    pub fn compute(field: &[f64], n: usize) -> Self {
        let riesz = RieszTransform::compute(field, n);

        let samples: Vec<MonogenicSample> = (0..n * n * n)
            .map(|i| {
                let f = field[i];
                let r1 = riesz.r1[i];
                let r2 = riesz.r2[i];
                let r3 = riesz.r3[i];

                let r_mag = (r1 * r1 + r2 * r2 + r3 * r3).sqrt();
                let amplitude = (f * f + r_mag * r_mag).sqrt();
                let phase = r_mag.atan2(f);

                let orientation = if r_mag > 1e-10 {
                    [r1 / r_mag, r2 / r_mag, r3 / r_mag]
                } else {
                    [0.0, 0.0, 0.0]
                };

                MonogenicSample {
                    f,
                    r1,
                    r2,
                    r3,
                    amplitude,
                    phase,
                    orientation,
                }
            })
            .collect();

        MonogenicField { samples, n }
    }

    /// Extract surface samples where amplitude is significant
    /// and phase indicates a specific feature type.
    pub fn surface_at_phase(&self, phase_center: f64, phase_tolerance: f64, min_amplitude: f64) -> Vec<(usize, &MonogenicSample)> {
        self.samples.iter().enumerate()
            .filter(|(_, s)| {
                s.amplitude > min_amplitude
                    && (s.phase - phase_center).abs() < phase_tolerance
            })
            .collect()
    }

    /// Extract edge-like features (phase ≈ π/2).
    pub fn edges(&self, min_amplitude: f64) -> Vec<(usize, &MonogenicSample)> {
        self.surface_at_phase(std::f64::consts::FRAC_PI_2, 0.4, min_amplitude)
    }

    /// Extract ridge-like features (phase ≈ 0).
    pub fn ridges(&self, min_amplitude: f64) -> Vec<(usize, &MonogenicSample)> {
        self.surface_at_phase(0.0, 0.4, min_amplitude)
    }

    /// Extract blob-like features (phase ≈ π).
    pub fn blobs(&self, min_amplitude: f64) -> Vec<(usize, &MonogenicSample)> {
        self.surface_at_phase(std::f64::consts::PI, 0.4, min_amplitude)
    }

    /// Group samples by orientation similarity.
    /// Returns clusters of grid indices that share a common direction.
    pub fn cluster_by_orientation(&self, indices: &[(usize, &MonogenicSample)], angle_threshold: f64) -> Vec<Vec<usize>> {
        let cos_thresh = angle_threshold.cos();
        let mut clusters: Vec<Vec<usize>> = Vec::new();
        let mut assigned = vec![false; indices.len()];

        for i in 0..indices.len() {
            if assigned[i] { continue; }
            let (idx_i, sample_i) = indices[i];
            let oi = sample_i.orientation;
            if oi[0] == 0.0 && oi[1] == 0.0 && oi[2] == 0.0 {
                continue;
            }

            let mut cluster = vec![idx_i];
            assigned[i] = true;

            for j in (i + 1)..indices.len() {
                if assigned[j] { continue; }
                let (idx_j, sample_j) = indices[j];
                let oj = sample_j.orientation;

                let dot = (oi[0] * oj[0] + oi[1] * oj[1] + oi[2] * oj[2]).abs();
                if dot > cos_thresh {
                    cluster.push(idx_j);
                    assigned[j] = true;
                }
            }

            if cluster.len() >= 3 {
                clusters.push(cluster);
            }
        }

        clusters
    }

    /// Index → (ix, iy, iz) grid coordinates.
    pub fn index_to_coord(&self, idx: usize) -> (usize, usize, usize) {
        let iz = idx / (self.n * self.n);
        let iy = (idx % (self.n * self.n)) / self.n;
        let ix = idx % self.n;
        (ix, iy, iz)
    }

    /// Grid coordinate → world position given bounds.
    pub fn coord_to_world(&self, ix: usize, iy: usize, iz: usize, min: [f64; 3], max: [f64; 3]) -> [f64; 3] {
        let step = [
            (max[0] - min[0]) / (self.n - 1).max(1) as f64,
            (max[1] - min[1]) / (self.n - 1).max(1) as f64,
            (max[2] - min[2]) / (self.n - 1).max(1) as f64,
        ];
        [
            min[0] + ix as f64 * step[0],
            min[1] + iy as f64 * step[1],
            min[2] + iz as f64 * step[2],
        ]
    }
}
