#!/usr/bin/env python3
import scriptconfig as scfg
import ubelt as ub
from torch.utils import data


if 1:
    # FIXME: put the vendored package into our namespace.
    try:
        PARENT_DPATH = ub.Path(__file__).parent
    except NameError:
        from watch.tasks.sam import predict as this_mod
        PARENT_DPATH = ub.Path(this_mod.__file__)
    TPL_DPATH = PARENT_DPATH / 'tpl'

    import sys
    import os
    sys.path.append(os.fspath(TPL_DPATH / 'segment_anything'))
    import segment_anything


class SAMConfig(scfg.DataConfig):
    input_kwcoco = scfg.Value(None, help='input kwcoco dataset')
    output_kwcoco = scfg.Value(None, help='output')

    weights_fpath = scfg.Value(None, help='path to pytorch model weights')

    window_dims = (1024, 1024)
    fixed_resolution = "10GSD"
    batch_size = 1
    window_overlap = 0.33333
    device = scfg.Value(0)
    select_images = None
    track_emissions = True
    data_workers = 2
    io_workers = 0
    assets_dname = 'sam'


def rgb_from_kwcoco_frame(frame):
    import torchvision.transforms.functional as F
    import kwimage
    import torch
    # import numpy as np
    modes = frame['modes']

    if 'red|green|blue' in modes:
        chw = modes['red|green|blue']
        is_nan = torch.isnan(chw)
        rgb_nan_frac = is_nan.sum() / is_nan.numel()
    else:
        rgb_nan_frac = 1.0

    # print('----')
    # print(f'rgb_nan_frac={rgb_nan_frac}')
    if rgb_nan_frac >= 0.0:
        # fallback on pan
        if 'pan' in modes:
            pan_chw = modes['pan']
            pan_is_nan = torch.isnan(pan_chw)
            pan_nan_frac = pan_is_nan.sum() / pan_is_nan.numel()
            # print(f'pan_nan_frac={pan_nan_frac}')
            if rgb_nan_frac > pan_nan_frac:
                pan_hwc = pan_chw.permute(1, 2, 0).cpu().numpy()
                pan_hwc = kwimage.atleast_3channels(pan_hwc)
                chw = torch.Tensor(pan_hwc).permute(2, 0, 1)

    hwc = chw.permute(1, 2, 0).cpu().numpy()
    normed_rgb = kwimage.normalize_intensity(hwc)
    chw = F.to_tensor(normed_rgb)
    return chw


class SAMWrapperDataset(data.Dataset):

    def __init__(self, subdset):
        self.subdset = subdset

    def __len__(self):
        return len(self.subdset)

    def __getitem__(self, index):
        import kwimage
        import torch
        import numpy as np

        torch_dataset = self.subdset

        item = torch_dataset[index]

        frame0 = item['frames'][0]
        chw = rgb_from_kwcoco_frame(frame0)
        hwc = chw.permute(1, 2, 0).cpu().numpy()
        norm01_resized, resize_info = kwimage.imresize(
            hwc, dsize=(1024, 1024), interpolation='lanczos',
            letterbox=True, border_value=np.nan, return_info=True)
        norm01_resized = norm01_resized.clip(0, 1)
        hwc_isnodata_mask = np.isnan(norm01_resized)

        # Get a nodata mask for the output we expect, Only mask an output block
        # if more than half of its input pixels were nodata.
        import einops
        factor = 16
        isnodat_mask = hwc_isnodata_mask.all(axis=2).astype(np.uint8)
        isnodata_blocks = einops.rearrange(isnodat_mask, '(h2 s1) (w2 s2) -> h2 w2 (s1 s2)', s1=factor, s2=factor)
        output_isnodata_mask = isnodata_blocks.sum(axis=2) > (factor * factor) // 2

        norm255_resized = (np.nan_to_num(norm01_resized) * 255).round().astype(np.uint8)
        input_image_torch = torch.as_tensor(norm255_resized)

        chw255 = input_image_torch.permute(2, 0, 1).contiguous()

        # coco_dset = torch_dataset.sampler.dset
        gid = frame0['gid']
        # coco_img = coco_dset.coco_image(gid)

        # SAM will always downsample the output by a factor of 16 after it has
        # been resized to 1024x1024. Munge the transforms to account for this.
        fakeoutspace_from_outspace = kwimage.Affine.scale(16)
        outspace_from_fakeoutspace = fakeoutspace_from_outspace.inv()

        # The dataloader is assuming the output will be the same shape as the
        # input but that's not true. We will call this assumed output space the
        # fake outspace. The real outspace is actually downsampled by a factor
        # of 16.

        fake_output_dsize = kwimage.Box.from_dsize(frame0['output_image_dsize'])
        fake_outspace_slice = kwimage.Box.from_slice(frame0['output_space_slice'])
        fakeoutspace_from_vid = kwimage.Affine.coerce(scale=frame0['scale_outspace_from_vid'])

        output_dsize = fake_output_dsize.warp(outspace_from_fakeoutspace).quantize()
        outspace_slice_ = fake_outspace_slice.warp(outspace_from_fakeoutspace)
        asset_slice = outspace_slice_.quantize().clip(0, 0, None, None).to_slice()
        asset_dsize = (output_dsize.width, output_dsize.height)

        vid_from_fakeoutspace = fakeoutspace_from_vid.inv()
        vid_from_outspace = vid_from_fakeoutspace @ fakeoutspace_from_outspace
        outspace_from_vid = vid_from_outspace.inv()

        resize_offset = resize_info['offset']
        resize_scale = resize_info['scale']
        undo_resize = kwimage.Affine.coerce(scale=resize_scale,
                                            offset=resize_offset).inv()
        if not undo_resize.isclose_identity():
            kwimage.warp_affine
            raise NotImplementedError('todo')

        item['chw255'] = chw255
        item['chw_isnodata_mask'] = torch.as_tensor(hwc_isnodata_mask.transpose(2, 1, 0)).contiguous()
        item['output_isnodata_mask'] = np.ascontiguousarray(output_isnodata_mask)
        item['gid'] = gid
        item['asset_slice'] = asset_slice
        item['asset_dsize'] = asset_dsize
        item['scale_asset_from_vid'] = outspace_from_vid.decompose()['scale']
        return item


