import sys, os
sys.path.append('/home/thebears/Source/task_runners/vision_v3/')
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
from cuda_objdet_clip.infer import *
enc_file_path = '/home/thebears/Source/ml/clip_test/kya_tail/Leopards1_00_20260226160534.mp4'
# %%

crop_embeds_path = file_names.get_cropped_embed_npz_path(enc_file_path)
enc_file_path = file_names.resolve_file_location(enc_file_path)

nvc_batch_size = 32
decoder.reconfigure_decoder(enc_file_path)

all_sc = list()

keep_reading_video = True

clip_src_tensor_list = []
det_src_tensor_list = []

c_frm_num = 0

st = time.time()
n_det_scores = 0
n_clip_scores = 0

transform_stats = None
thresh_frames = list()
thresh_scores = list()
thresh_labels = list()
thresh_boxes = list()
det_all_scored_frames = list()
clip_embeddings = list()
clip_frames = list()
while keep_reading_video:
    frames = decoder.get_batch_frames(nvc_batch_size)
    if len(frames) == 0:
        keep_reading_video = False

    for frame in frames:
        if (c_frm_num % clip_frame_skip_interval) == 0:
            clip_src_tensor_list.append(
                {"tensor": torch.from_dlpack(frame), "frame_number": c_frm_num}
            )
        c_frm_num += 1

    while len(clip_src_tensor_list) >= clip_model.batch_size:
        print('bump')
        clip_tens_pass = clip_src_tensor_list[0 : clip_model.batch_size]
        clip_src_tensor_list = clip_src_tensor_list[clip_model.batch_size :]     
        clip_frame_numbers = [x["frame_number"] for x in clip_tens_pass]
        clip_frames.append(clip_frame_numbers)
        clip_stacked = (
            torch.stack([x["tensor"] for x in clip_tens_pass], dim=0) / 255.0
        )

        movie_size = list(clip_stacked.shape[::-1][0:2])
        crop_div = 4
        olap_factor = 0.5
        
        crop_width = int(movie_size[0]/crop_div)
        crop_height = int(crop_width)
        
        spacing_width = int(crop_width * (1-olap_factor))
        spacing_height = int(crop_height * (1-olap_factor))

        starts_width = list(range(0,movie_size[0], spacing_width))
        starts_height = list(range(0, movie_size[1], spacing_height))
        crop_areas = set()
        for st_w in starts_width:
            for st_h in starts_height:
                en_w = st_w + crop_width
                en_h = st_h + crop_height

                if en_w > movie_size[0]:
                    st_w = movie_size[0] - crop_width
                    en_w = st_w + crop_width
                if en_h > movie_size[1]:
                    st_h = movie_size[1] - crop_height
                    en_h = st_h + crop_height            

                crop_area = ( st_w, st_h, en_w, en_h)
                crop_areas.add(crop_area)

        cropped_embeddings = dict()
        for crop_area in crop_areas:
            ( st_w, st_h, en_w, en_h) = crop_area
            clip_stack_score = clip_stacked[:,:, st_h:en_h, st_w:en_w]
            clip_res = clip_model.score(clip_transforms(clip_stack_score).data_ptr())
            cropped_embeddings[crop_area] = np.copy(clip_res["projected"])
            n_clip_scores += clip_res["projected"].shape[0]
        clip_embeddings.append(cropped_embeddings)


clip_dict = dict()
clip_dict["src_path"] = enc_file_path

cropped_embeddings = dict()
for c_clip in clip_embeddings:
    for k,v in c_clip.items():
        if k not in cropped_embeddings:
            cropped_embeddings[k] = list()
            
        cropped_embeddings[k].append(v)

clip_dict["frame_numbers"]  =np.concatenate(clip_frames)
clip_dict["embeds_cropped"] =  {k:np.concatenate(v) for k,v in cropped_embeddings.items()}
np.savez(crop_embeds_path, **clip_dict)