import os 
import sys
import json
import boto3
import subprocess

from tqdm import tqdm


bucket = "tri-ml-sandbox-16011-us-west-2-datasets"
prefix = f"cv_datasets/processed"

unused_labels = ['scene_flow_fwd', 'scene_flow_bwd']


def list_s3(s3_path, mode):
    proc = subprocess.Popen(f'aws s3 ls {s3_path}/', stdout=subprocess.PIPE, shell=True)
    (data, _) = proc.communicate()
    data = data.decode().replace('/\n', ' ').replace('\n', ' ').replace('  ', ' ').split(' ')
    if mode == 'json':
        out = [d for d in data if d.endswith('.json')]
    elif mode == 'folder':
        out = [data[i+1] for i in range(len(data)) if data[i] == 'PRE']
    proc.kill()
    return out


def delete_s3_splits(dataset):
    path = f's3://{bucket}/{prefix}/{dataset}'
    files = list_s3(path, mode='json')
    for f in files:
        command = f'aws s3 rm {path}/{f} --quiet'
        os.system(command)


def upload_s3(dataset):
    path = f's3://{bucket}/{prefix}/{dataset}'

    dataset_local = dataset.replace('/', '__')

    os.system(f'aws s3 cp {dataset_local}_info.json {path}/stats_info.json --quiet')
    os.system(f'aws s3 cp {dataset_local}_data.json {path}/stats_data.json --quiet')
    os.system(f'aws s3 cp {dataset_local}_check.json {path}/stats_check.json --quiet')
    os.system(f'aws s3 cp {dataset_local}_split.json {path}/split_all.json --quiet')

    os.system(f'rm {dataset_local}_info.json')
    os.system(f'rm {dataset_local}_data.json')
    os.system(f'rm {dataset_local}_check.json')
    os.system(f'rm {dataset_local}_split.json')


def get_pages(dataset, bucket, prefix):
    s3_client = boto3.client('s3')
    paginator = s3_client.get_paginator('list_objects_v2')
    return paginator.paginate(Bucket=bucket, Prefix=f'{prefix}/{dataset}')


def populate(dataset, pages):
    data = {}
    nseqs = 0
    nsamples = {}
    cams, labels = [], []
    pages = tqdm(pages)
    extra = '/'.join(dataset.split('/')[1:])
    print(f'\n#################### {dataset}')
    for page in pages:
        for obj in page['Contents']:
            try:
                file = obj['Key']
                if 'zbuffer' in file: file = file.replace('/zbuffer', '')
                if 'lidar'   in file: file = file.replace('/lidar', '')
                file = file.split('/')[2:]

                dset = file[0]
                lab, cam, name = file[-3:]

                if name.endswith('.json'):
                    continue

                seq = '/'.join(file[1:-3])
                seq = seq.replace(extra, '')
                if seq.startswith('/'): seq = seq[1:]

                if seq not in data:
                    data[seq] = {} 
                    nseqs += 1       
                if lab not in data[seq]:
                    data[seq][lab] = {}
                    if lab not in labels and lab not in unused_labels: labels.append(lab)
                if cam not in data[seq][lab]:
                    data[seq][lab][cam] = []
                    if cam not in cams: cams.append(cam)

                data[seq][lab][cam].append(name)

                if cam not in nsamples: nsamples[cam] = 0
                if lab == 'rgb': nsamples[cam] += 1
                total = sum(nsamples.values())

            except:
                pass

        str_cams = str(sorted(cams)).replace("'", "").replace(" ", "")
        str_labels = str(sorted(labels)).replace("'", "").replace(" ", "")
        str_nsamples = str(sorted(list(nsamples.values()))).replace("'", "").replace(" ", "")
        total = sum(nsamples.values())
        pages.set_description(
            f'### {dset} | T: {total} | S: {nseqs} | L: {str_labels} | C: {str_cams} |',
        )
    
    stats = dict(
        data=data, 
        info= dict(
            total=sum(nsamples.values()), 
            nsamples=nsamples,
            nseqs=nseqs, 
            labels=labels, 
            cameras=cams,
        )
    )

    return stats


