# Legacy webdataset utilities -- used by webdataset_per_camera.py (also legacy).
# Moved here to keep utils.py clean for web_local.py.

import io
import copy
import random
from types import SimpleNamespace

import numpy as np
import torch
import torchvision.transforms as T
from PIL import Image

from anydata.utils.types import is_tensor
from anydata.webdataset.utils import (
    rgb_to_byte_array, mask_to_byte_array,
    check_invalid_resolution, check_invalid_matrix,
)


def numpy_to_byte_array(mat):
    return mat.tobytes()


def add_length(data, key):
    if key not in data.keys():
        data[key] = 1
    else:
        data[key] += 1
    return data


def check_valid(lengths, thr=None):
    valid = True
    values = [val for key, val in lengths.items()
              if key not in ['action', 'language'] and not key.startswith('lat')]
    if len(values) == 0:
        return False
    for i in range(1, len(values)):
        valid = values[0] == values[i]
        if not valid:
            break
    if valid and thr is not None:
        valid = values[0] >= thr + 1
    return valid


def to_numpy(data):
    if isinstance(data, dict):
        return {key: to_numpy(val) for key, val in data.items()}
    if isinstance(data, list):
        return [to_numpy(val) for val in data]
    if isinstance(data, torch.Tensor):
        return data.cpu().float().numpy()
    return data


def _to_npz_compatible(val):
    if isinstance(val, np.ndarray):
        return val
    try:
        return np.array(val, dtype=object)
    except Exception:
        return np.array([val], dtype=object)


def _to_npz_tree(val):
    if isinstance(val, dict):
        return {k: _to_npz_tree(v) for k, v in val.items()}
    if isinstance(val, list):
        return [_to_npz_tree(v) for v in val]
    return _to_npz_compatible(val)


def _action_keys_are_time_only(action_dict):
    if not action_dict:
        return False
    first_key = next(iter(action_dict.keys()))
    return not isinstance(first_key, tuple)


def _language_keys_are_time_only(language_dict):
    if not language_dict:
        return False
    first_key = next(iter(language_dict.keys()))
    return not isinstance(first_key, tuple)


