#!/usr/bin/env python

# SPDX-License-Identifier: GPL-2.0-or-later
# Copyright (C) 2023, Jacopo Mondi - Ideas on Board Oy
#
# Parse Android .xml configuration file and extract the LSC tables.
#
# Print to standard output a "LensShadingCorrection" section, understandable by
# libcamera LSC algorithm, that can be pasted to the sensor configuration file.

import argparse
import string
import sys
import re
import xml.etree.ElementTree as et


def sanitize(name):
    return re.sub(r"[\n\t\s]*", "", name)


def split_table(table):
    values = ""
    for v in table.text.strip(' ').split():
        values += v.strip('[').strip(']') + ", "
    return values


def print_cell(cell):
    lsc_template = string.Template('''        #${name} - ${illuminant}
        - ct: ${ct}
          resolution: ${res}
          r: [${red}]
          gr: [${greenr}]
          gb: [${greenb}]
          b: [${blue}]''')

    illuminant = cell.find("illumination")
    ct = illuminant_to_ct(illuminant)

    template_dict = {
        'name': sanitize(cell.find("name").text),
        'illuminant': sanitize(illuminant.text),
        'ct': ct,
        'res': sanitize(cell.find("resolution").text)
    }

    red_table = cell.find("LSC_SAMPLES_red")
    greenr_table = cell.find("LSC_SAMPLES_greenR")
    greenb_table = cell.find("LSC_SAMPLES_greenB")
    blue_table = cell.find("LSC_SAMPLES_blue")

    if red_table is None or greenr_table is None or greenb_table is None or blue_table is None:
        return

    template_dict['red'] = split_table(red_table)
    template_dict['greenr'] = split_table(greenr_table)
    template_dict['greenb'] = split_table(greenb_table)
    template_dict['blue'] = split_table(blue_table)

    return lsc_template.substitute(template_dict)


def illuminant_to_ct(illuminant):
    # Standard CIE Illiminants to Color Temperature in Kelvin
    # https://en.wikipedia.org/wiki/Standard_illuminant
    #
    # Not included (and then ignored when parsing the configuration file):
    # - "Horizon" == D50 == 5003
    # - "BW" == ?
    # - "PREFLASH" == ?
    illuminants_dict = {
        'A': 2856,
        'D50': 5003,
        'D65': 6504,
        'D75': 7504,
        'F11_TL84': 4000,
        'F2_CWF': 4230,
    }

    ill_key = sanitize(illuminant.text)
    try:
        ct = illuminants_dict[ill_key]
    except KeyError:
        return None

    return ct


# Make sure the cell is well formed and return it
def filter_cells(cell, res, lsc_cells):
    name = cell.find("name")
    resolution = cell.find("resolution")
    illumination = cell.find("illumination")
    vignetting = cell.find("vignetting")

    if name is None or resolution is None or \
       illumination is None or vignetting is None:
        return

    # Skip tables for smaller sensor resolutions
    if res != sanitize(resolution.text):
        return

    # Skip tables for which we don't know how to translate the illuminant value
    ct = illuminant_to_ct(illumination)
    if ct is None:
        return

    # Only pick tables with vignetting == 70
    if sanitize(vignetting.text) != "[70]":
        return

    lsc_cells.append(cell)


# Get the "LSC" node
def find_lsc_table(root):
    sensor = root.find('sensor')
    if sensor is None:
        print("Failed to find \"sensor\" node in config file")
        raise Exception

    lsc = sensor.find('LSC')
    if lsc is None:
        print("Filed to find \"LSC\" node in config file")
        raise Exception

    return lsc

# libcamera LSC algorithm only operates on a single resolution.
# Find the largest sensor mode among the ones reported in the LSC tables


def parse_max_res(cells):
    max_res = ""
    max_size = 0

    for cell in cells:
        resolution = sanitize(cell.find("resolution").text)
        [w, h] = resolution.split('x')

        area = int(w) * int(h)
        if area > max_size:
            max_res = resolution

    return max_res


def main(argv):
    # Parse command line arguments.
    parser = argparse.ArgumentParser(
        description='Parse Android camera configuration file to extract LSC tables')
    parser.add_argument('--file', '-f', required=True,
                        help='Path to the Android .xml configuration file')
    args = parser.parse_args(argv[1:])

    root = et.parse(args.file).getroot()
    try:
        lsc_node = find_lsc_table(root)
    except Exception:
        return 1

    cells = lsc_node.findall("cell")

    max_res = parse_max_res(cells)
    if max_res == "":
        return

    lsc_cells = []
    for cell in cells:
        filter_cells(cell, max_res, lsc_cells)

    lsc_section = '''  - LensShadingCorrection:
      x-size: [ 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625 ]
      y-size: [ 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625 ]
      sets:
'''

    for cell in lsc_cells:
        lsc_section += print_cell(cell) + "\n"

    print(lsc_section)


if __name__ == '__main__':
    sys.exit(main(sys.argv))