# SPDX-License-Identifier: BSD-2-Clause # # Copyright (C) 2019, Raspberry Pi Ltd # # camera tuning tool for CCM (colour correction matrix) from ctt_image_load import * from ctt_awb import get_alsc_patches import colors from scipy.optimize import minimize from ctt_visualise import visualise_macbeth_chart import numpy as np """ takes 8-bit macbeth chart values, degammas and returns 16 bit """ ''' This program has many options from which to derive the color matrix from. The first is average. This minimises the average delta E across all patches of the macbeth chart. Testing across all cameras yeilded this as the most color accurate and vivid. Other options are avalible however. Maximum minimises the maximum Delta E of the patches. It iterates through till a minimum maximum is found (so that there is not one patch that deviates wildly.) This yields generally good results but overall the colors are less accurate Have a fiddle with maximum and see what you think. The final option allows you to select the patches for which to average across. This means that you can bias certain patches, for instance if you want the reds to be more accurate. ''' matrix_selection_types = ["average", "maximum", "patches"] typenum = 0 # select from array above, 0 = average, 1 = maximum, 2 = patches test_patches = [1, 2, 5, 8, 9, 12, 14] ''' Enter patches to test for. Can also be entered twice if you would like twice as much bias on one patch. ''' def degamma(x): x = x / ((2 ** 8) - 1) # takes 255 and scales it down to one x = np.where(x < 0.04045, x / 12.92, ((x + 0.055) / 1.055) ** 2.4) x = x * ((2 ** 16) - 1) # takes one and scales up to 65535, 16 bit color return x def gamma(x): # Take 3 long array of color values and gamma them return [((colour / 255) ** (1 / 2.4) * 1.055 - 0.055) * 255 for colour in x] """ FInds colour correction matrices for list of images """ def ccm(Cam, cal_cr_list, cal_cb_list, grid_size): global matrix_selection_types, typenum imgs = Cam.imgs """ standard macbeth chart colour values """ m_rgb = np.array([ # these are in RGB [116, 81, 67], # dark skin [199, 147, 129], # light skin [91, 122, 156], # blue sky [90, 108, 64], # foliage [130, 128, 176], # blue flower [92, 190, 172], # bluish green [224, 124, 47], # orange [68, 91, 170], # purplish blue [198, 82, 97], # moderate red [94, 58, 106], # purple [159, 189, 63], # yellow green [230, 162, 39], # orange yellow [35, 63, 147], # blue [67, 149, 74], # green [180, 49, 57], # red [238, 198, 20], # yellow [193, 84, 151], # magenta [0, 136, 170], # cyan (goes out of gamut) [245, 245, 243], # white 9.5 [200, 202, 202], # neutral 8 [161, 163, 163], # neutral 6.5 [121, 121, 122], # neutral 5 [82, 84, 86], # neutral 3.5 [49, 49, 51] # black 2 ]) """ convert reference colours from srgb to rgb """ m_srgb = degamma(m_rgb) # now in 16 bit color. # Produce array of LAB values for ideal color chart m_lab = [colors.RGB_to_LAB(color / 256) for color in m_srgb] """ reorder reference values to match how patches are ordered """ m_srgb = np.array([m_srgb[i::6] for i in range(6)]).reshape((24, 3)) m_lab = np.array([m_lab[i::6] for i in range(6)]).reshape((24, 3)) m_rgb = np.array([m_rgb[i::6] for i in range(6)]).reshape((24, 3)) """ reformat alsc correction tables or set colour_cals to None if alsc is deactivated """ if cal_cr_list is None: colour_cals = None else: colour_cals = {} for cr, cb in zip(cal_cr_list, cal_cb_list): cr_tab = cr['table'] cb_tab = cb['table'] """ normalise tables so min value is 1 """ cr_tab = cr_tab / np.min(cr_tab) cb_tab = cb_tab / np.min(cb_tab) colour_cals[cr['ct']] = [cr_tab, cb_tab] """ for each image, perform awb and alsc corrections. Then calculate the colour correction matrix for that image, recording the ccm and the colour tempertaure. """ ccm_tab = {} for Img in imgs: Cam.log += '\nProcessing image: ' + Img.name """ get macbeth patches with alsc applied if alsc enabled. Note: if alsc is disabled then colour_cals will be set to None and no the function will simply return the macbeth patches """ r, b, g = get_alsc_patches(Img, colour_cals, grey=False, grid_size=grid_size) """ do awb Note: awb is done by measuring the macbeth chart in the image, rather than from the awb calibration. This is done so the awb will be perfect and the ccm matrices will be more accurate. """ r_greys, b_greys, g_greys = r[3::4], b[3::4], g[3::4] r_g = np.mean(r_greys / g_greys) b_g = np.mean(b_greys / g_greys) r = r / r_g b = b / b_g """ normalise brightness wrt reference macbeth colours and then average each channel for each patch """ gain = np.mean(m_srgb) / np.mean((r, g, b)) Cam.log += '\nGain with respect to standard colours: {:.3f}'.format(gain) r = np.mean(gain * r, axis=1) b = np.mean(gain * b, axis=1) g = np.mean(gain * g, axis=1) """ calculate ccm matrix """ # ==== All of below should in sRGB ===## sumde = 0 ccm = do_ccm(r, g, b, m_srgb) # This is the initial guess that our optimisation code works with. original_ccm = ccm r1 = ccm[0] r2 = ccm[1] g1 = ccm[3] g2 = ccm[4] b1 = ccm[6] b2 = ccm[7] ''' COLOR MATRIX LOOKS AS BELOW R1 R2 R3 Rval Outr G1 G2 G3 * Gval = G B1 B2 B3 Bval B Will be optimising 6 elements and working out the third element using 1-r1-r2 = r3 ''' x0 = [r1, r2, g1, g2, b1, b2] ''' We use our old CCM as the initial guess for the program to find the optimised matrix ''' result = minimize(guess, x0, args=(r, g, b, m_lab), tol=0.01) ''' This produces a color matrix which has the lowest delta E possible, based off the input data. Note it is impossible for this to reach zero since the input data is imperfect ''' Cam.log += ("\n \n Optimised Matrix Below: \n \n") [r1, r2, g1, g2, b1, b2] = result.x # The new, optimised color correction matrix values optimised_ccm = [r1, r2, (1 - r1 - r2), g1, g2, (1 - g1 - g2), b1, b2, (1 - b1 - b2)] # This is the optimised Color Matrix (preserving greys by summing rows up to 1) Cam.log += str(optimised_ccm) Cam.log += "\n Old Color Correction Matrix Below \n" Cam.log += str(ccm) formatted_ccm = np.array(original_ccm).reshape((3, 3)) ''' below is a whole load of code that then applies the latest color matrix, and returns LAB values for color. This can then be used to calculate the final delta E ''' optimised_ccm_rgb = [] # Original Color Corrected Matrix RGB / LAB optimised_ccm_lab = [] formatted_optimised_ccm = np.array(optimised_ccm).reshape((3, 3)) after_gamma_rgb = [] after_gamma_lab = [] for RGB in zip(r, g, b): ccm_applied_rgb = np.dot(formatted_ccm, (np.array(RGB) / 256)) optimised_ccm_rgb.append(gamma(ccm_applied_rgb)) optimised_ccm_lab.append(colors.RGB_to_LAB(ccm_applied_rgb)) optimised_ccm_applied_rgb = np.dot(formatted_optimised_ccm, np.array(RGB) / 256) after_gamma_rgb.append(gamma(optimised_ccm_applied_rgb)) after_gamma_lab.append(colors.RGB_to_LAB(optimised_ccm_applied_rgb)) ''' Gamma After RGB / LAB - not used in calculations, only used for visualisation We now want to spit out some data that shows how the optimisation has improved the color matrices ''' Cam.log += "Here are the Improvements" # CALCULATE WORST CASE delta e old_worst_delta_e = 0 before_average = transform_and_evaluate(formatted_ccm, r, g, b, m_lab) new_worst_delta_e = 0 after_average = transform_and_evaluate(formatted_optimised_ccm, r, g, b, m_lab) for i in range(24): old_delta_e = deltae(optimised_ccm_lab[i], m_lab[i]) # Current Old Delta E new_delta_e = deltae(after_gamma_lab[i], m_lab[i]) # Current New Delta E if old_delta_e > old_worst_delta_e: old_worst_delta_e = old_delta_e if new_delta_e > new_worst_delta_e: new_worst_delta_e = new_delta_e Cam.log += "Before color correction matrix was optimised, we got an average delta E of " + str(before_average) + " and a maximum delta E of " + str(old_worst_delta_e) Cam.log += "After color correction matrix was optimised, we got an average delta E of " + str(after_average) + " and a maximum delta E of " + str(new_worst_delta_e) visualise_macbeth_chart(m_rgb, optimised_ccm_rgb, after_gamma_rgb, str(Img.col) + str(matrix_selection_types[typenum])) ''' The program will also save some visualisations of improvements. Very pretty to look at. Top rectangle is ideal, Left square is before optimisation, right square is after. ''' """ if a ccm has already been calculated for that temperature then don't overwrite but save both. They will then be averaged later on """ # Now going to use optimised color matrix, optimised_ccm if Img.col in ccm_tab.keys(): ccm_tab[Img.col].append(optimised_ccm) else: ccm_tab[Img.col] = [optimised_ccm] Cam.log += '\n' Cam.log += '\nFinished processing images' """ average any ccms that share a colour temperature """ for k, v in ccm_tab.items(): tab = np.mean(v, axis=0) tab = np.where((10000 * tab) % 1 <= 0.05, tab + 0.00001, tab) tab = np.where((10000 * tab) % 1 >= 0.95, tab - 0.00001, tab) ccm_tab[k] = list(np.round(tab, 5)) Cam.log += '\nMatrix calculated for colour temperature of {} K'.format(k) """ return all ccms with respective colour temperature in the correct format, sorted by their colour temperature """ sorted_ccms = sorted(ccm_tab.items(), key=lambda kv: kv[0]) ccms = [] for i in sorted_ccms: ccms.append({ 'ct': i[0], 'ccm': i[1] }) return ccms def guess(x0, r, g, b, m_lab): # provides a method of numerical feedback for the optimisation code [r1, r2, g1, g2, b1, b2] = x0 ccm = np.array([r1, r2, (1 - r1 - r2), g1, g2, (1 - g1 - g2), b1, b2, (1 - b1 - b2)]).reshape((3, 3)) # format the matrix correctly return transform_and_evaluate(ccm, r, g, b, m_lab) def transform_and_evaluate(ccm, r, g, b, m_lab): # Transforms colors to LAB and applies the correction matrix # create list of matrix changed colors realrgb = [] for RGB in zip(r, g, b): rgb_post_ccm = np.dot(ccm, np.array(RGB) / 256) # This is RGB values after the color correction matrix has been applied realrgb.append(colors.RGB_to_LAB(rgb_post_ccm)) # now compare that with m_lab and return numeric result, averaged for each patch return (sumde(realrgb, m_lab) / 24) # returns an average result of delta E def sumde(listA, listB): global typenum, test_patches sumde = 0 maxde = 0 patchde = [] # Create array of the delta E values for each patch. useful for optimisation of certain patches for listA_item, listB_item in zip(listA, listB): if maxde < (deltae(listA_item, listB_item)): maxde = deltae(listA_item, listB_item) patchde.append(deltae(listA_item, listB_item)) sumde += deltae(listA_item, listB_item) ''' The different options specified at the start allow for the maximum to be returned, average or specific patches ''' if typenum == 0: return sumde if typenum == 1: return maxde if typenum == 2: output = sum([patchde[test_patch] for test_patch in test_patches]) # Selects only certain patches and returns the output for them return output """ calculates the ccm for an individual image. ccms are calculated in rgb space, and are fit by hand. Although it is a 3x3 matrix, each row must add up to 1 in order to conserve greyness, simplifying calculation. The initial CCM is calculated in RGB, and then optimised in LAB color space This simplifies the initial calculation but then gets us the accuracy of using LAB color space. """ def do_ccm(r, g, b, m_srgb): rb = r-b gb = g-b rb_2s = (rb * rb) rb_gbs = (rb * gb) gb_2s = (gb * gb) r_rbs = rb * (m_srgb[..., 0] - b) r_gbs = gb * (m_srgb[..., 0] - b) g_rbs = rb * (m_srgb[..., 1] - b) g_gbs = gb * (m_srgb[..., 1] - b) b_rbs = rb * (m_srgb[..., 2] - b) b_gbs = gb * (m_srgb[..., 2] - b) """ Obtain least squares fit """ rb_2 = np.sum(rb_2s) gb_2 = np.sum(gb_2s) rb_gb = np.sum(rb_gbs) r_rb = np.sum(r_rbs) r_gb = np.sum(r_gbs) g_rb = np.sum(g_rbs) g_gb = np.sum(g_gbs) b_rb = np.sum(b_rbs) b_gb = np.sum(b_gbs) det = rb_2 * gb_2 - rb_gb * rb_gb """ Raise error if matrix is singular... This shouldn't really happen with real data but if it does just take new pictures and try again, not much else to be done unfortunately... """ if det < 0.001: raise ArithmeticError r_a = (gb_2 * r_rb - rb_gb * r_gb) / det r_b = (rb_2 * r_gb - rb_gb * r_rb) / det """ Last row can be calculated by knowing the sum must be 1 """ r_c = 1 - r_a - r_b g_a = (gb_2 * g_rb - rb_gb * g_gb) / det g_b = (rb_2 * g_gb - rb_gb * g_rb) / det g_c = 1 - g_a - g_b b_a = (gb_2 * b_rb - rb_gb * b_gb) / det b_b = (rb_2 * b_gb - rb_gb * b_rb) / det b_c = 1 - b_a - b_b """ format ccm """ ccm = [r_a, r_b, r_c, g_a, g_b, g_c, b_a, b_b, b_c] return ccm def deltae(colorA, colorB): return ((colorA[0] - colorB[0]) ** 2 + (colorA[1] - colorB[1]) ** 2 + (colorA[2] - colorB[2]) ** 2) ** 0.5 # return ((colorA[1]-colorB[1]) * * 2 + (colorA[2]-colorB[2]) * * 2) * * 0.5 # UNCOMMENT IF YOU WANT TO NEGLECT LUMINANCE FROM CALCULATION OF DELTA E 9'>349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141
/* SPDX-License-Identifier: GPL-2.0-or-later */
/*
 * Copyright (C) 2019, Google Inc.
 *
 * options.cpp - cam - Options parsing
 */

