"""Local classifier for the CHRIS carbon-free ceramic fallback.

The model is intentionally narrow: it estimates whether a carbon-free,
otherwise-unclassified material belongs above the historical 700 C ceramic
melting-point gate.  It is not a general melting-point model.

Runtime data files (version-controlled under ``data/``):

* ``ceramic_classifier_training.csv``
    Trusted material formulas and reference MP ranges used to train the model.
* ``magpie_element_properties.csv``
    Magpie elemental-property table copied from CBFV.  Elemental melting-point
    columns are excluded at runtime because the validated production model did
    not need them.

The model trains lazily on first use and is cached for the lifetime of the
process.  ``warm_classifier()`` may be called during application startup if
explicit eager training is preferred.
"""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass
import logging
import os
from pathlib import Path
import re
import threading
from typing import Optional

import numpy as np
import pandas as pd
import chemicals
from rdkit.Chem import AllChem as Chem
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline


_logger = logging.getLogger(__name__)

MODEL_VERSION = "magpie700-v1"
TRAINING_DATA_VERSION = "2026-10-07"
REFERENCE_CUTOFF_C = 700.0
DEFAULT_THRESHOLD = 0.5
RANDOM_STATE = 20261006
N_TREES = 300
EXPECTED_FEATURE_COUNT = 126

_DATA_DIR = Path(__file__).resolve().parent / "data"
_DEFAULT_TRAINING_PATH = _DATA_DIR / "ceramic_classifier_training.csv"
_DEFAULT_MAGPIE_PATH = _DATA_DIR / "magpie_element_properties.csv"

FORMULA_TOKEN_RE = re.compile(r"([A-Z][a-z]?)(\d*)")
ISOTOPE_RE = re.compile(r"\[(?:\d+)?([A-Z][a-z]?)\]")

# Deliberately narrow warnings copied from the material-formula preparation
# workflow.  They prevent molecular-database formula artifacts from being used
# as material stoichiometry for unseen compounds.
_MISSING_H_NAME_RE = re.compile(
    r"(?:\bborohydride\b|^\s*(?:di|tri|tetra|penta|hexa|hepta|octa|nona|deca)?borane\b)",
    flags=re.IGNORECASE,
)
_ANHYDROUS_SOLID_RE = re.compile(
    r"(?:\b(?:oxide|sulfide|selenide|telluride|silicide|nitride|carbide|boride|"
    r"phosphide|arsenide|fluoride|chloride|bromide|iodide|silicate|aluminate|"
    r"titanate|zirconate|ferrite|chromite|spinel)\b)",
    flags=re.IGNORECASE,
)
_HYDROGEN_ALLOWED_RE = re.compile(
    r"(?:(?:mono|di|tri|tetra|penta|hexa|hepta|octa|nona|deca)?hydrate\b|"
    r"\bhemihydrate\b|\bhydroxide\b|\bhydride\b|\bhydrogen\b|"
    r"\bdeuterium\b|\btritium\b|\bammonium\b|\bphosphonium\b|"
    r"\bphosphine\b|\barsine\b|\bstibine\b|\bsilane\b|"
    r"\bacid\b|\bborohydride\b)",
    flags=re.IGNORECASE,
)


@dataclass(frozen=True)
class CeramicPrediction:
    """Result returned by :func:`predict_high_mp`."""

    is_high_mp: bool
    probability: float
    threshold: float
    formula: str
    model_version: str = MODEL_VERSION
    reference_cutoff_C: float = REFERENCE_CUTOFF_C

    @property
    def origin(self) -> str:
        return f"{self.model_version}/classifier p={self.probability:.3f}"


@dataclass
class _ModelState:
    pipeline: Pipeline
    feature_columns: list[str]
    magpie_table: pd.DataFrame
    trusted_formula_by_cas: dict[str, str]
    n_training_rows: int


_MODEL_STATE: Optional[_ModelState] = None
_MODEL_ERROR: Optional[str] = None
_MODEL_LOCK = threading.Lock()


def _disabled() -> bool:
    return os.environ.get("DISABLE_CERAMIC_CLASSIFIER", "").strip().lower() in {
        "1",
        "true",
        "yes",
    }


