import logging
import math as m
from typing import Optional, Sequence, Tuple, Union, cast

import numpy as np
import pandas as pd
import sqlalchemy as sqla
import torch as th
import torch.nn.functional as F
import zarr
from numpy.typing import ArrayLike
from skimage.measure import regionprops_table
from sqlalchemy.orm import Session
from tqdm import tqdm
from zarr.storage import StoreLike

from ultrack.config import MainConfig
from ultrack.core.database import NodeDB
from ultrack.utils.array import create_zarr
from ultrack.utils.constants import ULTRACK_DEBUG
from ultrack.utils.cuda import import_module, torch_default_device, xp

logging.basicConfig()
logging.getLogger("sqlachemy.engine").setLevel(logging.INFO)

LOG = logging.getLogger(__name__)


_ALIGN_CORNERS = True


def _interpolate(tensor: th.Tensor, antialias: bool = False, **kwargs) -> th.Tensor:
    """Interpolates tensor.

    Parameters
    ----------
    tensor : th.Tensor
        Input 4 or 5-dim tensor.
    antialias : bool, optional
        When true applies a gaussian filter with 0.5 * downscale factor.
        Ignored if upscaling.

    Returns
    -------
    th.Tensor
        Interpolated tensor.
    """
    mode = "trilinear" if tensor.ndim == 5 else "bilinear"
    if antialias:
        scale_factor = kwargs.get("scale_factor")
        if scale_factor is None:
            raise ValueError(
                "`_interpolate` with `antialias=True` requires `scale_factor` parameter."
            )

        if scale_factor < 1.0:
            ndi = import_module("scipy", "ndimage")
            orig_shape = tensor.shape

            if xp == np:  # xp ideally should be cupy in this case
                LOG.warning(
                    "Cupy not found, using scipy for interpolation. This may slow down ultrack."
                )
                array = xp.asarray(tensor.cpu().squeeze().contiguous().numpy())
            else:
                array = xp.asarray(tensor.squeeze().contiguous())
            blurred = ndi.gaussian_filter(
                array,
                sigma=0.5 / scale_factor,
                output=array,
            )
            tensor = th.as_tensor(blurred, device=tensor.device)
            tensor = tensor.reshape(orig_shape)
            LOG.info(f"Antialiasing with sigma = {0.5 / scale_factor}.")

    return F.interpolate(tensor, **kwargs, mode=mode, align_corners=_ALIGN_CORNERS)


def total_variation_loss(tensor: th.Tensor) -> th.Tensor:
    loss = 0.0
    for i in range(1, tensor.ndim - 1):
        idx = th.arange(tensor.shape[i], device=tensor.device)
        # tv = th.square(
        #     2 * th.index_select(tensor, i, idx[1:-1]) \
        #       - th.index_select(tensor, i, idx[2:]) \
        #       - th.index_select(tensor, i, idx[:-2])
        # )  # second derivative
        tv = th.square(
            th.index_select(tensor, i, idx[:-1]) - th.index_select(tensor, i, idx[1:])
        )  # first derivative

        loss = loss + tv.sum()
    return loss


def identity_grid(shape: tuple[int, ...]) -> th.Tensor:
    """Grid equivalent to a identity vector field (no flow).

    Parameters
    ----------
    shape : tuple[int, ...]
        Grid shape.

    Returns
    -------
    th.Tensor
        Tensor of shape (Z, Y, X, D)
    """
    ndim = len(shape)
    T = th.zeros((1, ndim, ndim + 1))
    T[:, :, :-1] = th.eye(ndim)
    grid_shape = (1, 1) + shape
    return F.affine_grid(T, grid_shape, align_corners=_ALIGN_CORNERS)