#include <assert.h>
#include <getopt.h>
#include <iomanip>
#include <iostream>
#include <string.h>

#include "options.h"

/**
 * \enum OptionArgument
 * \brief Indicate if an option takes an argument
 *
 * \var OptionArgument::ArgumentNone
 * \brief The option doesn't accept any argument
 *
 * \var OptionArgument::ArgumentRequired
 * \brief The option requires an argument
 *
 * \var OptionArgument::ArgumentOptional
 * \brief The option accepts an optional argument
 */

/**
 * \enum OptionType
 * \brief The type of argument for an option
 *
 * \var OptionType::OptionNone
 * \brief No argument type, used for options that take no argument
 *
 * \var OptionType::OptionInteger
 * \brief Integer argument type, with an optional base prefix (`0` for base 8,
 * `0x` for base 16, none for base 10)
 *
 * \var OptionType::OptionString
 * \brief String argument
 *
 * \var OptionType::OptionKeyValue
 * \brief key=value list argument
 */

/* -----------------------------------------------------------------------------
 * Option
 */

/**
 * \struct Option
 * \brief Store metadata about an option
 *
 * \var Option::opt
 * \brief The option identifier
 *
 * \var Option::type
 * \brief The type of the option argument
 *
 * \var Option::name
 * \brief The option name
 *
 * \var Option::argument
 * \brief Whether the option accepts an optional argument, a mandatory
 * argument, or no argument at all
 *
 * \var Option::argumentName
 * \brief The argument name used in the help text
 *
 * \var Option::help
 * \brief The help text (may be a multi-line string)
 *
 * \var Option::keyValueParser
 * \brief For options of type OptionType::OptionKeyValue, the key-value parser
 * to parse the argument
 *
 * \var Option::isArray
 * \brief Whether the option can appear once or multiple times
 *
 * \var Option::parent
 * \brief The parent option
 *
 * \var Option::children
 * \brief List of child options, storing all options whose parent is this option
 *
 * \fn Option::hasShortOption()
 * \brief Tell if the option has a short option specifier (e.g. `-f`)
 * \return True if the option has a short option specifier, false otherwise
 *
 * \fn Option::hasLongOption()
 * \brief Tell if the option has a long option specifier (e.g. `--foo`)
 * \return True if the option has a long option specifier, false otherwise
 */
