"""
End-to-end test of run_multistream.py with the_well.data.WellDataset
mocked at the point env.py imports it. This exercises the FULL real path
(streaming -> adapter -> per-domain normalizer -> channel-filtered
replay -> training) together.
"""
from __future__ import annotations
import os
import sys
from unittest.mock import patch, MagicMock

import pytest
import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.model import MultiScaleEncoder, HierarchicalHyperbolicPredictor
from src.continual import ReplayBuffer, DiagonalEWC, fit_normalizer_for_domain
from src.config import BEST_HPARAMS as BEST
from src.provenance import DataLoadError


def make_fake_well_dataset(n_samples, Ti, To, H, W, C):
    class FakeWellDataset:
        def __init__(self, **kwargs):
            self._n = n_samples
        def __len__(self):
            return self._n
        def __getitem__(self, idx):
            return {
                "input_fields": torch.randn(Ti, H, W, C),
                "output_fields": torch.randn(To, H, W, C),
            }
    return FakeWellDataset


def test_load_real_domain_hard_fails_when_stream_unavailable():
    from src.run_multistream import load_real_domain
    with patch("the_well.data.WellDataset", side_effect=ConnectionError("no network")):
        with pytest.raises(DataLoadError) as exc_info:
            load_real_domain("nonexistent_dataset", "train", max_samples=8)
    assert exc_info.value.outcome_code == "STREAM_FAILED"


def test_load_real_domain_succeeds_with_mocked_stream():
    from src.run_multistream import load_real_domain
    fake_cls = make_fake_well_dataset(n_samples=16, Ti=4, To=4, H=16, W=16, C=2)
    with patch("the_well.data.WellDataset", fake_cls):
        ds = load_real_domain("fake_dataset", "train", max_samples=16)
    assert ds.provenance == "REAL_STREAMED"
    item = ds[0]
    assert item["fields"].shape == (8, 2, 16, 16)  # Ti+To, C, H, W


def test_full_multistream_pipeline_two_domains_different_channels():
    from src.run_multistream import load_real_domain, train_one_domain

    fake_c2 = make_fake_well_dataset(n_samples=16, Ti=4, To=4, H=16, W=16, C=2)
    fake_c11 = make_fake_well_dataset(n_samples=16, Ti=4, To=4, H=16, W=16, C=11)

    enc = MultiScaleEncoder(hidden=32, out_dim=8)
    model = HierarchicalHyperbolicPredictor(enc, c=1.0, pred_steps=2, levels=2)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    replay = ReplayBuffer(capacity=64)
    ewc = DiagonalEWC(model, lambda_ewc=100.0)

    with patch("the_well.data.WellDataset", fake_c2):
        ds1 = load_real_domain("gray_scott_reaction_diffusion", "train", max_samples=16)
    assert ds1.provenance == "REAL_STREAMED"
    norm1 = fit_normalizer_for_domain(ds1, max_fit=8)
    before = model.encoder.stem.weight.detach().clone()
    train_one_domain(model, norm1, ds1, opt, "cpu", epochs=1,
                      replay=replay, ewc=None, teacher=None, mix_replay=0.0)
    after_domain1 = model.encoder.stem.weight.detach().clone()
    assert not torch.allclose(before, after_domain1)
    assert replay.counts_by_channels().get(2, 0) > 0

    with patch("the_well.data.WellDataset", fake_c11):
        ds2 = load_real_domain("active_matter", "train", max_samples=16)
    assert ds2.provenance == "REAL_STREAMED"
    norm2 = fit_normalizer_for_domain(ds2, max_fit=8)
    train_one_domain(model, norm2, ds2, opt, "cpu", epochs=1,
                      replay=replay, ewc=ewc, teacher=None, mix_replay=0.5)
    after_domain2 = model.encoder.stem.weight.detach().clone()
    assert not torch.allclose(after_domain1, after_domain2)
    assert replay.counts_by_channels().get(11, 0) > 0
    assert set(replay.counts_by_channels().keys()) == {2, 11}

    assert norm1.mean.shape == (2,)
    assert norm2.mean.shape == (11,)


if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v"]))