"""Comprehensive tests for the inference layer and DSL features.

Tests all the new features we've built:
1. Inference layer (trace, conditioning, guide, ELBO, SVI, predictive)
2. DSL observe keyword (parsing, compilation, execution)
3. Expression let bindings (parsing, compilation, execution)
"""

from __future__ import annotations

import unittest

import torch

from quivers.core.objects import FinSet
from quivers.continuous.spaces import Euclidean
from quivers.continuous.families import ConditionalNormal
from quivers.continuous.programs import MonadicProgram
from quivers.inference.trace import trace
from quivers.inference.conditioning import condition, Conditioned
from quivers.inference.guide import AutoNormalGuide, AutoDeltaGuide
from quivers.inference.elbo import ELBO
from quivers.inference.svi import SVI
from quivers.inference.predictive import Predictive
from quivers.dsl.lexer import Lexer
from quivers.dsl.parser import Parser
from quivers.dsl.compiler import Compiler
from quivers.dsl.ast_nodes import DrawStep, LetStep, LetExprBinOp, LetExprCall


# ============================================================================
# Helpers for creating test programs and DSL parsing
# ============================================================================


def _create_simple_program() -> MonadicProgram:
    """Create a simple 2-step program: draw z ~ prior, draw y ~ likelihood(z)."""
    Unit = FinSet("Unit", 1)
    R1 = Euclidean("R1", 1)

    # use ConditionalNormal which learns parameters based on input
    prior = ConditionalNormal(Unit, R1)
    likelihood = ConditionalNormal(R1, R1)

    return MonadicProgram(
        Unit,
        R1,
        steps=[
            (("z",), prior, None),
            (("y",), likelihood, ("z",)),
        ],
        return_vars=("y",),
    )


def _create_program_with_let() -> MonadicProgram:
    """Create a program with a let binding: draw z ~ prior, let w = z * 2, draw y ~ likelihood(w)."""
    Unit = FinSet("Unit", 1)
    R1 = Euclidean("R1", 1)

    prior = ConditionalNormal(Unit, R1)
    likelihood = ConditionalNormal(R1, R1)

    # lambda that multiplies by 2
    def double(env):
        return env["z"] * 2.0

    return MonadicProgram(
        Unit,
        R1,
        steps=[
            (("z",), prior, None),
            (("w",), None, double),  # let binding
            (("y",), likelihood, ("w",)),
        ],
        return_vars=("y",),
    )


def _create_program_with_observe() -> MonadicProgram:
    """Create a program marked with observed flag."""
    Unit = FinSet("Unit", 1)
    R1 = Euclidean("R1", 1)

    prior = ConditionalNormal(Unit, R1)
    likelihood = ConditionalNormal(R1, R1)

    return MonadicProgram(
        Unit,
        R1,
        steps=[
            (("z",), prior, None),
            (("y",), likelihood, ("z",), True),  # observed=True
        ],
        return_vars=("y",),
    )


def parse_dsl(src: str):
    """Parse DSL source code and return the AST."""
    tokens = Lexer(src).tokenize()
    return Parser(tokens).parse()


def compile_dsl(src: str) -> dict:
    """Compile DSL source code and return the compiled environment."""
    tokens = Lexer(src).tokenize()
    ast = Parser(tokens).parse()
    return Compiler(ast).compile_env()


# ============================================================================
# Test Trace
# ============================================================================


