# Copyright 2024 University of Calgary
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import aacgmv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.path import Path
from ..classes.keogram import Keogram
from ...data.ucalgary import Skymap


# Helper function that returns all array indices within
# the polygon defined by input vertices.
def __indices_in_polygon(vertices, image_shape):
    """
    Function to obtain all indices of an array/image, within the 
    polygon defined by input list of ordered vertices.
    """
    # Create a Path object using the vertices
    path = Path(vertices)

    # Create a grid of points covering the image shape
    x, y = np.meshgrid(np.arange(image_shape[1]), np.arange(image_shape[0]))
    points = np.vstack((x.flatten(), y.flatten())).T

    # Create a mask indicating which points are inside the polygon
    mask = path.contains_points(points).reshape(image_shape)

    # Get all indices inside or on the boundary of the polygon
    indices_inside = np.argwhere(mask)

    return indices_inside


def __haversine_distances(target_lat, target_lon, lat_array, lon_array):
    """
    Computes the distance on the globe between target lat/lon,
    and all points defined by lat/lon arrays.
    """
    # Earth radius (meters)
    r = 6371000.0

    # Convert degrees to radians
    phi1 = np.radians(target_lat)
    phi2 = np.radians(lat_array)
    delta_phi = np.radians(lat_array - target_lat)
    delta_lambda = np.radians(lon_array - target_lon)

    # Haversine formula
    a = np.sin(delta_phi / 2.0)**2 + np.cos(phi1) * np.cos(phi2) * np.sin(delta_lambda / 2.0)**2
    c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1 - a))

    return r * c


# Helper function for handling magnetic lat/lon inputs
def __convert_latlon_to_ccd(lon_locs, lat_locs, timestamp, skymap: Skymap, altitude_km, magnetic):
    """
    Function for handling lat/lon inputs. Takes the lat/lon
    locations, and uses a skymap to convert those locations
    to CCD (image index) coordinates.
    """
    # Obtain lat/lon arrays from skymap
    if (altitude_km * 1000.0 in skymap.full_map_altitude):
        altitude_idx = np.where(altitude_km * 1000.0 == skymap.full_map_altitude)

        lats = np.squeeze(skymap.full_map_latitude[altitude_idx, :, :])
        lons = np.squeeze(skymap.full_map_longitude[altitude_idx, :, :])
        lons[np.where(lons > 180)] -= 360.0

    else:
        # Make sure altitude is in range that can be interpolated
        if (altitude_km * 1000.0 < skymap.full_map_altitude[0]) or (altitude_km * 1000.0 > skymap.full_map_altitude[2]):
            raise ValueError("Altitude " + str(altitude_km) + " outside valid range of " +
                             str((skymap.full_map_altitude[0] / 1000.0, skymap.full_map_altitude[2] / 1000.0)))

        # Initialize empty lat/lon arrays
        lats = np.full(np.squeeze(skymap.full_map_latitude[0, :, :]).shape, np.nan, dtype=skymap.full_map_latitude[0, :, :].dtype)
        lons = lats.copy()

        # Interpolate lats and lons at desired altitude
        for i in range(skymap.full_map_latitude.shape[1]):
            for j in range(skymap.full_map_latitude.shape[2]):
                lats[i, j] = np.interp(altitude_km * 1000.0, skymap.full_map_altitude, skymap.full_map_latitude[:, i, j])
                lons[i, j] = np.interp(altitude_km * 1000.0, skymap.full_map_altitude, skymap.full_map_longitude[:, i, j])

        lons[np.where(lons > 180)] -= 360.0  # Fix skymap to be in (-180,180) format

    # Convert skymap to magnetic coords if necessary
    if magnetic:
        mag_lats, mag_lons, _ = aacgmv2.convert_latlon_arr(
            lats.flatten(),
            lons.flatten(),
            (lons * 0.0).flatten(),
            timestamp[0],
            method_code="G2A",
        )
        lats = np.reshape(mag_lats, lats.shape)
        lons = np.reshape(mag_lons, lons.shape)

    # Iterate through each target point
    x_locs = []
    y_locs = []
    for target_lat, target_lon in zip(lat_locs, lon_locs):

        # Make sure lat/lon falls within skymap
        if target_lat < np.nanmin(lats) or target_lat > np.nanmax(lats):
            continue
            # raise ValueError(f"Latitude {target_lat} is outside this skymap's valid range of {(np.nanmin(lats),np.nanmax(lats))}.")
        if target_lon < np.nanmin(lons) or target_lon > np.nanmax(lons):
            continue
            # raise ValueError(f"Longitude {target_lon} is outside this skymap's valid range of {(np.nanmin(lons),np.nanmax(lons))}.")

        # Compute haversine distance between all points in skymap
        haversine_diff = __haversine_distances(target_lat, target_lon, lats, lons)

        # Obtain the skymap indices of the nearest point
        nan_mask = np.isnan(haversine_diff)
        masked_data = np.ma.masked_array(haversine_diff, mask=nan_mask)
        nearest_indices = np.unravel_index(np.ma.argmin(masked_data), haversine_diff.shape)

        # Convert indices to CCD Coordinates
        y_loc = nearest_indices[0] - 1
        if y_loc < 0:
            y_loc = 0
        x_loc = nearest_indices[1] - 1
        if x_loc < 0:
            x_loc = 0

        x_locs.append(x_loc)
        y_locs.append(y_loc)

    x_locs = np.array(x_locs)
    y_locs = np.array(y_locs)

    return x_locs, y_locs


