import cv2
import numpy as np
import glob
from dataclasses import dataclass
from typing import List, Tuple

@dataclass
class MeasureBound:
    x_start: int
    x_end: int
    y_top: int
    y_bottom: int

class GridParser:
    def __init__(self, frame: np.ndarray):
        self.frame = frame
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        _, self.binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
        self.h, self.w = self.binary.shape
        
    def find_staff_y_bounds(self) -> Tuple[int, int]:
        h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (100, 1))
        h_lines = cv2.morphologyEx(self.binary, cv2.MORPH_OPEN, h_kernel)
        row_sums = np.sum(h_lines, axis=1) / 255
        
        staff_rows = np.where(row_sums > self.w * 0.4)[0]
        if len(staff_rows) == 0: return 0, 0
            
        y_top = int(staff_rows[0])
        y_bottom = y_top
        
        for y in staff_rows:
            if y - y_bottom > 150: break
            y_bottom = int(y)
            
        return max(0, y_top - 5), min(self.h, y_bottom + 5)
        
    def find_measure_bounds(self) -> List[MeasureBound]:
        y_top, y_bottom = self.find_staff_y_bounds()
        if y_bottom - y_top < 20: return []
        staff_height = y_bottom - y_top
        
        # Isolate all vertical linear structures at least 30px tall (ignores almost all hand features and note heads)
        v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 30))
        roi = self.binary[y_top:y_bottom, :]
        v_lines = cv2.morphologyEx(roi, cv2.MORPH_OPEN, v_kernel, iterations=1)
        
        # Aggregate the vertical structures. Measure bars will have a high column density.
        col_sums = np.sum(v_lines, axis=0) / 255
        
        # We expect a measure bar to cross both staves, totaling maybe 50% of the ROI height
        bar_cols = np.where(col_sums > staff_height * 0.4)[0]
        
        clean_bars = []
        for x in bar_cols:
            if not clean_bars or x - clean_bars[-1] > 20:
                clean_bars.append(int(x))
                
        if not clean_bars or clean_bars[0] > 50:
            clean_bars.insert(0, 0)
        if clean_bars[-1] < self.w - 50:
            clean_bars.append(self.w)
            
        measures = []
        for i in range(len(clean_bars) - 1):
            x1 = clean_bars[i]
            x2 = clean_bars[i+1]
            if x2 - x1 < 40: continue
            measures.append(MeasureBound(x1, x2, y_top, y_bottom))
            
        return measures

if __name__ == "__main__":
    videos = glob.glob('output/*.mp4')
    cap = cv2.VideoCapture(videos[0])
    cap.set(cv2.CAP_PROP_POS_FRAMES, 500)
    ret, frame = cap.read()
    cap.release()
    
    parser = GridParser(frame)
    measures = parser.find_measure_bounds()
    print(f"Measures: {[(m.x_start, m.x_end) for m in measures]}")
