"""Own from-scratch course ML-Agents training, export, independent ONNX evaluation.
Run inside the compatible Linux environment prepared by colab-setup.py.
"""
from pathlib import Path
from datetime import datetime,timezone
import os,sys,json,subprocess,zipfile,urllib.request,hashlib,shutil,time
import numpy as np
import yaml
ROOT=Path('/content/vlad-deep-rl-unity');ROOT.mkdir(exist_ok=True)
os.environ['OMP_NUM_THREADS']='1';os.environ['MKL_NUM_THREADS']='1'
SEED=20261101

def download(url,target):
 print('Downloading official course environment',target.name,flush=True)
 urllib.request.urlretrieve(url,target)
def unzip(path,target):
 target.mkdir(exist_ok=True)
 with zipfile.ZipFile(path) as z:
  for item in z.infolist():
   if not (target/item.filename).resolve().is_relative_to(target.resolve()):raise ValueError('Unsafe zip member')
  z.extractall(target)
def evaluate_onnx(model,executable,output):
 from mlagents_envs.environment import UnityEnvironment
 from mlagents_envs.base_env import ActionTuple
 from mlagents_envs.side_channel.engine_configuration_channel import EngineConfigurationChannel
 import onnxruntime as ort
 opts=ort.SessionOptions();opts.intra_op_num_threads=1;opts.inter_op_num_threads=1
 session=ort.InferenceSession(str(model),sess_options=opts,providers=['CPUExecutionProvider'])
 print('ONNX inputs',[(i.name,i.shape) for i in session.get_inputs()],flush=True)
 channel=EngineConfigurationChannel();channel.set_configuration_parameters(time_scale=20)
 env=UnityEnvironment(file_name=str(executable),seed=1000000,no_graphics=True,side_channels=[channel],worker_id=10,timeout_wait=120)
 totals={};records=[]
 try:
  env.reset()
  for iteration in range(30000):
   for behavior in env.behavior_specs:
    decision,terminal=env.get_steps(behavior)
    for agent_id in terminal.agent_id:
     step=terminal[agent_id];key=(behavior,int(agent_id));total=totals.pop(key,0.0)+float(step.reward)+float(step.group_reward)
     records.append({'behavior':behavior,'agent_id':int(agent_id),'reward':total,'interrupted':bool(step.interrupted)})
    for agent_id in decision.agent_id:
     step=decision[agent_id];key=(behavior,int(agent_id));totals[key]=totals.get(key,0.0)+float(step.reward)+float(step.group_reward)
    if len(decision):
     inputs={}
     for spec in session.get_inputs():
      if spec.name.startswith('obs_'):inputs[spec.name]=decision.obs[int(spec.name.split('_')[-1])].astype(np.float32)
      elif spec.name=='action_masks':
       inputs[spec.name]=np.concatenate([1-m.astype(np.float32) for m in decision.action_mask],axis=1)
      else:raise ValueError('Unexpected ONNX input '+spec.name)
     values=session.run(None,inputs);values={spec.name:value for spec,value in zip(session.get_outputs(),values)}
     action=ActionTuple()
     if 'continuous_actions' in values:action.add_continuous(values['continuous_actions'].astype(np.float32))
     if 'discrete_actions' in values:action.add_discrete(values['discrete_actions'].astype(np.int32))
     env.set_actions(behavior,action)
   if len(records)>=40:break
   env.step()
 finally:env.close()
 if not records:raise RuntimeError('Independent evaluation completed no episodes')
 rewards=[r['reward'] for r in records];mean=float(np.mean(rewards));std=float(np.std(rewards))
 result={'mean_reward':mean,'std_reward':std,'mean_minus_std':mean-std,'n_eval_agent_episodes':len(rewards),'episode_rewards':rewards,'records':records,'evaluation_seed':1000000,'eval_datetime':datetime.now(timezone.utc).isoformat(),'evaluation_kind':'separate Unity process, ONNX policy, completed agent episodes (soccer episodes within a match are correlated)'}
 output.write_text(json.dumps(result,indent=2));print('EVALUATED',model.name,mean,std,flush=True);return result