struct Option {
	int opt;
	OptionType type;
	const char *name;
	OptionArgument argument;
	const char *argumentName;
	const char *help;
	KeyValueParser *keyValueParser;
	bool isArray;
	Option *parent;
	std::list<Option> children;

	bool hasShortOption() const { return isalnum(opt); }
	bool hasLongOption() const { return name != nullptr; }
	const char *typeName() const;
	std::string optionName() const;
};

/**
 * \brief Retrieve a string describing the option type
 * \return A string describing the option type
 */
const char *Option::typeName() const
{
	switch (type) {
	case OptionNone:
		return "none";

	case OptionInteger:
		return "integer";

	case OptionString:
		return "string";

	case OptionKeyValue:
		return "key=value";
	}

	return "unknown";
}

/**
 * \brief Retrieve a string describing the option name, with leading dashes
 * \return A string describing the option name, as a long option identifier
 * (double dash) if the option has a name, or a short option identifier (single
 * dash) otherwise
 */
std::string Option::optionName() const
{
	if (name)
		return "--" + std::string(name);
	else
		return "-" + std::string(1, opt);
}

/* -----------------------------------------------------------------------------
 * OptionBase<T>
 */

/**
 * \class template<typename T> OptionBase
 * \brief Container to store the values of parsed options
 * \tparam T The type through which options are identified
 *
 * The OptionsBase class is generated by a parser (either OptionsParser or
 * KeyValueParser) when parsing options. It stores values for all the options
 * found, and exposes accessor functions to retrieve them. The options are
 * accessed through an identifier to type \a T, which is an int referencing an
 * Option::opt for OptionsParser, or a std::string referencing an Option::name
 * for KeyValueParser.
 */

