//! Tier-C statistics over a fully materialized logical selection.

use std::cmp::Ordering;
use std::path::Path;

use memmap2::MmapMut;

use crate::query::types::{Operation, OperationPreviewFields, ReadPlan, TetError};

use crate::query::decode::indexing::coords_from_linear_row_major;
use crate::query::engine::budget::MemoryStrategy;
use crate::query::fold::partial_geometry::{partial_axis_layout, reduced_index};

use super::int;
use super::{LogicalF32Backing, LogicalF64Backing, MaterializedLogical};

fn median_f64(values: &mut [f64]) -> Result<f64, TetError> {
    if values.is_empty() {
        return Err(TetError::Validation(
            "median requires at least one element".into(),
        ));
    }
    let cmp = |a: &f64, b: &f64| a.partial_cmp(b).unwrap_or(Ordering::Equal);
    let n = values.len();
    let mid = n / 2;
    if n.is_multiple_of(2) {
        values.select_nth_unstable_by(mid, cmp);
        let hi = values[mid];
        values.select_nth_unstable_by(mid - 1, cmp);
        Ok(f64::midpoint(values[mid - 1], hi))
    } else {
        values.select_nth_unstable_by(mid, cmp);
        Ok(values[mid])
    }
}

#[allow(
    clippy::cast_precision_loss,
    clippy::cast_possible_truncation,
    clippy::cast_sign_loss
)]
fn quantile_f64(values: &mut [f64], q: f64) -> Result<f64, TetError> {
    if values.is_empty() {
        return Err(TetError::Validation(
            "quantile requires at least one element".into(),
        ));
    }
    if values.len() == 1 {
        return Ok(values[0]);
    }
    let pos = q * (values.len() - 1) as f64;
    let lo = pos.floor() as usize;
    let hi = pos.ceil() as usize;
    let cmp = |a: &f64, b: &f64| a.partial_cmp(b).unwrap_or(Ordering::Equal);
    values.select_nth_unstable_by(lo, cmp);
    if lo == hi {
        return Ok(values[lo]);
    }
    values.select_nth_unstable_by(hi, cmp);
    let lower = values[lo];
    let upper = values[hi];
    Ok(lower + (upper - lower) * (pos - lo as f64))
}

#[allow(
    clippy::cast_precision_loss,
    clippy::cast_possible_truncation,
    clippy::cast_sign_loss
)]
fn histogram_f64(values: &[f64], bins: u32) -> Result<(Vec<f64>, Vec<f64>), TetError> {
    if values.is_empty() {
        return Err(TetError::Validation(
            "histogram requires at least one element".into(),
        ));
    }
    let bins = usize::try_from(bins)
        .map_err(|_| TetError::Validation("histogram bins overflow".into()))?;
    let min = values.iter().copied().reduce(f64::min).unwrap_or(f64::NAN);
    let max = values.iter().copied().reduce(f64::max).unwrap_or(f64::NAN);
    let mut counts = vec![0.0; bins];
    let mut edges = vec![0.0; bins + 1];
    if (max - min).abs() < f64::EPSILON {
        counts[0] = values.len() as f64;
        for (i, e) in edges.iter_mut().enumerate() {
            *e = if i == 0 { min } else { max };
        }
        return Ok((counts, edges));
    }
    let width = (max - min) / bins as f64;
    for (i, e) in edges.iter_mut().enumerate() {
        *e = min + width * i as f64;
    }
    if let Some(last) = edges.last_mut() {
        *last = max;
    }
    for &v in values {
        let mut idx = ((v - min) / width).floor() as usize;
        if idx >= bins {
            idx = bins - 1;
        }
        counts[idx] += 1.0;
    }
    Ok((counts, edges))
}

fn median_f32(values: &mut [f32]) -> Result<f64, TetError> {
    let mut tmp: Vec<f64> = values.iter().map(|&v| f64::from(v)).collect();
    median_f64(&mut tmp)
}

