import numpy as np
import cv2

# Segmented Absolute Difference (SAD)
# Compares two frames, highlights differences in segments, and returns rectangles of changed areas

def HighlightDifferences(BaseFrame: np.ndarray, NextFrame: np.ndarray, Columns: int = 20, Rows: int = 12, Threshold: float = 10, Padding: int = 1):
    FrameHeight, FrameWidth = BaseFrame.shape[:2]
    SegmentWidth = FrameWidth // Columns
    SegmentHeight = FrameHeight // Rows
    HighlightedFrame = BaseFrame.copy()
    TotalSegments = 0
    SimilarSegments = 0
    DifferentSegments = 0
    DifferentSegmentMask = np.zeros((Rows, Columns), dtype=bool)
    for Row in range(Rows):
        for Col in range(Columns):
            Y = Row * SegmentHeight
            X = Col * SegmentWidth
            Y2 = FrameHeight if Row == Rows - 1 else Y + SegmentHeight
            X2 = FrameWidth if Col == Columns - 1 else X + SegmentWidth
            TotalSegments += 1
            SegmentBase = BaseFrame[Y:Y2, X:X2]
            SegmentNext = NextFrame[Y:Y2, X:X2]
            GreyBase = cv2.cvtColor(SegmentBase, cv2.COLOR_BGR2GRAY)
            GreyNext = cv2.cvtColor(SegmentNext, cv2.COLOR_BGR2GRAY)
            BlurredBase = cv2.GaussianBlur(GreyBase, (5, 5), 0)
            BlurredNext = cv2.GaussianBlur(GreyNext, (5, 5), 0)
            AbsDiff = cv2.absdiff(BlurredBase, BlurredNext)
            MeanDiff = np.mean(AbsDiff) # type: ignore
            if MeanDiff > Threshold:
                DifferentSegments += 1
                DifferentSegmentMask[Row, Col] = True
            else:
                SimilarSegments += 1
    PaddedMask = DifferentSegmentMask.copy()
    for Row in range(Rows):
        for Col in range(Columns):
            if DifferentSegmentMask[Row, Col]:
                for PR in range(max(0, Row - Padding), min(Rows, Row + Padding + 1)):
                    for PC in range(max(0, Col - Padding), min(Columns, Col + Padding + 1)):
                        PaddedMask[PR, PC] = True
    for Row in range(Rows):
        for Col in range(Columns):
            Y = Row * SegmentHeight
            X = Col * SegmentWidth
            Y2 = FrameHeight if Row == Rows - 1 else Y + SegmentHeight
            X2 = FrameWidth if Col == Columns - 1 else X + SegmentWidth
            SegmentBase = BaseFrame[Y:Y2, X:X2]
            if PaddedMask[Row, Col]:
                HighlightedFrame[Y:Y2, X:X2] = cv2.addWeighted(
                    HighlightedFrame[Y:Y2, X:X2], 0.5,
                    np.full_like(SegmentBase, (0, 0, 255)), 0.2, 0
                )
            else:
                HighlightedFrame[Y:Y2, X:X2] = cv2.addWeighted(
                    HighlightedFrame[Y:Y2, X:X2], 0.5,
                    np.full_like(SegmentBase, (0, 255, 0)), 0.2, 0
                )
    SimilarityPercentage = (SimilarSegments / TotalSegments) * 100
    TileCoords = []
    for Row in range(Rows):
        for Col in range(Columns):
            if PaddedMask[Row, Col]:
                TileCoords.append((Col, Row))
    return HighlightedFrame, DifferentSegments, SimilarityPercentage, TileCoords

def GetRectanglesFromTiles(TileMask: np.ndarray, MinDifferentRatio: float = 0.8):
    Height, Width = TileMask.shape
    Visited = np.zeros_like(TileMask, dtype=bool)
    Rectangles = []
    for Y in range(Height):
        for X in range(Width):
            if TileMask[Y, X] and not Visited[Y, X]:
                W = 1
                H = 1
                Expand = True
                while Expand:
                    Expand = False
                    if X + W < Width:
                        NewCol = TileMask[Y:Y+H, X+W] & ~Visited[Y:Y+H, X+W]
                        if np.any(NewCol):
                            NewRect = TileMask[Y:Y+H, X:X+W+1] & ~Visited[Y:Y+H, X:X+W+1]
                            Total = NewRect.size
                            Diff = np.count_nonzero(NewRect)
                            Ratio = Diff / Total
                            if Ratio >= MinDifferentRatio and not np.any(Visited[Y:Y+H, X+W]):
                                W += 1
                                Expand = True
                    if Y + H < Height:
                        NewRow = TileMask[Y+H, X:X+W] & ~Visited[Y+H, X:X+W]
                        if np.any(NewRow):
                            NewRect = TileMask[Y:Y+H+1, X:X+W] & ~Visited[Y:Y+H+1, X:X+W]
                            Total = NewRect.size
                            Diff = np.count_nonzero(NewRect)
                            Ratio = Diff / Total
                            if Ratio >= MinDifferentRatio and not np.any(Visited[Y+H, X:X+W]):
                                H += 1
                                Expand = True
                Visited[Y:Y+H, X:X+W] = True
                Rectangles.append((X, Y, W, H))
    return Rectangles

def GetDifferenceRectangles(BaseFrame, NextFrame, Columns=20, Rows=12, Threshold=5, Padding=1):
    HighlightedFrame, DifferentSegments, SimilarPercentage, TileCoords = HighlightDifferences(
        BaseFrame, NextFrame, Columns=Columns, Rows=Rows, Threshold=Threshold, Padding=Padding
    )
    TileMask = np.zeros((Rows, Columns), dtype=bool)
    for Col, Row in TileCoords:
        if Row < TileMask.shape[0] and Col < TileMask.shape[1]:
            TileMask[Row, Col] = True
    Rectangles = GetRectanglesFromTiles(TileMask, MinDifferentRatio=0.7)
    return {
        'HighlightedFrame': HighlightedFrame,
        'Rectangles': Rectangles,
        'SimilarPercentage': SimilarPercentage,
        'TileCoords': TileCoords,
        'Columns': Columns,
        'Rows': Rows
    }