#!/usr/bin/env python3
"""AgentTailor GPU reproduction script."""
import sys, os, json, torch, numpy as np, random

os.system('pip install "setuptools<70" astunparse wikipedia aiohttp class-registry --quiet 2>&1 | tail -1')
os.system('git clone https://github.com/Pt3Y/AgentTailor.git /tmp/AgentTailor 2>&1 | tail -3')
import class_registry.entry_points
with open(class_registry.entry_points.__file__, 'w') as f:
    f.write('')
sys.path.insert(0, '/tmp/AgentTailor')
os.environ['HF_ENDPOINT'] = 'https://huggingface.co'
os.environ['HF_HUB_ENDPOINT'] = 'https://huggingface.co'
os.chdir('/tmp/AgentTailor')

SEED = 888
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Device: {device}')

from AgentTailor.ATNetwork.Actor import Actor
from AgentTailor.ATNetwork.Critics import Critics, EPN, Encoder
from AgentTailor.ATNetwork.ExpBuffer import ExperienceBuffer
from AgentTailor.agents.agent_registry import AgentRegistry

results = {}
gpu_info = {'available': torch.cuda.is_available()}
if torch.cuda.is_available():
    gpu_info['name'] = torch.cuda.get_device_name(0)
    gpu_info['count'] = torch.cuda.device_count()

# Test 1: EPN
epn = EPN(dims=[1920, 1], dropout=0.1, temperature=5.0).to(device)
loss_fn = torch.nn.MSELoss()
opt = torch.optim.Adam(epn.parameters(), lr=1e-3)
losses = []
for step in range(50):
    pred = epn(torch.randn(16, 1920).to(device))
    target = torch.rand(16, 1).to(device)
    loss = loss_fn(pred, target)
    opt.zero_grad(); loss.backward(); opt.step()
    losses.append(float(loss))
results['epn_training'] = {'final_loss': losses[-1], 'loss_dropped': losses[-1] < losses[0]}
print(f'1/6 EPN: final_loss={losses[-1]:.6f}')

# Test 2: Actor
actor = Actor(domain='gsm8k', llm_name='gpt-4o',
              agent_names=['MathSolver','AnalyzeAgent','AnalyzeAgent','AnalyzeAgent','AnalyzeAgent'],
              decision_method='FinalRefer', optimized_spatial=True, optimized_temporal=True)
actor.construct_spatial_connection(temperature=1.0)
results['actor_graph'] = {'nodes': actor.num_nodes, 'potential_edges': len(actor.potential_spatial_edges)}
print(f'2/6 Actor: {actor.num_nodes} nodes, {len(actor.potential_spatial_edges)} edges')

# Test 3: Critics
critics = Critics(epn_dims=[1920, 1], model_name='all-MiniLM-L6-v2',
                  lock_threshold=0.01, temperature=5.0, dropout=0.0)
val = critics.run_differentiated(
    'MathSolver', 'solved x+2=5 -> x=3', 'What is 2+2?',
    'AnalyzeAgent', 'verified answer is 3')
results['critics_diff'] = float(val.item())
print(f'3/6 Critics: diff_value={val.item():.4f}')

# Test 4: Self-locking
critics.epn.train()
opt_c = torch.optim.Adam(critics.epn.parameters(), lr=1e-2)
for _ in range(100):
    pred = critics.epn(torch.randn(8, 1920).to(critics.device))
    tgt = torch.sigmoid(torch.randn(8, 1).to(critics.device) * 0.3 + 0.5)
    l = ((pred - tgt) ** 2).mean()
    opt_c.zero_grad(); l.backward(); opt_c.step()
critics.lock_critic()
results['self_locking'] = {'is_locked': critics.is_locked, 'confidence': critics.lock_confidence}
print(f'4/6 Self-lock: locked={critics.is_locked}')

# Test 5: Edge pruning
new_s, new_t = actor.apply_pruning(k_spatial=5, k_temporal=3)
results['edge_pruning'] = {'spatial_active': int(new_s.sum().item()), 'temporal_active': int(new_t.sum().item())}
print(f'5/6 Edge pruning: spatial={int(new_s.sum().item())}, temporal={int(new_t.sum().item())}')

# Test 6: Semantic edge evaluation
good_val = critics.run_differentiated(
    'MathSolver: math expert', 'solved equation correctly',
    'Tom has 8 marbles, loses 3, how many left?',
    'AnalyzeAgent: verifier', 'confirmed answer is 5')
bad_val = critics.run_differentiated(
    'AdversarialAgent: misleading', 'generated weather facts',
    'Tom has 8 marbles, loses 3, how many left?',
    'CodeWriting: unrelated code', 'wrote sorting algorithm')
results['semantic_separation'] = {'good': float(good_val.item()), 'bad': float(bad_val.item()),
                                   'diff': float(good_val.item() - bad_val.item())}
print(f'6/6 Semantic: good={good_val.item():.4f}, bad={bad_val.item():.4f}, diff={good_val.item()-bad_val.item():.4f}')

output = {'gpu': gpu_info, 'device': device, 'results': results, 'all_passed': all([
    results['epn_training']['loss_dropped'],
    results['self_locking']['is_locked'],
    results['semantic_separation']['diff'] > 0
])}
print(json.dumps(output, indent=2))