def _runtime_path(env_name: str, default: Path) -> Path:
    value = os.environ.get(env_name)
    return Path(value).expanduser().resolve() if value else default


def _training_path() -> Path:
    return _runtime_path("CERAMIC_CLASSIFIER_TRAINING_PATH", _DEFAULT_TRAINING_PATH)


def _magpie_path() -> Path:
    return _runtime_path("CERAMIC_CLASSIFIER_MAGPIE_PATH", _DEFAULT_MAGPIE_PATH)


def _formula_counts(formula: object) -> Counter[str] | None:
    if formula is None or (isinstance(formula, float) and np.isnan(formula)):
        return None
    text = str(formula).strip()
    if not text:
        return None

    text = ISOTOPE_RE.sub(r"\1", text)
    text = re.sub(r"[+-]\d*$", "", text)
    text = re.sub(r"[+-]$", "", text)

    counts: Counter[str] = Counter()
    pos = 0
    for match in FORMULA_TOKEN_RE.finditer(text):
        if match.start() != pos:
            return None
        symbol, number = match.groups()
        if symbol in {"D", "T"}:
            symbol = "H"
        counts[symbol] += int(number) if number else 1
        pos = match.end()

    if pos != len(text) or not counts:
        return None
    return counts


def _canonical_formula(counts: Counter[str] | None) -> str | None:
    if not counts:
        return None
    symbols = list(counts)
    if "C" in counts:
        order = ["C"]
        if "H" in counts:
            order.append("H")
        order.extend(sorted(symbol for symbol in symbols if symbol not in {"C", "H"}))
    else:
        order = sorted(symbols)

    parts: list[str] = []
    for symbol in order:
        n = int(counts[symbol])
        parts.append(symbol)
        if n != 1:
            parts.append(str(n))
    return "".join(parts)


def _heavy_counts(counts: Counter[str] | None) -> Counter[str] | None:
    if counts is None:
        return None
    return Counter({symbol: n for symbol, n in counts.items() if symbol != "H" and n})


def _smiles_heavy_counts(smiles: object) -> Counter[str] | None:
    if not isinstance(smiles, str) or not smiles.strip():
        return None
    try:
        mol = Chem.MolFromSmiles(smiles)
    except Exception:
        mol = None
    if mol is None:
        return None
    return Counter(
        atom.GetSymbol()
        for atom in mol.GetAtoms()
        if atom.GetSymbol() != "H"
    )


def _safe_database_formula(identifier: str | None, name: str | None, smiles: str | None) -> str | None:
    """Return a conservative material-formula candidate from local metadata.

    The chemicals database formula is accepted only when its heavy-element
    stoichiometry exactly matches the atom nodes in the supplied SMILES and it
    does not trip known hydrogen-artifact warnings.  No implicit RDKit hydrogen
    is ever added here.
    """

    raw_formula = None
    for query in (identifier, name):
        if not query:
            continue
        try:
            record = chemicals.search_chemical(query)
            raw_formula = record.formula
            if raw_formula:
                break
        except Exception:
            continue
    if not raw_formula:
        return None

    counts = _formula_counts(raw_formula)
    if counts is None:
        return None

    smiles_heavy = _smiles_heavy_counts(smiles)
    if smiles_heavy is None or _heavy_counts(counts) != smiles_heavy:
        return None

    text = name or ""
    h_count = int(counts.get("H", 0))
    if h_count == 0 and _MISSING_H_NAME_RE.search(text):
        return None
    if (
        h_count > 0
        and _ANHYDROUS_SOLID_RE.search(text)
        and not _HYDROGEN_ALLOWED_RE.search(text)
    ):
        return None

    return _canonical_formula(counts)


