import warnings
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Any

import numpy as np
import torch

from idtrackerai import Blob, ListOfBlobs, ListOfFragments, Session
from idtrackerai.base.network import DEVICE, IdentifierBase, get_onthefly_dataloader
from idtrackerai.utils import conf, create_dir, save_trajectories, track

from .assign_them_all import close_trajectories_gaps
from .correct_impossible_jumps import correct_impossible_velocity_jumps


@torch.inference_mode()
def compute_identity_probabilities(
    identifier_model: IdentifierBase,
    list_of_blobs: ListOfBlobs,
    id_images_file_paths: Sequence[Path],
) -> None:
    blobs_to_predict: list[Blob] = []
    blob_images: list[tuple[int, int]] = []
    identities: list[int] = []
    for blob in list_of_blobs.all_blobs:
        ids = list(blob.final_identities)
        if len(ids) != 1 or ids[0] in (None, 0):
            continue
        try:
            blob_images.append((blob.id_image_index, blob.episode))
        except AttributeError:
            continue
        else:
            identities.append(ids[0])
            blobs_to_predict.append(blob)

    probabilities = []
    identifier_model.to(DEVICE)
    identifier_model.eval()
    dataloader = get_onthefly_dataloader(
        blob_images, id_images_file_paths, np.asarray(identities) - 1
    )
    for images, labels in track(dataloader, "Predicting final identity certainties"):
        probs = identifier_model(images.to(DEVICE))
        probabilities.append(probs[torch.arange(len(labels)), labels].numpy(force=True))

    for blob, probability in zip(blobs_to_predict, np.concatenate(probabilities)):
        blob.identity_certainty = probability


def trajectories_API(
    session: Session,
    list_of_blobs: ListOfBlobs,
    list_of_fragments: ListOfFragments,
    identifier_model: IdentifierBase | None,
) -> None:
    session.velocity_threshold = _get_velocity_threshold(list_of_fragments)

    if (
        session.track_wo_identities
        or session.single_animal
        or len(list_of_fragments) < 2
    ):
        session.estimated_accuracy = 1.0
    else:
        with session.new_timer("Impossible jumps correction"):
            correct_impossible_velocity_jumps(session, list_of_fragments)
            session.individual_fragments_stats = list_of_fragments.get_stats()
            session.estimated_accuracy = compute_estimated_accuracy(list_of_fragments)
            list_of_fragments.update_blobs(list_of_blobs.all_blobs)
        with session.new_timer("Crossings solver"):
            close_trajectories_gaps(session, list_of_blobs, list_of_fragments)

    if identifier_model is not None:
        compute_identity_probabilities(
            identifier_model, list_of_blobs, list_of_fragments.id_images_file_paths
        )

    create_dir(session.trajectories_folder, remove_existing=True)

    output_dict = produce_output_dict(list_of_blobs.blobs_in_video, session)
    save_trajectories(
        session.trajectories_folder, output_dict, session.trajectories_formats
    )


def _get_velocity_threshold(list_of_fragments: ListOfFragments) -> float:
    distances = np.concatenate(
        [
            fragment.frame_by_frame_velocity
            for fragment in list_of_fragments.individual_fragments
        ]
    )
    if distances.size == 0:
        return np.nan
    return float(
        2 * np.percentile(distances, conf.VEL_PERCENTILE, overwrite_input=True)
    )


def produce_output_dict(
    blobs_in_video: list[list[Blob]],
    session: Session,
    progress_bar=None,
    abort: Callable = lambda: False,
) -> dict[str, Any]:
    """Outputs the dictionary with keys: trajectories, git_commit, video_path,
    frames_per_second

    Parameters
    ----------
    blobs_in_video : list
        List of all blob objects (see :class:`Blob`) generated by
        considering all the blobs segmented from the video
    session : <Session object>
        See :class:`~video.Video`

    Returns
    -------
    dict
        Output dictionary containing trajectories as values
    """

    trajectories = np.full((session.number_of_frames, session.n_animals, 2), np.nan)
    id_probabilities = np.full((session.number_of_frames, session.n_animals, 1), np.nan)
    areas = np.full((session.number_of_frames, session.n_animals), np.nan)

    for frame_number, blobs_in_frame in enumerate(
        track(blobs_in_video, "Producing trajectories")
    ):
        if abort():
            return {}
        if progress_bar:
            progress_bar.emit(frame_number)
        for blob in blobs_in_frame:
            for identity, centroid in blob.final_ids_and_centroids:
                if identity not in (None, 0):
                    trajectories[blob.frame_number, identity - 1] = centroid
            blob_final_identities = list(blob.final_identities)

            if (
                blob.is_an_individual
                and len(blob_final_identities) == 1
                and blob_final_identities[0] not in (None, 0)
            ):
                identity = blob_final_identities[0]
                areas[blob.frame_number, identity - 1] = blob.area
                id_probabilities[blob.frame_number, identity - 1] = (
                    blob.identity_certainty
                )
    # id_probabilities for missing trajectories
    id_probabilities[np.isnan(trajectories[..., 0])] = 0

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)  # mean of empty slice
        return {
            "trajectories": trajectories,
            "id_probabilities": id_probabilities,
            "version": session.version,
            "height": session.height,
            "width": session.width,
            "video_paths": list(map(str, session.video_paths)),
            "frames_per_second": session.frames_per_second,
            "body_length": session.median_body_length,
            "estimated_accuracy": session.estimated_accuracy,
            "areas": {
                "mean": np.nanmean(areas, axis=0),
                "median": np.nanmedian(areas, axis=0),
                "std": np.nanstd(areas, axis=0),
            },
            "setup_points": session.setup_points,
            "identities_labels": session.identities_labels
            or [str(i + 1) for i in range(session.n_animals)],
            "identities_groups": {
                key: list(value) for key, value in session.identities_groups.items()
            },
            "length_unit": session.length_unit,
            "silhouette_score": session.silhouette_score,
            "fragment_connectivity": session.fragment_connectivity,
            "fraction_identified": np.mean(np.isfinite(trajectories)),
        }


def compute_estimated_accuracy(list_of_fragments: ListOfFragments) -> float:
    weighted_P2 = 0
    number_of_individual_blobs = 0

    for fragment in list_of_fragments.individual_fragments:
        if fragment.assigned_identities[0] not in (0, None):
            assert fragment.P2_vector is not None
            weighted_P2 += (
                fragment.P2_vector[fragment.assigned_identities[0] - 1]
                * fragment.n_images
            )
        number_of_individual_blobs += fragment.n_images
    return float(weighted_P2 / number_of_individual_blobs)
