"""
test_docker.py -- Test the Hypothesis Lab environment using the OpenEnv client.

Run with: python test_docker.py
Requires: the Docker container running on localhost:8000
"""

import asyncio

from client import HypothesisLabEnv


async def main():
    async with HypothesisLabEnv(base_url="http://localhost:8000") as env:

        # 1. Reset
        result = await env.reset(noise_level="low", domain="system_alpha", seed=42)
        obs = result.observation
        print("=== Episode Started ===")
        print(f"Message:   {obs.system_message[:80]}...")
        print(f"Variables: {obs.available_variables}")
        print(f"Budget:    {obs.budget_remaining}")
        print()

        v = obs.available_variables
        cause, effect = v[0], v[1]

        # 2. Intervention -- set one variable, observe the other
        result = await env.run_intervention(cause, 5.0, effect)
        obs = result.observation
        print(f"[Intervention] Set {cause}=5.0 -> {effect}={obs.result_value:.4f}")
        print(f"  Info gain: {obs.info_gain_reward}  Budget left: {obs.budget_remaining}")
        print()

        # 3. Try the reverse direction to check causality
        result = await env.run_intervention(effect, 5.0, cause)
        obs = result.observation
        print(f"[Reverse]      Set {effect}=5.0 -> {cause}={obs.result_value:.4f}")
        print(f"  Info gain: {obs.info_gain_reward}  Budget left: {obs.budget_remaining}")
        print()

        # 4. Correlation sweep -- see the shape of the relationship
        result = await env.run_correlation(cause, [0.5, 20.0, 8], effect)
        obs = result.observation
        print(f"[Correlation]  Swept {cause} from 0.5 to 20.0:")
        for x, y in obs.result_value:
            print(f"  {cause}={x:.1f} -> {effect}={y:.4f}")
        print(f"  Info gain: {obs.info_gain_reward}  Budget left: {obs.budget_remaining}")
        print()

        # 5. Counterfactual -- what if cause changes by +3?
        result = await env.run_counterfactual(cause, 3.0, effect)
        obs = result.observation
        cf = obs.result_value
        print(f"[Counterfactual] If {cause} changes by +3.0:")
        print(f"  Baseline: {cause}={cf['baseline_x']:.2f} -> {effect}={cf['baseline_y_noisy']:.4f}")
        print(f"  After:    {cause}={cf['counterfactual_x']:.2f} -> {effect}={cf['counterfactual_y_noisy']:.4f}")
        print(f"  Direction: {cf['direction']}")
        print(f"  Info gain: {obs.info_gain_reward}  Budget left: {obs.budget_remaining}")
        print()

        # 6. Passive observation -- observe a variable without touching anything
        result = await env.run_passive(effect)
        obs = result.observation
        print(f"[Passive]  {effect} at rest = {obs.result_value:.4f}")
        print(f"  Info gain: {obs.info_gain_reward}  Budget left: {obs.budget_remaining}")
        print()

        # 7. Submit hypothesis
        result = await env.submit_hypothesis(
            hypothesis_text=f"{effect} decays exponentially as {cause} increases.",
            hypothesis_equations=[f"{effect} = 1.1 * exp(-0.16 * {cause})"],
            confidence=0.65,
        )
        obs = result.observation
        print("=== Episode Finished ===")
        print(f"Accuracy:      {obs.accuracy_score}")
        print(f"Precision:     {obs.precision_bonus}")
        print(f"Calibration:   {obs.calibration_score}")
        print(f"Efficiency:    {obs.efficiency_bonus}")
        print(f"Contradiction: {obs.contradiction_penalty}")
        print(f"TOTAL REWARD:  {obs.total_episode_reward}")
        print()
        print(f"Ground truth:\n{obs.ground_truth_revealed}")


asyncio.run(main())