def prep_stats(dataset):

    if os.path.exists(f'{dataset}.json'):
        with open(f'{dataset}.json', "r") as json_file:
            stats = json.load(json_file)
    else:
        pages = get_pages(dataset, bucket, prefix)
        stats = populate(dataset, pages)

    stats, split = process_stats(stats)

    print(f'################################################ {dataset}')
    print(f'### Sequences: {stats["info"]["nseqs_valids"]} / {stats["info"]["nseqs"]}')
    print(f'### Samples: {stats["info"]["nsamples_valids"]} / {list(stats["info"]["nsamples"].values())}')
    print(f'### Total: {stats["info"]["total_valids"]} / {stats["info"]["total"]}')
    print(f'### Cameras ({len(stats["info"]["cameras"])}): {stats["info"]["cameras"]}')
    print(f'### Labels: {stats["info"]["labels"]}')
    print(f'################################################ {dataset}')

    dataset_local = dataset.replace('/', '__')
    with open(f'{dataset_local}_info.json', "w") as json_file:
        json.dump(stats['info'], json_file, indent=4)
    with open(f'{dataset_local}_data.json', "w") as json_file:
        json.dump(stats['data'], json_file, indent=4)
    with open(f'{dataset_local}_check.json', "w") as json_file:
        json.dump(stats['check'], json_file, indent=4)
    with open(f'{dataset_local}_split.json', "w") as json_file:
        json.dump(split, json_file, indent=4)

    delete_s3_splits(dataset)
    upload_s3(dataset)
    
def process_stats(stats):

    info = stats['info']
    data = stats['data']

    same_name = True

    valids, invalids = [], []
    for key_seq, val_seq in data.items():
        labels = list(val_seq.keys())
        labels = [l for l in labels if l not in unused_labels]
        valid = labels == info['labels']
        if not valid: invalids.append(['labels', key_seq]); continue
        lengths, names = [], []
        for key_lab, val_lab in val_seq.items():
            if not valid: continue
            if key_lab not in labels: continue
            cams = list(val_lab.keys())
            valid = cams == info['cameras']
            if not valid: invalids.append(['cameras', key_seq]); continue
            for key_cam, val_cam in val_lab.items():
                lengths.append(len(val_cam))
                name = [n.split('.')[0] for n in val_cam]
                names.append(name)
        if valid:
            valid = not any([length != lengths[0] for length in lengths])
            if not valid: invalids.append(['length', key_seq]); continue
            valids.append(key_seq)

        if same_name:
            same_name = not any([name != names[0] for name in names])

    if same_name:
        split = {key: val['rgb'][cams[0]] for key, val in data.items() if key in valids}
        samples = sum([len(val) for val in split.values()])
    else:
        split = {key: {cam: val['rgb'][cam] for cam in cams} for key, val in data.items() if key in valids}
        samples = {cam: sum([len(val[cam]) for val in split.values()]) for cam in cams}
        samples = list(samples.values())[0]
    nseqs, ncams = len(split), len(cams)
    total = samples * ncams

    stats['info']['total_valids'] = total
    stats['info']['nsamples_valids'] = samples
    stats['info']['nseqs_valids'] = nseqs

    stats['check'] = {
        'valids': valids,
        'invalids': invalids,
    }

    return stats, split


def rec_prep_stats(dataset, i, levels):
    if i == max(levels): return
    path = f's3://{bucket}/{prefix}/{dataset}'
    folders = list_s3(path, mode='folder')
    for folder in folders:
        rec_prep_stats(f'{dataset}/{folder}', i + 1, levels)
        if i + 1 in levels:
            # print('PREP STATS', i + 1, levels, f'{dataset}/{folder}')
            prep_stats(f'{dataset}/{folder}')
    if i == 0 and i in levels:
        # print('PREP STATS', i, levels, dataset)
        prep_stats(dataset)


if __name__ == "__main__":

    dataset = sys.argv[1]
    if len(sys.argv) > 2:
        levels = [int(v) for v in sys.argv[2:]]
        rec_prep_stats(dataset, 0, levels)
    else:
        prep_stats(dataset)