#!/usr/bin/env python3

import bpy


def setDeviceCPU(context, cpref):
    cpref.compute_device_type = 'NONE'
    return True


def isDisplayDevice(device):
    return "(Display)" in device.name


def setUseRequestedDevice(context,
                          cpref,
                          device_type,
                          requested_device,
                          device_single):
    import _cycles
    # Empty device type means we'l ltry to render on a single card,
    # preferably non-display one.
    if requested_device == "":
        device_found = False
        display_device = None
        # Try to enable first non-display card.
        for device in cpref.devices:
            if device.type != device_type:
                device.use = False
                continue
            if isDisplayDevice(device):
                device.use = False
                display_device = device
            elif not device_found or not device_single:
                # Enable first non-display GPU.
                device_found = True
                device.use = True
            else:
                # Keep disabling rest of GPUs.
                device.use = False
        if not device_found:
            # Only display devices, enable first of them.
            display_device.use = True
            device_found = True
    else:
        device_found = False
        for device in cpref.devices:
            device_name = device.name.replace(" (Display)", "")
            if device_name == requested_device:
                if not device_found or not device_single:
                    device.use = True
                    device_found = True
                else:
                    device.use = False
            else:
                device.use = False
    return device_found


def setDeviceCUDA(context, cpref, requested_device, device_single):
    cpref.compute_device_type = 'CUDA'
    return setUseRequestedDevice(context,
                                 cpref,
                                 'CUDA',
                                 requested_device,
                                 device_single)


def setDeviceOpenCL(context, cpref, requested_device, device_single):
    cpref.compute_device_type = 'OPENCL'
    return setUseRequestedDevice(context,
                                 cpref,
                                 'OPENCL',
                                 requested_device,
                                 device_single)


def setDeviceGPU(context, cpref, requested_device, device_single):
    import _cycles
    has_cuda = has_opencl = False
    for device in _cycles.available_devices():
        if device[1] == 'CUDA':
            has_cuda = True
        if device[1] == 'OPENCL':
            has_opencl = True

    if has_cuda:
        return setDeviceCUDA(context, cpref, requested_device, device_single)
    if has_opencl:
        return setDeviceOpenCL(context, cpref, requested_device, device_single)
    return False


def logComputeDevices(cpref):
    device_type = cpref.compute_device_type
    if device_type == 'NONE':
        device_type = 'CPU'
    print("Compute device type: {}"  . format(device_type))
    if device_type == 'CPU':
        import _cycles
        for device in _cycles.available_devices():
            if device[1] == 'CPU':
                print("Using compute device: {}" . format(device[0]))
    else:
        for device in cpref.devices:
            if device.type != device_type:
                continue
            if device.use:
                print("Using compute device: {}" . format(device.name))


def logSystemInfo(cpref):
    import json
    info_devices = []
    for device in cpref.devices:
        info_device = {
            "name": device.name.replace(" (Display)", ""),
            "type": device.type,
        }
        info_devices.append(info_device)
    print("Benchmark Devices: {}" . format(json.dumps(info_devices)))


def main():
    import argparse
    import sys

    argv = sys.argv
    if "--" not in argv:
        return

    parser = argparse.ArgumentParser(
        description="Cycles benchmark helper script.")
    parser.add_argument("--benchmark-warmup",
                        help="Do quick warm-up render pass",
                        action='store_true',
                        default=False)
    parser.add_argument("--benchmark-device-type",
                        help="Device type to be rendered on",
                        default="CPU")
    parser.add_argument("--benchmark-device",
                        help="Device to be rendered on",
                        default="")
    parser.add_argument("--benchmark-device-single",
                        help="Use single device when multiple are found",
                        action='store_true',
                        default=False)
    parser.add_argument("--benchmark-system-info",
                        help="Dump whole system information",
                        action='store_true',
                        default=False)

    argv = argv[argv.index("--") + 1:]
    args = parser.parse_args(argv)

    context = bpy.context
    cpref = context.user_preferences.addons['cycles'].preferences

    # Adjust samples so we render real quick.
    if args.benchmark_warmup:
        for scene in bpy.data.scenes:
            scene.cycles.samples = 1
            scene.cycles.aa_samples = 1

    # Cycles device.
    device = 'CPU' if args.benchmark_device_type == 'CPU' else 'GPU'
    tile_size = 16 if device == 'CPU' else 512
    for scene in bpy.data.scenes:
        scene.cycles.device = device
        scene.render.tile_x = scene.render.tile_y = tile_size

    # Report number of CPU threads used.
    for scene in bpy.data.scenes:
        print("CPU threads used: {}" . format(scene.render.threads))
        # TODO(sergey): What to do if there are multiple scenes?
        break

    # Configure the compute device.
    if args.benchmark_device_type == 'CPU':
        device_ok = setDeviceCPU(context, cpref)
    elif args.benchmark_device_type == 'CUDA':
        device_ok = setDeviceCUDA(context,
                                  cpref,
                                  args.benchmark_device,
                                  args.benchmark_device_single)
    elif args.benchmark_device_type == 'OPENCL':
        device_ok = setDeviceOpenCL(context,
                                    cpref,
                                    args.benchmark_device,
                                    args.benchmark_device_single)
    elif args.benchmark_device_type == 'GPU':
        device_ok = setDeviceGPU(context,
                                 cpref,
                                 args.benchmark_device,
                                 args.benchmark_device_single)
    if not device_ok:
        sys.exit(1)

    if args.benchmark_system_info:
        logSystemInfo(cpref)
    else:
        logComputeDevices(cpref)


if __name__ == "__main__":
    main()