class DenseFeaturePredictor:
    """
    Base class for computing per-image dense features
    """
    short_code = NotImplemented
    chan_code = NotImplemented
    proc_name = NotImplemented
    WrapperDsetClass = NotImplemented

    def __init__(self, config):
        self.config = config

    def setup(self):
        from watch.tasks.fusion.datamodules import kwcoco_datamodule
        from watch.tasks.fusion.coco_stitcher import CocoStitchingManager
        from watch.utils import util_parallel
        from watch.utils import process_context

        config = self.config
        datamodule = kwcoco_datamodule.KWCocoVideoDataModule(
            test_dataset=config.input_kwcoco,
            batch_size=config.batch_size,
            fixed_resolution=config.fixed_resolution,
            num_workers=config.data_workers,
            window_dims=config.window_dims,
            select_images=config.select_images,
            window_overlap=config.window_overlap,
            force_bad_frames=True,
            resample_invalid_frames=0,
            time_steps=1,
            channels="red|green|blue,pan",
        )
        datamodule.setup('test')
        self.datamodule = datamodule

        torch_dataset = self.datamodule.torch_datasets['test']
        coco_dset = torch_dataset.sampler.dset
        coco_dset.clear_annotations()
        self.out_dset = coco_dset.copy()
        self.out_dset.fpath = ub.Path(config.output_kwcoco)
        # TODO: graceful bundle changes
        self.out_dset.reroot(absolute=True)

        self.proc_context = process_context.ProcessContext(
            name=self.proc_name,
            config=dict(config),
            track_emissions=config.track_emissions,
        )

        # Not sure why targets are not in image-first order, but let's ensure
        # they are so we can fix it.
        torch_dataset.new_sample_grid['targets'] = sorted(
            torch_dataset.new_sample_grid['targets'],
            key=lambda x: (x['video_id'], x['main_gid']))

        self.wrapper_dset = self.WrapperDsetClass(torch_dataset)

        writer_queue = util_parallel.BlockingJobQueue(max_workers=config.io_workers)
        self.coco_stitcher = CocoStitchingManager(
            self.out_dset,
            short_code=self.short_code,
            chan_code=self.chan_code,
            # stiching_space='image',
            stiching_space='video',  # stitch in video space
            writer_queue=writer_queue,
            expected_minmax=(0, 1),
            assets_dname=config.assets_dname or self.short_code,
        )

        self.device = self.config.device
        sam_model = segment_anything.sam_model_registry['vit_h'](checkpoint=config.weights_fpath)
        sam_model = sam_model.eval()
        sam_model = sam_model.to(self.device)
        self.sam_model = sam_model

    def run(self):
        import rich
        from torch.utils import data as torch_data
        from watch.utils import util_progress
        from kwcoco.util import util_json
        import torch

        out_fpath = ub.Path(self.out_dset.fpath)
        pred_dpath = out_fpath.parent
        pred_dpath.ensuredir()

        self.proc_context.add_disk_info(pred_dpath)
        self.proc_context.start()

        loader = torch_data.DataLoader(
            self.wrapper_dset, batch_size=self.config.batch_size,
            num_workers=self.config.data_workers,
            shuffle=False, pin_memory=False,
            # worker_init_fn=worker_init_fn,
            collate_fn=ub.identity,  # disable collation
        )
        batch_iter = iter(loader)
        # torch.set_grad_enabled(False)

        pman = util_progress.ProgressManager()
        with pman, torch.no_grad():
            prog = pman.progiter(batch_iter, total=len(loader),
                                  desc=f'Predict {self.short_code}')
            for batch in prog:
                self.predict_batch(batch)

        for gid in self.coco_stitcher.managed_image_ids():
            self.coco_stitcher.submit_finalize_image(gid)
        self.coco_stitcher.writer_queue.wait_until_finished()

        obj = self.proc_context.stop()
        obj = util_json.ensure_json_serializable(obj)
        self.out_dset.dataset['info'].append(obj)
        self.out_dset._ensure_json_serializable()

        print(f'out_fpath={out_fpath}')
        rich.print(f'Pred Dpath: [link={pred_dpath}]{pred_dpath}[/link]')
        self.out_dset.dump(indent='    ')
        print('Finished predict')

    def predict_batch(self, batch):
        raise NotImplementedError


