import logging
import os
from typing import List, Optional, Tuple

import brainglobe_space as bgs
import numpy as np
import pandas as pd
import tifffile
from brainglobe_atlasapi import BrainGlobeAtlas

from brainglobe_utils.image_io import get_size_image_from_file_paths


def transform_points_from_downsampled_to_atlas_space(
    downsampled_points: np.ndarray,
    atlas: BrainGlobeAtlas,
    deformation_field_paths: List[os.PathLike],
    output_filename: Optional[os.PathLike] = None,
    warn_out_of_bounds: bool = True,
) -> Tuple[np.ndarray, List]:
    """
    Transform points from the "downsampled" space in brainreg (i.e. raw data
    transformed to the same orientation and scale as the atlas) to the atlas
    coordinate space.


    Parameters
    ----------
    downsampled_points :
        Points already downsampled to the atlas resolution.
    atlas :
        Target BrainGlobe atlas.
    deformation_field_paths :
        File paths to the deformation fields as generated by brainreg.
    output_filename :
        File to save transformed points to. Points are saved as a HDF file
        using pandas.
    warn_out_of_bounds :
        Whether to warn if points fall outside of the atlas space.

    Returns
    -------
    transformed_points
        Points transformed to the atlas space.

    points_out_of_bounds
        Points that fall outside of the transformed atlas space.
    """

    field_scales = [int(1000 / resolution) for resolution in atlas.resolution]
    points: List[List] = [[], [], []]
    points_out_of_bounds = []
    for axis, deformation_field_path in enumerate(deformation_field_paths):
        deformation_field = tifffile.imread(deformation_field_path)
        for point in downsampled_points:
            try:
                point = [int(round(p)) for p in point]
                points[axis].append(
                    int(
                        round(
                            field_scales[axis]
                            * deformation_field[point[0], point[1], point[2]]
                        )
                    )
                )
            except IndexError:
                if point not in points_out_of_bounds:
                    points_out_of_bounds.append(point)
                    if warn_out_of_bounds:
                        logging.info(
                            f"Ignoring point: {point} "
                            f"as it falls outside the atlas."
                        )

    transformed_points = np.array(points).T

    if output_filename is not None:
        df = pd.DataFrame(transformed_points)
        df.to_hdf(output_filename, key="df", mode="w")

    return transformed_points, points_out_of_bounds


def get_anatomical_space_from_image_planes(
    image_plane: np.ndarray, orientation: str, voxel_sizes: List[float]
) -> bgs.AnatomicalSpace:
    """
    Determine the brainglobe anatomical space from an image plane.

    Parameters
    ----------
    image_plane : np.ndarray
        A numpy-like array representing a single image plane from a 3D image.
    orientation : str
        The orientation of the image following the brainglobe-space
        three-letter convention (e.g., 'asr', 'psl').
    voxel_sizes : List[float]
        A list of floats representing the voxel sizes (e.g., [5, 2, 2]).

    Returns
    -------
    bgs.AnatomicalSpace
        An object representing the anatomical space.

    """

    shape = tuple(get_size_image_from_file_paths(image_plane).values())

    space = bgs.AnatomicalSpace(
        orientation,
        shape=(shape[2], shape[1], shape[0]),
        resolution=[float(i) for i in voxel_sizes],
    )
    return space


def transform_points_from_raw_to_downsampled_space(
    points: np.ndarray,
    source_image_plane: np.ndarray,
    orientation: str,
    voxel_sizes: List[float],
    downsampled_space: bgs.AnatomicalSpace,
    output_filename: Optional[os.PathLike] = None,
) -> np.ndarray:
    """
    Transform points from raw data space to the "downsampled" space as
    defined in brainreg (same orientation and voxel size as the atlas)

    Parameters
    ----------
    points : np.ndarray
        Points in the original space.
    source_image_plane : np.ndarray
        A numpy-like array representing a single image.
    orientation : str
        The orientation of the image following the brainglobe-space
    three letter convention (e.g. 'asr', 'psl')
    voxel_sizes : List[float]
        A list of floats representing the voxel sizes (e.g. [5, 2, 2])
    downsampled_space : bgs.AnatomicalSpace
        Target anatomical space.
    output_filename : Optional[os.PathLike], optional
        File to save downsampled points to. Points are saved as a HDF file
        using pandas.

    Returns
    -------
    np.ndarray
        Points transformed to the downsampled space.

    """
    source_space = get_anatomical_space_from_image_planes(
        source_image_plane, orientation, voxel_sizes
    )
    points = source_space.map_points_to(downsampled_space, points)
    if output_filename is not None:
        df = pd.DataFrame(points)
        df.to_hdf(output_filename, key="df", mode="w")
    return points


def transform_points_to_atlas_space(
    points: np.ndarray,
    source_image_plane: np.ndarray,
    orientation: str,
    voxel_sizes: List[float],
    downsampled_space: bgs.AnatomicalSpace,
    atlas: BrainGlobeAtlas,
    deformation_field_paths: List[os.PathLike],
    downsampled_points_path: Optional[os.PathLike] = None,
    atlas_points_path: Optional[os.PathLike] = None,
) -> np.ndarray:
    """
    Transform points from raw data to an atlas space.

    The points are first downsampled to the atlas resolution, and then
    transformed to the atlas space.

    Parameters
    ----------
    points : np.ndarray
        Points in the original space.
    source_image_plane : np.ndarray
        A numpy-like array representing a single image.
    orientation : str
        The orientation of the image following the brainglobe-space
    three letter convention (e.g. 'asr', 'psl')
    voxel_sizes : List[float]
        A list of floats representing the voxel sizes (e.g. [5, 2, 2])
    downsampled_space : bgs.AnatomicalSpace
        The BrainGlobe anatomical space representing the downsampled space
    atlas : BrainGlobeAtlas
        The BrainGlobe atlas object used for registration
    deformation_field_paths : List[os.PathLike]
        File paths to the deformation fields as generated by brainreg.
    downsampled_points_path : Optional[os.PathLike], optional
        The file path where the downsampled points
        should be saved (default is None).
    atlas_points_path : Optional[os.PathLike], optional
        The file path where the atlas-transformed points
        should be saved (default is None).

    Returns
    -------
    np.ndarray
        An array of points transformed into the atlas space.

    """
    downsampled_points = transform_points_from_raw_to_downsampled_space(
        points,
        source_image_plane,
        orientation,
        voxel_sizes,
        downsampled_space,
        output_filename=downsampled_points_path,
    )
    points = transform_points_from_downsampled_to_atlas_space(
        downsampled_points,
        atlas,
        deformation_field_paths,
        output_filename=atlas_points_path,
    )

    return points