/**
 * \fn OptionsBase::OptionsBase()
 * \brief Construct an OptionsBase instance
 *
 * The constructed instance is initially invalid, and will be populated by the
 * options parser.
 */

/**
 * \brief Tell if the stored options list is empty
 * \return True if the container is empty, false otherwise
 */
template<typename T>
bool OptionsBase<T>::empty() const
{
	return values_.empty();
}

/**
 * \brief Tell if the options parsing completed successfully
 * \return True if the container is returned after successfully parsing
 * options, false if it is returned after an error was detected during parsing
 */
template<typename T>
bool OptionsBase<T>::valid() const
{
	return valid_;
}

/**
 * \brief Tell if the option \a opt is specified
 * \param[in] opt The option to search for
 * \return True if the \a opt option is set, false otherwise
 */
template<typename T>
bool OptionsBase<T>::isSet(const T &opt) const
{
	return values_.find(opt) != values_.end();
}

/**
 * \brief Retrieve the value of option \a opt
 * \param[in] opt The option to retrieve
 * \return The value of option \a opt if found, an empty OptionValue otherwise
 */
template<typename T>
const OptionValue &OptionsBase<T>::operator[](const T &opt) const
{
	static const OptionValue empty;

	auto it = values_.find(opt);
	if (it != values_.end())
		return it->second;
	return empty;
}

/**
 * \brief Mark the container as invalid
 *
 * This function can be used in a key-value parser's override of the
 * KeyValueParser::parse() function to mark the returned options as invalid if
 * a validation error occurs.
 */
template<typename T>
void OptionsBase<T>::invalidate()
{
	valid_ = false;
}

template<typename T>
bool OptionsBase<T>::parseValue(const T &opt, const Option &option,
				const char *arg)
{
	OptionValue value;

	switch (option.type) {
	case OptionNone:
		break;

	case OptionInteger:
		unsigned int integer;

		if (arg) {
			char *endptr;
			integer = strtoul(arg, &endptr, 0);
			if (*endptr != '\0')
				return false;
		} else {
			integer = 0;
		}

		value = OptionValue(integer);
		break;

	case OptionString:
		value = OptionValue(arg ? arg : "");
		break;

	case OptionKeyValue:
		KeyValueParser *kvParser = option.keyValueParser;
		KeyValueParser::Options keyValues = kvParser->parse(arg);
		if (!keyValues.valid())
			return false;

		value = OptionValue(keyValues);
		break;
	}

	if (option.isArray)
		values_[opt].addValue(value);
	else
		values_[opt] = value;

	return true;
}

template class OptionsBase<int>;
template class OptionsBase<std::string>;

/* -----------------------------------------------------------------------------
 * KeyValueParser
 */

/**
 * \class KeyValueParser
 * \brief A specialized parser for list of key-value pairs
 *
 * The KeyValueParser is an options parser for comma-separated lists of
 * `key=value` pairs. The supported keys are added to the parser with
 * addOption(). A given key can only appear once in the parsed list.
 *
 * Instances of this class can be passed to the OptionsParser::addOption()
 * function to create options that take key-value pairs as an option argument.
 * Specialized versions of the key-value parser can be created by inheriting
 * from this class, to pre-build the options list in the constructor, and to add
 * custom validation by overriding the parse() function.
 */

/**
 * \class KeyValueParser::Options
 * \brief An option list generated by the key-value parser
 *
 * This is a specialization of OptionsBase with the option reference type set to
 * std::string.
 */

KeyValueParser::KeyValueParser() = default;
KeyValueParser::~KeyValueParser() = default;

/**
 * \brief Add a supported option to the parser
 * \param[in] name The option name, corresponding to the key name in the
 * key=value pair. The name shall be unique.
 * \param[in] type The type of the value in the key=value pair
 * \param[in] help The help text
 * \param[in] argument Whether the value is optional, mandatory or not allowed.
 * Shall be ArgumentNone if \a type is OptionNone.
 *
 * \sa OptionsParser
 *
 * \return True if the option was added successfully, false if an error
 * occurred.
 */
