"""
Verify that regen_road actually produces different tracks.

Method: connect to sim, regen road with 3 seeds, drive straight ahead for
100 steps on each, record final CTE. Different roads curve differently,
so a straight-ahead policy will accumulate CTE in different directions/amounts.

Also prints the first few node positions from the handler if accessible.
"""
import sys
import time
import numpy as np

sys.path.insert(0, '/home/paulh/projects/donkeycar-rl-autoresearch/agent')

import gymnasium as gym
import gym_donkeycar  # noqa: F401 — registers donkey envs

HOST     = 'localhost'
PORT     = 9091
TRACK_ID = 'donkey-generated-roads-v0'
SEEDS    = [1111, 55555, 99999]
STEPS    = 500      # drive lane-following steps per seed
THROTTLE = 0.3
STEER_GAIN = 0.8    # proportional: steer = -cte * gain
WAIT     = 3.5      # seconds after regen before reset


def get_handler(env):
    return env.unwrapped.viewer.handler


def regen_road(env, seed):
    get_handler(env).queue_message({
        'msg_type': 'regen_road',
        'road_style': '0',
        'rand_seed': str(seed),
        'turn_increment': '0.0',
    })
    time.sleep(WAIT)


print('Connecting to sim...')
env = gym.make(TRACK_ID, conf={'host': HOST, 'port': PORT})
print(f'  Connected. obs={env.observation_space.shape}')

results = {}

for seed in SEEDS:
    print(f'\n── Seed {seed} ──────────────────────')
    print(f'  Regenerating road...')
    regen_road(env, seed)

    obs, info = env.reset()
    cte_values = []
    pos_values = []

    for step in range(STEPS):
        # Lane-following: steer proportional to CTE so the car stays on road.
        # Different road geometries will produce different CTE histories.
        last_cte = cte_values[-1] if cte_values else 0.0
        steer = float(np.clip(-last_cte * STEER_GAIN, -1.0, 1.0))
        action = np.array([steer, THROTTLE], dtype=np.float32)

        obs, reward, terminated, truncated, info = env.step(action)
        cte = info.get('cte', 0.0)
        pos = info.get('pos', None)
        cte_values.append(float(cte))
        if pos is not None:
            pos_values.append(list(pos)[:3])
        if terminated or truncated:
            print(f'  Episode ended at step {step+1}')
            break

    final_cte   = cte_values[-1] if cte_values else 0.0
    mean_cte    = float(np.mean(cte_values)) if cte_values else 0.0
    max_abs_cte = float(np.max(np.abs(cte_values))) if cte_values else 0.0
    final_pos   = pos_values[-1] if pos_values else None

    results[seed] = {
        'final_cte':   final_cte,
        'mean_cte':    mean_cte,
        'max_abs_cte': max_abs_cte,
        'steps':       len(cte_values),
        'final_pos':   final_pos,
    }

    print(f'  Steps driven : {len(cte_values)}')
    print(f'  Final CTE    : {final_cte:+.3f}m  (+ = right of centre, - = left)')
    print(f'  Mean CTE     : {mean_cte:+.3f}m')
    print(f'  Max |CTE|    : {max_abs_cte:.3f}m')
    if final_pos:
        print(f'  Final pos    : x={final_pos[0]:.2f}  y={final_pos[1]:.2f}  z={final_pos[2]:.2f}')

env.close()

print('\n' + '='*50)
print('SUMMARY — same straight-ahead policy, different seeds:')
print('='*50)
for seed, r in results.items():
    p = r['final_pos']
    pos_str = f'x={p[0]:.1f} z={p[2]:.1f}' if p else 'N/A'
    print(f'  Seed {seed:6d}:  CTE={r["final_cte"]:+.3f}m  steps={r["steps"]}  pos={pos_str}')

ctes = [r['final_cte'] for r in results.values()]
spread = max(ctes) - min(ctes)
print(f'\nCTE spread across seeds: {spread:.3f}m')
if spread > 0.3:
    print('✅  ROADS ARE DIFFERENT — CTE spread > 0.3m confirms different road geometries')
else:
    print('❌  ROADS MAY BE THE SAME — CTE spread is small, road gen may not be working')