def _interpolate_grid(grid: th.Tensor, out_dim: int = -1, **kwargs) -> th.Tensor:
    """Interpolate a grid.

    Parameters
    ----------
    grid : th.Tensor
        Grid tensor of shape (z, y, x, D).
    out_dim : int, optional
        Output dimension, by default -1.

    Returns
    -------
    th.Tensor
        Tensor of shape (Z, Y, X, D) when out_dim=1, D position is subject to `out_dim` value.
    """
    return th.stack(
        [_interpolate(grid[None, ..., d], **kwargs)[0] for d in range(grid.shape[-1])],
        dim=out_dim,
    )


def flow_field(
    source: th.Tensor,
    target: th.Tensor,
    im_factor: int = 4,
    grid_factor: int = 4,
    num_iterations: int = 1000,
    lr: float = 1e-2,
    n_scales: int = 3,
    init_grid: Optional[th.Tensor] = None,
) -> th.Tensor:
    """
    Compute the flow vector field `T` that minimizes the
    mean squared error between `T(source)` and `target`.

    Parameters
    ----------
    source : torch.Tensor
        Source image (C, Z, Y, X).
    target : torch.Tensor
        Target image (C, Z, Y, X).
    im_factor : int, optional
        Image space down scaling factor, by default 4.
    grid_factor : int, optional
        Grid space down scaling factor, by default 4.
        Grid dimensions will be divided by both `im_factor` and `grid_factor`.
    num_iterations : int, optional
        Number of gradient descent iterations, by default 1000.
    lr : float, optional
        Learning rate (gradient descent step), by default 1e-2
    n_scales : int, optional
        Number of scales used for multi-scale optimization, by default 3.

    Returns
    -------
    torch.Tensor
        Vector field array with shape (D, (Z / factor), Y / factor, X / factor)
    """
    assert source.shape == target.shape
    assert n_scales > 0
    ndim = source.ndim - 1

    if init_grid is not None:
        if init_grid.ndim != ndim + 2:
            raise ValueError(
                f"Grid must have ndim dimensions ({ndim + 2}), (B, (Z), Y, X, D). Found {init_grid.ndim}."
            )

        if init_grid.shape[-1] != ndim:
            raise ValueError(
                f"Grid last channel must match input dimensions. Found {init_grid.shape[-1]} expected {ndim}."
            )

    source = source.unsqueeze(0)
    target = target.unsqueeze(0)

    device = source.device
    scales = np.flip(np.power(2, np.arange(n_scales)))
    grid = init_grid
    norm_factor = None
    constant = 1_000  # scale of input affects the results, 1_000 is a good value

    for scale in scales:
        with th.no_grad():
            scaled_im_factor = im_factor * scale
            scaled_source = _interpolate(
                source, scale_factor=1 / scaled_im_factor, antialias=True
            )
            scaled_target = _interpolate(
                target, scale_factor=1 / scaled_im_factor, antialias=True
            )
            if norm_factor is None:
                norm_factor = (
                    2.0
                    * constant
                    / (scaled_source.quantile(0.9999) + scaled_target.quantile(0.9999))
                )

            scaled_source *= norm_factor
            scaled_target *= norm_factor

            grid_shape = tuple(
                m.ceil(s / grid_factor) for s in scaled_source.shape[-ndim:]
            )

            if grid is None:
                grid = identity_grid(grid_shape).to(device)
                grid0 = identity_grid(grid_shape).to(device)
            else:
                grid = _interpolate_grid(grid, size=grid_shape)
                grid0 = identity_grid(grid_shape).to(device)

        grid.requires_grad_(True).retain_grad()

        LOG.info(f"scale: {scale}")
        LOG.info(f"image shape: {scaled_source.shape}")
        LOG.info(f"grid shape: {grid.shape}")

        for i in range(num_iterations):
            if grid.grad is not None:
                grid.grad.zero_()

            large_grid = _interpolate_grid(grid, size=scaled_source.shape[-ndim:])

            im2hat = F.grid_sample(
                scaled_target, large_grid, align_corners=_ALIGN_CORNERS
            )

            loss = F.l1_loss(im2hat, scaled_source)
            loss = loss + total_variation_loss(grid - grid0)
            loss.backward()

            if i % 10 == 0:
                LOG.info(f"iter. {i} MSE: {loss:0.4f}")

            with th.no_grad():
                grid -= lr * grid.grad

    LOG.info(f"image size: {source.shape}")
    LOG.info(f"image factor: {im_factor}")
    LOG.info(f"grid factor: {grid_factor}")

    with th.no_grad():
        grid = cast(th.Tensor, grid)
        grid = grid - grid0
        grid = th.flip(grid, (-1,))  # x, y, z -> z, y, x

        # divided by 2.0 because the range is -1 to 1 (length = 2.0)
        grid /= 2.0
        grid = _interpolate_grid(grid, out_dim=1, size=scaled_source.shape[-ndim:])[0]

        LOG.info(f"vector field shape: {grid.shape}")

    if ULTRACK_DEBUG:
        import napari

        viewer = napari.Viewer()
        viewer.add_image(
            scaled_source.cpu().numpy(),
            name="im1",
            blending="additive",
            colormap="blue",
            visible=False,
        )
        viewer.add_image(
            scaled_target.cpu().numpy(),
            name="im2",
            blending="additive",
            colormap="green",
        )
        viewer.add_image(
            im2hat.detach().cpu().numpy(),
            name="im2hat",
            blending="additive",
            colormap="red",
        )
        viewer.add_image(
            grid.detach().cpu().numpy(),
            name="grid",
            colormap="turbo",
            visible=False,
        )
        napari.run()

    return grid