def card(folder,name,env_tag,result,config):
 meta={'license':'apache-2.0','tags':['reinforcement-learning','ml-agents',env_tag,'deep-rl-class'],'model-index':[{'name':name,'results':[{'task':{'type':'reinforcement-learning','name':'Reinforcement Learning'},'dataset':{'name':env_tag,'type':env_tag},'metrics':[{'type':'mean_reward','name':'mean_reward','value':f"{result['mean_reward']:.4f} +/- {result['std_reward']:.4f}"}]}]}]}
 text='---\n'+yaml.safe_dump(meta,sort_keys=False)+'---\n\n# '+name+'\n\n'
 text+='Trained from random initialization for Vlad Kolomiiets (`Vladprime`) using the official Hugging Face Deep RL Course Unity environment and Unity ML-Agents 1.1.0. No pretrained weights were used. This is an initial training run; it is not claimed to have converged. Units 5 and 7 require publishing a trained model without a reward threshold.\n\n'
 text+=f"Training: {config['trainer_type']}, {config['max_steps']} configured environment steps, seed {SEED}. Export: genuine ML-Agents ONNX plus PyTorch checkpoint and trainer configuration. Separate ONNX evaluation process: seed 1000000; {result['n_eval_agent_episodes']} completed agent episodes, mean {result['mean_reward']:.4f}, standard deviation {result['std_reward']:.4f}. Original environment, action and observation spaces were retained. See results.json for individual returns and evaluation limitations.\n\n"
 if 'Soccer' in env_tag:text+='Offline MA-POCA self-play training and local evaluation only. No ranked AI-vs-AI challenge result is claimed; the course online leaderboard is unavailable. Player episodes within a soccer match are correlated.\n\n'
 text+='The .onnx file is compatible with the course Unity browser viewer. Training/evaluation code is in train-unity.py; the exact run configuration is config.yaml.\n'
 (folder/'README.md').write_text(text)
 (folder/'SHA256SUMS.json').write_text(json.dumps({str(p.relative_to(folder)):hashlib.sha256(p.read_bytes()).hexdigest() for p in folder.rglob('*') if p.is_file() and p.name!='SHA256SUMS.json'},indent=2))

def main():
 sources=[('SnowballTarget','https://github.com/huggingface/Snowball-Target/raw/main/SnowballTarget.zip','SnowballTarget/SnowballTarget.x86_64','SnowballTarget.yaml','ML-Agents-SnowballTarget'),('Pyramids','https://huggingface.co/spaces/unity/ML-Agents-Pyramids/resolve/main/Pyramids.zip','Pyramids/Pyramids','PyramidsRND-local.yaml','ML-Agents-Pyramids'),('SoccerTwos','gdrive:1KuqBKYiXiIcU4kNMqEzhgypuFP5_45CL','SoccerTwos.x86_64','SoccerTwos-local.yaml','ML-Agents-SoccerTwos')]
 results=[]
 for name,url,relative,config_name,env_tag in sources:
  zip_path=ROOT/(name+'.zip')
  if not zip_path.exists():
   if url.startswith('gdrive:'):
    import gdown;gdown.download(id=url.split(':')[1],output=str(zip_path),quiet=False)
   else:download(url,zip_path)
  target=ROOT/('env-'+name);unzip(zip_path,target);executable=target/relative;executable.chmod(executable.stat().st_mode|0o111)
  config_path=ROOT/config_name;config=yaml.safe_load(config_path.read_text());behavior=next(iter(config['behaviors'].values()));run_id='vlad-'+name
  command=[str(Path(sys.executable).parent/'mlagents-learn'),str(config_path),'--env='+str(executable),'--run-id='+run_id,'--no-graphics','--seed='+str(SEED),'--torch-device=cpu','--results-dir='+str(ROOT/'training')]
  print('TRAINING',name,' '.join(command),flush=True)
  subprocess.run(command,check=True,env={**os.environ,'OMP_NUM_THREADS':'1','MKL_NUM_THREADS':'1'})
  trained=ROOT/'training'/run_id;model=list(trained.rglob('*.onnx'))
  # Final exported model is at the run root; step checkpoints live in behavior subdirectories.
  model=next((p for p in model if p.parent==trained),model[-1])
  output=ROOT/('publish-'+name);output.mkdir(exist_ok=True)
  shutil.copy2(model,output/(name+'.onnx'))
  for p in trained.rglob('*.pt'):
   if p.name=='checkpoint.pt':shutil.copy2(p,output/(name+'-checkpoint.pt'));break
  shutil.copy2(config_path,output/'config.yaml');shutil.copy2(__file__,output/'train-unity.py')
  logs=list(trained.rglob('events.out.tfevents*'))
  if logs:shutil.copy2(logs[0],output/logs[0].name)
  result=evaluate_onnx(output/(name+'.onnx'),executable,output/'results.json');card(output,name,env_tag,result,behavior)
  shutil.make_archive(str(ROOT/('publish-'+name)),'zip',output)
  results.append({'name':name,**{k:result[k] for k in ['mean_reward','std_reward','mean_minus_std','n_eval_agent_episodes']}})
 (ROOT/'summary.json').write_text(json.dumps(results,indent=2));print('ALL_UNITY_TRAINING_COMPLETE',json.dumps(results),flush=True)
if __name__=='__main__':main()