def _load_magpie_table() -> pd.DataFrame:
    path = _magpie_path()
    if not path.exists():
        raise FileNotFoundError(
            f"Ceramic classifier Magpie table not found: {path}"
        )

    table = pd.read_csv(path)
    if "element" not in table.columns:
        raise ValueError(f"Magpie table {path} has no 'element' column")

    table = table.copy().set_index("element", drop=True)
    for column in table.columns:
        table[column] = pd.to_numeric(table[column], errors="coerce")

    # The validated production candidate excluded elemental melting-point
    # descriptors.  Keep the source table unmodified and drop them here.
    melting_columns = [column for column in table.columns if "melt" in column.lower()]
    table = table.drop(columns=melting_columns)

    feature_count = len(table.columns) * 6
    if feature_count != EXPECTED_FEATURE_COUNT:
        raise ValueError(
            "Unexpected Magpie schema after excluding elemental melting-point "
            f"columns: expected {EXPECTED_FEATURE_COUNT} features, got {feature_count}. "
            "Use the same Magpie table version used for validation."
        )
    return table


def _magpie_feature_columns(table: pd.DataFrame) -> list[str]:
    return [
        f"{stat}_{prop}"
        for stat in ("avg", "dev", "range", "max", "min", "mode")
        for prop in table.columns
    ]


def _magpie_vector(counts: Counter[str], table: pd.DataFrame) -> np.ndarray | None:
    symbols = list(counts)
    available = set(table.index.astype(str))
    if not set(symbols).issubset(available):
        return None

    atom_counts = np.asarray([counts[symbol] for symbol in symbols], dtype=float)
    fractions = atom_counts / atom_counts.sum()
    values = table.loc[symbols].to_numpy(dtype=float)

    with np.errstate(invalid="ignore", divide="ignore"):
        avg = np.sum(values * fractions[:, None], axis=0)
        dev = np.sum(np.abs(values - avg[None, :]) * fractions[:, None], axis=0)
        value_range = np.ptp(values, axis=0)
        value_max = np.max(values, axis=0)
        value_min = np.min(values, axis=0)
        prominent = np.isclose(fractions, np.max(fractions))
        mode = np.min(values[prominent, :], axis=0)

    return np.concatenate([avg, dev, value_range, value_max, value_min, mode])


def _build_feature_frame(formulas: pd.Series, table: pd.DataFrame) -> tuple[pd.DataFrame, pd.Series]:
    columns = _magpie_feature_columns(table)
    rows: list[np.ndarray | None] = []
    supported: list[bool] = []

    for formula in formulas:
        counts = _formula_counts(formula)
        vector = _magpie_vector(counts, table) if counts is not None else None
        rows.append(vector)
        supported.append(vector is not None)

    data = np.full((len(rows), len(columns)), np.nan, dtype=float)
    for i, vector in enumerate(rows):
        if vector is not None:
            data[i, :] = vector

    return (
        pd.DataFrame(data, index=formulas.index, columns=columns),
        pd.Series(supported, index=formulas.index, dtype=bool),
    )


def _train_model() -> _ModelState:
    training_path = _training_path()
    if not training_path.exists():
        raise FileNotFoundError(
            f"Ceramic classifier training data not found: {training_path}"
        )

    df = pd.read_csv(training_path, dtype={"CASRN": str})
    required = {"CASRN", "material_formula", "mp_min_C", "mp_max_C"}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(
            f"Ceramic classifier training file is missing columns: {sorted(missing)}"
        )

    df = df.copy()
    df["mp_min_C"] = pd.to_numeric(df["mp_min_C"], errors="coerce")
    df["mp_max_C"] = pd.to_numeric(df["mp_max_C"], errors="coerce")
    df = df.dropna(subset=["material_formula", "mp_min_C", "mp_max_C"])

    # Formula provenance is independent of whether reference MP sources cross
    # the cutoff, so keep all trusted formulas available for runtime lookup.
    trusted_formula_by_cas = {
        str(cas).strip(): str(formula).strip()
        for cas, formula in zip(df["CASRN"], df["material_formula"])
        if pd.notna(cas) and pd.notna(formula)
    }

    # Conservative reference labels: only rows whose complete reference range
    # is on one side of 700 C are used for training.
    low = df["mp_max_C"] <= REFERENCE_CUTOFF_C
    high = df["mp_min_C"] > REFERENCE_CUTOFF_C
    stable = low | high
    df = df.loc[stable].copy()
    y = high.loc[df.index].astype(int).to_numpy()

    if len(df) < 100 or len(np.unique(y)) != 2:
        raise ValueError(
            "Ceramic classifier training data are incomplete: expected a substantial "
            "set containing both LOW and HIGH classes."
        )

    table = _load_magpie_table()
    X, supported = _build_feature_frame(df["material_formula"], table)
    if not bool(supported.all()):
        unsupported = int((~supported).sum())
        raise ValueError(
            f"{unsupported} training formulas are unsupported by the Magpie table"
        )

    pipeline = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="median", keep_empty_features=True)),
            (
                "classifier",
                ExtraTreesClassifier(
                    n_estimators=N_TREES,
                    min_samples_leaf=2,
                    max_features="sqrt",
                    class_weight="balanced",
                    random_state=RANDOM_STATE,
                    n_jobs=1,
                ),
            ),
        ]
    )
    pipeline.fit(X, y)

    _logger.info(
        "Ceramic classifier trained: version=%s rows=%d features=%d threshold=%.2f",
        MODEL_VERSION,
        len(df),
        X.shape[1],
        DEFAULT_THRESHOLD,
    )

    return _ModelState(
        pipeline=pipeline,
        feature_columns=list(X.columns),
        magpie_table=table,
        trusted_formula_by_cas=trusted_formula_by_cas,
        n_training_rows=len(df),
    )


