import json
import numpy as np
from scipy.ndimage import binary_erosion, label
from boundary_radial import radial_loss


def disk(shape, center, radius):
    yy, xx = np.indices(shape)
    return ((xx-center[0])**2 + (yy-center[1])**2 <= radius**2)


def boundary_component_radii(mask, center):
    """Extract connected one-pixel boundaries and summarize each by median radius."""
    boundary = mask & ~binary_erosion(mask, structure=np.ones((3, 3), bool))
    lab, n = label(boundary, structure=np.ones((3, 3), bool))
    yy, xx = np.indices(mask.shape)
    out = []
    for k in range(1, n+1):
        pix = lab == k
        out.append(float(np.median(np.sqrt((xx[pix]-center[0])**2 +
                                           (yy[pix]-center[1])**2))))
    return sorted(out)


def dice(a, b):
    den = int(a.sum() + b.sum())
    return float(2*(a & b).sum()/den) if den else 1.0


def main():
    shape, center = (64, 64), (32, 32)
    target = disk(shape, center, 14)
    shifted = disk(shape, (35, 32), 14)
    # Put the extra component away from the target but retain a comparable area.
    extra = target | disk(shape, (51, 48), 5)
    rows = []
    for name, pred in [('shifted', shifted), ('extra_component', extra)]:
        tr = boundary_component_radii(target, center)
        pr = boundary_component_radii(pred, center)
        rows.append({'case': name, 'dice': dice(pred, target),
                     'target_radii': tr, 'pred_radii': pr,
                     'radial_loss_lambda_10': radial_loss(pr, tr, unmatched=10.0)})
    result = {'center': center, 'rows': rows}
    print(json.dumps(result, indent=2))
    with open('segmentation_results.json', 'w') as f:
        json.dump(result, f, indent=2)


if __name__ == '__main__':
    main()