#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""This program checks the accuracy of Lensfun's autoscaling algorithm.
Lensfun calculates the autoscaling by testing all four corner points and the
four edge centres.  For each of these 8 points, it calculates the scaling
necessary, and then takes the maximal value.

It calculates the scaling necessary by determining the distance of the point of
the corrected image edge which lies in the same direction as the corner point.
(Because these two points should lie in top of each other after the
autoscaling.)  Then, it divides the distance of the corner point by the
distance of the corrected edge point.  “Distance” means “distance from origin”.

The choice of those 8 points is the critical compromise here.  In most cases,
the maximal scaling is indeed necessary for one of these points.  For example,
a corrected pincushion distortion pulls the corners most closely to the centre,
so they need the maximal scaling.  On the other hand, a corrected barrel
distorion leaves the edge centres most closely to the origin, so their position
is critical.  In case of a nasty moustache distortion, the maximal scaling may
be inbetween, and in these cases, Lensfun's approach fails.

This program shows for which lenses it fails, and how badly.  For
non-rectilinear lenses, the projection is changed to rectilinear first.
Currenly (as of April 2015), there is nothing to worry about.  Only three
lenses fail, and all of them exhibit only very slim black areas after the
correction.  If one leaves the original projection, two fisheyes join the
pack.

The simple reason why the current approach works well is that the used
distortion polynomials usually have monotonic behaviour in the important
region, so the maximal values are at the borders of the region.  This is even
more true for projection changes.

For problematic lenses, this program prints the number of sampling points on
the longer edges necessary to get a decent autoscaling (< 1‰).  Lensfun itself
only uses 2.  If the problem of black areas increases, one can increase the
number of edge sampling points to the maximal value this program prints.

The only safe approach is to test *all* edge pixels.  Really.  This is most
wasteful in almost all cases, so Lensfun doesn't do it (for now).  On the other
hand, its complexity is O(size).  All other operations of Lensfun are O(size²).
"""

import math, glob, multiprocessing, argparse, os.path
from xml.etree import ElementTree
from scipy.optimize import brentq


parser = argparse.ArgumentParser(description="Check autoscaling algorithm.")
parser.add_argument("database_path", metavar="database_path", help="The full path to the Lensfun database.")
parser.add_argument("--original-geometry", dest="original_geometry", action="store_true",
                    help="Use original geometry of fisheyes rather than transforming to rectilinear.")
args = parser.parse_args()


class Calibration:

    def __init__(self, f_per_half_height, a=None, b=None, c=None, k1=None, type_="rectilinear", aspect_ratio=1.5):
        self.f_per_half_height, self.a, self.b, self.c, self.k1, self.type_, self.aspect_ratio = \
                f_per_half_height, a, b, c, k1, type_, aspect_ratio

    def de_ptlens(self, r_):
        try:
            return brentq(lambda r: r * (self.a * r**3 + self.b * r**2 + self.c * r + 1 - self.a - self.b - self.c) - r_, 0, 2.1)
        except ValueError:
            return float("inf")

    def de_poly3(self, r_):
        try:
            return brentq(lambda r: r * (self.k1 * r**2 + 1 - self.k1) - r_, 0, 2)
        except ValueError:
            return float("inf")

    def de_fisheye(self, r):
        angle = r / self.f_per_half_height
        if angle < math.pi / 2:
            return self.f_per_half_height * math.tan(angle)
        else:
            return float("inf")

    def de_equisolidangle(self, r):
        try:
            return self.f_per_half_height * math.tan(2 * math.asin(r / self.f_per_half_height / 2))
        except ValueError:
            return float("inf")

    def de_stereographic(self, r):
        angle = 2 * math.atan(r / self.f_per_half_height / 2)
        if angle < math.pi / 2:
            return self.f_per_half_height * math.tan(angle)
        else:
            return float("inf")

    def get_scaling(self, x, y):
        r_ = math.hypot(x, y)
        r = self.de_poly3(r_) if self.a is None else self.de_ptlens(r_)
        if self.type_ == "rectilinear" or args.original_geometry:
            pass
        elif self.type_ == "fisheye":
            r = self.de_fisheye(r)
        elif self.type_ == "equisolid":
            r = self.de_equisolidangle(r)
        elif self.type_ == "stereographic":
            r = self.de_stereographic(r)
        return r_ / r

    def get_autoscale(self, samples):
        scalings = []
        for i in range(samples):
            x = self.aspect_ratio * (-1 + i / samples * 2)
            scalings.append(self.get_scaling(x, +1))
        samples = int(math.ceil(samples / self.aspect_ratio / 2)) * 2
        for i in range(samples):
            y = -1 + i / samples * 2
            scalings.append(self.get_scaling(+ self.aspect_ratio, y))
        return max(scalings) or 1

    def get_perfect_sample_number(self):
        perfect_autoscale = self.get_autoscale(100)
        samples = 2
        scalings = []
        while True:
            scalings.append(self.get_autoscale(samples))
            if all(abs(scaling - perfect_autoscale) / perfect_autoscale < 1e-3 for scaling in scalings[-3:]):
                return samples
            samples += 2


def process_calibration(distortion, lens, crop_factor, type_, aspect_ratio):
    f_per_half_height = float(distortion.attrib.get("focal")) / (12 / crop_factor)
    model = distortion.attrib.get("model")
    if model == "ptlens":
        calibration = Calibration(f_per_half_height,
                                  a=float(distortion.attrib.get("a", 0)),
                                  b=float(distortion.attrib.get("b", 0)),
                                  c=float(distortion.attrib.get("c", 0)), type_=type_, aspect_ratio=aspect_ratio)
    elif model == "poly3":
        calibration = Calibration(f_per_half_height, k1=float(distortion.attrib.get("k1", 0)),
                                  type_=type_, aspect_ratio=aspect_ratio)
    else:
        return None, None, None
    return lens.find("model").text, distortion.attrib.get("focal"), calibration.get_perfect_sample_number()

results = set()
pool = multiprocessing.Pool()
for filename in glob.glob(os.path.join(args.database_path, "*.xml")):
    tree = ElementTree.parse(filename)
    for lens in tree.findall("lens"):
        type_element = lens.find("type")
        type_ = "rectilinear" if type_element is None else type_element.text.strip()
        crop_factor = float(lens.find("cropfactor").text)
        aspect_ratio_element = lens.find("aspect-ratio")
        if aspect_ratio_element is None:
            aspect_ratio = 1.5
        else:
            numerator, __, denominator = aspect_ratio_element.text.partition(":")
            if denominator:
                aspect_ratio = float(numerator) / float(denominator)
            else:
                aspect_ratio = float(numerator)
        calibration_element = lens.find("calibration")
        if calibration_element is not None:
            for distortion in calibration_element.findall("distortion"):
                results.add(pool.apply_async(process_calibration, (distortion, lens, crop_factor, type_, aspect_ratio)))
pool.close()
pool.join()

problematic_lenses = {}
for result in results:
    model, focal_length, perfect_sample_number = result.get()
    if perfect_sample_number and perfect_sample_number > 2:
        problematic_lenses.setdefault(perfect_sample_number, set()).add((model, focal_length))
if problematic_lenses:
    print("Number of equidistant samples on the long edge:\n")
    for perfect_sample_number, lenses in problematic_lenses.items():
        print(",\n".join("{} at {}mm".format(*lens) for lens in lenses), ": ", perfect_sample_number)