bool KeyValueParser::addOption(const char *name, OptionType type,
			       const char *help, OptionArgument argument)
{
	if (!name)
		return false;
	if (!help || help[0] == '\0')
		return false;
	if (argument != ArgumentNone && type == OptionNone)
		return false;

	/* Reject duplicate options. */
	if (optionsMap_.find(name) != optionsMap_.end())
		return false;

	optionsMap_[name] = Option({ 0, type, name, argument, nullptr,
				     help, nullptr, false, nullptr, {} });
	return true;
}

/**
 * \brief Parse a string containing a list of key-value pairs
 * \param[in] arguments The key-value pairs string to parse
 *
 * If a parsing error occurs, the parsing stops and the function returns an
 * invalid container. The container is populated with the options successfully
 * parsed so far.
 *
 * \return A valid container with the list of parsed options on success, or an
 * invalid container otherwise
 */
KeyValueParser::Options KeyValueParser::parse(const char *arguments)
{
	Options options;

	for (const char *pair = arguments; *arguments != '\0'; pair = arguments) {
		const char *comma = strchrnul(arguments, ',');
		size_t len = comma - pair;

		/* Skip over the comma. */
		arguments = *comma == ',' ? comma + 1 : comma;

		/* Skip to the next pair if the pair is empty. */
		if (!len)
			continue;

		std::string key;
		std::string value;

		const char *separator = static_cast<const char *>(memchr(pair, '=', len));
		if (!separator) {
			key = std::string(pair, len);
			value = "";
		} else {
			key = std::string(pair, separator - pair);
			value = std::string(separator + 1, comma - separator - 1);
		}

		/* The key is mandatory, the value might be optional. */
		if (key.empty())
			continue;

		if (optionsMap_.find(key) == optionsMap_.end()) {
			std::cerr << "Invalid option " << key << std::endl;
			return options;
		}

		OptionArgument arg = optionsMap_[key].argument;
		if (value.empty() && arg == ArgumentRequired) {
			std::cerr << "Option " << key << " requires an argument"
				  << std::endl;
			return options;
		} else if (!value.empty() && arg == ArgumentNone) {
			std::cerr << "Option " << key << " takes no argument"
				  << std::endl;
			return options;
		}

		const Option &option = optionsMap_[key];
		if (!options.parseValue(key, option, value.c_str())) {
			std::cerr << "Failed to parse '" << value << "' as "
				  << option.typeName() << " for option " << key
				  << std::endl;
			return options;
		}
	}

	options.valid_ = true;
	return options;
}

unsigned int KeyValueParser::maxOptionLength() const
{
	unsigned int maxLength = 0;

	for (auto const &iter : optionsMap_) {
		const Option &option = iter.second;
		unsigned int length = 10 + strlen(option.name);
		if (option.argument != ArgumentNone)
			length += 1 + strlen(option.typeName());
		if (option.argument == ArgumentOptional)
			length += 2;

		if (length > maxLength)
			maxLength = length;
	}

	return maxLength;
}

void KeyValueParser::usage(int indent)
{
	for (auto const &iter : optionsMap_) {
		const Option &option = iter.second;
		std::string argument = std::string("          ") + option.name;

		if (option.argument != ArgumentNone) {
			if (option.argument == ArgumentOptional)
				argument += "[=";
			else
				argument += "=";
			argument += option.typeName();
			if (option.argument == ArgumentOptional)
				argument += "]";
		}

		std::cerr << std::setw(indent) << argument;

		for (const char *help = option.help, *end = help; end;) {
			end = strchr(help, '\n');
			if (end) {
				std::cerr << std::string(help, end - help + 1);
				std::cerr << std::setw(indent) << " ";
				help = end + 1;
			} else {
				std::cerr << help << std::endl;
			}
		}
	}
}

/* -----------------------------------------------------------------------------
 * OptionValue
 */

/**
 * \class OptionValue
 * \brief Container to store the value of an option
 *
 * The OptionValue class is a variant-type container to store the value of an
 * option. It supports empty values, integers, strings, key-value lists, as well
 * as arrays of those types. For array values, all array elements shall have the
 * same type.
 *
 * OptionValue instances are organized in a tree-based structure that matches
 * the parent-child relationship of the options added to the parser. Children
 * are retrieved with the children() function, and are stored as an
 * OptionsBase<int>.
 */

/**
 * \enum OptionValue::ValueType
 * \brief The option value type
 *
 * \var OptionValue::ValueType::ValueNone
 * \brief Empty value
 *
 * \var OptionValue::ValueType::ValueInteger
 * \brief Integer value (int)
 *
 * \var OptionValue::ValueType::ValueString
 * \brief String value (std::string)
 *
 * \var OptionValue::ValueType::ValueKeyValue
 * \brief Key-value list value (KeyValueParser::Options)
 *
 * \var OptionValue::ValueType::ValueArray
 * \brief Array value
 */

/**
 * \brief Construct an empty OptionValue instance
 *
 * The value type is set to ValueType::ValueNone.
 */
OptionValue::OptionValue()
	: type_(ValueNone), integer_(0)
{
}

/**
 * \brief Construct an integer OptionValue instance
 * \param[in] value The integer value
 *
 * The value type is set to ValueType::ValueInteger.
 */
OptionValue::OptionValue(int value)
	: type_(ValueInteger), integer_(value)
{
}

/**
 * \brief Construct a string OptionValue instance
 * \param[in] value The string value
 *
 * The value type is set to ValueType::ValueString.
 */
OptionValue::OptionValue(const char *value)
	: type_(ValueString), integer_(0), string_(value)
{
}