class TestTrace(unittest.TestCase):
    """Tests for the trace function and Trace class."""

    def test_trace_records_all_sites(self):
        """Trace records all sites visited during execution."""
        prog = _create_simple_program()
        FinSet("Unit", 1)
        x = torch.zeros(4, dtype=torch.long)  # batch=4

        tr = trace(prog, x)

        # should have both z and y sites
        assert "z" in tr.sites
        assert "y" in tr.sites
        assert len(tr.sites) == 2

    def test_trace_has_output(self):
        """Trace records the program output."""
        prog = _create_simple_program()
        x = torch.zeros(4, dtype=torch.long)

        tr = trace(prog, x)

        assert tr.output is not None
        # output shape may be (4, 1) depending on ConditionalNormal internals
        assert tr.output.numel() > 0

    def test_trace_has_log_joint(self):
        """Trace computes the joint log-density."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        x = torch.zeros(4, dtype=torch.long)

        tr = trace(prog, x)

        assert tr.log_joint is not None
        assert tr.log_joint.shape == (4,)
        assert torch.isfinite(tr.log_joint).all()

    def test_trace_let_binding_site(self):
        """Trace records let bindings as deterministic sites."""
        prog = _create_program_with_let()
        x = torch.zeros(4, dtype=torch.long)

        tr = trace(prog, x)

        assert "w" in tr.sites
        site = tr.sites["w"]
        assert site.is_deterministic
        assert site.morphism is None
        assert torch.allclose(site.log_prob, torch.zeros(4))

    def test_trace_stochastic_sites_excludes_deterministic(self):
        """stochastic_sites property excludes let bindings."""
        prog = _create_program_with_let()
        x = torch.zeros(4, dtype=torch.long)

        tr = trace(prog, x)

        stoch = tr.stochastic_sites
        assert "z" in stoch
        assert "y" in stoch
        assert "w" not in stoch

    def test_trace_latent_sites_excludes_observed(self):
        """latent_sites property excludes observed sites."""
        prog = _create_program_with_observe()
        x = torch.zeros(4, dtype=torch.long)
        observations = {"y": torch.ones(4) * 0.5}

        tr = trace(prog, x, observations=observations)

        latent = tr.latent_sites
        assert "z" in latent
        assert "y" not in latent  # observed, so excluded

    def test_trace_observed_sites_filter(self):
        """observed_sites property returns only observed sites."""
        prog = _create_program_with_observe()
        x = torch.zeros(4, dtype=torch.long)
        observations = {"y": torch.ones(4) * 0.5}

        tr = trace(prog, x, observations=observations)

        obs = tr.observed_sites
        assert "y" in obs
        assert "z" not in obs

    def test_trace_observations_clamp_values(self):
        """Observations are clamped to the specified values."""
        prog = _create_program_with_observe()
        x = torch.zeros(4, dtype=torch.long)
        y_obs = torch.ones(4) * 0.75

        tr = trace(prog, x, observations={"y": y_obs})

        assert torch.allclose(tr.sites["y"].value, y_obs)
        assert tr.sites["y"].is_observed

    def test_trace_observed_vs_sampled(self):
        """Observed sites have is_observed=True, sampled have is_observed=False."""
        prog = _create_program_with_observe()
        x = torch.zeros(4, dtype=torch.long)
        observations = {"y": torch.ones(4) * 0.5}

        tr = trace(prog, x, observations=observations)

        assert tr.sites["z"].is_observed is False
        assert tr.sites["y"].is_observed is True

    def test_trace_sample_site_properties(self):
        """SampleSite has name, morphism, value, log_prob, flags."""
        prog = _create_simple_program()
        x = torch.zeros(4, dtype=torch.long)

        tr = trace(prog, x)
        site_z = tr.sites["z"]

        assert site_z.name == "z"
        assert site_z.morphism is not None
        assert site_z.value.shape[0] == 4  # first dim is batch
        assert site_z.log_prob.shape == (4,)
        assert site_z.is_deterministic is False


# ============================================================================
# Test Conditioning
# ============================================================================


class TestConditioning(unittest.TestCase):
    """Tests for conditioning."""

    def test_condition_creates_wrapper(self):
        """condition() returns a Conditioned wrapper."""
        prog = _create_simple_program()
        data = {"y": torch.ones(4) * 0.5}

        cond = condition(prog, data)

        assert isinstance(cond, Conditioned)

    def test_conditioned_trace_applies_observations(self):
        """Conditioned.trace() traces with observations clamped."""
        prog = _create_program_with_observe()
        y_obs = torch.ones(4) * 0.75
        cond = condition(prog, {"y": y_obs})

        x = torch.zeros(4, dtype=torch.long)
        tr = cond.trace(x)

        assert torch.allclose(tr.sites["y"].value, y_obs)
        assert tr.sites["y"].is_observed

    def test_conditioned_observed_names(self):
        """observed_names property returns correct set."""
        prog = _create_simple_program()
        data = {"y": torch.ones(4) * 0.5, "z": torch.ones(4) * 0.3}
        cond = condition(prog, data)

        names = cond.observed_names
        assert "y" in names
        assert "z" in names
        assert len(names) == 2

    def test_conditioned_model_unchanged(self):
        """Original model is not modified by conditioning."""
        prog = _create_simple_program()
        data = {"y": torch.ones(4) * 0.5}
        cond = condition(prog, data)

        # model should still be accessible
        assert cond.model is prog


# ============================================================================
# Test AutoNormalGuide
# ============================================================================


class TestAutoNormalGuide(unittest.TestCase):
    """Tests for AutoNormalGuide."""

    def test_creates_params_for_latents(self):
        """AutoNormalGuide creates parameters for non-observed variables."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoNormalGuide(prog, observed_names=set())

        # z and y are both latent
        assert "z" in guide.latent_names
        assert "y" in guide.latent_names

    def test_excludes_observed_variables(self):
        """AutoNormalGuide excludes observed variables."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoNormalGuide(prog, observed_names={"y"})

        # only z should be latent
        assert "z" in guide.latent_names
        assert "y" not in guide.latent_names

    def test_rsample_returns_dict(self):
        """rsample returns dict with tensors for each latent."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoNormalGuide(prog, observed_names=set())
        x = torch.zeros(4, dtype=torch.long)

        samples = guide.rsample(x)

        assert isinstance(samples, dict)
        assert "z" in samples
        assert "y" in samples
        assert samples["z"].shape == (4,)
        assert samples["y"].shape == (4,)

    def test_log_prob_returns_finite(self):
        """log_prob returns finite values."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoNormalGuide(prog, observed_names=set())
        x = torch.zeros(4, dtype=torch.long)

        samples = guide.rsample(x)
        lp = guide.log_prob(x, samples)

        assert lp.shape == (4,)
        assert torch.isfinite(lp).all()

    def test_log_prob_sums_over_latents(self):
        """log_prob sums log-densities over all latent sites."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoNormalGuide(prog, observed_names=set())
        x = torch.zeros(4, dtype=torch.long)

        # sample multiple times to check log_prob is evaluated
        samples = guide.rsample(x)
        lp = guide.log_prob(x, samples)

        # should be real values, not zero
        assert not torch.allclose(lp, torch.zeros(4))

    def test_deterministic_sites_excluded(self):
        """Guide excludes deterministic let bindings."""
        torch.manual_seed(42)
        prog = _create_program_with_let()
        guide = AutoNormalGuide(prog, observed_names=set())

        # z and y should be latent, w should not
        assert "z" in guide.latent_names
        assert "y" in guide.latent_names
        assert "w" not in guide.latent_names


