import logging
import unittest.mock as mock
from pathlib import Path
from typing import Any

import numpy as np
import pytest
from astropy.table import Table
from torchvision.transforms.v2 import CenterCrop, Lambda

from hyrax.data_sets.hsc_data_set import FitsImageDataSet, HSCDataSet

test_dir = Path(__file__).parent / "test_data" / "dataloader"


@pytest.fixture(scope="module")
def hsc_test_mode():
    """Turn on unit test mode for HSC and FitsImage datasets. This allows them to run without metadata"""
    HSCDataSet._called_from_test = True
    FitsImageDataSet._called_from_test = True


class FakeFitsFS:
    """
    Mocks the only operations on a directory full of fits files that
    the dataloader should care about for initialization:

    1) Globbing fits files with Path.glob

    2) Reading the shape of the first table out of the fits header using
    astropy.io.fits.open

    If we end up doing anything more sophisticated this will need to be changed
    We are planning to scan 1M - 10M fits files so its unlikely we want to do
    more filesystem operations without a really good reason.
    """

    def __init__(self, test_files: dict, filter_catalog: dict | None = None):
        self.patchers: list[mock._patch[Any]] = []

        self.test_files = test_files

        mock_paths = [Path(x) for x in list(test_files.keys())]
        target = "hyrax.data_sets.hsc_data_set.Path.iterdir"
        self.patchers.append(mock.patch(target, return_value=mock_paths))

        mock_fits_open = mock.Mock(side_effect=self._open_file)
        self.patchers.append(mock.patch("astropy.io.fits.open", mock_fits_open))

        if filter_catalog is not None:
            mock_read_filter_catalog = mock.patch(
                "hyrax.data_sets.hsc_data_set.HSCDataSet._read_filter_catalog",
                lambda x, y: Table(
                    {"object_id": ["notreal"], "filter": ["notreal"], "filename": ["notreal"]}
                ),
            )
            self.patchers.append(mock_read_filter_catalog)
            mock_parse_filter_catalog = mock.patch(
                "hyrax.data_sets.hsc_data_set.FitsImageDataSet._parse_filter_catalog",
                lambda x, y: filter_catalog,
            )
            self.patchers.append(mock_parse_filter_catalog)

    def _open_file(self, filename: Path, **kwargs) -> mock.Mock:
        shape = self.test_files[filename.name]
        mock_open_ctx = mock.Mock()
        mock_open_ctx.__enter__ = mock.Mock(return_value=["", np.zeros(shape)])
        mock_open_ctx.__exit__ = mock.Mock()
        return mock_open_ctx

    def __enter__(self):
        for patcher in self.patchers:
            patcher.start()

    def __exit__(self, *exc):
        for patcher in self.patchers:
            patcher.stop()


def mkconfig(
    crop_to=False,
    filters=False,
    filter_catalog=False,
    use_cache=False,
    transform="tanh",
):
    """Makes a configuration that points at nonexistent path so HSCDataSet.__init__ will create an object,
    and our FakeFitsFS shim can be called.
    """
    return {
        "general": {"data_dir": "thispathdoesnotexist"},
        "data_set": {
            "crop_to": crop_to,
            "filters": filters,
            "filter_catalog": filter_catalog,
            "use_cache": use_cache,
            "transform": transform,
            "preload_cache": False,  # Don't run the preloading in unit tests, because it needs real data.
            "object_id_column_name": False,
            "filter_column_name": False,
            "filename_column_name": False,
        },
    }


def generate_files(
    num_objects=10, num_filters=5, shape=(100, 100), offset=0, infill_str="all_filters"
) -> dict:
    """Generates a dictionary to pass in to FakeFitsFS.

    This generates a dict from filename->shape tuple for a set of uniform fake fits files
    corresponding to the naming convention used by the HSC data loader.

    Completely uniform data sets with all filters for all files with every file the same
    size are generated by this function. Sequential object_ids are used starting from 0
    by default.

    Parameters
    ----------
    num_objects : int, optional
        How many objects are represented in the data set, by default 10
    num_filters : int, optional
        How many filters does each object have, by default 5. If you provide a number greater than 5,
        only 5 filters will be output.
    shape : tuple, optional
        What are the dimensions of the image in each fits file, by default (100,100)
    offset : int, optional
        What is the first object_id to start with, by default 0
    infill_str: str, optional
        What to put in the fake filename in between the object ID and filter name. By default "all_filters"

    Returns
    -------
    dict
        Dictionary from filename -> shape appropriate for FakeFitsFS
    """
    filters = ["HSC-G", "HSC-R", "HSC-I", "HSC-Z", "HSC-Y"][:num_filters]
    test_files = {}
    for object_id in range(offset, num_objects + offset):
        for filter in filters:
            test_files[f"{object_id:017d}_{infill_str}_{filter}.fits"] = shape

    return test_files


