import pathlib, urllib.request, torch, sys

ROOT = pathlib.Path(__file__).resolve().parents[3]
MODEL_DIR = ROOT / "models"
DEPTH_REPO = MODEL_DIR / "Depth-Anything-V2"
DEPTH_W = ROOT / "weights" / "depth_anything_v2" / "depth_anything_v2_vits.pth"
SAM_REPO = MODEL_DIR / "segment-anything"
SAM_W    = ROOT / "weights" / "segment-anything" / "sam_vit_b_01ec64.pth"

sys.path.insert(0, str(DEPTH_REPO))
sys.path.insert(0, str(SAM_REPO))
from depth_anything_v2.dpt import DepthAnythingV2
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator

def _download(url: str, dest: pathlib.Path):
    dest.parent.mkdir(parents=True, exist_ok=True)
    if not dest.exists():
        urllib.request.urlretrieve(url, dest)

def load_depth(device):
    _download("https://huggingface.co/Depth-Anything/Depth-Anything-V2-Small/resolve/main/depth_anything_v2_vits.pth", DEPTH_W)
    
    # Model configuration for ViT-S
    model_configs = {
        'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
    }
    
    net = DepthAnythingV2(**model_configs['vits'])
    state_dict = torch.load(DEPTH_W, map_location=device)
    net.load_state_dict(state_dict, strict=True)
    net = net.to(device).eval()
    return net

def load_sam(device):
    _download("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth", SAM_W)
    sam = sam_model_registry["vit_b"](checkpoint=str(SAM_W)).to(device)
    return SamAutomaticMaskGenerator(sam)

def predict_depth(net, img_rgb, device):
    from agent.io import pad_to_multiple_of_14, safe_resize
    import torch, numpy as np
    img = pad_to_multiple_of_14(safe_resize(img_rgb))
    ten = torch.from_numpy(img).permute(2,0,1).unsqueeze(0).float() / 255
    with torch.no_grad():
        d = net(ten.to(device)).squeeze().cpu().numpy()
    return d

def generate_masks(mask_gen, img_rgb, use_goal1_approach=True):
    from agent.io import safe_resize, pad_to_multiple_of_14
    
    if use_goal1_approach:
        # Goal1 approach: Just resize to 1024px, NO padding (this works!)
        img_processed = safe_resize(img_rgb, long_side=1024)
    else:
        # Goal2 approach: Use same preprocessing as depth (may misalign but consistent)
        img_processed = pad_to_multiple_of_14(safe_resize(img_rgb))
    
    return mask_gen.generate(img_processed)

def predict_depth_and_masks(depth_net, mask_gen, img_rgb, device, approach="goal1_sam"):
    """
    Generate depth and masks with different alignment strategies.
    
    Args:
        approach: 
            - "goal1_sam": Use Goal1's SAM approach (1024px, no padding) - better SAM results
            - "aligned": Use identical preprocessing for both - perfect alignment
    """
    from agent.io import pad_to_multiple_of_14, safe_resize
    import torch
    
    if approach == "goal1_sam":
        # Goal1 approach: Different preprocessing for better SAM results
        depth_img = pad_to_multiple_of_14(safe_resize(img_rgb))  # For depth (960px + padding)
        sam_img = safe_resize(img_rgb, long_side=1024)           # For SAM (1024px, no padding)
        
        # Generate depth map
        ten = torch.from_numpy(depth_img).permute(2,0,1).unsqueeze(0).float() / 255
        with torch.no_grad():
            depth = depth_net(ten.to(device)).squeeze().cpu().numpy()
        
        # Generate masks with Goal1 approach
        masks = mask_gen.generate(sam_img)
        
        return depth, masks, depth_img, sam_img
        
    else:  # "aligned"
        # Goal2 approach: Same preprocessing for perfect alignment
        img_processed = pad_to_multiple_of_14(safe_resize(img_rgb))
        
        # Generate depth map
        ten = torch.from_numpy(img_processed).permute(2,0,1).unsqueeze(0).float() / 255
        with torch.no_grad():
            depth = depth_net(ten.to(device)).squeeze().cpu().numpy()
        
        # Generate masks on the same processed image
        masks = mask_gen.generate(img_processed)
        
        return depth, masks, img_processed