fn gather_partial_in_memory<V: Copy>(
    values: &[V],
    shape: &[u64],
    layout: &crate::query::fold::partial_geometry::PartialAxisLayout,
) -> Result<Vec<Vec<V>>, TetError> {
    let mut cells = vec![Vec::new(); layout.out_len];
    for (li, &value) in values.iter().enumerate() {
        let coords = coords_from_linear_row_major(li, shape)?;
        let oi = reduced_index(&coords, &layout.axis_set, &layout.out_shape)?;
        cells[oi].push(value);
    }
    Ok(cells)
}

fn gather_partial_in_memory_f64_from_f32(
    values: &[f32],
    shape: &[u64],
    layout: &crate::query::fold::partial_geometry::PartialAxisLayout,
) -> Result<Vec<Vec<f64>>, TetError> {
    let mut cells = vec![Vec::new(); layout.out_len];
    for (li, &value) in values.iter().enumerate() {
        let coords = coords_from_linear_row_major(li, shape)?;
        let oi = reduced_index(&coords, &layout.axis_set, &layout.out_shape)?;
        cells[oi].push(f64::from(value));
    }
    Ok(cells)
}

fn gather_partial_f64<F>(
    n: usize,
    shape: &[u64],
    layout: &crate::query::fold::partial_geometry::PartialAxisLayout,
    mut read_at: F,
) -> Result<Vec<Vec<f64>>, TetError>
where
    F: FnMut(usize) -> Result<f64, TetError>,
{
    let mut cells = vec![Vec::new(); layout.out_len];
    for li in 0..n {
        let coords = coords_from_linear_row_major(li, shape)?;
        let oi = reduced_index(&coords, &layout.axis_set, &layout.out_shape)?;
        cells[oi].push(read_at(li)?);
    }
    Ok(cells)
}

fn read_f64_at_spill_f32(path: &Path, li: usize) -> Result<f64, TetError> {
    let file = std::fs::File::open(path)
        .map_err(|e| TetError::Validation(format!("temp spill read failed: {e}")))?;
    let mmap = unsafe {
        memmap2::Mmap::map(&file)
            .map_err(|e| TetError::Validation(format!("temp spill mmap failed: {e}")))?
    };
    Ok(f64::from(bytemuck::cast_slice::<_, f32>(&mmap)[li]))
}

fn read_f64_at_spill_f64(path: &Path, li: usize) -> Result<f64, TetError> {
    let file = std::fs::File::open(path)
        .map_err(|e| TetError::Validation(format!("temp spill read failed: {e}")))?;
    let mmap = unsafe {
        memmap2::Mmap::map(&file)
            .map_err(|e| TetError::Validation(format!("temp spill mmap failed: {e}")))?
    };
    Ok(bytemuck::cast_slice::<_, f64>(&mmap)[li])
}

fn median_f32_spill_file(path: &Path, n: usize) -> Result<f64, TetError> {
    let file = std::fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(path)
        .map_err(|e| TetError::Validation(format!("temp spill open failed: {e}")))?;
    let mut mmap = unsafe {
        MmapMut::map_mut(&file)
            .map_err(|e| TetError::Validation(format!("temp spill mmap mut failed: {e}")))?
    };
    let slice = bytemuck::cast_slice_mut(mmap.as_mut());
    median_f32(&mut slice[..n])
}

fn median_f64_spill_file(path: &Path, n: usize) -> Result<f64, TetError> {
    let file = std::fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(path)
        .map_err(|e| TetError::Validation(format!("temp spill open failed: {e}")))?;
    let mut mmap = unsafe {
        MmapMut::map_mut(&file)
            .map_err(|e| TetError::Validation(format!("temp spill mmap mut failed: {e}")))?
    };
    let slice = bytemuck::cast_slice_mut(mmap.as_mut());
    median_f64(&mut slice[..n])
}