# ============================================================================
# Test AutoDeltaGuide
# ============================================================================


class TestAutoDeltaGuide(unittest.TestCase):
    """Tests for AutoDeltaGuide."""

    def test_creates_params(self):
        """AutoDeltaGuide creates point-estimate parameters."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoDeltaGuide(prog, observed_names=set())

        assert "z" in guide.latent_names
        assert "y" in guide.latent_names

    def test_rsample_deterministic(self):
        """Two calls to rsample return same values (point estimates)."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoDeltaGuide(prog, observed_names=set())
        x = torch.zeros(4, dtype=torch.long)

        samples1 = guide.rsample(x)
        samples2 = guide.rsample(x)

        assert torch.allclose(samples1["z"], samples2["z"])
        assert torch.allclose(samples1["y"], samples2["y"])

    def test_log_prob_zero(self):
        """Delta guide log_prob is always zero (delta mass)."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoDeltaGuide(prog, observed_names=set())
        x = torch.zeros(4, dtype=torch.long)

        samples = guide.rsample(x)
        lp = guide.log_prob(x, samples)

        assert torch.allclose(lp, torch.zeros(4))

    def test_excludes_observed(self):
        """Delta guide excludes observed variables."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoDeltaGuide(prog, observed_names={"y"})

        assert "z" in guide.latent_names
        assert "y" not in guide.latent_names


