from pathlib import Path
import subprocess,json,shutil
import cv2
import numpy as np
from PIL import Image,ImageDraw
import argparse
parser=argparse.ArgumentParser(description='WorldGrid 接缝与镜头处理')
parser.add_argument('--input-dir',type=Path,required=True,help='包含未叠加后期效果的25秒与15秒视频的目录')
parser.add_argument('--output-dir',type=Path,required=True,help='输出目录')
args=parser.parse_args()
d=args.input_dir.resolve();w=args.output_dir.resolve();w.mkdir(parents=True,exist_ok=True)
source=d/'WorldGrid_T1H1_T1H2_展示成片_25s.mp4'
standard=d/'WorldGrid_T1H1_T1H2_标准序列_15s.mp4'
for f in (source,standard):
 if not f.is_file():parser.error(f'找不到输入视频：{f}')
enc=['-an','-c:v','libx264','-preset','medium','-crf','17','-pix_fmt','yuv420p','-r','24','-video_track_timescale','12288']
def run(a):subprocess.run(['ffmpeg','-y','-v','error']+a,check=True)
def read(f,n):
 cap=cv2.VideoCapture(str(f));cap.set(cv2.CAP_PROP_POS_FRAMES,n);ok,a=cap.read();cap.release();assert ok;return a

def camera(frame,index):
 t=index/24
 angle=0.0021*np.sin(2*np.pi*t/5)
 dx=9.0*np.sin(2*np.pi*t/5)
 dy=5.4*np.sin(2*np.pi*t/5+0.7)
 h,ww=frame.shape[:2]
 mat=cv2.getRotationMatrix2D(((ww-1)/2,(h-1)/2),angle*180/np.pi,2048/1920)
 mat[0,2]-=dx;mat[1,2]-=dy
 return cv2.warpAffine(frame,mat,(ww,h),flags=cv2.INTER_CUBIC,borderMode=cv2.BORDER_REFLECT_101)

def encoder(out):
 return subprocess.Popen(['ffmpeg','-y','-v','error','-f','rawvideo','-pix_fmt','bgr24','-s','1920x1080','-r','24','-i','-']+enc+[str(out)],stdin=subprocess.PIPE)

def body_clip(src,start,end,out):
 cap=cv2.VideoCapture(str(src));cap.set(cv2.CAP_PROP_POS_FRAMES,start)
 proc=encoder(out)
 for n in range(start,end):
  ok,frame=cap.read();assert ok
  proc.stdin.write(camera(frame,n).tobytes())
 cap.release();proc.stdin.close();assert proc.wait()==0

def patch(f,b,out,half):
 # Interpolate 2*half frames between the surrounding source frames.
 a=read(f,b-half-1);z=read(f,b+half);h,ww=a.shape[:2]
 gray=lambda x:cv2.cvtColor(cv2.resize(x,(ww//2,h//2)),cv2.COLOR_BGR2GRAY)
 dis=cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_MEDIUM)
 ab=cv2.resize(dis.calc(gray(a),gray(z),None),(ww,h))*2
 ba=cv2.resize(dis.calc(gray(z),gray(a),None),(ww,h))*2
 x,y=np.meshgrid(np.arange(ww,dtype=np.float32),np.arange(h,dtype=np.float32))
 proc=encoder(out)
 for i in range(1,2*half+1):
  t=i/(2*half+1)
  aa=cv2.remap(a,x-t*ab[:,:,0],y-t*ab[:,:,1],cv2.INTER_LINEAR,borderMode=cv2.BORDER_REFLECT_101)
  bb=cv2.remap(z,x-(1-t)*ba[:,:,0],y-(1-t)*ba[:,:,1],cv2.INTER_LINEAR,borderMode=cv2.BORDER_REFLECT_101)
  mix=t*t*(3-2*t)
  frame=cv2.addWeighted(aa,1-mix,bb,mix,0)
  proc.stdin.write(camera(frame,b-half+i-1).tobytes())
 proc.stdin.close();assert proc.wait()==0

def build(src,bounds,total,filename):
 clips=[];start=0
 for i,b in enumerate(bounds):
  half=6 if b==(240 if total==600 else 120) else 10
  body=w/f'{total}_body_{i}.mp4'
  body_clip(src,start,b-half,body);clips.append(body)
  fix=w/f'{total}_seam_{b}.mp4';patch(src,b,fix,half);clips.append(fix);start=b+half
 body=w/f'{total}_body_end.mp4';body_clip(src,start,total,body);clips.append(body)
 manifest=w/f'{total}_concat.txt';manifest.write_text(''.join("file '"+str(x)+"'\n" for x in clips))
 target=w/filename;run(['-f','concat','-safe','0','-i',str(manifest),'-c','copy','-movflags','+faststart',str(target)])
 info=json.loads(subprocess.check_output(['ffprobe','-v','error','-show_streams','-of','json',str(target)]));assert int(info['streams'][0]['nb_frames'])==total
 run(['-i',str(target),'-f','null','-'])
 stats=[]
 for b in bounds:
  canvas=Image.new('RGB',(1440,620),'#121b25');draw=ImageDraw.Draw(canvas)
  # Same instants before and after, including both sides of the cut.
  for row,f in enumerate([src,target]):
   for col,n in enumerate([b-3,b-1,b,b+1,b+3,b+5]):
    im=Image.fromarray(cv2.cvtColor(read(f,n),cv2.COLOR_BGR2RGB)).resize((240,135))
    canvas.paste(im,(col*240,row*310+30));draw.text((col*240+6,row*310+8),f'{"BEFORE" if row==0 else "AFTER"} {n/24:.3f}s',fill='white')
    # Train crop for inspection.
    crop=read(f,n)[350:800,100:1800];ci=Image.fromarray(cv2.cvtColor(crop,cv2.COLOR_BGR2RGB));ci.thumbnail((240,125));canvas.paste(ci,(col*240,row*310+180))
  canvas.save(w/f'qa_{total}_{b}.jpg')
  mae=lambda f:float(np.mean(np.abs(read(f,b).astype(float)-read(f,b-1).astype(float))))
  stats.append({'time':b/24,'old_cut_mean_abs_delta':mae(src),'new_cut_mean_abs_delta':mae(target)})
 print(filename,stats,flush=True)
 (w/f'qa_{total}.json').write_text(json.dumps(stats,indent=2))
 return target
if __name__=='__main__':
 build(source,[120,240,360,480],600,'WorldGrid_T1H1_T1H2_展示成片_25s_亚像素增强二版.mp4')
 build(standard,[120,240],360,'WorldGrid_T1H1_T1H2_标准序列_15s_亚像素增强二版.mp4')