def generate_filter_catalog(test_files: dict, config: dict, tmp_path: str) -> dict:
    """Generates a filter catalog dict for use with FakeFitsFS from a filesystem dictionary
    created by generate_files.

    This allows tests to alter the parsed filter_catalog, and interrogate what decisions HSCDataSet makes
    when a manifest or filter_catalog file contains corrupt information.

    Caveats:
    Always Run this prior to creating any mocks, or using a FakeFitsFS

    This function calls the HSCDataSet parsing code to actually assemble the files array, which means:
    1) This test setup is fairly tightly coupled to the internals of HSCDataSet
    2) When writing tests using this, ensure any functionality you depend on from the slow file-scan
       initialization codepath is tested indepently.

    Parameters
    ----------
    test_files : dict
        Test files dictionary created with generate_files.
    config : dict
        The config you will use to initialize the test HSCDataSet object in your test.
        Create this with mkconfig()
    tmp_path : str
        A temporary path to pass to HSCDataSet as the data_directory.

    Returns
    -------
    dict
        Dictionary from ObjectID -> (Dictionary from Filter -> Filename)
    """
    # Ensure the filter_catalog codepath won't be triggered
    config["data_set"]["filter_catalog"] = False

    # Use our initialization code to create a parsed files object
    with FakeFitsFS(test_files):
        return HSCDataSet(config, data_location=tmp_path).files


def test_load(caplog, hsc_test_mode, tmp_path):
    """Test to ensure loading a perfectly regular set of files works"""
    caplog.set_level(logging.WARNING)
    test_files = generate_files(num_objects=10, num_filters=5, shape=(262, 263))
    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(), data_location=tmp_path)

        # 10 objects should load
        assert len(a) == 10

        # The number of filters, and image dimensions should be correct and square
        assert a.shape() == (5, 262, 262)

        # No warnings should be printed
        assert caplog.text == ""


def test_load_duplicate(caplog, hsc_test_mode, tmp_path):
    """Test to ensure duplicate fits files that reference the same object id and filter create the
    appropriate error messages.
    """
    caplog.set_level(logging.ERROR)
    test_files = generate_files(num_objects=10, num_filters=5, shape=(262, 263))
    duplicate_files = generate_files(num_objects=10, num_filters=5, shape=(262, 263), infill_str="duplicate")
    test_files.update(duplicate_files)
    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(), data_location=tmp_path)

        # Only 10 objects should load
        assert len(a) == 10

        # The number of filters, and image dimensions should be correct and square
        assert a.shape() == (5, 262, 262)

        # We should get duplicate object errors
        assert "Duplicate object ID" in caplog.text

        # We should get errors that include the duplicate filenames
        assert "_duplicate_" in caplog.text

        # The duplicate files should not be in the data set
        for filepath in a._all_files():
            assert "_duplicate_" not in str(filepath)


def test_prune_warn_1_percent(caplog, hsc_test_mode, tmp_path):
    """Test to ensure when >1% of loaded objects are missing a filter, that is a warning
    and that the resulting dataset drops the objects that are missing filters
    """
    caplog.set_level(logging.WARNING)

    # Generate two files which
    test_files = generate_files(num_objects=98, num_filters=2, shape=(100, 100))
    # Object 101 is missing the HSC-G filter, we only provide the R filter
    test_files["00000000000000101_missing_g_HSC-R.fits"] = (100, 100)

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(), data_location=tmp_path)

        # We should have the correct number of objects
        assert len(a) == 98

        # Object 2 should not be loaded
        assert "00000000000000101" not in a

        # We should Error log because greater than 5% of the objects were pruned
        assert "Greater than 1% of objects in the data directory were pruned." in caplog.text

        # We should warn that we dropped an object explicitly
        assert "Dropping object" in caplog.text


