"""
Acceptance tests: verify that recordMetadata stays in sync with the actual
data items throughout the full pyPhases pipeline – including datafold with
multiple datasets, shuffle, and deterministic re-loading.

Scenario
--------
- 3 datasets via ``datafold.datasets``, each contributing 5 recordings
  (recordIds "d0-0"…"d0-4", "d1-0"…"d1-4", "d2-0"…"d2-4").
- Each recording's X data embeds the *flat record index* so we can tell
  which record a data item came from.
- A custom DataManipulation subclass (injected via Swappable.setClass)
  adds a ``checkMetadata`` step that, for every item, records what
  ``self.recordMetadata[self.currentIndex]`` returns versus the recordId
  derivable from the data payload.
- We pull data through ``project.getData("dataset-training")`` — the
  same code path the real Training phase uses.
"""

from unittest.mock import MagicMock

import numpy as np
import pandas as pd

from pyPhases.test.Mocks import OverwriteConfig
from pyPhases.test.TestCase import TestCase
from pyPhasesRecordloader import RecordLoader

from SleePyPhases.DataManipulation import DataManipulation
from SleePyPhases.phases.BuildDataset import BuildDataset


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

# Accumulates (index, metadataRecordId, dataValue) triples that the custom
# manipulation step records while the pipeline iterates.
_metadata_checks: list = []


class CheckingDataManipulation(DataManipulation):
    """DataManipulation subclass with an extra ``checkMetadata`` step.

    For every item it reads ``self.recordMetadata[self.currentIndex]`` and
    compares the ``recordId`` to the flat-index value embedded in the X
    data (after ``addBatchDimension`` + ``fixedSize`` centering to pos 5).
    """

    def checkMetadata(self, X, Y):
        meta = self.recordMetadata[self.currentIndex]
        metadataRecordId = meta["recordId"]
        dataValue = int(X[0, 5, 0])   # flat index encoded in the signal
        _metadata_checks.append(
            (self.currentIndex, metadataRecordId, dataValue)
        )
        return X, Y


# Per-dataset record ids: 3 datasets × 5 records
ALL_IDS = [[f"d{d}-{r}" for r in range(7)] for d in range(3)]
FLAT_IDS = [rid for ds in ALL_IDS for rid in ds]          # 15 ids in flat order


# ---------------------------------------------------------------------------
# Test class
# ---------------------------------------------------------------------------