/**
 * \brief Construct a string OptionValue instance
 * \param[in] value The string value
 *
 * The value type is set to ValueType::ValueString.
 */
OptionValue::OptionValue(const std::string &value)
	: type_(ValueString), integer_(0), string_(value)
{
}

/**
 * \brief Construct a key-value OptionValue instance
 * \param[in] value The key-value list
 *
 * The value type is set to ValueType::ValueKeyValue.
 */
OptionValue::OptionValue(const KeyValueParser::Options &value)
	: type_(ValueKeyValue), integer_(0), keyValues_(value)
{
}

/**
 * \brief Add an entry to an array value
 * \param[in] value The entry value
 *
 * This function can only be called if the OptionValue type is
 * ValueType::ValueNone or ValueType::ValueArray. Upon return, the type will be
 * set to ValueType::ValueArray.
 */
void OptionValue::addValue(const OptionValue &value)
{
	assert(type_ == ValueNone || type_ == ValueArray);

	type_ = ValueArray;
	array_.push_back(value);
}

/**
 * \fn OptionValue::type()
 * \brief Retrieve the value type
 * \return The value type
 */

/**
 * \fn OptionValue::empty()
 * \brief Check if the value is empty
 * \return True if the value is empty (type set to ValueType::ValueNone), or
 * false otherwise
 */

/**
 * \brief Cast the value to an int
 * \return The option value as an int, or 0 if the value type isn't
 * ValueType::ValueInteger
 */
OptionValue::operator int() const
{
	return toInteger();
}

/**
 * \brief Cast the value to a std::string
 * \return The option value as an std::string, or an empty string if the value
 * type isn't ValueType::ValueString
 */
OptionValue::operator std::string() const
{
	return toString();
}

/**
 * \brief Retrieve the value as an int
 * \return The option value as an int, or 0 if the value type isn't
 * ValueType::ValueInteger
 */
int OptionValue::toInteger() const
{
	if (type_ != ValueInteger)
		return 0;

	return integer_;
}

/**
 * \brief Retrieve the value as a std::string
 * \return The option value as a std::string, or an empty string if the value
 * type isn't ValueType::ValueString
 */
std::string OptionValue::toString() const
{
	if (type_ != ValueString)
		return std::string();

	return string_;
}

/**
 * \brief Retrieve the value as a key-value list
 *
 * The behaviour is undefined if the value type isn't ValueType::ValueKeyValue.
 *
 * \return The option value as a KeyValueParser::Options
 */
const KeyValueParser::Options &OptionValue::toKeyValues() const
{
	assert(type_ == ValueKeyValue);
	return keyValues_;
}

/**
 * \brief Retrieve the value as an array
 *
 * The behaviour is undefined if the value type isn't ValueType::ValueArray.
 *
 * \return The option value as a std::vector of OptionValue
 */
const std::vector<OptionValue> &OptionValue::toArray() const
{
	assert(type_ == ValueArray);
	return array_;
}

/**
 * \brief Retrieve the list of child values
 * \return The list of child values
 */
const OptionsParser::Options &OptionValue::children() const
{
	return children_;
}

/* -----------------------------------------------------------------------------
 * OptionsParser
 */

/**
 * \class OptionsParser
 * \brief A command line options parser
 *
 * The OptionsParser class is an easy to use options parser for POSIX-style
 * command line options. Supports short (e.g. `-f`) and long (e.g. `--foo`)
 * options, optional and mandatory arguments, automatic parsing arguments for
 * integer types and comma-separated list of key=value pairs, and multi-value
 * arguments. It handles help text generation automatically.
 *
 * An OptionsParser instance is initialized by adding supported options with
 * addOption(). Options are specified by an identifier and a name. If the
 * identifier is an alphanumeric character, it will be used by the parser as a
 * short option identifier (e.g. `-f`). The name, if specified, will be used as
 * a long option identifier (e.g. `--foo`). It should not include the double
 * dashes. The name is optional if the option identifier is an alphanumeric
 * character and mandatory otherwise.
 *
 * An option has a mandatory help text, which is used to print the full options
 * list with the usage() function. The help text may be a multi-line string.
 * Correct indentation of the help text is handled automatically.
 *
 * Options accept arguments when created with OptionArgument::ArgumentRequired
 * or OptionArgument::ArgumentOptional. If the argument is required, it can be
 * specified as a positional argument after the option (e.g. `-f bar`,
 * `--foo bar`), collated with the short option (e.g. `-fbar`) or separated from
 * the long option by an equal sign (e.g. `--foo=bar`'). When the argument is
 * optional, it must be collated with the short option or separated from the
 * long option by an equal sign.
 *
 * If an option has a required or optional argument, an argument name must be
 * set when adding the option. The argument name is used in the help text as a
 * place holder for an argument value. For instance, a `--write` option that
 * takes a file name as an argument could set the argument name to `filename`,
 * and the help text would display `--write filename`. This is only used to
 * clarify the help text and has no effect on option parsing.
 *
 * The option type tells the parser how to process the argument. Arguments for
 * string options (OptionType::OptionString) are stored as-is without any
 * processing. Arguments for integer options (OptionType::OptionInteger) are
 * converted to an integer value, using an optional base prefix (`0` for base 8,
 * `0x` for base 16, none for base 10). Arguments for key-value options are
 * parsed by a KeyValueParser given to addOption().
 *
 * By default, a given option can appear once only in the parsed command line.
 * If the option is created as an array option, the parser will accept multiple
 * instances of the option. The order in which identical options are specified
 * is preserved in the values of an array option.
 *
 * After preparing the parser, it can be used any number of times to parse
 * command line options with the parse() function. The function returns an
 * Options instance that stores the values for the parsed options. The
 * Options::isSet() function can be used to test if an option has been found,
 * and is the only way to access options that take no argument (specified by
 * OptionType::OptionNone and OptionArgument::ArgumentNone). For options that
 * accept an argument, the option value can be access by Options::operator[]()
 * using the option identifier as the key. The order in which different options
 * are specified on the command line isn't preserved.
 *
 * Options can be created with parent-child relationships to organize them as a
 * tree instead of a flat list. When parsing a command line, the child options
 * are considered related to the parent option that precedes them. This is
 * useful when the parent is an array option. The Options values list generated
 * by the parser then turns into a tree, which each parent value storing the
 * values of child options that follow that instance of the parent option.
 * For instance, with a `capture` option specified as a child of a `camera`
 * array option, parsing the command line
 *
 * `--camera 1 --capture=10 --camera 2 --capture=20`
 *
 * will return an Options instance containing a single OptionValue instance of
 * array type, for the `camera` option. The OptionValue will contain two
 * entries, with the first entry containing the integer value 1 and the second
 * entry the integer value 2. Each of those entries will in turn store an
 * Options instance that contains the respective children. The first entry will
 * store in its children a `capture` option of value 10, and the second entry a
 * `capture` option of value 20.
 *
 * The command line
 *
 * `--capture=10 --camera 1`
 *
 * would result in a parsing error, as the `capture` option has no preceding
 * `camera` option on the command line.
 */