def test_prune_error_5_percent(caplog, hsc_test_mode, tmp_path):
    """Test to ensure when >5% of loaded objects are missing a filter, that is an error
    and that the resulting dataset drops the objects that are missing filters
    """
    caplog.set_level(logging.ERROR)

    # Generate two files which
    test_files = generate_files(num_objects=18, num_filters=2, shape=(100, 100))
    # Object 20 is missing the HSC-G filter, we only provide the R filter
    test_files["00000000000000020_missing_g_HSC-R.fits"] = (100, 100)

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(), data_location=tmp_path)

        # We should have two objects, having dropped one.
        assert len(a) == 18

        # Object 20 should not be loaded
        assert "00000000000000020" not in a

        # We should Error log because greater than 5% of the objects were pruned
        assert "Greater than 5% of objects in the data directory were pruned." in caplog.text


def test_crop(caplog, hsc_test_mode, tmp_path):
    """Test to ensure that in the presence of heterogenous sizes within 1px of a central size
    We load all images and crop to the smallest dimenensions without any logs
    """
    caplog.set_level(logging.WARNING)
    test_files = {}
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 100), offset=0))
    # Add some images with dimensions 1 px larger
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(101, 100), offset=10))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 101), offset=20))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(101, 101), offset=30))
    # Add some images with dimensions 1 px smaller
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(99, 100), offset=40))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 99), offset=50))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(99, 99), offset=60))

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(), data_location=tmp_path)

        assert len(a) == 70
        assert a.shape() == (5, 99, 99)

        # No warnings should be printed since we're within 1px of the mean size
        assert caplog.text == ""


def test_crop_warn_2px_larger(caplog, hsc_test_mode, tmp_path):
    """Test to ensure that in the presence of heterogenous sizes within 2px of a central size
    We load all images and crop to the smallest dimenensions and warn the user of the issue
    """
    caplog.set_level(logging.WARNING)
    test_files = {}
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 100), offset=0))
    # Add some images with dimensions 2 px larger
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(102, 100), offset=10))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 102), offset=20))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(102, 102), offset=30))
    # Add some images with dimensions 1 px smaller
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(99, 100), offset=40))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 99), offset=50))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(99, 99), offset=60))

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(), tmp_path)

        assert len(a) == 70
        assert a.shape() == (5, 99, 99)

        # We should warn that images differ
        assert "Some images differ" in caplog.text


def test_crop_warn_2px_smaller(caplog, hsc_test_mode, tmp_path):
    """Test to ensure that in the presence of heterogenous sizes within 2px of a central size
    We load all images and crop to the smallest dimenensions and warn the user of the issue
    """
    caplog.set_level(logging.WARNING)
    test_files = {}
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 100), offset=0))
    # Add some images with dimensions 1 px larger
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(101, 100), offset=10))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 101), offset=20))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(101, 101), offset=30))
    # Add some images with dimensions 2 px smaller
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(98, 100), offset=40))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 98), offset=50))
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(98, 98), offset=60))

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(), data_location=tmp_path)

        assert len(a) == 70
        assert a.shape() == (5, 98, 98)

        # We should warn that images differ
        assert "Some images differ" in caplog.text


def test_prune_size(caplog, hsc_test_mode, tmp_path):
    """Test to ensure images that are too small will be pruned from the data set when a custom size is
    passed."""
    caplog.set_level(logging.WARNING)
    test_files = {}
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 100), offset=0))
    # Add some images with dimensions 1 px larger
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(101, 101), offset=20))
    # Add some images with dimensions 2 px smaller
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(98, 98), offset=30))

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(crop_to=(99, 99)), data_location=tmp_path)

        assert len(a) == 20
        assert a.shape() == (5, 99, 99)

        # We should warn that we are dropping objects and the reason
        assert "Dropping object" in caplog.text
        assert "too small" in caplog.text


def test_prune_filter_size_mismatch(caplog, hsc_test_mode, tmp_path):
    """Test to ensure images with different sizes per filter will be dropped"""
    caplog.set_level(logging.WARNING)
    test_files = {}
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 100), offset=0))
    test_files["00000000000000000_all_filters_HSC-R.fits"] = (99, 99)

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(crop_to=(99, 99)), data_location=tmp_path)

        assert len(a) == 9
        assert a.shape() == (5, 99, 99)

        # We should warn that we are dropping objects and the reason
        assert "Dropping object" in caplog.text
        assert "first filter" in caplog.text