@th.no_grad()
def apply_field(field: th.Tensor, image: th.Tensor) -> th.Tensor:
    """
    Transform image using vector field.
    Image will be scaled to the field size.

    Parameters
    ----------
    field : th.Tensor
        Vector field (D, z, y, x)
    image : th.Tensor
        Original image used to compute the vector field.

    Returns
    -------
    th.Tensor
        Transformed image (z, y, x)
    """
    assert image.ndim == 4

    field = th.flip(field, (0,))  # z, y, x -> x, y, z
    field = field.movedim(0, -1)[None]

    field = field * 2.0  # mapping range from image shape to -1 to 1
    field = identity_grid(field.shape[1:-1]).to(field.device) - field

    transformed_image = F.grid_sample(image[None], field, align_corners=_ALIGN_CORNERS)

    return transformed_image[0]


@th.no_grad()
def advenct_field(
    field: ArrayLike,
    sources: th.Tensor,
    shape: Optional[tuple[int, ...]] = None,
    invert: bool = True,
) -> th.Tensor:
    """
    Advenct points from sources through the provided field.
    Shape indicates the original shape (space) and sources.
    Useful when field is down scaled from the original space.

    Parameters
    ----------
    field : ArrayLike
        Field array with shape T x D x (Z) x Y x X
    sources : th.Tensor
        Array of sources N x D
    shape : tuple[int, ...]
        When provided scales field accordingly, D-dimensional tuple.
    invert : bool
        When true flow is multiplied by -1, resulting in reversal of the flow.

    Returns
    -------
    th.Tensor
        Trajectories of sources N x T x D
    """
    ndim = field.ndim - 2
    device = sources.device
    orig_shape = th.tensor(shape, device=device)
    field_shape = th.tensor(field.shape[2:], device=device)

    if orig_shape is None:
        scales = th.ones(ndim, device=device)
    else:
        scales = (field_shape - 1) / (orig_shape - 1)

    trajectories = [sources]

    zero = th.zeros(1, device=device)

    # ignore last (first) frame since it will be empty
    iterator = range(1, field.shape[0]) if invert else range(0, field_shape[0] - 1)

    for t in iterator:
        current = th.as_tensor(field[t]).to(device=device, non_blocking=True)

        int_sources = th.round(trajectories[-1] * scales)
        int_sources = th.maximum(int_sources, zero)
        int_sources = th.minimum(int_sources, field_shape - 1).long()
        spatial_idx = tuple(
            t.T[0] for t in th.tensor_split(int_sources, len(orig_shape), dim=1)
        )
        idx = (slice(None), *spatial_idx)

        movement = current[idx].T * orig_shape

        if invert:
            sources = sources - movement
        else:
            sources = sources + movement

        trajectories.append(sources)

    trajectories = th.stack(trajectories, dim=1)

    return trajectories


