"""
Test the hypothesis-lab environment running locally on Docker (port 8000).
Uses the standard OpenEnv interface: env.reset() and env.step(action).
"""
import asyncio
import sys
sys.path.insert(0, "/Users/sbhimraj/Documents/Openenv/lab-experiment")

from client import HypothesisLabEnv
from models import HypLabAction, ActionType, ExperimentType


async def main():
    print("=== Testing local hypothesis-lab environment ===\n")

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

        # Reset episode
        result = await env.reset(noise_level="low")
        obs = result.observation
        variables = obs.available_variables
        print(f"[reset]")
        print(f"  variables:  {variables}")
        print(f"  budget:     {obs.budget_remaining}")
        print(f"  message:    {obs.system_message[:120]}\n")

        v1, v2 = variables[0], variables[1]

        # Step 1 — intervention experiment
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.INTERVENTION,
            control_variable=v1,
            control_value=2.5,
            target_variable=v2,
        )
        result = await env.step(action)
        obs = result.observation
        print(f"[step 1: intervention]")
        print(f"  {v1}=2.5 → {v2} = {obs.result_value}")
        print(f"  info_gain:   {obs.info_gain_reward}")
        print(f"  budget left: {obs.budget_remaining}")
        print(f"  done:        {result.done}\n")

        # Step 2 — correlation experiment
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.CORRELATION,
            control_variable=v1,
            control_range=[0.0, 1.0, 2.0, 3.0, 4.0],
            target_variable=v2,
        )
        result = await env.step(action)
        obs = result.observation
        print(f"[step 2: correlation]")
        print(f"  {v1} swept → {v2} = {obs.result_value}")
        print(f"  info_gain:   {obs.info_gain_reward}")
        print(f"  budget left: {obs.budget_remaining}")
        print(f"  done:        {result.done}\n")

        # Step 3 — submit hypothesis
        action = HypLabAction(
            action_type=ActionType.SUBMIT,
            hypothesis_text=f"{v2} increases linearly with {v1} by approximately 2.5 units",
            confidence=0.7,
        )
        result = await env.step(action)
        obs = result.observation
        print(f"[step 3: submit]")
        print(f"  reward:              {result.reward}")
        print(f"  accuracy_score:      {obs.accuracy_score}")
        print(f"  precision_bonus:     {obs.precision_bonus}")
        print(f"  calibration_score:   {obs.calibration_score}")
        print(f"  efficiency_bonus:    {obs.efficiency_bonus}")
        print(f"  done:                {result.done}")
        if obs.ground_truth_revealed:
            print(f"  ground_truth:        {obs.ground_truth_revealed[:300]}")

    print("\n=== Done — environment is working correctly ===")


asyncio.run(main())