def test_prune_bad_filename(caplog, hsc_test_mode, tmp_path):
    """Test to ensure images with filenames set wrong will be dropped"""
    caplog.set_level(logging.WARNING)
    test_files = {}
    test_files.update(generate_files(num_objects=10, num_filters=5, shape=(100, 100), offset=0))

    config = mkconfig(crop_to=(99, 99), filter_catalog="notarealfile.fits")

    # Create a filter catalog with wrong file information.
    filter_catalog = generate_filter_catalog(test_files, config, tmp_path)
    filters = list(filter_catalog["00000000000000000"].keys())
    filter_catalog["00000000000000000"][filters[0]] = filter_catalog["00000000000000001"][filters[0]]

    with FakeFitsFS(test_files, filter_catalog):
        # Initialize HSCDataset exercising the filter_catalog provided initialization pathway
        a = HSCDataSet(config, data_location=tmp_path)

        # Verify that the broken object has been dropped
        assert len(a) == 9

        # Verify the shape is correct.
        assert a.shape() == (5, 99, 99)

        # We should warn that we are dropping objects and the correct reason
        assert "Dropping object" in caplog.text
        assert "manifest is likely corrupt" in caplog.text


def test_partial_filter(caplog, hsc_test_mode, tmp_path):
    """Test to ensure when we only load some of the filters, only those filters end up in the dataset"""
    caplog.set_level(logging.WARNING)
    test_files = generate_files(num_objects=10, num_filters=5, shape=(262, 263))
    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(filters=["HSC-G", "HSC-R"]), data_location=tmp_path)

        # 10 objects should load
        assert len(a) == 10

        # The number of filters, and image dimensions should be correct and square
        assert a.shape() == (2, 262, 262)

        # No warnings should be printed
        assert caplog.text == ""


def test_partial_filter_prune_warn_1_percent(caplog, hsc_test_mode, tmp_path):
    """Test to ensure when a the user supplies a filter list and >1% of loaded objects are
    missing a filter, that is a warning and that the resulting dataset drops the objects that
    are missing filters.
    """
    caplog.set_level(logging.WARNING)

    # Generate two files which
    test_files = generate_files(num_objects=98, num_filters=3, shape=(100, 100))
    # Object 101 is missing the HSC-G and HSC-I filters, we only provide the R filter
    test_files["00000000000000101_missing_g_HSC-R.fits"] = (100, 100)

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(filters=["HSC-R", "HSC-I"]), data_location=tmp_path)

        # We should have the correct number of objects
        assert len(a) == 98

        # Object 101 should not be loaded
        assert "00000000000000101" not in a

        # We should Error log because greater than 5% of the objects were pruned
        assert "Greater than 1% of objects in the data directory were pruned." in caplog.text

        # We should warn that we dropped an object explicitly
        assert "Dropping object" in caplog.text


def test_valid_transform_string(caplog, hsc_test_mode, tmp_path):
    """Test to ensure that a valid string passed to transform
    will map to a numpy function"""

    caplog.set_level(logging.ERROR)
    test_files = generate_files(num_objects=10, num_filters=5, shape=(262, 263))

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(transform="arcsinh"), data_location=tmp_path)

        # transform always has CenterCrop in the beginning followed by the user
        # defined transform
        lambda_transform = [t for t in a.transform.transforms if isinstance(t, Lambda)][0]
        assert lambda_transform.lambd == np.arcsinh

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(transform="tanh"), data_location=tmp_path)

        # transform always has CenterCrop in the beginning followed by the user
        # defined transform
        lambda_transform = [t for t in a.transform.transforms if isinstance(t, Lambda)][0]
        assert lambda_transform.lambd == np.tanh


def test_invalid_transform_string(caplog, hsc_test_mode, tmp_path):
    """Test to ensure that an invalid string passed to transform will raise an error"""

    caplog.set_level(logging.ERROR)
    test_files = generate_files(num_objects=10, num_filters=5, shape=(262, 263))

    with FakeFitsFS(test_files):
        with pytest.raises(RuntimeError):
            HSCDataSet(mkconfig(transform="invalid_function"), data_location=tmp_path)


def test_false_transform(caplog, hsc_test_mode, tmp_path):
    """Test to ensure that false passed to transform behaves as expected"""

    caplog.set_level(logging.ERROR)
    test_files = generate_files(num_objects=10, num_filters=5, shape=(262, 263))

    with FakeFitsFS(test_files):
        a = HSCDataSet(mkconfig(transform=False), data_location=tmp_path)

        # When transform is False; only a CenterCrop should be applied
        # automatically with a size conforming to test_files above
        expected_transform = CenterCrop(size=(np.int64(262), np.int64(262)))
        actual_transform = a.transform
        assert isinstance(actual_transform, CenterCrop)
        assert actual_transform.size == expected_transform.size