def advenct_field_from_labels(
    field: ArrayLike,
    label: ArrayLike,
    invert: bool = True,
) -> ArrayLike:
    """
    Advenct points from segmentation labels centroid.

    Parameters
    ----------
    field : ArrayLike
        Field array with shape T x D x (Z) x Y x X
    label : ArrayLike
        Label image.
    invert : bool
        When true flow is multiplied by -1, resulting in reversal of the flow.

    Returns
    -------
    ArrayLike
        Trajectories of sources N x T x D
    """
    if isinstance(label, th.Tensor):
        device = label.device
        label = label.numpy()
    else:
        device = torch_default_device()

    centroids = np.stack(
        [v for v in regionprops_table(label, properties=("centroid",)).values()], axis=1
    )
    centroids = th.as_tensor(centroids, device=device)

    trajectories = advenct_field(field, centroids, shape=label.shape, invert=invert)

    trajectories = th.minimum(trajectories, th.as_tensor(label.shape, device=device))
    trajectories = th.maximum(trajectories, th.zeros(len(label.shape), device=device))

    return trajectories


def advenct_from_quasi_random(
    field: ArrayLike,
    img_shape: Tuple[int, ...],
    n_samples: int,
    invert: bool = True,
    device: Optional[th.device] = None,
) -> ArrayLike:
    """
    Advenct points from quasi random uniform distribution.

    Parameters
    ----------
    field : ArrayLike
        Field array with shape T x D x (Z) x Y x X
    img_shape : Tuple[int, ...]
        Must be D-dimensional.
    n_samples : int
        Number of samples.
    invert : bool
        When true flow is multiplied by -1, resulting in reversal of the flow.
    device : Optional[th.device]
        Torch device, by default uses last GPU if available or mps.

    Returns
    -------
    ArrayLike
        Trajectories of sources N x T x D
    """

    if field.shape[1] != len(img_shape):
        raise ValueError(
            f"Field dimension {field.shape[1]} does not match image shape {len(img_shape)}."
        )

    if device is None:
        device = torch_default_device()

    sources = th.quasirandom.SobolEngine(field.shape[1], seed=42).draw(n_samples)
    sources = sources.to(device)
    sources *= th.as_tensor(img_shape, device=device)

    trajectories = advenct_field(field, sources, img_shape, invert=invert)

    trajectories = th.minimum(trajectories, th.as_tensor(img_shape, device=device))
    trajectories = th.maximum(trajectories, th.zeros(len(img_shape), device=device))

    return trajectories


def trajectories_to_tracks(trajectories: th.Tensor) -> np.ndarray:
    """Converts trajectories to napari tracks format.

    Parameters
    ----------
    trajectories : th.Tensor
        Input N x T x D trajectories.

    Returns
    -------
    np.ndarray
        Napari tracks (N x T) x (2 + D) array.
    """
    trajectories = trajectories.cpu().numpy()
    N, T, D = trajectories.shape

    track_ids = np.repeat(np.arange(N), T)[..., None]
    time_pts = np.tile(np.arange(T), N)[..., None]
    coordinates = trajectories.reshape(-1, D)

    return np.concatenate((track_ids, time_pts, coordinates), axis=1)


