prdutils.py

#! /usr/bin/env python3
"""Utilities for parsing and manipulating Primitive Root Diffuser results from
http://www.oliverprime.com/prd/
"""
import argparse
from pathlib import Path
import dataclasses
from dataclasses import dataclass, field
from typing import Union, List, Tuple, Dict, Optional
import json

import requests
from bs4 import BeautifulSoup

try:
    import bpy
except ImportError:
    bpy = None

Tag = 'bs4.element.Tag'
TableResult = List[List[int]]
Pathlike = Union[str, Path]


@dataclass
class ParseResult:
    """A calculated set of results
    """

    speed_of_sound: int #: Speed of sound (form field)
    low_freq: int #: Lowest frequency in Hz (form field)
    hi_freq: int #: Highest frequency in Hz (form field)
    prime_num: int #: Prime number P (form field)
    prime_root: int #: Primitive root of :attr:`P <prime_num>` (form field)
    num_cols: int #: Number of well columns (form field)
    num_rows: int #: Number of well rows (form field)
    well_width: Optional[float] = None
    """The calculated width of each well in the grid"""

    table: TableResult = field(default_factory=lambda: [[]])
    """The well heights of each element in the grid (in cm)

    The table is laid out with rows as the first dimension and the columns
    in the second
    """

    def calc_well_counts(self) -> Dict[int,int]:
        """Count how many of each well height are contained in the :attr:`table`
        """
        counts = {}
        for row in self.table:
            for w in row:
                if w not in counts:
                    counts[w] = 0
                counts[w] += 1
        return {k:counts[k] for k in sorted(counts)}

    def offset_well_heights(self, offset: int):
        """Offset the well heights by the given number

        The offset will be added to each of the elements in the :attr:`table`
        in place.  This may be useful if a height of zero (0) is not desired.
        """
        for row in self.table:
            for i in range(self.num_cols):
                row[i] += offset

    @classmethod
    def load(cls, filename: Pathlike) -> 'ParseResult':
        """Load a previously saved :class:`ParseResult` from the given filename
        """
        if not isinstance(filename, Path):
            filename = Path(filename)
        return cls.from_json(filename.read_text())

    @classmethod
    def from_json(cls, json_str: str) -> 'ParseResult':
        kw = json.loads(json_str)
        return cls(**kw)

    def serialize(self) -> Dict:
        return dataclasses.asdict(self)

    def save(self, filename: Pathlike):
        """Save the results to the given filename (JSON formatted)
        """
        if not isinstance(filename, Path):
            filename = Path(filename)
        filename.write_text(self.to_json())

    def to_json(self) -> str:
        d = self.serialize()
        return json.dumps(d)

    def pprint(self) -> str:
        """Format the results for human readable output
        """
        lines = []
        horiz_sep = '+'.join(['----' for _ in range(self.num_cols)])
        lines.append(horiz_sep)
        for row in self.table:
            line = '|'.join([f'{w:^4d}' for w in row])
            lines.append(line)
            lines.append(horiz_sep)

        lines.append('')
        counts = self.calc_well_counts()
        total_length = 0
        for cm, n in counts.items():
            if cm == 0:
                continue
            lines.append(f'{cm:2d} cm: {n}')
            total_length += cm * n
        lines.append('')
        lines.append(f'Total length: {total_length} cm')
        lines.append(f'Lf: {self.low_freq}Hz, Hf: {self.hi_freq}Hz')
        lines.append(f'P: {self.prime_num}, Prime Root: {self.prime_root}')
        lines.append(f'Shape: ({self.num_cols}, {self.num_rows})')
        return '\n'.join(lines)


def parse_table(table_el: Tag) -> TableResult:
    """Parse the Well Heights table from HTML
    """
    result = []
    for tr in table_el.tbody.find_all('tr'):
        row_values = []
        for td in tr.find_all('td'):
            value = int(td.get_text())
            row_values.append(value)
        result.append(row_values)
    return result