fn run_tier_c_median_f32(
    backing: &LogicalF32Backing,
    plan: &ReadPlan,
    axes: &[String],
    n: usize,
    shape: &[u64],
) -> Result<OperationPreviewFields, TetError> {
    if axes.is_empty() {
        let median = match backing {
            LogicalF32Backing::InMemory(v) => median_f32(&mut v.clone())?,
            LogicalF32Backing::TempSpill(temp) => median_f32_spill_file(temp.path(), n)?,
        };
        return Ok(OperationPreviewFields {
            element_count: Some(n),
            median: Some(median),
            ..OperationPreviewFields::default()
        });
    }
    let layout = partial_axis_layout(plan, axes)?;
    match backing {
        LogicalF32Backing::TempSpill(temp) => {
            let f64_cells = gather_partial_f64(n, shape, &layout, |li| {
                read_f64_at_spill_f32(temp.path(), li)
            })?;
            let medians: Vec<f64> = f64_cells
                .iter()
                .map(|c| median_f64(&mut c.clone()))
                .collect::<Result<Vec<_>, _>>()?;
            Ok(OperationPreviewFields {
                element_count: Some(n),
                reduced_shape: Some(layout.out_shape),
                reduced_median: Some(medians),
                ..OperationPreviewFields::default()
            })
        }
        LogicalF32Backing::InMemory(v) => {
            let mut cells = gather_partial_in_memory(v, shape, &layout)?;
            let medians: Vec<f64> = cells
                .iter_mut()
                .map(|c| median_f32(c))
                .collect::<Result<Vec<_>, _>>()?;
            Ok(OperationPreviewFields {
                element_count: Some(n),
                reduced_shape: Some(layout.out_shape),
                reduced_median: Some(medians),
                ..OperationPreviewFields::default()
            })
        }
    }
}

fn run_tier_c_median_f64(
    backing: &LogicalF64Backing,
    plan: &ReadPlan,
    axes: &[String],
    n: usize,
    shape: &[u64],
) -> Result<OperationPreviewFields, TetError> {
    if axes.is_empty() {
        let median = match backing {
            LogicalF64Backing::InMemory(v) => median_f64(&mut v.clone())?,
            LogicalF64Backing::TempSpill(temp) => median_f64_spill_file(temp.path(), n)?,
        };
        return Ok(OperationPreviewFields {
            element_count: Some(n),
            median: Some(median),
            ..OperationPreviewFields::default()
        });
    }
    let layout = partial_axis_layout(plan, axes)?;
    let mut cells = match backing {
        LogicalF64Backing::InMemory(v) => gather_partial_in_memory(v, shape, &layout)?,
        LogicalF64Backing::TempSpill(temp) => gather_partial_f64(n, shape, &layout, |li| {
            read_f64_at_spill_f64(temp.path(), li)
        })?,
    };
    let medians: Vec<f64> = cells
        .iter_mut()
        .map(|c| median_f64(c))
        .collect::<Result<Vec<_>, _>>()?;
    Ok(OperationPreviewFields {
        element_count: Some(n),
        reduced_shape: Some(layout.out_shape),
        reduced_median: Some(medians),
        ..OperationPreviewFields::default()
    })
}

