# SPDX-License-Identifier: BSD-2-Clause
#
# Copyright (C) 2023, Raspberry Pi Ltd
#
# ctt_cac.py - CAC (Chromatic Aberration Correction) tuning tool

from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm

from ctt_dots_locator import find_dots_locations


# This is the wrapper file that creates a JSON entry for you to append
# to your camera tuning file.
# It calculates the chromatic aberration at different points throughout
# the image and uses that to produce a martix that can then be used
# in the camera tuning files to correct this aberration.


def pprint_array(array):
    # Function to print the array in a tidier format
    array = array
    output = ""
    for i in range(len(array)):
        for j in range(len(array[0])):
            output += str(round(array[i, j], 2)) + ", "
        # Add the necessary indentation to the array
        output += "\n                   "
    # Cut off the end of the array (nicely formats it)
    return output[:-22]


def plot_shifts(red_shifts, blue_shifts):
    # If users want, they can pass a command line option to show the shifts on a graph
    # Can be useful to check that the functions are all working, and that the sample
    # images are doing the right thing
    Xs = np.array(red_shifts)[:, 0]
    Ys = np.array(red_shifts)[:, 1]
    Zs = np.array(red_shifts)[:, 2]
    Zs2 = np.array(red_shifts)[:, 3]
    Zs3 = np.array(blue_shifts)[:, 2]
    Zs4 = np.array(blue_shifts)[:, 3]

    fig, axs = plt.subplots(2, 2)
    ax = fig.add_subplot(2, 2, 1, projection='3d')
    ax.scatter(Xs, Ys, Zs, cmap=cm.jet, linewidth=0)
    ax.set_title('Red X Shift')
    ax = fig.add_subplot(2, 2, 2, projection='3d')
    ax.scatter(Xs, Ys, Zs2, cmap=cm.jet, linewidth=0)
    ax.set_title('Red Y Shift')
    ax = fig.add_subplot(2, 2, 3, projection='3d')
    ax.scatter(Xs, Ys, Zs3, cmap=cm.jet, linewidth=0)
    ax.set_title('Blue X Shift')
    ax = fig.add_subplot(2, 2, 4, projection='3d')
    ax.scatter(Xs, Ys, Zs4, cmap=cm.jet, linewidth=0)
    ax.set_title('Blue Y Shift')
    fig.tight_layout()
    plt.show()


def shifts_to_yaml(red_shift, blue_shift, image_dimensions, output_grid_size=9):
    # Convert the shifts to a numpy array for easier handling and initialise other variables
    red_shifts = np.array(red_shift)
    blue_shifts = np.array(blue_shift)
    # create a grid that's smaller than the output grid, which we then interpolate from to get the output values
    xrgrid = np.zeros((output_grid_size - 1, output_grid_size - 1))
    xbgrid = np.zeros((output_grid_size - 1, output_grid_size - 1))
    yrgrid = np.zeros((output_grid_size - 1, output_grid_size - 1))
    ybgrid = np.zeros((output_grid_size - 1, output_grid_size - 1))

    xrsgrid = []
    xbsgrid = []
    yrsgrid = []
    ybsgrid = []
    xg = np.zeros((output_grid_size - 1, output_grid_size - 1))
    yg = np.zeros((output_grid_size - 1, output_grid_size - 1))

    # Format the grids - numpy doesn't work for this, it wants a
    # nice uniformly spaced grid, which we don't know if we have yet, hence the rather mundane setup
    for x in range(output_grid_size - 1):
        xrsgrid.append([])
        yrsgrid.append([])
        xbsgrid.append([])
        ybsgrid.append([])
        for y in range(output_grid_size - 1):
            xrsgrid[x].append([])
            yrsgrid[x].append([])
            xbsgrid[x].append([])
            ybsgrid[x].append([])

    image_size = (image_dimensions[0], image_dimensions[1])
    gridxsize = image_size[0] / (output_grid_size - 1)
    gridysize = image_size[1] / (output_grid_size - 1)

    # Iterate through each dot, and it's shift values and put these into the correct grid location
    for red_shift in red_shifts:
        xgridloc = int(red_shift[0] / gridxsize)
        ygridloc = int(red_shift[1] / gridysize)
        xrsgrid[xgridloc][ygridloc].append(red_shift[2])
        yrsgrid[xgridloc][ygridloc].append(red_shift[3])

    for blue_shift in blue_shifts:
        xgridloc = int(blue_shift[0] / gridxsize)
        ygridloc = int(blue_shift[1] / gridysize)
        xbsgrid[xgridloc][ygridloc].append(blue_shift[2])
        ybsgrid[xgridloc][ygridloc].append(blue_shift[3])

    # Now calculate the average pixel shift for each square in the grid
    for x in range(output_grid_size - 1):
        for y in range(output_grid_size - 1):
            xrgrid[x, y] = np.mean(xrsgrid[x][y])
            yrgrid[x, y] = np.mean(yrsgrid[x][y])
            xbgrid[x, y] = np.mean(xbsgrid[x][y])
            ybgrid[x, y] = np.mean(ybsgrid[x][y])

    # Next, we start to interpolate the central points of the grid that gets passed to the tuning file
    input_grids = np.array([xrgrid, yrgrid, xbgrid, ybgrid])
    output_grids = np.zeros((4, output_grid_size, output_grid_size))

    # Interpolate the centre of the grid
    output_grids[:, 1:-1, 1:-1] = (input_grids[:, 1:, :-1] + input_grids[:, 1:, 1:] + input_grids[:, :-1, 1:] + input_grids[:, :-1, :-1]) / 4

    # Edge cases:
    output_grids[:, 1:-1, 0] = ((input_grids[:, :-1, 0] + input_grids[:, 1:, 0]) / 2 - output_grids[:, 1:-1, 1]) * 2 + output_grids[:, 1:-1, 1]