from lxml import etree


def create_original_alto(filename, width, height):
    """Create a base template for the canonical ALTOXML file. This
    file is supposed to contain the OCR of all individual cutouts.

    :param filename: Path to the original JPG file.
    :param width: Width of the original image.
    :param height: Height of the original image.
    :returns: XML object tree of template.
    """
    schema_url = 'http://www.loc.gov/standards/alto/ns-v4# http://www.loc.gov/alto/v4/alto-4-4.xsd'
    root = etree.Element("alto",
                         nsmap={
                             "xlink": "http://www.w3.org/1999/xlink",
                             "xsi": "http://www.w3.org/2001/XMLSchema-instance"
                         })
    root.set("{http://www.w3.org/2001/XMLSchema-instance}schemaLocation",
             schema_url)
    description = etree.SubElement(root, "Description")
    measurement_unit = etree.SubElement(description, "MeasurementUnit")
    measurement_unit.text = "pixel"
    source_image_information = etree.SubElement(description,
                                                "sourceImageInformation")
    file_name = etree.SubElement(source_image_information, "fileName")
    file_name.text = filename
    ocr_processing = etree.SubElement(description, "OCRProcessing", ID="OCR_0")
    ocr_processing_step = etree.SubElement(ocr_processing, "ocrProcessingStep")
    processing_software = etree.SubElement(ocr_processing_step,
                                           "processingSoftware")
    software_name = etree.SubElement(processing_software, "softwareName")
    software_name.text = "tesseract 5.3.2"
    layout = etree.SubElement(root, "Layout")
    page = etree.SubElement(layout,
                            "Page",
                            WIDTH=str(width),
                            HEIGHT=str(height),
                            PHYSICAL_IMG_NR="0",
                            ID="page_0")
    print_space = etree.SubElement(page,
                                   "PrintSpace",
                                   HPOS="0",
                                   VPOS="0",
                                   WIDTH=str(width),
                                   HEIGHT=str(height))

    return root


def modify_cutout(filename, bbox, angle):
    """Modifies an ALTOXML file for a cutout such that the translation
    and rotation during preprocessing is undone. The resulting file
    must show the OCR on the original image in the correct positions.

    :param filename: Path to the ALTOXML file of the cutout.
    :param bbox: Bbox of the cutout in x, y, wx, wy format.
    :param angle: Angle of the rotation in counter-clockwise degrees.
    :returns: XML object tree of modified cutout.
    """
    with open(filename, 'rb') as f:
        cutout = etree.parse(f)

    ns = {'ns0': 'http://www.loc.gov/standards/alto/ns-v3#'}

    # Find the PrintSpace element with xpath
    print_space = cutout.xpath('//ns0:PrintSpace', namespaces=ns)[0]

    # Find all tags with HPOS attributes
    elements = print_space.xpath('//*[@HPOS]')
    for elem in elements:
        hpos = int(elem.get('HPOS'))
        vpos = int(elem.get('VPOS'))
        # Add bbox coords and subtract 10 px padding
        # Padding was added during preprocessing
        elem.set('HPOS', str(hpos + bbox[0] - 10))
        elem.set('VPOS', str(vpos + bbox[1] - 10))

    # Find all ComposedBlock tags and set the ROTATION attribute
    composed_blocks = print_space.xpath('//ns0:ComposedBlock', namespaces=ns)
    for block in composed_blocks:
        block.set('ROTATION', str(-angle))

    return cutout


def copy_blocks_to_original(root, orig_width, orig_height, cutout):
    """Copy ComposedBlocks from ALTOXML of cutout to canonical
    ALTOXML.

    :param root: XML object of DOM of canonical ALTOXML.
    :param orig_width: Width of original image in px.
    :param orig_height: Height of original image in px.
    :param cutout: XML object of DOM of cutout ALTOXML.
    :returns: New root object with added ComposedBlocks.
    """
    ns = {'ns0': 'http://www.loc.gov/standards/alto/ns-v3#'}
    # Find PrintSpace in original root XML
    print_space_original = root.find('.//PrintSpace')

    # Find all ComposedBlock tags in cutout
    composed_blocks = cutout.findall('.//ns0:ComposedBlock', namespaces=ns)
    for block in composed_blocks:
        set_tag_to_localname(block)
        print_space_original.append(block)

    return root


def set_tag_to_localname(elem):
    """Recursively rename all tags in ELEM to not contain the
    namespace prefix. All changes happen in place!

    :param elem: The XML object to rename.
    :returns: None
    """
    elem.tag = etree.QName(elem).localname
    for child in elem:
        set_tag_to_localname(child)


def recompute_ids(root):
    """Make all IDs of the following tags unique in ascending order:

    - ComposedBlock (cblock_X)
    - TextBlock (block_X)
    - TextLine (line_X)
    - String (string_X)

    :param root: The XML object to renumber.
    :returns: The modified XML object tree.
    """
    elems = root.xpath(
        "//*[self::ComposedBlock or self::TextBlock or self::TextLine or self::String]"
    )

    cblock_ctr = 0
    block_ctr = 0
    line_ctr = 0
    string_ctr = 0
    for elem in elems:
        # cblocks
        if elem.get('ID').startswith('c'):
            elem.set('ID', f'cblock_{cblock_ctr}')
            cblock_ctr += 1
        # blocks
        elif elem.get('ID').startswith('b'):
            elem.set('ID', f'block_{block_ctr}')
            block_ctr += 1
        # lines
        elif elem.get('ID').startswith('l'):
            elem.set('ID', f'line_{line_ctr}')
            line_ctr += 1
        # strings
        elif elem.get('ID').startswith('s'):
            elem.set('ID', f'string_{string_ctr}')
            string_ctr += 1
        else:
            continue

    return root


if __name__ == '__main__':
    root = create_original_alto('00000000.jpg', 2088, 7024)
    cutout_one = modify_cutout(
        '../artifacts/predictions/117CA97B/00000000/0-deskewed.alto.xml',
        (622, 5519, 1084, 789), 4)
    cutout_two = modify_cutout(
        '../artifacts/predictions/117CA97B/00000000/1-deskewed.alto.xml',
        (618, 1129, 1104, 653), 0)
    new_root = copy_blocks_to_original(root, 2088, 7024, cutout_one)
    new_root = copy_blocks_to_original(new_root, 2088, 7024, cutout_two)
    new_root = recompute_ids(new_root)

    xml_string = etree.tostring(new_root,
                                pretty_print=True,
                                xml_declaration=True,
                                encoding="utf-8")
    with open("original.xml", "wb") as f:
        f.write(xml_string)
