# -*- coding: utf-8 -*-
"""
Created on Sun Aug 17 17:44:17 2014

@author: worker
"""

import numpy as np
from numpy.fft import fftshift,ifftshift,fft,ifft 
from numpy.fft import fft2 as fast_fft2
from numpy.fft import ifft2 as fast_ifft2
from numpy.fft import fftn as fast_fftn
from numpy.fft import ifftn as fast_ifftn
#from myutils.myplotlib import imshow
from packinon.pg_circle import calc_points 
from myutils.transformations import matrix_between_vectors 
from   skimage.measure import label #regionprops
import os
from   scipy import misc 
from   myutils.utils import tprint
from   scipy.signal import medfilt

def medfilt_ims(refs,medker,mediters):
    nrefs  = np.prod(refs.shape[:-2])
    mrefs  = np.zeros(refs.shape,np.float32)
    for k in range(nrefs):        
        mref = refs[k].copy() 
        #mref -= np.median(mref)
        #mref[mref < 0.0] = 0.0        
        for m in range(mediters):        
            mref = medfilt(mref,kernel_size=medker)
        mrefs[k] = np.float32(mref)  
    return mrefs      

def rad2res(rad,win,psize):
    return psize*win/rad  

def radfilters(N,rexp):
    nf   = rexp.size
    r    = np.arange(N,dtype=np.float32)+1.0    
    rw   = np.zeros((nf,N),dtype=np.float32)
    for k in range(nf):
        rw[k]  = np.power(r,rexp[k]) 
    #rw /= rw[:,-1][:,None]        
    return rw

def logfilters(win,nscales,minres,**kwargs):
    # filter locations
    width = kwargs.pop('width',0.33)
    N     = win//2
    minr  = int(np.round(win/minres))
    fidxs = np.int32(np.ceil(np.logspace(np.log10(minr),np.log10(N),nscales))) 
    sigs  = width*fidxs 
    x     = np.arange(N,dtype=np.float32)
    #b     = np.exp(bfact*(x**2)/(4.0*N**2))[None,:]  
    rw    = np.float32([np.exp(-((x-fidxs[k])**2)/(2.0*(sigs[k]**2))) for k in range(nscales)])
    # extend low pass
    rw[0,:fidxs[0]] = 1.0
    # extend high-pass
    rw[-1,fidxs[-1]:] = 1.0    
    #rw    = np.float32(rw/rw.sum(axis=0)[None,:]) 
    return np.ascontiguousarray(rw)

def logwidthfilters(win,nscales,**kwargs):
    # filter locations
    width = kwargs.pop('width',0.33)
    N     = win//2
    minr  = 1 #int(np.round(win/minres))
    fidxs = np.int32(np.ceil(np.logspace(np.log10(minr),np.log10(N),nscales+2))) 
    sigs  = width*(N-fidxs[1:-1]) 
    x     = np.arange(N,dtype=np.float32)
    #b     = np.exp(bfact*(x**2)/(4.0*N**2))[None,:]  
    rw    = np.float32([np.exp(-((x-N)**2)/(2.0*(sigs[k]**2))) for k in range(nscales)])
    # extend low pass
    #rw[0,:fidxs[0]] = 1.0
    # extend high-pass
    #rw[-1,fidxs[-1]:] = 1.0    
    #rw    = np.float32(rw/rw.sum(axis=1)[:,None]) 
    #rw    = np.float32(rw/rw[:,0][:,None]) 
    return np.ascontiguousarray(rw)

def pdf2mat(pdf):
    pdf   = pos(pdf)
    pdf   = pdf/pdf.sum()
    #nmats = np.prod(pdf.shape[:-2])
    sz    = pdf.shape[-2:]
    x,y   = cart_coords2D(sz)
    x,y   = x[None,:,:],y[None,:,:]
    xx    = ((x**2)*pdf).sum(axis=(-1,-2))
    yy    = ((y**2)*pdf).sum(axis=(-1,-2))
    xy    = (x*y*pdf).sum(axis=(-1,-2))    
    A     = np.float32([[[lxx,lxy],[lxy,lyy]] for [lxx,lyy,lxy] in zip(xx,yy,xy)])
    d     = [lxx*lyy-lxy*lxy for [lxx,lyy,lxy] in zip(xx,yy,xy)]
    AI    = np.float32([[[lyy/ld,-lxy/ld],[-lxy/ld,lxx/ld]] for [lxx,lyy,lxy,ld] in zip(xx,yy,xy,d)])    
    return A,AI

    
def rectmask(sz,msz):
    ''' return rectangular mask of size sz and inner size msz '''
    msk = np.ones(msz,np.float32)
    #assert(len(msz.shape)==2)
    return uncrop2D(msk,sz)

def myframe2img(frame):
    frame = frame.reformat(format='gray16le')        
    img   = np.frombuffer(frame.planes[0], np.dtype('<u2'))
    return np.reshape(img,(frame.height, img.size/frame.height))[:,:frame.width]

