#!/usr/bin/env python3

from pathlib import Path
import subprocess
import numpy as np
import math
import click
import struct

from preprocess import Preprocessor

def read_edge_info(filename):
    """
    Read the (private) binary edge information produced by MTF Mapper

    Parameters
    ----------
    filename: string
        full path to "serialized_edges.bin"

    Returns
    -------
    header_dict
        A dictionary containing the global metadata pertaining to all edges. Noteworthy fields
            count: number of edges analyzed
            pixel_pitch: the --pixelsize option passed in the mtfmapper call
            mtf_contrast: the --mtf option passed in mtfmapper call
            meta_bayer: the --bayer option passed in mtfmaper call (integer enum)
            meta-channels: the number of colour channels detected in the input image
            img_width: input image width in pixels
            img_height: input image height in pixels
            nodata: special constant to indicate that a field does not have a value, e.g., 
                the ca_red field could be equal to header_dict['nodata'] if "--ca" was not
                passed in mtfmapper call
            x_scale: the spatial frequencies corresponding to the data points in the extracted SFR
                see 'sfr' field of 'edges' below; note that the units are always specified in cycles
                per pixel, even if the --pixelsize option was passed in the mtfmapper call, so
                if you need lp/mm you will have to perform your own conversion

    edges
        A list of dictionaries containing the MTF information extracted for each edge. 
        Noteworhy fields in each element of the list
            centroid_x: x-axis pixel location of this edge
            centroid_y: y-axis pixel location of this edge
            angle: edge orientation angle in radians
            mtf50: the MTF50 value of this edge, i.e., spatial frequency at which SFR contrast is 50%
            cnr: contrast-to-noise-ratio, indication of image noise which affects measurement accuracy
                try to keep cnr above 30, but generally higher is better
            oversampling: the oversampling factor of this edge, related to edge orientation angle relative
                to image sensor axes; only keep edges with oversampling >= 4 for best results
            ca_red: chromatic aberration between red and green channels, only if "--ca" was passed
                in mtfmapper call, otherwise value will be set to header_dict['nodata']
            ca_blue: chromatic aberration between blue and green channels, only if "--ca" was passed
                in mtfmapper call, otherwise value will be set to header_dict['nodata']
            edge_length: length of slanted edge that was analyzed in pixels
                recommendation is to discard edges with edge_length < 25 pixels (need bigger test chart targets)
            geometric_edge_length: similar to edge_length, but measures actual corner-to-corner length of the
                edge as opposed to the length of the edge included in the slanted-edge analysis (edge_length)
            sfr: spatial frequency response, an array of contrast values, nominally in range 
                [0.0,1.0] (but can exceed 1.0), representing contrast at each spatial frequency in 'x_scale'
                field of 'header_dict'; use zip(header_dict['x_scale'], edges[i]['sfr']) to obtain complete
                SFR

    """
    edge_file = open(filename, "rb")
    edge_bin = edge_file.read()
    header_field_names = ("count", "pixel_pitch", "mtf_contrast", "meta_bayer", "meta_channels", "img_width", "img_height")
    header_format = "=IddIIII"
    raw_header = struct.unpack(header_format, edge_bin[0:struct.calcsize(header_format)])
    offset = struct.calcsize(header_format)
    header_dict = dict(zip(header_field_names, raw_header))
    header_dict['nodata'] = -1073741824.0
    edges = []
    for edge in range(header_dict["count"]):
        entry_field_names = (\
            "centroid_x", "centroid_y", "angle", "mtf50", "cnr", "oversampling", \
            "ca_red", "ca_blue", "edge_length", "geometric_edge_length" \
        )
        record_format = "=10d"
        record_size = struct.calcsize(record_format)
        raw_values = struct.unpack(record_format, edge_bin[offset:(offset + record_size)])
        dict_values = dict(zip (entry_field_names, raw_values))
        offset += record_size

        sfr_count = struct.unpack("=I", edge_bin[offset:(offset + 4)])[0]
        offset += 4

        fmt = "=" + str(sfr_count) + "f"
        fmt_size = struct.calcsize(fmt)
        sfr_values = struct.unpack(fmt, edge_bin[offset:(offset + fmt_size)])
        offset += fmt_size
        dict_values["sfr"] = np.array(sfr_values)

        # skip over ESF data for now
        esf_count = struct.unpack("=I", edge_bin[offset:(offset + 4)])[0]
        offset += 4
        offset += struct.calcsize("f") * esf_count

        edges.append(dict_values)

    if len(edges[0]["sfr"]) > 64:
        header_dict["x_scale"] = np.linspace(0.0, 2.0, num=128, endpoint=False)
    else:
        header_dict["x_scale"] = np.linspace(0.0, 1.0, num=64, endpoint=False)

    return header_dict, edges

def reduce_angle(angle_radians):
    quad1 = abs(divmod(angle_radians, math.pi/2.0)[1])
    if quad1 > math.pi/4.0:
        quad1 = math.pi/2.0 - quad1
    quad1 = quad1 / math.pi * 180
    return quad1

def mean_mtf(header, edges):
    """
    An example of how to process or aggregate the whole list of MTF measurements produced by MTF Mapper
    In this case, we simply compute the mean of all MTF50 values that pass some recommended quality checks,
    but you could filter values based on location within the image (e.g., only the edges near center of the image)
    """
    mtf50s = []
    for edge in edges:
        # basic quality checks to ensure only more reliable measurements are considered
        if edge["edge_length"] >= 25 and edge["oversampling"] >= 4 and edge["cnr"] > 2 and reduce_angle(edge["angle"]) > 0.5 and edge["mtf50"] < 0.99999:
            mtf50s.append(edge["mtf50"])
    
    return np.mean(mtf50s), len(mtf50s)

def mtfmapper_call(fname, threshold, extra_args="", verbose=False):
    out_dir = Path("/tmp/output")
    out_dir.mkdir(parents=True, exist_ok=True)
    args = ["/usr/local/bin/mtf_mapper", "-q", "-v", "2",
        # additional mtf_mapper arguments can be added here
        fname, str(out_dir)]
    args.extend(extra_args.split())

    cp = subprocess.run(args, capture_output=(not verbose))
    if cp.returncode == 0:
        header, edges = read_edge_info(str(out_dir / "serialized_edges.bin"))
        return header, edges
    
    return None, None

def run_mtfmapper(input, verbose, args):
    try:
        pp = Preprocessor(input, verbose)
        _, pp_input = pp.process()

        if not Path(pp_input).is_file():
            raise Exception(f"file {pp_input} not found")

        header, edges = mtfmapper_call(pp_input, args, verbose=verbose)
        if header is None:
            raise Exception(f"MTF Mapper found no edges")
        
        agg_mtf, nedges = mean_mtf(header, edges)
        print(f"mean mtf = {agg_mtf:.4f} over {nedges} edges")
    except Exception as e:
        print(f"exception: {e}")
        
@click.command()
@click.option('--input', required=True, help="Specify the path to the input image")
@click.option('--verbose', '-v', is_flag=True, help="Show the MTF Mapper output")
@click.option('--args', default="", help="Additional arguments to pass to MTF Mapper, quoted as '-t 0.3 --bayer green', for example")
def main(input, verbose, args):
    run_mtfmapper(input, verbose, args)

if __name__ == "__main__":
    main()
