# This file is part of idtracker.ai a multiple animals tracking system
# described in [1].
# Copyright (C) 2017- Francisco Romero Ferrero, Mattia G. Bergomi,
# Francisco J.H. Heras, Robert Hinz, Gonzalo G. de Polavieja and the
# Champalimaud Foundation.
#
# idtracker.ai is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details. In addition, we require
# derivatives or applications to acknowledge the authors by citing [1].
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
# For more information please send an email (idtrackerai@gmail.com) or
# use the tools available at https://gitlab.com/polavieja_lab/idtrackerai.git.
#
# [1] Romero-Ferrero, F., Bergomi, M.G., Hinz, R.C., Heras, F.J.H.,
# de Polavieja, G.G., Nature Methods, 2019.
# idtracker.ai: tracking all individuals in small or large collectives of
# unmarked animals.
# (F.R.-F. and M.G.B. contributed equally to this work.
# Correspondence should be addressed to G.G.d.P:
# gonzalo.polavieja@neuro.fchampalimaud.org)

import logging
import os

import numpy as np


class GroundTruthBlob:
    """Lighter blob objects.
    Attributes:
        identity (preferring the one assigned by the user, if it is not None)
        centroid
        pixels (pixels is stored to check the groundtruth in crossings)
    """

    attributes = [
        "identity",
        "assigned_identity",
        "used_for_training",
        "accumulation_step",
        "centroid",
        "pixels",
        "frame_number",
        "is_an_individual",
        "is_a_crossing",
        "blob_index",
        "fragment_identifier",
    ]

    def get_attribute(self, blob):
        for attribute in self.attributes:
            if attribute == "identity":
                setattr(self, attribute, blob.final_identity)
            else:
                setattr(self, attribute, getattr(blob, attribute))


class IndividualGroundTruth:
    def __init__(
        self,
        video,
        individual_blobs_in_video,
        start=None,
        end=None,
        validated_identity=None,
    ):
        self.video = video
        self.individual_blobs_in_video = individual_blobs_in_video
        self.start = start
        self.end = end
        self.validated_identity = validated_identity

    def save(self):
        gt_name = "_individual_" + str(self.validated_identity) + "_groundtruth.npy"
        path_to_save_groundtruth = os.path.join(self.video.video_folder, gt_name)
        logging.info("saving ground truth at %s" % path_to_save_groundtruth)
        np.save(path_to_save_groundtruth, self)
        logging.info("done")


def generate_individual_groundtruth(
    video,
    blobs_in_video=None,
    start=None,
    end=None,
    validated_identity=None,
    save_gt=True,
):
    """Generates a list of light blobs_in_video, given a video object corresponding to a
    tracked video
    """
    individual_blobs_in_video_groundtruth = []

    for blobs_in_frame in blobs_in_video:
        for blob in blobs_in_frame:
            if blob.final_identity == validated_identity:
                gt_blob = GroundTruthBlob()
                gt_blob.get_attribute(blob)
                individual_blobs_in_video_groundtruth.append(gt_blob)

    groundtruth = IndividualGroundTruth(
        video=video,
        individual_blobs_in_video=individual_blobs_in_video_groundtruth,
        start=start,
        end=end,
        validated_identity=validated_identity,
    )
    if save_gt:
        groundtruth.save()
    return groundtruth