def video2dir(vidname,seqpath):
    import av    
    container = av.open(vidname)
    video = next(s for s in container.streams if s.type == b'video')    
    for packet in container.demux(video):
        for frame in packet.decode():        
            imname = os.path.join(seqpath,'frame-%04d.png' % frame.index)
            frame  = frame.reformat(format='gray16le')        
            img    = np.frombuffer(frame.planes[0], np.dtype('<u2'))
            im     = np.reshape(img,(frame.height, img.size/frame.height))[:,:frame.width]            
            #im  = myframe2img(frame)
            img = misc.toimage(im,high=im.max(), low=im.min(),mode='I')
            tprint('Saving %s' % (imname,))
            img.save(imname)

def props2bw(regions,sz):
    bw = np.zeros(sz, dtype='int32')
    for props in regions:
        coords = props['coords']
        idxs   = np.ravel_multi_index((coords[:,0],coords[:,1]),sz)
        bw.flat[idxs] = 1
    return bw 

def bw2coords(bw,im,min_obj_area=1.0):
    lb  = label(bw,background=0.0)
    sz  = im.shape
    x,y = np.mgrid[0:sz[0],0:sz[1]]  
    # collect detected coordinates
    coords = list()
    for l in range(1,lb.max()+1):
        # locate all pixels that belong to this index
        idxs = np.where(lb.flat == l)[0]
        if len(idxs) < min_obj_area:
            continue
        # find centroid coordinate
        ims  = im.flat[idxs].sum()
        imf  = im.flat[idxs]
        cx   = (x.flat[idxs]*imf).sum()/ims
        cy   = (y.flat[idxs]*imf).sum()/ims
        coords.append([int(cx),int(cy)])
    return coords   

def norm0_1(im):
    im -= im.min()
    return im/im.max()    
    
def norm0_255(im):
    return np.round(norm0_1(im)*255).astype(np.uint8)
    
def alpha_compose_red_green(I1,ir,ig,alpha):
    I1  = norm0_255(I1)
    I2  = I1.copy()
    I3  = I1.copy()
    nr  = ir > 0.1*ir.max()
    ng  = ig > 0.1*ig.max()
    ir  = norm0_255(ir)    
    ig  = norm0_255(ig)    
    # combine red channel
    I1[nr]  = np.uint8((1-alpha)*I1[nr] + alpha*(ir[nr]))   
    # weaken other channels 
    I2[nr] = np.uint8((1-alpha)*I2[nr])   
    I3[nr] = np.uint8((1-alpha)*I3[nr])   
    # combine green channel
    I2[ng]  = np.uint8((1-alpha)*I2[ng] + alpha*(ig[ng]))   
    # weaken other channels 
    I1[ng] = np.uint8((1-alpha)*I1[ng])   
    I3[ng] = np.uint8((1-alpha)*I3[ng])   
    return np.dstack((I1,I2,I3))   

def subim_mask(sz,subsz,xk,yk):
    sx,sy   = ndgrid(np.arange(subsz[0]),np.arange(subsz[1]))
    sx = sx + x
k - subsz[0]//2
    sy = sy + yk - subsz[1]//2
    submaskx = np.logical_and(sx >= 0, sx < sz[0]) 
    submasky = np.logical_and(sy >= 0, sy < sz[1]) 
    submask  = np.logical_and(submaskx,submasky)
    #mask     = np.zeros(sz,dtype='bool')
    #mask.flat[sx[submask]*sz[1]+sy[submask]] = True
    return submask,sx,sy   

def inside(pc,sz,bsize):
    ''' Determine whether box can be picked from the micrograph '''        
    bsize2   = bsize//np.sqrt(2.0)

    res      = pc[:,0] > bsize2
    res      = np.logical_and(res,pc[:,1] > bsize2)
    res      = np.logical_and(res,pc[:,0] < sz-bsize2)
    res      = np.logical_and(res,pc[:,1] < sz-bsize2)
    return res

def mad(x):
    return np.median(np.abs(x-np.median(x)))   
    
def mad_detect(x, thresh):
    mn = np.median(x)
    md = mad(x)    
    return x > mn+thresh*md       
    
def mad_detect_masked(x, mask, thresh):
    mn = np.median(x[mask])
    md = mad(x[mask])    
    return x > mn+thresh*md    

def psize2s1d(sz,psize):
    slen = np.ceil(np.sqrt(3.0)*(sz/2.0))+1
    return np.arange(slen,dtype=np.float32)/(sz*psize)

def pa1Dfrom3D(V,psize):
    ''' 1D power amplitude from a 3D volume '''
    sz      = V.shape
    s1d     = psize2s1d(sz[0],psize)[:(sz[0]//2)]       
    VF      = np.fft.fftshift(np.abs(np.fft.fftn(V)))/np.prod(sz)
    VF      = VF**2
    # radially average
    X,Y,Z   = cart_coords3D(sz)
    r       = np.sqrt(X*X + Y*Y + Z*Z)
    rmax    = sz[0]//2
    pa      = np.zeros(rmax,dtype='float32')
    for k in range(rmax):
        pa[k] = np.sqrt(np.mean(VF[r==k]))    
    return pa,s1d

def scale2lp_res(scale):
    ''' Calculates resolution needed to filter to in order to downsize an image
        2.0 for Nyquist, 1.1 for spare '''
    return 2.0*1.0/float(scale)

def sample_sphere(ang_inc, max_lat, max_long, pole_vec=np.float32([0,0,1])):
    ''' ang_inc - angle increments, max_ang = [0,pi]'''
    points = calc_points(ang_inc, max_lat, max_long)
    points.