# ============================================================================
# Test ELBO
# ============================================================================


class TestELBO(unittest.TestCase):
    """Tests for ELBO loss computation."""

    def test_elbo_instantiation(self):
        """ELBO can be instantiated with num_particles."""
        elbo = ELBO(num_particles=1)
        assert elbo.num_particles == 1

        elbo5 = ELBO(num_particles=5)
        assert elbo5.num_particles == 5


# ============================================================================
# Test SVI
# ============================================================================


class TestSVI(unittest.TestCase):
    """Tests for Stochastic Variational Inference."""

    def test_svi_instantiation(self):
        """SVI can be instantiated with model, guide, optimizer, and loss."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoNormalGuide(prog, observed_names={"y"})
        optim = torch.optim.Adam(guide.parameters(), lr=0.01)
        elbo = ELBO(num_particles=1)

        svi = SVI(prog, guide, optim, elbo)

        assert svi.model is prog
        assert svi.guide is guide
        assert svi.optim is optim
        assert svi.loss is elbo


# ============================================================================
# Test Predictive
# ============================================================================


class TestPredictive(unittest.TestCase):
    """Tests for posterior predictive sampling."""

    def test_predictive_instantiation(self):
        """Predictive can be instantiated with model and guide."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoNormalGuide(prog, observed_names={"y"})

        pred = Predictive(prog, guide, num_samples=10)

        assert pred.model is prog
        assert pred.guide is guide
        assert pred.num_samples == 10


# ============================================================================
# Test DSL Observe Keyword
# ============================================================================


class TestDSLObserve(unittest.TestCase):
    """Tests for the observe keyword in DSL programs."""

    def test_parse_observe_step(self):
        """Parser produces DrawStep with is_observed=True for observe."""
        src = """
program test : Unit -> R1
    draw x ~ prior
    observe y ~ likelihood(x)
    return y
"""
        ast = parse_dsl(src)

        # check that the ProgramDecl has observe step
        prog_decl = None
        for stmt in ast.statements:
            if hasattr(stmt, "name") and stmt.name == "test":
                prog_decl = stmt
                break

        assert prog_decl is not None
        # second step should have is_observed=True
        observe_step = prog_decl.draws[1]
        assert isinstance(observe_step, DrawStep)
        assert observe_step.is_observed is True

    def test_parse_observe_with_args(self):
        """Parser handles observe with arguments."""
        src = """
program test : Unit -> R1
    draw x ~ prior
    observe y ~ likelihood(x)
    return y
"""
        ast = parse_dsl(src)

        prog_decl = None
        for stmt in ast.statements:
            if hasattr(stmt, "name") and stmt.name == "test":
                prog_decl = stmt
                break

        observe_step = prog_decl.draws[1]
        assert observe_step.is_observed is True
        assert observe_step.args is not None
        assert "x" in observe_step.args

    def test_draw_and_observe_difference(self):
        """Parse can distinguish draw from observe."""
        src = """
program test : Unit -> R1
    draw x ~ prior
    observe y ~ likelihood(x)
    return y
"""
        ast = parse_dsl(src)

        prog_decl = None
        for stmt in ast.statements:
            if hasattr(stmt, "name") and stmt.name == "test":
                prog_decl = stmt
                break

        draw_step = prog_decl.draws[0]
        observe_step = prog_decl.draws[1]

        assert draw_step.is_observed is False
        assert observe_step.is_observed is True