def _ensure_model() -> _ModelState | None:
    global _MODEL_STATE, _MODEL_ERROR

    if _disabled():
        return None
    if _MODEL_STATE is not None:
        return _MODEL_STATE
    if _MODEL_ERROR is not None:
        return None

    with _MODEL_LOCK:
        if _MODEL_STATE is not None:
            return _MODEL_STATE
        if _MODEL_ERROR is not None:
            return None
        try:
            _MODEL_STATE = _train_model()
        except Exception as exc:
            _MODEL_ERROR = f"{type(exc).__name__}: {exc}"
            _logger.warning(
                "Ceramic classifier unavailable; falling back to MP lookups: %s",
                _MODEL_ERROR,
            )
            return None
    return _MODEL_STATE


def warm_classifier() -> bool:
    """Train/load the classifier now instead of waiting for first prediction."""

    return _ensure_model() is not None


def resolve_material_formula(
    *,
    cas: str | None = None,
    name: str | None = None,
    smiles: str | None = None,
    formula: str | None = None,
) -> str | None:
    """Resolve a conservative material formula for classifier input.

    Explicit ``formula`` wins.  Otherwise a trusted formula from the bundled
    training table is used when the CAS is known.  For unseen compounds, a
    local chemicals-database formula is accepted only after conservative
    heavy-stoichiometry and hydrogen-artifact checks.
    """

    if formula:
        counts = _formula_counts(formula)
        return _canonical_formula(counts)

    state = _ensure_model()
    if state is None:
        return None

    clean_cas = None if cas in {None, "", "Not found"} else str(cas).strip()
    clean_name = None if name in {None, "", "Not found"} else str(name).strip()

    if clean_cas and clean_cas in state.trusted_formula_by_cas:
        return state.trusted_formula_by_cas[clean_cas]

    return _safe_database_formula(clean_cas, clean_name, smiles)


def predict_high_mp(
    *,
    cas: str | None = None,
    name: str | None = None,
    smiles: str | None = None,
    formula: str | None = None,
    threshold: float = DEFAULT_THRESHOLD,
) -> CeramicPrediction | None:
    """Predict whether an unresolved carbon-free material is above 700 C.

    Returns ``None`` when the model or a trustworthy material formula is not
    available.  Callers should then fall back to experimental/tabulated MP
    lookup and ultimately to the conservative non-ceramic default.
    """

    state = _ensure_model()
    if state is None:
        return None

    material_formula = resolve_material_formula(
        cas=cas,
        name=name,
        smiles=smiles,
        formula=formula,
    )
    if material_formula is None:
        return None

    counts = _formula_counts(material_formula)
    if counts is None:
        return None
    vector = _magpie_vector(counts, state.magpie_table)
    if vector is None:
        return None

    X = pd.DataFrame([vector], columns=state.feature_columns)
    probability = float(state.pipeline.predict_proba(X)[0, 1])
    return CeramicPrediction(
        is_high_mp=probability >= float(threshold),
        probability=probability,
        threshold=float(threshold),
        formula=material_formula,
    )