def encode_latents(data, time_st, time_fn, cams, label, vae, b):
    labs = []
    for cam in cams:
        lab = data[label]
        lab = {key: val.to(device='cuda:0', dtype=torch.bfloat16, non_blocking=True)
               for key, val in lab.items() if time_st <= key[0] < time_fn and key[1] == cam}
        lab = [val for val in lab.values()]
        while len(lab) < b + 1:
            lab.append(lab[-1].clone())
        lab = torch.stack(lab, 1)
        labs.append(lab)

    lab = torch.stack(labs, 0)
    with torch.no_grad():
        if label == 'rgb':
            latents = vae.encode_single_video_rgb(lab)
        elif label == 'depth':
            latents = vae.encode_single_video_depth(lab)
        else:
            raise ValueError(f"Unsupported latent label {label}")

    latents = torch.split(latents, latents.shape[0] // len(labs), 0)

    latents_dict = {}
    for i in range(len(latents)):
        latents_i = latents[i]
        for j in range(latents_i.shape[0]):
            latents_dict[j + time_st, cams[i]] = latents_i[j]

    return latents_dict


def save_data_single(args, data, n, prefix, file=None, time_st=-1e9, time_fn=1e9, pad_after=1e9,
                     vae=None, lat_language=None, store_raw=False, store_latents=True):
    sample = {'__key__': prefix}
    lengths = {}
    cams = list(set([time_cam[1] for time_cam in data['rgb'].keys()]))
    base_key = (0, 0)
    to_pil = T.ToPILImage()
    if lat_language is None:
        lat_language = {}

    use_lowdim = store_raw and isinstance(data.get('lowdim', None), dict) and len(data.get('lowdim', {})) > 0
    if use_lowdim and 'language' in data:
        language_dict = data['language']
        if _language_keys_are_time_only(language_dict):
            base_val = language_dict.get(0, {})
        else:
            base_val = language_dict.get(base_key, {})
        language = {key: val.cpu().numpy() if is_tensor(val) else val for key, val in base_val.items()}
        language = {k: _to_npz_compatible(v) for k, v in language.items()}
        for key2, entry in data['lowdim'].items():
            if not isinstance(entry, dict):
                continue
            if 'language' not in entry:
                entry['language'] = language

    if store_latents:
        if 'rgb' in data.keys():
            data['lat_rgb'] = encode_latents(data, time_st, time_fn, cams, 'rgb', vae, args.buffer)
        if 'depth' in data.keys():
            data['lat_depth'] = encode_latents(data, time_st, time_fn, cams, 'depth', vae, args.buffer)
        if 'language' in data.keys():
            language_dict = data['language']
            if _language_keys_are_time_only(language_dict):
                prompt = language_dict.get(0, {}).get('prompt', None)
            else:
                prompt = language_dict.get(base_key, {}).get('prompt', None)
            if prompt is not None:
                if isinstance(prompt, list):
                    prompt = random.sample(prompt, 1)
                if prompt[0] not in lat_language:
                    lat_language[prompt[0]] = encoded = vae.encode_text(prompt)
                else:
                    encoded = lat_language[prompt[0]]
                encoded = {key: val[0] for key, val in encoded.items()}
                encoded['prompt'] = prompt
                data['lat_language'] = {(time_st, 0): encoded}

    cams = cams + ['ego']
    for cam in cams:
        time_cam_dict = {}
        timestep_dict = {}
        action_dict = {}
        pose_dict = {}
        intrinsics_dict = {}
        for key1 in data.keys():
            if use_lowdim and key1 in ['intrinsics', 'extrinsics', 'pose', 'action', 'timestep', 'time_cam']:
                continue
            for time in range(time_st, time_fn):
                key2, key3 = (time, cam), (time - 1, cam)
                if key1 in ['rgb'] and store_raw:
                    if key2 in data[key1].keys():
                        if check_invalid_resolution(data[key1][key2], args.check_resolution):
                            raise ValueError(f'Invalid rgb resolution {data[key1][key2].shape} {args.check_resolution}')
                        sample[f'{key1}.{key2}.jpg'.replace(' ', '')] = rgb_to_byte_array(to_pil(data[key1][key2]))
                        add_length(lengths, key1)
                    elif time >= pad_after:
                        sample[f'{key1}.{key2}.jpg'.replace(' ', '')] = sample[f'{key1}.{key3}.jpg'.replace(' ', '')]
                        add_length(lengths, key1)
                elif key1.startswith('lat'):
                    if key2 in data[key1].keys():
                        if key1.endswith('language'):
                            val = to_numpy(data[key1][key2])
                            if isinstance(val, dict):
                                val = {k: _to_npz_compatible(v) for k, v in val.items()}
                            sample[f'{key1}.{key2}.npz'.replace(' ', '')] = val
                        else:
                            sample[f'{key1}.{key2}.npz'.replace(' ', '')] = dict(data=data[key1][key2].float().cpu().numpy())
                        add_length(lengths, key1)
                elif key1.startswith('mask'):
                    if key2 in data[key1].keys():
                        sample[f'{key1}.{key2}.png'.replace(' ', '')] = mask_to_byte_array(data[key1][key2])
                        add_length(lengths, key1)
                    elif time >= pad_after:
                        sample[f'{key1}.{key2}.png'.replace(' ', '')] = sample[f'{key1}.{key3}.png'.replace(' ', '')]
                        add_length(lengths, key1)
                elif key1 in ['lowdim']:
                    if key2 in data[key1].keys():
                        entry = data[key1][key2]
                        if not isinstance(entry, dict):
                            entry = {'data': entry}
                        entry = {k: _to_npz_compatible(v) for k, v in entry.items()}
                        sample[f'{key1}.{key2}.npz'.replace(' ', '')] = entry
                        add_length(lengths, key1)
                    elif time >= pad_after:
                        sample[f'{key1}.{key2}.npz'.replace(' ', '')] = sample[f'{key1}.{key3}.npz'.replace(' ', '')]
                        add_length(lengths, key1)
                elif key1 in ['depth'] and store_raw:
                    if key2 in data[key1].keys():
                        if check_invalid_resolution(data[key1][key2], args.check_resolution):
                            raise ValueError(f'Invalid depth resolution {data[key1][key2].shape} {args.check_resolution}')
                        sample[f'{key1}.{key2}.npz'.replace(' ', '')] = dict(depth=data[key1][key2][0].cpu().numpy())
                        add_length(lengths, key1)
                    elif time >= pad_after:
                        sample[f'{key1}.{key2}.npz'.replace(' ', '')] = sample[f'{key1}.{key3}.npz'.replace(' ', '')]
                        add_length(lengths, key1)
                elif key1 in ['intrinsics', 'extrinsics', 'pose']:
                    if key2 in data[key1].keys():
                        value = np.asarray(data[key1][key2])
                        if check_invalid_matrix(value):
                            raise ValueError(f'Invalid key {key1} {value}')
                        if key1 == 'intrinsics':
                            intrinsics_dict[key2] = value
                        if key1 == 'pose':
                            pose_dict[key2] = value
                        if key2[1] != 'ego':
                            add_length(lengths, key1)
                    elif time >= pad_after:
                        value = np.array(value)
                        if key1 == 'intrinsics':
                            intrinsics_dict[key2] = value
                        if key1 == 'pose':
                            pose_dict[key2] = value
                        add_length(lengths, key1)
                elif key1 in ['action']:
                    if cam != 0:
                        continue
                    action_dict_in = data[key1]
                    action_time_only = _action_keys_are_time_only(action_dict_in)
                    if action_time_only:
                        key2_base, key3_base = key2[0], key3[0]
                    else:
                        key2_base, key3_base = (key2[0], 0), (key3[0], 0)
                    if key2_base in action_dict_in.keys():
                        value = to_numpy(action_dict_in[key2_base])
                        if check_invalid_matrix(value):
                            raise ValueError(f'Invalid key {key1} {value}')
                        if not isinstance(value, dict):
                            value = {'action': value}
                        length = value['action'].shape[0]
                        off = pad_after - time
                        if off < length:
                            for key_action1 in value.keys():
                                if isinstance(value[key_action1], dict):
                                    for key_action2 in value[key_action1].keys():
                                        value[key_action1][key_action2][off + 1:] = value[key_action1][key_action2][off]
                                else:
                                    value[key_action1][off + 1:] = value[key_action1][off]
                        action_dict[key2] = to_numpy(value)
                        add_length(lengths, key1)
                    elif time >= pad_after:
                        value = copy.deepcopy(value)
                        for key_action1 in value.keys():
                            if isinstance(value[key_action1], dict):
                                for key_action2 in value[key_action1].keys():
                                    value[key_action1][key_action2][1:] = value[key_action1][key_action2][0]
                            else:
                                value[key_action1][1:] = value[key_action1][0]
                        action_dict[key2] = to_numpy(value)
                        add_length(lengths, key1)
                elif key1 in ['language'] and store_raw:
                    if use_lowdim:
                        continue
                    language_dict = data[key1]
                    if _language_keys_are_time_only(language_dict):
                        base_val = language_dict.get(0, {})
                    else:
                        base_val = language_dict.get(base_key, {})
                    language = {key: val.cpu().numpy() if is_tensor(val) else val for key, val in base_val.items()}
                    language = {k: _to_npz_compatible(v) for k, v in language.items()}
                    sample[f'{key1}.{key2}.npz'.replace(' ', '')] = language
                    add_length(lengths, key1)
                elif key1 in ['timestep']:
                    if key2 in data[key1].keys():
                        value = np.asarray(data[key1][key2])
                        timestep_dict[key2] = value
                        add_length(lengths, key1)
                    elif time >= pad_after:
                        value = np.array(value)
                        value += 1
                        timestep_dict[key2] = value
                        add_length(lengths, key1)
                elif key1 in ['time_cam']:
                    if key2 in data[key1].keys():
                        value = np.asarray(data[key1][key2])
                        time_cam_dict[key2] = value
                        add_length(lengths, key1)
                    elif time >= pad_after:
                        value = np.array(value)
                        value[0] += 1
                        time_cam_dict[key2] = value
                        add_length(lengths, key1)
                elif key1 in ['optical_flow']:
                    if key2 in data[key1].keys():
                        for key3 in data[key1][key2].keys():
                            sample[f'{key1}.{key2}.{key3}.npz'.replace(' ', '')] = dict(
                                data=data[key1][key2][key3].cpu().numpy()
                            )
                            add_length(lengths, key1)

        if len(time_cam_dict) > 0:
            sample[f'time_cam.({time_st},{cam}).npz'] = dict(data=_to_npz_compatible(time_cam_dict))
        if len(timestep_dict) > 0:
            sample[f'timestep.({time_st},{cam}).npz'] = dict(data=_to_npz_compatible(timestep_dict))
        if len(action_dict) > 0:
            sample[f'action.({time_st},{cam}).npz'] = dict(data=_to_npz_compatible(action_dict))
        if len(intrinsics_dict) > 0:
            sample[f'intrinsics.({time_st},{cam}).npz'] = dict(data=_to_npz_compatible(intrinsics_dict))
        if len(pose_dict) > 0:
            sample[f'pose.({time_st},{cam}).npz'] = dict(data=_to_npz_compatible(pose_dict))

    length = len(sample.keys()) - 1
    sample_length = {'__key__': prefix}
    for key, val in sample.items():
        if not key.startswith('__'):
            key = f'{key[:-4]}.{length}{key[-4:]}'
            sample_length[key] = val
    sample = sample_length
    valid = check_valid(lengths)
    return valid, (file, sample), lat_language


def prep_buffer(args):

    if args.zone == 'west':
        args.s3_path = args.s3_path.replace('tri-ml-datasets', 'tri-ml-sandbox-16011-us-west-2-datasets')

    if args.context_sample is not None and '/' not in args.context_sample:
        args.context_sample = int(args.context_sample)

    if args.buffer is not None:
        args.samples_per_file = 1
        args.buffer_string = args.buffer

        args.buffer = args.buffer.replace('s', '_f')
        args.buffer = args.buffer.replace('pf', '_pf')
        args.buffer = args.buffer.replace('pl', '_pl')
        args.buffer = args.buffer.replace('cf', '_cf')

        buffer = args.buffer.split('_')

        args.buffer = int(buffer[0])
        if args.buffer == -1:
            args.buffer = int(1e9)
        args.buffer_pad = False
        args.buffer_fixed = False
        args.buffer_overlap = None
        args.buffer_stride = None
        args.buffer_random = None

        args.buffer_pad_first = None
        args.buffer_pad_last = None
        args.buffer_crop_first = None
        args.buffer_crop_last = None

        for buf in buffer[1:]:
            if buf.startswith('f'):
                args.buffer_overlap = int(buf[1:])
                args.buffer_pad = False
                args.buffer_fixed = False
            elif buf.startswith('pf'):
                args.buffer_pad_first = int(buf[2:])
            elif buf.startswith('cf'):
                args.buffer_crop_first = int(buf[2:])
            elif buf.startswith('pl'):
                if len(buf) == 2:
                    args.buffer_pad_last = 'pad'
                else:
                    args.buffer_pad_last = int(buf[2:])
            elif buf.startswith('cl'):
                args.buffer_crop_last = int(buf[2:])
            elif buf.startswith('p'):
                val = int(buf[1:])
                args.buffer += val
                buf = f's{val}--{val}'
                buf = buf[1:].split('--')
                if len(buf) == 1:
                    buf.append(0)
                else:
                    buf[1] = int(buf[1])
                if '-' in buf[0]:
                    buf0 = buf[0].split('-')
                    buf[0] = [int(buf0[0]),int(buf0[1])]
                else:
                    buf[0] = int(buf[0])
                args.buffer_stride = buf
                args.buffer_pad = True
                args.buffer_fixed = True
            elif buf.startswith('r'):
                args.buffer_random = int(buf[1:])
            elif buf == 'pad':
                args.buffer_pad = True
            elif buf == 'fixed':
                args.buffer_fixed = True

    return args


class CfgWrapper(SimpleNamespace):
    def has(self, key, default=None):
        if hasattr(self, key):
            return getattr(self, key)
        return default


def _wrap_cfg(cfg_dict):
    def _wrap(obj):
        if isinstance(obj, dict):
            return CfgWrapper(**{k: _wrap(v) for k, v in obj.items()})
        if isinstance(obj, list):
            return [ _wrap(v) for v in obj ]
        return obj
    return _wrap(cfg_dict)


def unwrap_cfg(cfg_obj):
    if isinstance(cfg_obj, CfgWrapper):
        return {k: unwrap_cfg(v) for k, v in cfg_obj.__dict__.items()}
    if isinstance(cfg_obj, list):
        return [unwrap_cfg(v) for v in cfg_obj]
    return cfg_obj


def get_name(name, cfg, dataset, args):
    if args.buffer is not None:
        name = name + f'_buf[{args.buffer_string}]'

    context = cfg.context
    if len(context) == 0:
        name = name + f'-ctx[]'
    elif context[0] == 'sequence':
        name = name + f'-ctx[s]'
    else:
        name = name + f'-ctx{context}'

    if cfg.has('context_sample'):
        name = name + f'_temp{cfg.context_sample}'.replace('/', ',')
    if cfg.has('cameras_context_sample'):
        name = name + f'_spat{cfg.cameras_context_sample}'.replace('/', ',')
    if cfg.has('context_interval'):
        name = name + f'_interv{cfg.context_interval}'
    if cfg.has('temporal_proximity'):
        name = name + '_T'
    if cfg.has('spatial_proximity'):
        name = name + '_S'
    if cfg.has('spatiotemporal_proximity'):
        name = name + '_ST'
    if args.min_context is not None:
        name = name + f'_MC{args.min_context}'

    cameras = cfg.cameras
    if len(cameras) == 1:
        name = name + f'_cam{cameras}'
    else:
        name = name + f'_cam[ALL]'

    if cfg.has('ego_camera'):
        name = name + f'_ego[{cfg.ego_camera}]'

    depth_type = cfg.has('depth_type', None)
    if depth_type is not None:
        name = name + f'_{depth_type}'

    resolution = cfg.has('resolution')
    if resolution:
        name = name + f'_reso{cfg.resolution}'

    resize = cfg.has('augmentation') and cfg.augmentation.has('resize')
    if resize:
        name = name + f'_resi{cfg.augmentation.resize}'

    if cfg.has('augmentation'):
        if cfg.augmentation.has('crop_borders'):
            name = name + f'_cropb{cfg.augmentation.crop_borders}'
        if cfg.augmentation.has('crop_random'):
            name = name + f'_cropr{cfg.augmentation.crop_random}'
        if cfg.augmentation.has('fill_depth'):
            name = name + f'_fill'

    if not args.buffer:
        length = len(dataset)
        if cfg.has('repeat'):
            name = name + f'_{length // cfg.repeat}x{cfg.repeat}'
        else:
            name = name + f'_{length}x1'

    if args.split is not None:
        name = name + f'-{args.split}'
    if args.split == 'test':
        name = name + f'/{args.world_size}'
    if args.check_resolution is not None:
        name = name + f'_rc{args.check_resolution}'

    return name.replace(' ', '').replace("'", "")


def make_cfg_name(args, cfg_dict):
    cfg = copy.deepcopy(cfg_dict)
    if 'augmentation' not in cfg:
        cfg['augmentation'] = {}

    if args.suffix is not None:
        cfg['path'] = [f"{c}/{args.suffix}" for c in cfg['path']]

    if args.fill_depth:
        cfg['augmentation']['fill_depth'] = True
    if args.repeat is not None:
        cfg['repeat'] = args.repeat
    if args.context is not None:
        cfg['context'] = args.context
    if args.context_sample is not None:
        cfg['context_sample'] = args.context_sample
    if args.cameras_context_sample is not None:
        cfg['cameras_context_sample'] = args.cameras_context_sample
    if args.resize is not None:
        cfg['augmentation']['resize'] = args.resize
    if args.resolution is not None and len(args.resolution) > 1:
        cfg['resolution'] = args.resolution
    if args.cameras is not None:
        cfg['cameras'] = args.cameras
    if args.resize_ratio is not None and args.resize_ratio != '':
        if 'resize' in cfg['augmentation'] and len(cfg['augmentation']['resize']) > 1:
            cfg['augmentation']['resize'][1] = int(args.resize_ratio)

    return _wrap_cfg(cfg)