"""Feature names and tensor layouts for HiCon-GNN."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Tuple


NODE_TYPE_AGGREGATE = 0
NODE_TYPE_MORTAR = 1

EDGE_TYPE_ITZ = 0
EDGE_TYPE_MORTAR_MORTAR = 1
EDGE_TYPE_AGGREGATE_AGGREGATE = 2

MORTAR_NODE_TYPE_SAND = 0
MORTAR_NODE_TYPE_PASTE = 1


def _with_placeholders(names: Tuple[str, ...], n: int = 4) -> Tuple[str, ...]:
    return tuple(names) + tuple(f"_placeholder_{i}" for i in range(n))


AGGREGATE_FEATURES: Tuple[str, ...] = _with_placeholders(
    (
        "aggregate_size_mm",
        "shape_descriptor",
        "crushing_strength_MPa",
        "toughness_J",
        "abrasion_resistance",
        "moisture_content",
        "water_absorption_pct",
        "surface_texture",
        "specific_gravity",
        "bulk_density_kg_m3",
        "area_fraction",
        # Fraction of coarse aggregate replaced by recycled aggregate.
        "rca_replacement_fraction",
        # Index into categoricals.COARSE_AGGREGATE_TYPES.
        "coarse_aggregate_type_id",
    ),
    n=1,
)

MORTAR_FEATURES: Tuple[str, ...] = (
    "mortar_compressive_strength_MPa",
    "mortar_tensile_strength_MPa",
    "mortar_flexural_strength_MPa",
    "mortar_elastic_modulus_GPa",
    "mortar_permeability",
    "mortar_diffusivity",
    "mortar_porosity",
    "mortar_creep_coefficient",
    # Fibre descriptors; zeros denote an unreinforced mix.
    "fibre_content_kg_m3",
    "fibre_aspect_ratio",
    "fibre_tensile_strength_MPa",
    "fibre_modulus_GPa",
    # Voronoi-cell area fraction.
    "mortar_area_fraction",
    # Index into categoricals.FIBRE_TYPES.
    "fibre_type_id",
)

ITZ_FEATURES: Tuple[str, ...] = _with_placeholders(
    (
        "itz_thickness_um",
        "itz_porosity",
        "itz_strength_MPa",
        "itz_elastic_modulus_GPa",
        "itz_permeability",
        "itz_diffusivity",
    )
)

MORTAR_EDGE_FEATURES: Tuple[str, ...] = _with_placeholders(
    (
        "edge_distance_mm",
        "shared_boundary_length_mm",
        "centroid_dx_mm",
        "centroid_dy_mm",
    ),
    n=2,
)

AGGREGATE_EDGE_FEATURES: Tuple[str, ...] = _with_placeholders(
    (
        "agg_agg_surface_gap_mm",
        "agg_agg_centroid_distance_mm",
    ),
    n=2,
)

GLOBAL_FEATURES: Tuple[str, ...] = (
    "relative_humidity",
    "temperature_C",
    "curing_age_days",
    "_placeholder_0",
    "_placeholder_1",
)

CONCRETE_TARGETS: Tuple[str, ...] = (
    "compressive_strength_MPa",
    "tensile_strength_MPa",
    "flexural_strength_MPa",
)

MORTAR_TARGETS: Tuple[str, ...] = (
    "mortar_compressive_strength_MPa",
    "mortar_tensile_strength_MPa",
    "mortar_flexural_strength_MPa",
    "mortar_elastic_modulus_GPa",
    "mortar_porosity",
    "mortar_creep_coefficient",
)

ITZ_TARGETS: Tuple[str, ...] = (
    "itz_thickness_um",
    "itz_porosity",
    "itz_strength_MPa",
    "itz_elastic_modulus_GPa",
)

# Paste properties used for supervision and mortar-scale injection.
PASTE_TARGETS: Tuple[str, ...] = (
    "paste_compressive_strength_MPa",
    "paste_tensile_strength_MPa",
    "paste_elastic_modulus_GPa",
    "paste_fracture_energy_J_m2",
    "paste_porosity",
)

SAND_FEATURES: Tuple[str, ...] = _with_placeholders(
    (
        "sand_size_mm",
        "sand_gradation",
        "sand_fineness_modulus",
        "sand_water_absorption_pct",
        "sand_compressive_strength_MPa",
        "sand_elastic_modulus_GPa",
        "sand_area_fraction",
        "sand_specific_gravity",
    ),
    n=2,
)

CEMENT_PASTE_FEATURES: Tuple[str, ...] = (
    "cement_type_id",
    "cement_CaO_pct",
    "cement_SiO2_pct",
    "cement_Al2O3_pct",
    "cement_Fe2O3_pct",
    "cement_content_kg_m3",
    "scm_type_id",
    "scm_CaO_pct",
    "scm_SiO2_pct",
    "scm_content_kg_m3",
    "water_to_binder_ratio",
    "admixture_dosage",
    "paste_area_fraction",
    # Alkali is reported as Na2O equivalent.
    "cement_MgO_pct",
    "cement_SO3_pct",
    "cement_alkali_pct",
    # Additional SCM chemistry.
    "scm_Al2O3_pct",
    "scm_Fe2O3_pct",
    "scm_MgO_pct",
    "scm_LOI_pct",
    # Bogue phases and SCM fineness.
    "cement_C3S_pct",
    "cement_C2S_pct",
    "cement_C3A_pct",
    "cement_C4AF_pct",
    "scm_specific_surface_m2_kg",
    # Glass-powder dosage and fly-ash class.
    "glass_powder_kg_m3",
    "fly_ash_type_id",
    # Cement strength class.
    "cement_grade_mpa",
)

PASTE_FEATURES_WITH_PROPS: Tuple[str, ...] = CEMENT_PASTE_FEATURES + PASTE_TARGETS

MORTAR_GLOBAL_FEATURES: Tuple[str, ...] = (
    "mortar_curing_relative_humidity",
    "mortar_curing_temperature_C",
    "mortar_curing_age_days",
    # Fibre descriptors are global because paste and sand nodes do not carry them.
    "mortar_fibre_content_kg_m3",
    "mortar_fibre_length_mm",
    "mortar_fibre_diameter_mm",
    "mortar_fibre_tensile_strength_MPa",
    "mortar_fibre_modulus_GPa",
    "mortar_fibre_type_id",
    "_placeholder_0",
    "_placeholder_1",
)


@dataclass(frozen=True)
class SchemaSpec:
    """Bundle of schema vectors with derived size helpers."""

    aggregate: Tuple[str, ...] = field(default=AGGREGATE_FEATURES)
    mortar: Tuple[str, ...] = field(default=MORTAR_FEATURES)
    itz: Tuple[str, ...] = field(default=ITZ_FEATURES)
    mortar_edge: Tuple[str, ...] = field(default=MORTAR_EDGE_FEATURES)
    aggregate_edge: Tuple[str, ...] = field(default=AGGREGATE_EDGE_FEATURES)
    glob: Tuple[str, ...] = field(default=GLOBAL_FEATURES)
    concrete_targets: Tuple[str, ...] = field(default=CONCRETE_TARGETS)
    mortar_targets: Tuple[str, ...] = field(default=MORTAR_TARGETS)
    itz_targets: Tuple[str, ...] = field(default=ITZ_TARGETS)
    paste_targets: Tuple[str, ...] = field(default=PASTE_TARGETS)
    sand: Tuple[str, ...] = field(default=SAND_FEATURES)
    paste: Tuple[str, ...] = field(default=PASTE_FEATURES_WITH_PROPS)
    mortar_global: Tuple[str, ...] = field(default=MORTAR_GLOBAL_FEATURES)

    @property
    def node_pad_dim(self) -> int:
        return max(len(self.aggregate), len(self.mortar))

    @property
    def edge_pad_dim(self) -> int:
        return max(len(self.itz), len(self.mortar_edge), len(self.aggregate_edge))

    @property
    def mortar_node_pad_dim(self) -> int:
        return max(len(self.sand), len(self.paste))

    @property
    def global_dim(self) -> int:
        return len(self.glob)

    @property
    def mortar_global_dim(self) -> int:
        return len(self.mortar_global)

    @property
    def num_node_types(self) -> int:
        return 2

    @property
    def num_edge_types(self) -> int:
        return 3

    @property
    def num_mortar_node_types(self) -> int:
        return 2


DEFAULT_SCHEMA = SchemaSpec()