pub(crate) fn run_tier_c_operation(
    materialized: &MaterializedLogical,
    plan: &ReadPlan,
    op: &Operation,
) -> Result<OperationPreviewFields, TetError> {
    if matches!(
        materialized,
        MaterializedLogical::I32 { .. } | MaterializedLogical::I64 { .. }
    ) {
        let vec = int::materialized_logical_as_f64(materialized)?;
        let synthetic = MaterializedLogical::F64 {
            backing: LogicalF64Backing::InMemory(vec),
            total_bytes_read_from_disk: 0,
            strategy: MemoryStrategy::InMemoryMaterialize,
        };
        return run_tier_c_operation(&synthetic, plan, op);
    }

    let n = plan.logical_f32_element_count;
    let shape = &plan.logical_selection_shape;
    let axes = op.axes();

    match (materialized, op) {
        (MaterializedLogical::F32 { backing, .. }, Operation::Median { .. }) => {
            run_tier_c_median_f32(backing, plan, axes, n, shape)
        }
        (MaterializedLogical::F64 { backing, .. }, Operation::Median { .. }) => {
            run_tier_c_median_f64(backing, plan, axes, n, shape)
        }
        (MaterializedLogical::F32 { backing, .. }, Operation::Quantile { q, .. }) => {
            run_quantile_f32(backing, plan, axes, n, shape, *q)
        }
        (MaterializedLogical::F64 { backing, .. }, Operation::Quantile { q, .. }) => {
            run_quantile_f64(backing, plan, axes, n, shape, *q)
        }
        (MaterializedLogical::F32 { backing, .. }, Operation::Histogram { bins, .. }) => {
            run_histogram_f32(backing, plan, axes, n, shape, *bins)
        }
        (MaterializedLogical::F64 { backing, .. }, Operation::Histogram { bins, .. }) => {
            run_histogram_f64(backing, plan, axes, n, shape, *bins)
        }
        _ => Err(TetError::Validation(format!(
            "unsupported materialize-required operation: {op:?}"
        ))),
    }
}

fn run_quantile_f64(
    backing: &LogicalF64Backing,
    plan: &ReadPlan,
    axes: &[String],
    n: usize,
    shape: &[u64],
    q: f64,
) -> Result<OperationPreviewFields, TetError> {
    if axes.is_empty() {
        let mut cell = match backing {
            LogicalF64Backing::InMemory(v) => v.clone(),
            LogicalF64Backing::TempSpill(temp) => (0..n)
                .map(|li| read_f64_at_spill_f64(temp.path(), li))
                .collect::<Result<Vec<_>, _>>()?,
        };
        let val = quantile_f64(&mut cell, q)?;
        return Ok(OperationPreviewFields {
            element_count: Some(n),
            quantile: Some(val),
            ..OperationPreviewFields::default()
        });
    }
    let layout = partial_axis_layout(plan, axes)?;
    let mut cells = match backing {
        LogicalF64Backing::InMemory(v) => gather_partial_in_memory(v, shape, &layout)?,
        LogicalF64Backing::TempSpill(temp) => gather_partial_f64(n, shape, &layout, |li| {
            read_f64_at_spill_f64(temp.path(), li)
        })?,
    };
    let qs: Vec<f64> = cells
        .iter_mut()
        .map(|c| quantile_f64(c, q))
        .collect::<Result<Vec<_>, _>>()?;
    Ok(OperationPreviewFields {
        element_count: Some(n),
        reduced_shape: Some(layout.out_shape),
        reduced_quantile: Some(qs),
        ..OperationPreviewFields::default()
    })
}

fn run_quantile_f32(
    backing: &LogicalF32Backing,
    plan: &ReadPlan,
    axes: &[String],
    n: usize,
    shape: &[u64],
    q: f64,
) -> Result<OperationPreviewFields, TetError> {
    if axes.is_empty() {
        let mut cell: Vec<f64> = match backing {
            LogicalF32Backing::InMemory(v) => v.iter().map(|&x| f64::from(x)).collect(),
            LogicalF32Backing::TempSpill(temp) => (0..n)
                .map(|li| read_f64_at_spill_f32(temp.path(), li))
                .collect::<Result<Vec<_>, _>>()?,
        };
        let val = quantile_f64(&mut cell, q)?;
        return Ok(OperationPreviewFields {
            element_count: Some(n),
            quantile: Some(val),
            ..OperationPreviewFields::default()
        });
    }
    let layout = partial_axis_layout(plan, axes)?;
    let mut cells = match backing {
        LogicalF32Backing::InMemory(v) => gather_partial_in_memory_f64_from_f32(v, shape, &layout)?,
        LogicalF32Backing::TempSpill(temp) => gather_partial_f64(n, shape, &layout, |li| {
            read_f64_at_spill_f32(temp.path(), li)
        })?,
    };
    let qs: Vec<f64> = cells
        .iter_mut()
        .map(|c| quantile_f64(c, q))
        .collect::<Result<Vec<_>, _>>()?;
    Ok(OperationPreviewFields {
        element_count: Some(n),
        reduced_shape: Some(layout.out_shape),
        reduced_quantile: Some(qs),
        ..OperationPreviewFields::default()
    })
}

