import os
import tempfile
import numpy as np
import pytest

from pwml.utilities import filehelpers as fh


def test_save_load_ndarray_roundtrip():
    """load_ndarray must return the saved array (was silently returning None)."""
    arr = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
    with tempfile.TemporaryDirectory() as tmpdir:
        path = os.path.join(tmpdir, 'test')
        fh.save_ndarray(path, arr)
        loaded = fh.load_ndarray(path + '.npz')
        assert loaded is not None, "load_ndarray returned None (missing return statement)"
        np.testing.assert_array_equal(arr, loaded)


def test_save_load_ndarray_1d():
    arr = np.arange(10, dtype=np.int64)
    with tempfile.TemporaryDirectory() as tmpdir:
        path = os.path.join(tmpdir, 'arr1d')
        fh.save_ndarray(path, arr)
        loaded = fh.load_ndarray(path + '.npz')
        np.testing.assert_array_equal(arr, loaded)


def test_save_load_pickle_roundtrip():
    data = {'key': [1, 2, 3], 'nested': {'a': 'b'}}
    with tempfile.TemporaryDirectory() as tmpdir:
        path = os.path.join(tmpdir, 'data.pkl')
        fh.save_to_pickle(path, data)
        loaded = fh.load_from_pickle(path)
        assert loaded == data


def test_save_to_file():
    with tempfile.TemporaryDirectory() as tmpdir:
        path = os.path.join(tmpdir, 'out.txt')
        fh.save_to_file(path, 'hello world')
        with open(path) as f:
            assert f.read() == 'hello world'


def test_save_load_cache_dict_file_roundtrip():
    """save_cache_dict_file/load_cache_dict_file must round-trip a dict (B1 fix)."""
    data = {'key1': [1, 2, 3], 'key2': {'nested': True}, 'key3': 42}
    with tempfile.TemporaryDirectory() as tmpdir:
        path = os.path.join(tmpdir, 'cache.pkl')
        fh.save_cache_dict_file(path, data)
        loaded = fh.load_cache_dict_file(path)
        assert loaded == data


def test_save_load_cache_dict_file_with_numpy():
    """save_cache_dict_file must handle dicts containing numpy arrays."""
    arr = np.array([1.0, 2.0, 3.0])
    data = {'arr': arr, 'label': 'test'}
    with tempfile.TemporaryDirectory() as tmpdir:
        path = os.path.join(tmpdir, 'cache.pkl')
        fh.save_cache_dict_file(path, data)
        loaded = fh.load_cache_dict_file(path)
        np.testing.assert_array_equal(loaded['arr'], arr)
        assert loaded['label'] == 'test'
