"""Background Difference Segmentation — Free Pixel-Level GT
=============================================================
Fixed cameras = static background. Compute per-camera background via
median of 100 images, then get goat masks via |frame - background|.

Output: Data/Detection_dataset/labels/train_seg/ with binary masks.
Runs on CPU.
"""
import sys,os
import numpy as np
from PIL import Image
from tqdm import tqdm
from collections import defaultdict

PROJECT_DIR='/home/user/goat'
os.chdir(PROJECT_DIR);sys.path.insert(0,PROJECT_DIR)

def compute_background(camera, img_dir, n_samples=100):
    """Compute median background for one camera."""
    files=[f for f in os.listdir(img_dir) if camera in f and f.endswith('.jpg')]
    files=sorted(files)[:n_samples]
    if not files: return None

    # Load all images
    imgs=[]
    for f in tqdm(files, desc=f'{camera} bg load'):
        img=np.array(Image.open(os.path.join(img_dir,f)).convert('L'),dtype=np.float32)
        imgs.append(img)

    # Median per pixel
    stack=np.stack(imgs,axis=0)
    bg=np.median(stack,axis=0).astype(np.uint8)
    return bg

def main():
    img_dir='Data/Detection_dataset/images/train'
    cameras=['EastLeft','EastRight','WestLeft','WestRight']

    # Compute backgrounds
    backgrounds={}
    for cam in cameras:
        bg=compute_background(cam, img_dir)
        if bg is not None:
            backgrounds[cam]=bg
            # Save
            out_path=f'Data/background_{cam}.png'
            Image.fromarray(bg).save(out_path)
            print(f'{cam} background saved: {out_path} ({bg.shape})')

    # Generate masks for all training images
    out_dir='Data/Detection_dataset/labels/train_seg'
    os.makedirs(out_dir,exist_ok=True)

    all_files=sorted([f for f in os.listdir(img_dir) if f.endswith('.jpg')])
    total_masks=0

    for f in tqdm(all_files, desc='Generating masks'):
        # Determine camera
        cam=None
        for c in cameras:
            if c in f: cam=c;break
        if cam is None or cam not in backgrounds: continue

        img=np.array(Image.open(os.path.join(img_dir,f)).convert('L'),dtype=np.float32)
        bg=backgrounds[cam].astype(np.float32)

        # |frame - background|
        diff=np.abs(img-bg)

        # Adaptive threshold: mean + 2*std of difference
        thresh=diff.mean()+2.0*diff.std()
        mask=(diff>thresh).astype(np.uint8)*255

        # Clean up: remove small noise
        from scipy import ndimage
        mask=ndimage.binary_fill_holes(mask).astype(np.uint8)*255
        # Remove objects smaller than 100 pixels
        labeled,_=ndimage.label(mask)
        sizes=ndimage.sum(mask,labeled,range(1,labeled.max()+1))
        clean_mask=np.zeros_like(mask)
        for i,sz in enumerate(sizes,1):
            if sz>100:
                clean_mask[labeled==i]=255
        mask=clean_mask

        # Save
        out_name=f.replace('.jpg','.png')
        Image.fromarray(mask).save(os.path.join(out_dir,out_name))

        # Count pixels
        goat_pixels=(mask>0).sum()
        total_masks+=1

    print(f'\nSegmentation masks generated: {total_masks} images')
    print(f'Saved to: {out_dir}')
    return out_dir

if __name__=='__main__':
    main()