from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import pytest

PARQUET_PATH = Path(__file__).parent.parent / "data" / "trajectories.parquet"
JSONL_PATH = Path(__file__).parent.parent / "data" / "synthetic_dataset_v2.jsonl"


def verify_drydown(record: dict) -> bool:
    """
    Check that a trajectory represents a physically plausible dry-down:
      - non-empty trajectory with a finite time span
      - final time is near the 8-hour horizon (or film is depleted earlier)
      - no negative or NaN mole fractions in the trajectory
    """
    traj = record.get("trajectory", [])
    if not traj:
        return False
    if not all(math.isfinite(step["t_s"]) for step in traj):
        return False
    if traj[0]["t_s"] != pytest.approx(0.0, abs=1.0):
        return False
    if traj[-1]["t_s"] < 1000.0:
        return False
    for step in traj:
        x_vals = list(step.get("x_liquid", {}).values())
        if any(v < -1e-6 or not math.isfinite(v) for v in x_vals):
            return False
    return True


def verify_stability(record: dict) -> bool:
    """
    Check thermodynamic/ODE stability:
      - status is 'passed' or 'depleted'
      - depletion rates are in [0, 1]
      - no NaN/Inf in depletion rates
    """
    status = record.get("status", "")
    if status not in ("passed", "depleted"):
        return False
    rates = record.get("depletion_rates", {})
    if not rates:
        return False
    for v in rates.values():
        if not math.isfinite(v) or v < -1e-3 or v > 1.0 + 1e-3:
            return False
    return True


@pytest.fixture(scope="module")
def parquet_records():
    if not PARQUET_PATH.exists():
        pytest.skip(f"Parquet dataset artifact not found: {PARQUET_PATH}")
    table = pq.read_table(PARQUET_PATH)
    rows = table.to_pydict()
    records = []
    for i in range(table.num_rows):
        records.append({k: json.loads(rows[k][i]) if k in ("formula", "depletion_rates", "trajectory") else rows[k][i] for k in rows})
    return records


def test_parquet_exists_and_has_records():
    if not PARQUET_PATH.exists():
        pytest.skip(f"Parquet dataset artifact not found: {PARQUET_PATH}")
    table = pq.read_table(PARQUET_PATH)
    assert table.num_rows >= 5000


def test_parquet_genre_stratification(parquet_records):
    strategies = {}
    for rec in parquet_records:
        strategies[rec["generation_strategy"]] = strategies.get(rec["generation_strategy"], 0) + 1
    expected_genres = {"citrus_cologne", "fougere", "floral_woody", "amber_oriental", "wildcard"}
    assert expected_genres.issubset(strategies.keys())
    for genre in expected_genres:
        assert strategies[genre] >= 900, f"Genre {genre} under-represented: {strategies[genre]}"


def test_verify_drydown_all_records(parquet_records):
    failures = []
    for rec in parquet_records:
        if not verify_drydown(rec):
            failures.append(rec["formula_id"])
    assert not failures, f"verify_drydown failed for {len(failures)} records: {failures[:10]}"


def test_verify_stability_all_records(parquet_records):
    failures = []
    for rec in parquet_records:
        if not verify_stability(rec):
            failures.append(rec["formula_id"])
    assert not failures, f"verify_stability failed for {len(failures)} records: {failures[:10]}"