class SAMFeaturePredictor(DenseFeaturePredictor):
    short_code = 'sam'
    proc_name = 'watch.tasks.sam.predict'
    chan_code = 'sam.0:256'

    WrapperDsetClass = SAMWrapperDataset

    def __init__(self, config):
        super().__init__(config)

    def predict_batch(self, batch):
        import torch
        batch_chw255 = []
        batch_chw_isnodata = []
        for item in batch:
            chw255 = item['chw255']
            chw_isnodata = item['chw_isnodata_mask']
            batch_chw255.append(chw255)
            batch_chw_isnodata.append(chw_isnodata)
        bchw255 = torch.stack(batch_chw255, dim=0)
        bchw_isnodata = torch.stack(batch_chw_isnodata, dim=0)

        bchw_isnodata = bchw_isnodata.to(self.device)
        bchw255 = bchw255.to(self.device)
        input_image = self.sam_model.preprocess(bchw255)
        input_image[bchw_isnodata] = 0
        embedding_bchw = self.sam_model.image_encoder(input_image)

        feat_bchw = embedding_bchw.detach().cpu().numpy()

        for item, feat_chw in zip(batch, feat_bchw):
            data = feat_chw.transpose(1, 2, 0)
            gid = item['gid']
            asset_dsize = item['asset_dsize']
            asset_slice = item['asset_slice']

            # If the input was mostly nodata, mask the output
            output_isnodata_mask = item['output_isnodata_mask']
            data[output_isnodata_mask] = float('nan')

            scale_asset_from_vid = item['scale_asset_from_vid']
            self.coco_stitcher.accumulate_image(
                gid, asset_slice, data, asset_dsize=asset_dsize,
                scale_asset_from_stitchspace=scale_asset_from_vid)