/**
 * \class OptionsParser::Options
 * \brief An option list generated by the options parser
 *
 * This is a specialization of OptionsBase with the option reference type set to
 * int.
 */

OptionsParser::OptionsParser() = default;
OptionsParser::~OptionsParser() = default;

/**
 * \brief Add an option to the parser
 * \param[in] opt The option identifier
 * \param[in] type The type of the option argument
 * \param[in] help The help text (may be a multi-line string)
 * \param[in] name The option name
 * \param[in] argument Whether the option accepts an optional argument, a
 * mandatory argument, or no argument at all
 * \param[in] argumentName The argument name used in the help text
 * \param[in] array Whether the option can appear once or multiple times
 * \param[in] parent The identifier of the parent option (optional)
 *
 * \return True if the option was added successfully, false if an error
 * occurred.
 */
bool OptionsParser::addOption(int opt, OptionType type, const char *help,
			      const char *name, OptionArgument argument,
			      const char *argumentName, bool array, int parent)
{
	/*
	 * Options must have at least a short or long name, and a text message.
	 * If an argument is accepted, it must be described by argumentName.
	 */
	if (!isalnum(opt) && !name)
		return false;
	if (!help || help[0] == '\0')
		return false;
	if (argument != ArgumentNone && !argumentName)
		return false;

	/* Reject duplicate options. */
	if (optionsMap_.find(opt) != optionsMap_.end())
		return false;

	/*
	 * If a parent is specified, create the option as a child of its parent.
	 * Otherwise, create it in the parser's options list.
	 */
	Option *option;

	if (parent) {
		auto iter = optionsMap_.find(parent);
		if (iter == optionsMap_.end())
			return false;

		Option *parentOpt = iter->second;
		parentOpt->children.push_back({
			opt, type, name, argument, argumentName, help, nullptr,
			array, parentOpt, {}
		});
		option = &parentOpt->children.back();
	} else {
		options_.push_back({ opt, type, name, argument, argumentName,
				     help, nullptr, array, nullptr, {} });
		option = &options_.back();
	}

	optionsMap_[opt] = option;

	return true;
}

/**
 * \brief Add a key-value pair option to the parser
 * \param[in] opt The option identifier
 * \param[in] parser The KeyValueParser for the option value
 * \param[in] help The help text (may be a multi-line string)
 * \param[in] name The option name
 * \param[in] array Whether the option can appear once or multiple times
 *
 * \sa Option
 *
 * \return True if the option was added successfully, false if an error
 * occurred.
 */
bool OptionsParser::addOption(int opt, KeyValueParser *parser, const char *help,
			      const char *name, bool array, int parent)
{
	if (!addOption(opt, OptionKeyValue, help, name, ArgumentRequired,
		       "key=value[,key=value,...]", array, parent))
		return false;

	optionsMap_[opt]->keyValueParser = parser;
	return true;
}

/**
 * \brief Parse command line arguments
 * \param[in] argc The number of arguments in the \a argv array
 * \param[in] argv The array of arguments
 *
 * If a parsing error occurs, the parsing stops, the function prints an error
 * message that identifies the invalid argument, prints usage information with
 * usage(), and returns an invalid container. The container is populated with
 * the options successfully parsed so far.
 *
 * \return A valid container with the list of parsed options on success, or an
 * invalid container otherwise
 */