def create_custom(images, timestamp, coordinate_system, width, x_locs, y_locs, preview, skymap, altitude_km, metric):
    # If using CCD coordinates we don't need a skymao or altitude
    if (coordinate_system == "ccd") and (skymap is not None or altitude_km is not None):
        raise ValueError("Conflict in passing in a skymap or altitude when working in CCD coordinates. This value will be ignored.")

    # convert any lists to np.arrays  and check shape
    x_locs = np.array(x_locs)
    y_locs = np.array(y_locs)

    # determine if we are single or 3 channel
    n_channels = 1
    if (len(images.shape) == 3):
        # single channel
        n_channels = 1
    elif (len(images.shape) == 4):
        # three channel
        n_channels = 3
    else:
        ValueError("Unable to determine number of channels based on the supplied images. Make sure you are supplying a " +
                   "[rows,cols,images] or [rows,cols,channels,images] sized array.")

    # Initialize empty keogram array
    keo_arr = np.squeeze(np.full((x_locs.shape[0] - 1, len(timestamp), n_channels), 0))

    if len(x_locs.shape) != 1:
        raise ValueError(f"X coordinates may not be multidimensional. Sequence passed with shape {x_locs.shape}")
    if len(y_locs.shape) != 1:
        raise ValueError(f"Y coordinates may not be multidimensional. Sequence passed with shape {y_locs.shape}")
    if len(x_locs.shape) != len(y_locs.shape):
        raise ValueError(f"X and Y coordinates must have same length. Sequences passed with shapes {x_locs.shape} and {y_locs.shape}")

    # Convert lat/lon coordinates to CCD
    if coordinate_system == "mag":
        if (skymap is None or altitude_km is None):
            raise ValueError("When magnetic coordinates, a Skymap object and Altitude must be passed in through the skymap argument.")
        x_locs, y_locs = __convert_latlon_to_ccd(x_locs, y_locs, timestamp, skymap, altitude_km, magnetic=True)
    elif coordinate_system == "geo":
        if (skymap is None or altitude_km is None):
            raise ValueError("When geographic coordinates, a Skymap object and Altitude must be passed in through the skymap argument.")
        x_locs, y_locs = __convert_latlon_to_ccd(x_locs, y_locs, timestamp, skymap, altitude_km, magnetic=False)

    # We will use the first image as a preview
    if n_channels == 1:
        preview_img = images.copy()[:, :, 0]
    else:
        preview_img = images.copy()[:, :, :, 0]

    # Now working in CCD Coordinates
    x_max = images.shape[1] - 1
    y_max = images.shape[0] - 1

    # Remove any points that are not within the image CCD
    parsed_x_locs = []
    parsed_y_locs = []
    for i in range(x_locs.shape[0]):
        x = x_locs[i]
        y = y_locs[i]

        if x < 0 or x > x_max:
            continue
        if y < 0 or y > y_max:
            continue

        parsed_x_locs.append(x)
        parsed_y_locs.append(y)
    x_locs = np.array(parsed_x_locs)
    y_locs = np.array(parsed_y_locs)

    # Make sure all supplied points are within image bounds # NOTE: This should be good to remove as it shouldn't hit, testing needed
    if len(np.where(np.logical_or(x_locs < 0, x_locs > x_max))[0]) != 0:
        raise ValueError("The following CCD coordinates passed in through x_locs are outside of the CCD image range: " +
                         str(x_locs[np.where(np.logical_or(x_locs < 0, x_locs > x_max))]))
    if len(np.where(np.logical_or(y_locs < 0, y_locs > y_max))[0]) != 0:
        raise ValueError("The following CCD coordinates passed in through y_locs are outside of the CCD image range: " +
                         str(y_locs[np.where(np.logical_or(y_locs < 0, y_locs > y_max))]))

    # Iterate points in pairs of two
    path_counter = 0
    for i in range(x_locs.shape[0] - 1):

        # Points of concern for this iteration
        x_0 = x_locs[i]
        x_1 = x_locs[i + 1]
        y_0 = y_locs[i]
        y_1 = y_locs[i + 1]

        # Compute the unit vector between the two points
        dx = x_1 - x_0
        dy = y_1 - y_0
        length = np.sqrt(dx**2 + dy**2)
        if length == 0:
            continue

        dx /= length
        dy /= length

        # Compute orthogonal unit vector
        perp_dx = -dy
        perp_dy = dx

        # Calculate (+/-) offsets for each perpendicular direction
        offset1_x = perp_dx * width / 2
        offset1_y = perp_dy * width / 2
        offset2_x = -perp_dx * width / 2
        offset2_y = -perp_dy * width / 2

        # Calculate vertices in correct order for this polygon
        vertex1 = (int(x_0 + offset1_x), int(y_0 + offset1_y))
        vertex2 = (int(x_1 + offset1_x), int(y_1 + offset1_y))
        vertex3 = (int(x_1 + offset2_x), int(y_1 + offset2_y))
        vertex4 = (int(x_0 + offset2_x), int(y_0 + offset2_y))

        # Append vertices in the correct order to form a closed polygon
        vertices = [vertex1, vertex2, vertex3, vertex4]

        # Obtain the indexes into the image of this polygon
        indices_inside = __indices_in_polygon(vertices, (images.shape[0], images.shape[1]))

        if np.any(np.array(indices_inside.shape) == 0):
            continue

        row_idx, col_idx = zip(*indices_inside)

        if n_channels == 1:
            # Update the preview image
            preview_img[row_idx, col_idx] = np.iinfo(preview_img.dtype).max

            # Extract metric from all images and add to keogram
            if metric == "median":
                # median metric
                pixel_keogram = np.median(images[row_idx, col_idx, :], axis=0)
            elif metric == "mean":
                # mean metric
                pixel_keogram = np.mean(images[row_idx, col_idx, :], axis=0)
            else:
                # sum metric
                pixel_keogram = np.sum(images[row_idx, col_idx, :], axis=0)
            keo_arr[i, :] = pixel_keogram
        elif n_channels == 3:
            # Update the preview image
            preview_img[row_idx, col_idx, :] = np.iinfo(preview_img.dtype).max

            # Extract metric from all images and add to keogram
            if metric == "median":
                r_pixel_keogram = np.floor(np.median(images[row_idx, col_idx, 0, :], axis=0))
                g_pixel_keogram = np.floor(np.median(images[row_idx, col_idx, 1, :], axis=0))
                b_pixel_keogram = np.floor(np.median(images[row_idx, col_idx, 2, :], axis=0))
            elif metric == "mean":
                r_pixel_keogram = np.floor(np.mean(images[row_idx, col_idx, 0, :], axis=0))
                g_pixel_keogram = np.floor(np.mean(images[row_idx, col_idx, 1, :], axis=0))
                b_pixel_keogram = np.floor(np.mean(images[row_idx, col_idx, 2, :], axis=0))
            else:
                # sum metric
                r_pixel_keogram = np.floor(np.sum(images[row_idx, col_idx, 0, :], axis=0))
                g_pixel_keogram = np.floor(np.sum(images[row_idx, col_idx, 1, :], axis=0))
                b_pixel_keogram = np.floor(np.sum(images[row_idx, col_idx, 2, :], axis=0))

            keo_arr[i, :, 0] = r_pixel_keogram
            keo_arr[i, :, 1] = g_pixel_keogram
            keo_arr[i, :, 2] = b_pixel_keogram

        path_counter += 1

    if path_counter == 0:
        raise ValueError("Could not form keogram path. First ensure that coordinates are within image range. Then " +
                         "try increasing 'width' or decreasing number of points in input coordinates.")

    # Create keogram object
    keo_obj = Keogram(data=keo_arr, timestamp=timestamp, instrument_type="asi")

    # show preview
    if preview:
        plt.figure()
        plt.imshow(preview_img, cmap="gray", origin="lower")
        plt.axis("off")
        plt.title("Keogram Domain Preview")
        plt.show()

    # return
    return keo_obj