def main(cmdline=1, **kwargs):
    """

    CommandLine:
        DVC_DATA_DPATH=$(geowatch_dvc --tags='phase2_data' --hardware=auto)
        DVC_EXPT_DPATH=$(geowatch_dvc --tags='phase2_expt' --hardware=auto)

    CommandLine:
        TEST_SAM=1 xdoctest -m watch.tasks.sam.predict main

    Example:
        >>> # xdoctest: +REQUIRES(env:TEST_SAM)
        >>> from watch.tasks.sam.predict import *  # NOQA
        >>> import watch
        >>> dvc_expt_dpath = watch.find_dvc_dpath(tags='phase2_expt', hardware='auto')
        >>> dvc_data_dpath = watch.find_dvc_dpath(tags='phase2_data', hardware='auto')
        >>> weights_fpath = dvc_expt_dpath / 'models/sam/sam_vit_h_4b8939.pth'
        >>> input_kwcoco = dvc_data_dpath / 'Drop6-MeanYear10GSD-V2/imgonly-KR_R001.kwcoco.zip'
        >>> kwargs = dict(weights_fpath=weights_fpath, input_kwcoco=input_kwcoco, output_kwcoco=(dvc_data_dpath / 'Drop6-MeanYear10GSD-V2/_test.kwcoco.zip'))
        >>> cmdline = 0
        >>> main(cmdline, **kwargs)

    Ignore:
        from watch.tasks.sam.predict import *  # NOQA
        import watch
        dvc_expt_dpath = watch.find_dvc_dpath(tags='phase2_expt', hardware='auto')
        dvc_data_dpath = watch.find_dvc_dpath(tags='phase2_data', hardware='auto')
        weights_fpath = dvc_expt_dpath / 'models/sam/sam_vit_h_4b8939.pth'
        input_kwcoco = dvc_data_dpath / 'Drop6-MeanYear10GSD-V2/imganns-AE_R001.kwcoco.zip'
        kwargs = dict(weights_fpath=weights_fpath, input_kwcoco=input_kwcoco, output_kwcoco=(dvc_data_dpath / 'Drop6-MeanYear10GSD-V2/imganns-AE_R001_sam.kwcoco.zip'))
        cmdline = 0
        config = SAMConfig.cli(cmdline=cmdline, data=kwargs, strict=True)
        import rich
        rich.print('config = ' + ub.urepr(config, nl=1))
        self = predictor = SAMFeaturePredictor(config)
        predictor.setup()
        # main(cmdline, **kwargs)

    Example:
        >>> # xdoctest: +SKIP
        >>> cmdline = 0
        >>> kwargs = dict()
        >>> main(cmdline=cmdline, **kwargs)
    """
    config = SAMConfig.cli(cmdline=cmdline, data=kwargs, strict=True)

    import rich
    rich.print('config = ' + ub.urepr(config, nl=1))

    predictor = SAMFeaturePredictor(config)
    predictor.setup()
    predictor.run()

    # if 0:
    #     import kwcoco
    #     coco_dset = kwcoco.CocoDataset.coerce(config.input_kwcoco)

    #     coco_img = coco_dset.images().coco_images[0]
    #     delayed = coco_img.imdelay(channels='red|green|blue')
    #     data = delayed.finalize()

    #     import kwimage
    #     import numpy as np
    #     norm01 = kwimage.normalize_intensity(data)
    #     norm255 = kwimage.ensure_uint255(np.nan_to_num(norm01))

    #     sam_model = segment_anything.sam_model_registry['vit_h'](checkpoint=config.weights_fpath)
    #     sam_model = sam_model.eval()
    #     device = 0
    #     sam_model = sam_model.to(device)

    #     import torch
    #     torch.set_grad_enabled(False)

    #     if 0:
    #         sam_predictor = segment_anything.SamPredictor(sam_model)
    #         sam_predictor.set_image(norm255)
    #         embedding_bchw_v1 = sam_predictor.get_image_embedding()
    #         embedding_hwc_v1 = embedding_bchw_v1[0].permute(1, 2, 0).detach().cpu().numpy()
    #         embedding_viz_v1 = kwimage.normalize_intensity(embedding_hwc_v1[..., 0:3])

    #     if 1:
    #         import torch
    #         norm01_resized, info = kwimage.imresize(
    #             norm01, dsize=(1024, 1024), interpolation='lanczos',
    #             letterbox=True, border_value=np.nan, return_info=True)
    #         norm01_resized = norm01_resized.clip(0, 1)
    #         # isnan_resized = np.isnan(norm01_resized)
    #         norm255_resized = (norm01_resized * 255)
    #         input_image_torch = torch.as_tensor(norm255_resized, device=sam_model.device)
    #         input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]
    #         input_image = sam_model.preprocess(input_image_torch)
    #         # input_image[:, :, 0:h, 0:w][nan_mask_bchw] = 0
    #         input_image.nan_to_num_()
    #         print(f'input_image_torch.shape={input_image_torch.shape}')
    #         print(f'input_image.shape={input_image.shape}')
    #         embedding_bchw = sam_model.image_encoder(input_image)
    #         embedding_hwc = embedding_bchw[0].permute(1, 2, 0).detach().cpu().numpy()
    #         embedding_viz = kwimage.normalize_intensity(embedding_hwc[..., 0:3])

    #     if 0:
    #         import kwplot
    #         kwplot.autompl()
    #         kwplot.figure(fnum=1, doclf=True)
    #         kwplot.imshow(norm255, pnum=(1, 3, 1))
    #         kwplot.imshow(embedding_viz, pnum=(1, 3, 2))
    #         kwplot.imshow(embedding_viz_v1, pnum=(1, 3, 3))


if __name__ == '__main__':
    """

    CommandLine:
        python ~/code/watch/watch/tasks/sam/predict.py
        python -m predict

    Ignore:

        python -m watch.tasks.sam.predict \
            --input_kwcoco "/home/joncrall/remote/toothbrush/data/dvc-repos/smart_data_dvc-ssd/Drop6-MeanYear10GSD-V2/imganns-AE_R001.kwcoco.zip" \
            --output_kwcoco "/home/joncrall/remote/toothbrush/data/dvc-repos/smart_data_dvc-ssd/Drop6-MeanYear10GSD-V2/imganns-AE_R001_sam.kwcoco.zip" \
            --weights_fpath "/home/joncrall/remote/toothbrush/data/dvc-repos/smart_expt_dvc/models/sam/sam_vit_h_4b8939.pth" \
            --window_overlap=0.33333 \
            --data_workers="2" \
            --io_workers 0 \
            --assets_dname="teamfeats"

        smartwatch visualize /home/joncrall/remote/toothbrush/data/dvc-repos/smart_data_dvc-ssd/Drop6-MeanYear10GSD-V2/imganns-AE_R001_sam.kwcoco.zip \
            --channels "red|green|blue,pan,sam.0:3,sam.3:6,sam.6:9" --smart

    """
    main()