OptionsParser::Options OptionsParser::parse(int argc, char **argv)
{
	OptionsParser::Options options;

	/*
	 * Allocate short and long options arrays large enough to contain all
	 * options.
	 */
	char shortOptions[optionsMap_.size() * 3 + 2];
	struct option longOptions[optionsMap_.size() + 1];
	unsigned int ids = 0;
	unsigned int idl = 0;

	shortOptions[ids++] = ':';

	for (const auto [opt, option] : optionsMap_) {
		if (option->hasShortOption()) {
			shortOptions[ids++] = opt;
			if (option->argument != ArgumentNone)
				shortOptions[ids++] = ':';
			if (option->argument == ArgumentOptional)
				shortOptions[ids++] = ':';
		}

		if (option->hasLongOption()) {
			longOptions[idl].name = option->name;

			switch (option->argument) {
			case ArgumentNone:
				longOptions[idl].has_arg = no_argument;
				break;
			case ArgumentRequired:
				longOptions[idl].has_arg = required_argument;
				break;
			case ArgumentOptional:
				longOptions[idl].has_arg = optional_argument;
				break;
			}

			longOptions[idl].flag = 0;
			longOptions[idl].val = option->opt;
			idl++;
		}
	}

	shortOptions[ids] = '\0';
	memset(&longOptions[idl], 0, sizeof(longOptions[idl]));

	opterr = 0;

	while (true) {
		int c = getopt_long(argc, argv, shortOptions, longOptions, nullptr);

		if (c == -1)
			break;

		if (c == '?' || c == ':') {
			if (c == '?')
				std::cerr << "Invalid option ";
			else
				std::cerr << "Missing argument for option ";
			std::cerr << argv[optind - 1] << std::endl;

			usage();
			return options;
		}

		const Option &option = *optionsMap_[c];
		if (!parseValue(option, optarg, &options)) {
			usage();
			return options;
		}
	}

	if (optind < argc) {
		std::cerr << "Invalid non-option argument '" << argv[optind]
			  << "'" << std::endl;
		usage();
		return options;
	}

	options.valid_ = true;
	return options;
}

/**
 * \brief Print usage text to std::cerr
 *
 * The usage text list all the supported option with their arguments. It is
 * generated automatically from the options added to the parser. Caller of this
 * function may print additional usage information for the application before
 * the list of options.
 */
void OptionsParser::usage()
{
	unsigned int indent = 0;

	for (const auto &opt : optionsMap_) {
		const Option *option = opt.second;
		unsigned int length = 14;
		if (option->hasLongOption())
			length += 2 + strlen(option->name);
		if (option->argument != ArgumentNone)
			length += 1 + strlen(option->argumentName);
		if (option->argument == ArgumentOptional)
			length += 2;
		if (option->isArray)
			length += 4;

		if (length > indent)
			indent = length;

		if (option->keyValueParser) {
			length = option->keyValueParser->maxOptionLength();
			if (length > indent)
				indent = length;
		}
	}

	indent = (indent + 7) / 8 * 8;

	std::cerr << "Options:" << std::endl;

	std::ios_base::fmtflags f(std::cerr.flags());
	std::cerr << std::left;

	usageOptions(options_, indent);

	std::cerr.flags(f);
}

void OptionsParser::usageOptions(const std::list<Option> &options,
				 unsigned int indent)
{
	std::vector<const Option *> parentOptions;

	for (const Option &option : options) {
		std::string argument;
		if (option.hasShortOption())
			argument = std::string("  -")
				 + static_cast<char>(option.opt);
		else
			argument = "    ";

		if (option.hasLongOption()) {
			if (option.hasShortOption())
				argument += ", ";
			else
				argument += "  ";
			argument += std::string("--") + option.name;
		}

		if (option.argument != ArgumentNone) {
			if (option.argument == ArgumentOptional)
				argument += "[=";
			else
				argument += " ";
			argument += option.argumentName;
			if (option.argument == ArgumentOptional)
				argument += "]";
		}

		if (option.isArray)
			argument += " ...";

		std::cerr << std::setw(indent) << argument;

		for (const char *help = option.help, *end = help; end; ) {
			end = strchr(help, '\n');
			if (end) {
				std::cerr << std::string(help, end - help + 1);
				std::cerr << std::setw(indent) << " ";
				help = end + 1;
			} else {
				std::cerr << help << std::endl;
			}
		}

		if (option.keyValueParser)
			option.keyValueParser->usage(indent);

		if (!option.children.empty())
			parentOptions.push_back(&option);
	}

	if (parentOptions.empty())
		return;

	for (const Option *option : parentOptions) {
		std::cerr << std::endl << "Options valid in the context of "
			  << option->optionName() << ":" << std::endl;
		usageOptions(option->children, indent);
	}
}

std::tuple<OptionsParser::Options *, const Option *>
OptionsParser::childOption(const Option *parent, Options *options)
{
	/*
	 * The parent argument points to the parent of the leaf node Option,
	 * and the options argument to the root node of the Options tree. Use
	 * recursive calls to traverse the Option tree up to the root node while
	 * traversing the Options tree down to the leaf node:
	 */

	/*
	 * - If we have no parent, we've reached the root node of the Option
	 *   tree, the options argument is what we need.
	 */
	if (!parent)
		return { options, nullptr };

	/*
	 * - If the parent has a parent, use recursion to move one level up the
	 *   Option tree. This returns the Options corresponding to parent, or
	 *   nullptr if a suitable Options child isn't found.
	 */
	if (parent->parent) {
		const Option *error;
		std::tie(options, error) = childOption(parent->parent, options);

		/* Propagate the error all the way back up the call stack. */
		if (!error)
			return { options, error };
	}

	/*
	 * - The parent has no parent, we're now one level down the root.
	 *   Return the Options child corresponding to the parent. The child may
	 *   not exist if options are specified in an incorrect order.
	 */
	if (!options->isSet(parent->opt))
		return { nullptr, parent };

	/*
	 * If the child value is of array type, children are not stored in the
	 * value .children() list, but in the .children() of the value's array
	 * elements. Use the last array element in that case, as a child option
	 * relates to the last instance of its parent option.
	 */
	const OptionValue *value = &(*options)[parent->opt];
	if (value->type() == OptionValue::ValueArray)
		value = &value->toArray().back();

	return { const_cast<Options *>(&value->children()), nullptr };
}

bool OptionsParser::parseValue(const Option &option, const char *arg,
			       Options *options)
{
	const Option *error;

	std::tie(options, error) = childOption(option.parent, options);
	if (error) {
		std::cerr << "Option " << option.optionName() << " requires a "