import torch
from tqdm import tqdm
from pyiqa.archs.musiq_arch import MUSIQ
from vbench.utils import load_video, load_dimension_info

def transform(images):
    return images / 255.

def technical_quality(model, video_list, device):
    video_results = []
    for video_path in tqdm(video_list):
        images = load_video(video_path)
        images = transform(images)
        acc_score_video = 0.
        for i in range(len(images)):
            frame = images[i].unsqueeze(0).to(device)
            score = model(frame)
            acc_score_video += float(score)
        video_results.append({'video_path': video_path, 'video_results': acc_score_video/len(images)})
    average_score = sum([o['video_results'] for o in video_results]) / len(video_results)
    average_score = average_score / 100.
    return average_score, video_results


def compute_imaging_quality(json_dir, device, submodules_list):
    model_path = submodules_list['model_path']

    model = MUSIQ(pretrained_model_path=model_path)
    model.to(device)
    model.training = False
    
    video_list, _ = load_dimension_info(json_dir, dimension='imaging_quality', lang='en')
    all_results, video_results = technical_quality(model, video_list, device)
    return all_results, video_results