# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import random
import shutil
import argparse

from glob import glob
from copy import deepcopy
from tqdm import tqdm

from anydata.utils.read import read_json
from anydata.utils.write import write_json
from anydata.utils.data import merge_dicts
from anydata.sync.sync_utils import aws_s3_cp, aws_s3_sync
from anydata.converters.utils import crawl, write_stats, get_shared_info, get_num_frames
from anydata.dataloaders.Base import normalize_metadata

#######################################################

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=str)
    parser.add_argument("--name", type=str, default='split_all')
    parser.add_argument("--unified_folder", type=str, default='cv_unified')
    parser.add_argument("--webbed_folder", type=str, default='cv_webbed')
    parser.add_argument("--local_path", type=str, default=os.environ.get('ANYDATA_LOCAL_ROOT', '/data'))
    parser.add_argument("--s3_path", type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets')
    parser.add_argument('--download', action='store_true')  # Download dataset for processing
    parser.add_argument('--upload', action='store_true')    # Upload split to s3
    parser.add_argument('--delete', action='store_true')    # Delete local dataset afterwards
    parser.add_argument('--official', action='store_true')  # If it's not official, use "_debug"
    parser.add_argument('--subfolder', type=str, default=None)
    parser.add_argument('--subset', type=str, default=None)

    parser.add_argument('--num_frames', type=str, default=None)
    parser.add_argument('--num_cameras', type=str, default=None)

    parser.add_argument('--with_cameras', type=str, nargs='+', default=None)
    parser.add_argument('--with_only_cameras', type=str, nargs='+', default=None)
    parser.add_argument('--without_cameras', type=str, nargs='+', default=None)

    parser.add_argument('--with_labels', type=str, nargs='+', default=None)
    parser.add_argument('--with_only_labels', type=str, nargs='+', default=None)
    parser.add_argument('--without_labels', type=str, nargs='+', default=None)

    parser.add_argument('--with_tags', type=str, nargs='+', default=None)
    parser.add_argument('--without_tags', type=str, nargs='+', default=None)

    parser.add_argument('--path_with', type=str, nargs='+', default=None)
    parser.add_argument('--path_without', type=str, nargs='+', default=None)

    parser.add_argument('--resolution', type=str, default=None)

    parser.add_argument('--subsample', type=str, default=None)             # --subset 0-10
    parser.add_argument('--subsample_slice', type=str, default=None)       # --subset_slice 0-10
    parser.add_argument('--subsample_random', type=str, default=None)      # --subset_random 42/0-10 (seed also)

    parser.add_argument('--keys_equal', type=str, default=None)

    parser.add_argument('--break_cameras', action='store_true')
    parser.add_argument('--webbed', action='store_true')
    parser.add_argument('--quiet', action='store_true')

    sto = parser.add_mutually_exclusive_group(required=True)
    sto.add_argument('--frames', action='store_const', dest='storage', const='frames')
    sto.add_argument('--videos', action='store_const', dest='storage', const='videos')

    args = parser.parse_args()

    # Get unified or webbed folder
    args.folder = args.webbed_folder if args.webbed else args.unified_folder

    # Determine storage mode
    args.folder += f'/{args.storage}'

    # Go from string to None
    if args.subset == 'None':
        args.subset = None

    # Make it debug if not official
    if not args.official:
        args.folder += '_debug'

    args.local_path = f'{args.local_path}/{args.folder}'
    args.s3_path = f'{args.s3_path}/{args.folder}'

    if args.path.startswith('/'): # If path is absolute, use directly
        args.src = args.path
    else:
        args.src = f'{args.local_path}/{args.path}'

    # Splits start with "split_"
    if not args.name.startswith('split_'):
        args.name = f'split_{args.name}'

    # Remove .json extension for now
    if args.name.endswith('.json'):
        args.name = args.name[:-len('.json')]
    # Create task from a dataset subfolder (e.g., LBM task), so use its name
    if args.subfolder is not None:
        args.name = f'{args.subfolder}/{args.name}'

    return args

#######################################################

def parse_create_split(files=None, metadata_all=None, suc=None, err=None):
    args = parse_args()
    return create_split(args, files=files, metadata_all=metadata_all, suc=suc, err=err)

#######################################################

def create_split(args, files=None, metadata_all=None, suc=None, err=None):

    if args.download:
        ### Download dataset (RGB only)
        s3_path = f'{args.s3_path}/{args.path}'
        local_path = f'{args.local_path}/{args.path}'
        if args.subfolder is not None:
            s3_path += f'/{args.subfolder}'
            local_path += f'/{args.subfolder}'
        aws_s3_sync(s3_path, local_path, extras='--exclude "*" --include "*.json"')

    # Find all relevant files
    if files is None:
        files = crawl(args.src, 'metadata.json')
    assert len(files) > 0, f'{args.src} not found!'

    # Get shared metadata
    metadata_path = f'{args.src}/metadata_shared.json'
    shared = read_json(metadata_path) if os.path.exists(metadata_path) else dict()

    # Dictionary to filter metadata and create specific splits (start from webbed filters if available)
    if 'webbed' in shared and 'filters' in shared['webbed']:
        filters = shared['webbed']['filters']
    else:
        filters = dict()

    # Add specific filters for this split
    if args.name != 'split_all':
        if args.with_cameras is not None:
            filters['with_cameras'] = args.with_cameras
        if args.with_labels is not None:
            filters['with_labels'] = args.with_labels
        if args.with_tags is not None:
            filters['with_tags'] = args.with_tags

        if args.with_only_cameras is not None:
            filters['with_only_cameras'] = args.with_only_cameras
        if args.with_only_labels is not None:
            filters['with_only_labels'] = args.with_only_labels

        if args.without_cameras is not None:
            filters['without_cameras'] = args.without_cameras
        if args.without_labels is not None:
            filters['without_labels'] = args.without_labels
        if args.without_tags is not None:
            filters['without_tags'] = args.without_tags

        if args.num_frames is not None:
            filters['num_frames'] = args.num_frames
        if args.num_cameras is not None:
            filters['num_cameras'] = args.num_cameras
        if args.resolution is not None:
            if '/' in args.resolution:
                resolution = args.resolution.split('/')
                filters['resolution'] = {resolution[0]: [int(v) for v in resolution[1].split('x')]}
            else:
                filters['resolution'] = [int(v) for v in args.resolution.split('x')]
        if args.keys_equal is not None:
            filters['keys_equal'] = args.keys_equal

        if args.path_with is not None:
            filters['path_with'] = args.path_with
        if args.path_without is not None:
            filters['path_without'] = args.path_without

    seqs = dict()
    n_frames, n_samples = 0, 0
    all_tags, all_cameras, all_labels = [], [], []
    all_metadata = {}

    if suc is not None:
        suc = set(suc.keys())

    for file in tqdm(files, ncols=96, leave=False, desc='Creating split'):

        # Open metadata file
        metadata_raw = read_json(file) if metadata_all is None else metadata_all[file]
        if metadata_raw is None: continue

        # Get relative file path
        rel_file = os.path.dirname(os.path.relpath(file, args.src))

        # Skip if it's not in the success list
        if suc is not None and not rel_file in suc:
            continue

        # Merge with shared and normalize to batch format
        metadata = normalize_metadata(merge_dicts(deepcopy(metadata_raw), shared))

        # Filter based on subfolder (e.g., LBM task)
        if args.subfolder is not None:
            valid = metadata.get('task', metadata.get('conventions', {}).get('language', {}).get('task')) == args.subfolder
            if not valid:
                continue

        # Optional filtering
        if len(filters) > 0:
            valid = True
            for key, val in filters.items():
                if key in ['num_cameras']:
                    if val.endswith('+'):
                        valid = len(metadata['cameras']) >= int(val[:-1])
                    elif val.endswith('-'):
                        valid = len(metadata['cameras']) <= int(val[:-1])
                    else:
                        valid = len(metadata['cameras']) == int(val)
                elif key in ['num_frames']:
                    num_frames = get_num_frames(metadata[key], metadata['cameras'])
                    if val.endswith('+'):
                        valid = num_frames >= int(val[:-1])
                    elif val.endswith('-'):
                        valid = num_frames <= int(val[:-1])
                    else:
                        valid = num_frames == int(val)
                elif key.startswith('with_only_'):
                    key_meta = key[len('with_only_'):]
                    valid = val == metadata[key_meta]
                elif key.startswith('with_'):
                    key_meta = key[len('with_'):]
                    if key.endswith('tags'): # Tags are at top level (normalized)
                        valid = all([v in metadata[key_meta] for v in val])
                    else: # Others are not
                        valid = all([v in metadata[key_meta] for v in val])
                elif key.startswith('without_'):
                    key_meta = key[len('without_'):]
                    if key.endswith('tags'): # Tags are at top level (normalized)
                        valid = all([v not in metadata[key_meta] for v in val])
                    else: # Others are not
                        valid = all([v not in metadata[key_meta] for v in val])
                elif key in ['path_with']:
                    valid = all([v in file for v in val])
                elif key in ['path_without']:
                    valid = any([v not in file for v in val])
                elif key in ['keys_equal']:
                    k, v = val.split('=')
                    k = k.split('/')
                    m = metadata
                    while len(k) > 0:
                        m, k = m[k[0]], k[1:]
                    valid = m == v
                else:
                    if isinstance(val, dict):  # Dict, compare specific keys
                        valid = all([v in metadata['cameras'] for v in val.keys()])  # Check if cameras exist
                        if valid:  # If cameras exist, continue  checking
                            if not isinstance(metadata[key], dict):  # Resolution is the same for this sequence
                                valid = all([metadata[key] == v for v in val.values()])  
                            else:  # Resolution is not the same
                                valid = all([v in metadata[key].keys() for v in val.keys()])  # Check if cameras exist
                                if valid:
                                    valid = all([metadata[key][k] == val[k] for k in val.keys()])
                    else:  # Not dict, compare directly
                        valid = metadata[key] == val
                if not valid:
                    break
            if not valid:
                continue

        # Get relevant information from metadata
        tags = metadata['tags']
        labels = metadata['labels']
        cameras = metadata['cameras']
        length = metadata['num_frames']

        # num_frames may be a nested {cam: {label: n}} dict when counts differ
        # across cameras; unwrap to a scalar by taking the first leaf value
        while isinstance(length, dict):
            if not length:
                break
            length = next(iter(length.values()))
        if isinstance(length, dict) or not length:
            continue

        # Get sequence name and add to the sequence dictionary
        key = os.path.dirname(file.replace(f'{args.src}/', ''))

        if not args.webbed: # Store sequence length
            seqs[key] = length
        else: # Store tarfile list
            seqs[key] = metadata['tarfiles']

        # Keep track of number of files
        n_samples += length
        n_frames += length * len(cameras)

        # Update list of all available cameras
        all_tags.extend(tags)
        all_tags = list(set(all_tags))

        # Update list of all available cameras
        all_cameras.extend(cameras)
        all_cameras = list(set(all_cameras))

        # Update list of all available labels
        all_labels.extend(labels)
        all_labels = list(set(all_labels))

        all_metadata[rel_file] = metadata_raw

    # Needs to have at least one sequence
    if len(seqs) == 0:
        print('\nEMPTY SPLIT!!! Exiting...\n')
        return 'failed'
    
    # Filter found sequences based on subset success
    if suc is not None:
        assert suc == set(all_metadata.keys()), 'Sequences and Successes dont match'

    # Subset based on skips
    if args.subsample is not None:
        keys = list(seqs.keys())
        st, fn = args.subsample.split('-')
        selected = []
        for k in range(int(st), int(fn)):
            selected.extend(keys[k::100])
        selected = sorted(list(set(selected)))
        seqs = {key: seqs[key] for key in selected}
    # Subset based on slices
    if args.subsample_slice is not None:
        keys = list(seqs.keys())
        st, fn = args.subsample_slice.split('-')
        st = float(st) / 100 * len(keys)
        fn = float(fn) / 100 * len(keys)
        selected = sorted(keys[int(st):int(fn)])
        seqs = {key: seqs[key] for key in selected}
    # Subset based on random sampling
    if args.subsample_random is not None:
        seed, perc = args.subsample_random.split('/')
        random.seed(int(seed))
        keys = list(seqs.keys())
        random.shuffle(keys)
        st, fn = perc.split('-')
        st = float(st) / 100 * len(keys)
        fn = float(fn) / 100 * len(keys)
        selected = sorted(keys[int(st):int(fn)])
        seqs = {key: seqs[key] for key in selected}

    # Store split size (number of sequences, samples, frames, cameras)
    size = dict(
        seqs=len(seqs),
        samples=n_samples,
        frames=n_frames,
        cameras=len(all_cameras),
    )

    removed = len(files) - len(seqs)
    if suc is None: size['removed'] = removed

    # Count how many failed sequences were logged
    tmp_folder = f'{args.path}/tmp'
    if os.path.exists(tmp_folder):
        failed = glob(f'{tmp_folder}/*.txt')
        failed = [os.path.basename(f)[:-len('.txt')] for f in failed]
        if err is not None:
            err_flat = [e.replace('/', '__') for e in err.keys()]
            if args.subset is not None:
                failed = [f for f in failed if f in err_flat]
            for f in failed:
                if f not in err_flat:
                    err[f] = 'DID NOT FINISH'
            err_flat = [e.replace('/', '__') for e in err.keys()]
            assert set(failed) == set(err_flat), f'Failed and Errors dont match:  {len(failed)} x {len(err_flat)}'
        failed = len(failed)
    else:
        assert err is None or len(err) == 0, f'Failed and Errors dont match:  {len(err)}'
        failed = 0
    size['failed'] = failed

    # Store split information (tags, cameras and labels)
    shared_tags, shared_labels, shared_cameras = get_shared_info(all_metadata)
    info = dict(
        tags=sorted(all_tags),
        labels=sorted(all_labels),
        cameras=sorted(all_cameras),
        shared_tags=sorted(shared_tags),
        shared_labels=sorted(shared_labels),
        shared_cameras=sorted(shared_cameras),
    )

    if args.webbed: # Get webbed store value (e.g., 41b4) from path name
        store = args.path.split('/')[-1].split('__')[1].split('_')[1]
        info['webbed'] = store

    # Remove subfolder prefix from the sequence, so it can be loaded
    if args.subfolder is not None:
        seqs = {key.replace(f'{args.subfolder}/', ''): val for key, val in seqs.items()}
    seqs = {key: seqs[key] for key in sorted(seqs.keys())}

    # Write split to disk
    split = dict(
        filters=filters,
        size=size,
        info=info,
        sequences=seqs,
    )
    write_json(f'{args.src}/{args.name}.json', split)

    if args.upload: # Upload split to s3
        aws_s3_cp(f'{args.src}/{args.name}.json', f'{args.s3_path}/{args.path}/{args.name}.json')

    if args.subset is None:
        all_metadata = {key: all_metadata[key] for key in seqs.keys()}
        stats_name = write_stats(args, metadata_path, all_metadata, shared, errors=err, name=args.name, removed=removed)
        stats_name = os.path.basename(stats_name)
        if args.upload:
            aws_s3_cp(f'{args.src}/{stats_name}', f'{args.s3_path}/{args.path}/{stats_name}')

    if args.delete: # Delete dataset afterwards
        shutil.rmtree(args.src)

    # Display relevant data
    if not args.quiet:
        print(f'#################################################')
        print(f'################### SPLITTING ###################')
        print(f'#################################################')
        print(f'### NAME: {args.src}/{args.name}.json')
        print(f'#################################################')
        print(f'### SIZE:')
        for key, val in size.items():
                print(f'###### {key} : {val}')
        print(f'#################################################')
        print(f'### FILTERS:')
        for key, val in filters.items():
            if isinstance(val, dict):
                print(f'###### {key}')
                for key2, val2 in val.items():
                    print(f'######### {key2} : {val2}')
            else:
                print(f'###### {key} : {val}')
        print(f'#################################################')
        if shared is not None:
            print(f'### SHARED:')
            for key, val in shared.items():
                if isinstance(val, dict):
                    print(f'###### {key}')
                    for key2, val2 in val.items():
                        print(f'######### {key2} : {val2}')
                else:
                    print(f'###### {key} : {val}')
        print(f'#################################################')

    return files

#######################################################

if __name__ == '__main__':
    args = parse_args()
    if args.break_cameras:
        shared = read_json(f'{args.src}/metadata_shared.json')
        cameras = shared.get('cameras', shared.get('cameras_available', []))
        name = args.name
        files = None
        for i, cam in enumerate(cameras):
            args.with_cameras = [cam]
            args.name = f'{name}__{cam}'
            args.quiet = i < len(cameras) - 1
            files = create_split(args, files=files)
    else:
        create_split(args)

#######################################################