"""
Author: Wenyu Ouyang
Date: 2024-09-15 11:23:28
LastEditTime: 2025-04-19 14:11:24
LastEditors: Wenyu Ouyang
Description: Test the Resulter class
FilePath: /torchhydro/tests/test_resulter.py
Copyright (c) 2023-2024 Wenyu Ouyang. All rights reserved.
"""

import os
import pytest
import xarray as xr
import numpy as np
import pandas as pd
from torchhydro.trainers.resulter import Resulter


@pytest.fixture
def mock_cfgs(tmpdir):
    return {
        "data_cfgs": {
            "case_dir": str(tmpdir),
            "object_ids": [f"basin_{i}" for i in range(10)],
            "target_cols": ["streamflow", "total_evaporation_hourly"],
        },
        "evaluation_cfgs": {
            "model_loader": {"load_way": "latest"},
            "metrics": ["RMSE", "NSE"],
            "fill_nan": "no",
            "explainer": "none",
        },
        "training_cfgs": {"epochs": 10},
    }


@pytest.fixture
def mock_resulter(mock_cfgs):
    return Resulter(mock_cfgs)


def create_mock_netcdf(file_path, data):
    ds = xr.Dataset({"data": (("basin", "time"), data)})
    ds.to_netcdf(file_path)


def test_load_result(mock_resulter, tmpdir):
    # Create mock prediction and observation NetCDF files
    pred_data = np.random.rand(10, 10)
    obs_data = np.random.rand(10, 10)

    pred_file = os.path.join(tmpdir, "epoch10flow_pred.nc")
    obs_file = os.path.join(tmpdir, "epoch10flow_obs.nc")

    create_mock_netcdf(pred_file, pred_data)
    create_mock_netcdf(obs_file, obs_data)

    # Load the results using the method
    pred, obs = mock_resulter.load_result()

    # Check if the loaded data matches the mock data
    np.testing.assert_array_equal(pred["data"].values, pred_data)
    np.testing.assert_array_equal(obs["data"].values, obs_data)


def test_eval_result(mock_resulter, tmpdir):
    # Create mock prediction and observation xarray datasets
    basins = mock_resulter.cfgs["data_cfgs"]["object_ids"]
    preds_data = {
        "streamflow": (("time", "basin"), np.random.rand(10, 10)),
        "total_evaporation_hourly": (("time", "basin"), np.random.rand(10, 10)),
    }
    obss_data = {
        "streamflow": (("time", "basin"), np.random.rand(10, 10)),
        "total_evaporation_hourly": (("time", "basin"), np.random.rand(10, 10)),
    }
    preds_xr = xr.Dataset(preds_data, coords={"basin": basins, "time": pd.date_range("2000-01-01", periods=10)})
    obss_xr = xr.Dataset(obss_data, coords={"basin": basins, "time": pd.date_range("2000-01-01", periods=10)})

    # Evaluate the results - this should create CSV files
    mock_resulter.eval_result(preds_xr, obss_xr)

    # Check if the evaluation CSV files were created
    for col in mock_resulter.cfgs["data_cfgs"]["target_cols"]:
        output_file = os.path.join(tmpdir, f"metric_{col}.csv")
        assert os.path.exists(output_file)
        # Check content of the CSV
        df = pd.read_csv(output_file)
        assert "basin_id" in df.columns
        assert len(df) == len(basins)
        for metric in mock_resulter.cfgs["evaluation_cfgs"]["metrics"]:
            assert metric in df.columns