//! Parquet ingestion: loading and joining the properties + postcode parquet
//! files, validating coordinates, spatially sorting rows and building the
//! quantized row-major feature matrix plus all side tables.

use anyhow::{bail, Context};
use polars::lazy::frame::LazyFrame;
use polars::prelude::*;
use rayon::prelude::*;
use std::path::Path;

use rustc_hash::{FxHashMap, FxHashSet};

use crate::consts::{NAN_U16, QUANT_SCALE};
use crate::features::{self, Bounds};

use super::address_search::{
    build_address_prefix_index, is_address_candidate_token, merged_address_search_tokens,
};
use super::poi_metrics::{PostcodePoiMetrics, NO_POI_METRIC_ROW};
use super::stats::{column_to_f32_vec, compute_feature_stats, FeatureStats, Histogram};
use super::{HistoricalPrice, PropertyData, RenovationEvent, TenureEvent};

const MISSING_COORDINATE_SAMPLE_LIMIT: usize = 10;
const COUNTRY_COLUMN_CANDIDATES: &[&str] = &[
    "ctry25cd",
    "ctry24cd",
    "ctry23cd",
    "ctry22cd",
    "country_code",
    "Country code",
    "country",
    "Country",
];
const ENGLAND_COUNTRY_VALUES: &[&str] = &["E92000001", "England", "ENGLAND", "england"];

fn is_numeric_dtype(dtype: &DataType) -> bool {
    matches!(
        dtype,
        DataType::Int8
            | DataType::Int16
            | DataType::Int32
            | DataType::Int64
            | DataType::UInt8
            | DataType::UInt16
            | DataType::UInt32
            | DataType::UInt64
            | DataType::Float32
            | DataType::Float64
            | DataType::Datetime(_, _)
            | DataType::Date
    )
}

fn is_datetime_dtype(dtype: &DataType) -> bool {
    matches!(dtype, DataType::Datetime(_, _) | DataType::Date)
}

fn find_country_column(schema: &Schema) -> Option<String> {
    COUNTRY_COLUMN_CANDIDATES
        .iter()
        .find_map(|&name| match schema.get(name) {
            Some(dtype) if matches!(dtype, DataType::String) || dtype.is_categorical() => {
                Some(name.to_string())
            }
            _ => None,
        })
}

fn england_country_expr(country_column: &str) -> Expr {
    ENGLAND_COUNTRY_VALUES.iter().skip(1).fold(
        col(country_column)
            .cast(DataType::String)
            .eq(lit(ENGLAND_COUNTRY_VALUES[0])),
        |expr, value| expr.or(col(country_column).cast(DataType::String).eq(lit(*value))),
    )
}

#[derive(Debug, PartialEq, Eq)]
struct MissingCoordinateSummary {
    row_count: usize,
    unique_postcode_count: usize,
    sample_postcodes: Vec<String>,
}

fn summarize_missing_coordinate_postcodes<'a>(
    postcodes: impl IntoIterator<Item = Option<&'a str>>,
) -> MissingCoordinateSummary {
    let mut row_count = 0usize;
    let mut unique_postcodes = FxHashSet::default();

    for postcode in postcodes {
        row_count += 1;
        if let Some(postcode) = postcode {
            let trimmed = postcode.trim();
            if !trimmed.is_empty() {
                unique_postcodes.insert(trimmed.to_string());
            }
        }
    }

    let mut sample_postcodes: Vec<String> = unique_postcodes.iter().cloned().collect();
    sample_postcodes.sort_unstable();
    sample_postcodes.truncate(MISSING_COORDINATE_SAMPLE_LIMIT);

    MissingCoordinateSummary {
        row_count,
        unique_postcode_count: unique_postcodes.len(),
        sample_postcodes,
    }
}

fn missing_england_coordinates_error(
    summary: &MissingCoordinateSummary,
    country_column: &str,
) -> String {
    let samples = if summary.sample_postcodes.is_empty() {
        "none".to_string()
    } else {
        summary.sample_postcodes.join(", ")
    };
    format!(
        "England property rows missing postcode coordinates after joining postcode data: {} rows across {} postcodes (country column '{}'). Sample postcodes: {}",
        summary.row_count, summary.unique_postcode_count, country_column, samples
    )
}

fn validate_no_england_rows_missing_coordinates(
    combined_lf: &LazyFrame,
    schema: &Schema,
) -> anyhow::Result<()> {
    let Some(country_column) = find_country_column(schema) else {
        bail!(
            "Postcode feature parquet has no reliable country column; cannot verify that rows with missing coordinates are outside England. Regenerate postcode.parquet with pipeline.transform.merge so it includes ctry25cd."
        );
    };

    let missing_coordinates = col("lat").is_null().or(col("lon").is_null());
    let offending_df = combined_lf
        .clone()
        .filter(missing_coordinates.and(england_country_expr(&country_column)))
        .select([col("Postcode")])
        .collect()
        .context("Failed to validate missing postcode coordinates")?;

    let postcode_column = offending_df
        .column("Postcode")
        .context("Joined frame missing 'Postcode' during coordinate validation")?
        .str()
        .context("'Postcode' column is not a string during coordinate validation")?;
    let summary = summarize_missing_coordinate_postcodes(postcode_column);

    if summary.row_count > 0 {
        bail!(
            "{}",
            missing_england_coordinates_error(&summary, &country_column)
        );
    }

    Ok(())
}

