import json

import numpy as np
import xarray as xr

from dr_gcm.inventory import build_inventory, has_failures, inventory_file


def _write_field(path, variable="t2m", lon=(105.0, 105.25), lat=(20.0, 19.75)):
    time = np.array(["2000-01-01", "2000-02-01", "2000-04-01"], dtype="datetime64[ns]")
    values = np.arange(12, dtype=float).reshape(3, 2, 2)
    values[0, 0, 0] = np.nan
    xr.Dataset(
        {variable: (("time", "latitude", "longitude"), values, {"units": "degC"})},
        coords={"time": time, "latitude": list(lat), "longitude": list(lon)},
    ).to_netcdf(path)


def test_inventory_reports_grid_time_gap_and_statistics(tmp_path):
    path = tmp_path / "sample.nc"
    _write_field(path)
    report = inventory_file(path, expected_variables=["t2m", "pr"], statistics=True)
    assert report["status"] == "ok"
    assert report["coordinates"]["time"]["first"].startswith("2000-01-01")
    assert report["missing_expected_variables"] == ["pr"]
    assert report["coordinates"]["latitude"]["monotonic"] == "decreasing"
    assert report["coordinates"]["time"]["estimated_missing_steps"] == 1
    assert report["variables"]["t2m"]["value_statistics"]["missing_count"] == 1


def test_config_inventory_compares_grids_and_strict_failures(tmp_path):
    first, second, mask = (tmp_path / name for name in ("pr.nc", "t2m.nc", "mask.nc"))
    _write_field(first, variable="r")
    _write_field(second, variable="t2m")
    xr.Dataset(
        {"reg": (("latitude", "longitude"), np.ones((2, 2)))},
        coords={"latitude": [20.0, 19.75], "longitude": [105.0, 105.25]},
    ).to_netcdf(mask)
    config = tmp_path / "inventory.json"
    config.write_text(json.dumps({"datasets": [
        {"name": "rainfall_observation", "paths": [str(first)], "expected_variables": ["r"]},
        {"name": "temperature_observation", "paths": [str(second)], "expected_variables": ["t2m"]},
        {"name": "vietnam_climate_mask", "paths": [str(mask)], "expected_variables": ["reg"]},
    ]}), encoding="utf-8")
    report = build_inventory(config)
    assert report["compatibility"]["rainfall_temperature_grid_match"] is True
    assert report["compatibility"]["rainfall_mask_grid_match"] is True
    assert has_failures(report) is False


def _write_initialized_field(path, initial_month):
    init = np.datetime64(initial_month + "-01T00:00:00")
    xr.Dataset(
        {"t2m": (("forecast_reference_time", "number", "latitude", "longitude"),
                 np.ones((1, 2, 1, 1)))},
        coords={"forecast_reference_time": [init], "number": [0, 1],
                "latitude": [20.0], "longitude": [105.0]},
    ).to_netcdf(path)


def test_inventory_checks_continuous_and_paired_initializations(tmp_path):
    for group, months in {"instant": ("2000-01", "2000-03"),
                          "accum": ("2000-01", "2000-02", "2000-03")}.items():
        for month in months:
            _write_initialized_field(tmp_path / f"{group}_{month}.nc", month)
    config = tmp_path / "inventory.json"
    config.write_text(json.dumps({
        "datasets": [
            {"name": "instant", "paths": [str(tmp_path / "instant_*.nc")],
             "expected_variables": ["t2m"]},
            {"name": "accum", "paths": [str(tmp_path / "accum_*.nc")],
             "expected_variables": ["t2m"]},
        ],
        "paired_initializations": [{"name": "raw_pair", "datasets": ["instant", "accum"]}],
    }), encoding="utf-8")
    report = build_inventory(config)
    instant = report["datasets"][0]["summary"]["initialization_months"]
    assert instant["missing_months"] == ["2000-02"]
    pair = report["compatibility"]["paired_initializations"][0]
    assert pair["match"] is False
    assert pair["missing_by_dataset"] == {"instant": ["2000-02"]}
    assert has_failures(report) is True