from django.db import models
import numpy as np

def find_segments_in_prediction_array(prediction_array: np.array, min_frame_len: int):
    """
    Expects a prediction array of shape (num_frames) and a minimum frame length.
    Returns a list of tuples (start_frame_number, end_frame_number) that represent the segments.
    """
    # Add False to the beginning and end to detect changes at the array boundaries
    padded_prediction = np.pad(prediction_array, (1, 1), 'constant', constant_values=False)
    
    # Find the start points and end points of the segments
    diffs = np.diff(padded_prediction.astype(int))
    segment_starts = np.where(diffs == 1)[0]
    segment_ends = np.where(diffs == -1)[0]
    
    # Filter segments based on min_frame_len
    segments = [(start, end) for start, end in zip(segment_starts, segment_ends) if end - start >= min_frame_len]
    
    return segments

class AbstractLabelVideoSegment(models.Model):
    video = None # Placeholder for the video field, to be defined in derived classes
    prediction_meta = None # Placeholder for the prediction_meta field, to be defined in derived classes
    start_frame_number = models.IntegerField()
    end_frame_number = models.IntegerField()
    source = models.ForeignKey("InformationSource", on_delete=models.CASCADE)
    label = models.ForeignKey("Label", on_delete=models.CASCADE)

    class Meta:
        abstract = True

    def __str__(self):
        return self.video.file.path + " Label - " + self.label.name + " - " + str(self.start_frame_number) + " - " + str(self.end_frame_number)
    
    def get_frames(self):
        return self.video.get_frame_range(self.start_frame_number, self.end_frame_number)

class LegacyLabelVideoSegment(AbstractLabelVideoSegment):
    video = models.ForeignKey("LegacyVideo", on_delete=models.CASCADE)
    prediction_meta = models.ForeignKey("LegacyVideoPredictionMeta", on_delete=models.CASCADE, related_name="video_segments")