def parse_form(form_el: Tag) -> ParseResult:
    """Parse the form values from HTML
    """
    form_name_map = {
        'C':'speed_of_sound', 'F_LOW':'low_freq', 'F_HI':'hi_freq',
        'P':'prime_num', 'PR':'prime_root',
        'COLS':'num_cols', 'ROWS':'num_rows',
    }
    parse_kwargs = {}
    for input_el in form_el.find_all('input'):
        parse_attr = form_name_map.get(input_el['name'])
        if parse_attr is None:
            continue
        val = int(input_el['value'])
        parse_kwargs[parse_attr] = val
    return ParseResult(**parse_kwargs)

def parse_doc(doc: 'bs4.BeautifulSoup') -> ParseResult:
    """Parse a results HTML document
    """
    result = None
    result_table = None
    for tbl in doc.find_all('table'):
        if tbl.parent.name == 'form':
            result = parse_form(tbl.parent)
        else:
            result_table = tbl
    if result_table is not None:
        result.table = parse_table(result_table)
    diff_prop_h = doc.find('h4', string='Diffuser properties')
    span = diff_prop_h.next_sibling
    for s in span.strings:
        if 'Well length/width' in s:
            w = s.split('=')[2]
            w = w.split('cm')[0]
            w = float(w)
            result.well_width = w
    return result

def get_from_str(s: str) -> ParseResult:
    """Parse the results page from the given string
    """
    doc = BeautifulSoup(s, 'html5lib')
    return parse_doc(doc)

def get_from_url(url: str) -> ParseResult:
    """Parse the results page from the given url

    The url can be copied from the "permalink" on the page. Since the calculated
    results require a form submission, this function will make two http requests.
    """
    r = requests.get(url)
    if not r.ok:
        r.raise_for_status()

    doc = BeautifulSoup(r.content, 'html5lib')
    form_el = doc.find('form')
    form_data = {}
    for input_el in form_el.find_all('input'):
        form_data[input_el['name']] = input_el['value']
    base_url = url.split('?')[0]
    r = requests.post(base_url, form_data)
    if not r.ok:
        r.raise_for_status()
    return get_from_str(r.content)


def get_from_file(filename: Pathlike) -> ParseResult:
    """Parse the results page from an html file
    """
    if not isinstance(filename, Path):
        filename = Path(filename)
    return get_from_str(filename.read_text())