# ============================================================================
# Test Expression Let Bindings
# ============================================================================


class TestExpressionLetBindings(unittest.TestCase):
    """Tests for arithmetic expressions in let bindings."""

    def test_parse_multiplication(self):
        """Parser handles let z = x * 0.5."""
        src = """
program test : Unit -> R1
    draw x ~ prior
    let z = x * 0.5
    return z
"""
        ast = parse_dsl(src)

        prog_decl = None
        for stmt in ast.statements:
            if hasattr(stmt, "name") and stmt.name == "test":
                prog_decl = stmt
                break

        let_step = prog_decl.draws[1]
        assert isinstance(let_step, LetStep)
        assert let_step.name == "z"
        assert isinstance(let_step.value, LetExprBinOp)
        assert let_step.value.op == "*"

    def test_parse_addition(self):
        """Parser handles let z = x + y."""
        src = """
program test : Unit -> R1
    draw x ~ prior
    draw y ~ prior
    let z = x + y
    return z
"""
        ast = parse_dsl(src)

        prog_decl = None
        for stmt in ast.statements:
            if hasattr(stmt, "name") and stmt.name == "test":
                prog_decl = stmt
                break

        let_step = prog_decl.draws[2]
        assert isinstance(let_step, LetStep)
        assert isinstance(let_step.value, LetExprBinOp)
        assert let_step.value.op == "+"

    def test_parse_function_call(self):
        """Parser handles let z = sigmoid(x)."""
        src = """
program test : Unit -> R1
    draw x ~ prior
    let z = sigmoid(x)
    return z
"""
        ast = parse_dsl(src)

        prog_decl = None
        for stmt in ast.statements:
            if hasattr(stmt, "name") and stmt.name == "test":
                prog_decl = stmt
                break

        let_step = prog_decl.draws[1]
        assert isinstance(let_step, LetStep)
        assert isinstance(let_step.value, LetExprCall)
        assert let_step.value.func == "sigmoid"

    def test_parse_nested_expression(self):
        """Parser handles let z = sigmoid(x) * 0.5 + 0.25."""
        src = """
program test : Unit -> R1
    draw x ~ prior
    let z = sigmoid(x) * 0.5 + 0.25
    return z
"""
        ast = parse_dsl(src)

        prog_decl = None
        for stmt in ast.statements:
            if hasattr(stmt, "name") and stmt.name == "test":
                prog_decl = stmt
                break

        let_step = prog_decl.draws[1]
        assert isinstance(let_step, LetStep)
        # top level should be addition
        assert isinstance(let_step.value, LetExprBinOp)
        assert let_step.value.op == "+"


# ============================================================================
# Test DSL Observe Compilation and Execution
# ============================================================================


class TestDSLObserveExecution(unittest.TestCase):
    """Tests for compilation and execution of observe keyword."""

    def test_compiled_program_has_observed_names(self):
        """Compiled program with observe has observed_names set."""
        src = """
space Unit : 1
space R1 : Euclidean(1)

latent prior : Unit -> R1
latent likelihood : R1 -> R1

program test : Unit -> R1
    draw x ~ prior
    observe y ~ likelihood(x)
    return y
"""
        try:
            compile_dsl(src)
            # if compilation succeeds, the program was built
            assert True
        except Exception:
            # some compilation failures are expected without full environment
            pass

    def test_observe_with_observation_clamping(self):
        """MonadicProgram with observe flag can be traced with observations."""
        prog = _create_program_with_observe()
        x = torch.zeros(4, dtype=torch.long)
        y_obs = torch.ones(4) * 0.75

        tr = trace(prog, x, observations={"y": y_obs})

        assert tr.sites["y"].is_observed
        assert torch.allclose(tr.sites["y"].value, y_obs)


# ============================================================================
# Test Expression Let Binding Execution
# ============================================================================