def _to_tensor(
    image: ArrayLike,
    channel_axis: Optional[int],
    device: th.device,
) -> th.Tensor:
    """Expands image to 4D tensor with channel axis first and converts to tensor

    Parameters
    ----------
    image : ArrayLike
        Image array.
    channel_axis : Optional[int]
        Channel axis, if None image is expanded to 4D tensor with channel axis first.
    device : th.device
        Torch device.

    Returns
    -------
    th.Tensor
        (C, (Z), Y, X)-tensor.
    """
    image = np.asarray(image)

    if channel_axis is None:
        image = image[np.newaxis, ...]

    elif channel_axis > 0:
        image = np.transpose(
            image,
            (channel_axis, *range(channel_axis), *range(channel_axis + 1, image.ndim)),
        )

    return th.as_tensor(image.astype(np.float32), device=device)


def timelapse_flow(
    images: ArrayLike,
    store_or_path: Optional[StoreLike] = None,
    chunks: Optional[Tuple[int]] = None,
    channel_axis: Optional[int] = None,
    im_factor: int = 4,
    grid_factor: int = 4,
    num_iterations: int = 1000,
    lr: float = 1e-2,
    n_scales: int = 3,
    device: Optional[th.device] = None,
) -> zarr.Array:
    """Compute vector field from timelapse.

    Parameters
    ----------
    images : ArrayLike
        Timelapse images shape as (T, ...).
    store_or_path : Optional[StoreLike], optional
        Zarr storage or output path, if not provided a temporary store is used.
    chunks : Optional[Tuple[int]], optional
        Chunk size, if not provided it chunks time with 1 and the spatial dimensions as big as possible.
    channel_axis : Optional[int], optional
        Channel axis EXCLUDING TIME (first axis), e.g (T, C, Y, X) would have `channel_axis=0`.
        If not provided assumes first axis after time.
    im_factor : int, optional
        Image space down scaling factor, by default 4.
    grid_factor : int, optional
        Grid space down scaling factor, by default 4.
        Grid dimensions will be divided by both `im_factor` and `grid_factor`.
    num_iterations : int, optional
        Number of gradient descent iterations, by default 2000.
    lr : float, optional
        Learning rate (gradient descent step), by default 1e-4
    n_scales : int, optional
        Number of scales used for multi-scale optimization, by default 3.
    device : Optional[th.device], optional
        Torch device, by default uses last GPU if available or mps.

    Returns
    -------
    zarr.Array
        Vector field array with shape (T, D, (Z), Y, X).
    """
    if device is None:
        device = torch_default_device()

    imgs_shape = list(images.shape[1:])
    if channel_axis is not None:
        imgs_shape.pop(channel_axis)

    resized_shape = [max(1, int(s // im_factor)) for s in imgs_shape]
    shape = (images.shape[0], len(imgs_shape), *resized_shape)  # (T, D, (Z), Y, X)
    dtype = np.float16

    output = create_zarr(
        shape,
        dtype=dtype,
        store_or_path=store_or_path,
        chunks=chunks,
        default_store_type=None,
    )

    target = _to_tensor(images[0], channel_axis, device)

    for t in tqdm(range(1, len(images)), "Computing flow"):
        source = _to_tensor(images[t], channel_axis, device)

        output[t] = (
            flow_field(
                source,
                target,
                im_factor=im_factor,
                grid_factor=grid_factor,
                lr=lr,
                n_scales=n_scales,
                num_iterations=num_iterations,
            )
            .cpu()
            .numpy()
        )

        target = source

    return output


def add_flow(
    config: MainConfig,
    vector_field: Union[ArrayLike, Sequence[ArrayLike]],
) -> None:
    """
    Adds vector field (coordinate shift) data into nodes.
    If there are fewer vector fields than dimensions, the last dimensions from (z,y,x) have priority.
    For example, if 2 vector fields are provided for a 3D data, only (y, x) are updated.
    Vector field shape, except `t`, can be different from the original image.
    When this happens, the indexing is done mapping the position and rounding.

    Parameters
    ----------
    data_config : DataConfig
        Data configuration parameters.
    vector_field : Sequence[ArrayLike]
        Vector field arrays. Each array per coordinate or a single (T, D, (Z), Y, X)-array.
    """
    LOG.info("Adding shift (vector field) to nodes")

    if isinstance(vector_field, Sequence):
        assert all(v.shape == vector_field[0].shape for v in vector_field)
        ndim = vector_field[0].ndim
        nvecs = len(vector_field)
        vec_shape = np.asarray(vector_field[0].shape)
        is_sequence = True
    else:
        ndim = vector_field.ndim - 1
        nvecs = vector_field.shape[1]
        vec_shape = np.asarray([vector_field.shape[0], *vector_field.shape[2:]])
        is_sequence = False

    shape = np.asarray(config.data_config.metadata["shape"])
    LOG.info(f"Image shape {shape}, ndim={ndim}")
    LOG.info(f"Vec. shape {shape}")
    LOG.info(f"Num. Vec. {nvecs}")

    if len(shape) != ndim:
        raise ValueError(
            "Original data shape and vector field must have same number of dimensions (ignoring channels)."
            f" Found {len(shape)} and {ndim}."
        )

    if np.any(np.abs(vector_field[0]) > 1):
        raise ValueError(
            "Vector field values must be normalized. "
            f"Found {vector_field[0].min()} and {vector_field[0].max()}."
        )

    columns = ["x_shift", "y_shift", "z_shift"]
    coords_scaling = (vec_shape[1:] - 1) / (shape[1:] - 1)
    coordinate_columns = ["z", "y", "x"][-len(coords_scaling) :]
    vec_index_iterator = list(reversed(range(nvecs)))
    # vec_scaling varies depending on the number of vector field and image dimensions
    vec_scaling = np.asarray(shape[1 + len(coords_scaling) - nvecs :])[::-1]

    LOG.info(f"Coordinate_columns {coordinate_columns}")
    LOG.info(f"Vector field scaling of {vec_scaling}")

    engine = sqla.create_engine(config.data_config.database_path)

    for t in tqdm(range(shape[0])):
        with Session(engine) as session:
            query = session.query(NodeDB.id, NodeDB.z, NodeDB.y, NodeDB.x).where(
                NodeDB.t == t
            )
            df = pd.read_sql_query(query.statement, session.bind, index_col="id")

        if len(df) == 0:
            LOG.warning(f"No node found at time point {t}.")
            continue

        coords = df[coordinate_columns].to_numpy()
        coords = np.round(coords * coords_scaling).astype(int)
        coords = np.minimum(
            np.maximum(0, coords), vec_shape[1:] - 1
        )  # truncating boundary
        coords = tuple(coords.T)

        # default value
        df[columns] = 0.0
        # reversed because z could be missing
        for v, colname in zip(vec_index_iterator, columns):
            if is_sequence:
                df[colname] = np.asarray(vector_field[v][t])[
                    coords
                ]  # asarray due lazy loading formats (e.g. dask)
            else:
                df[colname] = np.asarray(vector_field[t, v])[coords]

        # vector field is between -0.5 to 0.5, so it's scaled to the original image size
        # ours is the torch convection divided by 2.0
        # reference
        # https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html
        df[columns[:nvecs]] *= vec_scaling  # x, y, z

        df["node_id"] = df.index
        with Session(engine) as session:
            statement = (
                sqla.update(NodeDB)
                .where(NodeDB.id == sqla.bindparam("node_id"))
                .values(
                    z_shift=sqla.bindparam("z_shift"),
                    y_shift=sqla.bindparam("y_shift"),
                    x_shift=sqla.bindparam("x_shift"),
                )
            )
            session.connection().execute(
                statement,
                df[["node_id"] + columns].to_dict("records"),
                execution_options={"synchronize_session": False},
            )
            session.commit()