if bpy is not None:
    from bpy_extras import io_utils

    def move_to_collection(obj, coll):
        to_remove = []
        if coll not in obj.users_collection:
            coll.objects.link(obj)
        for oth_coll in obj.users_collection:
            if oth_coll is coll:
                continue
            to_remove.append(oth_coll)
        for oth_coll in to_remove:
            oth_coll.objects.unlink(obj)

    def clear_collection_objects(coll):
        to_remove = list(coll.objects.values())
        for obj in to_remove:
            coll.objects.unlink(obj)



    class PrdSceneProps(bpy.types.PropertyGroup):
        base_coll_name: bpy.props.StringProperty(
            name='Base Collection Name',
            description='Base Collection Name',
            default='PrdBase',
        )
        obj_coll_name: bpy.props.StringProperty(
            name='Object Collection Name',
            description='Object Collection Name',
            default='PrdObjects',
        )
        base_coll: bpy.props.PointerProperty(
            type=bpy.types.Collection,
            name='Base Collection',
            description='Collection to place the base mesh in',
        )
        obj_coll: bpy.props.PointerProperty(
            type=bpy.types.Collection,
            name='Object Collection',
            description='Collection to store all instanced meshes',
        )

        well_width: bpy.props.IntProperty(
            name='Well Width',
            description='The width/height (in cm) of each well',
        )
        array_shape: bpy.props.IntVectorProperty(
            name='Array Shape',
            description='The number of columns(x) and rows(y) of the well array',
            subtype='XYZ_LENGTH',
        )
        array_dimensions: bpy.props.FloatVectorProperty(
            name='Array Dimensions',
            description='Overall Dimensions (in scene space) of the well array',
            subtype='XYZ_LENGTH',
        )


    class PrdWellProps(bpy.types.PropertyGroup):
        row: bpy.props.IntProperty(default=-1)
        column: bpy.props.IntProperty(default=-1)
        height: bpy.props.IntProperty(default=-1)

    class PrdBuilderProps(bpy.types.PropertyGroup):
        import_url: bpy.props.StringProperty(
            name='Import Url',
            description='Url to import data from',
        )
        json_str: bpy.props.StringProperty()
        instance_mode_options = [
            ('COLLECTION', 'Collection', 'Create wells as collection instances'),
            ('OBJECT', 'Object', 'Duplicate each well using the same data (mesh)'),
            ('OBJECT_DATA', 'Object/Data', 'Duplicate each well and its data (mesh)'),
        ]
        instance_mode: bpy.props.EnumProperty(
            items=instance_mode_options,
            name='Instance Mode',
            description='Method of creating the individual well objects',
            default='COLLECTION',
        )
        well_offset: bpy.props.IntProperty(
            name='Well Offset',
            description='Amount to offset well heights',
            default=1,
        )
        @classmethod
        def check(cls, value):
            return value in [opt[0] for opt in cls.instance_mode_options]

    class PrdBuilderOp(bpy.types.Operator):
        """Build an array of cubes matching the wells from a PRD table
        """
        bl_idname = "prdutils.build"
        bl_label = 'Build PRD Scene'

        def execute(self, context):
            build_settings = context.scene.prd_data.builder_props
            assert len(build_settings.json_str)
            parsed = ParseResult.from_json(build_settings.json_str)
            if build_settings.well_offset != 0:
                parsed.offset_well_heights(build_settings.well_offset)
            build_settings.json_str = ''
            self.build_collections(context)
            self.setup_scene_props(context, parsed)
            self.build_objects(context, parsed)
            return {'FINISHED'}

        def build_collections(self, context):
            scene_props = context.scene.prd_data


            for attr in ['base_coll', 'obj_coll']:
                coll = getattr(scene_props, attr)
                if coll is not None:
                    continue
                coll_name = getattr(scene_props, f'{attr}_name')
                bpy.ops.collection.create(name=coll_name)
                coll = bpy.data.collections[-1]
                setattr(scene_props, attr, coll)
                clear_collection_objects(coll)
                context.scene.collection.children.link(coll)

            scene_props.base_coll.hide_render = True

        def setup_scene_props(self, context, parsed: ParseResult):
            scene_props = context.scene.prd_data

            width = scene_props.well_width = parsed.well_width
            scene_props.array_shape[0] = parsed.num_cols
            scene_props.array_shape[1] = parsed.num_rows
            scene_props.array_dimensions.x = parsed.num_cols * width
            scene_props.array_dimensions.y = parsed.num_rows * width
            tbl = parsed.table
            max_well = max([y for x in tbl for y in x])
            scene_props.array_dimensions.z = max_well


        def build_objects(self, context, parsed: ParseResult):
            scene_props = context.scene.prd_data
            build_settings = context.scene.prd_data.builder_props

            width = scene_props.well_width
            half_width = width / 2
            total_y = scene_props.array_shape[1]

            base_coll, obj_coll = scene_props.base_coll, scene_props.obj_coll

            bpy.ops.mesh.primitive_cube_add(size=width)
            base_cube = context.active_object
            move_to_collection(base_cube, base_coll)

            context.scene.cursor.location = [0, 0, 0]
            base_cube.location.z = half_width
            bpy.ops.object.origin_set(type='ORIGIN_CURSOR')
            base_cube.dimensions.z = 1
            bpy.ops.object.transform_apply(location=False, properties=False)

            instance_mode = build_settings.instance_mode

            for i, row in enumerate(parsed.table):
                y = i * -width + total_y + half_width
                for j, well_height in enumerate(row):
                    x = width * j + half_width
                    if instance_mode == 'COLLECTION':
                        bpy.ops.object.collection_instance_add(collection=base_coll.name)
                        obj = context.active_object
                    elif instance_mode == 'OBJECT':
                        bpy.ops.object.duplicate(linked=True)
                        obj = context.active_object
                    elif instance_mode == 'OBJECT_DATA':
                        bpy.ops.object.duplicate(linked=False)
                        obj = context.active_object
                    move_to_collection(obj, obj_coll)

                    obj.location.x = x
                    obj.location.y = y
                    obj.scale.z = well_height
                    obj.prd_data.row = i
                    obj.prd_data.column = j
                    obj.prd_data.height = well_height

    class PrdImportOp(bpy.types.Operator, io_utils.ImportHelper):
        """Import a saved :class:`ParseResult` into the current scene
        """
        bl_idname = 'prdutils.json_import'
        bl_label = 'Import PRD Data from json'

        filename_ext = '.json'
        filter_glob: bpy.props.StringProperty(
            default='*.json',
            options={'HIDDEN'},
            maxlen=255,
        )

        @classmethod
        def poll(cls, context):
            return context.scene is not None

        def execute(self, context):
            build_settings = context.scene.prd_data.builder_props
            parsed = ParseResult.load(self.filepath)
            build_settings.json_str = Path(self.filepath).read_text()

            bpy.ops.prdutils.build()
            return {'FINISHED'}

        def invoke(self, context, event):
            context.window_manager.fileselect_add(self)
            return {'RUNNING_MODAL'}

    class PrdFromUrlOp(bpy.types.Operator):
        bl_idname = 'prdutils.from_url'
        bl_label = 'Import PRD Data from Url'

        def execute(self, context):
            build_settings = context.scene.prd_data.builder_props
            parsed = get_from_url(build_settings.import_url)
            json_str = parsed.to_json()
            build_settings.json_str = json_str
            bpy.ops.prdutils.build()
            return {'FINISHED'}


    class PrdFromUrlPanel(bpy.types.Panel):
        bl_idname = 'VIEW_3D_PT_prd_from_url'
        bl_owner_id = 'prdutils.from_url'
        bl_label = 'PRD Utils'
        bl_space_type = 'VIEW_3D'
        bl_region_type = 'UI'
        bl_category = 'Tool'

        def draw(self, context):
            layout = self.layout
            scene_props = context.scene.prd_data
            build_settings = scene_props.builder_props
            main_box = layout.box()
            main_box.label(text='PRD Utils')

            box = main_box.box()
            box.label(text='From Url')
            box.prop(build_settings, 'import_url')
            box.operator('prdutils.from_url')

            box = main_box.box()
            box.label(text='From File')
            box.operator('prdutils.json_import')

            box = main_box.box()
            box.label(text='Import Settings')
            box.prop(build_settings, 'instance_mode')
            box.prop(build_settings, 'well_offset')
            box.prop(scene_props.base_coll)
            box.prop(scene_props.obj_coll)


    def menu_func(self, context):
        self.layout.operator_context = 'INVOKE_DEFAULT'
        self.layout.operator(PrdImportOp.bl_idname, text=f'Text {PrdImportOp.bl_label}')

    bl_classes = [
        PrdSceneProps, PrdWellProps, PrdBuilderProps,
        PrdBuilderOp, PrdImportOp, PrdFromUrlOp, PrdFromUrlPanel,
    ]

    def register():
        for cls in bl_classes:
            bpy.utils.register_class(cls)
        PrdSceneProps.builder_props = bpy.props.PointerProperty(type=PrdBuilderProps)
        bpy.types.Scene.prd_data = bpy.props.PointerProperty(type=PrdSceneProps)
        bpy.types.Object.prd_data = bpy.props.PointerProperty(type=PrdWellProps)

        bpy.types.TOPBAR_MT_file_import.append(menu_func)

    def unregister():
        del bpy.types.Scene.prd_data.builder_props
        del bpy.types.Object.prd_data
        del bpy.types.Scene.prd_data
        del bpy.ops.prdutils.from_url
        for cls in reversed(bl_classes):
            bpy.utils.unregister_class(cls)

        bpy.types.TOPBAR_MT_file_import.remove(menu_func)

    if __name__ == '__main__':
        register()