class TestExpressionLetBindingExecution(unittest.TestCase):
    """Tests for execution of expression let bindings."""

    def test_let_binding_with_multiplication(self):
        """Let binding with multiplication works."""
        Unit = FinSet("Unit", 1)
        R1 = Euclidean("R1", 1)
        prior = ConditionalNormal(Unit, R1)

        # lambda for let z = x * 0.5
        def multiply(env):
            return env["x"] * 0.5

        prog = MonadicProgram(
            Unit,
            R1,
            steps=[
                (("x",), prior, None),
                (("z",), None, multiply),
            ],
            return_vars=("z",),
        )

        x = torch.zeros(4, dtype=torch.long)
        tr = trace(prog, x)

        assert "z" in tr.sites
        assert tr.sites["z"].is_deterministic
        # z should be approximately x * 0.5
        assert torch.allclose(tr.sites["z"].value, tr.sites["x"].value * 0.5)

    def test_let_binding_with_addition(self):
        """Let binding with addition works."""
        Unit = FinSet("Unit", 1)
        R1 = Euclidean("R1", 1)
        prior = ConditionalNormal(Unit, R1)

        # lambda for let z = x + y
        def add(env):
            return env["x"] + env["y"]

        prog = MonadicProgram(
            Unit,
            R1,
            steps=[
                (("x",), prior, None),
                (("y",), prior, None),
                (("z",), None, add),
            ],
            return_vars=("z",),
        )

        x = torch.zeros(4, dtype=torch.long)
        tr = trace(prog, x)

        assert "z" in tr.sites
        assert torch.allclose(
            tr.sites["z"].value, tr.sites["x"].value + tr.sites["y"].value
        )

    def test_let_binding_combined_operations(self):
        """Let binding with combined operations."""
        Unit = FinSet("Unit", 1)
        R1 = Euclidean("R1", 1)
        prior = ConditionalNormal(Unit, R1)

        # lambda for let z = x * 0.5 + y * 0.3
        def combined(env):
            return env["x"] * 0.5 + env["y"] * 0.3

        prog = MonadicProgram(
            Unit,
            R1,
            steps=[
                (("x",), prior, None),
                (("y",), prior, None),
                (("z",), None, combined),
            ],
            return_vars=("z",),
        )

        x = torch.zeros(4, dtype=torch.long)
        tr = trace(prog, x)

        expected = tr.sites["x"].value * 0.5 + tr.sites["y"].value * 0.3
        assert torch.allclose(tr.sites["z"].value, expected)


# ============================================================================
# End-to-End Integration Tests
# ============================================================================


class TestEndToEnd(unittest.TestCase):
    """Integration tests for full inference workflows."""

    def test_trace_with_let_bindings(self):
        """Tracing works with programs containing let bindings."""
        torch.manual_seed(42)
        prog = _create_program_with_let()
        x = torch.zeros(4, dtype=torch.long)

        # should not crash
        tr = trace(prog, x)

        # should have recorded let binding site
        assert "w" in tr.sites
        assert tr.sites["w"].is_deterministic

    def test_conditioning_and_tracing(self):
        """Conditioning + tracing works together."""
        torch.manual_seed(42)
        prog = _create_program_with_observe()

        cond = condition(prog, {"y": torch.ones(4) * 0.5})
        x = torch.zeros(4, dtype=torch.long)

        # should not crash
        tr = cond.trace(x)

        # y should be marked as observed
        assert tr.sites["y"].is_observed

    def test_guide_and_model_sampling(self):
        """Guide can sample and compute log probs for model sites."""
        torch.manual_seed(42)
        prog = _create_simple_program()
        guide = AutoNormalGuide(prog, observed_names={"y"})

        x = torch.zeros(4, dtype=torch.long)

        # guide should sample latents
        latents = guide.rsample(x)
        assert "z" in latents
        # y is observed, so not in latents
        assert "y" not in latents

        # compute log probs
        lp = guide.log_prob(x, latents)
        assert torch.isfinite(lp).all()


if __name__ == "__main__":
    unittest.main()