impl PropertyData {
    pub fn load(properties_path: &Path, postcode_features_path: &Path) -> anyhow::Result<Self> {
        crate::data::run_polars_io(|| Self::load_inner(properties_path, postcode_features_path))
    }

    fn load_inner(properties_path: &Path, postcode_features_path: &Path) -> anyhow::Result<Self> {
        // Load postcode.parquet
        tracing::info!(
            "Loading postcode features from {:?}",
            postcode_features_path
        );
        let postcode_features_path = PlRefPath::try_from_path(postcode_features_path)
            .context("Failed to normalize postcode parquet path")?;
        let postcode_df = LazyFrame::scan_parquet(postcode_features_path, Default::default())
            .context("Failed to scan postcode parquet")?
            .collect()
            .context("Failed to read postcode parquet")?;
        tracing::info!(rows = postcode_df.height(), "Postcode features loaded");

        let mut poi_metric_names: Vec<String> = postcode_df
            .get_column_names()
            .iter()
            .map(|name| name.as_str())
            .filter(|&name| features::is_dynamic_poi_feature(name))
            .map(str::to_string)
            .collect();
        poi_metric_names.sort_by_key(|name| features::dynamic_poi_feature_sort_key(name));

        let poi_metric_by_postcode: FxHashMap<String, u32> = if poi_metric_names.is_empty() {
            FxHashMap::default()
        } else {
            let postcode_column = postcode_df
                .column("Postcode")
                .context("Postcode feature parquet missing 'Postcode' column")?
                .str()
                .context("'Postcode' column in postcode feature parquet is not a string")?;
            postcode_column
                .into_iter()
                .enumerate()
                .filter_map(|(idx, postcode)| {
                    postcode.map(|postcode| (postcode.to_string(), idx as u32))
                })
                .collect()
        };
        let mut poi_metrics = PostcodePoiMetrics::from_postcode_df(&postcode_df, poi_metric_names)?;

        // Load properties.parquet and join with postcode data lazily so the
        // wide combined frame is never fully materialized — projection is
        // pushed down into the join, keeping peak memory bounded.
        tracing::info!("Loading properties from {:?}", properties_path);
        let properties_path = PlRefPath::try_from_path(properties_path)
            .context("Failed to normalize properties parquet path")?;
        let properties_lf = LazyFrame::scan_parquet(properties_path, Default::default())
            .context("Failed to scan properties parquet")?;
        let combined_lf = properties_lf.join(
            postcode_df.lazy(),
            [col("Postcode")],
            [col("Postcode")],
            JoinArgs::new(JoinType::Left),
        );

        // Get configured feature/enum names in config order. Dynamic POI
        // metrics live in a postcode-level side table so they do not widen the
        // hot row-major property feature matrix.
        let configured_numeric_names = features::all_numeric_feature_names();
        let enum_names = features::all_enum_feature_names();

        let schema = combined_lf
            .clone()
            .collect_schema()
            .context("Failed to collect joined schema")?;
        validate_no_england_rows_missing_coordinates(&combined_lf, &schema)?;
        let numeric_names: Vec<String> = configured_numeric_names
            .iter()
            .map(|name| (*name).to_string())
            .collect();

        for name in &numeric_names {
            match schema.get(name.as_str()) {
                Some(dtype) if is_numeric_dtype(dtype) => {}
                Some(dtype) => bail!(
                    "Configured numeric feature '{}' has non-numeric type {:?}",
                    name,
                    dtype
                ),
                None => bail!(
                    "Configured numeric feature '{}' not found in combined schema",
                    name
                ),
            }
        }
        for name in &enum_names {
            match schema.get(name) {
                Some(dtype) if matches!(dtype, DataType::String) || dtype.is_categorical() => {}
                Some(dtype) => bail!(
                    "Configured enum feature '{}' has unexpected type {:?}",
                    name,
                    dtype
                ),
                None => bail!(
                    "Configured enum feature '{}' not found in combined schema",
                    name
                ),
            }
        }

        // Combine numeric and enum feature names (numeric first, then enum)
        let feature_names: Vec<String> = numeric_names
            .iter()
            .map(|name| name.to_string())
            .chain(enum_names.iter().map(|name| name.to_string()))
            .collect();
        let num_features = feature_names.len();
        let num_numeric = numeric_names.len();
        tracing::info!(
            numeric = num_numeric,
            enums = enum_names.len(),
            total = num_features,
            "Feature columns from config"
        );

        // Build select expressions for the combined DataFrame
        let mut select_exprs: Vec<polars::prelude::Expr> = vec![];
        select_exprs.push(col("lat").cast(DataType::Float32));
        select_exprs.push(col("lon").cast(DataType::Float32));

        // Select numeric features as Float32 (datetime columns → fractional year)
        for name in &numeric_names {
            if is_datetime_dtype(schema.get(name.as_str()).unwrap()) {
                select_exprs.push(
                    (col(name.as_str()).dt().year().cast(DataType::Float32)
                        + (col(name.as_str()).dt().month().cast(DataType::Float32) - lit(1.0f32))
                            / lit(12.0f32))
                    .alias(name.as_str()),
                );
            } else {
                select_exprs.push(col(name.as_str()).cast(DataType::Float32));
            }
        }

        // String columns for address/postcode and property metadata
        for &string_col_name in &[
            "Address per Property Register",
            "Address per EPC",
            "Postcode",
            "Property sub-type",
            "Price qualifier",
        ] {
            if schema.get(string_col_name).is_some() {
                select_exprs.push(col(string_col_name).cast(DataType::String));
            }
        }

        // Enum features as String
        for &name in &enum_names {
            select_exprs.push(col(name).cast(DataType::String));
        }

        // Optional columns
        let has_approx_col = schema.get("Is construction date approximate").is_some();
        if has_approx_col {
            select_exprs.push(col("Is construction date approximate").cast(DataType::Float32));
        }
        let has_renovation_history = schema.get("renovation_history").is_some();
        if has_renovation_history {
            select_exprs.push(col("renovation_history"));
        }
        let has_historical_prices = schema.get("historical_prices").is_some();
        if has_historical_prices {
            select_exprs.push(col("historical_prices"));
        }
        let has_tenure_history = schema.get("tenure_history").is_some();
        if has_tenure_history {
            select_exprs.push(col("tenure_history"));
        }
        let df = combined_lf
            .filter(col("lat").is_not_null().and(col("lon").is_not_null()))
            .select(select_exprs)
            .collect()
            .context("Failed to select columns from joined frame")?;

        let row_count = df.height();
        if row_count == 0 {
            bail!("No property rows have usable coordinates after joining postcode data");
        }
        tracing::info!(rows = row_count, "Combined data selected");

        let lat_series = df
            .column("lat")
            .context("Missing 'lat' column")?
            .cast(&DataType::Float32)
            .context("Failed to cast 'lat' to Float32")?;
        let lat: Vec<f32> = lat_series
            .f32()
            .context("Failed to read 'lat' as f32")?
            .into_iter()
            .map(|value| value.context("Missing 'lat' value after coordinate filter"))
            .collect::<anyhow::Result<Vec<_>>>()?;

        let lon_series = df
            .column("lon")
            .context("Missing 'lon' column")?
            .cast(&DataType::Float32)
            .context("Failed to cast 'lon' to Float32")?;
        let lon: Vec<f32> = lon_series
            .f32()
            .context("Failed to read 'lon' as f32")?
            .into_iter()
            .map(|value| value.context("Missing 'lon' value after coordinate filter"))
            .collect::<anyhow::Result<Vec<_>>>()?;

        for (row, (&latitude, &longitude)) in lat.iter().zip(&lon).enumerate() {
            if !(-90.0..=90.0).contains(&latitude) || !(-180.0..=180.0).contains(&longitude) {
                bail!("Invalid coordinates at row {row}: lat={latitude}, lon={longitude}");
            }
        }

        tracing::info!("Extracting numeric feature columns");
        let numeric_col_major: Vec<Vec<f32>> = numeric_names
            .par_iter()
            .map(|name| {
                let column = df
                    .column(name.as_str())
                    .with_context(|| format!("Missing feature column '{name}'"))?;
                column_to_f32_vec(column)
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        tracing::info!("Computing histograms for numeric features");
        let numeric_feature_stats: Vec<FeatureStats> = numeric_col_major
            .par_iter()
            .enumerate()
            .map(|(feat_index, vals)| {
                let name = numeric_names[feat_index].as_str();
                let bounds = features::bounds_for(name)
                    .with_context(|| format!("No bounds config for feature '{}'", name))?;
                let stats = compute_feature_stats(vals, &bounds, features::has_integer_bins(name));
                tracing::debug!(
                    feature = %name,
                    slider_min = format_args!("{:.2}", stats.slider_min),
                    slider_max = format_args!("{:.2}", stats.slider_max),
                    bins = stats.histogram.counts.len(),
                    "Feature stats"
                );
                Ok(stats)
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        // Compute quantization parameters from feature stats (numeric features).
        // For features with Fixed bounds, use those bounds so the full configured range
        // is representable — the histogram refinement can narrow min/max to exclude
        // "outliers" that are actually valid data (e.g. ethnicity percentages).
        // For Percentile-bounded features, use the (possibly refined) histogram range
        // so extreme outliers don't destroy precision for the main distribution.
        let mut quant_min = Vec::with_capacity(num_features);
        let mut quant_range = Vec::with_capacity(num_features);
        for (feat_idx, stats) in numeric_feature_stats.iter().enumerate() {
            let (min, max) = match features::bounds_for(numeric_names[feat_idx].as_str()) {
                Some(Bounds::Fixed { min, max }) => (min, max),
                _ => (stats.histogram.min, stats.histogram.max),
            };
            quant_min.push(min);
            quant_range.push(if max > min { max - min } else { 0.0 });
        }

        tracing::info!("Extracting string columns");
        let extract_string_col = |df: &DataFrame, name: &str| -> anyhow::Result<Vec<String>> {
            let column = df
                .column(name)
                .with_context(|| format!("Required column '{name}' not found in parquet"))?;
            let string_column = column
                .str()
                .with_context(|| format!("Column '{name}' is not a string column"))?;
            string_column
                .into_iter()
                .enumerate()
                .map(|(row, value)| {
                    value
                        .map(ToString::to_string)
                        .with_context(|| format!("Required column '{name}' has null at row {row}"))
                })
                .collect()
        };

        let postcode_raw = extract_string_col(&df, "Postcode")?;

        // Extract optional string columns
        let extract_optional_string_col =
            |df: &DataFrame, name: &str| -> anyhow::Result<Vec<Option<String>>> {
                if let Ok(column) = df.column(name) {
                    let string_column = column
                        .str()
                        .with_context(|| format!("Column '{name}' is not a string column"))?;
                    Ok(string_column
                        .into_iter()
                        .map(|value| {
                            value.and_then(|s| {
                                let trimmed = s.trim();
                                if trimmed.is_empty() {
                                    None
                                } else {
                                    Some(trimmed.to_string())
                                }
                            })
                        })
                        .collect())
                } else {
                    Ok(vec![None; row_count])
                }
            };
        let property_sub_type_raw = extract_optional_string_col(&df, "Property sub-type")?;
        let price_qualifier_raw = extract_optional_string_col(&df, "Price qualifier")?;
        // Display + search addresses. The price-paid form ("Address per Property
        // Register") is the primary display address; the EPC form ("Address per
        // EPC") is an alternative spelling indexed alongside it so a property is
        // findable by either. For EPC-only dwellings (no Land Registry sale) the
        // price-paid form is null, so the EPC form becomes the display address.
        let pp_address_raw = extract_optional_string_col(&df, "Address per Property Register")?;
        let epc_address_raw = extract_optional_string_col(&df, "Address per EPC")?;

        tracing::info!("Building enum features");
        // enum_col_major: Vec<(values_list, encoded_as_f32)>
        let enum_col_major: Vec<(Vec<String>, Vec<f32>)> = enum_names
            .par_iter()
            .map(|&name| -> anyhow::Result<(Vec<String>, Vec<f32>)> {
                let column_data = df
                    .column(name)
                    .with_context(|| format!("Required enum column '{name}' not found"))?;
                let string_column = column_data
                    .str()
                    .with_context(|| format!("Enum column '{name}' is not a string column"))?;
                let unique_set: std::collections::HashSet<String> = string_column
                    .into_iter()
                    .filter_map(|value| {
                        let text = value?.trim();
                        (!text.is_empty()).then(|| text.to_string())
                    })
                    .collect();

                // Use configured order if available, otherwise alphabetical
                let unique: Vec<String> = if let Some(order) = features::order_for(name) {
                    let mut ordered: Vec<String> = Vec::new();
                    for &ordered_value in order {
                        if unique_set.contains(ordered_value) {
                            ordered.push(ordered_value.to_string());
                        }
                    }
                    // Append any values not in the configured order, alphabetically
                    // Use HashSet for O(1) contains instead of O(n) slice search
                    let order_set: rustc_hash::FxHashSet<&str> = order.iter().copied().collect();
                    let mut remainder: Vec<String> = unique_set
                        .iter()
                        .filter(|value| !order_set.contains(value.as_str()))
                        .cloned()
                        .collect();
                    remainder.sort();
                    ordered.extend(remainder);
                    ordered
                } else {
                    let mut sorted: Vec<String> = unique_set.into_iter().collect();
                    sorted.sort();
                    sorted
                };

                let value_to_idx: std::collections::HashMap<&str, f32> = unique
                    .iter()
                    .enumerate()
                    .map(|(index, value)| (value.as_str(), index as f32))
                    .collect();

                let encoded: Vec<f32> = string_column
                    .into_iter()
                    .enumerate()
                    .map(|(row, value)| {
                        let Some(text) = value.map(str::trim).filter(|text| !text.is_empty())
                        else {
                            return Ok(f32::NAN);
                        };
                        value_to_idx.get(text).copied().with_context(|| {
                            format!("Enum column '{name}' has unknown value '{text}' at row {row}")
                        })
                    })
                    .collect::<anyhow::Result<Vec<_>>>()?;

                tracing::debug!(column = %name, unique_values = unique.len(), "Enum feature encoded as f32");
                Ok((unique, encoded))
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        // Extract is_approx_build_date: 0.0 = exact, anything else (1.0/NaN) = approximate
        let is_approx_build_date_raw: Vec<bool> = if has_approx_col {
            let column_data = df
                .column("Is construction date approximate")
                .context("Missing 'Is construction date approximate' column")?;
            let float_series = column_data
                .cast(&DataType::Float32)
                .context("Failed to cast 'Is construction date approximate' to Float32")?;
            let chunked = float_series
                .f32()
                .context("Failed to read 'Is construction date approximate' as f32")?;
            chunked
                .into_iter()
                .map(|value| match value {
                    Some(0.0) => false,
                    _ => true, // 1.0 or NaN → approximate
                })
                .collect()
        } else {
            vec![true; row_count] // default: all approximate
        };

        // Extract renovation_history: List<Struct{year: i32, event: str}>
        let mut renovation_raw: FxHashMap<u32, Vec<RenovationEvent>> = if has_renovation_history {
            tracing::info!("Extracting renovation history");
            let reno_col = df
                .column("renovation_history")
                .context("Missing renovation_history column")?;
            let list_ca = reno_col
                .list()
                .context("renovation_history is not a list column")?;

            let mut history: FxHashMap<u32, Vec<RenovationEvent>> = FxHashMap::default();
            for old_row in 0..row_count {
                if let Some(inner) = list_ca.get_as_series(old_row) {
                    if inner.is_empty() {
                        continue;
                    }
                    let structs = inner
                        .struct_()
                        .context("renovation_history inner is not a struct")?;
                    let years = structs
                        .field_by_name("year")
                        .context("Missing 'year' field in renovation_history struct")?;
                    let events = structs
                        .field_by_name("event")
                        .context("Missing 'event' field in renovation_history struct")?;

                    let mut row_events = Vec::new();
                    for idx in 0..inner.len() {
                        let year = years.get(idx).context("Failed to get year value")?;
                        let event = events.get(idx).context("Failed to get event value")?;
                        if let (AnyValue::Int32(yr), AnyValue::String(ev)) = (&year, &event) {
                            row_events.push(RenovationEvent {
                                year: *yr,
                                event: ev.to_string(),
                            });
                        }
                    }
                    if !row_events.is_empty() {
                        history.insert(old_row as u32, row_events);
                    }
                }
            }
            tracing::info!(
                properties_with_events = history.len(),
                "Renovation history extracted"
            );
            history
        } else {
            FxHashMap::default()
        };

        // Extract historical_prices: List<Struct{year: i32, month: u8, price: i64}>
        let mut historical_prices_raw: FxHashMap<u32, Vec<HistoricalPrice>> =
            if has_historical_prices {
                tracing::info!("Extracting historical prices");
                let prices_col = df
                    .column("historical_prices")
                    .context("Missing historical_prices column")?;
                let list_ca = prices_col
                    .list()
                    .context("historical_prices is not a list column")?;

                let mut history: FxHashMap<u32, Vec<HistoricalPrice>> = FxHashMap::default();
                for old_row in 0..row_count {
                    if let Some(inner) = list_ca.get_as_series(old_row) {
                        if inner.is_empty() {
                            continue;
                        }
                        let structs = inner
                            .struct_()
                            .context("historical_prices inner is not a struct")?;
                        let years = structs
                            .field_by_name("year")
                            .context("Missing 'year' field in historical_prices struct")?;
                        let months = structs
                            .field_by_name("month")
                            .context("Missing 'month' field in historical_prices struct")?;
                        let prices = structs
                            .field_by_name("price")
                            .context("Missing 'price' field in historical_prices struct")?;

                        let mut row_prices = Vec::new();
                        for idx in 0..inner.len() {
                            let year = years.get(idx).context("Failed to get year value")?;
                            let month = months.get(idx).context("Failed to get month value")?;
                            let price = prices.get(idx).context("Failed to get price value")?;
                            let AnyValue::Int32(year_i32) = year else {
                                bail!("historical_prices.year is not Int32 at row {old_row}, got {year:?}");
                            };
                            let AnyValue::UInt8(month_u8) = month else {
                                bail!("historical_prices.month is not UInt8 at row {old_row}, got {month:?}");
                            };
                            let AnyValue::Int64(price_i64) = price else {
                                bail!("historical_prices.price is not Int64 at row {old_row}, got {price:?}");
                            };
                            row_prices.push(HistoricalPrice {
                                year: year_i32,
                                month: month_u8,
                                price: price_i64,
                            });
                        }
                        if !row_prices.is_empty() {
                            history.insert(old_row as u32, row_prices);
                        }
                    }
                }
                tracing::info!(
                    properties_with_prices = history.len(),
                    "Historical prices extracted"
                );
                history
            } else {
                FxHashMap::default()
            };

        // Extract tenure_history: List<Struct{year: i32, status: str}>
        let mut tenure_raw: FxHashMap<u32, Vec<TenureEvent>> = if has_tenure_history {
            tracing::info!("Extracting tenure history");
            let tenure_col = df
                .column("tenure_history")
                .context("Missing tenure_history column")?;
            let list_ca = tenure_col
                .list()
                .context("tenure_history is not a list column")?;

            let mut history: FxHashMap<u32, Vec<TenureEvent>> = FxHashMap::default();
            for old_row in 0..row_count {
                if let Some(inner) = list_ca.get_as_series(old_row) {
                    if inner.is_empty() {
                        continue;
                    }
                    let structs = inner
                        .struct_()
                        .context("tenure_history inner is not a struct")?;
                    let years = structs
                        .field_by_name("year")
                        .context("Missing 'year' field in tenure_history struct")?;
                    let statuses = structs
                        .field_by_name("status")
                        .context("Missing 'status' field in tenure_history struct")?;

                    let mut row_events = Vec::new();
                    for idx in 0..inner.len() {
                        let year = years.get(idx).context("Failed to get year value")?;
                        let status = statuses.get(idx).context("Failed to get status value")?;
                        if let (AnyValue::Int32(yr), AnyValue::String(st)) = (&year, &status) {
                            row_events.push(TenureEvent {
                                year: *yr,
                                status: st.to_string(),
                            });
                        }
                    }
                    if !row_events.is_empty() {
                        history.insert(old_row as u32, row_events);
                    }
                }
            }
            tracing::info!(
                properties_with_tenure_changes = history.len(),
                "Tenure history extracted"
            );
            history
        } else {
            FxHashMap::default()
        };

        // Free the projected joined frame before building the row-major matrix.
        drop(df);

        // Sort all rows by spatial locality so that grid queries access
        // contiguous memory (sequential reads instead of random DRAM accesses).
        tracing::info!("Sorting rows by spatial locality");
        let grid_cell_size = 0.01_f32;
        let min_lat_val = lat.iter().cloned().fold(f32::INFINITY, f32::min) - grid_cell_size;
        let min_lon_val = lon.iter().cloned().fold(f32::INFINITY, f32::min) - grid_cell_size;
        let max_lon_val = lon.iter().cloned().fold(f32::NEG_INFINITY, f32::max) + grid_cell_size;
        let grid_cols = ((max_lon_val - min_lon_val) / grid_cell_size).ceil() as u64 + 1;

        let mut perm: Vec<u32> = (0..row_count as u32).collect();
        perm.par_sort_unstable_by_key(|&perm_index| {
            let grid_row = ((lat[perm_index as usize] - min_lat_val) / grid_cell_size) as u64;
            let grid_col = ((lon[perm_index as usize] - min_lon_val) / grid_cell_size) as u64;
            grid_row * grid_cols + grid_col
        });

        let lat: Vec<f32> = perm
            .iter()
            .map(|&perm_index| lat[perm_index as usize])
            .collect();
        let lon: Vec<f32> = perm
            .iter()
            .map(|&perm_index| lon[perm_index as usize])
            .collect();
        let last_known_price_raw: Vec<f32> = numeric_names
            .iter()
            .position(|name| name == "Last known price")
            .map(|price_idx| {
                perm.iter()
                    .map(|&perm_index| numeric_col_major[price_idx][perm_index as usize])
                    .collect()
            })
            .context("Required numeric column 'Last known price' not configured")?;

        // Build contiguous address buffer and address search index (permuted).
        // Each row's posting lists cover BOTH address spellings we hold — the
        // price-paid form and the EPC form — so a property is findable by either;
        // the display address prefers the price-paid form, falling back to the EPC
        // form for never-sold (EPC-only) dwellings whose price-paid form is null.
        tracing::info!("Building interned strings");
        // Display address for a row, preferring the price-paid form and falling
        // back to the EPC form (so EPC-only dwellings still have a display string).
        let coalesced_display = |old_row: usize| -> &str {
            match pp_address_raw[old_row].as_deref() {
                Some(pp) if !pp.is_empty() => pp,
                _ => epc_address_raw[old_row].as_deref().unwrap_or(""),
            }
        };
        let total_addr_bytes: usize = (0..row_count).map(|row| coalesced_display(row).len()).sum();
        let mut address_buffer = String::with_capacity(total_addr_bytes);
        let mut address_offsets = Vec::with_capacity(row_count);
        let mut address_lengths = Vec::with_capacity(row_count);
        let mut address_token_index: FxHashMap<String, Vec<u32>> = FxHashMap::default();
        let mut address_search_rodeo = lasso::Rodeo::default();
        let mut address_search_token_keys: Vec<lasso::Spur> = Vec::new();
        let mut address_search_token_offsets = Vec::with_capacity(row_count);
        let mut address_search_token_lengths = Vec::with_capacity(row_count);
        for (new_row, &perm_index) in perm.iter().enumerate() {
            let old_row = perm_index as usize;
            let display = coalesced_display(old_row);
            let offset = address_buffer.len() as u32;
            let length = display.len().min(u16::MAX as usize) as u16;
            address_offsets.push(offset);
            address_lengths.push(length);
            address_buffer.push_str(&display[..length as usize]);

            // Tokens cover the display address plus the EPC-form spelling when
            // distinct, so the property is findable by either form (deduped so
            // each token posts this row exactly once).
            let search_tokens =
                merged_address_search_tokens(display, epc_address_raw[old_row].as_deref());
            let token_offset = address_search_token_keys.len() as u32;
            let token_length = search_tokens.len().min(u16::MAX as usize) as u16;
            address_search_token_offsets.push(token_offset);
            address_search_token_lengths.push(token_length);

            for token in search_tokens.iter().take(token_length as usize) {
                let key = address_search_rodeo.get_or_intern(token);
                address_search_token_keys.push(key);

                if is_address_candidate_token(token) {
                    address_token_index
                        .entry(token.clone())
                        .or_default()
                        .push(new_row as u32);
                }
            }
        }
        // Keep every distinctive token: common road words ("high", "church", "station") are
        // exactly what people search, and dropping them made those roads unsearchable while a
        // prefix fallback surfaced the wrong street ("Highbury" for "High"). The candidate scan
        // is bounded per query instead (ADDRESS_SEARCH_CANDIDATE_LIMIT), and stop words are
        // already excluded from the index, so the largest posting lists stay modest.
        let max_postings = address_token_index
            .values()
            .map(Vec::len)
            .max()
            .unwrap_or(0);
        let address_prefix_index = build_address_prefix_index(&address_token_index);
        let address_search_interner = address_search_rodeo.into_reader();
        let address_postings_count: usize = address_token_index.values().map(Vec::len).sum();
        tracing::info!(
            tokens = address_token_index.len(),
            prefixes = address_prefix_index.len(),
            max_postings_per_token = max_postings,
            postings = address_postings_count,
            row_tokens = address_search_token_keys.len(),
            "Address search index built"
        );

        // Intern postcodes (permuted)
        let mut postcode_rodeo = lasso::Rodeo::default();
        let mut postcode_keys: Vec<lasso::Spur> = Vec::with_capacity(row_count);
        let mut postcode_row_index: FxHashMap<lasso::Spur, Vec<u32>> = FxHashMap::default();
        for (new_row, &perm_index) in perm.iter().enumerate() {
            let key = postcode_rodeo.get_or_intern(&postcode_raw[perm_index as usize]);
            postcode_keys.push(key);
            postcode_row_index
                .entry(key)
                .or_default()
                .push(new_row as u32);
        }
        let postcode_interner = postcode_rodeo.into_reader();

        let row_to_poi_metric_idx: Vec<u32> = if poi_metrics.is_empty() {
            vec![NO_POI_METRIC_ROW; row_count]
        } else {
            perm.iter()
                .map(|&old_row| {
                    poi_metric_by_postcode
                        .get(postcode_raw[old_row as usize].as_str())
                        .copied()
                        .unwrap_or(NO_POI_METRIC_ROW)
                })
                .collect()
        };
        poi_metrics.set_row_mapping(row_to_poi_metric_idx);

        // Pack is_approx_build_date into a bitvec (8 bools per byte)
        let num_bytes = row_count.div_ceil(8);
        let mut approx_build_date_bits = vec![0u8; num_bytes];
        for (new_row, &old_row) in perm.iter().enumerate() {
            if is_approx_build_date_raw[old_row as usize] {
                approx_build_date_bits[new_row / 8] |= 1 << (new_row % 8);
            }
        }

        // Re-key renovation_history by permuted row index
        let renovation_history: FxHashMap<u32, Vec<RenovationEvent>> = {
            let mut map =
                FxHashMap::with_capacity_and_hasher(renovation_raw.len(), Default::default());
            for (new_row, &old_row) in perm.iter().enumerate() {
                if let Some(events) = renovation_raw.remove(&old_row) {
                    map.insert(new_row as u32, events);
                }
            }
            map
        };

        // Re-key historical_prices by permuted row index
        let historical_prices: FxHashMap<u32, Vec<HistoricalPrice>> = {
            let mut map = FxHashMap::with_capacity_and_hasher(
                historical_prices_raw.len(),
                Default::default(),
            );
            for (new_row, &old_row) in perm.iter().enumerate() {
                if let Some(prices) = historical_prices_raw.remove(&old_row) {
                    map.insert(new_row as u32, prices);
                }
            }
            map
        };

        // Re-key tenure_history by permuted row index
        let tenure_history: FxHashMap<u32, Vec<TenureEvent>> = {
            let mut map = FxHashMap::with_capacity_and_hasher(tenure_raw.len(), Default::default());
            for (new_row, &old_row) in perm.iter().enumerate() {
                if let Some(events) = tenure_raw.remove(&old_row) {
                    map.insert(new_row as u32, events);
                }
            }
            map
        };

        // Permute optional string columns into sparse HashMaps
        let property_sub_type: FxHashMap<u32, String> = {
            let mut map = FxHashMap::default();
            for (new_row, &old_row) in perm.iter().enumerate() {
                if let Some(ref s) = property_sub_type_raw[old_row as usize] {
                    map.insert(new_row as u32, s.clone());
                }
            }
            map
        };
        let price_qualifier: FxHashMap<u32, String> = {
            let mut map = FxHashMap::default();
            for (new_row, &old_row) in perm.iter().enumerate() {
                if let Some(ref s) = price_qualifier_raw[old_row as usize] {
                    map.insert(new_row as u32, s.clone());
                }
            }
            map
        };

        // Build enum_values map: feature_index -> list of string values
        // and enum_counts map: feature_index -> per-value global counts
        let mut enum_values: rustc_hash::FxHashMap<usize, Vec<String>> =
            rustc_hash::FxHashMap::default();
        let mut enum_counts: rustc_hash::FxHashMap<usize, Vec<u64>> =
            rustc_hash::FxHashMap::default();
        for (enum_idx, (values, encoded)) in enum_col_major.iter().enumerate() {
            let feature_idx = num_numeric + enum_idx;
            enum_values.insert(feature_idx, values.clone());
            let mut counts = vec![0u64; values.len()];
            for &val in encoded {
                if val.is_finite() {
                    let idx = val as usize;
                    if idx < counts.len() {
                        counts[idx] += 1;
                    }
                }
            }
            enum_counts.insert(feature_idx, counts);
        }

        // Build feature_stats: numeric stats + placeholder stats for enums
        let mut feature_stats = numeric_feature_stats;
        for (values, _) in &enum_col_major {
            // For enum features, slider range is 0 to num_values-1
            let num_values = values.len();
            let max_val = num_values as f32;
            feature_stats.push(FeatureStats {
                slider_min: 0.0,
                slider_max: (num_values.saturating_sub(1)) as f32,
                histogram: Histogram {
                    min: 0.0,
                    max: max_val,
                    p1: 0.0,
                    p99: max_val,
                    counts: vec![0; num_values.max(1)],
                },
            });
            // Enum features: not quantized, stored directly as u16
            quant_min.push(0.0);
            quant_range.push(0.0);
        }
        let dequant_a: Vec<f32> = quant_range
            .iter()
            .map(|&r| if r > 0.0 { r / QUANT_SCALE } else { 0.0 })
            .collect();

        // Transpose to row-major AND apply spatial permutation in one pass.
        // Combines numeric and enum features into a single feature_data array, quantized to u16.
        tracing::info!("Transposing to row-major layout (spatially sorted, quantized to u16)");
        let mut feature_data = vec![NAN_U16; row_count * num_features];
        feature_data
            .par_chunks_mut(num_features)
            .enumerate()
            .for_each(|(new_row, row_slice)| {
                let old_index = perm[new_row] as usize;
                // Numeric features: quantize to u16
                for (feat_idx, col_vec) in numeric_col_major.iter().enumerate() {
                    let value = col_vec[old_index];
                    row_slice[feat_idx] = if value.is_finite() {
                        let range = quant_range[feat_idx];
                        if range > 0.0 {
                            let normalized = (value - quant_min[feat_idx]) / range;
                            (normalized * QUANT_SCALE).round().clamp(0.0, QUANT_SCALE) as u16
                        } else {
                            0
                        }
                    } else {
                        NAN_U16
                    };
                }
                // Enum features: store as u16 directly
                for (enum_idx, (_, encoded)) in enum_col_major.iter().enumerate() {
                    let value = encoded[old_index];
                    row_slice[num_numeric + enum_idx] = if value.is_finite() {
                        value as u16
                    } else {
                        NAN_U16
                    };
                }
            });

        tracing::info!("Data loading complete");

        Ok(PropertyData {
            lat,
            lon,
            feature_names,
            num_features,
            num_numeric,
            feature_data,
            dequant_a,
            quant_min,
            quant_range,
            feature_stats,
            poi_metrics,
            last_known_price_raw,
            address_buffer,
            address_offsets,
            address_lengths,
            postcode_interner,
            postcode_keys,
            postcode_row_index,
            address_token_index,
            address_prefix_index,
            address_search_interner,
            address_search_token_keys,
            address_search_token_offsets,
            address_search_token_lengths,
            enum_values,
            enum_counts,
            approx_build_date_bits,
            renovation_history,
            historical_prices,
            tenure_history,
            property_sub_type,
            price_qualifier,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn country_column_detection_prefers_reliable_country_code() {
        let df = df!(
            "Postcode" => &["SW1A 1AA"],
            "ctry25cd" => &["E92000001"],
            "lat" => &[Some(51.501_f64)],
            "lon" => &[Some(-0.141_f64)],
        )
        .expect("test dataframe should build");

        assert_eq!(
            find_country_column(df.schema()).as_deref(),
            Some("ctry25cd")
        );
    }

    #[test]
    fn missing_coordinate_summary_counts_rows_and_distinct_samples() {
        let summary = summarize_missing_coordinate_postcodes([
            Some("SW1A 1AA"),
            Some("SW1A 1AA"),
            Some("E14 2DG"),
            Some(""),
            None,
        ]);

        assert_eq!(
            summary,
            MissingCoordinateSummary {
                row_count: 5,
                unique_postcode_count: 2,
                sample_postcodes: vec!["E14 2DG".to_string(), "SW1A 1AA".to_string()],
            }
        );
    }

    #[test]
    fn missing_england_coordinates_error_includes_counts_and_samples() {
        let summary = MissingCoordinateSummary {
            row_count: 3,
            unique_postcode_count: 2,
            sample_postcodes: vec!["E14 2DG".to_string(), "SW1A 1AA".to_string()],
        };

        let message = missing_england_coordinates_error(&summary, "ctry25cd");

        assert!(message.contains("3 rows across 2 postcodes"));
        assert!(message.contains("country column 'ctry25cd'"));
        assert!(message.contains("E14 2DG, SW1A 1AA"));
    }

    #[test]
    fn coordinate_validation_errors_for_england_rows() {
        let lf = df!(
            "Postcode" => &["SW1A 1AA", "E14 2DG"],
            "ctry25cd" => &["E92000001", "E92000001"],
            "lat" => &[Some(51.501_f64), None],
            "lon" => &[Some(-0.141_f64), Some(-0.001_f64)],
        )
        .expect("test dataframe should build")
        .lazy();
        let schema = lf.clone().collect_schema().expect("schema should collect");

        let err = validate_no_england_rows_missing_coordinates(&lf, &schema)
            .expect_err("England row with missing coordinate should error");
        let message = err.to_string();

        assert!(message.contains("1 rows across 1 postcodes"));
        assert!(message.contains("E14 2DG"));
    }

    #[test]
    fn coordinate_validation_allows_non_england_rows() {
        let lf = df!(
            "Postcode" => &["CF10 1AA", "SW1A 1AA"],
            "ctry25cd" => &["W92000004", "E92000001"],
            "lat" => &[None, Some(51.501_f64)],
            "lon" => &[Some(-3.179_f64), Some(-0.141_f64)],
        )
        .expect("test dataframe should build")
        .lazy();
        let schema = lf.clone().collect_schema().expect("schema should collect");

        validate_no_england_rows_missing_coordinates(&lf, &schema)
            .expect("non-England row with missing coordinate should be skipped");
    }

    #[test]
    fn coordinate_validation_requires_country_column() {
        let lf = df!(
            "Postcode" => &["SW1A 1AA"],
            "lat" => &[None::<f64>],
            "lon" => &[Some(-0.141_f64)],
        )
        .expect("test dataframe should build")
        .lazy();
        let schema = lf.clone().collect_schema().expect("schema should collect");

        let err = validate_no_england_rows_missing_coordinates(&lf, &schema)
            .expect_err("missing country provenance should error");

        assert!(err.to_string().contains("no reliable country column"));
    }

    #[test]
    fn enum_value_counting() {
        let values = vec![0.0_f32, 1.0, 1.0, 2.0, f32::NAN, 3.0, 1.0];
        let enum_count = 4;

        let mut counts = vec![0u64; enum_count];
        for &v in &values {
            if v.is_finite() {
                let idx = v as usize;
                if idx < enum_count {
                    counts[idx] += 1;
                }
            }
        }

        assert_eq!(counts[0], 1);
        assert_eq!(counts[1], 3);
        assert_eq!(counts[2], 1);
        assert_eq!(counts[3], 1);
    }
}