"""Run PRIMA on ordered image groups and text prompts."""
import argparse
import json
from pathlib import Path

from tools.config import parse_with_config, validate_paths


def parse_args():
    from tools.model_loader import DEFAULT_HUB_MODEL, add_model_arguments

    parser = argparse.ArgumentParser(description=__doc__)
    add_model_arguments(parser, merged_checkpoint=DEFAULT_HUB_MODEL)
    parser.add_argument('--image_paths', nargs='+', help='One image group for CLI inference.')
    parser.add_argument('--prompts', nargs='+', help='Prompts paired with image groups in the config.')
    parser.add_argument('--output_save_path', default='outputs/inference')
    parser.add_argument('--max_new_tokens', type=int, default=1024)
    return parse_with_config(parser)


def main():
    args = parse_args()
    validate_paths(args)
    if not args.image_paths or not args.prompts:
        raise ValueError('Provide image_paths and prompts in the config or CLI.')
    groups = args.image_paths
    if all(isinstance(item, str) for item in groups):
        groups = [groups]
    if len(groups) != len(args.prompts):
        raise ValueError('Each image group must have exactly one prompt.')
    if any(not isinstance(group, list) or not 1 <= len(group) <= args.max_num_images for group in groups):
        raise ValueError('Each image group must be a nonempty list within max_num_images.')
    for group in groups:
        for path in group:
            if not Path(path).is_file():
                raise ValueError(f'Image does not exist: {path}')
    import numpy as np
    import torch
    from model.llava import conversation
    from tools.chat import Chat
    from tools.model_loader import load_bundle
    from tools.outputs import image_mask_slots

    model, tokenizer, qformer_tokenizer = load_bundle(args)
    chat = Chat(model, tokenizer, qformer_tokenizer, args, device=f'cuda:{args.local_rank}')
    output = Path(args.output_save_path)
    output.mkdir(parents=True, exist_ok=True)
    for index, (images, question) in enumerate(zip(groups, args.prompts)):
        turn = output / f'turn{index}'
        turn.mkdir(exist_ok=True)
        conv = conversation.conv_templates[args.conv_type].copy()
        encoded = []
        for image in images:
            chat.upload_img(image, conv, encoded)
        chat.ask(question, conv, num_images=len(images))
        with torch.inference_mode():
            answer, masks, _ = chat.answer(conv, encoded, question,
                                           max_new_tokens=args.max_new_tokens,
                                           max_length=args.model_max_length + args.max_new_tokens)
        mask_paths = []
        for image_index, predicted in enumerate(image_mask_slots(masks, len(images))):
            paths = []
            if predicted is not None:
                for mask_index, mask in enumerate(predicted):
                    name = f'image{image_index + 1}_mask{mask_index + 1}.npy'
                    np.save(turn / name, (mask > 0).cpu().numpy())
                    paths.append(name)
            mask_paths.append(paths)
        record = {'images': images, 'prompt': question, 'response': answer, 'masks': mask_paths}
        (turn / 'dialogue.json').write_text(json.dumps(record, indent=2) + '\n')
        print(answer)


if __name__ == '__main__':
    try:
        main()
    except ValueError as error:
        raise SystemExit(str(error))