; NAME:
;   maskfix
;
; Author:
;   Will Foxall (Bristol University, BSc Physics 2011)
; 
; PURPOSE:
;      Where a mask array of values 1 & 0 is equal to 0, the function replaces
;   the like-indexed value in a data array of the same dimensions with
;   the average of the surrounding array values, including only the values
;   for which the corresponding index mask value is 1.
;      The function is iterative such that it prioritizes for array values for
;   which the surrounding values in the mask array have the most 1s.
;      Designed to average over bad pixels in an image based on a generated
;   mask where a 1 is a good pixel and 0 is a bad pixel.
; 
; INPUTS:
;   dataarray   - Data array to be corrected.
;   maskarray   - Mask array of 1s and 0s.
;
; OUTPUTS:
;   result - The corrected array

function maskfix,dataarray,maskarray
;
; duplicate data and mask so as to make function non-destructive
data=dataarray
mask=maskarray
;
; Obtains the required dimension of the inputs
    sda = size( data )
    sma = size( mask )
    Lx = sda[1]-1
    Ly = sda[2]-1
;
; Check arrays match in dimensions
    if (sda[0] NE 2) OR (sda[4] LE 9) $
    OR (sma[1] ne sda[1]) OR (sma[2] ne sda[2]) then begin $
       message,"inputs must be images (matrices) of matching dimensions",/INFO & $
;       return,data & $
    endif
;
; Check mask is just 0s and 1s
    if (MEAN(mask(where(mask ne 0)) - 1) ne 0) then begin $
       message,"maskarray must be composed of 0 and 1 values only",/INFO & $
;       return,data & $
    endif
;
; Obtain the number of pixels to fix
    nbad = sma[4] - TOTAL(mask)
    if (nbad eq 0) OR (nbad gt 0.75*sma[4]) then begin $
       message,"mask 0% OR >75% 0s (perfect mask OR highly inaccurate result)" & $
;       return,data & $
    endif
;
; Create array for bad pixel information (x,y coords and availability)
    info=intarr(3,nbad)
;
; Store x y coords of bad pixels
    i=0 & j=0 & $
    n=double(0) & $
    for i=0,lx do begin $
    for j=0,ly do begin $
       if (mask(i,j) eq 0) then begin $
          info(0,n)=i & $
          info(1,n)=j & $
          n++ & $
       endif & $
    endfor & $
    endfor
;
; Make an oversized data/mask array to prevent operations off edge of array
odata=fltarr(lx+3,ly+3)
for i=1,lx+1 do begin $
for j=1,ly+1 do begin $
odata(i,j)=data(i-1,j-1) & $
endfor & $
endfor
;
omask=fltarr(lx+3,ly+3)
for i=1,lx+1 do begin $
for j=1,ly+1 do begin $
omask(i,j)=mask(i-1,j-1) & $
endfor & $
endfor
;
corrected=fltarr(nbad)
c=0L
;
;
WHILE (sma[4]-total(mask) gt 0) do begin
 
; Write availability of pixel
    for n=0L,(nbad-1) do begin $
    region=fltarr(3,3) & $
       if (info(2,n) ne -1) then begin $
          x=info(0,n) & $
          y=info(1,n) & $
          region(0,0)=omask(x,y) & $
          region(1,0)=omask(x+1,y) & $
          region(2,0)=omask(x+2,y) & $
          region(0,1)=omask(x,y+1) & $
          region(2,1)=omask(x+2,y+1) & $
          region(0,2)=omask(x,y+2) & $
          region(1,2)=omask(x+1,y+2) & $
          region(2,2)=omask(x+2,y+2) & $
          info(2,n)=TOTAL(region) & $
       endif & 
       n=n+1 & $
    endfor & $
;
; Index order of collumn three in decending order (highest avail. first)

    order = REVERSE(SORT(info(2,*))) & $
    info = info(*,order) & $
;
; Start with maximum availability
    s = 8 & $
    While (max(info(2,*)) lt s AND s gt 0) do begin $
        s = s - 1 & $
    Endwhile & $
;
; Correct pixels for highest and second highest availability
    n = 0L & $
    region=fltarr(3,3) & $
    While (info(2,n) eq s) OR (info(2,n) eq s-1) do begin $
       x=info(0,n) & $
       y=info(1,n) & $
       region(0,0)=odata(x,y)*omask(x,y) & $
       region(1,0)=odata(x+1,y)*omask(x+1,y) & $
       region(2,0)=odata(x+2,y)*omask(x+2,y) & $
       region(0,1)=odata(x,y+1)*omask(x,y+1) & $
       region(2,1)=odata(x+2,y+1)*omask(x+2,y+1) & $
       region(0,2)=odata(x,y+2)*omask(x,y+2) & $
       region(1,2)=odata(x+1,y+2)*omask(x+1,y+2) & $
       region(2,2)=odata(x+2,y+2)*omask(x+2,y+2) & $
       odata(x+1,y+1)=TOTAL(region)/info(2,n) & $
       n=n+1 & $
       corrected(c)=odata(x+1,y+1) & $
       c++ & $
    endwhile & $
;
; Correct mask
    n = 0L & $
    While (info(2,n) eq s) OR (info(2,n) eq s-1) do begin $
       x=info(0,n) & $
       y=info(1,n) & $
       mask(x,y)=1 & $
       omask(x+1,y+1)=1 & $
       info(2,n)=-1 & $
       n++ & $
    endwhile & $
;
; Resort info array
    order = REVERSE(sort(info(2,*))) & $
    info = info(*,order) & $
endwhile

for i=0,lx do begin & $
for j=0,ly do begin & $
data(i,j) = odata(i+1,j+1) & $
endfor & $
endfor

return,data
end