class TestBuildDatasetMetadataSync(TestCase):
    phase = BuildDataset()

    def config(self):
        self.project.setConfig("dataversion.split", {})
        return {
            "useLoader": "test",
            "dataversion": {
                "seed": 2023,
                "recordIds": FLAT_IDS,
                "split": {"training": ["0:15"]},
                "groupBy": None,
            },
            "segmentManipulation": [
                {"name": "addBatchDimension"},
                {"name": "fixedSize", "position": "center", "size": 10},
                {"name": "checkMetadata"},
            ],
            "segmentManipulationEval": [
                {"name": "addBatchDimension"},
                {"name": "fixedSize", "position": "center", "size": 10},
                {"name": "checkMetadata"},
            ],
            "evalOn": None,
            "datafold": {
                "shuffle": False,
                "seed": 42,
                "datasets": [
                    {"config": {"dataversion": {"recordIds": ALL_IDS[0], "seed": 1, "split": {"training": ["0:5"], "validation": ["5:7"]}, "groupBy": None}}},
                    {"config": {"dataversion": {"recordIds": ALL_IDS[1], "seed": 2, "split": {"training": ["0:5"], "validation": ["5:7"]}, "groupBy": None}}},
                    {"config": {"dataversion": {"recordIds": ALL_IDS[2], "seed": 3, "split": {"training": ["0:5"], "validation": ["5:7"]}, "groupBy": None}}},
                ],
            },
        }

    def setUp(self):
        super().setUp()
        DataManipulation.setClass(CheckingDataManipulation)

        metadata = pd.DataFrame([{"recordId": rid} for rid in FLAT_IDS])

        RecordLoader.getRecordList = MagicMock(return_value=FLAT_IDS)

        # Register data for each datafold dataset context — mirrors what
        # buildDataRecords() does with ``with self.project:`` + updateConfig.
        # Each dataset gets its own data arrays (like separate memmaps in prod).
        datasets = self.project.getConfig("datafold.datasets", [{"config": {}}])
        for dsIdx, ds in enumerate(datasets):
            overwriteConfig = ds["config"]
            with self.project:
                if "dataversion" in overwriteConfig:
                    self.project.setConfig("dataversion", {})
                self.project.updateConfig(overwriteConfig)

                dsRecordIds = self.project.getConfig("dataversion.recordIds")
                # Each dataset's memmap only contains its own records.
                # Values are the *global flat index* so we can trace which
                # record a data item originally belonged to.
                dsStart = dsIdx * 7
                dsProcessed = np.arange(dsStart, dsStart + len(dsRecordIds)).reshape(-1, 1, 1)
                dsFeatures = np.arange(dsStart, dsStart + len(dsRecordIds)).reshape(-1, 1, 1)

                self.project.registerData("allDBRecordIds", {rid: [rid] for rid in dsRecordIds})
                self.project.registerData("metadata", metadata)
                self.project.registerData("data-processed", dsProcessed)
                self.project.registerData("data-features", dsFeatures)

    def tearDown(self):
        DataManipulation.setClass(None)
        super().tearDown()

    # -- helpers ----------------------------------------------------------

    def _drain(self, dataset):
        """Iterate the full dataset, return list of (x, y)."""
        return [item for item in dataset]

    def _data_order(self, items):
        """Return the list of flat-index values read from the X payloads."""
        return [int(x[0, 5, 0]) for x, _ in items]

    def _assert_metadata_matches_data(self, label=""):
        """Assert every recorded (metadata, data) pair agrees."""
        mismatches = []
        for idx, metaId, dataVal in _metadata_checks:
            expected_flat = FLAT_IDS.index(metaId)
            if expected_flat != dataVal:
                mismatches.append(
                    f"  index {idx}: metadata={metaId} (flat {expected_flat}), data={dataVal}"
                )
        if mismatches:
            self.fail(
                f"{label}{len(mismatches)} metadata/data mismatches:\n"
                + "\n".join(mismatches)
            )

    # =====================================================================
    # 1) No shuffle — metadata must match data for every item
    # =====================================================================

    def testNoShuffle_MetadataMatchesData(self):
        """Without shuffle the metadata recordId at each index must equal
        the recordId derivable from the data payload."""
        _metadata_checks.clear()
        items = self._drain(self.getData("dataset-training"))
        self.assertEqual(len(items), 15)
        self._assert_metadata_matches_data("no-shuffle: ")

    def testNoShuffle_MetadataMatchesDataValidation(self):
        """Without shuffle the metadata recordId at each index must equal
        the recordId derivable from the data payload."""
        _metadata_checks.clear()
        items = self._drain(self.getData("dataset-validation"))
        self.assertEqual(len(items), 6)
        self._assert_metadata_matches_data("no-shuffle: ")

    # =====================================================================
    # 2) No shuffle — deterministic across two loads
    # =====================================================================

    def testNoShuffle_DeterministicOrder(self):
        """Two consecutive loads without shuffle must return data in
        identical order."""
        items1 = self._drain(self.getData("dataset-training"))
        order1 = self._data_order(items1)

        items2 = self._drain(self.getData("dataset-training"))
        order2 = self._data_order(items2)

        self.assertEqual(order1, order2)

    def testNoShuffle_DeterministicOrderValidation(self):
        """Two consecutive loads without shuffle must return validation data in
        identical order."""
        items1 = self._drain(self.getData("dataset-validation"))
        order1 = self._data_order(items1)

        items2 = self._drain(self.getData("dataset-validation"))
        order2 = self._data_order(items2)

        self.assertEqual(order1, order2)

    # =====================================================================
    # 3) With shuffle — load twice, verify deterministic order via seed
    # =====================================================================

    @OverwriteConfig({"datafold": {"shuffle": True, "seed": 42}})
    def testShuffle_DeterministicOrder(self):
        """Loading the dataset twice with the same seed must yield
        identical item order both times."""
        items1 = self._drain(self.getData("dataset-training"))
        order1 = self._data_order(items1)

        items2 = self._drain(self.getData("dataset-training"))
        order2 = self._data_order(items2)

        self.assertEqual(order1, order2,
                         "Same seed must produce the same item order")

    @OverwriteConfig({"datafold": {"shuffle": True, "seed": 42}})
    def testShuffle_DeterministicOrderValidation(self):
        """Loading the validation dataset twice with the same seed must yield
        identical item order both times."""
        items1 = self._drain(self.getData("dataset-validation"))
        order1 = self._data_order(items1)

        items2 = self._drain(self.getData("dataset-validation"))
        order2 = self._data_order(items2)

        self.assertEqual(order1, order2,
                         "Same seed must produce the same item order")

    # =====================================================================
    # 4) With shuffle — metadata must still match data for every item
    # =====================================================================

    @OverwriteConfig({"datafold": {"shuffle": True, "seed": 42}})
    def testShuffle_MetadataMatchesData(self):
        """After shuffle, recordMetadata[i].recordId must describe the
        record that data item i actually contains."""
        _metadata_checks.clear()
        items = self._drain(self.getData("dataset-training"))
        self.assertEqual(len(items), 15)
        self._assert_metadata_matches_data("shuffled: ")

    @OverwriteConfig({"datafold": {"shuffle": True, "seed": 42}})
    def testShuffle_MetadataMatchesDataValidation(self):
        """After shuffle, recordMetadata[i].recordId must describe the
        record that validation data item i actually contains."""
        _metadata_checks.clear()
        items = self._drain(self.getData("dataset-validation"))
        self.assertEqual(len(items), 6)
        self._assert_metadata_matches_data("shuffled-validation: ")