fn run_histogram_f64(
    backing: &LogicalF64Backing,
    plan: &ReadPlan,
    axes: &[String],
    n: usize,
    shape: &[u64],
    bins: u32,
) -> Result<OperationPreviewFields, TetError> {
    if axes.is_empty() {
        let values = match backing {
            LogicalF64Backing::InMemory(v) => v.clone(),
            LogicalF64Backing::TempSpill(temp) => (0..n)
                .map(|li| read_f64_at_spill_f64(temp.path(), li))
                .collect::<Result<Vec<_>, _>>()?,
        };
        let (counts, edges) = histogram_f64(&values, bins)?;
        return Ok(OperationPreviewFields {
            element_count: Some(n),
            histogram_counts: Some(counts),
            histogram_edges: Some(edges),
            ..OperationPreviewFields::default()
        });
    }
    let layout = partial_axis_layout(plan, axes)?;
    let cells = match backing {
        LogicalF64Backing::InMemory(v) => gather_partial_in_memory(v, shape, &layout)?,
        LogicalF64Backing::TempSpill(temp) => gather_partial_f64(n, shape, &layout, |li| {
            read_f64_at_spill_f64(temp.path(), li)
        })?,
    };
    let mut flat = Vec::with_capacity(cells.len() * bins as usize);
    for cell in &cells {
        let (counts, _) = histogram_f64(cell, bins)?;
        flat.extend(counts);
    }
    Ok(OperationPreviewFields {
        element_count: Some(n),
        reduced_shape: Some(layout.out_shape),
        reduced_histogram_counts: Some(flat),
        ..OperationPreviewFields::default()
    })
}

fn run_histogram_f32(
    backing: &LogicalF32Backing,
    plan: &ReadPlan,
    axes: &[String],
    n: usize,
    shape: &[u64],
    bins: u32,
) -> Result<OperationPreviewFields, TetError> {
    if axes.is_empty() {
        let values: Vec<f64> = match backing {
            LogicalF32Backing::InMemory(v) => v.iter().map(|&x| f64::from(x)).collect(),
            LogicalF32Backing::TempSpill(temp) => (0..n)
                .map(|li| read_f64_at_spill_f32(temp.path(), li))
                .collect::<Result<Vec<_>, _>>()?,
        };
        let (counts, edges) = histogram_f64(&values, bins)?;
        return Ok(OperationPreviewFields {
            element_count: Some(n),
            histogram_counts: Some(counts),
            histogram_edges: Some(edges),
            ..OperationPreviewFields::default()
        });
    }
    let layout = partial_axis_layout(plan, axes)?;
    let cells = match backing {
        LogicalF32Backing::InMemory(v) => gather_partial_in_memory_f64_from_f32(v, shape, &layout)?,
        LogicalF32Backing::TempSpill(temp) => gather_partial_f64(n, shape, &layout, |li| {
            read_f64_at_spill_f32(temp.path(), li)
        })?,
    };
    let mut flat = Vec::with_capacity(cells.len() * bins as usize);
    for cell in &cells {
        let (counts, _) = histogram_f64(cell, bins)?;
        flat.extend(counts);
    }
    Ok(OperationPreviewFields {
        element_count: Some(n),
        reduced_shape: Some(layout.out_shape),
        reduced_histogram_counts: Some(flat),
        ..OperationPreviewFields::default()
    })
}
