#!/usr/bin/env python3
from argparse import ArgumentParser, Namespace
from pathlib import Path
from natsort import natsorted
import numpy as np
import pandas as pd
import pims
import sys
from vre import VRE
from vre.utils import FakeVideo, image_read
from vre.representations import build_representations_from_cfg
from vre.logger import vre_logger as logger
from omegaconf import OmegaConf

import imp

def get_args() -> Namespace:
    """cli args"""
    parser = ArgumentParser()
    parser.add_argument("video_path", help="Path to the scene video we are processing")
    parser.add_argument("--output_path", "-o", required=True,
                        help="Path to the output directory where representations are stored")
    parser.add_argument("--cfg_path", required=True, help="Path to global YAML cfg file")
    parser.add_argument("--representations", nargs="+", help="Subset of representations to run VRE for")
    # run parameters
    parser.add_argument("--start_frame", type=int, default=0, help="The first frame. If not set, defaults to first.")
    parser.add_argument("--end_frame", type=int, help="End frame. If not set, defaults to length of video")
    parser.add_argument("--output_dir_exists_mode", choices=["overwrite", "skip_computed", "raise"])
    parser.add_argument("--n_threads_data_storer", type=int, default=0)
    # other parameters
    parser.add_argument("--frame_rate", type=int, default=-1, help="For a dir of images. Required for optical flow")
    parser.add_argument("--make_video_collage", action="store_true")
    parser.add_argument("--collage_fps", type=int)

    args = parser.parse_args()
    args.video_path = Path(args.video_path).absolute()
    args.output_path = Path(args.output_path).absolute()
    args.cfg_path = Path(args.cfg_path).absolute()
    if args.make_video_collage:
        assert args.collage_fps is not None and isinstance(args.collage_fps, int) and args.collage_fps > 0
    return args

def get_cfg():
    """cfg and cli args"""
    args = get_args()
    cfg = OmegaConf.load(open(args.cfg_path, "r"))
    if args.representations is not None:
        assert (d := set(args.representations).difference(r := cfg.representations.keys())) == set(), f"{d} not in {r}"
        logger.info(f"Keeping only {args.representations} from the config")
        cfg.representations = {k: v for k, v in cfg.representations.items() if k in args.representations}
    return args, cfg

def main():
    args, cfg = get_cfg()
    if args.video_path.is_dir():
        raw_data = [image_read(f) for f in natsorted(args.video_path.glob("*.png"), key=lambda p: p.name)]
        assert all(x.shape == raw_data[0].shape for x in raw_data), f"Images shape differ in '{args.video_path}'"
        logger.info(f"--video_path is a directory. Assuming a directory of images. Found {len(raw_data)} images.")
        video = FakeVideo(np.array(raw_data, dtype=np.uint8), frame_rate=args.frame_rate)
    else:
        video = pims.Video(args.video_path)
    representations = build_representations_from_cfg(cfg.representations, cfg.get("default_compute_parameters"),
                                                     cfg.get("default_learned_parameters"))
    vre = VRE(video, representations)
    vre_metadata = vre.run(args.output_path, start_frame=args.start_frame, end_frame=args.end_frame,
                           output_dir_exists_mode=args.output_dir_exists_mode,
                           n_threads_data_storer=args.n_threads_data_storer)
    vre_run_stats = pd.DataFrame(vre_metadata["run_stats"], index=range(*vre_metadata["runtime_args"]["frames"]))
    print(vre_run_stats if len(vre_run_stats) > 1 else vre_run_stats.T)

    if args.make_video_collage:
        vre_collage = imp.load_source("vre_collage", (Path(__file__).parent / "vre_collage").__str__())
        vre_collage.main(Namespace(vre_export_dir=args.output_path, output_path=args.output_path/"collage",
                                   video=True, fps=args.collage_fps, output_resolution=None))

if __name__ == "__main__":
    main()
