#!/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):
    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
        # Try to enable first non-display card.
        for device in cpref.devices:
            if isDisplayDevice(device):
                device.use = False
            elif not device_found:
                # Enable first non-display GPU.
                device_found = True
                device.use = True
            else:
                # Keep disanling rest of GPUs.
                device.use = False
        if not device_found:
            # Only display devices, enable first of them.
            cpref.devices[0].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:
                device.use = True
                device_found = True
            else:
                device.use = False
    return device_found


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


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


def setDeviceGPU(context, cpref, requested_device):
    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)
    if has_opencl:
        return setDeviceOpenCL(context, cpref, requested_device)
    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.rstrip(" (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-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

    # 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)
    elif args.benchmark_device_type == 'OPENCL':
        device_ok = setDeviceOpenCL(context, cpref, args.benchmark_device)
    elif args.benchmark_device_type == 'GPU':
        device_ok = setDeviceGPU(context, cpref, args.benchmark_device)
    if not device_ok:
        sys.exit(1)

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


if __name__ == "__main__":
    main()