import torch
import numpy as np
from dreamer import AgenticForecaster, ReplayBuffer

def run_test():
    device = torch.device("cpu")
    obs_dim = 10
    action_dim = 5
    batch_size = 8
    seq_len = 16
    
    agent = AgenticForecaster(obs_dim=obs_dim, action_dim=action_dim, device=device)
    buffer = ReplayBuffer(capacity=100, seq_len=seq_len)
    
    # 1. Collect some dummy data
    for _ in range(20):
        # random episode of length 30
        obs = np.random.randn(30, obs_dim)
        acts = np.abs(np.random.randn(30, action_dim))
        acts = acts / acts.sum(axis=-1, keepdims=True) # Normalize as portfolio weights
        rews = np.random.randn(30)
        buffer.add_episode(obs, acts, rews)
        
    # 2. Sample batch
    obs_batch, act_batch, rew_batch = buffer.sample_batch(batch_size)
    obs_batch = obs_batch.to(device)
    act_batch = act_batch.to(device)
    rew_batch = rew_batch.to(device)
    
    # 3. Train World Model
    wm_logs = agent.train_world_model(obs_batch, act_batch, rew_batch)
    print(f"World Model Training Logs: {wm_logs}")
    
    # 4. Train Actor Critic
    # Create arbitrary start state from posterior for imagination
    start_state = agent.rssm.initial_state(batch_size, device)
    ac_logs = agent.train_actor_critic(start_state, horizon=10)
    print(f"Actor-Critic Training Logs: {ac_logs}")
    
    print("Multi-module Dreamer integration passed successfully.")

if __name__ == "__main__":
    run_test()